import gradio as gr
import torch
import torch.nn.functional as F
import os
import zipfile
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from peft import PeftModel

# ==========================================
# 1. Model Loading
# ==========================================
# [수정] LoRA 어댑터(adapter_config.json)를 확인해보면 실제 base_model_name_or_path는
#        "monologg/koelectra-small-v3-discriminator" 이며, task_type도 "SEQ_CLS"(분류)입니다.
#        Qwen 생성형 LLM이 아니라 KoELECTRA 기반 "이진 분류(정상/오정보)" 모델입니다.
BASE_MODEL = "monologg/koelectra-small-v3-discriminator"

LORA_WEIGHTS = ""

NUM_LABELS = 2  # safetensors의 classifier.out_proj shape이 (2, 256) -> 클래스 2개
LABEL_MAP = {
    0: "신뢰 가능 (기후 오정보로 판단되지 않음)",
    1: "기후 오정보 가능성 있음 (주의 필요)"
}

print("Loading model and tokenizer...")
try:
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
    base_model = AutoModelForSequenceClassification.from_pretrained(
        BASE_MODEL,
        num_labels=NUM_LABELS,
        torch_dtype=torch.float32,
    )
    model = PeftModel.from_pretrained(base_model, LORA_WEIGHTS)
    model.eval()
    MODEL_LOADED = True
    print("LoRA Model successfully loaded!")
except Exception as e:
    print(f"Warning: Model loading failed ({e}). Running in Simulation Mode.")
    MODEL_LOADED = False


# ==========================================
# 2. Pipeline Algorithm & Inference Engine
# ==========================================
def run_pipeline(user_input):
    """
    1. Text Preprocessing
    2. Tokenizing
    3. LoRA Model Inference (분류)
    4. XAI / Decision Guide Extraction
    5. Attention / Keyword Map Analysis
    6. Output Generation & Limitations
    """
    if not user_input.strip():
        return "입력 텍스트를 입력해주세요.", "N/A", "N/A", "N/A"

    # Step 1: Preprocessing
    cleaned_input = user_input.strip()

    if MODEL_LOADED:
        # Step 2: Tokenizing
        inputs = tokenizer(cleaned_input, return_tensors="pt", truncation=True, max_length=512)

        # Step 3: LoRA Model Inference (+ attention 출력 요청)
        with torch.no_grad():
            outputs = model(**inputs, output_attentions=True)

        logits = outputs.logits
        probs = F.softmax(logits, dim=-1)[0]
        pred_id = int(torch.argmax(probs).item())
        confidence = float(probs[pred_id].item())

        generated_text = (
            f"[분석 결과]\n"
            f"판단: {LABEL_MAP.get(pred_id, f'클래스 {pred_id}')}\n"
            f"신뢰도(확률): {confidence * 100:.1f}%\n\n"
            f"클래스별 확률 -> 정상: {probs[0]*100:.1f}% / 오정보 의심: {probs[1]*100:.1f}%"
        )

        # Step 4: Decision Guide (xAI) - 실제 확률값 기반
        decision_guide = (
            "🔍 **xAI 판단 가이드 추출**\n"
            f"- 모델 예측 클래스: {pred_id} ({LABEL_MAP.get(pred_id, '')})\n"
            f"- 예측 신뢰도: {confidence * 100:.1f}%\n"
            f"- 클래스 확률 분포: 정상 {probs[0]*100:.1f}% / 오정보 의심 {probs[1]*100:.1f}%"
        )

        # Step 5: Attention Map Analysis - 실제 어텐션 가중치 기반 (CLS -> 각 토큰)
        try:
            # 마지막 레이어의 모든 헤드 평균, [CLS] 토큰이 각 토큰에 준 가중치 사용
            last_layer_attn = outputs.attentions[-1][0]          # (num_heads, seq_len, seq_len)
            cls_attn = last_layer_attn.mean(dim=0)[0]             # (seq_len,) CLS -> 각 토큰
            tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])

            scored = [
                (tok, float(w)) for tok, w in zip(tokens, cls_attn.tolist())
                if tok not in tokenizer.all_special_tokens
            ]
            scored.sort(key=lambda x: x[1], reverse=True)
            top_k = scored[:5]

            lines = "\n".join(f"{i+1}. '{tok}' -> 가중치 {w:.3f}" for i, (tok, w) in enumerate(top_k))
            attention_analysis = "📊 **어텐션 맵 핵심 가중치 (Attention Weights)**\n" + (lines or "표시할 토큰이 없습니다.")
        except Exception as e:
            attention_analysis = f"📊 어텐션 분석 중 오류 발생: {e}"

    else:
        # 모델 로딩 자체가 실패했을 때만 사용되는 시뮬레이션 출력
        generated_text = (
            "[시뮬레이션 결과 - 모델 미로딩]\n"
            "현재 서버에 LoRA 가중치 또는 base 모델을 불러오지 못해 예시 출력을 표시합니다.\n"
            "requirements.txt 및 모델 경로 설정을 확인해주세요."
        )
        decision_guide = "🔍 모델이 로딩되지 않아 xAI 가이드를 생성할 수 없습니다."
        attention_analysis = "📊 모델이 로딩되지 않아 어텐션 분석을 수행할 수 없습니다."

    # Step 6: Disclaimer & Limitations
    disclaimer = (
        "⚠️ **Model Card 한계 고지 및 사전 안내 (Disclaimer)**\n"
        "• **책임 선언**: 본 모델(Team 3, 2026)은 기후변화 오정보 탐지 목적으로 개발되었습니다.\n"
        "• **의도된 사용**: 일반적인 기후 정보 탐색 및 교양 판단 참고용.\n"
        "• **사용 금지**: 의학적 판단, 법적 규제 근거, 자동 차단/제재 시스템의 단독 근거로 사용 불가.\n"
        "• **약점 및 한계**: 공신력 있는 기관의 학술 자료가 아닌 소셜 미디어 특유의 신조어, 극단적 비유, 텍스트 형태가 왜곡된 조건에서는 오탐률이 상승할 수 있습니다.\n"
        "• **이의 제기**: 결과에 대한 이의 신청 및 오류 제보는 개발팀(Team 3) 통로를 이용해 주세요."
    )

    return generated_text, decision_guide, attention_analysis, disclaimer


