"""Conditional Flow Matching model with Transformer backbone.

Predicts velocity field v_θ(x_t, t, cond) that transports noise to data.
- x_t: (B, T, 3) noisy trajectory at time t
- t: (B,) interpolation time ∈ [0,1]
- cond: (B, 3) condition [dist_norm, log_dist, log_dur]
- Output: (B, T, 3) velocity field
"""
from __future__ import annotations

import math

import torch
import torch.nn as nn
from torch.distributions import Normal


# ---------------------------------------------------------------------------
# Sinusoidal time embedding
# ---------------------------------------------------------------------------


class SinusoidalTimeEmbedding(nn.Module):
    """Map scalar timestep t ∈ [0,1] to a d_model-dimensional vector."""

    def __init__(self, d_model: int):
        super().__init__()
        self.d_model = d_model

    def forward(self, t: torch.Tensor) -> torch.Tensor:
        """
        Args:
            t: (B,) scalar timesteps

        Returns:
            (B, d_model) embedding vectors
        """
        half = self.d_model // 2
        freqs = torch.exp(
            -math.log(10000.0) * torch.arange(half, device=t.device, dtype=t.dtype) / half
        )
        args = t.unsqueeze(-1) * freqs.unsqueeze(0)  # (B, half)
        return torch.cat([torch.sin(args), torch.cos(args)], dim=-1)  # (B, d_model)


# ---------------------------------------------------------------------------
# TrajectoryFlowModel
# ---------------------------------------------------------------------------


class TrajectoryFlowModel(nn.Module):
    """Conditional Flow Matching model for mouse trajectory generation.

    Architecture:
    - SinusoidalTimeEmbedding: t scalar → d_model vector
    - Input projection: 3 → d_model
    - Learned positional embedding: (1, seq_len, d_model)
    - Time embed + Condition embed added as bias to all tokens
    - 4-layer TransformerEncoder (pre-norm, GELU, batch_first=True)
    - Output projection: d_model → 3

    Args:
        seq_len: Number of time steps (default 64).
        d_model: Transformer hidden dimension (default 128).
        nhead: Number of attention heads (default 4).
        num_layers: Number of transformer layers (default 4).
        dim_feedforward: Feedforward hidden size (default 256).
        dropout: Dropout rate (default 0.1).
        cond_dim: Condition vector size (default 3).
    """

    def __init__(
        self,
        seq_len: int = 64,
        d_model: int = 128,
        nhead: int = 4,
        num_layers: int = 4,
        dim_feedforward: int = 256,
        dropout: float = 0.1,
        cond_dim: int = 3,
    ):
        super().__init__()
        self.seq_len = seq_len
        self.d_model = d_model
        self.cond_dim = cond_dim

        # Input projection: (forward, lateral, log_dt) → d_model
        self.input_proj = nn.Linear(3, d_model)

        # Learned positional embedding
        self.pos_embed = nn.Parameter(torch.randn(1, seq_len, d_model) * 0.02)

        # Time embedding
        self.time_embed = SinusoidalTimeEmbedding(d_model)
        self.time_mlp = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
            nn.Linear(d_model, d_model),
        )

        # Condition embedding
        self.cond_mlp = nn.Sequential(
            nn.Linear(cond_dim, d_model),
            nn.GELU(),
            nn.Linear(d_model, d_model),
        )

        # Transformer encoder (pre-norm via norm_first=True)
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=nhead,
            dim_feedforward=dim_feedforward,
            dropout=dropout,
            activation="gelu",
            batch_first=True,
            norm_first=True,
        )
        self.transformer = nn.TransformerEncoder(
            encoder_layer, num_layers=num_layers, enable_nested_tensor=False
        )

        # Output projection: d_model → 3
        self.output_proj = nn.Linear(d_model, 3)

    def forward(
        self,
        x_t: torch.Tensor,
        t: torch.Tensor,
        cond: torch.Tensor,
    ) -> torch.Tensor:
        """Predict velocity field.

        Args:
            x_t: (B, T, 3) noisy trajectory at interpolation time t
            t: (B,) interpolation timestep ∈ [0,1]
            cond: (B, 3) condition vector [dist_norm, log_dist, log_dur]

        Returns:
            (B, T, 3) predicted velocity field
        """
        # Project input tokens
        h = self.input_proj(x_t)  # (B, T, d_model)

        # Add positional embedding
        h = h + self.pos_embed  # (B, T, d_model)

        # Time embedding → bias added to all tokens
        t_emb = self.time_mlp(self.time_embed(t))  # (B, d_model)
        h = h + t_emb.unsqueeze(1)  # broadcast over T

        # Condition embedding → bias added to all tokens
        c_emb = self.cond_mlp(cond)  # (B, d_model)
        h = h + c_emb.unsqueeze(1)  # broadcast over T

        # Transformer
        h = self.transformer(h)  # (B, T, d_model)

        # Output projection
        return self.output_proj(h)  # (B, T, 3)


# ---------------------------------------------------------------------------
# Legacy JointCVAE — kept for backward compatibility with generator.py
# ---------------------------------------------------------------------------


class JointCVAE(nn.Module):
    """Joint Conditional VAE for mouse trajectory generation (legacy).

    Kept for backward compatibility with the existing generator.
    See TrajectoryFlowModel for the new approach.
    """

    def __init__(
        self,
        seq_len: int = 64,
        latent_dim: int = 32,
        hidden: int = 128,
        cond_dim: int = 3,
    ):
        super().__init__()
        self.seq_len = seq_len
        self.latent_dim = latent_dim
        self.hidden = hidden
        self.cond_dim = cond_dim
        self.feat_dim = 3

        self.enc_gru = nn.GRU(
            input_size=self.feat_dim + cond_dim,
            hidden_size=hidden,
            num_layers=2,
            batch_first=True,
            bidirectional=True,
        )
        self.enc_mu = nn.Linear(hidden * 2, latent_dim)
        self.enc_logvar = nn.Linear(hidden * 2, latent_dim)

        self.dec_h0 = nn.Linear(latent_dim + cond_dim, hidden * 2)
        self.dec_gru = nn.GRU(
            input_size=latent_dim + cond_dim,
            hidden_size=hidden,
            num_layers=2,
            batch_first=True,
        )
        self.dec_out = nn.Linear(hidden, self.feat_dim)

    def encode(self, seq: torch.Tensor, cond: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        B, T, _ = seq.shape
        c_exp = cond.unsqueeze(1).expand(B, T, self.cond_dim)
        x_in = torch.cat([seq, c_exp], dim=-1)
        _, h_n = self.enc_gru(x_in)
        h_fwd = h_n[-2]
        h_bwd = h_n[-1]
        h_cat = torch.cat([h_fwd, h_bwd], dim=-1)
        return self.enc_mu(h_cat), self.enc_logvar(h_cat)

    def decode(self, z: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
        B = z.shape[0]
        zc = torch.cat([z, cond], dim=-1)
        h0_flat = self.dec_h0(zc)
        h0 = h0_flat.view(B, 2, self.hidden).permute(1, 0, 2).contiguous()
        inp = zc.unsqueeze(1).expand(B, self.seq_len, -1)
        out, _ = self.dec_gru(inp, h0)
        return self.dec_out(out)

    def reparameterise(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
        std = torch.exp(0.5 * logvar)
        return Normal(mu, std).rsample()

    def forward(
        self, seq: torch.Tensor, cond: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        mu, logvar = self.encode(seq, cond)
        z = self.reparameterise(mu, logvar)
        recon = self.decode(z, cond)
        return recon, mu, logvar