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

"""
Main training entry point for AXIOM-1.

Handles Phase 1, 2, and 3 based on configuration.
Always uses bfloat16 autocast (never float16).
Gradient clipping at 1.0 every step (no exceptions — Rule 7).

Usage:
    python -m axiom1.training.train --phase 1 --data train_tokens.pt
"""

import argparse
import math
import os
import time
from pathlib import Path

import torch
import torch.nn.functional as F
from torch.cuda.amp import GradScaler

from axiom1.model.config import AXIOM1Config
from axiom1.model.model import AXIOM1
from axiom1.training.optimizer import build_optimizer, get_lr, get_ponder_tau
from axiom1.training.loss import compute_ce_loss, compute_consistency_loss
from axiom1.training.data import build_dataloader
from axiom1.training.phases import (
    get_phase_config, PhaseConfig,
    save_checkpoint, load_checkpoint,
    transition_phase1_to_phase2,
)
from axiom1.evals.sanity_check import run_sanity_check
from axiom1.evals.perplexity import evaluate


def train(
    model: AXIOM1,
    train_dataloader: torch.utils.data.DataLoader,
    val_dataloader: torch.utils.data.DataLoader | None,
    cfg: AXIOM1Config,
    phase_cfg: PhaseConfig,
    device: str = 'cuda'
) -> None:
    """
    Main training loop for AXIOM-1.

    Handles Phase 1, 2, and 3 based on phase_cfg settings.
    Features:
      - bfloat16 autocast (never float16)
      - Gradient accumulation
      - Gradient clipping at 1.0
      - Warmup + cosine decay LR schedule
      - Periodic validation and checkpointing
      - Consistency loss (Phase 2 only)
      - Ponder cost (Phase 3 only)

    Args:
        model: AXIOM1 model instance
        train_dataloader: Training DataLoader
        val_dataloader: Optional validation DataLoader
        cfg: AXIOM1Config instance
        phase_cfg: PhaseConfig for the current training phase
        device: Device to train on
    """
    model = model.to(device)

    # Build optimizer with correct param groups
    optimizer = build_optimizer(model, cfg)
    scaler = GradScaler()

    step = 0
    accum_loss = 0.0
    t0 = time.time()

    model.train()

    print(f"\n{'='*60}")
    print(f"[AXIOM-1] Phase {phase_cfg.phase} Training Start")
    print(f"  Loops: {phase_cfg.num_loops} | Seq: {phase_cfg.seq_len} | "
          f"LR: {phase_cfg.peak_lr:.0e}")
    print(f"  Steps: {phase_cfg.total_steps} | Batch: {phase_cfg.batch_size} | "
          f"Accum: {phase_cfg.grad_accum}")
    print(f"  Consistency: {phase_cfg.consistency_weight} | "
          f"ADC: {phase_cfg.use_adaptive_depth}")
    print(f"{'='*60}\n")

    # Override num_loops for this phase
    original_num_loops = model.cfg.num_loops
    model.cfg.num_loops = phase_cfg.num_loops

    epoch = 0
    while step < phase_cfg.total_steps:
        epoch += 1
        for batch in train_dataloader:
            input_ids = batch['input_ids'].to(device)

            # ── Forward pass ────────────────────────────────────────
            with torch.autocast('cuda', dtype=torch.bfloat16):

                if phase_cfg.consistency_weight > 0.0 and model.training:
                    logits, intermediates = model(
                        input_ids, return_intermediates=True
                    )
                    ce_loss = compute_ce_loss(logits, input_ids)
                    cons_loss = compute_consistency_loss(
                        model, intermediates, input_ids,
                        start_loop=cfg.consistency_start_loop
                    )
                    loss = ce_loss + phase_cfg.consistency_weight * cons_loss

                elif phase_cfg.use_adaptive_depth and model.training:
                    # Phase 3 — ADC with ponder cost
                    h = model.pre_loop_norm(
                        model.token_embed(input_ids) * (cfg.d_model ** 0.5)
                    )
                    # ADC forward would go here when implemented
                    logits = model(input_ids)
                    ce_loss = compute_ce_loss(logits, input_ids)
                    tau = get_ponder_tau(step, cfg)
                    loss = ce_loss  # + tau * ponder_cost when ADC is implemented

                else:
                    logits = model(input_ids)
                    loss = compute_ce_loss(logits, input_ids)

                loss = loss / phase_cfg.grad_accum

            # ── Backward ────────────────────────────────────────────
            scaler.scale(loss).backward()
            accum_loss += loss.item()

            # ── Optimizer step (every grad_accum steps) ─────────────
            if (step + 1) % phase_cfg.grad_accum == 0:
                scaler.unscale_(optimizer)
                grad_norm = torch.nn.utils.clip_grad_norm_(
                    model.parameters(), cfg.grad_clip
                )

                # Update LR
                lr = get_lr(
                    step, phase_cfg.warmup_steps, phase_cfg.total_steps,
                    phase_cfg.peak_lr, phase_cfg.min_lr
                )
                for pg in optimizer.param_groups:
                    pg['lr'] = lr

                scaler.step(optimizer)
                scaler.update()
                optimizer.zero_grad()

                # ── Logging ─────────────────────────────────────────
                if step % 100 == 0:
                    elapsed = time.time() - t0
                    print(
                        f"[AXIOM-1] step={step:>6} | "
                        f"loss={accum_loss * phase_cfg.grad_accum:.4f} | "
                        f"grad_norm={grad_norm:.3f} | "
                        f"lr={lr:.2e} | "
                        f"t={elapsed:.1f}s"
                    )
                    accum_loss = 0.0
                    t0 = time.time()

                # ── Validation ──────────────────────────────────────
                if step % 2000 == 0 and val_dataloader is not None:
                    val_loss = evaluate(model, val_dataloader, cfg)
                    val_ppl = math.exp(min(val_loss, 100.0))
                    print(
                        f"[AXIOM-1] VAL step={step} | "
                        f"loss={val_loss:.4f} | ppl={val_ppl:.1f}"
                    )
                    model.train()

                # ── Checkpointing ───────────────────────────────────
                if step % 10_000 == 0 and step > 0:
                    save_checkpoint(
                        model, optimizer, step,
                        accum_loss * phase_cfg.grad_accum, cfg,
                        f'axiom1/checkpoints/phase{phase_cfg.phase}/step{step}.pt'
                    )

            step += 1
            if step >= phase_cfg.total_steps:
                break

    # Restore original num_loops
    model.cfg.num_loops = original_num_loops

    # Final checkpoint
    save_checkpoint(
        model, optimizer, step,
        accum_loss * phase_cfg.grad_accum, cfg,
        f'axiom1/checkpoints/phase{phase_cfg.phase}/final.pt'
    )

    print(f"\n[AXIOM-1] Phase {phase_cfg.phase} training complete. "
          f"Total steps: {step}")


