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

"""
Phase transition logic for AXIOM-1 three-phase training protocol.

Phase 1 → Phase 2: Expand loop step embeddings from 4 to max_loops.
Phase 2 → Phase 3: Enable adaptive depth, adjust LR.

Phase ordering is mandatory. Never skip Phase 1 (Rule 2).
"""

import os
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import torch
import torch.nn as nn

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


@dataclass
class PhaseConfig:
    """Runtime configuration for a training phase."""
    phase: int
    num_loops: int
    seq_len: int
    peak_lr: float
    min_lr: float
    warmup_steps: int
    total_steps: int
    batch_size: int
    grad_accum: int
    consistency_weight: float
    use_adaptive_depth: bool


def get_phase_config(cfg: AXIOM1Config, phase: int) -> PhaseConfig:
    """
    Get the runtime configuration for a specific training phase.

    Args:
        cfg: AXIOM1Config instance
        phase: Phase number (1, 2, or 3)

    Returns:
        PhaseConfig with all phase-specific settings
    """
    if phase == 1:
        return PhaseConfig(
            phase=1,
            num_loops=cfg.phase1_loops,
            seq_len=cfg.phase1_seq_len,
            peak_lr=cfg.phase1_peak_lr,
            min_lr=cfg.phase1_min_lr,
            warmup_steps=cfg.phase1_warmup_steps,
            total_steps=cfg.phase1_total_steps,
            batch_size=cfg.phase1_batch_size,
            grad_accum=cfg.phase1_grad_accum,
            consistency_weight=0.0,          # CE only in Phase 1
            use_adaptive_depth=False,
        )
    elif phase == 2:
        return PhaseConfig(
            phase=2,
            num_loops=cfg.phase2_loops,
            seq_len=cfg.phase2_seq_len,
            peak_lr=cfg.phase2_peak_lr,
            min_lr=cfg.phase2_min_lr,
            warmup_steps=cfg.phase2_warmup_steps,
            total_steps=cfg.phase2_total_steps,
            batch_size=cfg.phase2_batch_size,
            grad_accum=cfg.phase2_grad_accum,
            consistency_weight=cfg.consistency_weight,  # 0.25
            use_adaptive_depth=False,
        )
    elif phase == 3:
        return PhaseConfig(
            phase=3,
            num_loops=cfg.phase3_loops,
            seq_len=cfg.phase3_seq_len,
            peak_lr=cfg.phase3_peak_lr,
            min_lr=cfg.phase3_min_lr,
            warmup_steps=cfg.phase3_warmup_steps,
            total_steps=cfg.phase3_total_steps,
            batch_size=cfg.phase3_batch_size,
            grad_accum=cfg.phase3_grad_accum,
            consistency_weight=0.0,          # Ponder cost replaces consistency
            use_adaptive_depth=True,
        )
    else:
        raise ValueError(f"Unknown phase {phase}. Must be 1, 2, or 3.")


def transition_phase1_to_phase2(
    model: AXIOM1,
    cfg: AXIOM1Config,
    checkpoint_path: str | Path
) -> AXIOM1:
    """
    Transition from Phase 1 to Phase 2.

    Key operations:
      1. Load Phase 1 checkpoint
      2. Expand loop_step_embed from 4 to max_loops entries
         - Preserve entries 0-3 from Phase 1
         - Initialize entries 4+ at std=0.001 (near-zero)
      3. Update num_loops from 4 to 12

    Args:
        model: AXIOM1 model instance
        cfg: AXIOM1Config instance
        checkpoint_path: Path to Phase 1 final checkpoint

    Returns:
        Model ready for Phase 2 training
    """
    # Load Phase 1 checkpoint
    checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=False)
    model.load_state_dict(checkpoint['model_state_dict'], strict=False)
    print(f"[AXIOM-1] Loaded Phase 1 checkpoint from {checkpoint_path}")

    # Expand loop step embeddings from 4 to max_loops
    old_lse = model.loop_step_embed.weight.data.clone()   # shape: (max_loops, 512)
    # But only first 4 were actively used in Phase 1
    phase1_lse = old_lse[:cfg.phase1_loops]                # shape: (4, 512)

    new_lse = nn.Embedding(cfg.max_loops, cfg.d_model)
    nn.init.zeros_(new_lse.weight)
    new_lse.weight.data[:cfg.phase1_loops] = phase1_lse    # preserve Phase 1 embeddings
    nn.init.normal_(new_lse.weight[cfg.phase1_loops:], std=0.001)  # tiny init for new steps

    model.loop_step_embed = new_lse.to(old_lse.device)

    # Update num_loops for Phase 2
    cfg.num_loops = cfg.phase2_loops  # 12
    model.cfg.num_loops = cfg.phase2_loops

    print(f"[AXIOM-1] Phase 1→2 transition complete:")
    print(f"  Loop step embeddings expanded: {cfg.phase1_loops} → {cfg.max_loops}")
    print(f"  num_loops updated: 4 → {cfg.phase2_loops}")
    print(f"  New steps (4-{cfg.max_loops-1}) initialized at std=0.001")

    return model


