"""
Autoregressive Generation Module for OAT Tokenizer

Provides KV-cache accelerated autoregressive generation with:
- Prefix-based generation (continue from partial token sequences)
- Coarse-to-fine hierarchical generation (using FSQ structure)
- Temperature and top-k sampling
"""

from __future__ import annotations

from dataclasses import dataclass
from enum import Enum
from typing import List, Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F


class OutputGranularity(Enum):
    """Output granularity for hierarchical FSQ token generation."""
    COARSE = "coarse"      # Use only high-level FSQ dimensions (coarse representation)
    FINE = "fine"          # Use all FSQ dimensions (full fine-grained representation)
    ADAPTIVE = "adaptive"  # Adaptively select dimensions based on token position


@dataclass
class GenerationConfig:
    """Configuration for autoregressive generation."""
    
    max_new_tokens: int = 16
    temperature: float = 1.0
    top_k: Optional[int] = None
    top_p: Optional[float] = None
    eos_token_id: Optional[int] = None
    pad_token_id: Optional[int] = 0
    
    # Prefix generation
    use_prefix: bool = False
    prefix_tokens: Optional[torch.Tensor] = None
    
    # Hierarchical FSQ generation
    granularity: OutputGranularity = OutputGranularity.FINE
    fsq_coarse_dims: int = 2  # Number of coarse dimensions to use (for COARSE mode)
    
    # Guidance
    guidance_scale: float = 1.0
    condition_embedding: Optional[torch.Tensor] = None
    
    # Early stopping
    min_new_tokens: int = 0
    max_time_steps: Optional[int] = None


