"""Make the demo eval log: ~1,500 prompts x 4 models, wide format, written to web/demo.csv.

    python data/make_demo.py            # writes web/demo.csv and data/demo_truth.json (the generator's hidden parameters)

The models are tiny / small / mid / frontier. Prices per million tokens span about 300x
(tiny $0.04 in / $0.16 out, frontier $8 / $32). Each prompt has a hidden difficulty shared by
all models, so hard prompts are hard for everyone (correlated errors). Target accuracies are about
62 / 74 / 84 / 92%. Each model also reports a confidence that is informative but imperfect, and
the small models are over-confident. Latency in milliseconds grows with model size and output length.
"""
from __future__ import annotations

import csv
import json
import math
import os

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)

MODELS = ["tiny", "small", "mid", "frontier"]
PRICE_IN = [0.04, 0.15, 1.0, 8.0]       # $ per million input tokens
PRICE_OUT = [0.16, 0.60, 4.0, 32.0]     # $ per million output tokens
SKILL = [0.8, 1.68, 2.45, 3.6]        # logit offsets; tuned for ~62/74/84/92% accuracy
OVERCONF = [0.9, 0.6, 0.35, 0.1]        # added to the confidence logit (small models brag)
CONF_NOISE = [0.9, 0.8, 0.7, 0.6]       # how noisy the self-report is
LAT_BASE = [120, 220, 450, 900]         # ms before the first token
LAT_PER_TOK = [1.5, 3.0, 8.0, 22.0]     # ms per output token

# prompt families: (template list, base difficulty, output tokens scale)
FAMILIES = {
    "arithmetic": (["What is {a} times {b}?", "Compute {a} plus {b} minus {c}.", "What is {a} divided by {d}, to two decimals?",
                    "Is {a} a prime number?", "What is {a} percent of {b}?"], -0.9, 20),
    "trivia": (["Who wrote {book}?", "What is the capital of {country}?", "In which year did {event} happen?",
                "Which element has the symbol {sym}?", "Who painted {painting}?"], -0.6, 25),
    "conversion": (["Convert {a} kilometres to miles.", "How many minutes are in {a} hours?", "Convert {a} degrees Celsius to Fahrenheit.",
                    "How many bytes are in {a} kilobytes?"], -0.7, 20),
    "rewrite": (["Rewrite this sentence in a formal tone: {sent}", "Fix the grammar: {sent}", "Make this shorter: {sent}",
                 "Translate to French: {sent}"], -0.1, 60),
    "extraction": (["Extract the date and the total from this receipt: {blob}", "List the people named in: {blob}",
                    "Turn this into JSON with keys name, city, amount: {blob}"], 0.3, 70),
    "code": (["Write a Python function that {task}.", "Fix the bug in this code: {code}", "Write a SQL query that {sqltask}.",
              "Explain what this regular expression matches: {regex}"], 0.6, 160),
    "reasoning": (["If {name1} is older than {name2} and {name2} is older than {name3}, who is youngest?",
                   "A train leaves at {h}:00 travelling {a} km/h. How far has it gone by {h2}:30?",
                   "{name1} has {a} apples, gives {c} to {name2} and buys {d} more. How many now?",
                   "Three boxes are labelled wrong. You may open one. Which do you open, and why?",
                   "Which weighs more, {a} kg of feathers or {c} kg of steel? Explain the trick."], 0.8, 90),
    "analysis": (["Summarize the main argument of this paragraph: {blob}", "What are the risks in this plan: {blob}",
                  "Compare the two options described here and recommend one: {blob}",
                  "Draft a polite reply declining this request: {blob}"], 1.1, 180),
    "math_proof": (["Prove that the sum of two even numbers is even.", "Show that sqrt({p}) is irrational.",
                    "Find all integer solutions of {a}x + {b}y = {c}.", "What is the derivative of x^{d} * sin(x)?",
                    "Solve for x: {a}x^2 - {b}x + {c} = 0"], 1.4, 120),
}
FAMILY_WEIGHTS = {"arithmetic": 0.13, "trivia": 0.14, "conversion": 0.08, "rewrite": 0.12, "extraction": 0.1,
                  "code": 0.12, "reasoning": 0.11, "analysis": 0.12, "math_proof": 0.08}

