from dataclasses import dataclass
from typing import Optional

import torch

from .config import DiaConfig

def create_attn_mask(
    q_padding_mask_1d: torch.Tensor,
    k_padding_mask_1d: torch.Tensor,
    device: torch.device,
    is_causal: bool = False,
) -> torch.Tensor:
    """
    Creates the attention mask (self or cross) mimicking JAX segment ID logic.
    """
    # B1, Tq = q_padding_mask_1d.shape
    # B2, Tk = k_padding_mask_1d.shape

    p_mask_q = q_padding_mask_1d.unsqueeze(2)  # Shape [B, Tq, 1]
    p_mask_k = k_padding_mask_1d.unsqueeze(1)  # Shape [B, 1, Tk]

    # Condition A: Non-padding query attends to non-padding key
    non_pad_attends_non_pad = p_mask_q & p_mask_k  # Shape [B, Tq, Tk]

    # Condition B: Padding query attends to padding key
    pad_attends_pad = (~p_mask_q) & (~p_mask_k)  # Shape [B, Tq, Tk]

    # Combine: True if padding status is compatible (both non-pad OR both pad)
    mask = non_pad_attends_non_pad | pad_attends_pad  # Shape [B, Tq, Tk]

    if is_causal:
        # assert Tq == Tk, "Causal mask requires query and key sequence lengths to be equal"
        causal_mask_2d = torch.tril(torch.ones_like(mask[0], dtype=torch.bool, device=device))  # Shape [B, Tq, Tk]
        causal_mask = mask & causal_mask_2d  # Shape [B, Tq, Tk]
        return causal_mask.unsqueeze(1)  # Shape [B, 1, Tq, Tk]
    else:
        return mask.unsqueeze(1)  # Shape [B, 1, Tq, Tk]

@dataclass
class EncoderInferenceState:
    """Parameters specifically for encoder inference."""

    max_seq_len: int
    device: torch.device
    positions: torch.Tensor
    padding_mask: torch.Tensor
    attn_mask: torch.Tensor

    @classmethod
    def new(cls, config: DiaConfig, cond_src: torch.Tensor) -> "EncoderInferenceState":
        """Creates EtorchrInferenceParams from DiaConfig and a device."""
        device = cond_src.device

        positions = torch.arange(
            config.encoder_config.max_position_embeddings, dtype=torch.float32, device=device
        ).unsqueeze(0)
        padding_mask = (cond_src.squeeze(1) != 0).to(device).repeat_interleave(2, dim=0)
        attn_mask = create_attn_mask(padding_mask, padding_mask, device, is_causal=False)

        return cls(
            max_seq_len=config.encoder_config.max_position_embeddings,
            device=device,
            positions=positions,
            padding_mask=padding_mask,
            attn_mask=attn_mask,
        )

