"""Heads for SOFIA v1 (StripNet+DCNv4 backbone).

Designed parallel to SOFIA v7.1 heads but lighter — basic GGeM pooling
instead of CHP, optional altitude-FiLM (UAV) and text-FiLM (both).

The heads produce un-normalized D-dim descriptors when `return_normalized=False`,
so that a fusion wrapper can blend with text via gated fusion before the final
L2 normalization.
"""

from __future__ import annotations

from typing import Optional

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

from src.models.sofia_v71.layers import AltitudeFiLM, GGeM, TextFiLM


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

    def __init__(
        self,
        in_channels: int,
        d_descriptor: int,
        return_normalized: bool = False,
        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(in_channels, text_dim=text_film_dim, hidden_dim=text_film_hidden)
        self.ggem = GGeM(in_channels)
        self.bn = nn.BatchNorm1d(in_channels, affine=False)
        self.proj = nn.Linear(in_channels, d_descriptor)

    def forward(
        self,
        x: torch.Tensor,
        text_emb: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        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 UAVHeadV1(nn.Module):
    """UAV head: AltitudeFiLM [+ TextFiLM] + GGeM + BN + Linear [+ L2]."""

    def __init__(
        self,
        in_channels: int,
        d_descriptor: int,
        use_film: bool = True,
        altitude_norm: float = 500.0,
        return_normalized: bool = False,
        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(in_channels, altitude_norm=altitude_norm)
        if use_text_film:
            self.text_film = TextFiLM(in_channels, text_dim=text_film_dim, hidden_dim=text_film_hidden)
        self.ggem = GGeM(in_channels)
        self.bn = nn.BatchNorm1d(in_channels, affine=False)
        self.proj = nn.Linear(in_channels, d_descriptor)

    def forward(
        self,
        x: torch.Tensor,
        altitude: Optional[torch.Tensor] = None,
        text_emb: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        if self.use_film:
            x = self.film(x, altitude)
        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
