import gradio as gr
import torch
import numpy as np
from pathlib import Path
import re
from Model import OmniPathWithInterTaskAttention
from transformers import AutoModelForCausalLM, AutoTokenizer
import transformers
import os
from threading import Thread
from transformers import TextIteratorStreamer
from utils_preprocessing import slide2tiles, tiles2features
import openslide
from PIL import Image


# 强制设置 Gradio 为英文环境
os.environ["GRADIO_LOCALE"] = "en"

# 设备设置
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")

def generate_preview(image_path, max_size=800):
    """Generate preview image without full-resolution loading"""
    try:
        ext = os.path.splitext(image_path)[1].lower()
        if ext in [".svs", ".tif", ".tiff"]:
            slide = openslide.OpenSlide(image_path)
            level = slide.level_count - 1  # lowest resolution
            img = slide.read_region((0, 0), level, slide.level_dimensions[level])
            img = img.convert("RGB")
        else:
            img = Image.open(image_path).convert("RGB")

        img.thumbnail((max_size, max_size))
        return img
    except Exception as e:
        print(f"Preview error: {e}")
        return None


def extract_feature_from_image(image_path):
    """
    输入: 癌症图像路径 (.svs / .tiff / .png / .jpg)
    输出: 生成的临时 npy 特征文件路径 (字符串)
    """

    # ---------------- NEW: 统一读取图像为 PIL.Image ---------------- #
    ext = os.path.splitext(image_path)[1].lower()

    if ext in [".svs", ".tif", ".tiff"]:
        slide = openslide.OpenSlide(image_path)
        w, h = slide.dimensions
        img = slide.read_region((0, 0), 0, (w, h)).convert("RGB")
    else:
        img = Image.open(image_path).convert("RGB")
    # --------------------------------------------------------------- #

    # 参数与 3main-ucec.py 保持一致
    mag_assumed = 40
    mag_selected = 20
    tile_size = 512
    mask_downsampling = 16
    edge_mag_thrsh = 15
    edge_fraction_thrsh = 0.5
    batch_size = 32
    model_name = "ctrans"   # 固定 CTransPath 特征提取

    tmp_out_dir = "generated_features"
    os.makedirs(tmp_out_dir, exist_ok=True)

    slide_name = Path(image_path).stem
    slide_file_name = os.path.basename(image_path)

    # ---------------- NEW: 改为传入 img，而不是让 slide2tiles 自己读文件 ---------------- #
    tiles_list = slide2tiles(
        pil_img=img,                     # 🔥 新增参数
        slide_name=slide_name,
        mag_assumed=mag_assumed,
        mag_selected=mag_selected,
        tile_size=tile_size,
        mask_downsampling=mask_downsampling,
        edge_mag_thrsh=edge_mag_thrsh,
        edge_fraction_thrsh=edge_fraction_thrsh,
        save_tile_file=False,
        path2mask=tmp_out_dir + "/",
        path2coordinates=tmp_out_dir + "/",
    )
    # ----------------------------------------------------------------------------- #

    # 特征提取
    npy = tiles2features(tiles_list, model_name, batch_size)

    npy_path = f"{tmp_out_dir}/{slide_name}.npy"
    np.save(npy_path, npy)

    return npy_path

# 预加载模型（避免重复加载）
@torch.no_grad()
def load_models():
    """Preload necessary models"""
    # 1. Load classification model
    ckpt_path = "best_model.pth"
    if not Path(ckpt_path).exists():
        raise FileNotFoundError(f"Model file not found: {ckpt_path}")
    
    ckpt = torch.load(ckpt_path, map_location=device)
    label_mappings = ckpt.get('label_mappings', None)
    if not label_mappings:
        raise ValueError("The checkpoint is missing label_mappings")
    
    ck_cfg = ckpt.get('config', {})
    feature_dim = 768  # Adjust according to your actual feature dimension
    hidden_dim = int(ck_cfg.get('hidden_dim', 256))
    dropout = float(ck_cfg.get('dropout', 0.3))
    use_inter_task_attention = bool(ck_cfg.get('use_inter_task_attention', True))
    inter_task_heads = int(ck_cfg.get('inter_task_heads', 4))

    classification_model = OmniPathWithInterTaskAttention(
        label_mappings=label_mappings,
        feature_dim=feature_dim,
        hidden_dim=hidden_dim,
        dropout=dropout,
        use_inter_task_attention=use_inter_task_attention,
        inter_task_heads=inter_task_heads
    ).to(device)
    classification_model.load_state_dict(ckpt['model_state_dict'], strict=False)
    classification_model.eval()
    
    # 2. Load text generation model
    llm_model_name = "Qwen/Qwen3-0.6B"
    # llm_model_name = "Qwen/QwQ-32B"
    tokenizer = AutoTokenizer.from_pretrained(llm_model_name)
    llm_model = AutoModelForCausalLM.from_pretrained(
        llm_model_name,
        dtype="auto",
        device_map="auto"
    )

    return classification_model, llm_model, tokenizer, label_mappings

