"""FlexiCT-2D slice encoder with class-query heads for 18-label CT classification.

Model = FlexiCT-2D (LoRA-adapted, patch 16, conv branch on the patch embedding)
        -> slice position encoding -> class-query pool -> class interaction
        -> per-class classifier.

Module and parameter names match the released checkpoints, so
`load_state_dict(strict=True)` works on a model built by `build_from_config`.
"""
from __future__ import annotations

import math
import re

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchio as tio

FLEXICT_ID = "ricklisz123/FlexiCT-2D"
FLEXICT_REVISION = "fcefac1e223c4a1ed72a6387c8e6f2d1f848fa90"


# --------------------------------------------------------------------------- #
# Patch embedding: conv branch and patch-size baking
# --------------------------------------------------------------------------- #
class ConvStem(nn.Module):
    """Small strided CNN that maps an image to the patch grid: (B, C, H, W) -> (B, D, H/p, W/p)."""

    def __init__(self, in_ch: int, embed_dim: int, patch, widths=(16, 32, 64, 128), groups: int = 8):
        super().__init__()
        patch = tuple(int(p) for p in patch)
        n_axis = [int(round(math.log2(p))) for p in patch]
        if any(2 ** n != p for n, p in zip(n_axis, patch)):
            raise ValueError(f"ConvStem needs power-of-two patches, got {patch}")
        layers, c_prev = [], in_ch
        for i in range(max(n_axis)):
            stride = tuple(2 if i < n else 1 for n in n_axis)
            c = widths[min(i, len(widths) - 1)]
            layers += [nn.Conv2d(c_prev, c, 3, stride=stride, padding=1, bias=False),
                       nn.GroupNorm(min(groups, c), c), nn.GELU()]
            c_prev = c
        layers += [nn.Conv2d(c_prev, embed_dim, 1, bias=True)]
        self.net = nn.Sequential(*layers)
        nn.init.trunc_normal_(self.net[-1].weight, std=0.02)
        nn.init.zeros_(self.net[-1].bias)

    def forward(self, x):
        return self.net(x)


class PatchEmbedPlusConv(nn.Module):
    """Runs the original patch projection and the conv branch in parallel.

    The branch output is scaled by a learned per-channel gain that starts near
    zero, so training starts from the pretrained projection.
    """

    def __init__(self, base_conv: nn.Conv2d, gamma_init: float = 1e-4, checkpoint: bool = False):
        super().__init__()
        self.base = base_conv
        self.conv_branch = ConvStem(base_conv.in_channels, base_conv.out_channels,
                                    tuple(base_conv.kernel_size))
        self.cb_gamma = nn.Parameter(torch.full((base_conv.out_channels,), float(gamma_init)))
        self.checkpoint = checkpoint

    # FlexiCT reads these attributes from `proj`
    @property
    def weight(self): return self.base.weight
    @property
    def bias(self): return self.base.bias
    @property
    def kernel_size(self): return self.base.kernel_size
    @property
    def stride(self): return self.base.stride
    @property
    def in_channels(self): return self.base.in_channels
    @property
    def out_channels(self): return self.base.out_channels

    def forward(self, x):
        y = self.base(x)
        if self.checkpoint and self.training:
            import torch.utils.checkpoint as cp
            b = cp.checkpoint(self.conv_branch, x, use_reentrant=False)
        else:
            b = self.conv_branch(x)
        return y + self.cb_gamma.view(1, -1, 1, 1) * b


def _find_patch_embed(model, hint="patch_embed_2D"):
    for name, module in model.named_modules():
        if hint in name and hasattr(module, "proj"):
            return name, module
    raise RuntimeError(f"no patch embedding named '{hint}'")


def add_conv_branch(model, checkpoint: bool = False):
    _, pe = _find_patch_embed(model)
    pe.proj = PatchEmbedPlusConv(pe.proj, checkpoint=checkpoint)
    return pe.proj


def bake_patch_size(model, new_patch: int):
    """Replace the 8x8 patch projection by a real 16x16 conv.

    FlexiCT resamples its kernel on the fly when the patch size differs from
    the pretrained one, and applies it functionally. That path would bypass
    LoRA, so the resampled kernel is computed once (FlexiCT's own resampling)
    and written into a conv module.
    """
    _, pe = _find_patch_embed(model)
    new_k = (new_patch, new_patch)
    if tuple(pe.proj.kernel_size) == new_k:
        return pe
    with torch.no_grad():
        w = pe._resample_conv_weight(pe.proj.weight, new_k)
    conv = nn.Conv2d(pe.proj.in_channels, pe.proj.out_channels, kernel_size=new_k,
                     stride=new_k, bias=pe.proj.bias is not None)
    with torch.no_grad():
        conv.weight.copy_(w)
        if pe.proj.bias is not None:
            conv.bias.copy_(pe.proj.bias)
    pe.proj = conv.to(w.device, w.dtype)
    pe.base_patch_size = new_k
    pe._runtime_patch_size = None
    pe._pinv_cache = {}
    return pe


