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

"""
Loss functions for AXIOM-1 training.

Two losses:
  1. CE loss — standard autoregressive cross-entropy (all phases)
  2. Consistency loss — penalizes second-half loops for not predicting well
     (Phase 2 only, from loop num_loops//2 = 6 onward, weight=0.25)
"""

import torch
import torch.nn.functional as F

from axiom1.model.model import AXIOM1


def compute_ce_loss(
    logits: torch.Tensor,
    input_ids: torch.Tensor,
    ignore_index: int = -100
) -> torch.Tensor:
    """
    Standard autoregressive cross-entropy loss.

    Shifts logits and labels by 1 position for next-token prediction.

    Args:
        logits: Model output logits (batch, seq_len, vocab_size)
        input_ids: Input token IDs (batch, seq_len)
        ignore_index: Label value to ignore in loss computation

    Returns:
        Scalar CE loss
    """
    shift_logits = logits[:, :-1, :].contiguous()
    shift_labels = input_ids[:, 1:].contiguous()
    return F.cross_entropy(
        shift_logits.view(-1, shift_logits.size(-1)),
        shift_labels.view(-1),
        ignore_index=ignore_index
    )


def compute_consistency_loss(
    model: AXIOM1,
    intermediates: list[torch.Tensor],
    input_ids: torch.Tensor,
    start_loop: int = 6
) -> torch.Tensor:
    """
    Loop consistency loss: penalizes second-half loops for not already predicting well.

    Encourages all loop steps to contribute, not just the final one.
    Applied from loop `start_loop` to end (default: loop 6 of 12).

    Each intermediate hidden state is passed through final_norm + lm_head
    and compared against the target labels. Later loops are weighted more
    heavily via linear weighting.

    Args:
        model: AXIOM1 model instance (for final_norm and lm_head)
        intermediates: List of hidden states from each loop step
        input_ids: Input token IDs (batch, seq_len)
        start_loop: First loop index to include in consistency loss

    Returns:
        Scalar consistency loss
    """
    shift_labels = input_ids[:, 1:].contiguous().view(-1)
    selected = intermediates[start_loop:]
    N = len(selected)

    if N == 0:
        return torch.tensor(0.0, device=input_ids.device)

    total = torch.tensor(0.0, device=input_ids.device)

    for i, h in enumerate(selected):
        weight = (i + 1) / N   # linear: later loops weighted more
        with torch.no_grad():
            h_normed = model.final_norm(h)
        logits_i = model.lm_head(h_normed)
        loss_i = F.cross_entropy(
            logits_i[:, :-1, :].contiguous().view(-1, logits_i.size(-1)),
            shift_labels,
            ignore_index=-100
        )
        total = total + weight * loss_i

    return total / N


def get_consistency_weight(
    step: int,
    full_weight: float,
    warmup_steps: int
) -> float:
    """
    Compute the consistency loss weight with linear warmup.

    Gradually ramps from 0 to full_weight over warmup_steps,
    preventing abrupt onset of the auxiliary loss which can
    destabilize early training.

    Args:
        step: Current training step
        full_weight: Target consistency weight (e.g., 0.25)
        warmup_steps: Number of steps for linear ramp (e.g., 5000)

    Returns:
        Current consistency weight (0.0 to full_weight)
    """
    if warmup_steps <= 0:
        return full_weight
    ramp = min(1.0, step / warmup_steps)
    return full_weight * ramp
