#!/usr/bin/env python
"""Precompute DPO reference log-probabilities from the frozen SFT checkpoint.

    # one no-grad pass per (arm, split), single GPU
    uv run python scripts/precompute_ref_logprobs.py configs/dpo_argilla.toml

    # the MANDATORY path-bias dual pass (needs the sharded forward, so torchrun)
    torchrun --nproc_per_node=4 scripts/precompute_ref_logprobs.py \
        configs/dpo_argilla.toml --path-bias

Its own heavy job: it takes the exclusive flock and is SERIALIZED before any
DPO run. Precompute and training must never be concurrent -- a data job
running against a live training job is exactly what killed the predecessor's
first SFT run.

What it produces, per arm, on /apps:
  * `<arm>-train.safetensors` and `<arm>-held.safetensors` -- pair_id ->
    (ref_logp_chosen, ref_logp_rejected, len_chosen, len_rejected), with the
    full six-component cache key in the header. The held-out split is computed
    in the SAME pass because held-out preference accuracy is a primary metric
    and a cache missing it produces an accuracy over a silently smaller
    denominator.
  * `<arm>-path_bias.json` -- the single-GPU-vs-sharded delta distribution and
    the decision derived from it. scripts/dpo.py refuses to launch without it,
    and refuses to launch if it fired and was not honoured.

Why the reference is cached at all: keeping a second 7B model resident would
cost 14.5GB of a 23GB card for no gradient. With dropout disabled a no-grad
pass over a fixed dataset is exactly reproducible, so the numbers can be
computed once and consumed as constants -- see nmo/train/ref_cache.py for the
six things the cache key has to cover for that to stay true.
"""

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("precompute_ref_logprobs")

