"""SOFIA v7.1 full model.

Architecture:
    stem → stage1 (StripDCN-lite) → stage2 (StripMixConv)
         → stage3 (MambaVision MV5) → stage4 (MambaVision MV1)
         → 1×1 neck projection → {Sat-Head, UAV-Head} + (training) Ring aux

Features:
- Siamese backbone with optional weight-sharing stages 1-2
- Asymmetric Sat/UAV heads (CVGL-Aware Head v7.1-α)
- Training-time Ring LPN aux for partial matching robustness
- KD taps on F2, F3, F4 for future teacher-student distillation
"""

from __future__ import annotations

from typing import Dict, List, Optional

import torch
import torch.nn as nn
import torch.nn.functional as F

from .config import SOFIAConfig
from .blocks import (
    Downsample,
    EVSSBridge,
    MambaVisionBlock,
    StripDCNLiteBlock,
    StripMixConvBlock,
)
from .layers import (
    AltitudeFiLM,
    CircularHarmonicPool,
    GGeM,
    LayerNorm2d,
    TextFiLM,
)


# ============================================================
# Stem
# ============================================================

class Stem(nn.Module):
    """Dual-conv FastViT-style stem: 3 → mid (s=2) → out (s=1).

    Downsampling ×2 total (input 256 → 128).
    """

    def __init__(self, in_channels: int, mid_channels: int, out_channels: int) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, mid_channels, 3, stride=2, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(mid_channels)
        self.act1 = nn.SiLU(inplace=True)
        self.conv2 = nn.Conv2d(mid_channels, out_channels, 3, stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.act2 = nn.SiLU(inplace=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.act1(self.bn1(self.conv1(x)))
        x = self.act2(self.bn2(self.conv2(x)))
        return x


# ============================================================
# Backbone (shared s1-2 optional, separate s3-4)
# ============================================================

class Backbone(nn.Module):
    """SOFIA v7.1 backbone: 4 stages + stem.

    Stage 1: StripDCN-lite blocks
    Stage 2: StripMixConv blocks
    Stage 3: MambaVision MV5 blocks (Mamba ∥ MHSA ∥ Strip)
    Stage 4: MambaVision MV1 blocks (Mamba ∥ MHSA)
    """

    def __init__(self, cfg: SOFIAConfig) -> None:
        super().__init__()
        self.cfg = cfg
        dims = cfg.embed_dims
        depths = cfg.depths

        # Stem: input → stem_out
        self.stem = Stem(cfg.in_channels, cfg.stem_mid, cfg.stem_out)

        # Stage 1: stem_out → dims[0], no downsample at block level
        # (stem already did 4× total; stage 1 operates at 64×64 for 256 input)
        self.ds1 = Downsample(cfg.stem_out, dims[0])
        self.stage1 = nn.Sequential(*[
            StripDCNLiteBlock(
                in_channels=dims[0],
                out_channels=dims[0],
                expand=cfg.mbconv_expand,
                kernel=cfg.strip_kernel_s1,
                stride=1,
                se_ratio=cfg.se_ratio,
                use_dcn=cfg.use_dcn_strip,
            )
            for _ in range(depths[0])
        ])

        # Stage 2
        self.ds2 = Downsample(dims[0], dims[1])
        self.stage2 = nn.Sequential(*[
            StripMixConvBlock(
                in_channels=dims[1],
                out_channels=dims[1],
                expand=cfg.mbconv_expand,
                strip_kernel=cfg.strip_kernel_s2,
                mix_kernels=cfg.mix_kernels,
                stride=1,
                se_ratio=cfg.se_ratio,
            )
            for _ in range(depths[1])
        ])

        # Stage 3 — optional EVSS bridge after downsample (B6-inspired refinement
        # to smooth DCN→MambaVision semantic gap)
        self.ds3 = Downsample(dims[1], dims[2])
        if cfg.use_evss_bridge and "pre_stage3" in cfg.evss_bridge_locations:
            self.bridge3 = EVSSBridge(
                channels=dims[2],
                mamba_d_state=cfg.mamba_d_state,
                mamba_dt_rank=cfg.mamba_dt_rank,
                mamba_variant="mamba1",    # lightweight refinement, not primary mixer
                mamba_backend=cfg.mamba_backend,
                se_reduction=cfg.se_ratio,
            )
        else:
            self.bridge3 = None
        self.stage3 = nn.Sequential(*[
            MambaVisionBlock(
                channels=dims[2],
                num_heads=cfg.num_heads_s3,
                d_state=cfg.mamba_d_state,
                dt_rank=cfg.mamba_dt_rank,
                use_strip_branch=cfg.use_strip_branch_s3,
                ffn_expand=cfg.ffn_expand,
                strip_kernel=cfg.strip_kernel_s1,
                mamba_backend=cfg.mamba_backend,
                mamba_variant=cfg.mamba_variant,
                mamba_extra_kwargs=cfg.mamba_extra_kwargs,
            )
            for _ in range(depths[2])
        ])

        # Stage 4 — optional EVSS bridge after downsample
        self.ds4 = Downsample(dims[2], dims[3])
        if cfg.use_evss_bridge and "pre_stage4" in cfg.evss_bridge_locations:
            self.bridge4 = EVSSBridge(
                channels=dims[3],
                mamba_d_state=cfg.mamba_d_state,
                mamba_dt_rank=cfg.mamba_dt_rank,
                mamba_variant="mamba1",
                mamba_backend=cfg.mamba_backend,
                se_reduction=cfg.se_ratio,
            )
        else:
            self.bridge4 = None
        self.stage4 = nn.Sequential(*[
            MambaVisionBlock(
                channels=dims[3],
                num_heads=cfg.num_heads_s4,
                d_state=cfg.mamba_d_state,
                dt_rank=cfg.mamba_dt_rank,
                use_strip_branch=cfg.use_strip_branch_s4,
                ffn_expand=cfg.ffn_expand,
                strip_kernel=cfg.strip_kernel_s1,
                mamba_backend=cfg.mamba_backend,
                mamba_variant=cfg.mamba_variant,
                mamba_extra_kwargs=cfg.mamba_extra_kwargs,
            )
            for _ in range(depths[3])
        ])

    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
        """Returns dict with F1-F4 taps for KD/multi-scale use."""
        s0 = self.stem(x)                  # stem_out, H/2
        f1 = self.stage1(self.ds1(s0))     # dims[0], H/4
        f2 = self.stage2(self.ds2(f1))     # dims[1], H/8

        # Stage 3 with optional EVSS bridge.
        ds3_out = self.ds3(f2)
        if self.bridge3 is not None:
            ds3_out = self.bridge3(ds3_out)
        f3 = self.stage3(ds3_out)          # dims[2], H/16

        # Stage 4 with optional EVSS bridge.
        ds4_out = self.ds4(f3)
        if self.bridge4 is not None:
            ds4_out = self.bridge4(ds4_out)
        f4 = self.stage4(ds4_out)          # dims[3], H/32 (= 8×8 for 256 input)

        return {"s0": s0, "f1": f1, "f2": f2, "f3": f3, "f4": f4}


# ============================================================
# Neck: ultra-lite 1×1 projection
# ============================================================

class UltraLiteNeck(nn.Module):
    """1×1 projection + BN + activation."""

    def __init__(self, in_channels: int, out_channels: int) -> None:
        super().__init__()
        self.proj = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.SiLU(inplace=True),
        )

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


# ============================================================
# Heads
# ============================================================

class SatHead(nn.Module):
    """Satellite view head: [TextFiLM] + GGeM + BN + Linear [+ L2].

    `text_emb` is optional; when None or `text_film` is disabled, behaves
    identically to the original head. `return_normalized=False` returns
    pre-L2 descriptors for late gated fusion in a wrapper.
    """

    def __init__(
        self,
        channels: int,
        d_descriptor: int,
        return_normalized: bool = True,
        use_text_film: bool = False,
        text_film_dim: int = 1024,
        text_film_hidden: int = 256,
    ) -> None:
        super().__init__()
        self.return_normalized = return_normalized
        self.use_text_film = use_text_film
        if use_text_film:
            self.text_film = TextFiLM(channels, text_dim=text_film_dim, hidden_dim=text_film_hidden)
        self.ggem = GGeM(channels)
        self.bn = nn.BatchNorm1d(channels, affine=False)
        self.proj = nn.Linear(channels, d_descriptor)

    def forward(
        self,
        x: torch.Tensor,
        text_emb: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """x: [B, C, H, W], text_emb: [B, D_txt] or None -> [B, d]"""
        if self.use_text_film:
            x = self.text_film(x, text_emb)
        g = self.ggem(x)
        g = self.bn(g)
        g = self.proj(g)
        if self.return_normalized:
            g = F.normalize(g, p=2, dim=-1)
        return g


class UAVHead(nn.Module):
    """UAV view head: FiLM(altitude) [+ TextFiLM] + CHP + BN + Linear [+ L2].

    Formally SO(2)-invariant via CHP. Altitude-aware via FiLM. Optional
    text-FiLM is applied AFTER altitude-FiLM (zero-init β so it starts as
    identity). `return_normalized=False` returns pre-L2 descriptors for
    late gated fusion in a wrapper.
    """

    def __init__(
        self,
        channels: int,
        d_descriptor: int,
        rings: int = 8,
        angles: int = 16,
        harmonics: int = 4,
        use_film: bool = True,
        altitude_norm: float = 500.0,
        return_normalized: bool = True,
        use_text_film: bool = False,
        text_film_dim: int = 1024,
        text_film_hidden: int = 256,
    ) -> None:
        super().__init__()
        self.return_normalized = return_normalized
        self.use_film = use_film
        self.use_text_film = use_text_film
        if use_film:
            self.film = AltitudeFiLM(channels, altitude_norm=altitude_norm)
        if use_text_film:
            self.text_film = TextFiLM(channels, text_dim=text_film_dim, hidden_dim=text_film_hidden)

        self.chp = CircularHarmonicPool(
            channels=channels,
            rings=rings,
            angles=angles,
            harmonics=harmonics,
        )
        chp_dim = channels * harmonics
        self.bn = nn.BatchNorm1d(chp_dim, affine=False)
        self.proj = nn.Linear(chp_dim, d_descriptor)

    def forward(
        self,
        x: torch.Tensor,
        altitude: Optional[torch.Tensor] = None,
        text_emb: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """x: [B, C, H, W], altitude: [B] or None, text_emb: [B, D_txt] or None -> [B, d]"""
        if self.use_film:
            x = self.film(x, altitude)
        if self.use_text_film:
            x = self.text_film(x, text_emb)
        g = self.chp(x)
        g = self.bn(g)
        g = self.proj(g)
        if self.return_normalized:
            g = F.normalize(g, p=2, dim=-1)
        return g


class RingAuxHead(nn.Module):
    """LPN Square-Ring pool + per-ring Linear + L2 (training-only auxiliary).

    Used for partial-matching robustness. Drop at inference.
    """

    def __init__(
        self,
        channels: int,
        rings: int = 4,
        d_per_ring: int = 128,
        feature_size: int = 8,
    ) -> None:
        super().__init__()
        self.rings = rings
        self.feature_size = feature_size
        self.d_per_ring = d_per_ring

        # Per-ring GGeM
        self.ggems = nn.ModuleList([GGeM(channels) for _ in range(rings)])
        self.projs = nn.ModuleList([
            nn.Linear(channels, d_per_ring) for _ in range(rings)
        ])

        # Precompute ring masks
        masks = self._make_ring_masks(rings, feature_size)    # [R, H, W]
        self.register_buffer("ring_masks", masks, persistent=False)

    @staticmethod
    def _make_ring_masks(R: int, S: int) -> torch.Tensor:
        """Concentric square rings on SxS feature map."""
        masks = torch.zeros(R, S, S)
        center = (S - 1) / 2.0
        # Define ring by Chebyshev distance thresholds
        for r in range(R):
            r_min = r * (S / (2 * R))
            r_max = (r + 1) * (S / (2 * R))
            for i in range(S):
                for j in range(S):
                    dist = max(abs(i - center), abs(j - center))
                    if r_min <= dist < r_max or (r == R - 1 and dist <= r_max):
                        masks[r, i, j] = 1.0
        # Normalize each ring to sum to 1 for pooling stability (not required but cleaner)
        masks = masks / masks.sum(dim=(1, 2), keepdim=True).clamp(min=1.0)
        return masks

    def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
        """
        Args:
            x: [B, C, H, W]
        Returns:
            List of R tensors each [B, d_per_ring], L2-normalized
        """
        B, C, H, W = x.shape
        outs = []
        for r in range(self.rings):
            mask = self.ring_masks[r].to(x.dtype)    # [H, W]
            # Apply mask
            x_masked = x * mask.unsqueeze(0).unsqueeze(0)    # [B, C, H, W]
            g = self.ggems[r](x_masked)    # [B, C]
            g = self.projs[r](g)    # [B, d_per_ring]
            g = F.normalize(g, p=2, dim=-1)
            outs.append(g)
        return outs


# ============================================================
# Full SOFIA v7.1 model
# ============================================================

class SOFIAv71(nn.Module):
    """Full SOFIA v7.1 model.

    Forward signature:
        forward(sat: [B,3,H,W], uav: [B,3,H,W], altitude: [B] = None,
                return_features: bool = False) -> dict
    """

    def __init__(self, cfg: SOFIAConfig) -> None:
        super().__init__()
        cfg.validate()
        self.cfg = cfg

        # Backbones (shared or separate depending on config)
        self.backbone_shared = Backbone(cfg)
        if cfg.share_stages_1_2:
            # Single backbone, use one for both
            self.backbone_sat = None
            self.backbone_uav = None
        else:
            # Fully separate backbones (rare case)
            self.backbone_sat = Backbone(cfg)
            self.backbone_uav = Backbone(cfg)
            # Clear shared
            del self.backbone_shared
            self.backbone_shared = None

        # Neck (shared)
        self.neck = UltraLiteNeck(cfg.embed_dims[-1], cfg.neck_channels)

        # Heads
        if cfg.use_asymmetric_heads:
            self.sat_head = SatHead(
                cfg.neck_channels, cfg.d_descriptor,
                return_normalized=cfg.return_normalized,
                use_text_film=cfg.use_text_film_sat,
                text_film_dim=cfg.text_film_dim,
                text_film_hidden=cfg.text_film_hidden,
            )
            self.uav_head = UAVHead(
                channels=cfg.neck_channels,
                d_descriptor=cfg.d_descriptor,
                rings=cfg.chp_rings,
                angles=cfg.chp_angles,
                harmonics=cfg.chp_harmonics,
                use_film=cfg.use_film_altitude,
                altitude_norm=cfg.altitude_norm,
                return_normalized=cfg.return_normalized,
                use_text_film=cfg.use_text_film_uav,
                text_film_dim=cfg.text_film_dim,
                text_film_hidden=cfg.text_film_hidden,
            )
        else:
            # Symmetric: use SatHead for both
            self.sat_head = SatHead(
                cfg.neck_channels, cfg.d_descriptor,
                return_normalized=cfg.return_normalized,
                use_text_film=cfg.use_text_film_sat,
                text_film_dim=cfg.text_film_dim,
                text_film_hidden=cfg.text_film_hidden,
            )
            self.uav_head = self.sat_head

        # Ring aux (training-only)
        # Resolution chain: input/2 (stem) → /2 (ds1) → /2 (ds2) → /2 (ds3) → /2 (ds4)
        # Total ×32 downsampling → input 256 gives 8×8 final feature map
        if cfg.use_ring_aux:
            feat_size = cfg.input_size // 32
            self.ring_head = RingAuxHead(
                channels=cfg.neck_channels,
                rings=cfg.ring_count,
                d_per_ring=cfg.d_descriptor // cfg.ring_count,
                feature_size=feat_size,
            )
        else:
            self.ring_head = None

    def _extract_features(self, x: torch.Tensor, view: str) -> Dict[str, torch.Tensor]:
        """Run backbone for sat or uav view."""
        if self.cfg.share_stages_1_2:
            return self.backbone_shared(x)
        else:
            bb = self.backbone_sat if view == "sat" else self.backbone_uav
            return bb(x)

    def forward(
        self,
        sat: Optional[torch.Tensor] = None,
        uav: Optional[torch.Tensor] = None,
        altitude: Optional[torch.Tensor] = None,
        text_emb_sat: Optional[torch.Tensor] = None,
        text_emb_uav: Optional[torch.Tensor] = None,
        return_features: bool = False,
    ) -> Dict[str, torch.Tensor]:
        """
        Args:
            sat: [B, 3, H, W] satellite image (or None to skip)
            uav: [B, 3, H, W] UAV image (or None to skip)
            altitude: [B] altitude in meters for UAV (or None = neutral)
            text_emb_sat: [B, D_txt] caption embedding for SatHead text-FiLM (or None)
            text_emb_uav: [B, D_txt] caption embedding for UAVHead text-FiLM (or None)
            return_features: if True, also return backbone features F2-F4 for KD

        Returns:
            dict with g_sat, g_uav (global descriptors), optional features,
            and optional rings (training only)
        """
        result: Dict[str, torch.Tensor] = {}

        # SAT path
        if sat is not None:
            feats_sat = self._extract_features(sat, view="sat")
            f4_sat = feats_sat["f4"]
            neck_sat = self.neck(f4_sat)
            g_sat = self.sat_head(neck_sat, text_emb=text_emb_sat)
            result["g_sat"] = g_sat
            if return_features:
                result["features_sat"] = feats_sat
                result["neck_sat"] = neck_sat
            if self.training and self.ring_head is not None:
                result["rings_sat"] = self.ring_head(neck_sat)

        # UAV path
        if uav is not None:
            feats_uav = self._extract_features(uav, view="uav")
            f4_uav = feats_uav["f4"]
            neck_uav = self.neck(f4_uav)
            if self.cfg.use_asymmetric_heads:
                g_uav = self.uav_head(neck_uav, altitude=altitude, text_emb=text_emb_uav)
            else:
                g_uav = self.uav_head(neck_uav, text_emb=text_emb_uav)
            result["g_uav"] = g_uav
            if return_features:
                result["features_uav"] = feats_uav
                result["neck_uav"] = neck_uav
            if self.training and self.ring_head is not None:
                result["rings_uav"] = self.ring_head(neck_uav)

        return result

    def count_parameters(self, only_trainable: bool = True) -> Dict[str, int]:
        """Count parameters per module."""
        counts = {}
        total = 0
        for name, module in self.named_children():
            if module is None:
                continue
            n = sum(p.numel() for p in module.parameters() if not only_trainable or p.requires_grad)
            counts[name] = n
            total += n
        counts["_total"] = total
        return counts


# ============================================================
# Build helper
# ============================================================

def build_sofia(preset: str = "M") -> SOFIAv71:
    """Build SOFIA v7.1 model from preset name."""
    from .config import sofia_m_config, sofia_l_config, sofia_tiny_config
    preset_map = {
        "M": sofia_m_config,
        "L": sofia_l_config,
        "Tiny": sofia_tiny_config,
    }
    if preset not in preset_map:
        raise ValueError(f"Unknown preset '{preset}'. Available: {list(preset_map)}")
    cfg = preset_map[preset]()
    return SOFIAv71(cfg)


if __name__ == "__main__":
    # Smoke test with Tiny preset (fast)
    print("Building SOFIA-Tiny for smoke test...")
    model = build_sofia("Tiny")
    model.eval()

    counts = model.count_parameters()
    total_m = counts["_total"] / 1e6
    print(f"Total params: {total_m:.2f} M")
    print("Per-module:")
    for k, v in counts.items():
        print(f"  {k}: {v / 1e6:.3f} M")

    # Dry forward
    sat = torch.randn(1, 3, 256, 256)
    uav = torch.randn(1, 3, 256, 256)
    alt = torch.tensor([150.0])
    with torch.no_grad():
        out = model(sat=sat, uav=uav, altitude=alt)
    for k, v in out.items():
        if isinstance(v, torch.Tensor):
            print(f"  out[{k}]: {tuple(v.shape)}")
        else:
            print(f"  out[{k}]: {type(v).__name__}")