def main() -> None:
    """Command-line entry point for training."""
    parser = argparse.ArgumentParser(
        description="AXIOM-1 Training | EGen Labs"
    )
    parser.add_argument(
        '--phase', type=int, required=True, choices=[1, 2, 3],
        help='Training phase (1=warm start, 2=full depth, 3=ADC)'
    )
    parser.add_argument(
        '--train-data', type=str, required=True,
        help='Path to pre-tokenized training data (.pt file)'
    )
    parser.add_argument(
        '--val-data', type=str, default=None,
        help='Path to pre-tokenized validation data (.pt file)'
    )
    parser.add_argument(
        '--checkpoint', type=str, default=None,
        help='Path to checkpoint to resume from (required for Phase 2+)'
    )
    parser.add_argument(
        '--device', type=str, default='cuda',
        help='Device to train on'
    )
    args = parser.parse_args()

    cfg = AXIOM1Config()
    phase_cfg = get_phase_config(cfg, args.phase)

    # Build model
    model = AXIOM1(cfg)

    # Handle phase transitions
    if args.phase == 2:
        if args.checkpoint is None:
            raise ValueError(
                "Phase 2 requires a Phase 1 checkpoint. "
                "Use --checkpoint checkpoints/phase1/final.pt"
            )
        model = transition_phase1_to_phase2(model, cfg, args.checkpoint)

    elif args.phase == 1 and args.checkpoint is not None:
        load_checkpoint(args.checkpoint, model)

    # Run sanity check (Rule 1: always, before every phase)
    print("\n[AXIOM-1] Running mandatory pre-training sanity check...")
    # Temporarily set num_loops for sanity check
    original_loops = model.cfg.num_loops
    model.cfg.num_loops = phase_cfg.num_loops
    run_sanity_check(model, cfg, device=args.device)
    model.cfg.num_loops = original_loops

    # Build dataloaders
    train_dl = build_dataloader(
        args.train_data, cfg,
        seq_len=phase_cfg.seq_len,
        batch_size=phase_cfg.batch_size
    )
    val_dl = None
    if args.val_data:
        val_dl = build_dataloader(
            args.val_data, cfg,
            seq_len=phase_cfg.seq_len,
            batch_size=phase_cfg.batch_size,
            shuffle=False
        )

    # Train
    train(model, train_dl, val_dl, cfg, phase_cfg, device=args.device)


if __name__ == "__main__":
    main()
