#!/usr/bin/env python3
"""
Micro-ARC (MARC-1 → MARC-3): a tiny, code-only benchmark for small language models.

Four task families, three difficulty tiers, zero data files. Every task is
generated from a seed, so the "dataset" is a function call and contamination
is impossible without reverse-engineering the grammar.

Leak-proofing invariants (verified by `--selftest`):
  (I1) query_in never equals any demo input   — no copy-the-demo shortcut.
  (I2) no demo is the identity (din == dout)  — no copy-the-input shortcut.
  (I3) pattern shifts are not multiples of the period — no prefix/suffix
       degradation of the pattern family.

Usage (single model):
    python marc.py --hf HuggingFaceTB/SmolLM2-360M
    python marc.py --hf Qwen/Qwen3-0.5B-Instruct --chat --n-samples 5 --temperature 0.7
    python marc.py --openai Qwen/Qwen3-32B --base-url http://localhost:8000/v1 --chat
    python marc.py --hf dotlabs/void.1 --trust-remote-code

Usage (multi-model bucket — run any mix of HF + OpenAI models, print a summary):
    python marc.py --hf HuggingFaceTB/SmolLM2-360M --hf Qwen/Qwen3-0.5B
    python marc.py --hf a --hf b --openai c --base-url http://localhost:8000/v1
    python marc.py --openai m1 --openai m2 \\
        --base-url http://h1/v1 --base-url http://h2/v1

Usage (verify leak invariants):
    python marc.py --selftest
"""
import argparse
import gc
import random
import re
import sys
import time
from dataclasses import dataclass
from typing import Callable

Grid = list[list[int]]


# ═══════════════════════════════════════════════════════════════════════════
# 1. Tasks — four generators, three difficulty tiers, all seeded
# ═══════════════════════════════════════════════════════════════════════════

@dataclass
class Task:
    demos: list[tuple[Grid, Grid]]
    query_in: Grid
    query_out: Grid
    family: str
    difficulty: int


def _gen_pattern(rng: random.Random, diff: int) -> Task:
    if diff == 1:
        periods, n_ins, n_out, syms = [2, 3], [3, 4], 3, [0, 1]
    elif diff == 2:
        periods, n_ins, n_out, syms = [2, 3], [4, 5], 4, [0, 1, 2]
    else:
        periods, n_ins, n_out, syms = [3, 4], [5, 6], 5, [0, 1, 2, 3]

    # Exclude n_in multiples of the period: those make the demo output a copy
    # of the input (or a trivial prefix), which teaches nothing and lets the
    # model echo input to score. See invariant (I2)/(I3).
    combos = [(p, n) for p in periods for n in n_ins if n % p != 0]
    if not combos:
        combos = [(p, n) for p in periods for n in n_ins]
    period, n_in = rng.choice(combos)

    # Reject bases whose minimal period is < period: otherwise two distinct
    # phases can yield identical sequences, breaking the "distinct phase ⇒
    # distinct sequence" guarantee used below.
    def has_min_period(s, p):
        return all(any(s[i] != s[(i + k) % p] for i in range(p))
                   for k in range(1, p))

    while True:
        base = [rng.choice(syms) for _ in range(period)]
        if len(set(base)) > 1 and has_min_period(base, period):
            break

    def seq(n, start):
        return [base[(start + i) % period] for i in range(n)]

    # Query phase first, demos from the complement ⇒ query_in ≠ every demo_in.
    s_query = rng.randint(0, period - 1)
    other_ss = [s for s in range(period) if s != s_query]
    rng.shuffle(other_ss)
    demo_ss = other_ss[:3]                        # 1 demo when period == 2

    demos = [([seq(n_in, s)], [seq(n_out, s + n_in)]) for s in demo_ss]
    return Task(demos, [seq(n_in, s_query)], [seq(n_out, s_query + n_in)],
                "pattern", diff)


