from __future__ import annotations

"""EdgeNeXt builder + transfer-learning regimes (freeze / partial / MONA).

Three fine-tuning regimes are exposed (see PROTOCOL §5):
* ``full``    — train the whole backbone + head (one or two LR groups).
* ``partial`` — freeze the first ``4 - n_unfrozen`` stages, train the rest + head;
                ``n_unfrozen=0`` is pure linear probing (feature extraction).
* ``mona``    — freeze the whole backbone, train only MONA adapters + head (PEFT).

``feature_extract_features`` returns L2-normalized pooled embeddings for the
UMAP analysis (see ``umap_analysis.py``).
"""

from dataclasses import dataclass

import timm
import torch
import torch.nn.functional as F
from torch import nn

from .mona import attach_mona_adapters


@dataclass
class ModelConfig:
    """Backbone + regime configuration (a subset is searched by Optuna)."""

    model_name: str = "edgenext_small.usi_in1k"
    num_classes: int = 13  # Naruto Sign: 12 seals + 'zero' (verify after download)
    pretrained: bool = True
    drop_rate: float = 0.0  # classifier dropout
    drop_path_rate: float = 0.0  # stochastic depth
    regime: str = "partial"  # {"full", "partial", "mona"}
    n_unfrozen_stages: int = 1  # used by "partial": 0..4
    mona_stages: tuple[int, ...] = (2, 3)  # used by "mona"
    mona_bottleneck: int = 64
    mona_kernels: tuple[int, ...] = (3, 5, 7)


def _set_norm_eval(module: nn.Module) -> None:
    """Put frozen BatchNorm/LayerNorm/GroupNorm in eval mode (freeze stats)."""
    for m in module.modules():
        if isinstance(m, (nn.BatchNorm2d, nn.LayerNorm, nn.GroupNorm)):
            m.eval()


def build_model(cfg: ModelConfig) -> nn.Module:
    """Create an EdgeNeXt and apply the requested fine-tuning regime.

    Args:
        cfg: Model configuration.

    Returns:
        The configured ``nn.Module``. Trainable parameters depend on ``regime``.
    """
    model = timm.create_model(
        cfg.model_name,
        pretrained=cfg.pretrained,
        num_classes=cfg.num_classes,
        drop_rate=cfg.drop_rate,
        drop_path_rate=cfg.drop_path_rate,
    )

    if cfg.regime == "full":
        return model

    # Freeze the whole backbone first; head stays trainable in every regime.
    for name, p in model.named_parameters():
        if not _is_head(name):
            p.requires_grad_(False)

    if cfg.regime == "partial":
        n = max(0, min(cfg.n_unfrozen_stages, len(model.stages)))
        if n > 0:
            for idx in range(len(model.stages) - n, len(model.stages)):
                for p in model.stages[idx].parameters():
                    p.requires_grad_(True)
        # final norm before the head (if present) follows the head.
        for name, p in model.named_parameters():
            if "norm_pre" in name or name.startswith("norm"):
                p.requires_grad_(True)
    elif cfg.regime == "mona":
        attach_mona_adapters(
            model,
            list(cfg.mona_stages),
            bottleneck=cfg.mona_bottleneck,
            kernels=cfg.mona_kernels,
        )  # adapters are created trainable by default
    else:
        raise ValueError(f"unknown regime: {cfg.regime!r}")

    return model


def _is_head(param_name: str) -> bool:
    """Heuristic: classifier-head parameters in timm models contain 'head'/'fc'."""
    return "head" in param_name or param_name.startswith("fc")


def freeze_consistency(model: nn.Module) -> None:
    """Set frozen norm layers to eval so their running stats are not updated.

    Call this **after** ``model.train()`` in every training step for the
    ``partial``/``mona`` regimes (see PROTOCOL §5.2 — the classic BN pitfall).
    """
    for module in model.modules():
        if isinstance(module, (nn.BatchNorm2d, nn.LayerNorm, nn.GroupNorm)):
            if not any(p.requires_grad for p in module.parameters(recurse=False)):
                module.eval()


def build_param_groups(
    model: nn.Module, base_lr: float, *, backbone_lr_mult: float = 0.1
) -> list[dict]:
    """Two LR groups: a smaller LR for backbone, the base LR for head/adapters.

    Discriminative learning rates (lower layers -> lower LR) are a standard
    transfer-learning trick (see PROTOCOL §5.3). Only parameters with
    ``requires_grad=True`` are included.

    Args:
        model: The configured model.
        base_lr: LR for the head and MONA adapters.
        backbone_lr_mult: Multiplier applied to ``base_lr`` for backbone params.

    Returns:
        A list of optimizer parameter-group dicts.
    """
    head, backbone = [], []
    for name, p in model.named_parameters():
        if not p.requires_grad:
            continue
        if _is_head(name) or "mona_adapters" in name:
            head.append(p)
        else:
            backbone.append(p)

    groups: list[dict] = []
    if head:
        groups.append({"params": head, "lr": base_lr})
    if backbone:
        groups.append({"params": backbone, "lr": base_lr * backbone_lr_mult})
    return groups


def count_trainable(model: nn.Module) -> tuple[int, int]:
    """Return ``(trainable, total)`` parameter counts."""
    total = sum(p.numel() for p in model.parameters())
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    return trainable, total


@torch.inference_mode()
def extract_features(
    model: nn.Module, x: torch.Tensor, *, l2_normalize: bool = True
) -> torch.Tensor:
    """Return pooled pre-logit embeddings ``[B, num_features]`` for UMAP.

    Uses timm's ``forward_features`` + ``forward_head(..., pre_logits=True)`` so
    the classifier is bypassed.

    Args:
        model: A timm model in eval mode.
        x: Input batch ``[B, 3, H, W]`` already normalized.
        l2_normalize: If ``True`` L2-normalize embeddings (cosine geometry).

    Returns:
        Float tensor on CPU of shape ``[B, num_features]``.
    """
    feats = model.forward_features(x)
    emb = model.forward_head(feats, pre_logits=True)
    if l2_normalize:
        emb = F.normalize(emb, dim=1)
    return emb.detach().cpu()