"""DR Grading Model Architecture — extracted from notebook v10."""
import torch
import torch.nn as nn
import torch.nn.functional as F
import timm
import numpy as np
from typing import Dict, Tuple, Optional, List

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 = ["#2ecc71", "#f1c40f", "#e67e22", "#e74c3c", "#8e44ad"]

class GeM(nn.Module):
    def __init__(self, p=3.0, eps=1e-6):
        super().__init__()
        self.p = nn.Parameter(torch.tensor(p))
        self.eps = eps
    def forward(self, x):
        return F.adaptive_avg_pool2d(x.clamp(min=self.eps).pow(self.p), 1).pow(1.0 / self.p)

class CoralHead(nn.Module):
    def __init__(self, in_features, num_classes):
        super().__init__()
        self.linear = nn.Linear(in_features, 1, bias=False)
        self.bias = nn.Parameter(torch.zeros(num_classes - 1))
    def forward(self, x):
        return self.linear(x) + self.bias

class DRHead(nn.Module):
    def __init__(self, in_features, num_classes=5, dropout=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):
        return self.net(x)

class DRModel(nn.Module):
    def __init__(self, backbone="tf_efficientnetv2_m", num_classes=5, pretrained=False):
        super().__init__()
        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)
        self.head_cls = DRHead(in_feat, num_classes)
        self.head_coral = CoralHead(in_feat, num_classes)
        self.num_classes = num_classes
    def forward(self, x):
        feat = self.backbone(x)
        feat = self.pool(feat).flatten(1)
        return {"logits": self.head_cls(feat), "coral": self.head_coral(feat), "feat": feat}

PRETRAINED_BACKBONE = "convnext_tiny.fb_in22k"

class FundusValidator(nn.Module):
    def __init__(self, pretrained=True, backbone_name=PRETRAINED_BACKBONE):
        super().__init__()
        self.backbone = timm.create_model(backbone_name, pretrained=pretrained, num_classes=0, global_pool="avg")
        for p in self.backbone.parameters():
            p.requires_grad = False
        self.backbone.eval()
        with torch.no_grad():
            feat_dim = self.backbone(torch.zeros(1, 3, 224, 224)).shape[1]
        self.feat_dim = feat_dim
        self.head = nn.Sequential(nn.Linear(feat_dim, 256), nn.GELU(), nn.Dropout(0.3), nn.Linear(256, 2))
    def extract_features(self, x):
        with torch.no_grad():
            return self.backbone(x)
    def forward(self, x):
        return self.head(self.extract_features(x))
    def train(self, mode=True):
        super().train(mode)
        self.backbone.eval()
        return self

class SingleLogitModel(nn.Module):
    def __init__(self, inner):
        super().__init__()
        self.inner = inner
        self.backbone = getattr(inner, "backbone", None)
    def forward(self, x):
        return self.inner(x)["logits"]

def safe_load(path, map_location="cpu"):
    try:
        return torch.load(path, map_location=map_location, weights_only=False)
    except Exception:
        return torch.load(path, map_location=map_location)