def _gen_object(rng: random.Random, diff: int) -> Task:
    size = {1: 3, 2: 4, 3: 5}[diff]
    kind = ("translate" if diff == 1 else
            rng.choice(["translate", "recolor"]) if diff == 2 else
            rng.choice(["translate", "recolor", "mirror_h", "mirror_v"]))

    if kind == "translate":
        dr, dc = rng.choice([(-1, 0), (1, 0), (0, -1), (0, 1)])
        r_lo, r_hi = max(0, -dr), size - 1 - max(0, dr)
        c_lo, c_hi = max(0, -dc), size - 1 - max(0, dc)
    else:
        r_lo, r_hi, c_lo, c_hi = 0, size - 1, 0, size - 1

    def grid_at(r, c):
        g = [[0] * size for _ in range(size)]
        g[r][c] = 1
        return g

    if kind == "translate":
        def transform(g):
            out = [[0] * size for _ in range(size)]
            for r in range(size):
                for c in range(size):
                    if g[r][c]:
                        out[r + dr][c + dc] = g[r][c]
            return out
    elif kind == "recolor":
        def transform(g):
            return [[2 if v == 1 else 0 for v in row] for row in g]
    elif kind == "mirror_h":
        def transform(g):
            return [row[::-1] for row in g]
    else:
        def transform(g):
            return g[::-1]

    cells = [(r, c) for r in range(r_lo, r_hi + 1) for c in range(c_lo, c_hi + 1)]
    rng.shuffle(cells)

    usable = [rc for rc in cells if transform(grid_at(*rc)) != grid_at(*rc)]
    if len(usable) < 4:                           # cannot happen for size ≥ 3
        return _gen_object(rng, diff)

    q_rc, d_rcs = usable[0], usable[1:4]
    demos = [(grid_at(r, c), transform(grid_at(r, c))) for r, c in d_rcs]
    qin = grid_at(*q_rc)
    return Task(demos, qin, transform(qin), "object", diff)


def _gen_counting(rng: random.Random, diff: int) -> Task:
    if diff == 1:
        size, max_n, mode = 3, 3, "number"
    elif diff == 2:
        size, max_n = 4, 5
        mode = rng.choice(["number", "bar"])
    else:
        size, max_n = 5, 6
        mode = rng.choice(["number", "bar", "compare"])

    def make_grid(n_ones, n_twos=0):
        g = [[0] * size for _ in range(size)]
        cells = rng.sample([(r, c) for r in range(size) for c in range(size)],
                           n_ones + n_twos)
        for r, c in cells[:n_ones]:
            g[r][c] = 1
        for r, c in cells[n_ones:]:
            g[r][c] = 2
        return g

    if mode == "compare":
        qn1, qn2 = rng.randint(1, 4), rng.randint(1, 4)
        while qn1 == qn2:
            qn2 = rng.randint(1, 4)
        qout = [[1]] if qn1 > qn2 else [[2]]
        all_pairs = [(a, b) for a in range(1, 5) for b in range(1, 5)
                     if a != b and (a, b) != (qn1, qn2)]
        rng.shuffle(all_pairs)
        demos = [(make_grid(a, b), [[1]] if a > b else [[2]])
                 for a, b in all_pairs[:3]]
        return Task(demos, make_grid(qn1, qn2), qout, "counting", diff)

    # number / bar: output is a function of the count, so query count must
    # differ from every demo count.
    q_count = rng.randint(1, max_n)
    other_counts = [c for c in range(1, max_n + 1) if c != q_count]
    if len(other_counts) >= 3:
        demo_counts = rng.sample(other_counts, 3)
    elif len(other_counts) == 2:
        demo_counts = other_counts + [rng.choice(other_counts)]
    else:
        demo_counts = other_counts * 3

    def out_for(count):
        return [[count]] if mode == "number" else [[2] * count]

    demos = [(make_grid(c), out_for(c)) for c in demo_counts]
    return Task(demos, make_grid(q_count), out_for(q_count), "counting", diff)


def _gen_compositional(rng: random.Random, diff: int) -> Task:
    size = {1: 3, 2: 4, 3: 5}[diff]
    n_ops = 2 if diff < 3 else 3

    ops = [("recolor", rng.choice([(1, 2), (2, 1), (1, 3), (3, 1)]))]
    spatial = ["translate", "mirror_h", "mirror_v"]
    rng.shuffle(spatial)
    for s in spatial:
        if len(ops) >= n_ops:
            break
        ops.append(("translate", rng.choice([(-1, 0), (1, 0), (0, -1), (0, 1)]))
                   if s == "translate" else (s, None))
    while len(ops) < n_ops:
        ops.append(("recolor", rng.choice([(1, 2), (2, 1)])))
    rng.shuffle(ops)

    def transform(g):
        out = [row[:] for row in g]
        for op, p in ops:
            if op == "translate":
                dr, dc = p
                new = [[0] * size for _ in range(size)]
                for r in range(size):
                    for c in range(size):
                        if out[r][c] and 0 <= r + dr < size and 0 <= c + dc < size:
                            new[r + dr][c + dc] = out[r][c]
                out = new
            elif op == "recolor":
                out = [[p[1] if v == p[0] else v for v in row] for row in out]
            elif op == "mirror_h":
                out = [row[::-1] for row in out]
            else:
                out = out[::-1]
        return out

    def grid_at(r, c):
        g = [[0] * size for _ in range(size)]
        g[r][c] = 1
        return g

    def nz(g):
        return sum(1 for row in g for v in row if v)

    cells = [(r, c) for r in range(size) for c in range(size)]
    rng.shuffle(cells)

    query_cands, demo_cands = [], []
    for rc in cells:
        g = grid_at(*rc)
        out = transform(g)
        if nz(out) != 1:
            continue
        query_cands.append(rc)
        if out != g:
            demo_cands.append(rc)

    if not query_cands or len(demo_cands) < 4:
        return _gen_compositional(rng, diff)

    q_rc = query_cands[0]
    d_rcs = [rc for rc in demo_cands if rc != q_rc][:3]
    if len(d_rcs) < 3:
        return _gen_compositional(rng, diff)

    qin = grid_at(*q_rc)
    demos = [(grid_at(r, c), transform(grid_at(r, c))) for r, c in d_rcs]
    return Task(demos, qin, transform(qin), "compositional", diff)


