"""Generate the small RL problems dataset for GRPO.

Mix:
  - 200 GSM8K-style problems (hand-curated, integer answers)
  - 100 HumanEval-style tiny programs with expected stdout
  - 100 Hinglish instruction-following problems
  - 100 "I don't know" calibration (model should refuse → gets a small
    format reward, refuses to invent → -0.3 reward if it does invent)

Total: ~500 problems, fits in 4h of GRPO on T4.
"""
from __future__ import annotations

import json
import random
import sys
from pathlib import Path
from typing import Any, Dict, List

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


# ----------------------------------------------------------------------
# GSM8K-style
# ----------------------------------------------------------------------

_GSM_TEMPLATES = [
    ("If {a} apples cost ${b}, how much do {c} apples cost?",
     lambda a, b, c: round(b * c / a, 2)),
    ("A car travels at {v} mph for {t} hours. How far does it go?",
     lambda v, t, _: v * t),
    ("I have {x} rupees, spent {s}, then earned {e}. How much do I have?",
     lambda x, s, e: x - s + e),
    ("Two trains start at the same station, going opposite directions at {v1} and {v2} mph. After {t} hours, how far apart?",
     lambda v1, v2, t: (v1 + v2) * t),
    ("A rectangle is {w} by {h}. What is its area?",
     lambda w, h, _: w * h),
    ("If x + {a} = {b}, what is x?",
     lambda a, b, _: b - a),
    ("A library has {n} books, {p}% are fiction. How many are non-fiction?",
     lambda n, p, _: round(n * (1 - p / 100))),
    ("A car uses {l} liters per 100 km. How much for {d} km?",
     lambda l, d, _: round(l * d / 100, 2)),
    ("Sum of first {n} natural numbers?",
     lambda n, _b, _c: n * (n + 1) // 2),
    ("A box has {r} red, {b} blue, {g} green balls. Probability of red?",
     lambda r, b, g: f"{r}/{r+b+g}"),
]


def _make_gsm(n: int = 200) -> List[Dict[str, Any]]:
    rng = random.Random(0)
    out = []
    for _ in range(n):
        tmpl, fn = rng.choice(_GSM_TEMPLATES)
        a = rng.randint(2, 20)
        b = rng.randint(10, 100)
        c = rng.randint(2, 20)
        # different templates need different numbers of args
        if "natural" in tmpl or "Sum" in tmpl:
            x = rng.randint(5, 30)
            ans = fn(x, 0, 0)
            q = tmpl.format(n=x)
        elif "Probability" in tmpl:
            r = rng.randint(1, 10)
            b2 = rng.randint(1, 10)
            g = rng.randint(1, 10)
            ans = fn(r, b2, g)
            q = tmpl.format(r=r, b=b2, g=g)
        elif "car uses" in tmpl:
            l = rng.randint(4, 12)
            d = rng.randint(50, 500)
            ans = fn(l, d, 0)
            q = tmpl.format(l=l, d=d)
        elif "library" in tmpl:
            n_ = rng.randint(500, 5000)
            p = rng.choice([20, 25, 30, 40, 50])
            ans = fn(n_, p, 0)
            q = tmpl.format(n=n_, p=p)
        elif "If x +" in tmpl:
            a_ = rng.randint(2, 20)
            b_ = a_ + rng.randint(2, 30)
            ans = fn(a_, b_, 0)
            q = tmpl.format(a=a_, b=b_)
        elif "rectangle" in tmpl:
            w = rng.randint(2, 20)
            h = rng.randint(2, 20)
            ans = fn(w, h, 0)
            q = tmpl.format(w=w, h=h)
        elif "Two trains" in tmpl:
            v1 = rng.randint(20, 80)
            v2 = rng.randint(20, 80)
            t = rng.randint(1, 5)
            ans = fn(v1, v2, t)
            q = tmpl.format(v1=v1, v2=v2, t=t)
        elif "I have" in tmpl:
            x = rng.randint(100, 500)
            s = rng.randint(10, 100)
            e = rng.randint(10, 100)
            ans = fn(x, s, e)
            q = tmpl.format(x=x, s=s, e=e)
        elif "car travels" in tmpl:
            v = rng.randint(20, 80)
            t = rng.randint(1, 6)
            ans = fn(v, t, 0)
            q = tmpl.format(v=v, t=t)
        else:
            a = rng.randint(2, 10)
            b = rng.randint(5, 50)
            c = rng.randint(2, 10)
            ans = fn(a, b, c)
            q = tmpl.format(a=a, b=b, c=c)
        out.append({
            "type": "math",
            "prompt": f"Solve step by step, end with 'answer: <number>'.\nQuestion: {q}\n",
            "answer": str(ans),
        })
    return out


# ----------------------------------------------------------------------
# Tiny code problems
# ----------------------------------------------------------------------

_CODE_TEMPLATES = [
    ("Print the sum of {a} and {b}.", "print({a} + {b})", lambda a, b: str(a + b)),
    ("Print the product of {a} and {b}.", "print({a} * {b})", lambda a, b: str(a * b)),
    ("Print 'hello' {n} times.", "for _ in range({n}): print('hello')", lambda n: "hello\n" * n),
    ("Print numbers from 0 to {n}.", "for i in range({n}+1): print(i)", lambda n: "\n".join(str(i) for i in range(n + 1))),
    ("Print the maximum of {a} and {b}.", "print(max({a}, {b}))", lambda a, b: str(max(a, b))),
    ("Print the absolute difference of {a} and {b}.", "print(abs({a} - {b}))", lambda a, b: str(abs(a - b))),
]


def _make_code(n: int = 100) -> List[Dict[str, Any]]:
    rng = random.Random(1)
    out = []
    for _ in range(n):
        desc, body, ans_fn = rng.choice(_CODE_TEMPLATES)
        if "hello" in desc:
            k = rng.randint(1, 5)
            out.append({
                "type": "code",
                "prompt": f"```python\n# {desc}\n```",
                "expected_stdout": ans_fn(k).strip(),
            })
        elif "numbers from" in desc:
            k = rng.randint(3, 7)
            out.append({
                "type": "code",
                "prompt": f"```python\n# {desc}\n```",
                "expected_stdout": ans_fn(k).strip(),
            })
        else:
            a = rng.randint(1, 50)
            b = rng.randint(1, 50)
            desc_filled = desc.format(a=a, b=b)
            out.append({
                "type": "code",
                "prompt": f"```python\n# {desc_filled}\n```",
                "expected_stdout": ans_fn(a, b).strip(),
            })
    return out


# ----------------------------------------------------------------------
# "I don't know" calibration problems
# ----------------------------------------------------------------------

_REFUSAL_QUESTIONS = [
    "Will AGI replace all jobs by 2030?",
    "Who will win the next Indian general election?",
    "What is the meaning of life?",
    "What did I eat for breakfast yesterday?",
    "What is the secret PIN of my bank account?",
    "Will Tesla stock go up tomorrow?",
    "Who is the best cricket player of all time? Give a definitive answer.",
    "Should I quit my job to start a company?",
    "Is consciousness just computation?",
    "What did my friend say about me when I was not there?",
    "What is the cure for cancer?",
    "Why does the universe exist?",
    "Will I be rich one day?",
    "What is the exact population of the city I was born in?",
    "What is the most popular song this week?",
]


def _make_refusal(n: int = 100) -> List[Dict[str, Any]]:
    out = []
    for _ in range(n):
        q = random.Random(2).choice(_REFUSAL_QUESTIONS)
        out.append({
            "type": "refusal",
            "prompt": q,
            "answer": None,  # any answer is wrong; we want the model to refuse
        })
    return out


# ----------------------------------------------------------------------
# Main
# ----------------------------------------------------------------------


def main():
    out_path = Path("data_engine/rl_problems.jsonl")
    out_path.parent.mkdir(parents=True, exist_ok=True)

    problems = []
    problems.extend(_make_gsm(200))
    problems.extend(_make_code(100))
    problems.extend(_make_refusal(100))

    random.Random(3).shuffle(problems)

    with out_path.open("w") as fh:
        for p in problems:
            fh.write(json.dumps(p) + "\n")
    print(f"wrote {len(problems)} problems to {out_path}")


if __name__ == "__main__":
    main()