import os
from huggingface_hub import login

login(token=os.getenv("HF_TOKEN"))

import gradio as gr
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models as models
import torchvision.transforms as transforms
import numpy as np
from PIL import Image
import json
import os
import pydicom
from pydicom.uid import ImplicitVRLittleEndian
from huggingface_hub import snapshot_download

# =========================
# 📁 MODEL CONFIG
# =========================

MODEL_REPO = "InfoBayAI/resnet18-ct-pathology-classifier"

SAVE_DIR = snapshot_download(repo_id=MODEL_REPO)

print("Model downloaded to:", SAVE_DIR) # same folder saved during training

with open(os.path.join(SAVE_DIR, "config.json"), "r") as f:
    config = json.load(f)

with open(os.path.join(SAVE_DIR, "labels.json"), "r") as f:
    label_names_raw = json.load(f)

# labels.json keys are strings ("0","1",...) → convert to int keys
label_names = {int(k): v for k, v in label_names_raw.items()}

IMG_SIZE    = config["img_size"]       # 224
NUM_CLASSES = config["num_classes"]    # 4

# =========================
# ⚙️ DEVICE
# =========================
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# =========================
# 🧠 LOAD MODEL
# — architecture MUST match training exactly —
# =========================
model = models.resnet18(pretrained=False)

# ✅ Same custom FC head used in training
model.fc = nn.Sequential(
    nn.Linear(model.fc.in_features, 256),
    nn.ReLU(),
    nn.Dropout(0.4),
    nn.Linear(256, NUM_CLASSES)
)

model.load_state_dict(
    torch.load(os.path.join(SAVE_DIR, "pytorch_model.bin"), map_location=device)
)
model.to(device)
model.eval()

print("✅ CT Pathology Model Loaded Successfully")

# =========================
# 🧼 PREPROCESS
# — same val_transform used during training —
# =========================
val_transform = transforms.Compose([
    transforms.Resize((IMG_SIZE, IMG_SIZE)),
    transforms.ToTensor(),
])

# =========================
# 🖼️ LOAD IMAGE (DICOM + normal formats)
# — exact same logic as training load_image() —
# =========================
def load_image_from_pil(pil_image):
    """Gradio passes a PIL image directly — just convert to RGB."""
    return pil_image.convert("RGB")


def load_dicom(path):
    """Load a DICOM file from disk path and return a PIL RGB image."""
    try:
        dcm = pydicom.dcmread(path, force=True)

        if not hasattr(dcm, "file_meta") or dcm.file_meta is None:
            dcm.file_meta = pydicom.dataset.FileMetaDataset()

        if not hasattr(dcm.file_meta, "TransferSyntaxUID"):
            dcm.file_meta.TransferSyntaxUID = ImplicitVRLittleEndian

        try:
            dcm.decompress()
        except:
            pass

        img = dcm.pixel_array.astype(np.float32)

        if img.max() == img.min():
            return None

        img = (img - img.min()) / (img.max() - img.min())
        img = (img * 255).astype(np.uint8)

        return Image.fromarray(img).convert("RGB")

    except Exception as e:
        print(f"❌ DICOM load failed | {e}")
        return None

# =========================
# 🔮 PREDICT
# =========================
def predict_ct(image):
    if image is None:
        return "⚠️ Please upload a CT scan image.", None

    img = load_image_from_pil(image)

    tensor = val_transform(img).unsqueeze(0).to(device)

    with torch.no_grad():
        output = model(tensor)
        probs  = F.softmax(output, dim=1)

    probs_np = probs.cpu().numpy()[0]

    sorted_indices = np.argsort(probs_np)[::-1]

    top_label = label_names[sorted_indices[0]]
    top_conf  = probs_np[sorted_indices[0]] * 100

    formatted_probs = "\n".join([
        f"  {label_names[i]:<25} — {probs_np[i]*100:.2f}%"
        for i in sorted_indices
    ])

    result_text = f"""
🫁 CT Scan Pathology Classification

📌 Prediction  : {top_label}
🔍 Confidence  : {top_conf:.2f}%

📊 All Class Probabilities:
{formatted_probs}

─────────────────────────────────────
⚠️  AI-assisted output — NOT a medical diagnosis.
    Always consult a qualified radiologist.
"""

    probs_dict = {
        label_names[i]: float(probs_np[i])
        for i in range(NUM_CLASSES)
    }

    return result_text, probs_dict

# =========================
# 🎨 GRADIO UI
# =========================
def create_interface():
    with gr.Blocks(
        theme=gr.themes.Soft(),
        css="""
        .gradio-container {
            max-width: 1600px !important;
            margin: auto;
        }
        .gr-image {
            min-height: 500px !important;
        }
        textarea {
            font-size: 15px !important;
            font-family: monospace !important;
        }
        h1, h2, h3 {
            text-align: center;
        }
        """
    ) as interface:

        gr.Markdown("# 🫁 CT Scan Pathology Classifier")
        gr.Markdown(
            "### Detects: **Hydropneumothorax · Brain Gliosis · Liver Cirrhosis · Liver Abscess**"
        )
        gr.Markdown("---")

        with gr.Row(equal_height=True):

            # ── LEFT: Image upload ──────────────────────────────
            with gr.Column(scale=1.2):
                image_input = gr.Image(
                    type="pil",
                    label="📤 Upload CT Scan (JPG / PNG / BMP)",
                    height=500,
                    sources=["upload"]   # webcam & clipboard disabled
                )

                with gr.Row():
                    predict_btn = gr.Button("🚀 Analyze CT Scan", variant="primary")
                    clear_btn   = gr.Button("🧹 Clear")

            # ── RIGHT: Results ──────────────────────────────────
            with gr.Column(scale=1.8):
                output_text = gr.Textbox(
                    label="📋 Classification Result",
                    lines=18
                )
                output_chart = gr.Label(
                    label="📊 Confidence Breakdown"
                )

        gr.Markdown("---")
        gr.Markdown(
            "💡 **Tips:** Use axial CT slices for best results. "
            "Convert DICOM to PNG/JPG before uploading if needed."
        )

        # ── Button actions ──────────────────────────────────────
        predict_btn.click(
            fn=predict_ct,
            inputs=image_input,
            outputs=[output_text, output_chart]
        )

        clear_btn.click(
            fn=lambda: (None, "", None),
            inputs=[],
            outputs=[image_input, output_text, output_chart]
        )

    return interface

# =========================
# 🚀 LAUNCH
# =========================
if __name__ == "__main__":
    app = create_interface()
    app.launch(share=True)