FAMILIES = {
    "pattern": _gen_pattern,
    "object": _gen_object,
    "counting": _gen_counting,
    "compositional": _gen_compositional,
}


def generate_task(family: str, difficulty: int, seed: str) -> Task:
    # string seeds are sha512-hashed by random.Random → deterministic across processes
    return FAMILIES[family](random.Random(seed), difficulty)


# ═══════════════════════════════════════════════════════════════════════════
# 2. Format — serialization, prompts, parsing
# ═══════════════════════════════════════════════════════════════════════════

def grid_to_text(g: Grid) -> str:
    return "\n".join(" ".join(str(v) for v in row) for row in g)


def _body(task: Task) -> str:
    parts = []
    for i, (inp, out) in enumerate(task.demos, 1):
        parts += [f"Example {i}:", "Input:", grid_to_text(inp),
                  "Output:", grid_to_text(out), ""]
    parts += ["Now solve this one:", "Input:", grid_to_text(task.query_in), "Output:"]
    return "\n".join(parts)


def prompt_for(task: Task, chat: bool):
    if not chat:
        return _body(task)
    return [
        {"role": "system", "content":
            "You are a pattern solver. Each example shows an input grid and its "
            "transformed output grid. Infer the rule and output ONLY the output "
            "grid for the final input, in the same format (one row per line, "
            "integers separated by spaces). Do not explain."},
        {"role": "user", "content": _body(task)},
    ]


_ROW = re.compile(r"^-?\d+(\s+-?\d+)*$")


def parse_grid(text: str) -> Grid | None:
    rows: Grid = []
    for raw in text.splitlines():
        line = raw.strip().strip("[]").replace(",", " ")
        if not line:
            if rows:
                break
            continue
        if not _ROW.match(line):
            if rows:
                break
            continue
        try:
            rows.append([int(t) for t in line.split()])
        except ValueError:
            if rows:
                break
    if not rows or any(len(r) != len(rows[0]) for r in rows):
        return None
    return rows


# ═══════════════════════════════════════════════════════════════════════════
# 3. Models — HF (base or instruct) + OpenAI-compatible endpoint
# ═══════════════════════════════════════════════════════════════════════════

class HFModel:
    def __init__(self, name: str, chat: bool, device: str = "cuda",
                 trust_remote_code: bool = False):
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer
        self.tok = AutoTokenizer.from_pretrained(
            name, trust_remote_code=trust_remote_code)
        if self.tok.pad_token is None:
            self.tok.pad_token = self.tok.eos_token
        self.tok.padding_side = "left"
        self.model = AutoModelForCausalLM.from_pretrained(
            name, torch_dtype=torch.bfloat16,
            trust_remote_code=trust_remote_code).to(device).eval()
        self.device, self.chat, self.torch = device, chat, torch

    def generate(self, prompts, n, temperature, max_new_tokens):
        out_all = []
        for p in prompts:
            text = (self.tok.apply_chat_template(p, tokenize=False, add_generation_prompt=True)
                    if self.chat else p)
            inp = self.tok(text, return_tensors="pt", truncation=True).to(self.device)
            with self.torch.no_grad():
                out = self.model.generate(
                    **inp, max_new_tokens=max_new_tokens,
                    do_sample=temperature > 0,
                    temperature=max(temperature, 1e-5),
                    num_return_sequences=n, pad_token_id=self.tok.pad_token_id)
            L = inp.input_ids.shape[1]
            out_all.append([self.tok.decode(o[L:], skip_special_tokens=True) for o in out])
        return out_all


