"""Shared helpers for the Qwen3-4B GRPO comparison scripts.

Exports:
    REASONING_START, REASONING_END, SOLUTION_START, SOLUTION_END, SYSTEM_PROMPT
    CHAT_TEMPLATE
    build_dataset(tokenizer, max_seq_length=2048)
    build_reward_funcs(tokenizer)
    build_grpo_kwargs(tokenizer, maximum_length, max_seq_length)
    StepTimer  (TrainerCallback recording per-step wall time, loss, reward)
    write_stats(path, backend, timer, ...)
    install_vllm_sampling_shim()

Keeps dataset loading, chat template, formatting rewards, and GRPO hparams
identical between the vLLM baseline script and the transformers-CB candidate.
"""

from __future__ import annotations

import json
import re
import time

import numpy as np
import pandas as pd
from datasets import Dataset, load_dataset


REASONING_START = "<start_working_out>"
REASONING_END = "<end_working_out>"
SOLUTION_START = "<SOLUTION>"
SOLUTION_END = "</SOLUTION>"

SYSTEM_PROMPT = (
    "You are given a problem.\n"
    "Think about the problem and provide your working out.\n"
    f"Place it between {REASONING_START} and {REASONING_END}.\n"
    f"Then, provide your solution between {SOLUTION_START}{SOLUTION_END}"
)


CHAT_TEMPLATE = (
    "{% if messages[0]['role'] == 'system' %}"
    "{{ messages[0]['content'] + eos_token }}"
    "{% set loop_messages = messages[1:] %}"
    "{% else %}"
    "{{ '%%%SYSTEM_PROMPT%%%' + eos_token }}"
    "{% set loop_messages = messages %}"
    "{% endif %}"
    "{% for message in loop_messages %}"
    "{% if message['role'] == 'user' %}"
    "{{ message['content'] }}"
    "{% elif message['role'] == 'assistant' %}"
    "{{ message['content'] + eos_token }}"
    "{% endif %}"
    "{% endfor %}"
    "{% if add_generation_prompt %}{{ '%%%REASONING_START%%%' }}"
    "{% endif %}"
)


def apply_chat_template_to_tokenizer(tokenizer):
    tmpl = CHAT_TEMPLATE.replace("%%%SYSTEM_PROMPT%%%", SYSTEM_PROMPT)
    tmpl = tmpl.replace("%%%REASONING_START%%%", REASONING_START)
    tokenizer.chat_template = tmpl
    return tokenizer


def build_dataset(tokenizer, *, max_seq_length: int = 2048):
    """Build the DAPO-Math-17k GRPO dataset with the prompt formatting from the
    notebook. Returns `(dataset, maximum_prompt_length)`.
    """
    ds = load_dataset("open-r1/DAPO-Math-17k-Processed", "en", split = "train")

    def _map_row(x):
        return {
            "prompt": [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": x["prompt"]},
            ],
            "answer": x["solution"],
        }

    ds = ds.map(_map_row)

    # Tokenize for length measurement (batched for speed).
    def _tokenize(batch):
        return {
            "tokens": tokenizer.apply_chat_template(
                batch["prompt"],
                add_generation_prompt = True,
                tokenize = True,
            )
        }

    tokenized = ds.map(_tokenize, batched = True)
    tokenized = tokenized.map(lambda x: {"L": len(x["tokens"])})
    lengths = np.array(tokenized["L"])
    maximum_length = int(np.quantile(lengths, 0.9))
    ds = ds.select(np.where(lengths <= maximum_length)[0])
    return ds, maximum_length


