"""How many classes can PixelGen absorb before something gives?

Synthesises N extra classes (random valid 16x8 templates + 3 unique nonsense
words each), rebuilds the model at each N, and measures:

  params / buffers / bytes       - the capacity cost
  build time                     - cost of "compiling" the weights
  generate latency (1 prompt)    - inference cost
  isolation                      - the real 21 classes still bit-identical?
  dense-layer share              - how much of the model is the word->class matrix

Nothing is trained. The synthetic classes exist only in this process.
"""

import random
import sys
import time

import numpy as np
import torch

import pixelgen.model as M
import pixelgen.templates as T
import pixelgen.tokenizer as K

REAL_TEMPLATES = dict(T.TEMPLATES)
REAL_SYN = dict(K.CLASS_SYNONYMS)
REAL_CLASSES = list(REAL_TEMPLATES)


def random_template(rng):
    rows = ["." * 8, "." * 8]
    for _ in range(12):
        row = ".." + "".join(rng.choice("##.?") for _ in range(6))
        rows.append(row)
    rows += ["." * 8, "." * 8]
    on = sum(r.count("#") for r in rows); free = sum(r.count("?") for r in rows)
    if on < 8 or on + free > 60:
        return random_template(rng)
    return rows


def build(n_extra, seed=0):
    rng = random.Random(seed)
    tmpl = dict(REAL_TEMPLATES); syn = dict(REAL_SYN)
    for i in range(n_extra):
        name = f"syn{i:05d}"
        tmpl[name] = random_template(rng)
        syn[name] = [name, f"{name}a", f"{name}b"]
    T.TEMPLATES = tmpl; T.CLASSES = list(tmpl)
    K.CLASS_SYNONYMS = syn; K.VOCAB = K.build_vocab()
    M.TEMPLATES, M.CLASSES, M.CLASS_SYNONYMS, M.VOCAB = tmpl, T.CLASSES, syn, K.VOCAB
    t0 = time.perf_counter(); m = M.PixelGen(); dt = time.perf_counter() - t0
    return m, dt


# reference outputs from the real model
ref_model, _ = build(0)
probe = REAL_CLASSES + ["red potion", "blue invader", "dragon", "", "golden key", "bone"]
ref_out = {p: ref_model.generate([p] * 2, [0, 1], scale=1) for p in probe}

Ns = [int(a) for a in sys.argv[1:]] or [0, 10, 30, 100, 300, 1000, 3000]
print(f"{'classes':>8} {'vocab':>6} {'params':>11} {'buffers':>9} {'MB':>7} {'dense%':>6} {'build s':>8} {'gen ms':>7} {'isolated':>8}")
for n in Ns:
    m, dt = build(n)
    P = sum(p.numel() for p in m.parameters()); B = sum(b.numel() for b in m.buffers())
    dense = m.word_to_class.weight.numel()
    t0 = time.perf_counter()
    for _ in range(5): m.generate(["red potion"], [0], scale=1)
    gen = (time.perf_counter() - t0) / 5 * 1000
    ok = all(np.array_equal(m.generate([p] * 2, [0, 1], scale=1), ref_out[p]) for p in probe)
    # also check a synthetic class actually resolves
    if n:
        nm = f"syn{n-1:05d}"
        got = m.generate([f"{nm}b"], [0], scale=1)[0].sum(-1) > 0
        want = m.generate([nm], [0], scale=1)[0].sum(-1) > 0
        ok = ok and np.array_equal(got, want)
    print(f"{len(m.classes):>8} {len(m.vocab):>6} {P:>11,} {B:>9,} {(P+B)*4/1e6:>7.2f} {100*dense/(P+B):>5.1f}% {dt:>8.2f} {gen:>7.2f} {str(ok):>8}")

print("\nparam model: word_to_class = V*C where V≈3C  ->  O(C^2); banks = 256*C  -> O(C)")
for C in (21, 100, 300, 1000, 3000, 10000):
    V = 59 - 0 + 3 * (C - 21) if C >= 21 else 59
    print(f"  C={C:>6}: dense {V*C:>12,}  banks {256*C:>10,}  dense/banks = {V*C/(256*C):.1f}x")