"""
Histopathologic Cancer Detection: Gradio Inference Space
Model: ConvNeXt-Tiny + Residual Head (locked-test AUC 0.9979)
"""

import torch
import numpy as np
import torch.nn as nn
import timm
import gradio as gr
from huggingface_hub import hf_hub_download
from PIL import Image
from torchvision import transforms

# ──────────────────────────────────────────────
# 1. Model Architecture (must match training)
# ──────────────────────────────────────────────

class ResidualHead(nn.Module):
    """Classification head with residual connection - MUST match training architecture."""
    def __init__(self, in_features: int, hidden: int = 256, dropout: float = 0.3):
        super().__init__()
        self.fc1 = nn.Linear(in_features, hidden)
        self.bn1 = nn.BatchNorm1d(hidden)
        self.fc2 = nn.Linear(hidden, hidden)
        self.bn2 = nn.BatchNorm1d(hidden)
        self.fc_out = nn.Linear(hidden, 1)
        self.shortcut = nn.Linear(in_features, hidden)
        self.gelu = nn.GELU()
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        identity = self.shortcut(x)

        out = self.fc1(x)
        out = self.bn1(out)
        out = self.gelu(out)
        out = self.dropout(out)

        out = self.fc2(out)
        out = self.bn2(out)
        out = out + identity  # Residual connection
        out = self.gelu(out)
        out = self.dropout(out)

        return self.fc_out(out)


class CancerClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.backbone = timm.create_model("convnext_tiny", pretrained=False, num_classes=0)
        num_features = self.backbone.num_features  # 768
        self.head = ResidualHead(num_features)

    def forward(self, x):
        features = self.backbone(x)
        return self.head(features)


# ──────────────────────────────────────────────
# 2. Load Model
# ──────────────────────────────────────────────

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = CancerClassifier()

# Serve the v2 model evaluated in the report: the one the reported metrics
# describe (SHA256 b08447ab760d6e0f…). The previous checkpoint was an epoch-50
# save from the v1 run, selected on v1 validation AUC; the locked-test figures
# do not apply to it. Same 198 tensor keys, so the architecture is unchanged;
# this file is a bare state_dict, handled by the third branch below.
ckpt_path = hf_hub_download(
    repo_id="hakim78/Cancer_pathos",
    filename="models/final_model_v2_locked.pth",
)

state = torch.load(ckpt_path, map_location=device, weights_only=False)

# Handle different checkpoint formats
if isinstance(state, dict) and "model_state_dict" in state:
    model.load_state_dict(state["model_state_dict"])
elif isinstance(state, dict) and "state_dict" in state:
    model.load_state_dict(state["state_dict"])
else:
    model.load_state_dict(state)

model.to(device)
model.eval()

# ──────────────────────────────────────────────
# 3. Preprocessing (must match training)
# ──────────────────────────────────────────────

transform = transforms.Compose([
    transforms.Resize((96, 96)),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225],
    ),
])


# ──────────────────────────────────────────────
# 4. Prediction function
# ──────────────────────────────────────────────

def _tta_views(t):
    """5 vues geometriques reproduisant get_tta_transforms() d'evaluate.py.

    Implementees en ops tensorielles (le Space n'a pas albumentations).
    Le patch est 96x96 (carre) : rot90/transpose ne changent pas la forme.
    """
    return [
        t,                                  # 1. original
        torch.flip(t, dims=[3]),            # 2. HorizontalFlip (axe largeur)
        torch.flip(t, dims=[2]),            # 3. VerticalFlip (axe hauteur)
        torch.rot90(t, k=1, dims=[2, 3]),   # 4. RandomRotate90 (rotation 90 deg)
        t.transpose(2, 3),                  # 5. Transpose (echange H/W)
    ]


@torch.no_grad()
def predict(image: Image.Image):
    if image is None:
        return {"Cancer": 0.0, "Non-Cancer": 0.0}

    img = image.convert("RGB")
    tensor = transform(img).unsqueeze(0).to(device)

    # TTA x5 : moyenne des probabilites sur les 5 vues (cf. evaluate.py)
    probs = [torch.sigmoid(model(view).squeeze()).item() for view in _tta_views(tensor)]
    prob = float(np.mean(probs))

    return {
        "Cancer": round(prob, 4),
        "Non-Cancer": round(1 - prob, 4),
    }


# ──────────────────────────────────────────────
# 5. Gradio Interface
# ──────────────────────────────────────────────

examples = []  # Add example image paths here if you upload them to the Space

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Histopathology Patch (96x96)"),
    outputs=gr.Label(num_top_classes=2, label="Prediction"),
    title="Histopathologic Cancer Detection",
    description=(
        "Upload a 96×96 histopathologic tissue patch to detect metastatic cancer. "
        "Model: ConvNeXt-Tiny + Residual Head, locked-test AUC 0.9979 "
        "(95% CI 0.9975–0.9982, n = 33,004). "
        "~0.8 s per patch on CPU, 5-fold TTA included."
    ),
    examples=examples or None,
    flagging_mode="never",
)

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