"""SOFIA v7.1 configuration system.

Two scale presets targeting Jetson Orin NX INT8 deployment:

| Preset   | Params   | INT8 size | FP16 size | Target latency |
|----------|---------:|----------:|----------:|---------------:|
| SOFIA-M  | ~500 M   | ~500 MB   | ~1 GB     | ~18 ms         |
| SOFIA-L  | ~1 B     | ~1 GB     | ~2 GB     | ~20 ms         |

Based on HYP_SOFIA_v7_UltraDeep_дизайн.md (Phases 2'–5' + Phase 4''
CVGL-Aware Head + revision to ultra-lite 1x1 neck).
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import List, Literal, Optional


@dataclass
class SOFIAConfig:
    """Configuration for SOFIA v7.1 student architecture.

    Controls backbone width/depth, neck, and CVGL-Aware Head.
    """

    # -------- Input --------
    input_size: int = 256
    in_channels: int = 3

    # -------- Stem (dual-conv FastViT-style) --------
    stem_mid: int = 64          # intermediate channels
    stem_out: int = 128         # output channels into stage 1

    # -------- Backbone dimensions per stage s1..s4 --------
    embed_dims: List[int] = field(
        default_factory=lambda: [256, 512, 1280, 1536]
    )
    depths: List[int] = field(default_factory=lambda: [3, 4, 15, 3])

    # -------- Stage 1–2 block parameters --------
    mbconv_expand: int = 4
    se_ratio: int = 16
    strip_kernel_s1: int = 7    # strip DW kernel size stage 1
    strip_kernel_s2: int = 5    # strip DW kernel size stage 2
    mix_kernels: List[int] = field(
        default_factory=lambda: [3, 5, 7]
    )   # MixConv DW kernel sizes stage 2
    use_dcn_strip: bool = True  # adaptive offset on horizontal strip

    # -------- Stage 3–4 (MambaVision) --------
    mamba_d_state: int = 16
    mamba_dt_rank: Optional[int] = None     # auto = max(1, C // 16)
    mamba_backend: Literal["auto", "torch", "mamba_ssm"] = "auto"
    # "auto" uses mamba_ssm if importable, else torch fallback

    # Mamba variant — one of:
    #   "mamba2"            (preferred, SSD dual form, 2-8x faster)
    #   "mamba1"            (original selective scan, mature)
    #   "efficient_vmamba"  (speed fallback, atrous scan ~2-3x speedup)
    mamba_variant: Literal["mamba1", "mamba2", "efficient_vmamba"] = "mamba2"
    # Per-variant tunables passed through to factory
    mamba_extra_kwargs: dict = field(default_factory=lambda: {
        "d_state_mamba2": 64,   # Mamba-2 typically uses N=64 (not 16)
        "headdim": 64,          # Mamba-2 head dim
        "expand": 2,            # Mamba-2 inner expansion factor
        "d_conv": 4,            # Mamba-2 local conv kernel
        "n_directions": 2,      # EfficientVMamba: 2 or 4 atrous directions
    })

    num_heads_s3: int = 8
    num_heads_s4: int = 8
    use_strip_branch_s3: bool = True    # MV5 with Strip branch
    use_strip_branch_s4: bool = False   # MV1 without Strip branch
    ffn_expand: int = 4

    # -------- EVSS-style bridge (B6-inspired, opt-in) --------
    # When True, inserts a within-resolution dual-path refinement block
    # right after each downsample to stage 3 and stage 4. Smooths semantic
    # gap between heterogeneous stages (DCN → MambaVision). Adds ~165K params
    # and ~30 MMAC per insertion at C=192, HW=256.
    use_evss_bridge: bool = False
    evss_bridge_locations: List[str] = field(
        default_factory=lambda: ["pre_stage3"]    # subset of {pre_stage3, pre_stage4}
    )

    # -------- Neck (ultra-lite 1x1 projection) --------
    neck_channels: int = 192    # C_n, output channels from neck

    # -------- CVGL-Aware Head v7.1-α --------
    d_descriptor: int = 512     # global descriptor dimensionality
    use_asymmetric_heads: bool = True   # Sat vs UAV different heads
    chp_rings: int = 8
    chp_angles: int = 16
    chp_harmonics: int = 4
    use_film_altitude: bool = True
    altitude_norm: float = 500.0    # divides altitude in meters
    ring_count: int = 4             # LPN rings auxiliary
    use_ring_aux: bool = True       # training-only ring aux branch

    # -------- Text fusion (extension for caption-conditioned heads) --------
    return_normalized: bool = True   # if False, heads return pre-L2 features (for late gated fusion)
    use_text_film_sat: bool = False  # text-FiLM modulation in SatHead before GGeM
    use_text_film_uav: bool = False  # text-FiLM modulation in UAVHead alongside altitude FiLM
    text_film_dim: int = 1024        # text embedding dim feeding FiLM (matches TextFusionMLP out_dim)
    text_film_hidden: int = 256

    # -------- Weight-sharing --------
    share_stages_1_2: bool = True   # sat ↔ UAV shared weights stages 1-2

    # -------- KD taps (enable for future teacher KD) --------
    enable_kd_taps: bool = True

    # -------- Deployment hints --------
    precision: Literal["fp32", "fp16", "int8_mixed"] = "fp16"

    def validate(self) -> None:
        assert len(self.embed_dims) == 4, (
            f"embed_dims must have 4 entries, got {len(self.embed_dims)}"
        )
        assert len(self.depths) == 4, (
            f"depths must have 4 entries, got {len(self.depths)}"
        )
        assert self.input_size % 32 == 0, (
            f"input_size must be divisible by 32 (4 downsamples × stem), "
            f"got {self.input_size}"
        )

    def summary(self) -> str:
        return (
            f"SOFIAConfig(stem={self.stem_mid}/{self.stem_out}, "
            f"dims={self.embed_dims}, depths={self.depths}, "
            f"neck={self.neck_channels}, d={self.d_descriptor}, "
            f"precision={self.precision})"
        )


# ============================================================
# Scale Presets
# ============================================================

def sofia_m_config() -> SOFIAConfig:
    """SOFIA-M: ~500 M params target (~500 MB INT8, ~1 GB FP16).

    Fits in 500 MB VRAM after INT8 quantization on Jetson Orin NX.
    Expected latency: ~18 ms.
    """
    return SOFIAConfig(
        stem_mid=64,
        stem_out=128,
        embed_dims=[256, 512, 1280, 1536],
        depths=[3, 4, 15, 3],
        neck_channels=192,
        d_descriptor=512,
    )


def sofia_l_config() -> SOFIAConfig:
    """SOFIA-L: ~1 B params target (~1 GB INT8, ~2 GB FP16).

    Fits in 1 GB VRAM after INT8 quantization on Jetson Orin NX.
    Expected latency: ~20 ms.
    """
    return SOFIAConfig(
        stem_mid=64,
        stem_out=128,
        embed_dims=[256, 512, 1536, 2048],
        depths=[3, 4, 20, 3],
        neck_channels=256,
        d_descriptor=1024,
    )


def sofia_tiny_config() -> SOFIAConfig:
    """SOFIA-Tiny: ~5 M params (matches original v7.1 spec).

    Reference for research comparisons. Not optimized for 500 MB INT8 target.
    `num_heads_*` is set to 4 so `head_dim` (channels // heads) is divisible
    by 4 — required by RoPE 2D in `MambaVisionBlock` (s3: 176/4=44, s4: 224/4=56).
    `mamba_extra_kwargs.headdim=16` because Mamba-2 requires channels % headdim == 0;
    176 and 224 are not divisible by the default 64.
    """
    return SOFIAConfig(
        stem_mid=16,
        stem_out=32,
        embed_dims=[48, 96, 176, 224],
        depths=[2, 3, 4, 2],
        num_heads_s3=4,
        num_heads_s4=4,
        neck_channels=128,
        d_descriptor=512,
        mamba_extra_kwargs={
            "d_state_mamba2": 64,
            "headdim": 16,
            "expand": 2,
            "d_conv": 4,
            "n_directions": 2,
        },
    )


# Default preset
DEFAULT_CONFIG = sofia_m_config


if __name__ == "__main__":
    for name, fn in [("M", sofia_m_config), ("L", sofia_l_config), ("Tiny", sofia_tiny_config)]:
        cfg = fn()
        cfg.validate()
        print(f"SOFIA-{name}: {cfg.summary()}")