def build_reward_funcs(tokenizer):
    """Return the 4 reward functions used in the notebook, wired to `tokenizer`."""
    solution_end_regex = (
        r"</SOLUTION>[\s]{0,}" + "(?:" + re.escape(tokenizer.eos_token) + ")?"
    )
    match_format = re.compile(
        rf"{REASONING_END}.*?"
        rf"{SOLUTION_START}(.+?){solution_end_regex}"
        rf"[\s]{{0,}}$",
        flags = re.MULTILINE | re.DOTALL,
    )
    match_numbers = re.compile(
        SOLUTION_START + r".*?[\s]{0,}([-]?[\d\.\,]{1,})",
        flags = re.MULTILINE | re.DOTALL,
    )

    def match_format_exactly(completions, **kwargs):
        scores = []
        for completion in completions:
            score = 0.0
            response = completion[0]["content"]
            if match_format.search(response) is not None:
                score += 3.0
            scores.append(score)
        return scores

    def match_format_approximately(completions, **kwargs):
        scores = []
        for completion in completions:
            score = 0.0
            response = completion[0]["content"]
            score += 0.5 if response.count(REASONING_END) == 1 else -1.0
            score += 0.5 if response.count(SOLUTION_START) == 1 else -1.0
            score += 0.5 if response.count(SOLUTION_END) == 1 else -1.0
            scores.append(score)
        return scores

    def check_answer(prompts, completions, answer, **kwargs):
        responses = [c[0]["content"] for c in completions]
        extracted = [
            guess.group(1) if (guess := match_format.search(r)) is not None else None
            for r in responses
        ]
        scores = []
        for guess, true_answer in zip(extracted, answer):
            score = 0.0
            if guess is None:
                scores.append(-2.0)
                continue
            if guess == true_answer:
                score += 5.0
            elif guess.strip() == true_answer.strip():
                score += 3.5
            else:
                try:
                    ratio = float(guess) / float(true_answer)
                    if 0.9 <= ratio <= 1.1:
                        score += 2.0
                    elif 0.8 <= ratio <= 1.2:
                        score += 1.5
                    else:
                        score -= 2.5
                except Exception:
                    score -= 4.5
            scores.append(score)
        return scores

    _printed_state = {"n": 0, "every": 5}

    def check_numbers(prompts, completions, answer, **kwargs):
        question = prompts[0][-1]["content"]
        responses = [c[0]["content"] for c in completions]
        extracted = [
            guess.group(1) if (guess := match_numbers.search(r)) is not None else None
            for r in responses
        ]
        if _printed_state["n"] % _printed_state["every"] == 0:
            print(
                "*" * 20 + f"Question:\n{question}",
                f"\nAnswer:\n{answer[0]}",
                f"\nResponse:\n{responses[0]}",
                f"\nExtracted:\n{extracted[0]}",
            )
        _printed_state["n"] += 1

        scores = []
        for guess, true_answer in zip(extracted, answer):
            if guess is None:
                scores.append(-2.5)
                continue
            try:
                t = float(true_answer.strip())
                g = float(guess.strip().replace(",", ""))
                scores.append(3.5 if g == t else -1.5)
            except Exception:
                scores.append(0.0)
        return scores

    return [
        match_format_exactly,
        match_format_approximately,
        check_answer,
        check_numbers,
    ]


def build_grpo_kwargs(
    tokenizer,
    maximum_length: int,
    *,
    max_seq_length: int = 2048,
    max_steps: int = 100,
    num_generations: int = 4,
    per_device_train_batch_size: int = 1,
    gradient_accumulation_steps: int = 1,
    output_dir: str = "outputs",
):
    """Return the shared dict of GRPOConfig kwargs used by both backends.

    Caller adds backend-specific keys (use_vllm / use_transformers_paged / etc).
    """
    max_prompt_length = maximum_length + 1
    max_completion_length = max_seq_length - max_prompt_length

    return dict(
        temperature = 1.0,
        top_p = 1.0,
        top_k = -1,
        min_p = 0.1,
        learning_rate = 5e-6,
        weight_decay = 0.001,
        warmup_ratio = 0.1,
        lr_scheduler_type = "linear",
        optim = "adamw_8bit",
        logging_steps = 1,
        per_device_train_batch_size = per_device_train_batch_size,
        gradient_accumulation_steps = gradient_accumulation_steps,
        num_generations = num_generations,
        max_prompt_length = max_prompt_length,
        max_completion_length = max_completion_length,
        max_steps = max_steps,
        save_steps = max_steps,
        report_to = "none",
        output_dir = output_dir,
        seed = 3407,
    )


import torch
from transformers import TrainerCallback


