from dataclasses import dataclass, asdict


@dataclass
class TinyLiquidConfig:
    # --- tokenizer ---
    vocab_size: int = 8192

    # --- architecture (our own "liquid" design, non-transformer) ---
    d_model: int = 320              # hidden width
    n_blocks: int = 6               # liquid blocks
    basis_n: int = 16               # basis blocks in expansion layer
    basis_b: int = 4                # dims per basis block (basis_n*basis_b = expanded width)
    mlp_ratio: int = 2              # dense MLP hidden = d_model * mlp_ratio
    num_experts: int = 0            # 0 = dense gated MLP; >0 = mixture-of-experts MLP
    num_experts_per_tok: int = 2
    expert_hidden: int = 0          # 0 => mlp_ratio*d_model//2 per expert
    num_personas: int = 3           # 0=none, 1=analyst, 2=skeptic (learned persona vectors)
    tower_d: int = 0               # 0=off; >0 = wide-head tower width (baseline-preserving growth)
    tower_blocks: int = 0          # number of identity-init blocks in the tower
    mtp_heads: int = 0             # 0=off; >0 = multi-token prediction aux heads (Meta MTP)

    # --- position / normalization ---
    rope_theta: float = 10000.0
    max_seq_len: int = 1024
    norm_eps: float = 1e-6
    tie_embeddings: bool = True

    def params_estimate(self):
        """Rough parameter count (ignores small heads)."""
        d, v = self.d_model, self.vocab_size
        n = self.basis_n * self.basis_b
        per_block = 2 * n * d                       # basis: in + forget
        if self.num_experts > 0:
            h = self.expert_hidden or (self.mlp_ratio * d // 2)
            per_mlp = 3 * d * h + d * h + d * self.num_experts
            per_block += self.num_experts * per_mlp
        else:
            h = self.mlp_ratio * d
            per_block += 4 * d * h                  # up, gate, down, forget
        params = v * d + self.num_personas * d + self.n_blocks * per_block
        if self.tower_d and self.tower_blocks:
            td, tn = self.tower_d, self.tower_blocks
            tn_ = tn  # basis rows in tower
            per_tower = 2 * self.basis_n * self.basis_b * td + 4 * (self.mlp_ratio * td) * td
            params += td * d + d * td + tn * per_tower
        if self.mtp_heads:
            params += self.mtp_heads * (d * d + d)   # SiLU MLP heads, tied output
        return params


CONFIGS = {
    "tiny10m": dict(d_model=320, n_blocks=6, basis_n=16, basis_b=4, mlp_ratio=2),
    "tiny10m-moe": dict(d_model=320, n_blocks=6, basis_n=16, basis_b=4, mlp_ratio=2,
                        num_experts=4, num_experts_per_tok=2),
    "micro6m": dict(d_model=256, n_blocks=5, basis_n=16, basis_b=4, mlp_ratio=2),
    "tiny13m": dict(d_model=320, n_blocks=12, basis_n=16, basis_b=4, mlp_ratio=2),
    "tiny16m": dict(d_model=448, n_blocks=7, basis_n=16, basis_b=5, mlp_ratio=2),
    "tiny20m": dict(d_model=512, n_blocks=8, basis_n=16, basis_b=5, mlp_ratio=2),
    "tiny28m": dict(d_model=600, n_blocks=8, basis_n=16, basis_b=6, mlp_ratio=2),
    "hybrid18m": dict(d_model=320, n_blocks=6, basis_n=16, basis_b=4, mlp_ratio=2,
                      tower_d=512, tower_blocks=4),
    "hybrid25m": dict(d_model=320, n_blocks=6, basis_n=16, basis_b=4, mlp_ratio=2,
                      tower_d=512, tower_blocks=8),
    "hybrid50m": dict(d_model=320, n_blocks=6, basis_n=16, basis_b=4, mlp_ratio=2,
                      tower_d=800, tower_blocks=8),
}