def transition_phase2_to_phase3(
    model: AXIOM1,
    cfg: AXIOM1Config,
    checkpoint_path: str | Path,
    val_loss: float,
    val_ppl: float
) -> AXIOM1:
    """
    Transition from Phase 2 to Phase 3 (ADC).

    Prerequisites (Rule 8):
      - Phase 2 val perplexity < 50
      - Phase 2 val loss < 4.5

    Args:
        model: AXIOM1 model instance
        cfg: AXIOM1Config instance
        checkpoint_path: Path to Phase 2 final checkpoint
        val_loss: Final Phase 2 validation loss
        val_ppl: Final Phase 2 validation perplexity

    Returns:
        Model ready for Phase 3 training

    Raises:
        RuntimeError: If Phase 2 convergence criteria are not met
    """
    # Check Phase 2 convergence (Rule 8)
    if val_ppl >= 50.0:
        raise RuntimeError(
            f"Phase 2 not converged: val_ppl={val_ppl:.1f} (need < 50). "
            f"Continue Phase 2 training before attempting Phase 3."
        )
    if val_loss >= 4.5:
        raise RuntimeError(
            f"Phase 2 not converged: val_loss={val_loss:.4f} (need < 4.5). "
            f"Continue Phase 2 training before attempting Phase 3."
        )

    # Load Phase 2 checkpoint
    checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=False)
    model.load_state_dict(checkpoint['model_state_dict'])
    print(f"[AXIOM-1] Loaded Phase 2 checkpoint from {checkpoint_path}")

    # Enable adaptive depth
    cfg.use_adaptive_depth = True
    cfg.num_loops = cfg.phase3_loops  # 16
    model.cfg.use_adaptive_depth = True
    model.cfg.num_loops = cfg.phase3_loops

    print(f"[AXIOM-1] Phase 2→3 transition complete:")
    print(f"  ADC enabled, max_loops={cfg.phase3_loops}")
    print(f"  Phase 2 final: loss={val_loss:.4f}, ppl={val_ppl:.1f}")

    return model


def save_checkpoint(
    model: AXIOM1,
    optimizer: torch.optim.Optimizer,
    step: int,
    loss: float,
    cfg: AXIOM1Config,
    path: str | Path
) -> None:
    """
    Save a training checkpoint.

    Args:
        model: AXIOM1 model instance
        optimizer: Optimizer state
        step: Current training step
        loss: Current loss value
        cfg: AXIOM1Config instance
        path: File path to save the checkpoint
    """
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)

    checkpoint = {
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'step': step,
        'loss': loss,
        'config': cfg,
        'model_id': cfg.model_id,
    }
    torch.save(checkpoint, path)
    print(f"[AXIOM-1] Checkpoint saved: {path} (step={step}, loss={loss:.4f})")


def load_checkpoint(
    path: str | Path,
    model: AXIOM1,
    optimizer: torch.optim.Optimizer | None = None
) -> tuple[int, float]:
    """
    Load a training checkpoint.

    Args:
        path: Path to checkpoint file
        model: AXIOM1 model instance to load weights into
        optimizer: Optional optimizer to restore state

    Returns:
        Tuple of (step, loss) from the checkpoint
    """
    checkpoint = torch.load(path, map_location='cpu', weights_only=False)
    model.load_state_dict(checkpoint['model_state_dict'])

    if optimizer is not None and 'optimizer_state_dict' in checkpoint:
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

    step = checkpoint.get('step', 0)
    loss = checkpoint.get('loss', 0.0)
    print(f"[AXIOM-1] Checkpoint loaded: {path} (step={step}, loss={loss:.4f})")
    return step, loss