class KVCache(torch.nn.Module):
    k: torch.Tensor
    v: torch.Tensor

    def __init__(
        self,
        batch_size: int,
        num_heads: int,
        max_len: int,
        head_dim: int,
        dtype: torch.dtype,
        device: torch.device,
        k: torch.Tensor | None = None,
        v: torch.Tensor | None = None,
    ):
        k = torch.zeros((2 * batch_size, num_heads, max_len, head_dim), dtype=dtype, device=device) if k is None else k
        v = torch.zeros((2 * batch_size, num_heads, max_len, head_dim), dtype=dtype, device=device) if v is None else v
        super().__init__()

        self.register_buffer("k", k)
        self.register_buffer("v", v)

    @classmethod
    def from_kv(cls, k: torch.Tensor, v: torch.Tensor) -> "KVCache":
        return cls(
            batch_size=k.shape[0] // 2,
            num_heads=k.shape[1],
            max_len=k.shape[2],
            head_dim=k.shape[3],
            dtype=k.dtype,
            device=k.device,
            k=k,
            v=v,
        )

    def update(self, k: torch.Tensor, v: torch.Tensor, current_idx: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        k_out, v_out = self.k, self.v
        k_out[:, :, current_idx, :] = k
        v_out[:, :, current_idx, :] = v
        return self.k, self.v

    def prefill(self, k: torch.Tensor, v: torch.Tensor):
        prefill_len = k.shape[2]
        self.k[:, :, :prefill_len, :] = k
        self.v[:, :, :prefill_len, :] = v

@dataclass
class DecoderInferenceState:
    """Parameters specifically for decoder inference."""

    device: torch.device
    dtype: torch.dtype
    enc_out: torch.Tensor
    enc_positions: torch.Tensor
    dec_positions: torch.Tensor
    self_attn_cache: list[KVCache]
    cross_attn_cache: list[KVCache]
    casual_attn_mask: torch.Tensor
    cross_attn_mask: torch.Tensor

    @classmethod
    def new(
        cls,
        config: DiaConfig,
        enc_state: EncoderInferenceState,
        enc_out: torch.Tensor,
        dec_cross_attn_cache: list[KVCache],
        compute_dtype: torch.dtype,
        max_generation_length: Optional[int] = None,
    ) -> "DecoderInferenceState":
        """Creates DecoderInferenceParams from DiaConfig and a device."""
        device = enc_out.device
        max_audio_len = max_generation_length or config.decoder_config.max_position_embeddings
        batch_size = enc_out.shape[0] // 2

        dec_positions = torch.full((2 * batch_size, 1), fill_value=0, dtype=torch.int32, device=device)
        causal_mask = torch.tril(torch.ones(max_audio_len, max_audio_len, dtype=torch.bool, device=device))
        dec_mask = torch.ones((2 * batch_size, 1), dtype=torch.bool, device=device)
        cross_attn_mask = create_attn_mask(dec_mask, enc_state.padding_mask, device, is_causal=False)

        self_attn_cache = [
            KVCache(
                batch_size,
                config.decoder_config.num_key_value_heads,
                max_audio_len,
                config.decoder_config.head_dim,
                compute_dtype,
                device,
            )
            for _ in range(config.decoder_config.num_hidden_layers)
        ]

        return cls(
            device=device,
            dtype=compute_dtype,
            enc_out=enc_out,
            enc_positions=enc_state.positions,
            dec_positions=dec_positions,
            self_attn_cache=self_attn_cache,
            cross_attn_cache=dec_cross_attn_cache,
            casual_attn_mask=causal_mask,
            cross_attn_mask=cross_attn_mask,
        )

    def prepare_step(self, step_from: int, step_to: int | None = None) -> None:
        if step_to is None:
            step_to = step_from + 1
        self.dec_positions = torch.arange(step_from, step_to, dtype=torch.int32, device=self.device).unsqueeze(0)

@dataclass

class DecoderOutput:
    generated_tokens: torch.Tensor
    prefill_steps: list[int]
    @classmethod
    def new(cls, batch_size: int, config: DiaConfig, device: torch.device) -> "DecoderOutput":
        max_audio_len = config.decoder_config.max_position_embeddings
        return cls(
            generated_tokens=torch.full(
                (batch_size, max_audio_len, config.decoder_config.num_channels),
                fill_value=-1,
                dtype=torch.int,
                device=device,
            ),
            prefill_steps=[],
        )
    def get_tokens_at(self, step_from: int, step_to: int | None = None) -> torch.Tensor:
        if step_to is None:
            step_to = step_from + 1
        return self.generated_tokens[:, step_from:step_to, :]
    def update_one(self, dec_out: torch.Tensor, step: int, apply_mask: bool = False):
        dec_out = dec_out.to(self.generated_tokens.dtype)
        if apply_mask:
            mask = self.generated_tokens[:, step, :] == -1
            self.generated_tokens[:, step, :] = torch.where(mask, dec_out, self.generated_tokens[:, step, :])
        else:
            self.generated_tokens[:, step, :] = dec_out
    def prefill(self, dec_out: torch.Tensor, prefill_steps: list[int]):
        length = dec_out.shape[1]
        self.generated_tokens[:, :length, :] = dec_out
        self.prefill_steps = prefill_steps

"""Configuration management module for the Dia model.

This module provides comprehensive configuration management for the Dia model,
utilizing Pydantic for validation. It defines configurations for data processing,
model architecture (encoder and decoder), and training settings.

Key components:
- DataConfig: Parameters for data loading and preprocessing.
- EncoderConfig: Architecture details for the encoder module.
- DecoderConfig: Architecture details for the decoder module.
- ModelConfig: Combined model architecture settings.
- TrainingConfig: Training hyperparameters and settings.
- DiaConfig: Master configuration combining all components.
"""

import os

from pydantic import BaseModel, Field

class EncoderConfig(BaseModel, frozen=True):
    """Configuration for the encoder component of the Dia model.

    Attributes:
        model_type: Type of the model, defaults to "dia_encoder".
        hidden_size: Size of the encoder layers, defaults to 1024.
        intermediate_size: Size of the "intermediate" (i.e., feed-forward) layer in the encoder, defaults to 4096.
        num_hidden_layers: Number of hidden layers in the encoder, defaults to 12.
        num_attention_heads: Number of attention heads in the encoder, defaults to 16.
        num_key_value_heads: Number of key-value heads in the encoder, defaults to 16.
        head_dim: Dimension of each attention head, defaults to 128.
        hidden_act: Activation function in the encoder, defaults to "silu".
        max_position_embeddings: Maximum number of position embeddings, defaults to 1024.
        initializer_range: Range for initializing weights, defaults to 0.02.
        norm_eps: Epsilon value for normalization layers, defaults to 1e-5.
        rope_theta: Theta value for RoPE, defaults to 10000.0.
        rope_scaling: Optional scaling factor for RoPE.
        vocab_size: Vocabulary size, defaults to 256.
    """

    head_dim: int = Field(default=128, gt=0)
    hidden_act: str = Field(default="silu")
    hidden_size: int = Field(default=1024, gt=0)
    initializer_range: float = Field(default=0.02)
    intermediate_size: int = Field(default=4096, gt=0)
    max_position_embeddings: int = Field(default=1024, gt=0)
    model_type: str = Field(default="dia_encoder")
    norm_eps: float = Field(default=1e-5)
    num_attention_heads: int = Field(default=16, gt=0)
    num_hidden_layers: int = Field(default=12, gt=0)
    num_key_value_heads: int = Field(default=16, gt=0)
    rope_scaling: float | None = Field(default=None)
    rope_theta: float = Field(default=10000.0)
    vocab_size: int = Field(default=256, gt=0)

class DecoderConfig(BaseModel, frozen=True):
    """Configuration for the decoder component of the Dia model.

    Attributes:
        model_type: Type of the model, defaults to "dia_decoder".
        hidden_size: Size of the decoder layers, defaults to 2048.
        intermediate_size: Size of the "intermediate" (i.e., feed-forward) layer in the decoder, defaults to 8192.
        num_hidden_layers: Number of hidden layers in the decoder, defaults to 18.
        num_attention_heads: Number of attention heads in the decoder, defaults to 16.
        num_key_value_heads: Number of key-value heads in the decoder, defaults to 4.
        head_dim: Dimension of each attention head, defaults to 128.
        cross_hidden_size: Size of the cross-attention layers, defaults to 1024.
        cross_num_attention_heads: Number of attention heads in the cross-attention mechanism, defaults to 16.
        cross_num_key_value_heads: Number of key-value heads in the cross-attention mechanism, defaults to 16.
        cross_head_dim: Dimension of each cross-attention head, defaults to 128.
        hidden_act: Activation function in the decoder, defaults to "silu".
        max_position_embeddings: Maximum number of position embeddings in the decoder, defaults to 3072.
        initializer_range: Range for initializing weights in the decoder, defaults to 0.02.
        norm_eps: Epsilon value for normalization layers in the decoder, defaults to 1e-5.
        rope_theta: Theta value for RoPE in the decoder, defaults to 10000.0.
        rope_scaling: Optional scaling factor for RoPE in the decoder.
        vocab_size: Vocabulary size for the decoder, defaults to 1028.
        num_channels: Number of channels in the decoder, defaults to 9.
    """

    cross_head_dim: int = Field(default=128, gt=0)
    cross_hidden_size: int = Field(default=1024, gt=0)
    cross_num_attention_heads: int = Field(default=16, gt=0)
    cross_num_key_value_heads: int = Field(default=16, gt=0)
    head_dim: int = Field(default=128, gt=0)
    hidden_act: str = Field(default="silu")
    hidden_size: int = Field(default=2048, gt=0)
    initializer_range: float = Field(default=0.02)
    intermediate_size: int = Field(default=8192, gt=0)
    max_position_embeddings: int = Field(default=3072, gt=0)
    model_type: str = Field(default="dia_decoder")
    norm_eps: float = Field(default=1e-5)
    num_attention_heads: int = Field(default=16, gt=0)
    num_channels: int = Field(default=9, gt=0)
    num_hidden_layers: int = Field(default=18, gt=0)
    num_key_value_heads: int = Field(default=4, gt=0)
    rope_scaling: float | None = Field(default=None)
    rope_theta: float = Field(default=10000.0)
    vocab_size: int = Field(default=1028, gt=0)

class DiaConfig(BaseModel, frozen=True):
    """Main configuration container for the Dia model architecture.

    Attributes:
        model_type: Type of the model, defaults to "dia".
        is_encoder_decoder: Flag indicating if the model is an encoder-decoder type, defaults to True.
        encoder: Configuration for the encoder component.
        decoder: Configuration for the decoder component.
        src_vocab_size: Size of the source (text) vocabulary.
        tgt_vocab_size: Size of the target (audio code) vocabulary.
        initializer_range: Range for initializing weights, defaults to 0.02.
        norm_eps: Epsilon value for normalization layers, defaults to 1e-5.
        torch_dtype: Data type for model weights in PyTorch, defaults to "float32".
        bos_token_id: Beginning-of-sequence token ID, defaults to 1026.
        eos_token_id: End-of-sequence token ID, defaults to 1024.
        pad_token_id: Padding token ID, defaults to 1025.
        rope_theta: Theta value for RoPE, defaults to 10000.0.
        rope_scaling: Optional scaling factor for RoPE.
        transformers_version: Version of the transformers library, defaults to "4.53.0.dev0".
        architectures: List of model architectures, defaults to ["DiaForConditionalGeneration"].
        delay_pattern: List of delay values for each audio channel, defaults to [0,8,9,10,11,12,13,14,15].
    """

    architectures: list[str] = Field(default_factory=lambda: ["DiaForConditionalGeneration"])
    bos_token_id: int = Field(default=1026)
    decoder_config: DecoderConfig
    delay_pattern: list[int] = Field(default_factory=lambda: [0, 8, 9, 10, 11, 12, 13, 14, 15])
    encoder_config: EncoderConfig
    eos_token_id: int = Field(default=1024)
    initializer_range: float = Field(default=0.02)
    is_encoder_decoder: bool = Field(default=True)
    model_type: str = Field(default="dia")
    norm_eps: float = Field(default=1e-5)
    pad_token_id: int = Field(default=1025)
    torch_dtype: str = Field(default="float32")
    transformers_version: str = Field(default="4.53.0.dev0")

    def save(self, path: str) -> None:
        """Save the current configuration instance to a JSON file.

        Ensures the parent directory exists and the file has a .json extension.

        Args:
            path: The target file path to save the configuration.

        Raises:
            ValueError: If the path is not a file with a .json extension.
        """
        os.makedirs(os.path.dirname(path), exist_ok=True)
        config_json = self.model_dump_json(indent=2)
        with open(path, "w") as f:
            f.write(config_json)

    @classmethod
    def load(cls, path: str) -> "DiaConfig | None":
        """Load and validate a Dia configuration from a JSON file.

        Args:
            path: The path to the configuration file.

        Returns:
            A validated DiaConfig instance if the file exists and is valid,
            otherwise None if the file is not found.

        Raises:
            ValueError: If the path does not point to an existing .json file.
            pydantic.ValidationError: If the JSON content fails validation against the DiaConfig schema.
        """
        try:
            with open(path, "r") as f:
                content = f.read()
            return cls.model_validate_json(content)
        except FileNotFoundError:
            return None

class DecoderOutput:
    def new(cls, batch_size:
        pass
    def get_tokens_at(self, step_from:
        pass
    def update_one(self, dec_out:
        pass
    def prefill(self, dec_out:
        pass

This code segment functions to: Manage and store sequential output data during a generation process, including initializing storage structures, updating specific elements with optional masking, and retrieving subsets of generated results.

[TASK_BEGIN]
```python
class DecoderOutput:
    def __init__(self, logits: torch.Tensor, tokens: torch.Tensor, current_idx: torch.Tensor):
        self.logits = logits
        self.tokens = tokens
        self.current_idx = current_idx

    @classmethod
    def new(cls, batch_size: int, max_length: int, vocab_size: int, device: torch.device, dtype: torch.dtype) -> "DecoderOutput":
        logits = torch.zeros((batch_size, max_length, vocab_size), dtype=dtype, device=device)
        tokens = torch.zeros((batch_size, max_length), dtype=torch.long, device=device)
        current_idx = torch.tensor(0, dtype=torch.int32, device=device)
        return cls(logits, tokens, current_idx)

    def get_tokens_at(self, step_from: int) -> torch.Tensor:
        return self.tokens[:, step_from]

    def update_one(self, dec_out: torch.Tensor) -> None:
        batch_indices = torch.arange(self.logits.size(0), device=self.logits.device)
        self.logits[batch_indices, self.current_idx] = dec_out
        self.tokens[batch_indices, self.current_idx] = dec_out.argmax(dim=-1)
        self.current_idx += 1

    def prefill(self, dec_out: torch.Tensor) -> None:
        prefill_length = dec_out.size(1)
        self.logits[:, :prefill_length] = dec_out
        self.tokens[:, :prefill_length] = dec_out.argmax(dim=-1)
        self.current_idx = torch.tensor(prefill_length, device=self.current_idx.device)
```
[TASK_END]