"""ARCHON-Brain Configuration — 300M Transformer with MTP=5."""
from dataclasses import dataclass

@dataclass
class ArchonBrainConfig:
    """300M parameter transformer decoder optimized for ARCHON."""
    # Model architecture
    vocab_size: int = 32_000       # Custom tokenizer vocabulary
    hidden_dim: int = 1024         # Hidden dimension
    num_layers: int = 18           # Transformer layers (reduced for ~300M)
    num_heads: int = 16            # Attention heads
    head_dim: int = 64             # hidden_dim / num_heads
    intermediate_dim: int = 3072   # FFN intermediate (3x hidden, saves params)
    max_seq_len: int = 4096        # Context window — matches family quartet SEQ_LEN

    # Multi-Token Prediction (MTP)
    mtp_heads: int = 5             # Predict 5 tokens ahead simultaneously
    mtp_loss_weights: tuple = (1.0, 0.5, 0.3, 0.2, 0.1)  # Decreasing weight per head

    # Normalization & Activation
    norm_eps: float = 1e-6         # RMSNorm epsilon
    rope_theta: float = 10_000.0   # RoPE base frequency

    # Regularization
    dropout: float = 0.0           # No dropout during training (modern practice)
    tie_word_embeddings: bool = True  # Share input/output embeddings (saves ~32M params)

    # Training
    dtype: str = "bfloat16"        # Training precision

    @property
    def param_count(self) -> int:
        """Estimate total parameter count."""
        embed = self.vocab_size * self.hidden_dim  # ~32M
        if self.tie_word_embeddings:
            head = 0  # Shared with embedding
        else:
            head = self.vocab_size * self.hidden_dim

        # Per layer: attention (4 projections) + FFN (gate + up + down) + 2 norms
        attn = 4 * self.hidden_dim * self.hidden_dim  # Q, K, V, O
        ffn = 3 * self.hidden_dim * self.intermediate_dim  # gate, up, down (SwiGLU)
        norms = 2 * self.hidden_dim  # 2 RMSNorm per layer
        per_layer = attn + ffn + norms

        # MTP heads: small projection per head
        mtp = self.mtp_heads * self.hidden_dim * self.vocab_size
        if self.tie_word_embeddings:
            mtp = self.mtp_heads * self.hidden_dim * self.hidden_dim  # Project to hidden, then shared embed

        total = embed + head + (per_layer * self.num_layers) + mtp
        return total

    @property
    def param_count_human(self) -> str:
        count = self.param_count
        if count >= 1e9:
            return f"{count/1e9:.1f}B"
        return f"{count/1e6:.0f}M"


@dataclass
class TrainingConfig:
    """Training hyperparameters for ARCHON-Brain."""
    # Optimization
    learning_rate: float = 3e-4        # Peak LR (Chinchilla optimal for 300M)
    min_lr: float = 3e-5               # Min LR (10% of peak)
    weight_decay: float = 0.1
    beta1: float = 0.9
    beta2: float = 0.95
    grad_clip: float = 1.0

    # Schedule
    warmup_steps: int = 1000
    total_steps: int = 50_000          # ~6B tokens at batch ~120K tokens/step
    lr_schedule: str = "cosine"

    # Batch
    micro_batch_size: int = 8          # Per-GPU batch
    gradient_accumulation: int = 8     # Effective batch = 8 * 8 = 64
    seq_len: int = 2048

    # Data
    num_workers: int = 4

    # Logging
    log_interval: int = 10
    eval_interval: int = 500
    save_interval: int = 2000

    # Paths (relative to working directory — set cwd to /workspace/ARCHON-BRAIN)
    output_dir: str = "checkpoints"
    data_dir: str = "data/processed"
    tokenizer_path: str = "tokenizer"

    @property
    def tokens_per_step(self) -> int:
        return self.micro_batch_size * self.gradient_accumulation * self.seq_len

    @property
    def total_tokens(self) -> str:
        total = self.tokens_per_step * self.total_steps
        return f"{total/1e9:.1f}B"


# Print config on import
if __name__ == "__main__":
    model = ArchonBrainConfig()
    train = TrainingConfig()
    print(f"ARCHON-Brain Model Config:")
    print(f"  Parameters: {model.param_count_human} ({model.param_count:,})")
    print(f"  Layers: {model.num_layers}")
    print(f"  Hidden: {model.hidden_dim}")
    print(f"  Heads: {model.num_heads}")
    print(f"  Vocab: {model.vocab_size:,}")
    print(f"  Seq Len: {model.max_seq_len}")
    print(f"  MTP Heads: {model.mtp_heads}")
    print(f"\nTraining Config:")
    print(f"  Total Steps: {train.total_steps:,}")
    print(f"  Tokens/Step: {train.tokens_per_step:,}")
    print(f"  Total Tokens: {train.total_tokens}")
    print(f"  LR: {train.learning_rate}")
    print(f"  Batch: {train.micro_batch_size} x {train.gradient_accumulation} = {train.micro_batch_size * train.gradient_accumulation}")