import math
from typing import Optional

import torch
import torch.nn as nn
from torch import Tensor


# -----------------------
# Positional encodings
# -----------------------
def sinusoidal_positional_encoding(
    length: int, dim: int, device: torch.device
) -> Tensor:
    assert dim % 2 == 0, "positional dim must be even"
    pos = torch.arange(length, device=device).float().unsqueeze(1)  # (L, 1)
    half = dim // 2
    div = torch.exp(
        torch.arange(0, half, device=device).float()
        * (-math.log(10000.0) / max(1, half - 1))
    )  # (half,)
    pe = torch.cat([torch.sin(pos * div), torch.cos(pos * div)], dim=1)
    return pe.unsqueeze(0)  # (1, L, dim)


class SinusoidalPosEmb(nn.Module):
    def __init__(self, dim: int):
        super().__init__()
        assert dim % 2 == 0
        self.dim = dim

    def forward(self, t: Tensor) -> Tensor:
        if t.dim() == 2 and t.size(1) == 1:
            t = t[:, 0]
        half = self.dim // 2
        device = t.device
        freq = torch.exp(
            -math.log(10000)
            * torch.arange(half, device=device).float()
            / (half - 1)
        )
        args = t.float().unsqueeze(1) * freq.unsqueeze(0)
        emb = torch.cat([torch.sin(args), torch.cos(args)], dim=-1)
        return emb  # (B, dim)


# -----------------------
# Transformer building blocks
# -----------------------
class TransformerLayer(nn.Module):
    """
    Minimal layer:
      - self-attn on x
      - cross-attn x -> cond (if provided)
      - MLP
    All with pre-norm and residuals.
    """

    def __init__(self, d_model: int, n_heads: int = 8, dropout: float = 0.0):
        super().__init__()
        self.norm1 = nn.LayerNorm(d_model)
        self.self_attn = nn.MultiheadAttention(
            embed_dim=d_model, num_heads=n_heads, dropout=dropout, batch_first=True
        )
        self.drop1 = nn.Dropout(dropout)

        self.norm2 = nn.LayerNorm(d_model)
        self.cross_attn = nn.MultiheadAttention(
            embed_dim=d_model, num_heads=n_heads, dropout=dropout, batch_first=True
        )
        self.drop2 = nn.Dropout(dropout)

        self.norm3 = nn.LayerNorm(d_model)
        self.mlp = nn.Sequential(
            nn.Linear(d_model, d_model * 4),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_model * 4, d_model),
            nn.Dropout(dropout),
        )

    def forward(self, x: Tensor, cond: Optional[Tensor]) -> Tensor:
        # x: (B, L, d), cond: (B, T, d) or None
        y = self.norm1(x)
        y, _ = self.self_attn(y, y, y, need_weights=False)
        x = x + self.drop1(y)

        if cond is not None:
            z = self.norm2(x)
            z, _ = self.cross_attn(z, cond, cond, need_weights=False)
            x = x + self.drop2(z)

        w = self.norm3(x)
        x = x + self.mlp(w)
        return x


