#!/usr/bin/env python
"""Process-supervised PPO entry: gates -> worker -> banner -> train -> stage manifest.

    CUDA_VISIBLE_DEVICES=0,1,2 torchrun --nproc_per_node=3 scripts/ppo.py \
        configs/ppo_stage.toml
    # tiny end-to-end smoke (3 prompts x 4 samples x 256 tokens x 4 iterations):
    CUDA_VISIBLE_DEVICES=0,1,2 torchrun --nproc_per_node=3 scripts/ppo.py \
        configs/ppo_stage.toml --smoke

Ported from mistral-prm-ppo `scripts/train_ppo.py` and re-seated on the align stack the way
scripts/grpo.py is: stage config -> PPOTrainModule -> the shared Trainer. What replaced what:

  peft LoRA            -> align.ppo.lora (named adapters on the SAME resident base)
  DeepSpeed ZeRO-3     -> FSDP2 (align.train.fsdp); ZeRO-3 + Adafactor silently did not train
  HF generate rollout  -> align.rollout.client over the file protocol (vLLM worker, own venv)
  transformers config  -> align.model.hf_config.read_config (own config.json reader)
  torch DDP of LoRA    -> not needed: FSDP shards the frozen base, adapters ride the same wrap

Topology: N FSDP train ranks + one dedicated engine GPU (default 3, NMO_ROLLOUT_GPU). Rank 0
owns all protocol I/O. Checkpoint saves ARE the weight sync (the engine reloads the latest
checkpoint), and a LoRA run saves MERGED weights (PPOTrainModule.full_model_state_dict), so the
engine loads an ordinary HF-layout dir. Save cadence is DERIVED from the rollout shape, never
typed: a save lands exactly at the end of every sync_every_iterations-th iteration.

Two arms share this one script and differ only in the config's `arm` and reward table:
  arm = "prm"  process supervision (PRM adapter on the shared backbone scores each step)
  arm = "orm"  the MANDATORY control: identical pipeline, process term off. Without it nothing
               is attributable to process supervision (Math-Shepherd's PRM-attributable delta
               over its own ORM-PPO control is +1.7pp MATH, ~2/3 of the headline is plain RL).

Launch gates, ordered by cost: config validation (struct mode, eta, lam_actor, cadence, canary)
-> ceiling receipt (0/100-perfect-policy trap) -> topology check -> eos/boundary derivation
(derived, never hardcoded) -> module construction refusals (reference, scorer) -> horizon
consistency (scheduler vs policy steps).
"""

from __future__ import annotations