#: Pairs pushed through both forwards for the path-bias measurement. 100 is
#: the plan's number: enough to see a systematic offset against the
#: early-margin scale, small enough to run inside the pre-launch window.
PATH_BIAS_EXAMPLES = 100


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("config", help="configs/dpo_argilla.toml | dpo_steppairs.toml")
    ap.add_argument(
        "--out-dir", default=None, help="defaults to the config's ref_cache_path directory"
    )
    ap.add_argument(
        "--batch-size",
        type=int,
        default=8,
        help="forward batch size. Regrouping perturbs the fp32 "
        "reduction order by ~1e-7 relative (measured); it is "
        "recorded in the receipt, not in the cache key.",
    )
    ap.add_argument(
        "--path-bias",
        action="store_true",
        help="run the dual single-GPU/sharded pass and write the "
        "decision report. Requires torchrun.",
    )
    ap.add_argument(
        "--sharded-cache",
        action="store_true",
        help="compute the caches under the FSDP2-SHARDED forward (requires "
        "torchrun). This is how a fired path-bias decision "
        "(rerun_precompute_sharded) is honoured: training consumes the "
        "sharded forward, so the reference must too.",
    )
    ap.add_argument(
        "--batch-invariance",
        action="store_true",
        help="measure batched-vs-single-row precompute deltas over "
        "100 pairs and write the receipt that decides whether "
        "batch_size belongs in the cache key. Single process.",
    )
    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="precompute_ref_logprobs",
        strict_host_check=args.strict_host_check,
    )

    import torch

    from align.data.loader import preference_dataset
    from align.model import checkpoint as ckpt
    from align.model.mistral import NmoMistralForCausalLM
    from align.train.dpo_stage import (
        is_hf_checkpoint,
        load_stage_config,
        read_sft_quality_gate,
        resolve_pair_split_dir,
        resolve_sft_checkpoint,
    )
    from align.train.ref_cache import (
        build_ref_cache_key,
        path_bias_report_cache_forward_impl,
        precompute_ref_logprobs,
        write_path_bias_report,
        write_ref_cache,
    )
    from align.utils.guard import assert_paths_off_root

    cfg = load_stage_config(args.config)
    out_dir = args.out_dir or os.path.dirname(cfg.ref_cache_path)
    assert_paths_off_root([out_dir])

    with open(cfg.manifest_path) as f:
        manifest = json.load(f)

    # Resolve the frozen SFT checkpoint the same way scripts/dpo.py does.
    # The trainer writes <run_dir>/checkpoints/micro_NNNNNNNN with no stable
    # `final` alias; the quality-gate receipt records which one passed. A
    # configured path that IS a checkpoint is an explicit pin; otherwise the
    # receipt is read, and a missing/failed receipt or a missing checkpoint
    # refuses LOUDLY here -- a cache silently computed from the wrong (or no)
    # reference would poison both arms' margins.
    if is_hf_checkpoint(cfg.sft_checkpoint):
        sft_checkpoint, ckpt_source = cfg.sft_checkpoint, "config"
    else:
        gate = read_sft_quality_gate(cfg.sft_quality_gate_path)
        sft_checkpoint, ckpt_source = resolve_sft_checkpoint(
            cfg.sft_checkpoint, gate.get("checkpoint")
        )
    logger.info("frozen SFT checkpoint: %s (source=%s)", sft_checkpoint, ckpt_source)

    # Shard split dirs by DISCOVERY (the receipts the store wrote), never by
    # hand-composed `<root>-<split>` paths -- the attempt-1 defect class.
    train_dir = resolve_pair_split_dir(cfg.shard_dir, f"{cfg.arm}-train")
    held_dir = resolve_pair_split_dir(cfg.heldout_shard_dir, f"{cfg.arm}-held")

    # The key, assembled from facts rather than from intent, by the SAME
    # function scripts/dpo.py's reader uses (build_ref_cache_key). `checkpoint_hash`
    # is the WEIGHT-BYTES hash -- one streamed read of the checkpoint -- not the
    # tensor-inventory hash this script used to write, which could not tell two
    # SFT checkpoints of one architecture apart (2026-10-12 review). Caches
    # written under the old key are refused by dpo.py: re-run this script.
    # `manifest_hash` covers the shard rebuild that nothing else would notice.
    key = build_ref_cache_key(
        sft_checkpoint,
        manifest,
        forward_impl="fsdp2_sharded" if (args.path_bias or args.sharded_cache) else "single_gpu",
    )
    logger.info("cache key: %s", key.to_json())

    # Rank-pinned device: a bare "cuda" is device 0 for EVERY torchrun rank,
    # and the --path-bias 4-rank pass put four ~5.5GB reference shards on
    # GPU 0 (the duplicate-GPU class scripts/sft.py already fixed). This
    # mode had never been run before 2026-08-18.
    if torch.cuda.is_available():
        from align.utils.dist import init_distributed, local_rank

        init_distributed()
        torch.cuda.set_device(local_rank())
        device = f"cuda:{local_rank()}"  # string: downstream does device.startswith(...)
    else:
        device = "cpu"
    model = _load_reference(NmoMistralForCausalLM, ckpt, sft_checkpoint, device)

    if args.sharded_cache:
        import torch.distributed as dist

        from align.train.fsdp import apply_fsdp

        if not (dist.is_available() and dist.is_initialized()):
            raise RuntimeError(
                "--sharded-cache computes the reference under the SHARDED "
                "forward and must run under torchrun; single-process would "
                "silently produce the single_gpu numbers under a sharded key."
            )
        model = apply_fsdp(model)
        model.eval()

    written = []
    # In --path-bias mode the caches are NOT recomputed: the dedicated
    # single-GPU run owns them. Before this guard, all four torchrun ranks
    # each recomputed and rewrote the same cache files concurrently -- and
    # stamped forward_impl="fsdp2_sharded" into caches computed single-GPU.
    cache_splits = (
        ()
        if args.path_bias
        else (
            ("train", train_dir, cfg.ref_cache_path),
            ("held", held_dir, cfg.heldout_ref_cache_path),
        )
    )
    for split, shard_dir, cache_path in cache_splits:
        rows = preference_dataset(shard_dir)
        cache = precompute_ref_logprobs(
            model,
            rows,
            key=key,
            pad_token_id=_pad_id(manifest),
            batch_size=args.batch_size,
            device=torch.device(device),
        )
        from align.utils.dist import is_rank0 as _is_rank0

        if _is_rank0():
            # Sharded forwards are collectives (every rank computes); the
            # file is written once.
            write_ref_cache(cache_path, cache)
        written.append({"split": split, "path": cache_path, "pairs": len(cache)})
        logger.info("%s/%s: %d reference pairs", cfg.arm, split, len(cache))

    if args.batch_invariance:
        from align.train.ref_cache import batch_invariance_report

        rows = preference_dataset(train_dir)[:PATH_BIAS_EXAMPLES]
        report = batch_invariance_report(
            model,
            rows,
            key=key,
            pad_token_id=_pad_id(manifest),
            beta=cfg.beta,
            early_margin_scale=cfg.path_bias_early_margin_scale,
            threshold=cfg.path_bias_threshold,
        )
        out_path = os.path.join(out_dir, f"{cfg.arm}-batch_invariance.json")
        with open(out_path, "w") as f:
            json.dump(report, f, indent=2)
        logger.info(
            "batch invariance -> %s (belongs_in_cache_key=%s)",
            out_path,
            report["belongs_in_cache_key"],
        )

    if args.path_bias:
        # OWNERSHIP TRANSFER: pass the model in a box and drop main's name.
        # A plain argument keeps main's reference alive, so the between-pass
        # free inside _path_bias freed nothing and the reload OOM'd with the
        # old 14.5GB replica still resident (measured 21.66GB, batch-size
        # independent).
        model_box = [model]
        del model
        report = _path_bias(
            model_box,
            cfg,
            manifest,
            key,
            args.batch_size,
            device,
            train_dir,
            reload_fn=lambda: _load_reference(NmoMistralForCausalLM, ckpt, sft_checkpoint, device),
        )
        from align.utils.dist import is_rank0

        if is_rank0():
            # Rank-0-only: four ranks renaming one report is the ledger race.
            # `cache_forward_impl` must describe the CACHE, not this run. The
            # old code wrote key.forward_impl, which is "fsdp2_sharded" for
            # every --path-bias run -- but this mode recomputes no cache, so the
            # field named the measuring run and contradicted the single-GPU
            # cache sitting on disk (2026-10-12 review). Read it off the cache.
            write_path_bias_report(
                cfg.path_bias_report_path,
                report,
                n_examples=PATH_BIAS_EXAMPLES,
                forward_impl=path_bias_report_cache_forward_impl(cfg.ref_cache_path),
            )
        logger.info("path bias: %s (ratio %s)", report["action"], report["ratio"])

    logger.info("precompute complete: %s", written)
    return 0


