# bclm/config.py
"""
Configuration dataclasses and shared constants for BCLM models.
"""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Tuple

# =============================================================================
# Special Token IDs
# =============================================================================

EOT_TOKEN_ID: int = 32001
PAD_TOKEN_ID: int = 32002
TURN_TOKEN_ID: int = 32010
USER_TOKEN_ID: int = 32011
ASSISTANT_TOKEN_ID: int = 32012
SYSTEM_TOKEN_ID: int = 32013
IGNORE_INDEX: int = -100

# Default Hugging Face organization for short-form model IDs
DEFAULT_HF_ORG: str = "bcml"

# Expected filenames inside a model directory
CONFIG_FILENAME: str = "config.json"
WEIGHTS_FILENAME: str = "model.safetensors"


# =============================================================================
# BCLM-1 Model Configuration
# =============================================================================

@dataclass
class BCLM1ModelConfig:
    """
    All architectural hyperparameters for a BCLM-1 model.

    Instances are typically created via ``from_dict`` when loading a
    model's ``config.json``, but can also be constructed directly for
    programmatic use.
    """

    vocab_size: int = 32768
    tokenizer: str = "tokenmonster:english-32000-consistent-v1"
    embed_dim: int = 384
    n_layers: int = 10
    max_seq_len: int = 16384
    dropout: float = 0.0
    attn_heads: int = 6
    attn_kv_heads: int = 2
    local_attn_layers: Tuple[int, ...] = (0, 2, 4, 6, 8)
    global_attn_layers: Tuple[int, ...] = ()
    attn_window_size: int = 1024
    conv_kernel_size: int = 4
    osc_n_pairs: int = 1
    osc_n_real: int = 16
    osc_clamp_min_decay: float = 1e-5
    bigram_table_factor: int = 5

    @classmethod
    def from_dict(cls, cfg: dict) -> "BCLM1ModelConfig":
        """
        Build a config from a dictionary (typically the ``"model"`` key
        inside a ``config.json``).

        Handles legacy key names and coerces list fields to tuples.
        """
        d = dict(cfg)

        # Legacy alias
        if "attn_window_size" not in d and "attn_block_size" in d:
            d["attn_window_size"] = d.pop("attn_block_size")

        # Coerce JSON lists → tuples
        for key in ("local_attn_layers", "global_attn_layers"):
            if key in d and isinstance(d[key], list):
                d[key] = tuple(int(x) for x in d[key])

        # Strip unknown keys so dataclass construction doesn't fail
        known = {f.name for f in cls.__dataclass_fields__.values()}
        d = {k: v for k, v in d.items() if k in known}

        return cls(**d)

    def to_dict(self) -> dict:
        """Serialize back to a plain dict (JSON-safe types)."""
        from dataclasses import asdict
        d = asdict(self)
        d["local_attn_layers"] = list(d["local_attn_layers"])
        d["global_attn_layers"] = list(d["global_attn_layers"])
        return d