# --------------------------------------------------------------------------- #
# Preprocessing (CPU)
# --------------------------------------------------------------------------- #
def body_crop_bbox(vol: np.ndarray, thr: float = -500.0, margin: int = 4):
    """In-plane bounding box of voxels above `thr` HU over all slices, plus a margin."""
    m = vol > thr
    xs, ys = m.any(axis=(1, 2)), m.any(axis=(0, 2))
    if not xs.any() or not ys.any():
        return 0, vol.shape[0], 0, vol.shape[1]
    x0, x1 = int(np.argmax(xs)), int(len(xs) - np.argmax(xs[::-1]))
    y0, y1 = int(np.argmax(ys)), int(len(ys) - np.argmax(ys[::-1]))
    return (max(0, x0 - margin), min(vol.shape[0], x1 + margin),
            max(0, y0 - margin), min(vol.shape[1], y1 + margin))


def tissue_start(t: torch.Tensor, n: int, thr: float = -500.0) -> int | None:
    """Start of an n-slice window centred on the slices that contain tissue; None if no crop is needed."""
    s = t.shape[0]
    if s <= n:
        return None
    idx = torch.nonzero((t > thr).any(dim=1).any(dim=1)).flatten()
    c = int((idx[0] + idx[-1]) // 2) if idx.numel() else s // 2
    return max(0, min(c - n // 2, s - n))


def fit_xy(sl: torch.Tensor, size: int, fill: torch.Tensor) -> torch.Tensor:
    """Centre-crop or pad (S, X, Y) to (S, size, size) with a per-slice fill value. No resizing."""
    s, x, y = sl.shape
    out = fill.expand(s, size, size).clone()

    def span(n, k):
        if n >= k:
            a = (n - k) // 2
            return a, a + k, 0, k
        a = (k - n) // 2
        return 0, n, a, a + n

    xs0, xs1, xd0, xd1 = span(x, size)
    ys0, ys1, yd0, yd1 = span(y, size)
    out[:, xd0:xd1, yd0:yd1] = sl[:, xs0:xs1, ys0:ys1]
    return out


def preprocess_ct(volume: np.ndarray, affine: np.ndarray, *, img_size: int, xy_spacing: float,
                  z_spacing: float, max_slices: int, hu_clip=(-1000.0, 1000.0),
                  out_dtype=torch.float16) -> torch.Tensor:
    """NIfTI array in HU -> (max_slices, 1, img_size, img_size)."""
    img = tio.ScalarImage(tensor=torch.from_numpy(np.asarray(volume))[None].float(), affine=affine)
    img = tio.ToCanonical()(img)                                     # RAS
    img = tio.Resample((xy_spacing, xy_spacing, z_spacing))(img)
    vol = img.data[0].numpy()                                        # (X, Y, Z)
    x0, x1, y0, y1 = body_crop_bbox(vol)
    vol = vol[x0:x1, y0:y1, :]
    sl = np.ascontiguousarray(vol.transpose(2, 0, 1)[:, ::-1, ::-1])  # (Z, X, Y), LPS
    t = torch.from_numpy(sl).float().clamp_(*hu_clip)
    st = tissue_start(t, max_slices)
    if st is not None:
        t = t[st:st + max_slices].contiguous()
    mean = t.mean(dim=(1, 2), keepdim=True)                          # per-slice z-score on the unpadded slice
    std = t.std(dim=(1, 2), keepdim=True).clamp(min=1e-6)
    t = (t - mean) / std
    air = (hu_clip[0] - mean) / std
    t = fit_xy(t, img_size, air)
    if t.shape[0] < max_slices:                                      # pad z with air slices
        pad = max_slices - t.shape[0]
        t = F.pad(t, (0, 0, 0, 0, pad // 2, pad - pad // 2), value=float(air.mean()))
    return t.unsqueeze(1).to(out_dtype).contiguous()


# --------------------------------------------------------------------------- #
# Encoder
# --------------------------------------------------------------------------- #
class FlexiCT2DEncoder(nn.Module):
    """Encodes every slice with FlexiCT-2D and returns all tokens: (B, S, 1, H, W) -> (B, S*T, D)."""

    modality = "2d_slices"

    def __init__(self, img_size: int = 480, patch_size: int = 16, max_slices: int = 240,
                 slice_microbatch: int = 240, conv_branch: bool = True,
                 grad_checkpointing: bool = False, model_id: str = FLEXICT_ID,
                 revision: str = FLEXICT_REVISION, load_pretrained: bool = True):
        super().__init__()
        from transformers import AutoConfig, AutoModel

        cfg = AutoConfig.from_pretrained(model_id, revision=revision, trust_remote_code=True)
        cfg.image_size = [img_size, img_size]
        self.model = AutoModel.from_config(cfg, trust_remote_code=True)
        if load_pretrained:
            # from_pretrained fails on FlexiCT's tied-weight keys, so load the file directly
            from huggingface_hub import hf_hub_download
            from safetensors.torch import load_file
            self.model.load_state_dict(
                load_file(hf_hub_download(model_id, "model.safetensors", revision=revision)), strict=False)
        bake_patch_size(self.model, patch_size)

        if grad_checkpointing:
            import torch.utils.checkpoint as cp
            depth = int(getattr(cfg, "depth", 16))
            blocks = next(m for m in self.model.modules()
                          if isinstance(m, nn.ModuleList) and len(m) == depth)
            for blk in blocks:
                blk.forward = (lambda f: lambda *a, **k: cp.checkpoint(f, *a, use_reentrant=False, **k))(blk.forward)
        if conv_branch:
            add_conv_branch(self.model, checkpoint=grad_checkpointing)

        self.size = int(img_size)
        self.max_slices = int(max_slices)
        self.slice_microbatch = int(slice_microbatch)
        self.embed_dim = int(getattr(cfg, "embed_dim", 864))
        grid = self.size // patch_size
        self.tokens_per_slice = 1 + grid * grid
        # LoRA targets: attention qkv/proj, the 2D patch projection and the conv branch (names contain "proj");
        # not the MLPs, the unused 3D patch embedding or the text tower
        self.lora_target_modules = [
            n for n, m in self.model.named_modules()
            if isinstance(m, (nn.Linear, nn.Conv2d, nn.Conv3d))
            and "patch_embed_3D" not in n and "text" not in n.lower() and "mlp" not in n
            and any(k in n for k in ("qkv", "proj", "q_proj", "k_proj", "v_proj"))
        ]

    def _encode_slices(self, flat):
        outs = []
        for i in range(0, flat.shape[0], self.slice_microbatch):
            o = self.model(pixel_values=flat[i:i + self.slice_microbatch])
            cls = o.cls_token
            if cls.dim() == 2:
                cls = cls.unsqueeze(1)
            outs.append(torch.cat([cls, o.patch_tokens], dim=1))     # (mb, T, D)
        return torch.cat(outs, dim=0)

    def forward(self, batch):
        b, s = batch.shape[:2]
        f = self._encode_slices(batch.flatten(0, 1))
        return f.reshape(b, s * f.shape[1], -1), None


class SlicePosEnc(nn.Module):
    """Factorised position code for (B, S*T, D) tokens: slice (z) table + in-plane x/y tables + CLS vector."""

    def __init__(self, dim: int, n_slices: int, tokens_per_slice: int, scale: float = 0.25):
        super().__init__()
        self.S, self.T, self.dim = int(n_slices), int(tokens_per_slice), int(dim)
        self.g = int(round((self.T - 1) ** 0.5))
        std = scale * dim ** -0.5
        self.z = nn.Parameter(torch.randn(self.S, dim) * std)
        self.x = nn.Parameter(torch.randn(self.g, dim) * std)
        self.y = nn.Parameter(torch.randn(self.g, dim) * std)
        self.cls = nn.Parameter(torch.randn(1, dim) * std)

    def forward(self, tokens):
        b, n, d = tokens.shape
        xy = (self.x.view(self.g, 1, d) + self.y.view(1, self.g, d)).reshape(self.g * self.g, d)
        per_slice = torch.cat([self.cls, xy], 0)
        pe = self.z.unsqueeze(1) + per_slice.unsqueeze(0)
        return tokens + pe.reshape(1, self.S * self.T, d).to(tokens.dtype)


# --------------------------------------------------------------------------- #
# Heads
# --------------------------------------------------------------------------- #
class MultiHeadQueryAttentionPool(nn.Module):
    """One learned query per class attends over all tokens, with several heads."""

    def __init__(self, dim: int, num_queries: int = 18, num_heads: int = 4,
                 d_attn: int | None = None, identity_out: bool = True):
        super().__init__()
        d_attn = d_attn or dim
        self.num_queries, self.num_heads, self.head_dim = num_queries, num_heads, d_attn // num_heads
        self.query = nn.Parameter(torch.empty(1, num_heads, num_queries, self.head_dim))
        self.logit_scale = nn.Parameter(torch.tensor(math.log(4.0)))
        nn.init.trunc_normal_(self.query, std=0.1)
        self.k = nn.Linear(dim, d_attn, bias=False)
        self.v = nn.Linear(dim, dim, bias=False)
        self.proj_out = nn.Linear(dim, dim) if num_heads > 1 else nn.Identity()
        if identity_out and isinstance(self.proj_out, nn.Linear):   # start as a plain attention pool
            nn.init.eye_(self.proj_out.weight)
            nn.init.zeros_(self.proj_out.bias)

    def forward(self, x, mask=None):
        b, n, d = x.shape
        h, q_n = self.num_heads, self.num_queries
        k = self.k(x).view(b, n, h, self.head_dim).transpose(1, 2)
        v = self.v(x).view(b, n, h, d // h).transpose(1, 2)
        gain = self.logit_scale.clamp(max=math.log(100.0)).exp()
        q = self.query.expand(b, -1, -1, -1) * ((self.head_dim ** -0.5) * gain)
        out = F.scaled_dot_product_attention(q, k, v, scale=1.0)
        return self.proj_out(out.transpose(1, 2).contiguous().view(b, q_n, d))


class GatedMaxTokenPool(nn.Module):
    """Attention pool plus the single highest-scoring token per class, scaled by tanh(gate).

    The gate starts at zero, so the head starts as a plain attention pool.
    """

    def __init__(self, dim: int, num_queries: int = 18, d_attn: int | None = None):
        super().__init__()
        d_attn = d_attn or dim
        self.query = nn.Parameter(torch.empty(1, num_queries, d_attn))
        self.logit_scale = nn.Parameter(torch.tensor(math.log(4.0)))
        nn.init.trunc_normal_(self.query, std=0.1)
        self.k = nn.Linear(dim, d_attn, bias=False)
        self.v = nn.Linear(dim, dim, bias=False)
        self.scale = d_attn ** -0.5
        self.max_gate = nn.Parameter(torch.zeros(1, num_queries, 1))

    def forward(self, x, mask=None):
        b, n, d = x.shape
        k, v = self.k(x), self.v(x)
        gain = self.logit_scale.clamp(max=math.log(100.0)).exp()
        scores = torch.bmm(self.query.expand(b, -1, -1) * (self.scale * gain), k.transpose(1, 2))
        out_mean = torch.bmm(F.softmax(scores, dim=-1), v)
        idx = scores.argmax(dim=-1).unsqueeze(-1).expand(-1, -1, d)
        out_max = torch.gather(v, 1, idx)
        return out_mean + torch.tanh(self.max_gate) * out_max


class ClassInteractionBlock(nn.Module):
    """Self-attention over the class dimension; both branches start near identity (layer scale)."""

    def __init__(self, dim: int, num_heads: int = 4, dropout: float = 0.1, layerscale: float = 1e-4):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.self_attn = nn.MultiheadAttention(dim, num_heads, dropout=dropout, batch_first=True)
        self.mlp = nn.Sequential(nn.Linear(dim, dim * 2), nn.GELU(), nn.Dropout(dropout),
                                 nn.Linear(dim * 2, dim), nn.Dropout(dropout))
        self.norm_mlp = nn.LayerNorm(dim)
        self.ci_gamma_attn = nn.Parameter(torch.full((dim,), float(layerscale)))
        self.ci_gamma_mlp = nn.Parameter(torch.full((dim,), float(layerscale)))

    def forward(self, x):
        xn = self.norm(x)
        x = x + self.ci_gamma_attn * self.self_attn(xn, xn, xn)[0]
        return x + self.ci_gamma_mlp * self.mlp(self.norm_mlp(x))


class ClasswiseHead(nn.Module):
    """Shared Linear-LayerNorm-GELU-Dropout block, then one weight vector and bias per class."""

    def __init__(self, dim: int, num_classes: int, dropout: float = 0.1):
        super().__init__()
        self.shared = nn.Sequential(nn.Linear(dim, dim), nn.LayerNorm(dim), nn.GELU(), nn.Dropout(dropout))
        self.cls_w = nn.Parameter(torch.empty(num_classes, dim))
        self.cls_b = nn.Parameter(torch.zeros(num_classes))
        nn.init.trunc_normal_(self.cls_w, std=0.02)

    def forward(self, pooled):
        return (self.shared(pooled) * self.cls_w.unsqueeze(0)).sum(-1) + self.cls_b


class RadiologyHead(nn.Module):
    """tokens (B, N, D) -> logits (B, K): class-query pool -> width reduction -> class interaction -> classifier."""

    def __init__(self, dim: int, num_classes: int = 18, pool: str = "mhead", num_heads: int = 4,
                 d_attn: int | None = None, dropout: float = 0.1, head_dim: int | None = None,
                 layerscale: float = 1e-4, identity_proj: bool = True):
        super().__init__()
        if pool == "mhead":
            self.pool = MultiHeadQueryAttentionPool(dim, num_classes, num_heads, d_attn, identity_proj)
        elif pool == "hybrid":
            self.pool = GatedMaxTokenPool(dim, num_classes, d_attn)
        else:
            raise ValueError(f"unknown pool {pool!r}")
        w = int(head_dim) if head_dim else dim
        self.reduce = nn.Sequential(nn.Linear(dim, w), nn.LayerNorm(w)) if w != dim else None
        self.class_interaction = ClassInteractionBlock(w, num_heads, dropout, layerscale)
        self.head = ClasswiseHead(w, num_classes, dropout)

    def forward(self, x, mask=None):
        pooled = self.pool(x, mask)
        if self.reduce is not None:
            pooled = self.reduce(pooled)
        return self.head(self.class_interaction(pooled))


class CTClassifier(nn.Module):
    """encoder (LoRA-wrapped) -> slice position encoding -> radiology head."""

    def __init__(self, encoder, dim: int, num_classes: int, slice_posenc: bool = True, **head_kw):
        super().__init__()
        self.encoder = encoder
        base = encoder.base_model.model if hasattr(encoder, "base_model") else encoder
        self.pos_enc = SlicePosEnc(base.embed_dim, base.max_slices, base.tokens_per_slice) if slice_posenc else None
        self.head = RadiologyHead(dim, num_classes, **head_kw)

    def forward(self, x):
        tokens, mask = self.encoder(x)
        if self.pos_enc is not None:
            tokens = self.pos_enc(tokens)
        return self.head(tokens, mask)


# Parameter-name prefixes that sit outside the frozen backbone (used for optimizer groups)
HEAD_PREFIXES = ("head.", "pos_enc.")
HEAD_SUBSTRINGS = ("conv_branch", "cb_gamma")


def head_kwargs(cfg: dict) -> dict:
    return dict(pool=cfg.get("pool", "mhead"), num_heads=int(cfg.get("head_heads") or 4),
                d_attn=int(cfg["pool_d_attn"]) if cfg.get("pool_d_attn") else None,
                dropout=float(cfg.get("head_dropout") or 0.1),
                head_dim=int(cfg["head_dim"]) if cfg.get("head_dim") else None,
                layerscale=float(cfg.get("head_layerscale") or 0.0),
                identity_proj=bool(cfg.get("head_identity_proj", True)))


_LORA_KEY = re.compile(r"^(?P<mod>.+)\.lora_(?:A|B)\.(?:[^.]+\.)?weight$")


def lora_targets_from_state_dict(sd) -> list[str]:
    """Module paths that carry LoRA weights in a checkpoint."""
    out = set()
    for k in sd:
        m = _LORA_KEY.match(k)
        if m:
            mod = m.group("mod")
            i = mod.rfind("base_model.model.")
            out.add(mod[i + len("base_model.model."):] if i >= 0 else mod)
    return sorted(out)


def build_from_config(cfg: dict, state_dict: dict, num_classes: int = 18) -> CTClassifier:
    """Rebuild a released model from its config (streams.json) and weights."""
    from peft import LoraConfig, get_peft_model

    encoder = FlexiCT2DEncoder(img_size=int(cfg["img_size"]), patch_size=int(cfg["patch_size"]),
                               max_slices=int(cfg["max_slices"]),
                               slice_microbatch=int(cfg.get("slice_microbatch") or cfg["max_slices"]),
                               conv_branch=bool(cfg.get("conv_branch")),
                               load_pretrained=False)
    dim = encoder.embed_dim
    if cfg.get("use_lora"):
        encoder = get_peft_model(encoder, LoraConfig(
            r=int(cfg["lora_r"]), lora_alpha=int(cfg["lora_alpha"]), lora_dropout=0.0,
            bias="none", target_modules=lora_targets_from_state_dict(state_dict)))
    model = CTClassifier(encoder, dim, num_classes, slice_posenc=bool(cfg.get("slice_posenc")),
                         **head_kwargs(cfg))
    model.load_state_dict(state_dict, strict=True)
    return model.eval()