"""Reinforcement learning with a verifiable reward, using TRL's GRPOTrainer.

Purpose: the course's reference GRPO run. Loads a prompt-only task set whose answers
    are correct by construction, samples a group of completions per prompt, scores
    them with the reward functions from rewards.py, and updates a LoRA adapter on the
    group-relative advantages. Writes the reward, KL and length curves to a CSV so
    they can be read rather than guessed at, and appends one run record to the lab
    notebook.
Platform: spark, nvidia (CUDA) and strix (ROCm) with in-process or vLLM rollouts;
    mac runs the same script on PyTorch's MPS backend in float32 with in-process
    rollouts and a small model. vLLM does not run on macOS, so --rollouts vllm-* is
    refused there; --rollouts llama-server is the served-model path for tracks
    without vLLM and carries the stale-weights caveat printed at start-up.
Minimum memory: 16 GB
Assumes: torch, transformers, trl, peft and datasets installed in the active
    environment; make-tasks.py has been run so that <tasks-dir>/train.jsonl exists;
    rewards.py and runlog.py sit next to this file.

Usage: python3 train-grpo.py --tasks-dir tasks --output-dir runs/grpo-qwen3-1.7b \
           --model Qwen/Qwen3-1.7B --labbook labbook.md
       python3 train-grpo.py --tasks-dir tasks --model Qwen/Qwen3-4B \
           --rollouts vllm-colocate --num-generations 8 --labbook labbook.md
       python3 train-grpo.py --tasks-dir tasks --model Qwen/Qwen3-1.7B \
           --rollouts llama-server --rollout-url http://127.0.0.1:8080 --labbook labbook.md
"""

from __future__ import annotations

import argparse
import csv
import json
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path
from typing import Any, Optional

import torch
from datasets import load_dataset
from peft import LoraConfig
from transformers import AutoTokenizer
from trl import GRPOConfig, GRPOTrainer

import rewards as reward_lib
import runlog

DEFAULT_TARGETS = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]

# The metrics worth plotting, in the order the lab reads them.
CURVE_KEYS = ["reward", "reward_std", "kl", "completions/mean_length", "loss",
              "clip_ratio/region_mean", "learning_rate", "epoch"]


def pick_device() -> str:
    """CUDA (or ROCm, which reports as cuda), then MPS, then the CPU."""
    if torch.cuda.is_available():
        return "cuda"
    mps = getattr(torch.backends, "mps", None)
    if mps is not None and mps.is_available():
        return "mps"
    return "cpu"


def use_bf16(device: str, requested: str) -> bool:
    if requested == "fp32":
        return False
    if requested == "bf16":
        return True
    return device == "cuda" and torch.cuda.is_bf16_supported()


# --------------------------------------------------------------------------------------
# Rollouts through an OpenAI-compatible llama-server
# --------------------------------------------------------------------------------------

STALE_WEIGHTS_WARNING = """\
--rollouts llama-server samples from a server holding a fixed GGUF file. TRL's vLLM
integration streams the updated weights into the rollout engine after every optimiser
step; llama.cpp's server README documents no equivalent endpoint, so after the first
update the completions come from an older policy than the one being trained. GRPO's
importance ratio corrects for a *small* mismatch, not an unbounded one. Use this path
for a short run whose purpose is to see the machinery work on a track without vLLM,
keep --num-iterations at 1, and record in the run log that the rollouts were stale.
The in-process path (--rollouts inproc) is always in sync and is the correct choice
whenever you can afford its speed.
"""


