# ─────────────────────────────────────────────────────────
# EGen Labs | ErebusTN
# AXIOM-1 — EGen-Core/AXIOM-1
# File: axiom1/training/optimizer.py
# ─────────────────────────────────────────────────────────

"""
Optimizer and learning rate schedule for AXIOM-1.

AdamW with separated decay/no-decay parameter groups.
β2 = 0.95 (NOT 0.999) — recurrent models need lower β2 for stability.
Warmup + cosine decay LR schedule.
"""

import math

import torch

from axiom1.model.model import AXIOM1
from axiom1.model.config import AXIOM1Config


def build_optimizer(
    model: AXIOM1,
    cfg: AXIOM1Config
) -> torch.optim.AdamW:
    """
    AdamW with separated decay/no-decay parameter groups.

    NO weight decay on: embeddings, norm weights, biases, loop step embeddings.
    Weight decay on: projection weights (q/k/v/o_proj, gate, up, down).

    Args:
        model: AXIOM1 model instance
        cfg: AXIOM1Config instance

    Returns:
        Configured AdamW optimizer
    """
    no_decay_names = {'token_embed', 'loop_step_embed', 'pre_loop_norm',
                      'final_norm', 'ln1', 'ln2'}

    decay: list[torch.nn.Parameter] = []
    no_decay: list[torch.nn.Parameter] = []

    for name, param in model.named_parameters():
        if not param.requires_grad:
            continue
        is_nodecay = (
            param.ndim < 2
            or any(nd in name for nd in no_decay_names)
        )
        (no_decay if is_nodecay else decay).append(param)

    return torch.optim.AdamW(
        [
            {'params': decay,    'weight_decay': cfg.weight_decay},
            {'params': no_decay, 'weight_decay': 0.0},
        ],
        lr=cfg.phase1_peak_lr,
        betas=(cfg.optimizer_beta1, cfg.optimizer_beta2),
        eps=cfg.optimizer_eps,
        fused=True   # Requires PyTorch 2.0+ — faster on CUDA
    )


def get_lr(
    step: int,
    warmup_steps: int,
    total_steps: int,
    peak_lr: float,
    min_lr: float
) -> float:
    """
    Warmup + cosine decay learning rate schedule.

    Args:
        step: Current training step
        warmup_steps: Number of linear warmup steps
        total_steps: Total training steps
        peak_lr: Peak learning rate
        min_lr: Minimum learning rate

    Returns:
        Learning rate for the current step
    """
    if step < warmup_steps:
        return peak_lr * (step / max(1, warmup_steps))
    if step > total_steps:
        return min_lr
    progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
    return min_lr + 0.5 * (peak_lr - min_lr) * (1.0 + math.cos(math.pi * progress))


def get_ponder_tau(step: int, cfg: AXIOM1Config) -> float:
    """
    Gradually ramp ponder cost from 0 to target over Phase 3.

    This avoids destabilizing the already-trained base model.
    Ramps linearly from 0 to 0.01 over the first 20% of Phase 3 steps.

    Args:
        step: Current training step within Phase 3
        cfg: AXIOM1Config instance

    Returns:
        Ponder cost weight (tau) for the current step
    """
    phase3_start = 0      # start of Phase 3 (step counter resets)
    ramp_end = int(cfg.phase3_total_steps * 0.2)   # ramp over first 20%
    target_tau = 0.01

    if step < phase3_start:
        return 0.0
    progress = min((step - phase3_start) / max(1, ramp_end), 1.0)
    return progress * target_tau
