"""Stage 2 settings, learning-rate schedule and checkpoint protection.

Pure Python (no torch), so every rule here can be unit-tested on a laptop.
Stage 2 is switched on only by environment variables (exported from run_args.sh);
when STAGE2_START is not set, everything behaves exactly like Stage 1.
"""
import math
import os
from dataclasses import dataclass

# end-of-text id of our tokenizer: the data files separate documents with it
EOT_ID = 151645


@dataclass
class Stage2Settings:
    start: int            # first Stage 2 step (Stage 1's last step is start - 1)
    warmup: int           # steps to climb from Stage 1's final LR to `peak`
    peak: float
    min_lr: float
    bins: list            # Stage 2 data files
    weights: list         # their mix weights (same order)

    @staticmethod
    def from_env():
        """Read Stage 2 settings from the environment; None means 'no Stage 2'."""
        if not os.environ.get("STAGE2_START"):
            return None
        bins = os.environ["S2_BIN"].split()
        weights = [float(w) for w in os.environ["S2_WEIGHTS"].split()]
        assert len(bins) == len(weights) > 0, "S2_BIN and S2_WEIGHTS must have the same length"
        assert all(w >= 0 for w in weights) and sum(weights) > 0, "bad S2_WEIGHTS"
        s = Stage2Settings(
            start=int(os.environ["STAGE2_START"]),
            warmup=int(os.environ.get("STAGE2_WARMUP", "1000")),
            peak=float(os.environ.get("STAGE2_PEAK", "1e-4")),
            min_lr=float(os.environ.get("STAGE2_MIN", "1e-5")),
            bins=bins,
            weights=weights,
        )
        assert s.start > 0 and s.warmup > 0 and s.peak > s.min_lr > 0, "bad Stage 2 settings"
        return s


def stage1_lr(step, cfg):
    """Stage 1 schedule, unchanged: linear warmup, then cosine to min_lr at lr_horizon."""
    if step < cfg.warmup_steps:
        return cfg.peak_lr * (step + 1) / cfg.warmup_steps
    horizon = getattr(cfg, "lr_horizon", cfg.max_steps)
    progress = min(1.0, (step - cfg.warmup_steps) / max(1, horizon - cfg.warmup_steps))
    return cfg.min_lr + 0.5 * (1 + math.cos(math.pi * progress)) * (cfg.peak_lr - cfg.min_lr)


def lr_at(step, cfg, s2=None):
    """Stage 1 schedule before s2.start; after it: warm up from Stage 1's final LR
    to s2.peak, then cosine down to s2.min_lr at the last step (cfg.max_steps)."""
    if s2 is None or step < s2.start:
        return stage1_lr(step, cfg)
    t = step - s2.start
    base = stage1_lr(s2.start - 1, cfg)          # where Stage 1 ended
    if t < s2.warmup:
        return base + (s2.peak - base) * (t + 1) / s2.warmup
    span = max(1, cfg.max_steps - 1 - s2.start - s2.warmup)
    progress = min(1.0, (t - s2.warmup) / span)
    return s2.min_lr + 0.5 * (1 + math.cos(math.pi * progress)) * (s2.peak - s2.min_lr)


def protected_ckpts():
    """Checkpoint folder names that rotation must never delete (env PROTECT_CKPTS)."""
    return set(os.environ.get("PROTECT_CKPTS", "").split())


def ckpts_to_delete(entries, is_complete, keep, protect):
    """Which ckpt_<step> folders rotation removes: all half-written ones, plus complete
    ones older than the newest `keep`. Protected names are never removed or counted."""
    entries = [d for d in entries if d.startswith("ckpt_") and d not in protect]
    good = sorted((d for d in entries if is_complete(d)), key=lambda d: int(d.split("_")[1]))
    bad = [d for d in entries if not is_complete(d)]
    return (good[:-keep] if keep > 0 else good) + bad