from __future__ import annotations

from functools import partial
from typing import List, Optional

import torch
import torch.nn as nn

from .oat_linear import LinearHead, LinearLayer
from .oat_transformer import Fp32LayerNorm
from .positional import PositionalEmbeddingAdder, PositionalQuery
from .token_dropout import MaskedNestedDropout


class SinglePassDecoder(nn.Module):
    def __init__(
        self,
        sample_dim: int,
        sample_horizon: int,
        emb_dim: int,
        head_dim: int,
        depth: int,
        pdropout: float,
        token_dropout_mode: str,
        use_causal_decoder: bool,
        latent_dim: int,
        latent_horizon: int,
    ):
        super().__init__()

        if emb_dim % head_dim != 0:
            raise ValueError(
                f"emb_dim ({emb_dim}) must be divisible by head_dim ({head_dim})"
            )

        self.query = PositionalQuery(emb_dim, sample_horizon)
        self.latent_pos = PositionalEmbeddingAdder(emb_dim, latent_horizon)
        self.nested_dropout = MaskedNestedDropout(
            emb_dim, size_sampling_mode=token_dropout_mode
        )

        layer = nn.TransformerDecoderLayer(
            d_model=emb_dim,
            nhead=emb_dim // head_dim,
            dim_feedforward=4 * emb_dim,
            dropout=pdropout,
            activation="gelu",
            batch_first=True,
            norm_first=True,
        )
        self.decoder = nn.TransformerDecoder(layer, num_layers=depth)
        self.latent_proj = LinearLayer(latent_dim, emb_dim, weight_init_style="xavier")
        self.head = LinearHead(
            emb_dim,
            sample_dim,
            weight_init_style="xavier",  # Changed from "zero" to allow gradient flow
            norm_layer=partial(Fp32LayerNorm, bias=False, elementwise_affine=False),
        )

        self.sample_horizon = sample_horizon
        self.latent_horizon = latent_horizon
        self.use_causal_decoder = use_causal_decoder

    def forward(
        self, latents: torch.Tensor, eval_keep_k: Optional[List[int]] = None
    ) -> torch.Tensor:
        b = latents.shape[0]
        dtype = latents.dtype
        device = latents.device

        tgt = self.query(b, self.sample_horizon, device=device, dtype=dtype)

        memory = self.latent_proj(latents)
        memory = self.latent_pos(memory)
        memory = self.nested_dropout(memory, eval_keep_k=eval_keep_k)

        if self.use_causal_decoder:
            causal_mask = nn.Transformer.generate_square_subsequent_mask(
                self.sample_horizon, device=device
            )
            out = self.decoder(
                tgt=tgt, memory=memory, tgt_mask=causal_mask, tgt_is_causal=True
            )
        else:
            out = self.decoder(tgt=tgt, memory=memory)

        return self.head(out)