# 预加载模型
classification_model, llm_model, tokenizer, label_mappings = load_models()

def analyze_npy_file(npy_file):
    """Analyze NPY file and return prediction results"""
    if npy_file is None:
        return None, "Please upload an NPY file first"
    
    try:
        # Read NPY file
        arr = np.load(npy_file.name, allow_pickle=False)
        if not isinstance(arr, np.ndarray) or arr.ndim != 2:
            return None, "Error: NPY file must be a two-dimensional feature matrix"
        
        features = torch.from_numpy(arr).float()
        
        # Extract short ID
        p = Path(npy_file.name)
        m = re.search(r'(TCGA-[A-Z0-9]{2}-[A-Z0-9]{4})', p.name.upper())
        short_id = m.group(1) if m else p.stem[:12]
        
        # Inference
        feat_batch = features.unsqueeze(0).to(device)
        outputs = classification_model(feat_batch)
        
        # Decode results
        pred_names, pred_scores = {}, {}
        for task_name, logits in outputs.items():
            probs = torch.softmax(logits[0], dim=-1)
            idx = int(torch.argmax(probs).item())
            classes = label_mappings[task_name]['classes']
            class_name = classes[idx] if 0 <= idx < len(classes) else str(idx)
            pred_names[task_name] = class_name
            pred_scores[task_name] = float(probs[idx].item())
        
        # Format results
        results_text = f"Patient ID: {short_id}\n\nPrediction Results:\n"
        for task, name in pred_names.items():
            results_text += f"- {task}: {name} (Confidence: {pred_scores.get(task, 0.0):.3f})\n"
        
        return {"pred_names": pred_names, "pred_scores": pred_scores, "patient_id": short_id}, results_text
        
    except Exception as e:
        return None, f"An error occurred during processing: {str(e)}"

def generate_response(message, chat_history, analysis_results):
    """Generate streamed LLM response"""
    if analysis_results is None:
        yield "Please upload an NPY file first to analyze the patient data.", chat_history
        return

    pred_names = analysis_results["pred_names"]
    pred_scores = analysis_results["pred_scores"]
    patient_id = analysis_results["patient_id"]

    context = f"Patient {patient_id} analysis results:\n"
    for task, name in pred_names.items():
        context += f"- {task}: {name} (confidence: {pred_scores.get(task, 0.0):.3f})\n"

    if "diagnosis" in message.lower() or "result" in message.lower():
        prompt = f"{context}\nBased on the above analysis results, provide a detailed diagnosis summary and interpretation."
    elif "treatment" in message.lower() or "therapy" in message.lower():
        prompt = f"{context}\nBased on the diagnosis, suggest appropriate treatment options and considerations."
    elif "prognosis" in message.lower() or "outlook" in message.lower():
        prompt = f"{context}\nDiscuss the prognosis and potential outcomes for this patient."
    elif "stage" in message.lower():
        prompt = f"{context}\nExplain the staging information and its clinical implications."
    elif "histology" in message.lower() or "type" in message.lower():
        prompt = f"{context}\nDescribe the histological characteristics and their significance."
    else:
        prompt = f"{context}\nUser question: {message}\nPlease provide a helpful response based on the analysis results."

    messages = [{"role": "user", "content": prompt}]
    text = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True, enable_thinking=False
    )

    model_inputs = tokenizer([text], return_tensors="pt").to(llm_model.device)

    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

    thread = Thread(
        target=lambda: llm_model.generate(
            **model_inputs,
            max_new_tokens=1024,  # 🚀 改成较小输出以提升速度
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            streamer=streamer
        )
    )
    thread.start()

    partial = ""
    for new_text in streamer:
        partial += new_text
        # 实时输出
        yield "", chat_history + [(message, partial)]

    # 完成后写回最终内容到历史
    chat_history.append((message, partial))
    yield "", chat_history