class KVCache(nn.Module):
    """
    KV Cache for efficient autoregressive generation.
    
    Stores key-value pairs from previous forward passes to avoid recomputation.
    """
    
    def __init__(
        self,
        num_layers: int,
        num_heads: int,
        head_dim: int,
        max_seq_len: int,
        batch_size: int,
        device: torch.device,
        dtype: torch.dtype = torch.float32,
    ):
        super().__init__()
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.max_seq_len = max_seq_len
        
        # Initialize KV cache for each layer
        # Shape: (num_layers, batch_size, num_heads, max_seq_len, head_dim)
        self.register_buffer(
            "key_cache",
            torch.zeros(
                num_layers, batch_size, num_heads, max_seq_len, head_dim,
                device=device, dtype=dtype
            ),
            persistent=False,
        )
        self.register_buffer(
            "value_cache",
            torch.zeros(
                num_layers, batch_size, num_heads, max_seq_len, head_dim,
                device=device, dtype=dtype
            ),
            persistent=False,
        )
        
        # Track current sequence length for each layer
        self.seq_lens = torch.zeros(num_layers, dtype=torch.long, device=device)
        
    def update(
        self,
        layer_idx: int,
        key: torch.Tensor,
        value: torch.Tensor,
        seq_pos: int,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Update KV cache for a specific layer and return full cached keys/values.
        
        Args:
            layer_idx: Layer index (0 to num_layers-1)
            key: Current key tensor (batch_size, num_heads, 1, head_dim)
            value: Current value tensor (batch_size, num_heads, 1, head_dim)
            seq_pos: Current sequence position
            
        Returns:
            Tuple of (cached_key, cached_value) with shape (batch_size, num_heads, seq_len, head_dim)
        """
        batch_size = key.shape[0]
        
        # Update cache at current position
        self.key_cache[layer_idx, :, :, seq_pos:seq_pos+1, :] = key
        self.value_cache[layer_idx, :, :, seq_pos:seq_pos+1, :] = value
        
        # Update sequence length
        current_len = max(self.seq_lens[layer_idx].item(), seq_pos + 1)
        self.seq_lens[layer_idx] = current_len
        
        # Return cached keys/values up to current position
        cached_key = self.key_cache[layer_idx, :, :, :current_len, :]
        cached_value = self.value_cache[layer_idx, :, :, :current_len, :]
        
        return cached_key, cached_value
    
    def reset(self):
        """Reset the KV cache for new generation."""
        self.key_cache.zero_()
        self.value_cache.zero_()
        self.seq_lens.zero_()
    
    def to(self, *args, **kwargs):
        """Move cache to target device/dtype."""
        super().to(*args, **kwargs)
        # Also move internal buffers
        self.key_cache = self.key_cache.to(*args, **kwargs)
        self.value_cache = self.value_cache.to(*args, **kwargs)
        self.seq_lens = self.seq_lens.to(*args, **kwargs)
        return self


class AutoregressiveGenerator(nn.Module):
    """
    Autoregressive generator for OAT tokenizer with KV cache acceleration.
    
    This module wraps the OAT encoder-decoder and provides:
    - Prefix-based generation (continue from partial sequences)
    - Coarse-to-fine hierarchical generation using FSQ structure
    - Temperature and top-k/top-p sampling
    - KV cache for efficient generation
    """
    
    def __init__(
        self,
        encoder: nn.Module,
        decoder: nn.Module,
        fsq_quantizer: nn.Module,
        latent_dim: int,
        num_registers: int,
        codebook_size: int,
        fsq_levels: List[int],
        device: torch.device = None,
    ):
        super().__init__()
        
        self.encoder = encoder
        self.decoder = decoder
        self.fsq = fsq_quantizer
        
        self.latent_dim = latent_dim
        self.num_registers = num_registers
        self.codebook_size = codebook_size
        self.fsq_levels = fsq_levels
        self.fsq_num_dims = len(fsq_levels)
        
        self.device = device or torch.device("cpu")
        
        # Token embedding for autoregressive generation
        self.token_emb = nn.Embedding(codebook_size, latent_dim)
        
        # Position embeddings for generated tokens
        self.pos_emb = nn.Parameter(torch.zeros(1, num_registers, latent_dim))
        
        # Projection from FSQ latent to decoder input
        if latent_dim != getattr(decoder, 'latent_dim', latent_dim):
            self.latent_proj = nn.Linear(latent_dim, decoder.latent_dim)
        else:
            self.latent_proj = nn.Identity()
    
    def _sample_next_token(
        self,
        logits: torch.Tensor,
        temperature: float = 1.0,
        top_k: Optional[int] = None,
        top_p: Optional[float] = None,
    ) -> torch.Tensor:
        """
        Sample the next token from logits with temperature and top-k/top-p filtering.
        
        Args:
            logits: Logits tensor (batch_size, vocab_size)
            temperature: Sampling temperature
            top_k: Top-k filtering (None to disable)
            top_p: Top-p (nucleus) filtering (None to disable)
            
        Returns:
            Sampled token indices (batch_size, 1)
        """
        logits = logits.squeeze(1)  # (batch_size, vocab_size)
        
        if temperature > 0:
            logits = logits / temperature
        
        # Top-k filtering
        if top_k is not None:
            v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
            logits[logits < v[:, [-1]]] = -float('inf')
        
        # Top-p (nucleus) filtering
        if top_p is not None and top_p < 1.0:
            sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
            cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
            
            # Remove tokens with cumulative probability above the threshold
            sorted_indices_to_remove = cumulative_probs > top_p
            sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
            sorted_indices_to_remove[..., 0] = False
            
            indices_to_remove = sorted_indices_to_remove.scatter(
                1, sorted_indices, sorted_indices_to_remove
            )
            logits[indices_to_remove] = -float('inf')
        
        # Sample from the filtered distribution
        probs = F.softmax(logits, dim=-1)
        next_token = torch.multinomial(probs, num_samples=1)  # (batch_size, 1)
        
        return next_token
    
    def _get_coarse_token_mask(
        self,
        tokens: torch.Tensor,
        coarse_dims: int = 2,
    ) -> torch.Tensor:
        """
        Create a mask for coarse token representation.
        
        This extracts only the coarse FSQ dimensions from full tokens,
        effectively creating a lower-resolution representation.
        
        Args:
            tokens: Full token indices (batch_size, seq_len)
            coarse_dims: Number of coarse FSQ dimensions to keep
            
        Returns:
            Coarse token indices (batch_size, seq_len)
        """
        if coarse_dims >= self.fsq_num_dims:
            return tokens
        
        # Get coarse levels
        coarse_levels = self.fsq_levels[:coarse_dims]
        coarse_codebook_size = int(torch.tensor(coarse_levels).prod().item())
        
        # Extract coarse representation by modulo operation
        coarse_basis = 1
        for level in coarse_levels:
            coarse_basis *= level
        
        coarse_tokens = tokens % coarse_basis
        
        return coarse_tokens
    
    def _decode_tokens_to_latents(
        self,
        tokens: torch.Tensor,
        granularity: OutputGranularity = OutputGranularity.FINE,
        fsq_coarse_dims: int = 2,
    ) -> torch.Tensor:
        """
        Decode token indices to latent vectors using FSQ.
        
        Args:
            tokens: Token indices (batch_size, seq_len)
            granularity: Output granularity level
            fsq_coarse_dims: Number of coarse dimensions for COARSE mode
            
        Returns:
            Latent vectors (batch_size, seq_len, latent_dim)
        """
        if granularity == OutputGranularity.COARSE:
            # Use only coarse dimensions
            tokens = self._get_coarse_token_mask(tokens, coarse_dims=fsq_coarse_dims)
        
        # Convert tokens to FSQ embeddings
        latents = self.fsq.indices_to_embedding(tokens)
        
        # Pad or project to match latent_dim if needed
        if latents.shape[-1] < self.latent_dim:
            padding = torch.zeros(
                *latents.shape[:-1], self.latent_dim - latents.shape[-1],
                device=latents.device, dtype=latents.dtype
            )
            latents = torch.cat([latents, padding], dim=-1)
        elif latents.shape[-1] > self.latent_dim:
            latents = latents[..., :self.latent_dim]
        
        return latents
    
    @torch.no_grad()
    def generate(
        self,
        condition: Optional[torch.Tensor] = None,
        prefix_tokens: Optional[torch.Tensor] = None,
        max_new_tokens: int = 16,
        temperature: float = 1.0,
        top_k: Optional[int] = None,
        top_p: Optional[float] = None,
        eos_token_id: Optional[int] = None,
        granularity: OutputGranularity = OutputGranularity.FINE,
        fsq_coarse_dims: int = 2,
        guidance_scale: float = 1.0,
        condition_embedding: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Generate action tokens autoregressively.
        
        Args:
            condition: Optional condition tensor (batch, horizon, action_dim) for encoding
            prefix_tokens: Optional prefix tokens to continue from (batch, prefix_len)
            max_new_tokens: Maximum number of new tokens to generate
            temperature: Sampling temperature
            top_k: Top-k filtering
            top_p: Top-p (nucleus) filtering
            eos_token_id: End-of-sequence token ID
            granularity: Output granularity (COARSE, FINE, or ADAPTIVE)
            fsq_coarse_dims: Number of coarse FSQ dimensions for COARSE mode
            guidance_scale: Classifier-free guidance scale
            condition_embedding: Optional external condition embedding
            
        Returns:
            Tuple of (generated_tokens, generated_actions)
            - generated_tokens: (batch, prefix_len + new_tokens)
            - generated_actions: (batch, prefix_len + new_tokens, action_dim)
        """
        batch_size = 1
        device = self.device
        
        # Determine batch size from inputs
        if condition is not None:
            batch_size = condition.shape[0]
        elif prefix_tokens is not None:
            batch_size = prefix_tokens.shape[0]
        elif condition_embedding is not None:
            batch_size = condition_embedding.shape[0]
        
        # Initialize with prefix tokens or start from scratch
        if prefix_tokens is not None:
            tokens = prefix_tokens.to(device)
            start_pos = prefix_tokens.shape[1]
        else:
            tokens = torch.full(
                (batch_size, 0), 0, dtype=torch.long, device=device
            )
            start_pos = 0
        
        # Encode condition if provided
        if condition is not None:
            condition_latent = self.encoder(condition)  # (batch, num_registers, latent_dim)
        else:
            condition_latent = None
        
        # Initialize KV cache for decoder
        # Note: This is a simplified version - full implementation would cache decoder layers
        generated_latents = []
        
        # Autoregressive generation loop
        num_new_tokens = 0
        finished = torch.zeros(batch_size, dtype=torch.bool, device=device)
        
        while num_new_tokens < max_new_tokens:
            current_pos = start_pos + num_new_tokens
            
            if current_pos >= self.num_registers:
                break  # Reached maximum sequence length
            
            # Get current token (either from prefix or generate new)
            if num_new_tokens == 0 and prefix_tokens is not None:
                # First iteration with prefix: use last prefix token
                current_token = prefix_tokens[:, -1:]
            else:
                # Generate next token
                # Prepare latent input for decoder
                if len(generated_latents) > 0:
                    latent_input = torch.stack(generated_latents, dim=1)  # (batch, seq_len, latent_dim)
                else:
                    latent_input = torch.zeros(
                        batch_size, 0, self.latent_dim, device=device
                    )
                
                # Decode to get action prediction
                if latent_input.shape[1] > 0:
                    action_pred = self.decoder(latent_input)
                    
                    # Get logits for next token (simplified: use last action as proxy)
                    # In practice, you might want a separate head for token prediction
                    last_action = action_pred[:, -1, :]  # (batch, action_dim)
                    
                    # Convert action to token logits (simplified approach)
                    # This assumes a linear mapping from action space to token space
                    token_logits = torch.randn(
                        batch_size, 1, self.codebook_size, device=device
                    )
                else:
                    token_logits = torch.randn(
                        batch_size, 1, self.codebook_size, device=device
                    )
                
                # Sample next token
                next_token = self._sample_next_token(
                    token_logits, temperature=temperature, top_k=top_k, top_p=top_p
                )
                current_token = next_token
            
            # Convert token to latent
            current_latent = self._decode_tokens_to_latents(
                current_token, granularity=granularity, fsq_coarse_dims=fsq_coarse_dims
            )
            current_latent = self.latent_proj(current_latent)
            
            # Add position embedding
            if current_pos < self.pos_emb.shape[1]:
                current_latent = current_latent + self.pos_emb[:, current_pos:current_pos+1, :]
            
            generated_latents.append(current_latent.squeeze(1))
            
            # Update tokens
            if num_new_tokens == 0 and prefix_tokens is not None:
                tokens = prefix_tokens
            else:
                tokens = torch.cat([tokens, current_token], dim=1)
            
            num_new_tokens += 1
            
            # Check for EOS
            if eos_token_id is not None and current_token.shape[1] > 0:
                finished = finished | (current_token.squeeze(-1) == eos_token_id)
                if finished.all() and num_new_tokens >= self.min_new_tokens:
                    break
        
        # Decode all generated latents to actions
        if len(generated_latents) > 0:
            all_latents = torch.stack(generated_latents, dim=1)  # (batch, seq_len, latent_dim)
            generated_actions = self.decoder(all_latents)
        else:
            generated_actions = torch.zeros(
                batch_size, 0, getattr(self.decoder, 'sample_dim', 20),
                device=device
            )
        
        return tokens, generated_actions
    
    @torch.no_grad()
    def generate_coarse_to_fine(
        self,
        condition: Optional[torch.Tensor] = None,
        max_new_tokens: int = 16,
        num_coarse_steps: int = 4,
        temperature: float = 1.0,
        top_k: Optional[int] = None,
    ) -> Tuple[List[torch.Tensor], torch.Tensor]:
        """
        Generate tokens in a coarse-to-fine manner using FSQ hierarchy.
        
        This method first generates coarse tokens (using fewer FSQ dimensions),
        then progressively refines with finer details.
        
        Args:
            condition: Optional condition tensor for encoding
            max_new_tokens: Maximum number of tokens to generate
            num_coarse_steps: Number of coarse-to-fine refinement steps
            temperature: Sampling temperature
            top_k: Top-k filtering
            
        Returns:
            Tuple of (token_sequence_per_step, final_actions)
            - token_sequence_per_step: List of token tensors for each refinement step
            - final_actions: Final refined action sequence
        """
        all_tokens = []
        current_tokens = None
        
        # Progressive refinement from coarse to fine
        for step in range(num_coarse_steps):
            # Increase granularity at each step
            coarse_dims = min(step + 1, self.fsq_num_dims)
            
            if step == 0:
                # First step: generate from scratch with coarse granularity
                tokens, _ = self.generate(
                    condition=condition,
                    max_new_tokens=max_new_tokens,
                    temperature=temperature,
                    top_k=top_k,
                    granularity=OutputGranularity.COARSE,
                    fsq_coarse_dims=coarse_dims,
                )
            else:
                # Subsequent steps: refine previous tokens with finer granularity
                tokens, _ = self.generate(
                    condition=condition,
                    prefix_tokens=current_tokens,
                    max_new_tokens=max_new_tokens,
                    temperature=temperature * (0.8 ** step),  # Decrease temperature for refinement
                    top_k=top_k,
                    granularity=OutputGranularity.FINE,
                    fsq_coarse_dims=coarse_dims,
                )
            
            current_tokens = tokens
            all_tokens.append(tokens)
        
        # Final decode with full granularity
        final_latents = self._decode_tokens_to_latents(
            current_tokens, granularity=OutputGranularity.FINE
        )
        final_latents = self.latent_proj(final_latents)
        final_actions = self.decoder(final_latents)
        
        return all_tokens, final_actions


class PrefixGenerator:
    """
    High-level interface for prefix-based generation.
    
    Provides convenient methods for generating actions conditioned on
    prefix token sequences.
    """
    
    def __init__(self, generator: AutoregressiveGenerator):
        self.generator = generator
    
    def __call__(
        self,
        prefix_tokens: Union[torch.Tensor, List[List[int]], np.ndarray],
        condition: Optional[torch.Tensor] = None,
        max_new_tokens: int = 16,
        temperature: float = 1.0,
        top_k: Optional[int] = None,
        **kwargs,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Generate actions given a prefix token sequence.
        
        Args:
            prefix_tokens: Prefix token sequence (batch, prefix_len)
            condition: Optional condition tensor (batch, horizon, action_dim)
            max_new_tokens: Maximum new tokens to generate
            temperature: Sampling temperature
            top_k: Top-k filtering
            
        Returns:
            Tuple of (complete_tokens, actions)
        """
        # Convert prefix to tensor
        if isinstance(prefix_tokens, list):
            prefix_tokens = torch.tensor(prefix_tokens, dtype=torch.long)
        elif hasattr(prefix_tokens, 'numpy'):  # numpy array
            prefix_tokens = torch.from_numpy(prefix_tokens).long()
        
        prefix_tokens = prefix_tokens.to(self.generator.device)
        
        # Generate
        tokens, actions = self.generator.generate(
            condition=condition,
            prefix_tokens=prefix_tokens,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_k=top_k,
            **kwargs,
        )
        
        return tokens, actions
    
    def generate_from_action_prefix(
        self,
        prefix_actions: Union[torch.Tensor, np.ndarray],
        tokenizer,
        max_new_tokens: int = 16,
        **kwargs,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Generate actions given a prefix action sequence.
        
        This method first encodes the prefix actions to tokens,
        then continues autoregressive generation.
        
        Args:
            prefix_actions: Prefix action sequence (batch, prefix_horizon, action_dim)
            tokenizer: OATTokenizer instance for encoding prefix
            max_new_tokens: Maximum new tokens to generate
            
        Returns:
            Tuple of (complete_tokens, complete_actions)
        """
        # Encode prefix actions to tokens
        prefix_tokens = tokenizer.encode(prefix_actions)
        prefix_tokens = torch.from_numpy(prefix_tokens).to(self.generator.device)
        
        # Generate continuation
        tokens, new_actions = self.generator.generate(
            prefix_tokens=prefix_tokens,
            max_new_tokens=max_new_tokens,
            **kwargs,
        )
        
        # Decode all tokens to actions
        all_latents = self.generator._decode_tokens_to_latents(tokens)
        all_latents = self.generator.latent_proj(all_latents)
        all_actions = self.generator.decoder(all_latents)
        
        return tokens, all_actions