BOOKS = ["Middlemarch", "Beloved", "Dune", "The Idiot", "Persuasion", "Kim", "Ulysses", "Emma", "Dracula", "Hamlet"]
COUNTRIES = ["Mongolia", "Peru", "Ghana", "Latvia", "Nepal", "Uruguay", "Fiji", "Oman", "Chile", "Malta"]
EVENTS = ["the fall of Constantinople", "the first Moon landing", "the Treaty of Westphalia", "the Great Fire of London",
          "the opening of the Suez Canal", "the discovery of penicillin"]
SYMS = ["W", "Sn", "Hg", "Pb", "K", "Na", "Ag", "Au", "Fe", "Cu"]
PAINTINGS = ["The Night Watch", "Guernica", "The Kiss", "Olympia", "The Scream", "Las Meninas"]
SENTS = ["we gotta ship it friday no matter what", "me and him went to the store yesterday", "the report were late again",
         "pls send the file asap thx", "this don't work on my machine", "their going to announce it tomorrow",
         "the meeting got moved, idk when", "can u review my PR when you get a sec"]
BLOBS = ["Invoice 2291, 14 March, 3 items, total 182.40 EUR, paid by card.", "Minutes: Priya proposed the move, Tom objected, Lena abstained.",
         "Order from Marta Ruiz, Seville, 349.99, delivered Tuesday.", "Plan: cut prices 20%, double ad spend, hire 4 reps by Q3.",
         "Option A costs less but ships in June; option B costs 30% more and ships next week.",
         "Hi, could your team cover our on-call rotation over the holidays? Thanks, Jon.",
         "The study followed 212 patients for two years and found a modest effect that vanished after adjusting for age.",
         "Receipt: Cafe Lumen, 02/11, 2 flat whites, 1 croissant, total 11.80."]
TASKS = ["reverses the words in a sentence", "returns the n-th Fibonacci number iteratively", "checks whether a string is a palindrome",
         "merges two sorted lists", "counts vowels in a string", "parses a date like 2024-03-14 and returns the weekday",
         "flattens a nested list", "finds the longest common prefix of a list of strings"]
CODES = ["for i in range(len(xs)): total =+ xs[i]", "if x = 3: print('three')", "def f(a, b=[]): b.append(a); return b",
         "while i < n: print(i)", "return sorted(d.items(), key=lambda kv: kv[0])[1]"]
SQLTASKS = ["lists the top 5 customers by total spend", "counts orders per month in 2024", "finds products never ordered",
            "returns the second highest salary"]
REGEXES = ["^\\d{3}-\\d{4}$", "^[a-z]+@[a-z]+\\.(com|org)$", "\\b(\\w+)\\s+\\1\\b", "^(?=.*\\d)(?=.*[A-Z]).{8,}$"]
NAMES = ["Ana", "Ben", "Chloe", "Dev", "Esme", "Femi", "Greta", "Hugo", "Iris", "Jonas"]


def fill(template: str, rng: np.random.Generator) -> str:
    def pick(xs):
        return xs[rng.integers(len(xs))]
    return template.format(
        a=int(rng.integers(2, 999)), b=int(rng.integers(2, 999)), c=int(rng.integers(1, 99)), d=int(rng.integers(2, 12)),
        p=pick([2, 3, 5, 7, 11, 13]), h=int(rng.integers(1, 11)), h2=int(rng.integers(12, 23)),
        book=pick(BOOKS), country=pick(COUNTRIES), event=pick(EVENTS), sym=pick(SYMS), painting=pick(PAINTINGS),
        sent=pick(SENTS), blob=pick(BLOBS), task=pick(TASKS), code=pick(CODES), sqltask=pick(SQLTASKS), regex=pick(REGEXES),
        name1=pick(NAMES), name2=pick(NAMES), name3=pick(NAMES))


