from __future__ import annotations

import math

import torch
from torch import nn


def fourier_loglam(loglam: torch.Tensor, n_freq: int = 32) -> torch.Tensor:
    # Center roughly around optical wavelengths to improve conditioning.
    x = loglam - math.log(6000.0)
    freqs = torch.logspace(0.0, math.log10(128.0), n_freq, device=loglam.device, dtype=loglam.dtype)
    angles = x.unsqueeze(-1) * freqs
    return torch.cat([torch.sin(angles), torch.cos(angles)], dim=-1)


class NativeSpecZMAE(nn.Module):
    def __init__(
        self,
        feature_dim: int = 6,
        target_aux_dim: int = 4,
        d_model: int = 256,
        nhead: int = 8,
        encoder_layers: int = 4,
        decoder_layers: int = 3,
        dim_feedforward: int | None = None,
        dropout: float = 0.1,
        fourier_freqs: int = 32,
        z_bins: int = 64,
        y_min: float = 0.0,
        y_max: float = math.log1p(6.0),
    ):
        super().__init__()
        self.fourier_freqs = fourier_freqs
        self.y_min = y_min
        self.y_max = y_max
        self.z_bins = z_bins
        pos_dim = fourier_freqs * 2
        ff = dim_feedforward or d_model * 4

        self.input_proj = nn.Sequential(
            nn.Linear(feature_dim + pos_dim, d_model),
            nn.LayerNorm(d_model),
            nn.GELU(),
            nn.Linear(d_model, d_model),
        )
        self.target_proj = nn.Sequential(
            nn.Linear(target_aux_dim + pos_dim, d_model),
            nn.LayerNorm(d_model),
            nn.GELU(),
            nn.Linear(d_model, d_model),
        )
        self.cls = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)
        self.z_query = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)

        enc_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=nhead,
            dim_feedforward=ff,
            dropout=dropout,
            batch_first=True,
            norm_first=True,
            activation="gelu",
        )
        self.encoder = nn.TransformerEncoder(enc_layer, num_layers=encoder_layers)
        dec_layer = nn.TransformerDecoderLayer(
            d_model=d_model,
            nhead=nhead,
            dim_feedforward=ff,
            dropout=dropout,
            batch_first=True,
            norm_first=True,
            activation="gelu",
        )
        self.decoder = nn.TransformerDecoder(dec_layer, num_layers=decoder_layers)
        self.rec_head = nn.Sequential(nn.LayerNorm(d_model), nn.Linear(d_model, d_model // 2), nn.GELU(), nn.Linear(d_model // 2, 1))
        self.z_head = nn.Sequential(nn.LayerNorm(d_model * 2), nn.Linear(d_model * 2, d_model), nn.GELU(), nn.Linear(d_model, 2))
        self.z_bin_head = nn.Sequential(nn.LayerNorm(d_model * 2), nn.Linear(d_model * 2, d_model), nn.GELU(), nn.Linear(d_model, z_bins))

    def forward(
        self,
        enc_features: torch.Tensor,
        enc_loglam: torch.Tensor,
        enc_padding: torch.Tensor,
        target_loglam: torch.Tensor,
        target_aux: torch.Tensor,
    ) -> dict[str, torch.Tensor]:
        bsz = enc_features.shape[0]
        enc_pos = fourier_loglam(enc_loglam, self.fourier_freqs)
        tokens = self.input_proj(torch.cat([enc_features, enc_pos], dim=-1))
        cls = self.cls.expand(bsz, -1, -1)
        z_query = self.z_query.expand(bsz, -1, -1)
        memory_in = torch.cat([cls, z_query, tokens], dim=1)
        special_padding = torch.zeros((bsz, 2), dtype=torch.bool, device=enc_padding.device)
        memory_padding = torch.cat([special_padding, enc_padding], dim=1)
        memory = self.encoder(memory_in, src_key_padding_mask=memory_padding)

        cls_out = memory[:, 0]
        z_out = memory[:, 1]
        z_feat = torch.cat([z_out, cls_out], dim=-1)
        z_params = self.z_head(z_feat)
        z_bins = self.z_bin_head(z_feat)

        tgt_pos = fourier_loglam(target_loglam, self.fourier_freqs)
        tgt = self.target_proj(torch.cat([target_aux, tgt_pos], dim=-1))
        decoded = self.decoder(tgt=tgt, memory=memory, memory_key_padding_mask=memory_padding)
        rec = self.rec_head(decoded).squeeze(-1)
        return {
            "rec": rec,
            "y_mu": z_params[:, 0],
            "y_logvar": torch.clamp(z_params[:, 1], -8.0, 4.0),
            "z_bin_logits": z_bins,
        }

    def y_to_bin(self, y: torch.Tensor) -> torch.Tensor:
        scaled = (y - self.y_min) / max(self.y_max - self.y_min, 1e-6)
        return torch.clamp((scaled * self.z_bins).long(), 0, self.z_bins - 1)