"""
GRPO training for E0 (quality-only baseline), E1 (div-grpo-individual),
E2 (div-grpo-group). The three arms differ ONLY by YAML config -- same code
path, same data, same seed -- so any difference between them is attributable
to the reward configuration and nothing else.

Aggregation is GDPO (arXiv 2601.05242, Liu et al., NVIDIA): group-wise
normalization per reward channel, then batch-wise advantage normalization.
TRL 1.10 implements this as multi_objective_aggregation="normalize_then_sum".
"""
from __future__ import annotations

import argparse
import json
import os
import sys
from dataclasses import asdict
from pathlib import Path

import yaml

ROOT = Path(__file__).resolve().parent.parent


def build_dataset(prompts, tokenizer):
    from datasets import Dataset
    from data import chat_messages
    return Dataset.from_list([
        {"prompt": chat_messages(p["prompt"]), "prompt_id": p["id"]}
        for p in prompts
    ])


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", required=True)
    ap.add_argument("--max-steps", type=int, default=None, help="override (smoke tests)")
    ap.add_argument("--smoke", action="store_true")
    args = ap.parse_args()

    cfg = yaml.safe_load(open(args.config))
    name = cfg["name"]
    if args.smoke:
        name = f"{name}-smoke"

    import torch
    import wandb
    from peft import LoraConfig
    from transformers import AutoTokenizer, TrainerCallback
    from trl import GRPOConfig, GRPOTrainer

    import logbook
    from data import load_prompts
    from judge import build_judge
    from rewards import RewardConfig, RewardEngine

    out_dir = ROOT / "outputs" / name
    out_dir.mkdir(parents=True, exist_ok=True)

    steps = args.max_steps or cfg["train"]["max_steps"]
    model_id = cfg["model"]
    G = cfg["train"]["num_generations"]

    # ---- wandb -----------------------------------------------------------
    run = None
    if cfg.get("wandb", True) and os.environ.get("WANDB_API_KEY"):
        run = wandb.init(
            project=os.environ.get("WANDB_PROJECT", "div-grpo"),
            name=name, config=cfg, reinit=True,
            mode=os.environ.get("WANDB_MODE", "online"),
        )

    logbook.note(f"START {name}",
                 f"```yaml\n{yaml.safe_dump(cfg, sort_keys=False)}```\n"
                 f"steps={steps} G={G} model={model_id}")

    # ---- reward engine ---------------------------------------------------
    rcfg = RewardConfig(
        arm=cfg["reward"]["arm"],
        alpha=cfg["reward"].get("alpha", 0.5),
        gamma=cfg["reward"].get("gamma", 0.5),
        tau=cfg["reward"].get("tau", 5.0),
    )
    judge = build_judge(
        model=cfg["judge"]["model"],
        cache_path=str(ROOT / "cache" / "judge.sqlite"),
        concurrency=cfg["judge"].get("concurrency", 12),
    )
    engine = RewardEngine(rcfg, judge, wandb_run=run, log_prefix="train")
    reward_funcs = engine.make_reward_funcs()
    weights = rcfg.weights()
    print(f"[arm {rcfg.arm}] channels={rcfg.channels()} weights={weights} tau={rcfg.tau}")

    # ---- data ------------------------------------------------------------
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    train_prompts = load_prompts("train", ROOT / "data")
    if args.smoke:
        train_prompts = train_prompts[:64]
    train_ds = build_dataset(train_prompts, tokenizer)

    # ---- LoRA ------------------------------------------------------------
    lora = LoraConfig(
        r=cfg["lora"]["r"],
        lora_alpha=cfg["lora"]["alpha"],
        lora_dropout=cfg["lora"].get("dropout", 0.0),
        target_modules=cfg["lora"]["target_modules"],
        task_type="CAUSAL_LM",
        bias="none",
    )

    gcfg = GRPOConfig(
        output_dir=str(out_dir),
        max_steps=steps,
        per_device_train_batch_size=cfg["train"]["per_device_train_batch_size"],
        gradient_accumulation_steps=cfg["train"]["gradient_accumulation_steps"],
        num_generations=G,
        max_completion_length=cfg["train"]["max_completion_length"],
        # TRL 1.10 dropped max_prompt_length; vLLM's window is the control now.
        vllm_max_model_length=cfg["train"].get("vllm_max_model_length", 2048),
        # NOT masking truncated completions: a truncated story is gated to the
        # bottom of every reward channel, and we want that negative gradient to
        # reach the policy. Masking would make truncation free.
        mask_truncated_completions=False,
        learning_rate=cfg["train"]["learning_rate"],
        lr_scheduler_type=cfg["train"].get("lr_scheduler_type", "constant_with_warmup"),
        warmup_steps=cfg["train"].get("warmup_steps", 10),
        beta=cfg["train"]["beta"],
        temperature=cfg["train"].get("temperature", 1.0),
        top_p=cfg["train"].get("top_p", 1.0),
        # GDPO: per-reward group normalization, then batch-level advantage norm
        multi_objective_aggregation="normalize_then_sum",
        reward_weights=weights,
        scale_rewards=cfg["train"].get("scale_rewards", "group"),
        bf16=True,
        gradient_checkpointing=True,
        # Liger fuses RMSNorm/SwiGLU/RoPE and the LM-head cross-entropy, which
        # is where the peak lives: the logits tensor is
        # micro_batch x seq x 151936 vocab, and it was the allocation that OOMed.
        use_liger_kernel=cfg["train"].get("use_liger_kernel", True),
        torch_empty_cache_steps=cfg["train"].get("torch_empty_cache_steps", 8),
        use_vllm=True,
        vllm_mode="colocate",
        vllm_gpu_memory_utilization=cfg["train"]["vllm_gpu_memory_utilization"],
        logging_steps=1,
        save_steps=cfg["train"].get("save_steps", 50),
        save_total_limit=cfg["train"].get("save_total_limit", 7),
        # Checkpoints exist only to EVALUATE intermediate policies (ckpt_study),
        # never to resume training. Without this, HF writes a 505MB optimizer.pt
        # beside a 253MB adapter -- 3x the disk for state we never read. This was
        # set in the YAML from E1 onward but not passed through until now.
        save_only_model=cfg["train"].get("save_only_model", True),
        seed=cfg.get("seed", 42),
        report_to=["wandb"] if run else [],
        run_name=name,
    )

    trainer = GRPOTrainer(
        model=model_id,
        reward_funcs=reward_funcs,
        args=gcfg,
        train_dataset=train_ds,
        peft_config=lora,
    )

    # ---- periodic reward-hacking guardrail -------------------------------
    class Guardrail(TrainerCallback):
        """Stop the run if diversity climbs while quality/validity collapses.

        The brief's guardrail: 'if reward hacking appears (deviation up, quality
        flat/down, or degenerate text passing gates), stop the run'. We compare a
        trailing window against the opening baseline rather than step-to-step,
        because GRPO reward traces are far too noisy for a point comparison.
        """
        WINDOW = 25

        def on_step_end(self, a, state, control, **kw):
            h = engine.history
            if len(h) < self.WINDOW * 2:
                return
            base = h[:self.WINDOW]
            recent = h[-self.WINDOW:]

            def mean(rows, f):
                return sum(f(r) for r in rows) / len(rows)

            gate0, gate1 = mean(base, lambda r: r.gate_pass), mean(recent, lambda r: r.gate_pass)
            q0, q1 = mean(base, lambda r: r.mean_quality_passing), mean(recent, lambda r: r.mean_quality_passing)
            d0, d1 = mean(base, lambda r: r.mean_deviation), mean(recent, lambda r: r.mean_deviation)

            msg = None
            if gate1 < 0.55 and gate1 < gate0 - 0.25:
                msg = f"gate pass collapsed {gate0:.2f}->{gate1:.2f}"
            elif d1 > d0 + 0.05 and q1 < q0 - 1.0:
                msg = f"reward hacking: deviation {d0:.3f}->{d1:.3f} while quality {q0:.2f}->{q1:.2f}"
            if msg:
                logbook.note(f"GUARDRAIL TRIP {name}", msg, level="ALERT")
                print(f"\n!!! GUARDRAIL: {msg} -- stopping at step {state.global_step}\n", flush=True)
                control.should_training_stop = True

    trainer.add_callback(Guardrail())

    print(f"\n── training {name}: {steps} steps ──", flush=True)
    trainer.train()

    final = out_dir / "final"
    trainer.save_model(str(final))
    tokenizer.save_pretrained(str(final))

    hist = [asdict(s) for s in engine.history]
    json.dump(hist, open(out_dir / "reward_history.json", "w"), indent=1)
    # TRL's own log history carries per-token policy entropy, KL and clip ratio.
    # Entropy is only present on the non-liger loss path (compute_liger_loss logs
    # just clip_ratio and kl), which is why use_liger_kernel is disabled.
    json.dump(trainer.state.log_history,
              open(out_dir / "trl_log_history.json", "w"), indent=1)
    ent = [h["entropy"] for h in trainer.state.log_history if "entropy" in h]
    print(f"entropy logged for {len(ent)} steps"
          + (f" | first={ent[0]:.4f} last={ent[-1]:.4f}" if ent else " -- MISSING!"))
    cost = judge.cost_estimate(cfg["judge"]["price_in"], cfg["judge"]["price_out"])
    json.dump(cost, open(out_dir / "judge_cost.json", "w"), indent=1)
    print("judge cost:", cost)

    logbook.note(f"DONE {name}",
                 f"adapter: `{final}`\n\njudge cost: `{json.dumps(cost)}`")
    logbook.checkpoint(f"after {name}")
    if run:
        run.finish()
    return 0


if __name__ == "__main__":
    sys.exit(main())