"""SOFIA v7.1 custom layers.

Includes:
- GGeM: Generalized Mean Pooling with per-channel learnable exponent (F11)
- CircularHarmonicPool: Formally SO(2)-invariant pooling via polar + FFT magnitude (NH2 novel)
- AltitudeFiLM: FiLM conditioning on UAV altitude (NH4 novel)
- RoPE2D: 2D Rotary Position Embedding for attention
- SqueezeExcite: standard SE block
- LayerNorm2d: channel-last LN wrapper for 2D features

All rotation-invariance and FiLM modules are NOVEL contributions of SOFIA v7.1-α.
"""

from __future__ import annotations

import math
from typing import Optional, Tuple

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


# ============================================================
# GGeM: Generalized Mean Pooling (F11)
# ============================================================

class GGeM(nn.Module):
    """Per-channel learnable Generalized Mean pooling.

    Formula:
        GGeM(F)_c = (1/HW · Σ F_{c,h,w}^{p_c})^{1/p_c}
        p_c = softplus(p_hat_c) ∈ (0, ∞)
    """

    def __init__(self, channels: int, init_p: float = 3.0, eps: float = 1e-6) -> None:
        super().__init__()
        # softplus^{-1}(init_p) = log(exp(init_p) - 1)
        hat_init = math.log(math.exp(init_p) - 1.0)
        self.hat_p = nn.Parameter(torch.full((channels,), hat_init))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """x: [B, C, H, W] -> [B, C]"""
        p = F.softplus(self.hat_p).view(1, -1, 1, 1)    # [1, C, 1, 1]
        x_clamped = x.clamp(min=self.eps)
        x_pow = x_clamped.pow(p)
        x_mean = x_pow.mean(dim=(2, 3), keepdim=True)
        out = x_mean.pow(1.0 / p)
        return out.flatten(1)


# ============================================================
# CHP: Circular Harmonic Pool (NH2 novel — formally SO(2)-invariant)
# ============================================================

