import gradio as gr
import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForMaskedLM

# 설정
MODEL_ID = "solonsophy/kf-deberta-gen"
MAX_LEN = 256
Q_MAX_LEN = 100

# 모델 로드
print("🔄 Loading model...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForMaskedLM.from_pretrained(MODEL_ID)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
model.eval()
print(f"✅ Model loaded on {device}")

MASK_ID = tokenizer.mask_token_id
PAD_ID = tokenizer.pad_token_id
CLS_ID = tokenizer.cls_token_id
SEP_ID = tokenizer.sep_token_id


def heun_generate(question, num_steps, temperature, top_k, max_answer_len, heun_weight):
    """Heun Method를 사용한 Diffusion 기반 답변 생성"""
    if not question.strip():
        return "질문을 입력해주세요."
    
    # 질문 토큰화
    q_tokens = tokenizer.encode(question, add_special_tokens=False)[:Q_MAX_LEN]
    
    # 초기: [CLS] Q [SEP] [MASK]*N
    input_ids = [CLS_ID] + q_tokens + [SEP_ID] + [MASK_ID] * max_answer_len
    input_ids = input_ids[:MAX_LEN]
    
    answer_start = len(q_tokens) + 2
    answer_end = len(input_ids)
    
    input_ids = torch.tensor([input_ids], device=device)
    attention_mask = torch.ones_like(input_ids)
    
    # Iterative denoising
    for step in range(num_steps):
        # 마스크 위치 찾기
        mask_positions = (input_ids[0, answer_start:answer_end] == MASK_ID).nonzero(as_tuple=True)[0]
        mask_positions = mask_positions + answer_start
        
        if len(mask_positions) == 0:
            break
        
        # 이번 스텝에서 unmask할 개수
        remaining_steps = num_steps - step
        tokens_per_step = max(1, len(mask_positions) // remaining_steps)
        
        # ============================================================
        # Predictor Step
        # ============================================================
        with torch.no_grad():
            outputs_pred = model(input_ids=input_ids, attention_mask=attention_mask)
            logits_pred = outputs_pred.logits
        
        mask_logits_pred = logits_pred[0, mask_positions] / temperature
        
        # Top-k filtering
        if top_k > 0:
            top_k_values, _ = torch.topk(mask_logits_pred, min(top_k, mask_logits_pred.size(-1)), dim=-1)
            threshold = top_k_values[:, -1].unsqueeze(-1)
            mask_logits_pred = torch.where(mask_logits_pred < threshold, float('-inf'), mask_logits_pred)
        
        # 샘플링
        probs_pred = F.softmax(mask_logits_pred, dim=-1)
        sampled_tokens_pred = torch.multinomial(probs_pred, num_samples=1).squeeze(-1)
        
        # ============================================================
        # Corrector Step
        # ============================================================
        input_ids_temp = input_ids.clone()
        input_ids_temp[0, mask_positions] = sampled_tokens_pred
        
        with torch.no_grad():
            outputs_corr = model(input_ids=input_ids_temp, attention_mask=attention_mask)
            logits_corr = outputs_corr.logits
        
        mask_logits_corr = logits_corr[0, mask_positions] / temperature
        
        # Top-k filtering
        if top_k > 0:
            top_k_values, _ = torch.topk(mask_logits_corr, min(top_k, mask_logits_corr.size(-1)), dim=-1)
            threshold = top_k_values[:, -1].unsqueeze(-1)
            mask_logits_corr = torch.where(mask_logits_corr < threshold, float('-inf'), mask_logits_corr)
        
        probs_corr = F.softmax(mask_logits_corr, dim=-1)
        
        # ============================================================
        # Heun Combination
        # ============================================================
        combined_probs = (1 - heun_weight) * probs_pred + heun_weight * probs_corr
        
        # 최종 샘플링
        sampled_tokens_final = torch.multinomial(combined_probs, num_samples=1).squeeze(-1)
        confidences_final = combined_probs.gather(1, sampled_tokens_final.unsqueeze(-1)).squeeze(-1)
        
        # Confidence 기반 unmask
        _, top_indices = torch.topk(confidences_final, min(tokens_per_step, len(confidences_final)))
        
        selected_positions = mask_positions[top_indices]
        selected_tokens = sampled_tokens_final[top_indices]
        input_ids[0, selected_positions] = selected_tokens
    
    # 결과 추출
    answer_tokens = input_ids[0, answer_start:answer_end]
    valid_mask = (answer_tokens != MASK_ID) & (answer_tokens != PAD_ID)
    answer_tokens = answer_tokens[valid_mask]
    
    answer = tokenizer.decode(answer_tokens, skip_special_tokens=True)
    return answer.strip() if answer.strip() else "(생성 실패)"


# Gradio UI
with gr.Blocks(title="kf-deberta-gen Demo", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🌀 kf-deberta-gen
    
    **Generative Diffusion BERT** - 한국어 Diffusion 기반 생성 언어 모델
    
    > ⚠️ 이 모델은 실험적 단계입니다.
    """)
    
    with gr.Row():
        with gr.Column(scale=2):
            question_input = gr.Textbox(
                label="질문",
                placeholder="질문을 입력하세요...",
                lines=3
            )
            submit_btn = gr.Button("🚀 생성", variant="primary", size="lg")
            
        with gr.Column(scale=1):
            gr.Markdown("### 생성 설정")
            num_steps = gr.Slider(10, 100, value=50, step=5, label="Steps", 
                                 info="생성 반복 횟수 (높을수록 품질↑, 시간↑)")
            temperature = gr.Slider(0.1, 2.0, value=0.5, step=0.1, label="Temperature",
                                   info="무작위성 (낮을수록 안정적)")
            top_k = gr.Slider(1, 50, value=10, step=1, label="Top-K",
                            info="후보 토큰 수")
            heun_weight = gr.Slider(0.0, 1.0, value=0.5, step=0.1, label="Heun Weight",
                                   info="Corrector 가중치 (0=Predictor만, 1=Corrector만)")
            max_len = gr.Slider(20, 150, value=80, step=10, label="Max Answer Length",
                              info="최대 답변 길이")
    
    output = gr.Textbox(label="답변", lines=5)
    
    gr.Markdown("""
    ### 📝 사용 팁
    - **Steps**: 30-50 추천 (품질과 속도의 균형)
    - **Temperature**: 0.5-0.8 추천 (너무 낮으면 반복적, 너무 높으면 무작위)
    - **Heun Weight**: 0.5 추천 (Predictor와 Corrector 균형)
    """)
    
    gr.Examples(
        examples=[
            ["안녕하세요, 반갑습니다!"],
            ["오늘 날씨 어때?"],
            ["파이썬을 배우려면 어떻게 해야 하나요?"],
            ["인공지능이란 무엇인가요?"],
            ["좋은 하루 되세요!"],
        ],
        inputs=question_input
    )
    
    submit_btn.click(
        fn=heun_generate,
        inputs=[question_input, num_steps, temperature, top_k, max_len, heun_weight],
        outputs=output
    )
    
    question_input.submit(
        fn=heun_generate,
        inputs=[question_input, num_steps, temperature, top_k, max_len, heun_weight],
        outputs=output
    )

if __name__ == "__main__":
    demo.launch()