import argparse
import json
import logging
import os
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from align.utils.guard import apply_resource_guards  # noqa: E402

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger("ppo")


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("config")
    ap.add_argument("--smoke", action="store_true", help="tiny shapes, 4 iterations")
    ap.add_argument("--band-file", default=None, help="override the configured band file")
    ap.add_argument("--strict-host-check", default=True, action=argparse.BooleanOptionalAction)
    ap.add_argument("--heavy-lock", default=True, action=argparse.BooleanOptionalAction)
    args = ap.parse_args()

    apply_resource_guards(
        cpu_threads=8,
        heavy=args.heavy_lock,
        label="ppo-stage",
        strict_host_check=args.strict_host_check,
    )

    from align.utils.dist import init_distributed
    from align.ppo.stage_config import load_stage_config as _early_stage_config

    # Run dir from the config (parsed before the process group: a bad TOML fails
    # in milliseconds, and the NCCL flight-recorder dump lands under it).
    init_distributed(run_dir=_early_stage_config(args.config, smoke=args.smoke).run_dir)

    from align.train.stage_gates import assert_distinct_devices

    assert_distinct_devices()

    import torch
    import torch.distributed as dist

    from align.data.schema import build_prompt, extract_answer, register_answer_suffixes
    from align.data.sources import DataSources, gsm8k_gold_answer, read_gsm8k_test
    from align.data.tokenization import load_boundary_spec
    from align.eval.verifier import OutcomeVerifier, VerifierConfig
    from align.model import checkpoint as ckpt
    from align.model.mistral import NmoMistralForCausalLM
    from align.ppo.lora import inject_lora, mark_only_adapter_trainable
    from align.ppo.stage_config import load_stage_config
    from align.ppo.trainer import (
        AdapterProbeView,
        PPOBatchSource,
        PPOTrainModule,
        sync_params_from_rank0,
    )
    from align.ppo.value_head import PRMConfig, ProcessRewardScorer, ValueHead, load_prm_state
    from align.rl.canary import TaskAccuracyCanary
    from align.rollout.client import RolloutClient, spawn_worker, terminate_worker
    from align.rollout.protocol import RolloutRequest
    from align.train.dpo_stage import resolve_sft_checkpoint
    from align.train.trainer import Trainer, WarmupCosineSchedule
    from align.utils.guard import assert_paths_off_root

    rank = dist.get_rank() if dist.is_initialized() else 0
    world = dist.get_world_size() if dist.is_initialized() else 1

    cfg = load_stage_config(args.config, smoke=args.smoke)
    if args.band_file:
        cfg.data.band_file = args.band_file
    assert_paths_off_root([cfg.run_dir])
    ceiling = cfg.validate_ceiling_receipt()
    cfg.validate_world(world)

    register_answer_suffixes("</s>")  # BEFORE any extraction (golden.py's incident)

    # -- data: the band's prompt rows ---------------------------------------
    prompt_rows = []
    with open(cfg.data.band_file) as f:
        for line in f:
            if line.strip():
                prompt_rows.append(json.loads(line))
    logger.info("band file %s: %d prompts", cfg.data.band_file, len(prompt_rows))

    # -- derived ids: the boundary and EOS are READ, never hardcoded ----------
    spec = load_boundary_spec(cfg.boundary_spec_path)
    sft_checkpoint, ckpt_source = resolve_sft_checkpoint(cfg.sft_checkpoint, None)
    logger.info("substrate: %s (%s)", sft_checkpoint, ckpt_source)
    from align.model.hf_config import read_config
    from align.tokenizer import load_tokenizer

    model_cfg = read_config(sft_checkpoint)
    tokenizer = load_tokenizer(sft_checkpoint)
    pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0
    with open(os.path.join(sft_checkpoint, "config.json")) as f:
        snapshot_eos = json.load(f).get("eos_token_id")
    if not (cfg.ppo.eos_token_id == snapshot_eos == tokenizer.eos_token_id):
        raise RuntimeError(
            f"EOS ids disagree: ppo.eos_token_id={cfg.ppo.eos_token_id}, substrate "
            f"config.json={snapshot_eos}, tokenizer={tokenizer.eos_token_id}. Truncation would be "
            "mis-flagged: GAE would bootstrap terminal rows (or treat cut-off rows as terminal). "
            "Derived-vs-hardcoded incident class."
        )

    # -- policy: LoRA on the shared base ---------------------------------------
    model = NmoMistralForCausalLM.from_state_dict(
        model_cfg,
        ckpt.hf_to_nmo(ckpt.load_hf_safetensors(sft_checkpoint)),
        attn_implementation=cfg.train.attn_implementation,
        device="cuda",
    ).to("cuda", dtype=torch.bfloat16)

    # Identical adapter init on every rank: FSDP2 shards what each rank already holds and does
    # NOT broadcast, so a per-rank random A would be a different adapter per shard.
    gen = torch.Generator().manual_seed(cfg.train.seed)
    inject_lora(model, cfg.lora, adapter="policy", generator=gen)
    scorer = None
    if cfg.arm == "prm":
        from align.ppo.lora import LoRAConfig

        inject_lora(
            model,
            LoRAConfig(
                r=cfg.prm.lora_r, alpha=cfg.prm.lora_alpha, target_modules=cfg.lora.target_modules
            ),
            adapter="prm",
            generator=gen,
        )
        scorer = ProcessRewardScorer(
            model,
            model_cfg.hidden_size,
            PRMConfig(
                head_type=cfg.prm.head_type,
                adapter="prm",
                label_token_ids=cfg.prm.label_token_ids,
            ),
        )
        load_prm_state(scorer, cfg.prm.checkpoint)
        scorer.to("cuda")
        for p in scorer.parameters():
            p.requires_grad_(False)
    n_trainable = mark_only_adapter_trainable(model, "policy")  # AFTER both adapters exist
    logger.info("policy adapter: %d trainable parameters", n_trainable)

    from align.train.fsdp import apply_activation_checkpointing, apply_fsdp

    model = apply_activation_checkpointing(model)
    if world > 1:
        model = apply_fsdp(model)
    logger.info("post-shard cuda allocated: %.2f GiB", torch.cuda.memory_allocated() / 2**30)

    value_head = ValueHead(
        model_cfg.hidden_size, detach_hidden=not cfg.train.value_shares_backbone
    ).to("cuda")
    sync_params_from_rank0(value_head)

    # -- rollout worker (rank 0 owns it) ------------------------------------------
    spool = os.path.join(cfg.run_dir, "spool")
    worker_proc = None
    if rank == 0:
        worker_proc = spawn_worker(
            spool,
            gpu=int(os.environ.get("NMO_ROLLOUT_GPU", "3")),
            log_path=os.path.join(cfg.run_dir, "worker.log"),
        )
        logger.info("rollout worker spawned (pid %d)", worker_proc.pid)
    client = RolloutClient(spool, timeout_s=cfg.rollout_timeout_s)

    # -- module + trainer ----------------------------------------------------------
    from align.train.offload_adamw import make_optimizer

    source = PPOBatchSource(
        prompt_rows,
        prompts_per_iteration=cfg.rollout.prompts_per_iteration,
        n_samples_per_prompt=cfg.rollout.n_samples_per_prompt,
        max_response_len=cfg.rollout.max_response_len,
        temperature=cfg.rollout.temperature,
        iterations=cfg.train.iterations,
        ppo_cfg=cfg.ppo,
        seed=cfg.train.seed,
        prompt_builder=build_prompt,
        world_size=world,
    )
    run_cfg = _run_config(cfg, model_cfg, world, source.micros_per_iteration)
    policy_params = [p for p in model.parameters() if p.requires_grad]
    optimizer = make_optimizer(run_cfg, policy_params)
    value_optimizer = torch.optim.AdamW(
        value_head.parameters(), lr=cfg.train.lr * cfg.train.value_lr_mult, weight_decay=0.0
    )

    def _sum_across_ranks(t):
        if dist.is_initialized() and world > 1:
            t = t.clone()
            dist.all_reduce(t, op=dist.ReduceOp.SUM)
        return t

    def _mean_across_ranks(g):
        if dist.is_initialized() and world > 1:
            g = g.clone()
            dist.all_reduce(g, op=dist.ReduceOp.AVG)
        return g

    score_fn = _make_score_fn(scorer, "cuda") if scorer is not None else None
    module = PPOTrainModule(
        model,
        value_head,
        optimizer,
        value_optimizer,
        cfg.ppo,
        score_steps_fn=score_fn,
        reduce_fn=_sum_across_ranks if world > 1 else None,
        value_grad_reduce=_mean_across_ranks if world > 1 else None,
    )
    # HORIZON: warmup iterations move only the critic, so the schedule spans POLICY steps.
    scheduler = WarmupCosineSchedule(
        base_lr=cfg.train.lr,
        total_steps=source.policy_optimizer_steps,
        warmup_frac=cfg.train.warmup_frac,
    )
    trainer = Trainer(
        run_cfg,
        AdapterProbeView(model, "policy"),
        source,
        module._ppo_loss_adapter,
        run_dir=cfg.run_dir,
        stage="ppo",
        train_module=module,
        scheduler=scheduler,
    )
    # The Trainer logs scheduler.lr() but (in this tree) does not write it into the optimizer; the
    # module applies it at every step so the logged LR is the LR that ran.
    module.lr_schedule = lambda: trainer.scheduler.lr()
    if trainer.scheduler.total_steps != source.policy_optimizer_steps:
        raise RuntimeError(
            f"horizon mismatch: scheduler={trainer.scheduler.total_steps} vs policy optimizer "
            f"steps={source.policy_optimizer_steps}"
        )

    # -- canary (R4, rank 0) -----------------------------------------------------------
    canary = None
    if rank == 0:
        sources = DataSources.from_local_hub()
        gsm_rows = read_gsm8k_test(sources)[: cfg.canary.n_problems]
        canary_problems = [
            {"prompt": build_prompt(r["question"]), "gold": gsm8k_gold_answer(r)} for r in gsm_rows
        ]

        def canary_generate(prompts: list[str]) -> list[str]:
            req = RolloutRequest(
                request_id=f"canary-{trainer.optimizer_step}-{os.getpid()}",
                weights_dir=_latest_weights(trainer, sft_checkpoint),
                prompts=prompts,
                n_samples_per_prompt=1,
                max_tokens=cfg.rollout.max_response_len,
                temperature=0.0,
            )
            client.submit(req)
            result = client.collect(req.request_id)
            return [texts[0] for texts in result.texts]

        canary = TaskAccuracyCanary(
            problems=canary_problems,
            generate_fn=canary_generate,
            verifier=OutcomeVerifier(VerifierConfig(workers=4)),
            extract_fn=extract_answer,
            every=cfg.canary.every,
        )

    source.bind(
        client=client,
        verifier=OutcomeVerifier(VerifierConfig(workers=4)) if rank == 0 else None,
        module=module,
        weights_dir_fn=lambda: _latest_weights(trainer, sft_checkpoint),
        newline_id=spec.newline_id,
        pad_token_id=pad_id,
        rank=rank,
        world_size=world,
        trainer=trainer,
        canary=canary,
        canary_log_path=os.path.join(cfg.run_dir, "canary.jsonl"),
    )
    trainer.loss_adapter = source.lazy_adapter(module)

    banner = {
        "stage": f"ppo_{cfg.arm}" + ("_smoke" if args.smoke else ""),
        "arm": cfg.arm,
        "substrate": sft_checkpoint,
        "band_file": cfg.data.band_file,
        "band_prompts": len(prompt_rows),
        "shape": {
            "prompts_per_iteration": cfg.rollout.prompts_per_iteration,
            "n_samples_per_prompt": cfg.rollout.n_samples_per_prompt,
            "max_response_len": cfg.rollout.max_response_len,
            "iterations": cfg.train.iterations,
            "ppo_epochs": cfg.ppo.ppo_epochs,
            "mini_batch_rows": cfg.ppo.mini_batch_rows,
            "micro_batch_rows": cfg.ppo.micro_batch_rows,
            "grad_accum": cfg.ppo.grad_accum,
            "micros_per_iteration": source.micros_per_iteration,
            "policy_optimizer_steps": source.policy_optimizer_steps,
            "value_warmup_iterations": cfg.ppo.value_warmup_iterations,
        },
        "reward": {
            "transform": cfg.ppo.reward.transform,
            "eta": cfg.ppo.reward.eta,
            "w_prm": cfg.ppo.reward.w_prm,
            "w_outcome": cfg.ppo.reward.w_outcome,
            "kl_coef(reward)": cfg.ppo.reward.kl_coef,
            "kl_loss_coef": cfg.ppo.kl_loss_coef,
            "lam_actor": cfg.ppo.reward.lam_actor,
            "lam_critic": cfg.ppo.reward.lam_critic,
        },
        "lora": {"r": cfg.lora.r, "alpha": cfg.lora.alpha, "trainable": n_trainable},
        "ceiling": {k: v["verify_correct_rate"] for k, v in ceiling["sources"].items()},
        "boundary": {"newline_id": spec.newline_id, "tokenizer_hash": spec.tokenizer_hash},
        "topology": {"train_ranks": world, "rollout_gpu": os.environ.get("NMO_ROLLOUT_GPU", "3")},
        "rectifications": (
            f"R3 clip={cfg.ppo.max_grad_norm:.1f}, R4 canary every={cfg.canary.every}, "
            "R5 pin-on-alarm, explicit minibatch denominators, ratio tripwire armed+self-tested"
        ),
    }
    logger.info("BANNER %s", json.dumps(banner, indent=1, default=str))

    # -- train ------------------------------------------------------------------------------
    try:
        # RESUME. Without this a relaunch silently restarts from the SFT substrate with a fresh
        # optimizer while the run dir holds the real checkpoints. Resume BEFORE the first
        # rollout: the rollout engine's weights come from `_latest_weights`, which reads the
        # pointer restored here.
        resumed = _resume(trainer)
        if resumed:
            logger.info("RESUMED PPO run from %s", resumed)
        summary = trainer.train()
        checkpoint = trainer.save_checkpoint()
    finally:
        if worker_proc is not None:
            terminate_worker(worker_proc)
    if trainer.stop_exit_code:
        # A coordinated stop (SIGTERM / train.max_wall_seconds) saved at an
        # optimizer-step boundary; the stage is NOT finished. Exit with the
        # documented status instead of evaluating / reporting a half-trained arm.
        raise SystemExit(trainer.stop_exit_code)

    result = {
        "stage": banner["stage"],
        "summary": summary,
        "checkpoint": checkpoint,
        "banner": banner,
        "canary_history": canary.history if canary else [],
        "alarms": source.alarms,
        "ratio_tripwire": module.ratio_tripwire.as_log_dict(),
        "ratio_tripwire_alarms": module.ratio_tripwire.alarms,
    }
    if rank == 0:
        path = os.path.join(cfg.run_dir, "stage_result.json")
        with open(path + ".partial", "w") as f:
            json.dump(result, f, indent=1, default=str)
        os.replace(path + ".partial", path)
        logger.info("stage result: %s", path)
    return 0