def sigmoid(x):
    return 1.0 / (1.0 + np.exp(-x))


def make(n: int = 1500, seed: int = 7):
    rng = np.random.default_rng(seed)
    fams = list(FAMILIES)
    w = np.array([FAMILY_WEIGHTS[f] for f in fams])
    w /= w.sum()
    rows, truth = [], []
    for k in range(n):
        fam = fams[rng.choice(len(fams), p=w)]
        templates, base, out_scale = FAMILIES[fam]
        prompt = fill(templates[rng.integers(len(templates))], rng)
        # hidden difficulty: family + length + idiosyncratic part shared by all models
        diff = base + 0.25 * (len(prompt) - 60) / 40.0 + rng.normal(0, 1.0)
        in_tok = max(8, int(round(len(prompt) / 3.6 + rng.normal(0, 3))))
        out_tok = max(4, int(round(out_scale * rng.lognormal(0, 0.45))))
        row = {"prompt_id": f"p{k + 1:04d}", "prompt": prompt}
        t = {"family": fam, "difficulty": float(diff)}
        for i, name in enumerate(MODELS):
            logit = SKILL[i] - 1.1 * diff
            p_right = sigmoid(logit)
            correct = int(rng.random() < p_right)
            # self-reported confidence: noisy view of the true logit, shifted up for small models,
            # a bit higher when actually right
            conf_logit = logit + OVERCONF[i] + 1.2 * (correct - 0.5) + rng.normal(0, CONF_NOISE[i])
            conf = float(np.clip(sigmoid(conf_logit), 0.01, 0.99))
            tok_out = max(3, int(round(out_tok * (0.7 + 0.15 * i) * rng.lognormal(0, 0.2))))
            cost = (in_tok * PRICE_IN[i] + tok_out * PRICE_OUT[i]) / 1e6
            lat = LAT_BASE[i] * rng.lognormal(0, 0.25) + LAT_PER_TOK[i] * tok_out
            row[f"{name}_correct"] = correct
            row[f"{name}_cost"] = f"{cost:.7f}"
            row[f"{name}_latency"] = int(round(lat))
            row[f"{name}_confidence"] = f"{conf:.3f}"
            t[f"{name}_p"] = float(p_right)
        rows.append(row)
        truth.append(t)
    return rows, truth


def main():
    rows, truth = make()
    out = os.path.join(ROOT, "web", "demo.csv")
    cols = ["prompt_id", "prompt"] + [f"{m}_{f}" for m in MODELS for f in ("correct", "cost", "latency", "confidence")]
    with open(out, "w", newline="", encoding="utf-8") as f:
        wr = csv.DictWriter(f, fieldnames=cols)
        wr.writeheader()
        wr.writerows(rows)
    with open(os.path.join(HERE, "demo_truth.json"), "w", encoding="utf-8") as f:
        json.dump({"models": MODELS, "price_in": PRICE_IN, "price_out": PRICE_OUT, "skill": SKILL,
                   "mean_p_right": {m: float(np.mean([t[f"{m}_p"] for t in truth])) for m in MODELS},
                   "families": {f: sum(t["family"] == f for t in truth) for f in FAMILIES}}, f, indent=1)
    acc = {m: np.mean([r[f"{m}_correct"] for r in rows]) for m in MODELS}
    cost = {m: np.mean([float(r[f"{m}_cost"]) for r in rows]) for m in MODELS}
    print(f"wrote {out}: {len(rows)} prompts")
    for m in MODELS:
        print(f"  {m:9s} accuracy {acc[m]:.3f}  mean cost ${cost[m]:.6f}  (${1000 * cost[m]:.3f} per 1k)")
    print(f"  cost span {cost['frontier'] / cost['tiny']:.0f}x; size {os.path.getsize(out) / 1024:.0f} KB")


if __name__ == "__main__":
    main()