class StepTimer(TrainerCallback):
    """Per-step wall time / loss / reward recorder.

    Shared verbatim across qwen3_grpo_{vllm,naive,tpaged}. Records step wall
    time in `self.step_wall`, and picks up `loss` / `reward` from the TRL log
    dict in `on_log`.
    """

    def __init__(self):
        self.t0 = None
        self.step_wall = []
        self.loss = []
        self.reward = []

    def on_step_begin(self, _args, state, control, **kwargs):
        torch.cuda.synchronize()
        self.t0 = time.perf_counter()

    def on_log(self, _args, state, control, logs = None, **kwargs):
        if logs is None:
            return
        if "loss" in logs:
            self.loss.append(float(logs["loss"]))
        if "reward" in logs:
            self.reward.append(float(logs["reward"]))

    def on_step_end(self, _args, state, control, **kwargs):
        if self.t0 is not None:
            torch.cuda.synchronize()
            self.step_wall.append(time.perf_counter() - self.t0)


def write_stats(
    path: str,
    backend: str,
    timer: StepTimer,
    *,
    train_wall_s: float,
    peak_memory_gb: float,
    max_prompt_length: int,
    max_completion_length: int,
    num_generations: int,
    max_steps: int,
    extra: dict | None = None,
) -> None:
    """Dump the per-step stats dict used by all three GRPO drivers.

    Schema matches the pre-refactor output exactly: `backend`, `train_wall_s`,
    `peak_memory_gb`, `step_wall_s`, `losses`, `rewards`, `max_prompt_length`,
    `max_completion_length`, `num_generations`, `max_steps`, plus any
    backend-specific keys passed in `extra` (e.g. `attn_impl`, `persistent_cb`).
    """
    stats = {
        "backend": backend,
        "train_wall_s": train_wall_s,
        "peak_memory_gb": peak_memory_gb,
        "step_wall_s": timer.step_wall,
        "losses": timer.loss,
        "rewards": timer.reward,
        "max_prompt_length": max_prompt_length,
        "max_completion_length": max_completion_length,
        "num_generations": num_generations,
        "max_steps": max_steps,
    }
    if extra:
        stats.update(extra)
    with open(path, "w") as f:
        json.dump(stats, f, indent = 2)


def maybe_compile_trainer_forwards(
    trainer, compile_mode, *, dynamic: bool = True, tag: str = ""
):
    """torch.compile wrap `trainer.model.forward` and (if present)
    `trainer.ref_model.forward`. No-op if `compile_mode` is falsy.

    Ported from the old `qwen3_grpo_unified.py` compile path, minus the
    out-of-tree `torch_debugging_utils` imports that were dev-only.
    """
    if not compile_mode:
        return
    import torch._dynamo

    torch._dynamo.config.cache_size_limit = 128
    try:
        torch._dynamo.config.allow_unspec_int_on_nn_module = True
    except AttributeError:
        pass
    prefix = f"[{tag}] " if tag else ""
    print(
        f"{prefix}Compiling trainer.model.forward (mode={compile_mode}, dynamic={dynamic})"
    )
    trainer.model.forward = torch.compile(
        trainer.model.forward,
        mode = compile_mode,
        dynamic = dynamic,
    )
    ref = getattr(trainer, "ref_model", None)
    if ref is not None:
        ref.forward = torch.compile(
            ref.forward,
            mode = compile_mode,
            dynamic = dynamic,
        )


def install_vllm_sampling_shim():
    """Shim `vllm.sampling_params.GuidedDecodingParams` for newer vLLM releases.

    TRL's `GRPOTrainer` imports `GuidedDecodingParams` from
    `vllm.sampling_params`; newer vLLM versions have moved or removed it. Inject
    a no-op class so the import succeeds even on the non-vLLM training paths
    (naive, tpaged). No-op if vLLM is not installed or already exposes the
    symbol.
    """
    try:
        import vllm.sampling_params as _vllm_sp
    except ImportError:
        return
    if hasattr(_vllm_sp, "GuidedDecodingParams"):
        return

    class _GuidedDecodingParamsShim:  # pragma: no cover - used only if TRL asks
        def __init__(self, *a, **kw):
            pass

    _vllm_sp.GuidedDecodingParams = _GuidedDecodingParamsShim