# -----------------------
# Just-a-Transformer "UNet1D" (no up/down), 3 layers
# Keeps UNet-like forward signature incl. CFG knobs.
# -----------------------
class UNet1DTransformerSimple(nn.Module):
    def __init__(
        self,
        seq_len: int = 64,
        in_channels: int = 1,
        out_channels: int = 1,
        d_model: int = 256,
        n_heads: int = 8,
        dropout: float = 0.0,
        cond_dim: int = 128,  # embedding is [B, T, 128]
        time_emb_dim: int = 128,
    ):
        super().__init__()
        self.seq_len = seq_len
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.d_model = d_model
        self.cond_dim = cond_dim

        # Input/output projections
        self.in_proj = nn.Linear(in_channels, d_model)
        self.out_proj = nn.Linear(d_model, out_channels)

        # Time embedding -> d_model bias
        self.time_pos_emb = SinusoidalPosEmb(time_emb_dim)
        self.time_mlp = nn.Sequential(
            nn.Linear(time_emb_dim, d_model * 4),
            nn.SiLU(),
            nn.Linear(d_model * 4, d_model),
        )

        # Cond projection to d_model (for cross-attention memory)
        self.cond_proj = nn.Linear(cond_dim, d_model)
        # Learnable null cond for classifier-free guidance
        self.null_cond = nn.Parameter(torch.zeros(cond_dim))

        # Exactly 3 transformer layers
        self.layers = nn.ModuleList(
            [TransformerLayer(d_model, n_heads, dropout) for _ in range(3)]
        )

    # Single pass (no CFG)
    def _forward_once(
        self, x: Tensor, time: Tensor, embedding: Optional[Tensor]
    ) -> Tensor:
        """
        x: (B, 1, L) or (B, L)
        time: (B,) or (B,1)
        embedding: (B, T, cond_dim) or (B, cond_dim) or None
        returns: (B, 1, L)
        """
        if x.dim() == 2:
            x = x.unsqueeze(1)  # (B, 1, L)
        B, C, L = x.shape
        assert C == self.in_channels, "x must be [B, 1, L]"
        assert L == self.seq_len, "sequence length mismatch"

        device = x.device

        # Tokens along time
        x_tok = x.transpose(1, 2)  # (B, L, 1) as tokens
        x_tok = self.in_proj(x_tok)  # (B, L, d)

        # Add positional encodings and time bias
        x_tok = x_tok + sinusoidal_positional_encoding(
            L, self.d_model, device
        )
        t_emb = self.time_pos_emb(time.float())
        t_vec = self.time_mlp(t_emb)  # (B, d)
        x_tok = x_tok + t_vec.unsqueeze(1)  # broadcast over L

        # Prepare cond memory (optional)
        cond_mem: Optional[Tensor] = None
        if embedding is not None:
            if embedding.dim() == 2:
                # (B, D) -> (B, 1, D)
                embedding = embedding.unsqueeze(1)
            # (B, T, cond_dim) -> proj -> add PE
            cond_mem = self.cond_proj(embedding)  # (B, T, d)
            cond_mem = cond_mem + sinusoidal_positional_encoding(
                cond_mem.size(1), self.d_model, device
            )

        # Pass through 3 transformer layers
        for layer in self.layers:
            x_tok = layer(x_tok, cond_mem)

        # Back to (B, C, L)
        out = self.out_proj(x_tok).transpose(1, 2)  # (B, C_out, L)
        return out  # keep (B, 1, L)

    def forward(
        self,
        x: Tensor,
        time: Tensor,
        embedding_mask_proba: float = 0.1,
        embedding: Optional[Tensor] = None,
        embedding_scale: float = 1.0,
    ) -> Tensor:
        """
        Keep UNet-like signature, but this is just a 3-layer transformer.

        x: (B, 1, 64) or (B, 64)
        time: (B,) or (B,1)
        embedding: (B, T, 128) or (B, 128) or None
        embedding_mask_proba: used only during training (CFG masking)
        embedding_scale: classifier-free guidance scale (eval-time)
        returns: (B, 1, 64)
        """
        B = x.shape[0]
        device = x.device

        # Classifier-free guidance masking (train-time)
        if (
            self.training
            and embedding is not None
            and embedding_mask_proba > 0.0
        ):
            mask = torch.rand(B, device=device) < embedding_mask_proba
            if mask.any():
                embedding = embedding.clone()
                if embedding.dim() == 2:
                    # (B, D)
                    embedding[mask] = self.null_cond
                elif embedding.dim() == 3:
                    # (B, T, D)
                    T = embedding.size(1)
                    null_seq = (
                        self.null_cond.view(1, 1, -1)
                        .expand(int(mask.sum().item()), T, -1)
                        .to(device)
                    )
                    embedding[mask] = null_seq
                else:
                    raise ValueError("Unsupported embedding dim for masking")

        # CFG (eval-time two passes)
        if (not self.training) and embedding is not None and (
            embedding_scale != 1.0
        ):
            out_cond = self._forward_once(x, time, embedding)
            # build null embedding matching shape
            if embedding.dim() == 2:
                null_emb = self.null_cond.unsqueeze(0).expand(B, -1)
            elif embedding.dim() == 3:
                T = embedding.size(1)
                null_emb = (
                    self.null_cond.unsqueeze(0)
                    .unsqueeze(1)
                    .expand(B, T, -1)
                )
            else:
                raise ValueError("Unsupported embedding dim for CFG")
            out_null = self._forward_once(x, time, null_emb)
            return out_null + (out_cond - out_null) * embedding_scale

        # Single pass
        return self._forward_once(x, time, embedding)


# # Quick check
# if __name__ == "__main__":
#     torch.manual_seed(0)
#     model = UNet1DTransformerSimple(
#         seq_len=64, in_channels=1, out_channels=1, d_model=256, n_heads=8
#     )
#     x = torch.randn(2, 1, 64)
#     t = torch.randint(0, 1000, (2,))
#     cond = torch.randn(2, 20, 128)
#     y = model(x, t, embedding=cond, embedding_scale=1.2)
#     print(y.shape)  # torch.Size([2, 1, 64])