"""
OAT (Ordered-Action Transformer) Model

Provides a unified model wrapper for OAT encoder and decoder.
"""

from __future__ import annotations

from typing import List, Optional, Tuple, Union

import torch
import torch.nn as nn

from .encoder import RegisterEncoder
from .decoder import SinglePassDecoder


class OATModel(nn.Module):
    """
    OAT Model wrapper that combines encoder and decoder.

    This provides a unified interface similar to VQVAE/UAT models,
    allowing the tokenizer to access .model attribute.
    """

    def __init__(
        self,
        encoder: RegisterEncoder,
        decoder: SinglePassDecoder,
        fsq_quantizer: Optional[nn.Module] = None,
        latent_dim: int = 64,
        num_registers: int = 16,
        codebook_size: int = 4096,
        fsq_levels: Optional[List[int]] = None,
    ):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.fsq_quantizer = fsq_quantizer
        self.latent_dim = latent_dim
        self.num_registers = num_registers
        self.codebook_size = codebook_size
        self.fsq_levels = fsq_levels or [8, 8, 8, 8]

        # Initialize autoregressive generator (lazy initialization)
        self._generator = None

    def forward(
        self,
        actions: Optional[torch.Tensor] = None,
        latents: Optional[torch.Tensor] = None,
        eval_keep_k: Optional[list[int]] = None,
        use_fsq_quantizer: bool = True,
        return_tokens: bool = False,
    ) -> Union[
        Tuple[torch.Tensor, Optional[torch.Tensor]],
        Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]],
    ]:
        """
        Forward pass supporting both encoding and decoding.

        Args:
            actions: Input actions for encoding (batch, horizon, dim)
            latents: Input latents for decoding (batch, num_registers, latent_dim)
            eval_keep_k: Optional list for nested dropout during eval
            return_tokens: If True, also return discrete tokens (only when actions provided)

        Returns:
            - If actions provided and return_tokens=False: (reconstructed_actions, latent_tokens)
            - If actions provided and return_tokens=True: (reconstructed_actions, latent_tokens, discrete_tokens)
            - If only latents provided: (reconstructed_actions, None)
        """
        latent_tokens = None
        discrete_tokens = None
        reconstructed_actions = None
        
        if use_fsq_quantizer and self.fsq_quantizer is None:
            raise ValueError("FSQ quantizer is enabled but not provided in the fsq quantizer model.")
        use_fsq_quantizer = use_fsq_quantizer and self.fsq_quantizer is not None

        if actions is not None:
            # Encode then decode (full forward pass for training)
            # encode() already applies FSQ quantization by default (use_quantizer=True)
            if return_tokens:
                latents, discrete_tokens = self.encode(
                    actions, return_tokens=True, use_quantizer=use_fsq_quantizer
                )
            else:
                latents = self.encode(
                    actions, return_tokens=False, use_quantizer=use_fsq_quantizer
                )
                # make sure it's quantized latents.

            reconstructed_actions = self.decode(latents, eval_keep_k=eval_keep_k)

            if return_tokens and discrete_tokens is not None:
                return reconstructed_actions, latents, discrete_tokens
            return reconstructed_actions, latents

        if latents is not None:
            # Decode only
            reconstructed_actions = self.decode(latents, eval_keep_k=eval_keep_k)
            return reconstructed_actions, None

        return None, None

    def encode(
        self,
        actions: torch.Tensor,
        return_tokens: bool = True,
        use_quantizer: bool = True,
    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
        """
        Encode actions to latent tokens.

        Args:
            actions: Input actions (batch, horizon, action_dim)
            return_tokens: If True, also return discrete tokens from FSQ quantizer
            use_quantizer: If True, apply FSQ quantization and return quantized latents
                          (matching original OAT behavior)

        Returns:
            - If return_tokens=False: Latent tokens (batch, num_registers, latent_dim)
            - If return_tokens=True: (quantized_latents, discrete_tokens)
              - quantized_latents: (batch, num_registers, latent_dim) - after FSQ quantization
              - discrete_tokens: (batch, num_registers)
        """
        latents = self.encoder(actions)

        # Apply FSQ quantization (matching original OAT implementation)
        if use_quantizer and self.fsq_quantizer is not None:
            quantized_latents, tokens = self.fsq_quantizer(latents)
            if return_tokens:
                return quantized_latents, tokens
            return quantized_latents

        return latents

    def decode(
        self,
        latents: torch.Tensor,
        eval_keep_k: Optional[list[int]] = None,
    ) -> torch.Tensor:
        """
        Decode latent tokens to actions.

        Args:
            latents: Latent tokens (batch, num_registers, latent_dim)
            eval_keep_k: Optional list for nested dropout

        Returns:
            Reconstructed actions (batch, horizon, action_dim)
        """
        return self.decoder(latents, eval_keep_k=eval_keep_k)

    def train(self, mode: bool = True) -> "OATModel":
        """Set training mode"""
        super().train(mode)
        return self

    def eval(self) -> "OATModel":
        """Set evaluation mode"""
        super().eval()
        return self