def _resume(trainer) -> str | None:
    """Resume from the newest checkpoint, if any. Returns its path (None = fresh run).

    THE STALE-SUBSTRATE DEFECT (closed in `Trainer.resume_from`, pinned here). `checkpoints_written`
    is a per-process list that `_latest_weights` keys off; it used to be empty after a resume, so
    the first post-resume rollout was generated by the SFT SUBSTRATE -- the policy before every
    update -- while the trainer computed ratios against the resumed policy. The Trainer now
    registers the resumed checkpoint itself (generic: GRPO has the same `_latest_weights` shape and
    keep-N pruning needed it too), so this wrapper no longer patches the list from the outside."""
    return trainer.resume_from_latest()


def _latest_weights(trainer, sft_checkpoint: str) -> str:
    """The engine's weights source: latest saved (merged, HF-layout) checkpoint, else the
    substrate. During value warmup the policy has not moved, so the substrate IS the policy."""
    if trainer.checkpoints_written:
        return os.path.join(trainer.ckpt_root, f"micro_{trainer.checkpoints_written[-1]:08d}")
    return sft_checkpoint


def _make_score_fn(scorer, device: str, chunk_rows: int = 8):
    """PRM scores for a `PPORollout`: P(step correct) `[B, S]`, chunked so the scoring forward
    stays inside the activation budget. The rollout stores delimiter ACTIONS in the label frame;
    the PRM reads the hidden state AT the delimiter token, one input position later."""
    import torch

    @torch.no_grad()
    def score(rollout):
        B = rollout.n_rows
        out = torch.zeros(rollout.step_positions.shape, dtype=torch.float32)
        for lo in range(0, B, chunk_rows):
            sl = slice(lo, min(B, lo + chunk_rows))
            T = int(rollout.attention_mask[sl].sum(-1).max())
            p = scorer.score_steps(
                rollout.input_ids[sl, :T].to(device),
                rollout.attention_mask[sl, :T].to(device),
                (rollout.step_positions[sl] + 1).to(device),
                rollout.step_mask[sl].to(device),
            )
            out[sl] = p.float().cpu()
        return out

    return score


