"""Configuration for the owned dual-attention decoder LM."""

from dataclasses import dataclass

from .attention_config import SUPPORTED_SEQUENCE_BOUNDARY_POLICIES


SUPPORTED_PE_TYPES = {"sinusoidal", "learned", "relative", "rope", "none"}
SUPPORTED_SYMBOL_RETRIEVAL = {"symbolic", "positional", "relative", "relsymbolic"}
SUPPORTED_RA_TYPES = {"ra", "rca", "disrca"}
SUPPORTED_RA_ACTIVATIONS = {"softmax", "identity", "relu", "tanh", "sigmoid", "gelu"}
SUPPORTED_FFN_ACTIVATIONS = {"gelu", "relu", "swiglu", "identity"}
SUPPORTED_FFN_HIDDEN_DIM_MODES = {"dff_factor", "swiglu_parameter_matched"}
SUPPORTED_INIT_SCHEMES = {"xavier_uniform", "normal_0_02_scaled_projection"}


@dataclass(frozen=True)
class DatLMConfig:
    vocab_size: int
    max_seq_len: int
    pe_type: str = "rope"
    hidden_dim: int = 256
    n_heads_sa: int = 2
    n_heads_ra: int = 2
    n_layers: int = 4
    dropout: float = 0.0
    dff_factor: int = 4
    ffn_hidden_dim_mode: str = "dff_factor"
    ffn_activation: str = "gelu"
    rope_theta: float = 10000.0
    max_rel_pos: int | None = None
    sequence_boundary_policy: str = "eos_document"
    segment_boundary_token_id: int | None = None
    init_range: float = 0.15
    init_scheme: str = "xavier_uniform"
    norm_type: str = "rmsnorm"
    norm_first: bool = True
    use_bias_qkv: bool = False
    use_bias_out: bool = True
    use_bias_ffn: bool = True
    tie_lm_head: bool = True
    symbol_dim: int | None = None
    n_symbols: int | None = None
    symbolic_attn_n_heads: int | None = None
    symbol_retrieval: str = "symbolic"
    symbolic_use_bias: bool = False
    shared_symbol_retriever: bool = True
    share_attn_params: bool = False
    positional_symbols_sinusoidal: bool = False
    relative_symbols_rope: bool = False
    relsymbolic_rel_n_heads: int = 4
    relsymbolic_symbolic_attn_n_heads: int = 4
    relsymbolic_neighborhood_size: int = 2
    relsymbolic_include_self: bool = False
    relsymbolic_normalize_rels: bool = True
    relsymbolic_trainable_symbols: bool = True
    relsymbolic_dropout: float = 0.0
    relsymbolic_rel_scale: float | None = None
    relsymbolic_symbolic_attn_scale: float | None = None
    relsymbolic_use_bias: bool = False
    ra_type: str = "ra"
    ra_n_relations: int | None = None
    ra_rel_activation: str = "identity"
    ra_symmetric_rels: bool = False
    pad_token_id: int = 0
    bos_token_id: int = 1
    eos_token_id: int = 2
    mlm_head_enabled: bool = False

    def __post_init__(self) -> None:
        if self.vocab_size <= 0:
            raise ValueError(f"vocab_size must be positive, got {self.vocab_size}")
        if self.max_seq_len <= 0:
            raise ValueError(f"max_seq_len must be positive, got {self.max_seq_len}")
        if self.pe_type not in SUPPORTED_PE_TYPES:
            raise ValueError(f"Unsupported pe_type: {self.pe_type}")
        if self.hidden_dim <= 0:
            raise ValueError(f"hidden_dim must be positive, got {self.hidden_dim}")
        if self.n_heads_sa <= 0:
            raise ValueError(f"n_heads_sa must be positive for DAT, got {self.n_heads_sa}")
        if self.n_heads_ra <= 0:
            raise ValueError(f"n_heads_ra must be positive for DAT, got {self.n_heads_ra}")
        total_heads = self.total_n_heads
        if self.hidden_dim % total_heads != 0:
            raise ValueError(
                f"hidden_dim ({self.hidden_dim}) must be divisible by total DAT heads "
                f"({total_heads} = {self.n_heads_sa} SA + {self.n_heads_ra} RA)"
            )
        if self.n_layers <= 0:
            raise ValueError(f"n_layers must be positive, got {self.n_layers}")
        if not 0.0 <= self.dropout < 1.0:
            raise ValueError(f"dropout must be in [0.0, 1.0), got {self.dropout}")
        if self.dff_factor <= 0:
            raise ValueError(f"dff_factor must be positive, got {self.dff_factor}")
        if self.ffn_hidden_dim_mode not in SUPPORTED_FFN_HIDDEN_DIM_MODES:
            raise ValueError(f"Unsupported ffn_hidden_dim_mode for DAT: {self.ffn_hidden_dim_mode}")
        if self.ffn_activation not in SUPPORTED_FFN_ACTIVATIONS:
            raise ValueError(f"Unsupported ffn_activation for DAT: {self.ffn_activation}")
        if self.ffn_hidden_dim_mode == "swiglu_parameter_matched" and self.ffn_activation != "swiglu":
            raise ValueError(
                "ffn_hidden_dim_mode='swiglu_parameter_matched' requires "
                f"ffn_activation='swiglu', got {self.ffn_activation}"
            )
        if self.rope_theta <= 0.0:
            raise ValueError(f"rope_theta must be positive, got {self.rope_theta}")
        if self.max_rel_pos is not None and self.max_rel_pos <= 0:
            raise ValueError(f"max_rel_pos must be positive when provided, got {self.max_rel_pos}")
        if self.sequence_boundary_policy not in SUPPORTED_SEQUENCE_BOUNDARY_POLICIES:
            raise ValueError(
                "Unsupported sequence_boundary_policy for DAT: "
                f"{self.sequence_boundary_policy}"
            )
        if self.sequence_boundary_policy == "segment_document":
            if self.segment_boundary_token_id is None:
                raise ValueError(
                    "segment_boundary_token_id is required when "
                    "sequence_boundary_policy='segment_document'"
                )
        if self.init_range <= 0.0:
            raise ValueError(f"init_range must be positive, got {self.init_range}")
        if self.init_scheme not in SUPPORTED_INIT_SCHEMES:
            raise ValueError(f"Unsupported init_scheme for DAT: {self.init_scheme}")
        if self.norm_type not in {"layernorm", "rmsnorm"}:
            raise ValueError(
                f"norm_type must be 'layernorm' or 'rmsnorm', got {self.norm_type}"
            )
        if self.pe_type == "rope" and self.head_dim % 2 != 0:
            raise ValueError(
                "RoPE requires even DAT head_dim, got "
                f"{self.head_dim} from hidden_dim={self.hidden_dim}, total_heads={total_heads}"
            )
        if self.pe_type == "sinusoidal" and self.hidden_dim % 2 != 0:
            raise ValueError(f"Sinusoidal encoding requires even hidden_dim, got {self.hidden_dim}")
        if self.share_attn_params and self.n_heads_sa != self.n_heads_ra:
            raise ValueError(
                "share_attn_params=True requires n_heads_sa == n_heads_ra, "
                f"got {self.n_heads_sa} and {self.n_heads_ra}"
            )
        if self.symbol_dim is not None and self.symbol_dim <= 0:
            raise ValueError(f"symbol_dim must be positive when provided, got {self.symbol_dim}")
        if self.symbolic_attn_n_heads is not None and self.symbolic_attn_n_heads <= 0:
            raise ValueError(
                "symbolic_attn_n_heads must be positive when provided, "
                f"got {self.symbolic_attn_n_heads}"
            )
        if self.symbol_retrieval not in SUPPORTED_SYMBOL_RETRIEVAL:
            raise ValueError(f"Unsupported symbol_retrieval for DAT: {self.symbol_retrieval}")
        if self.symbol_retrieval == "symbolic":
            symbolic_heads = self.resolved_symbolic_attn_n_heads
            if self.hidden_dim % symbolic_heads != 0:
                raise ValueError(
                    f"hidden_dim ({self.hidden_dim}) must be divisible by symbolic_attn_n_heads "
                    f"({symbolic_heads}) for symbolic retrieval"
                )
            if self.resolved_symbol_dim % symbolic_heads != 0:
                raise ValueError(
                    f"symbol_dim ({self.resolved_symbol_dim}) must be divisible by symbolic_attn_n_heads "
                    f"({symbolic_heads}) for symbolic retrieval"
                )
        if (
            self.symbol_retrieval == "positional"
            and self.positional_symbols_sinusoidal
            and self.resolved_symbol_dim % 2 != 0
        ):
            raise ValueError(
                "Sinusoidal positional symbols require even symbol_dim, "
                f"got {self.resolved_symbol_dim}"
            )
        if self.relative_symbols_rope:
            if self.symbol_retrieval != "relative":
                raise ValueError(
                    "relative_symbols_rope=True requires symbol_retrieval='relative', "
                    f"got {self.symbol_retrieval!r}"
                )
            if self.resolved_symbol_dim % 2 != 0:
                raise ValueError(
                    "RoPE relative symbols require even symbol_dim, "
                    f"got {self.resolved_symbol_dim}"
                )
        if self.positional_symbols_sinusoidal and self.symbol_retrieval != "positional":
            raise ValueError(
                "positional_symbols_sinusoidal=True requires symbol_retrieval='positional', "
                f"got {self.symbol_retrieval!r}"
            )
        if self.resolved_n_symbols <= 0:
            raise ValueError(f"resolved_n_symbols must be positive, got {self.resolved_n_symbols}")
        if self.symbol_retrieval == "relsymbolic":
            if self.relsymbolic_rel_n_heads <= 0:
                raise ValueError(
                    f"relsymbolic_rel_n_heads must be positive, got {self.relsymbolic_rel_n_heads}"
                )
            if self.hidden_dim % self.relsymbolic_rel_n_heads != 0:
                raise ValueError(
                    f"hidden_dim ({self.hidden_dim}) must be divisible by "
                    f"relsymbolic_rel_n_heads ({self.relsymbolic_rel_n_heads})"
                )
            if self.relsymbolic_symbolic_attn_n_heads <= 0:
                raise ValueError(
                    "relsymbolic_symbolic_attn_n_heads must be positive, "
                    f"got {self.relsymbolic_symbolic_attn_n_heads}"
                )
            if self.hidden_dim % self.relsymbolic_symbolic_attn_n_heads != 0:
                raise ValueError(
                    f"hidden_dim ({self.hidden_dim}) must be divisible by "
                    f"relsymbolic_symbolic_attn_n_heads ({self.relsymbolic_symbolic_attn_n_heads})"
                )
            if self.resolved_symbol_dim % self.relsymbolic_symbolic_attn_n_heads != 0:
                raise ValueError(
                    f"symbol_dim ({self.resolved_symbol_dim}) must be divisible by "
                    f"relsymbolic_symbolic_attn_n_heads ({self.relsymbolic_symbolic_attn_n_heads})"
                )
            if self.relsymbolic_neighborhood_size <= 0:
                raise ValueError(
                    "relsymbolic_neighborhood_size must be positive, "
                    f"got {self.relsymbolic_neighborhood_size}"
                )
            if not 0.0 <= self.relsymbolic_dropout < 1.0:
                raise ValueError(
                    f"relsymbolic_dropout must be in [0.0, 1.0), got {self.relsymbolic_dropout}"
                )
            if self.relsymbolic_rel_scale is not None and self.relsymbolic_rel_scale <= 0.0:
                raise ValueError(
                    "relsymbolic_rel_scale must be positive when provided, "
                    f"got {self.relsymbolic_rel_scale}"
                )
            if (
                self.relsymbolic_symbolic_attn_scale is not None
                and self.relsymbolic_symbolic_attn_scale <= 0.0
            ):
                raise ValueError(
                    "relsymbolic_symbolic_attn_scale must be positive when provided, "
                    f"got {self.relsymbolic_symbolic_attn_scale}"
                )
        if self.ra_type not in SUPPORTED_RA_TYPES:
            raise ValueError(f"Unsupported ra_type for DAT: {self.ra_type}")
        if self.ra_n_relations is not None and self.ra_n_relations <= 0:
            raise ValueError(
                f"ra_n_relations must be positive when provided, got {self.ra_n_relations}"
            )
        if self.ra_type != "ra" and self.ra_n_relations is not None:
            raise ValueError(f"ra_n_relations applies only to ra_type='ra', got {self.ra_type}")
        if self.ra_type != "ra" and self.ra_symmetric_rels:
            raise ValueError(f"ra_symmetric_rels applies only to ra_type='ra', got {self.ra_type}")
        n_relations = self.resolved_ra_n_relations
        if self.ra_type == "ra" and (self.head_dim * self.n_heads_ra) % n_relations != 0:
            raise ValueError(
                f"head_dim * n_heads_ra ({self.head_dim * self.n_heads_ra}) must be "
                f"divisible by ra_n_relations ({n_relations})"
            )
        if self.ra_rel_activation not in SUPPORTED_RA_ACTIVATIONS:
            raise ValueError(
                f"Unsupported ra_rel_activation for DAT: {self.ra_rel_activation}"
            )

    @property
    def total_n_heads(self) -> int:
        return self.n_heads_sa + self.n_heads_ra

    @property
    def head_dim(self) -> int:
        return self.hidden_dim // self.total_n_heads

    @property
    def resolved_symbol_dim(self) -> int:
        return self.hidden_dim if self.symbol_dim is None else self.symbol_dim

    @property
    def resolved_symbolic_attn_n_heads(self) -> int:
        return self.total_n_heads if self.symbolic_attn_n_heads is None else self.symbolic_attn_n_heads

    @property
    def resolved_ffn_hidden_dim(self) -> int:
        if self.ffn_hidden_dim_mode == "swiglu_parameter_matched":
            return int(8 / 3 * self.hidden_dim)
        return self.hidden_dim * self.dff_factor

    @property
    def resolved_n_symbols(self) -> int:
        return self.max_seq_len if self.n_symbols is None else self.n_symbols

    @property
    def resolved_ra_n_relations(self) -> int:
        return self.n_heads_ra if self.ra_n_relations is None else self.ra_n_relations