# ==========================================
# 3. Gradio Web Interface
# ==========================================
theme = gr.themes.Soft(
    primary_hue="teal",
    secondary_hue="slate",
)

custom_css = """
.model-card-box {
    background-color: #f8fafc;
    border: 1px solid #e2e8f0;
    border-radius: 8px;
    padding: 15px;
    margin-bottom: 15px;
}
.disclaimer-box {
    background-color: #fffbe2;
    border-left: 4px solid #f59e0b;
    padding: 12px;
    border-radius: 4px;
    margin-top: 10px;
}
"""

with gr.Blocks(theme=theme, css=custom_css, title="기후변화 오정보 탐지 AI 모델") as demo:
    gr.Markdown(
        """
        # 🌍 기후변화 오정보 탐지 AI 모델 (Climate Misinfo Detector)
        **Team 3 | 책임안전 AI 프로젝트 (High School Module 3-9)**  
        LoRA 어댑터 가중치가 적용된 KoELECTRA 기반 분류 AI 알고리즘 웹 인터페이스입니다.
        """
    )

    with gr.Row():
        # Left Column
        with gr.Column(scale=1):
            gr.Markdown("### 📥 사용자 입력 (Input)")
            user_input = gr.Textbox(
                lines=5,
                placeholder="검증하고 싶은 기후변화 관련 문장이나 주장을 입력하세요...\n 예: 지구온난화는 인간 활동 때문이 아니라 자연적인 주기일 뿐이다.",
                label="입력 텍스트 (Text Preprocessing & Tokenizing)"
            )

            submit_btn = gr.Button("🚀 AI 추론 및 분석 실행", variant="primary")

            gr.Markdown("### 📄 Model Card 개요")
            gr.Markdown(
                """
                <div class="model-card-box">
                <b>• 모델명:</b> Team 3 LoRA Climate Misinfo LLM (KoELECTRA-small-v3 기반)<br>
                <b>• 학습 데이터:</b> 기후 데이터베이스 (2024~2026, 200건 검수완료)<br>
                <b>• 성능 지표:</b> Acc 85%, F1 0.82 (TP:100, TN:100, FP:10, FN:20)<br>
                <b>• 윤리 체크:</b> 의학/법률 자동 판단 금지
                </div>
                """
            )

        # Right Column
        with gr.Column(scale=1):
            gr.Markdown("### 📤 추론 및 XAI 분석 결과 (Output)")

            output_text = gr.Textbox(
                label="최종 판별 결과 (Output Generation)",
                lines=6,
                interactive=False
            )

            with gr.Tabs():
                with gr.TabItem("🔍 판단 가이드 (xAI)"):
                    xai_output = gr.Markdown("분석 실행 후 표시됩니다.")
                with gr.TabItem("📊 어텐션 맵 분석"):
                    attention_output = gr.Markdown("분석 실행 후 표시됩니다.")

            gr.Markdown("### ⚠️ 한계 고지 및 책임 선언 (Model Card Section 5 & 6)")
            disclaimer_output = gr.Markdown(
                """
                <div class="disclaimer-box">
                <b>• 한계 사전 고지:</b> 본 모델은 스트레스 테스트 결과 오픈 상황 및 왜곡 문맥에서 오탐 위험이 존재합니다.<br>
                <b>• 이의 제기 통로:</b> 결과에 대한 문의 및 피드백은 공식 모듈 3-9 이의제기 창구를 통해 접수 가능합니다.
                </div>
                """,
                elem_classes=["disclaimer-box"]
            )

    submit_btn.click(
        fn=run_pipeline,
        inputs=[user_input],
        outputs=[output_text, xai_output, attention_output, disclaimer_output]
    )

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