"""
model_classes.py — Exact port of DR_Grading_v10_COMPLETE.ipynb cell 25.

CRITICAL: attribute names MUST match the saved checkpoint keys exactly:
  self.backbone   (NOT self.encoder)
  self.head_cls   (NOT self.head)
  self.head_coral (NOT self.coral_head)
  forward() returns Dict {"logits", "coral", "feat"}  (NOT a tuple)
"""

import torch, torch.nn as nn, torch.nn.functional as F
import timm, numpy as np
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple

NUM_CLASSES   = 5
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD  = [0.229, 0.224, 0.225]

GRADE_MAP = {
    0: "No DR",
    1: "Mild DR",
    2: "Moderate DR",
    3: "Severe DR",
    4: "Proliferative DR (PDR)",
}

GRADE_COLORS = ["#037f0c", "#7d6f00", "#8d4800", "#cc3300", "#d91515"]

GRADE_RISK = [
    "Low",
    "Low-Moderate",
    "Moderate",
    "High",
    "Very High",
]

GRADE_ACTION = [
    "Routine follow-up in 12 months",
    "Follow-up in 6–12 months",
    "Refer to ophthalmologist within 3–6 months",
    "Urgent referral to retina specialist",
    "Immediate ophthalmological intervention required",
]

SEVERITY_ICONS = ["✅", "🟡", "🟠", "🔴", "🚨"]


class GeM(nn.Module):
    """Generalized mean pooling with a learnable exponent."""
    def __init__(self, p: float = 3.0, eps: float = 1e-6):
        super().__init__()
        self.p   = nn.Parameter(torch.tensor(p))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return F.adaptive_avg_pool2d(
            x.clamp(min=self.eps).pow(self.p), 1
        ).pow(1.0 / self.p)

    def __repr__(self) -> str:
        return f"GeM(p={self.p.item():.2f})"


class CoralHead(nn.Module):
    """CORAL ordinal regression head — K-1 binary predictors."""
    def __init__(self, in_features: int, num_classes: int):
        super().__init__()
        self.num_classes = num_classes
        self.linear = nn.Linear(in_features, 1, bias=False)
        self.bias   = nn.Parameter(torch.zeros(num_classes - 1))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        logit = self.linear(x)      # (B, 1)
        return logit + self.bias    # (B, K-1)


class DRHead(nn.Module):
    """Softmax classifier head — small MLP on top of pooled features."""
    def __init__(self, in_features: int, num_classes: int = 5, dropout: float = 0.3):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_features, 512), nn.BatchNorm1d(512),
            nn.SiLU(inplace=True), nn.Dropout(dropout),
            nn.Linear(512, 256),         nn.BatchNorm1d(256),
            nn.SiLU(inplace=True), nn.Dropout(dropout / 2),
            nn.Linear(256, num_classes),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


class DRModel(nn.Module):
    """Backbone + GeM + {softmax head, CORAL head}.
    Attribute names MUST match notebook checkpoint keys exactly.
    """
    def __init__(self, backbone: str = "tf_efficientnetv2_m",
                 num_classes: int = NUM_CLASSES, pretrained: bool = False):
        super().__init__()
        # Key: "backbone" not "encoder" — matches checkpoint saved by notebook
        self.backbone = timm.create_model(
            backbone, pretrained=pretrained, num_classes=0, global_pool=""
        )
        in_feat = self.backbone.num_features
        self.pool       = GeM(p=3.0)
        # Key: "head_cls" not "head"
        self.head_cls   = DRHead(in_feat, num_classes)
        # Key: "head_coral" not "coral_head"
        self.head_coral = CoralHead(in_feat, num_classes)
        self._backbone_name = backbone
        self.num_classes    = num_classes

    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
        feat = self.backbone(x)
        feat = self.pool(feat).flatten(1)
        return {
            "logits": self.head_cls(feat),
            "coral":  self.head_coral(feat),
            "feat":   feat,
        }

    def freeze_backbone(self) -> None:
        for p in self.backbone.parameters():
            p.requires_grad_(False)

    def unfreeze_all(self) -> None:
        for p in self.parameters():
            p.requires_grad_(True)


class FundusValidator(nn.Module):
    def __init__(self, pretrained: bool = False):
        super().__init__()
        self.backbone = timm.create_model(
            "convnext_tiny.fb_in22k", pretrained=pretrained,
            num_classes=0, global_pool="avg"
        )
        self.head = nn.Sequential(
            nn.Linear(self.backbone.num_features, 256),
            nn.GELU(), nn.Dropout(0.3), nn.Linear(256, 1),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.head(self.backbone(x)).squeeze(1)


@dataclass
class CalibrationBundle:
    temperature: float
    isotonic:    list = field(default_factory=list)
    num_classes: int  = 5

    def apply(self, logits: np.ndarray) -> np.ndarray:
        """Temperature scale raw logits → calibrated probabilities.
        logits shape: (1, C) or (C,) — returns (1, C) or (C,).
        """
        squeeze = logits.ndim == 1
        if squeeze:
            logits = logits[np.newaxis]
        scaled = logits / max(float(self.temperature), 1e-6)
        ex     = np.exp(scaled - scaled.max(axis=1, keepdims=True))
        probs  = ex / ex.sum(axis=1, keepdims=True)
        if self.isotonic:
            cal = np.zeros_like(probs)
            for k, iso in enumerate(self.isotonic):
                cal[:, k] = iso.transform(np.clip(probs[:, k], 1e-6, 1 - 1e-6))
            probs = cal / (cal.sum(axis=1, keepdims=True) + 1e-12)
        return probs[0] if squeeze else probs


def coral_rank_from_probs_np(probs: np.ndarray) -> np.ndarray:
    return (probs > 0.5).sum(axis=-1).astype(np.int64)