def upload_image(image_path, chat_history, analysis_results):
    if image_path is None:
        return chat_history, analysis_results, "❗ No image uploaded."

    if isinstance(image_path, dict):
        image_path = image_path["name"]   # for gr.File type

    try:
        # 1️⃣ 提取 CTransPath 特征 → 生成 npy
        npy_path = extract_feature_from_image(image_path)

        # 2️⃣ 调用现有 npy 分析流程
        new_analysis_results, results_text = analyze_npy_file(Path(npy_path))

        if new_analysis_results is None:
            return chat_history, analysis_results, results_text

        # 3️⃣ 写入聊天历史
        chat_history.append(("System", f"Image analyzed successfully!\n{results_text}"))
        chat_history.append(("System", "You can now ask questions about this patient (diagnosis, treatment, prognosis, etc.)"))

        return chat_history, new_analysis_results, "Image analysis completed successfully!"

    except Exception as e:
        return chat_history, analysis_results, f"❌ Error during image processing: {str(e)}"


def upload_file(npy_file, chat_history, analysis_results):
    """Handle file upload and initial analysis"""
    if npy_file is None:
        return chat_history, analysis_results, "Please select a file to upload"
    
    new_analysis_results, results_text = analyze_npy_file(npy_file)
    
    if new_analysis_results is None:
        return chat_history, analysis_results, results_text
    
    # Add analysis results to chat
    chat_history.append(("System", f"File uploaded and analyzed successfully!\n{results_text}"))
    chat_history.append(("System", "You can now ask questions about this patient's diagnosis, treatment options, prognosis, etc."))
    
    return chat_history, new_analysis_results, "Analysis completed successfully!"

def example_click(example):
    """Handle example question click"""
    return example

# Create conversational interface
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🏥 Medical Pathology Diagnostic Chat Assistant
    
    Upload a pathology NPY file and chat with the AI assistant about the diagnosis, treatment options, prognosis, and more.
    """)
    
    # Store analysis results in session state
    analysis_results = gr.State(value=None)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Upload Patient Data")
    
            img_input = gr.File(
                label="Upload Slide Image (.svs / .tiff / .png / .jpg)",
                type="filepath",
                file_types=[".svs", ".tiff", ".tif", ".png", ".jpg", ".jpeg"]
            )
    
            file_input = gr.File(
                label="Upload NPY Feature File",
                file_types=[".npy"],
                type="filepath"
            )
            upload_btn = gr.Button("Upload & Analyze", variant="primary")
    
            status_output = gr.Textbox(
                label="Status",
                lines=2,
                interactive=False
            )
            
            svs_preview = gr.Image(label="Slide Preview", interactive=False)

        with gr.Column(scale=2):
            gr.Markdown("### Chat with Medical Assistant")
            chatbot = gr.Chatbot(
                label="Conversation",
                height=400
            )
            
            with gr.Row():
                msg = gr.Textbox(
                    label="Your Question",
                    placeholder="Ask about diagnosis, treatment, prognosis...",
                    lines=2,
                    scale=4
                )
                send_btn = gr.Button("Send", variant="primary", scale=1)
            
            with gr.Row():
                clear_btn = gr.Button("Clear Chat")
            
            gr.Markdown("### Suggested Questions")
            examples = gr.Examples(
                examples=[
                    "What is the diagnosis?",
                    "What treatment options are available?",
                    "What is the prognosis?",
                    "Explain the staging information",
                    "Describe the histological findings"
                ],
                inputs=msg,  # 将示例应用到消息输入框
                fn=example_click,  # 点击示例时的处理函数
                outputs=msg,  # 输出到消息输入框
                label="Click a question to use it"
            )

    img_input.change(
        fn=lambda p, c, a: (c, a, "Image uploaded — preview shown below") if p else (c, a, "No image uploaded"),
        inputs=[img_input, chatbot, analysis_results],
        outputs=[chatbot, analysis_results, status_output]
    )
    
    img_input.change(
        fn=lambda p: generate_preview(p) if p else None,
        inputs=img_input,
        outputs=svs_preview
    )

    # Event handlers
    upload_btn.click(
        upload_file,
        inputs=[file_input, chatbot, analysis_results],
        outputs=[chatbot, analysis_results, status_output]
    )
    
    send_btn.click(
        generate_response,
        inputs=[msg, chatbot, analysis_results],
        outputs=[msg, chatbot]
    )
    
    msg.submit(
        generate_response,
        inputs=[msg, chatbot, analysis_results],
        outputs=[msg, chatbot]
    )
    
    clear_btn.click(
        lambda: ([], None, "Chat cleared"),
        inputs=[],
        outputs=[chatbot, analysis_results, status_output]
    )

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