def _manifest_hash(manifest: dict) -> str:
    """Hash of the build's own content hashes. Moves on any shard rebuild --
    which is the point: the section-03 marker fix retokenizes every pair while
    leaving the checkpoint, tokenizer and max_length untouched."""
    from align.data.receipts import content_hash

    return content_hash([manifest["hashes"][k] for k in sorted(manifest.get("hashes", {}))])


def _max_length(manifest: dict) -> int:
    return int(manifest.get("pinned", {}).get("max_len", 2048))


def _pad_id(manifest: dict) -> int:
    """Pad id for collation. Fully masked out of every logprob (prompt and pad
    positions carry IGNORE_INDEX), so its value cannot reach a number -- but
    it is read from the build rather than assumed, because 'cannot reach a
    number' is the kind of claim that stops being true after a refactor."""
    return int(manifest.get("pinned", {}).get("pad_token_id", 0))


def _load_reference(model_cls, ckpt, path: str, device: str):
    """The FROZEN SFT checkpoint, bf16 on GPU, eval mode.

    eval() is asserted downstream rather than merely called here: Mistral-v0.3
    has attention_dropout 0.0, so a model left in train mode would produce the
    same numbers today and silently different ones after any config change.
    """
    import torch
    from align.eval.golden import resolve_model_config

    model_cfg = resolve_model_config(path)
    model = model_cls.from_state_dict(
        model_cfg,
        ckpt.hf_to_nmo(ckpt.load_hf_safetensors(path)),
        attn_implementation="sdpa",
        device=device,
    )
    model = model.to(
        device, dtype=torch.bfloat16 if str(device).startswith("cuda") else torch.float32
    )
    model.eval()
    for module in model.modules():
        if isinstance(module, torch.nn.Dropout) and module.p > 0:
            raise RuntimeError(
                f"an active Dropout(p={module.p}) is present in the reference "
                "model. Cached reference logprobs must be exactly "
                "reproducible; dropout makes them a sample."
            )
    return model


def _path_bias(
    model_box, cfg, manifest, key, batch_size: int, device: str, train_dir: str, reload_fn=None
) -> dict:
    """The mandatory dual pass: single-GPU forward vs FSDP2-sharded forward.

    Requires torchrun. Without a process group there IS no sharded forward to
    compare against, and the honest report is then no report -- `path_bias_
    decision` treats an empty delta set as firing, because not measuring is
    not passing and a stage that skipped this must not silently proceed on a
    single-GPU cache.
    """
    import torch
    import torch.distributed as dist

    from align.data.loader import preference_dataset
    from align.train.fsdp import apply_fsdp
    from align.train.ref_cache import path_bias_decision, precompute_ref_logprobs

    if not (dist.is_available() and dist.is_initialized()):
        raise RuntimeError(
            "--path-bias needs the SHARDED forward, so it must run under "
            "torchrun (--nproc_per_node=4). Running it single-process would "
            "compare the single-GPU forward with itself and record a "
            "reassuring zero, which is worse than not measuring: it would "
            "look like evidence."
        )

    rows = preference_dataset(train_dir)[:PATH_BIAS_EXAMPLES]
    model = model_box.pop()  # sole reference from here on
    single = precompute_ref_logprobs(
        model,
        rows,
        key=key,
        pad_token_id=_pad_id(manifest),
        batch_size=batch_size,
        device=torch.device(device),
    )
    # FREE the single-pass replica before sharding: wrapping a resident
    # 14.5GB model in-place stacks the wrap transients ON TOP of it and
    # OOMs every rank at 21.8GB (measured, 2026-08-18). Rebuilding alone
    # and then sharding is the stage launcher's proven peak.
    if reload_fn is None:
        raise RuntimeError("--path-bias needs reload_fn to rebuild the model between passes")
    del model
    import gc

    gc.collect()
    torch.cuda.empty_cache()
    model = reload_fn()
    sharded_model = apply_fsdp(model)
    sharded_model.eval()
    sharded = precompute_ref_logprobs(
        sharded_model,
        rows,
        key=key,
        pad_token_id=_pad_id(manifest),
        batch_size=batch_size,
        device=torch.device(device),
    )

    # The delta of the reference MARGIN, not of either logprob: the margin is
    # what enters the loss, and a bias common to both sides cancels there.
    deltas = [
        (sharded.entries[p][0] - sharded.entries[p][1])
        - (single.entries[p][0] - single.entries[p][1])
        for p in single.entries
    ]
    return path_bias_decision(
        deltas,
        beta=cfg.beta,
        early_margin_scale=cfg.path_bias_early_margin_scale,
        threshold=cfg.path_bias_threshold,
    )


if __name__ == "__main__":
    raise SystemExit(main())