"""
Vbai-2.6AD Model
================
Multimodal Alzheimer's classifier with REAL MRI<->biomarker pairing.

Streams:
  * MRI encoder    : 3D ResNet (CBAM/SE) + ASPP → 512-d
  * Tabular encoder: MLP on (values + missing-mask) → 256-d
  * Fusion         : bidirectional cross-attention + gated combine → 512-d

Heads:
  * mri_logits     : Stage-1 MRI-only prediction
  * tab_logits     : Tabular-only prediction (used as auxiliary)
  * fused_logits   : Final 3-way classification (CN/MCI/AD)
  * progression    : will_progress (sigmoid), time_to_conversion (months),
                     time_distribution (24 bins, 5-month resolution)

Training-time tricks (in dataset/loss, not here):
  * modality dropout
  * per-feature random masking
  * cross-modal contrastive loss
"""
from __future__ import annotations
import torch
import torch.nn as nn
import torch.nn.functional as F

import config as C


# ============================================================
# Attention modules (3D)
# ============================================================
class ChannelAttention3D(nn.Module):
    def __init__(self, ch, r=16):
        super().__init__()
        m = max(ch // r, 8)
        self.mlp = nn.Sequential(nn.Linear(ch, m), nn.ReLU(inplace=True), nn.Linear(m, ch))

    def forward(self, x):
        a = x.mean(dim=[2, 3, 4]); b = x.amax(dim=[2, 3, 4])
        attn = torch.sigmoid(self.mlp(a) + self.mlp(b))
        return x * attn[..., None, None, None]


class SpatialAttention3D(nn.Module):
    def __init__(self, k=7):
        super().__init__()
        self.conv = nn.Conv3d(2, 1, k, padding=k // 2, bias=False)

    def forward(self, x):
        avg = x.mean(dim=1, keepdim=True); mx = x.amax(dim=1, keepdim=True)
        attn = torch.sigmoid(self.conv(torch.cat([avg, mx], dim=1)))
        return x * attn


class CBAM3D(nn.Module):
    def __init__(self, ch, r=16):
        super().__init__()
        self.c = ChannelAttention3D(ch, r); self.s = SpatialAttention3D()
    def forward(self, x): return self.s(self.c(x))


class SEBlock3D(nn.Module):
    def __init__(self, ch, r=16):
        super().__init__()
        m = max(ch // r, 8)
        self.fc = nn.Sequential(nn.Linear(ch, m), nn.ReLU(True), nn.Linear(m, ch), nn.Sigmoid())
    def forward(self, x):
        s = x.mean(dim=[2, 3, 4]); s = self.fc(s)[..., None, None, None]
        return x * s


# ============================================================
# 3D residual building blocks
# ============================================================
class ResBlock3D(nn.Module):
    def __init__(self, in_ch, out_ch, stride=1, use_cbam=True, use_se=True, drop_path=0.0):
        super().__init__()
        self.conv1 = nn.Conv3d(in_ch, out_ch, 3, stride, 1, bias=False)
        self.bn1 = nn.BatchNorm3d(out_ch)
        self.conv2 = nn.Conv3d(out_ch, out_ch, 3, 1, 1, bias=False)
        self.bn2 = nn.BatchNorm3d(out_ch)
        self.act = nn.GELU()
        self.cbam = CBAM3D(out_ch) if use_cbam else nn.Identity()
        self.se = SEBlock3D(out_ch) if use_se else nn.Identity()
        self.drop_path = drop_path
        self.skip = nn.Identity() if (in_ch == out_ch and stride == 1) else nn.Sequential(
            nn.Conv3d(in_ch, out_ch, 1, stride, bias=False), nn.BatchNorm3d(out_ch))

    def _stochastic(self, x):
        if not self.training or self.drop_path == 0.0:
            return x
        keep = 1.0 - self.drop_path
        mask = torch.empty(x.shape[0], 1, 1, 1, 1, device=x.device).bernoulli_(keep)
        return x * mask / keep

    def forward(self, x):
        identity = self.skip(x)
        out = self.act(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out = self.cbam(out); out = self.se(out)
        out = self._stochastic(out)
        return self.act(out + identity)


class ASPP3D(nn.Module):
    def __init__(self, in_ch, out_ch, dilations=(1, 6, 12, 18)):
        super().__init__()
        per = out_ch // len(dilations)
        self.branches = nn.ModuleList([
            nn.Sequential(nn.Conv3d(in_ch, per, 3, padding=d, dilation=d, bias=False),
                          nn.BatchNorm3d(per), nn.GELU())
            for d in dilations
        ])
        self.gp = nn.Sequential(
            nn.AdaptiveAvgPool3d(1),
            nn.Conv3d(in_ch, per, 1, bias=False),
            nn.BatchNorm3d(per), nn.GELU())
        self.fuse = nn.Sequential(nn.Conv3d(per * (len(dilations) + 1), out_ch, 1, bias=False),
                                  nn.BatchNorm3d(out_ch), nn.GELU())

    def forward(self, x):
        feats = [b(x) for b in self.branches]
        g = self.gp(x)
        g = F.interpolate(g, size=x.shape[2:], mode="trilinear", align_corners=False)
        feats.append(g)
        return self.fuse(torch.cat(feats, dim=1))


# ============================================================
# MRI encoder
# ============================================================
class MRIEncoder3D(nn.Module):
    def __init__(self, mcfg: C.ModelConfig):
        super().__init__()
        ch = mcfg.mri_encoder_channels
        self.stem = nn.Sequential(
            nn.Conv3d(1, ch[0], 7, 2, 3, bias=False), nn.BatchNorm3d(ch[0]), nn.GELU(),
            nn.MaxPool3d(3, 2, 1))
        depths = [2, 2, 2, 2]
        dp = [0.0, 0.05, 0.1, 0.15]
        self.stage1 = self._make(ch[0], ch[0], depths[0], 1, mcfg, dp[0])
        self.stage2 = self._make(ch[0], ch[1], depths[1], 2, mcfg, dp[1])
        self.stage3 = self._make(ch[1], ch[2], depths[2], 2, mcfg, dp[2])
        self.stage4 = self._make(ch[2], ch[3], depths[3], 2, mcfg, dp[3])
        self.aspp = ASPP3D(ch[3], mcfg.mri_bottleneck_channels)
        self.pool = nn.AdaptiveAvgPool3d(1)
        self.proj = nn.Sequential(
            nn.Linear(mcfg.mri_bottleneck_channels, mcfg.mri_feature_dim),
            nn.GELU(), nn.Dropout(mcfg.mri_dropout))

    def _make(self, in_ch, out_ch, n, stride, mcfg, dp):
        layers = [ResBlock3D(in_ch, out_ch, stride, mcfg.use_cbam, mcfg.use_se_block, dp)]
        for _ in range(1, n):
            layers.append(ResBlock3D(out_ch, out_ch, 1, mcfg.use_cbam, mcfg.use_se_block, dp))
        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.stem(x)
        x = self.stage1(x); x = self.stage2(x); x = self.stage3(x); x = self.stage4(x)
        x = self.aspp(x); x = self.pool(x).flatten(1)
        return self.proj(x)


# ============================================================
# Tabular encoder
# ============================================================
class TabularEncoder(nn.Module):
    def __init__(self, mcfg: C.ModelConfig):
        super().__init__()
        prev = mcfg.num_tabular_inputs
        layers = []
        for h in mcfg.tabular_hidden_dims:
            layers += [nn.Linear(prev, h), nn.LayerNorm(h), nn.GELU(), nn.Dropout(mcfg.tabular_dropout)]
            prev = h
        layers += [nn.Linear(prev, mcfg.tabular_feature_dim)]
        self.net = nn.Sequential(*layers)

    def forward(self, x):                       # (B, num_tabular_inputs)
        return self.net(x)


# ============================================================
# Cross-modal fusion
# ============================================================
class CrossModalFusion(nn.Module):
    def __init__(self, mri_dim, tab_dim, fdim, heads=8, dropout=0.1):
        super().__init__()
        self.pm = nn.Linear(mri_dim, fdim); self.pt = nn.Linear(tab_dim, fdim)
        self.a_mt = nn.MultiheadAttention(fdim, heads, dropout=dropout, batch_first=True)
        self.a_tm = nn.MultiheadAttention(fdim, heads, dropout=dropout, batch_first=True)
        self.lnm = nn.LayerNorm(fdim); self.lnt = nn.LayerNorm(fdim)
        self.gate = nn.Sequential(nn.Linear(fdim * 2, fdim), nn.Sigmoid())
        self.out = nn.Sequential(nn.Linear(fdim * 2, fdim), nn.GELU(), nn.Dropout(dropout))

    def forward(self, m, t):
        m1 = self.pm(m).unsqueeze(1); t1 = self.pt(t).unsqueeze(1)
        ma, _ = self.a_mt(m1, t1, t1); ta, _ = self.a_tm(t1, m1, m1)
        m2 = self.lnm(m1 + ma).squeeze(1); t2 = self.lnt(t1 + ta).squeeze(1)
        cat = torch.cat([m2, t2], dim=-1)
        g = self.gate(cat); o = self.out(cat)
        return g * m2 + (1 - g) * t2 + o


# ============================================================
# Heads
# ============================================================
class ClsHead(nn.Module):
    def __init__(self, in_dim, num_classes, dropout=0.3):
        super().__init__()
        self.h = nn.Sequential(
            nn.Linear(in_dim, 256), nn.GELU(), nn.Dropout(dropout),
            nn.Linear(256, 128), nn.GELU(), nn.Dropout(dropout),
            nn.Linear(128, num_classes))
    def forward(self, x): return self.h(x)


class ProgressionHead(nn.Module):
    def __init__(self, in_dim, hidden=256, max_months=120, n_bins=24):
        super().__init__()
        self.max_months = float(max_months); self.n_bins = n_bins
        self.shared = nn.Sequential(nn.Linear(in_dim, hidden), nn.GELU(), nn.Dropout(0.3))
        self.binary = nn.Linear(hidden, 1)
        self.time = nn.Sequential(nn.Linear(hidden, 64), nn.GELU(), nn.Linear(64, 1))
        self.dist = nn.Linear(hidden, n_bins)

    def forward(self, x):
        h = self.shared(x)
        logits = self.binary(h).squeeze(-1)
        return {
            "will_progress_logits": logits,                       # raw for BCEWithLogits
            "will_progress": torch.sigmoid(logits),                # for inference convenience
            "time_to_conversion": torch.clamp(F.softplus(self.time(h)).squeeze(-1),
                                              min=0.0, max=self.max_months),
            "time_distribution": F.softmax(self.dist(h), dim=-1),
        }


# ============================================================
# Full model
# ============================================================
class Vbai26ADModel(nn.Module):
    def __init__(self, mcfg: C.ModelConfig | None = None):
        super().__init__()
        self.cfg = mcfg or C.ModelConfig()
        self.mri_encoder = MRIEncoder3D(self.cfg)
        self.tab_encoder = TabularEncoder(self.cfg)

        self.mri_classifier = ClsHead(self.cfg.mri_feature_dim, self.cfg.num_classes, self.cfg.mri_dropout)
        self.tab_classifier = ClsHead(self.cfg.tabular_feature_dim, self.cfg.num_classes, self.cfg.tabular_dropout)

        self.fusion = CrossModalFusion(
            self.cfg.mri_feature_dim, self.cfg.tabular_feature_dim,
            self.cfg.fusion_dim, self.cfg.fusion_num_heads, self.cfg.fusion_dropout)
        self.fused_classifier = ClsHead(self.cfg.fusion_dim, self.cfg.num_classes, self.cfg.fusion_dropout)

        self.progression_head = ProgressionHead(
            self.cfg.fusion_dim, self.cfg.progression_hidden_dim,
            self.cfg.max_progression_months, self.cfg.num_time_bins)

        # Contrastive projection heads (used only at training time)
        self.contrast_mri = nn.Sequential(nn.Linear(self.cfg.mri_feature_dim, 128))
        self.contrast_tab = nn.Sequential(nn.Linear(self.cfg.tabular_feature_dim, 128))

    def forward(self, mri=None, tab=None):
        out = {}
        m_feat = t_feat = None
        if mri is not None:
            m_feat = self.mri_encoder(mri)
            out["mri_features"] = m_feat
            out["mri_logits"] = self.mri_classifier(m_feat)
        if tab is not None:
            t_feat = self.tab_encoder(tab)
            out["tab_features"] = t_feat
            out["tab_logits"] = self.tab_classifier(t_feat)
        if m_feat is not None and t_feat is not None:
            f = self.fusion(m_feat, t_feat)
            out["fused_features"] = f
            out["fused_logits"] = self.fused_classifier(f)
            out["progression"] = self.progression_head(f)
            # Contrastive embeddings
            out["zm"] = F.normalize(self.contrast_mri(m_feat), dim=-1)
            out["zt"] = F.normalize(self.contrast_tab(t_feat), dim=-1)
        elif m_feat is not None:
            out["fused_logits"] = out["mri_logits"]
        elif t_feat is not None:
            out["fused_logits"] = out["tab_logits"]
        return out

    def get_param_groups(self, lr_backbone, lr_fusion):
        backbone = list(self.mri_encoder.parameters()) + list(self.tab_encoder.parameters())
        fusion = (list(self.fusion.parameters()) + list(self.fused_classifier.parameters())
                  + list(self.progression_head.parameters())
                  + list(self.mri_classifier.parameters())
                  + list(self.tab_classifier.parameters())
                  + list(self.contrast_mri.parameters()) + list(self.contrast_tab.parameters()))
        return [{"params": backbone, "lr": lr_backbone},
                {"params": fusion, "lr": lr_fusion}]

    @torch.no_grad()
    def predict(self, mri=None, tab=None):
        self.eval()
        out = self.forward(mri=mri, tab=tab)
        probs = F.softmax(out["fused_logits"], dim=-1)
        pred = probs.argmax(dim=-1)
        result = {"pred_class": pred, "class_probs": probs,
                  "class_names": [C.CLASS_NAMES[c] for c in pred.cpu().tolist()]}
        if "progression" in out:
            p = out["progression"]
            result["will_progress"] = p["will_progress"]
            result["time_to_conversion_months"] = p["time_to_conversion"]
            result["time_distribution"] = p["time_distribution"]
        return result


def count_params(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)