def _run_config(cfg, model_cfg, world_size: int, micros_per_iteration: int):
    """The Trainer's RunConfig, built FROM the stage config. Pure and import-light so the
    config-translation test calls it directly and pins every crossing field (nmo silent-failure
    #6: a field that fails to cross trains at the trainer's default while the banner shows the
    pre-registered number)."""
    from align.config import RunConfig

    run = RunConfig()
    run.model = model_cfg
    run.optimizer.name = cfg.train.optimizer_name
    run.optimizer.lr = cfg.train.lr
    run.optimizer.grad_accum = cfg.ppo.grad_accum
    run.train.world_size = world_size
    run.train.epochs = (
        1  # the horizon is iterations x epochs; the source enumerates its own handles
    )
    run.train.micro_batch_size = cfg.ppo.micro_batch_rows
    run.train.save_every_micro_steps = micros_per_iteration * cfg.train.sync_every_iterations
    run.train.keep_checkpoints = cfg.train.keep_checkpoints
    run.train.warmup_frac = cfg.train.warmup_frac
    run.train.full_ft = False  # LoRA: the offload-economics accum>=4 rule is a full-FT rule
    run.train.seed = cfg.train.seed
    run.train.max_wall_seconds = cfg.train.max_wall_seconds
    run.train.max_consecutive_skipped_steps = cfg.train.max_consecutive_skipped_steps
    run.train.max_skipped_fraction = cfg.train.max_skipped_fraction
    run.tripwires.armed = tuple(cfg.armed_tripwires)
    run.run_name = f"ppo_{cfg.arm}"
    return run


if __name__ == "__main__":
    from align.exit_codes import run_with_exit_codes

    raise SystemExit(run_with_exit_codes(main))