class OpenAIModel:
    def __init__(self, model: str, base_url: str, chat: bool, api_key: str = "EMPTY"):
        from openai import OpenAI
        self.client, self.model, self.chat = OpenAI(base_url=base_url, api_key=api_key), model, chat

    def generate(self, prompts, n, temperature, max_new_tokens):
        out_all = []
        for p in prompts:
            msgs = p if self.chat else [{"role": "user", "content": p}]
            outs = []
            for _ in range(n):
                r = self.client.chat.completions.create(
                    model=self.model, messages=msgs,
                    temperature=temperature, max_tokens=max_new_tokens)
                outs.append(r.choices[0].message.content or "")
            out_all.append(outs)
        return out_all


# ═══════════════════════════════════════════════════════════════════════════
# 4. Eval loop
# ═══════════════════════════════════════════════════════════════════════════

def _score(samples: list[str], target: Grid) -> dict:
    em = fmt = 0
    for s in samples:
        p = parse_grid(s)
        if p is not None:
            fmt += 1
            em += (p == target)
    n = len(samples)
    return {"em": em / n, "format": fmt / n, "pass": 1.0 if em else 0.0}


def run(model, chat, families, difficulties, n_per_cell, n_samples,
        temperature, max_new_tokens, seed, verbose=True):
    rows, t0 = [], time.time()
    for diff in difficulties:
        for fam in families:
            acc = {"em": 0.0, "format": 0.0, "pass": 0.0}
            for k in range(n_per_cell):
                task = generate_task(fam, diff, f"{seed}:{fam}:{diff}:{k}")
                prompt = prompt_for(task, chat)
                samples = model.generate([prompt], n_samples, temperature, max_new_tokens)[0]
                for m, v in _score(samples, task.query_out).items():
                    acc[m] += v
            cell = {m: acc[m] / n_per_cell for m in acc}
            cell.update({"family": fam, "difficulty": diff})
            rows.append(cell)
            if verbose:
                print(f"MARC-{diff}  {fam:14s}  EM={cell['em']:6.2%}  "
                      f"fmt={cell['format']:6.2%}  pass={cell['pass']:6.2%}")
    if verbose:
        print(f"\nElapsed: {time.time() - t0:.1f}s\n=== per-difficulty mean ===")
        for d in difficulties:
            sel = [r for r in rows if r["difficulty"] == d]
            print(f"  MARC-{d}:  "
                  f"EM={sum(r['em'] for r in sel) / len(sel):6.2%}  "
                  f"fmt={sum(r['format'] for r in sel) / len(sel):6.2%}")
    return rows


# ═══════════════════════════════════════════════════════════════════════════
# 5. Model bucket — run N models, print a side-by-side summary
# ═══════════════════════════════════════════════════════════════════════════

def build_models(args):
    """Return list of (label, factory) pairs. Factories build models lazily
    so we never hold more than one model in memory at a time."""
    specs = []

    for name in args.hf:
        specs.append((f"hf:{name}",
                      lambda n=name: HFModel(
                          n, chat=args.chat, device=args.device,
                          trust_remote_code=args.trust_remote_code)))

    if args.openai:
        if not args.base_url:
            sys.exit("--base-url is required with --openai")
        if len(args.base_url) == 1:
            urls = args.base_url * len(args.openai)
        elif len(args.base_url) == len(args.openai):
            urls = args.base_url
        else:
            sys.exit("--base-url: give one URL (applied to all) "
                     "or one URL per --openai")
        for name, url in zip(args.openai, urls):
            specs.append((f"openai:{name}",
                          lambda n=name, u=url: OpenAIModel(
                              n, u, chat=args.chat)))

    # de-duplicate labels (add _2, _3, ... if the same string appears twice)
    seen, unique = {}, []
    for label, factory in specs:
        if label in seen:
            seen[label] += 1
            label = f"{label} ({seen[label]})"
        else:
            seen[label] = 1
        unique.append((label, factory))
    return unique


def run_multi(specs, args):
    """Evaluate every model in `specs`, one at a time, then print a summary."""
    all_results = {}
    for i, (label, factory) in enumerate(specs, 1):
        bar = "=" * 78
        print(f"\n{bar}\n=== [{i}/{len(specs)}] {label}\n{bar}")
        model = factory()
        try:
            all_results[label] = run(
                model, args.chat, args.families, args.difficulties,
                args.n_per_cell, args.n_samples, args.temperature,
                args.max_new_tokens, args.seed)
        finally:
            del model
            gc.collect()
            try:
                import torch
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
            except ImportError:
                pass
    if len(all_results) > 1:
        print_summary(all_results, args)
    return all_results