class CircularHarmonicPool(nn.Module):
    """Formally SO(2) rotation-invariant pooling.

    Algorithm:
        1. Sample input feature map at polar grid (r, θ) via bilinear grid_sample
        2. Apply 1D real FFT along θ-axis
        3. Keep magnitudes of first N harmonics (invariant to shift = rotation)
        4. GGeM pool over rings r (per-channel-per-harmonic)
        5. Flatten to descriptor [B, C * N]

    Output is theoretically invariant to input rotation of any angle.

    See HYP Phase 4'' Section 4''.1 NH2 and Section 4''.5 for formal proof.
    """

    def __init__(
        self,
        channels: int,
        rings: int = 8,
        angles: int = 16,
        harmonics: int = 4,
        r_min: float = 0.1,
        r_max: float = 1.0,
    ) -> None:
        super().__init__()
        assert harmonics <= angles // 2 + 1, (
            f"harmonics {harmonics} cannot exceed angles//2+1 = {angles // 2 + 1}"
        )

        self.channels = channels
        self.rings = rings
        self.angles = angles
        self.harmonics = harmonics

        # GGeM over rings (per channel × per harmonic)
        self.ggem = GGeM(channels * harmonics, init_p=3.0)

        # Precompute polar grid in normalized [-1, 1] coords for grid_sample
        grid = self._make_polar_grid(rings, angles, r_min, r_max)    # [R, T, 2]
        self.register_buffer("polar_grid", grid, persistent=False)

    @staticmethod
    def _make_polar_grid(R: int, T: int, r_min: float, r_max: float) -> torch.Tensor:
        r_values = torch.linspace(r_min, r_max, R)    # [R]
        theta_values = torch.linspace(0.0, 2 * math.pi, T + 1)[:-1]    # [T]
        r_grid, theta_grid = torch.meshgrid(r_values, theta_values, indexing="ij")    # [R, T]
        x = r_grid * torch.cos(theta_grid)
        y = r_grid * torch.sin(theta_grid)
        grid = torch.stack([x, y], dim=-1)    # [R, T, 2]
        return grid

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: [B, C, H, W] feature map
        Returns:
            descriptor: [B, C * harmonics]
        """
        B, C, H, W = x.shape
        assert C == self.channels, (
            f"Input channels {C} != expected {self.channels}"
        )

        # 1. Polar sampling
        # grid: [B, R, T, 2]
        grid = self.polar_grid.unsqueeze(0).expand(B, -1, -1, -1)
        polar = F.grid_sample(
            x, grid,
            mode="bilinear",
            padding_mode="zeros",
            align_corners=True,
        )    # [B, C, R, T]

        # 2. 1D real FFT along angular axis
        polar_fft = torch.fft.rfft(polar, dim=-1)    # [B, C, R, T//2+1]
        polar_fft = polar_fft[..., : self.harmonics]    # [B, C, R, N]

        # 3. Magnitude (rotation invariant)
        magnitude = polar_fft.abs()    # [B, C, R, N]

        # 4. Reshape for GGeM: treat (C, N) as combined channel dim, rings as spatial
        # Shape: [B, C*N, R, 1]
        magnitude_reshaped = (
            magnitude
            .permute(0, 1, 3, 2)    # [B, C, N, R]
            .reshape(B, C * self.harmonics, self.rings, 1)
        )

        # 5. GGeM pool over rings (H=R, W=1)
        descriptor = self.ggem(magnitude_reshaped)    # [B, C*N]

        return descriptor


# ============================================================
# FiLM: Altitude-conditioned modulation (NH4 novel)
# ============================================================

class AltitudeFiLM(nn.Module):
    """FiLM modulation conditioned on scalar altitude.

    F' = γ(h) · F + β(h)  where γ,β ∈ R^C are produced by MLP.

    At altitude=None, produces identity (γ=1, β=0) via zero-init of final layer.
    """

    def __init__(
        self,
        channels: int,
        hidden_dim: int = 64,
        altitude_norm: float = 500.0,
    ) -> None:
        super().__init__()
        self.channels = channels
        self.altitude_norm = altitude_norm

        self.mlp = nn.Sequential(
            nn.Linear(1, hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, 2 * channels),
        )
        # Zero-init last layer → initial γ=0 before residual, β=0
        nn.init.zeros_(self.mlp[-1].weight)
        nn.init.zeros_(self.mlp[-1].bias)

    def forward(self, x: torch.Tensor, altitude: Optional[torch.Tensor] = None) -> torch.Tensor:
        """
        Args:
            x: [B, C, H, W]
            altitude: [B] or [B, 1] scalar altitude in meters (or None for neutral)
        Returns:
            [B, C, H, W]
        """
        B = x.shape[0]
        if altitude is None:
            altitude = torch.zeros(B, 1, device=x.device, dtype=x.dtype)
        elif altitude.dim() == 1:
            altitude = altitude.unsqueeze(-1)

        h_norm = altitude.to(x.dtype) / self.altitude_norm
        gamma_beta = self.mlp(h_norm)    # [B, 2C]
        gamma, beta = gamma_beta.chunk(2, dim=-1)
        # Residual form: γ = 1 + delta_γ (starts at identity)
        gamma = gamma.view(B, self.channels, 1, 1) + 1.0
        beta = beta.view(B, self.channels, 1, 1)
        return gamma * x + beta


# ============================================================
# TextFiLM: text-conditioned modulation (extension for caption fusion)
# ============================================================

class TextFiLM(nn.Module):
    """FiLM modulation conditioned on a text embedding.

    F' = γ(z) · F + β(z) where γ,β ∈ R^C are produced by an MLP
    from z ∈ R^{D_txt}. Identity at init via zero-init of last layer
    (so γ=1, β=0 before training shifts the residual).
    """

    def __init__(
        self,
        channels: int,
        text_dim: int = 1024,
        hidden_dim: int = 256,
    ) -> None:
        super().__init__()
        self.channels = channels
        self.text_dim = text_dim
        self.mlp = nn.Sequential(
            nn.Linear(text_dim, hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, 2 * channels),
        )
        nn.init.zeros_(self.mlp[-1].weight)
        nn.init.zeros_(self.mlp[-1].bias)

    def forward(self, x: torch.Tensor, text_emb: Optional[torch.Tensor] = None) -> torch.Tensor:
        """
        Args:
            x: [B, C, H, W]
            text_emb: [B, D_txt] or None for no-op (identity)
        Returns:
            [B, C, H, W]
        """
        if text_emb is None:
            return x
        B = x.shape[0]
        gamma_beta = self.mlp(text_emb.to(x.dtype))    # [B, 2C]
        gamma, beta = gamma_beta.chunk(2, dim=-1)
        gamma = gamma.view(B, self.channels, 1, 1) + 1.0
        beta = beta.view(B, self.channels, 1, 1)
        return gamma * x + beta


# ============================================================
# RoPE 2D
# ============================================================

class RoPE2D(nn.Module):
    """2D Rotary Position Embedding.

    Splits head_dim into two halves: first half gets x-position encoding,
    second half gets y-position encoding. For each half, applies standard
    1D RoPE rotation.

    Reference: RoFormer (B49) adapted for 2D.
    """

    def __init__(self, head_dim: int, max_resolution: int = 64, base: float = 10000.0) -> None:
        super().__init__()
        assert head_dim % 2 == 0, "head_dim must be even for RoPE"
        assert head_dim % 4 == 0, "head_dim must be divisible by 4 for 2D RoPE"
        self.head_dim = head_dim
        self.half_dim = head_dim // 2    # dedicated to each axis
        self.max_resolution = max_resolution

        # Frequencies for each axis (half_dim per axis, sin+cos pairs)
        freqs = 1.0 / (base ** (torch.arange(0, self.half_dim, 2).float() / self.half_dim))
        self.register_buffer("freqs", freqs, persistent=False)

    def _make_embeds(self, H: int, W: int, device: torch.device) -> Tuple[torch.Tensor, torch.Tensor]:
        """Produce cos/sin embeddings for HW tokens in raster order."""
        y_pos = torch.arange(H, device=device, dtype=torch.float32)
        x_pos = torch.arange(W, device=device, dtype=torch.float32)

        freqs_y = torch.einsum("i,j->ij", y_pos, self.freqs)    # [H, half_dim/2]
        freqs_x = torch.einsum("i,j->ij", x_pos, self.freqs)    # [W, half_dim/2]

        # Expand to full grid: [H, W, half_dim/2] each
        freqs_y = freqs_y.unsqueeze(1).expand(-1, W, -1)    # [H, W, half_dim/2]
        freqs_x = freqs_x.unsqueeze(0).expand(H, -1, -1)    # [H, W, half_dim/2]

        # Concatenate: x-axis freqs into first half, y-axis into second half
        # Each half is [H, W, half_dim/2]; we pair-up for complex rotation
        freqs_combined_x = torch.cat([freqs_x, freqs_x], dim=-1)    # [H, W, half_dim]
        freqs_combined_y = torch.cat([freqs_y, freqs_y], dim=-1)    # [H, W, half_dim]

        freqs_full = torch.cat([freqs_combined_x, freqs_combined_y], dim=-1)    # [H, W, head_dim]
        cos = freqs_full.cos().reshape(H * W, -1)
        sin = freqs_full.sin().reshape(H * W, -1)
        return cos, sin

    @staticmethod
    def _rotate_half(x: torch.Tensor) -> torch.Tensor:
        """Rotate: (x1, x2) -> (-x2, x1)."""
        x1, x2 = x.chunk(2, dim=-1)
        return torch.cat([-x2, x1], dim=-1)

    def forward(self, q: torch.Tensor, k: torch.Tensor, H: int, W: int) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            q, k: [B, heads, HW, head_dim]
            H, W: spatial dims for position computation
        Returns:
            q_rot, k_rot with positional encoding applied
        """
        cos, sin = self._make_embeds(H, W, q.device)
        cos = cos.to(q.dtype).unsqueeze(0).unsqueeze(0)    # [1, 1, HW, head_dim]
        sin = sin.to(q.dtype).unsqueeze(0).unsqueeze(0)
        q_rot = (q * cos) + (self._rotate_half(q) * sin)
        k_rot = (k * cos) + (self._rotate_half(k) * sin)
        return q_rot, k_rot


