"""SOFIA v1 model: StripNet+DCNv4 backbone + asymmetric Sat/UAV heads.

Architecture:
    img [B,3,256,256] --> StripNetDCN (4 stages) --> [B, C_4, 8, 8]
                                                        |
                                  ┌─────────────────────┴────────────────┐
                                  ▼                                      ▼
                              SatHeadV1                            UAVHeadV1
                              [TextFiLM]                          AltitudeFiLM
                              GGeM                                [TextFiLM]
                              BN + Linear                         GGeM
                              [L2]                                BN + Linear
                                                                  [L2]

Heads return un-normalized D-dim descriptors (for downstream gated text fusion).
"""

from __future__ import annotations

from typing import Dict, List, Optional

import torch
import torch.nn as nn

from .config import SOFIAv1Config
from .heads import SatHeadV1, UAVHeadV1
from .stripnet_model_dcn import VARIANT_MAP, StripNetDCN


class SOFIAv1(nn.Module):
    """SOFIA v1: shared StripNet+DCNv4 backbone with asymmetric heads."""

    def __init__(self, cfg: SOFIAv1Config) -> None:
        super().__init__()
        if cfg.variant not in VARIANT_MAP:
            raise ValueError(f"Unknown variant {cfg.variant!r}")
        self.cfg = cfg

        # Single shared backbone for sat + uav (saves params; v1 stays lightweight).
        self.backbone: StripNetDCN = VARIANT_MAP[cfg.variant](dcn_variant=cfg.dcn_variant)
        last_channels = self.backbone.embed_dims[-1]
        self.feature_channels = last_channels

        self.sat_head = SatHeadV1(
            in_channels=last_channels,
            d_descriptor=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 = UAVHeadV1(
            in_channels=last_channels,
            d_descriptor=cfg.d_descriptor,
            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,
        )

    def _extract_last(self, x: torch.Tensor) -> torch.Tensor:
        """Run backbone, return only the deepest feature map [B, C, 8, 8]."""
        feats: List[torch.Tensor] = self.backbone(x)
        return feats[-1]

    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]:
        result: Dict[str, torch.Tensor] = {}

        if sat is not None:
            f_sat = self._extract_last(sat)
            g_sat = self.sat_head(f_sat, text_emb=text_emb_sat)
            result["g_sat"] = g_sat
            if return_features:
                result["features_sat"] = f_sat

        if uav is not None:
            f_uav = self._extract_last(uav)
            g_uav = self.uav_head(f_uav, altitude=altitude, text_emb=text_emb_uav)
            result["g_uav"] = g_uav
            if return_features:
                result["features_uav"] = f_uav

        return result