def print_summary(all_results, args):
    """Side-by-side exact-match table across models."""
    models = list(all_results)
    keys = [(d, f) for d in args.difficulties for f in args.families]

    fam_w = max(15, max((len(f) for f in args.families), default=8)) + 1
    col_w = max(18, max(len(m) for m in models) + 2)

    head = f"{'MARC':6s}{'family':<{fam_w}s}" + "".join(f"{m:>{col_w}s}" for m in models)
    sep = "-" * len(head)

    print(f"\n{'=' * len(head)}")
    print("=== SUMMARY: exact match (%)")
    print(f"{'=' * len(head)}")
    print(head)
    print(sep)

    for d, f in keys:
        row = f"MARC-{d:<2d} {f:<{fam_w}s}"
        for m in models:
            em = next(r["em"] for r in all_results[m]
                      if r["difficulty"] == d and r["family"] == f)
            row += f"{em:>{col_w}.2%}"
        print(row)

    print(sep)
    row = f"{'MEAN':6s}{'':<{fam_w}s}"
    for m in models:
        em = sum(r["em"] for r in all_results[m]) / len(all_results[m])
        row += f"{em:>{col_w}.2%}"
    print(row)

    # per-difficulty breakdown
    print(f"\n=== SUMMARY: EM by difficulty ===")
    print(head)
    print(sep)
    for d in args.difficulties:
        row = f"MARC-{d:<2d} {'—':<{fam_w}s}"
        for m in models:
            sel = [r for r in all_results[m] if r["difficulty"] == d]
            em = sum(r["em"] for r in sel) / len(sel)
            row += f"{em:>{col_w}.2%}"
        print(row)


# ═══════════════════════════════════════════════════════════════════════════
# 6. Self-test — asserts the leak invariants (I1), (I2), (I3)
# ═══════════════════════════════════════════════════════════════════════════

def _self_test(trials: int = 300) -> None:
    for fam in FAMILIES:
        for diff in (1, 2, 3):
            for k in range(trials):
                t = generate_task(fam, diff, f"selftest:{fam}:{diff}:{k}")
                # (I1) no query input identical to any demo input
                demo_ins = [d[0] for d in t.demos]
                assert t.query_in not in demo_ins, (
                    f"{fam} diff={diff} k={k}: query_in collides with a demo input")
                # (I2) no identity demo
                for din, dout in t.demos:
                    assert din != dout, (
                        f"{fam} diff={diff} k={k}: identity demo")
                # sanity: every demo and the query have well-formed grids
                assert all(isinstance(v, int) for row in t.query_in for v in row)
                assert all(isinstance(v, int) for row in t.query_out for v in row)
    print(f"self-test passed: {trials} seeds × {len(FAMILIES)} families × 3 difficulties")


# ═══════════════════════════════════════════════════════════════════════════
# 7. CLI
# ═══════════════════════════════════════════════════════════════════════════

def main():
    ap = argparse.ArgumentParser(prog="marc", description="Micro-ARC SLM benchmark")
    ap.add_argument("--hf", action="append", default=[], metavar="NAME",
                    help="HF model name (repeatable for multi-model bucket)")
    ap.add_argument("--openai", action="append", default=[], metavar="NAME",
                    help="model name on an OpenAI-compatible endpoint (repeatable)")
    ap.add_argument("--selftest", action="store_true",
                    help="run the leak-invariant self-test and exit")
    ap.add_argument("--base-url", action="append", default=[], metavar="URL",
                    help="for --openai; give once (broadcast) or once per --openai")
    ap.add_argument("--chat", action="store_true", help="use chat template (instruct models)")
    ap.add_argument("--trust-remote-code", action="store_true",
                    help="pass trust_remote_code=True to HF tokenizer/model loaders")
    ap.add_argument("--device", default="cuda",
                    help="device for HF models (default: cuda)")
    ap.add_argument("--families", nargs="+", choices=list(FAMILIES), default=list(FAMILIES))
    ap.add_argument("--difficulties", nargs="+", type=int, choices=[1, 2, 3], default=[1, 2, 3])
    ap.add_argument("--n-per-cell", type=int, default=20)
    ap.add_argument("--n-samples", type=int, default=1)
    ap.add_argument("--temperature", type=float, default=0.0)
    ap.add_argument("--max-new-tokens", type=int, default=200)
    ap.add_argument("--seed", default="0")
    args = ap.parse_args()

    if args.selftest:
        _self_test()
        return

    if not args.hf and not args.openai:
        sys.exit("pick at least one of --hf, --openai, --selftest")

    specs = build_models(args)
    run_multi(specs, args)


if __name__ == "__main__":
    main()