# ============================================================
# SE: Squeeze-Excite
# ============================================================

class SqueezeExcite(nn.Module):
    """Standard Squeeze-Excite channel attention."""

    def __init__(self, channels: int, reduction: int = 16) -> None:
        super().__init__()
        hidden = max(1, channels // reduction)
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
            nn.Conv2d(channels, hidden, 1),
            nn.SiLU(inplace=True),
            nn.Conv2d(hidden, channels, 1),
            nn.Sigmoid(),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        s = self.pool(x)
        s = self.fc(s)
        return x * s


# ============================================================
# LayerNorm2d: LN over channels for (B, C, H, W) layout
# ============================================================

class LayerNorm2d(nn.Module):
    """LayerNorm over C dimension for 4D tensors."""

    def __init__(self, channels: int, eps: float = 1e-6) -> None:
        super().__init__()
        self.norm = nn.LayerNorm(channels, eps=eps)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # [B, C, H, W] → [B, H, W, C] → LN → [B, C, H, W]
        x = x.permute(0, 2, 3, 1)
        x = self.norm(x)
        x = x.permute(0, 3, 1, 2).contiguous()
        return x


if __name__ == "__main__":
    # Smoke test
    torch.manual_seed(0)

    # GGeM test
    g = GGeM(64)
    x = torch.randn(2, 64, 8, 8)
    out = g(x)
    print(f"GGeM out: {out.shape}")    # [2, 64]

    # CHP test — verify rotation invariance
    chp = CircularHarmonicPool(32, rings=8, angles=16, harmonics=4)
    x = torch.randn(1, 32, 16, 16)
    out1 = chp(x)
    # Rotate x by 90° and verify invariance (approximately)
    x_rot = torch.rot90(x, k=1, dims=(-2, -1))
    out2 = chp(x_rot)
    diff = (out1 - out2).abs().max().item()
    print(f"CHP rotation-invariance max diff: {diff:.4e} (should be small)")
    print(f"CHP out shape: {out1.shape}")    # [1, 128]

    # FiLM test
    film = AltitudeFiLM(64)
    x = torch.randn(2, 64, 8, 8)
    altitudes = torch.tensor([150.0, 300.0])
    out = film(x, altitudes)
    print(f"FiLM out: {out.shape}")    # [2, 64, 8, 8]

    # RoPE2D test
    rope = RoPE2D(32)
    q = torch.randn(2, 4, 64, 32)    # [B, heads, HW, head_dim]
    k = torch.randn(2, 4, 64, 32)
    q_r, k_r = rope(q, k, 8, 8)
    print(f"RoPE out: q={q_r.shape}, k={k_r.shape}")