def post_json(url: str, payload: dict, api_key: Optional[str], timeout: int) -> dict:
    body = json.dumps(payload).encode("utf-8")
    headers = {"Content-Type": "application/json"}
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"
    request = urllib.request.Request(url, data=body, headers=headers, method="POST")
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            return json.loads(response.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", "replace")[:400]
        raise RuntimeError(f"{url} returned HTTP {exc.code}: {detail}") from exc
    except urllib.error.URLError as exc:
        raise RuntimeError(f"cannot reach {url}: {exc.reason}") from exc


def token_ids_and_logprobs(body: dict) -> tuple[list[int], list[float]]:
    """Pull the sampled tokens and their log-probabilities out of a llama-server reply.

    The server README documents `n_probs`: "If greater than 0, the response also
    contains the probabilities of top N tokens for each generated token", returned in
    `completion_probabilities`, where "Each item in the array has a nested array
    `top_logprobs`". The exact key names inside those items have changed between
    builds, so this reads defensively and fails with a message naming the field rather
    than training on log-probabilities it guessed at.
    """
    items = body.get("completion_probabilities")
    if not items:
        raise SystemExit(
            "the server returned no completion_probabilities. Send n_probs greater than 0 "
            "on the /completion endpoint, and check that your llama.cpp build returns it."
        )
    ids: list[int] = []
    logprobs: list[float] = []
    for item in items:
        chosen_id = item.get("id", item.get("tok"))
        candidates = item.get("top_logprobs") or []
        picked = None
        for candidate in candidates:
            if chosen_id is not None and candidate.get("id") == chosen_id:
                picked = candidate
                break
        if picked is None and candidates:
            picked = candidates[0]
        if picked is None:
            raise SystemExit("a completion_probabilities item has no top_logprobs entry")
        token_id = picked.get("id", chosen_id)
        logprob = picked.get("logprob")
        if token_id is None or logprob is None:
            raise SystemExit(
                "a top_logprobs entry is missing 'id' or 'logprob'; this build's response "
                "shape differs from the server README read for this lab"
            )
        ids.append(int(token_id))
        logprobs.append(float(logprob))
    return ids, logprobs


def make_llama_server_rollout(base_url: str, api_key: Optional[str], tokenizer,
                              num_generations: int, max_new_tokens: int,
                              temperature: float, top_p: float, seed: int, timeout: int):
    """A rollout function that samples from llama-server's native /completion endpoint.

    GRPOTrainer's `rollout_func` is documented as experimental: "It receives the list
    of prompts allocated to the current process and the trainer instance. It must
    return a dict with `prompt_ids`, `completion_ids`, and `logprobs` fields... This
    feature is experimental and may change or be removed at any time without prior
    notice." Because the calling convention may change, the wrapper below accepts the
    prompts positionally or by keyword and says so if it cannot find them.
    """
    endpoint = base_url.rstrip("/") + "/completion"
    state = {"warned_decode": False, "calls": 0}

    def one_completion(prompt_text: str, sample_index: int) -> tuple[list[int], list[float]]:
        payload = {
            "prompt": prompt_text,
            "n_predict": max_new_tokens,
            "temperature": temperature,
            "top_p": top_p,
            "n_probs": 1,
            "cache_prompt": True,
            "seed": seed + sample_index + state["calls"] * num_generations,
        }
        body = post_json(endpoint, payload, api_key, timeout)
        ids, logprobs = token_ids_and_logprobs(body)
        text = body.get("content", "")
        if not state["warned_decode"]:
            decoded = tokenizer.decode(ids, skip_special_tokens=True)
            if decoded.strip() != text.strip():
                print("WARNING: the server's token ids do not decode to the text it returned "
                      "under the trainer's tokeniser. The GGUF file and the Hugging Face model "
                      "must be the same model with the same vocabulary.", file=sys.stderr)
            state["warned_decode"] = True
        return ids, logprobs

    def rollout(prompts: Any = None, trainer: Any = None, **call_kwargs: Any) -> dict:
        # TRL types rollout_func as Callable[[list[str], "GRPOTrainer"], dict], so the
        # prompts arrive first and the trainer second. The keyword fallback below exists
        # because the documentation marks the feature experimental.
        if prompts is None:
            prompts = call_kwargs.get("prompts")
        if prompts is None:
            raise SystemExit("rollout_func was called without a list of prompts; TRL's "
                             "experimental calling convention has changed")
        del trainer  # this path needs no trainer state; the server holds the weights

        prompt_ids: list[list[int]] = []
        completion_ids: list[list[int]] = []
        all_logprobs: list[list[float]] = []
        for prompt in prompts:
            text = prompt if isinstance(prompt, str) else tokenizer.apply_chat_template(
                prompt, tokenize=False, add_generation_prompt=True
            )
            ids = tokenizer(text, add_special_tokens=False)["input_ids"]
            for k in range(num_generations):
                completion, logprobs = one_completion(text, k)
                prompt_ids.append(list(ids))
                completion_ids.append(completion)
                all_logprobs.append(logprobs)
        state["calls"] += 1
        return {"prompt_ids": prompt_ids, "completion_ids": completion_ids, "logprobs": all_logprobs}

    return rollout


# --------------------------------------------------------------------------------------
# Curves
# --------------------------------------------------------------------------------------

def write_curves(history: list[dict], path: Path) -> dict[str, Any]:
    """One row per logged step, and the three numbers the lab asks you to compare."""
    rows = [row for row in history if "reward" in row or "loss" in row]
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.writer(handle)
        writer.writerow(["step"] + CURVE_KEYS)
        for row in rows:
            writer.writerow([row.get("step", "")] + [row.get(k, "") for k in CURVE_KEYS])

    rewarded = [row for row in rows if isinstance(row.get("reward"), (int, float))]
    lengths = [row["completions/mean_length"] for row in rows
               if isinstance(row.get("completions/mean_length"), (int, float))]
    kls = [row["kl"] for row in rows if isinstance(row.get("kl"), (int, float))]
    summary = {
        "logged_steps": len(rows),
        "first_reward": round(rewarded[0]["reward"], 4) if rewarded else None,
        "last_reward": round(rewarded[-1]["reward"], 4) if rewarded else None,
        "best_reward": round(max(r["reward"] for r in rewarded), 4) if rewarded else None,
        "first_mean_length": round(lengths[0], 1) if lengths else None,
        "last_mean_length": round(lengths[-1], 1) if lengths else None,
        "last_kl": round(kls[-1], 5) if kls else None,
        "max_kl": round(max(kls), 5) if kls else None,
    }
    return summary


# --------------------------------------------------------------------------------------

def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--model", default="Qwen/Qwen3-1.7B",
                        help="base model repository id or local path")
    parser.add_argument("--adapter", default=None,
                        help="an existing LoRA adapter to continue training, e.g. a Part 13 fine-tune")
    parser.add_argument("--tasks-dir", default="tasks", help="directory holding train.jsonl from make-tasks.py")
    parser.add_argument("--output-dir", default="runs/grpo")
    parser.add_argument("--reward-kind", default="maths", choices=["maths", "exact", "code"])
    parser.add_argument("--no-format-reward", action="store_true")
    parser.add_argument("--no-length-reward", action="store_true")
    parser.add_argument("--fallback-to-last-number", action="store_true",
                        help="grade the last number in the completion when no answer marker is found")
    parser.add_argument("--num-generations", type=int, default=8,
                        help="rollouts per prompt, G in the lesson; the effective batch must divide by it")
    parser.add_argument("--max-completion-length", type=int, default=512)
    parser.add_argument("--target-tokens", type=int, default=256, help="soft budget for the length reward")
    parser.add_argument("--beta", type=float, default=0.04,
                        help="KL coefficient; TRL's default is 0.0, which skips the reference model entirely")
    parser.add_argument("--epsilon", type=float, default=0.2)
    parser.add_argument("--epsilon-high", type=float, default=None,
                        help="upper clipping bound; the DAPO paper recommends 0.28")
    parser.add_argument("--loss-type", default="dapo", choices=["grpo", "dapo", "dr_grpo"])
    parser.add_argument("--scale-rewards", default="group", choices=["group", "batch"],
                        help="divide the advantage by the group or the batch standard deviation; "
                             "to drop the term entirely, use --loss-type dr_grpo")
    parser.add_argument("--temperature", type=float, default=1.0)
    parser.add_argument("--top-p", type=float, default=1.0)
    parser.add_argument("--lr", type=float, default=1e-5)
    parser.add_argument("--batch-size", type=int, default=8, help="prompts per device per step")
    parser.add_argument("--grad-accum", type=int, default=1)
    parser.add_argument("--num-iterations", type=int, default=1, help="policy updates per batch of rollouts")
    parser.add_argument("--max-steps", type=int, default=200)
    parser.add_argument("--rank", type=int, default=16)
    parser.add_argument("--alpha", type=int, default=32)
    parser.add_argument("--target-modules", nargs="+", default=DEFAULT_TARGETS)
    parser.add_argument("--precision", choices=["auto", "bf16", "fp32"], default="auto")
    parser.add_argument("--gradient-checkpointing", action="store_true")
    parser.add_argument("--rollouts", default="inproc",
                        choices=["inproc", "vllm-colocate", "vllm-server", "llama-server"])
    parser.add_argument("--rollout-url", default="http://127.0.0.1:8080",
                        help="llama-server base URL for --rollouts llama-server")
    parser.add_argument("--rollout-api-key", default=None)
    parser.add_argument("--rollout-timeout", type=int, default=600)
    parser.add_argument("--report-to", default="none", choices=["none", "tensorboard"])
    parser.add_argument("--curves", default=None, help="write the reward, KL and length curves here as CSV")
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--labbook", default=None)
    parser.add_argument("--notes", default=None)
    args = parser.parse_args()

    device = pick_device()
    bf16 = use_bf16(device, args.precision)
    print(f"device: {device}   precision: {'bfloat16' if bf16 else 'float32'}   rollouts: {args.rollouts}")

    if args.rollouts.startswith("vllm") and device == "mps":
        raise SystemExit("vLLM's documented GPU path does not cover macOS; use --rollouts inproc "
                         "or --rollouts llama-server on Track M")

    effective_batch = args.batch_size * args.grad_accum
    if effective_batch % args.num_generations:
        raise SystemExit(
            f"the effective batch ({args.batch_size} x {args.grad_accum} = {effective_batch}) must be "
            f"divisible by --num-generations ({args.num_generations}); TRL requires it"
        )

    train_path = Path(args.tasks_dir) / "train.jsonl"
    if not train_path.is_file():
        raise SystemExit(f"{train_path} is missing; run make-tasks.py --out-dir {args.tasks_dir} first")
    dataset = load_dataset("json", data_files={"train": str(train_path)})["train"]
    print(f"train prompts: {len(dataset)}   families: {sorted(set(dataset['family']))}")

    tokenizer = AutoTokenizer.from_pretrained(args.adapter or args.model)
    if tokenizer.chat_template is None:
        raise SystemExit(f"{args.model} has no chat template; pick an instruction-tuned model")

    reward_funcs, reward_weights = reward_lib.build_reward_functions(
        kind=args.reward_kind,
        use_format=not args.no_format_reward,
        use_length=not args.no_length_reward,
        target_tokens=args.target_tokens,
        fallback_to_last_number=args.fallback_to_last_number,
    )
    print("rewards: " + ", ".join(f"{f.__name__} x{w}" for f, w in zip(reward_funcs, reward_weights)))

    config_kwargs: dict[str, Any] = dict(
        output_dir=args.output_dir,
        max_steps=args.max_steps,
        per_device_train_batch_size=args.batch_size,
        gradient_accumulation_steps=args.grad_accum,
        num_generations=args.num_generations,
        max_completion_length=args.max_completion_length,
        num_iterations=args.num_iterations,
        beta=args.beta,
        epsilon=args.epsilon,
        loss_type=args.loss_type,
        scale_rewards=args.scale_rewards,
        temperature=args.temperature,
        top_p=args.top_p,
        reward_weights=reward_weights,
        learning_rate=args.lr,
        lr_scheduler_type="constant_with_warmup",
        warmup_steps=5,
        bf16=bf16,
        gradient_checkpointing=args.gradient_checkpointing,
        logging_steps=1,
        save_strategy="steps",
        save_steps=max(25, args.max_steps // 4),
        save_total_limit=2,
        report_to=args.report_to,
        seed=args.seed,
        data_seed=args.seed,
        model_init_kwargs={"dtype": torch.bfloat16 if bf16 else torch.float32},
    )
    if args.epsilon_high is not None:
        config_kwargs["epsilon_high"] = args.epsilon_high
    if args.rollouts == "vllm-colocate":
        config_kwargs["use_vllm"] = True
        config_kwargs["vllm_mode"] = "colocate"
    elif args.rollouts == "vllm-server":
        config_kwargs["use_vllm"] = True
        config_kwargs["vllm_mode"] = "server"

    trainer_kwargs: dict[str, Any] = {}
    if args.rollouts == "llama-server":
        print(STALE_WEIGHTS_WARNING, file=sys.stderr)
        trainer_kwargs["rollout_func"] = make_llama_server_rollout(
            base_url=args.rollout_url,
            api_key=args.rollout_api_key,
            tokenizer=tokenizer,
            num_generations=args.num_generations,
            max_new_tokens=args.max_completion_length,
            temperature=args.temperature,
            top_p=args.top_p,
            seed=args.seed,
            timeout=args.rollout_timeout,
        )

    if args.adapter:
        from peft import AutoPeftModelForCausalLM  # noqa: PLC0415 - only needed on this branch
        model = AutoPeftModelForCausalLM.from_pretrained(args.adapter, is_trainable=True)
        config_kwargs.pop("model_init_kwargs", None)
        print(f"continuing the adapter in {args.adapter}")
    else:
        model = args.model
        trainer_kwargs["peft_config"] = LoraConfig(
            r=args.rank, lora_alpha=args.alpha, lora_dropout=0.0,
            target_modules=args.target_modules, bias="none", task_type="CAUSAL_LM",
        )

    trainer = GRPOTrainer(
        model=model,
        args=GRPOConfig(**config_kwargs),
        reward_funcs=reward_funcs,
        train_dataset=dataset,
        processing_class=tokenizer,
        **trainer_kwargs,
    )

    started = time.time()
    trainer.train()
    elapsed = time.time() - started

    trainer.save_model(args.output_dir)
    tokenizer.save_pretrained(args.output_dir)

    curves_path = Path(args.curves or Path(args.output_dir) / "curves.csv")
    summary = write_curves(trainer.state.log_history, curves_path)
    summary["seconds"] = round(elapsed, 1)
    print(json.dumps(summary, indent=2))
    print(f"curves written to {curves_path}")
    print(f"adapter saved to {args.output_dir}")

    if args.labbook:
        record = runlog.record(
            labbook=args.labbook,
            lab="part-14/train-grpo",
            model=args.adapter or args.model,
            dataset={
                "path": str(train_path),
                "sha256": runlog.file_sha256(train_path),
                "train_prompts": len(dataset),
            },
            hyperparameters={
                "method": "grpo-lora",
                "rollouts": args.rollouts,
                "rollouts_in_sync": args.rollouts != "llama-server",
                "num_generations": args.num_generations,
                "max_completion_length": args.max_completion_length,
                "beta": args.beta,
                "epsilon": args.epsilon,
                "epsilon_high": args.epsilon_high,
                "loss_type": args.loss_type,
                "scale_rewards": args.scale_rewards,
                "temperature": args.temperature,
                "top_p": args.top_p,
                "learning_rate": args.lr,
                "batch_size": args.batch_size,
                "grad_accum": args.grad_accum,
                "num_iterations": args.num_iterations,
                "max_steps": args.max_steps,
                "rank": args.rank,
                "alpha": args.alpha,
                "reward_functions": [f.__name__ for f in reward_funcs],
                "reward_weights": reward_weights,
                "precision": "bfloat16" if bf16 else "float32",
            },
            seed=args.seed,
            losses=summary,
            scores={},
            config_path=__file__,
            notes=args.notes,
        )
        print(f"recorded run {record['run_id']} in {args.labbook}")


if __name__ == "__main__":
    main()