"""Replace the dense V x C word->class matrix with a hand-written binary code.

  word -> d-bit code of its class, emitted as +-1      Linear(V, d, bias=False)
  class rows = +-1 codes; score_c = d - 2*hamming      Linear(d, C)  (+0.5 bias on default)

Params drop from V*C to V*d + d*C. Price: multi-word prompts sum two codes
and can land on a third class. Both are measured here.
"""

import itertools
import sys

import numpy as np
import torch
import torch.nn as nn

sys.path.insert(0, ".")
from experiments.scale_classes import build, REAL_CLASSES, REAL_SYN  # noqa: E402
import pixelgen.model as M  # noqa: E402


def code_layer(m):
    C, V = len(m.classes), len(m.vocab)
    d = max(1, int(np.ceil(np.log2(C))))
    codes = torch.tensor([[(ci >> b) & 1 for b in range(d)] for ci in range(C)], dtype=torch.float32) * 2 - 1
    enc = nn.Linear(V, d, bias=False); dec = nn.Linear(d, C)
    with torch.no_grad():
        enc.weight.zero_(); dec.bias.zero_()
        for ci, cls in enumerate(m.classes):
            for w in M.CLASS_SYNONYMS[cls]:
                enc.weight[:, m.vocab[w]] = codes[ci]
        dec.weight.copy_(codes)
        dec.bias[m.classes.index(M.DEFAULT_CLASS)] = 0.5
    seq = nn.Sequential(enc, dec)
    for p in seq.parameters(): p.requires_grad_(False)
    return seq, d


def interference(m):
    """Fraction of 2-class-word prompts whose argmax is neither word's class."""
    words = [(c, M.CLASS_SYNONYMS[c][0]) for c in REAL_CLASSES]
    bad = 0; n = 0
    for (ca, wa), (cb, wb) in itertools.combinations(words, 2):
        toks = m.tokenize(f"{wa} {wb}")
        bag = torch.zeros(1, len(m.vocab)); bag[0, toks] += 1  # noqa (tokenize returns ids)
        cls = m.classes[m.word_to_class(bag).argmax(-1).item()]
        bad += cls not in (ca, cb); n += 1
    return bad / n


Ns = [int(a) for a in sys.argv[1:]] or [0, 100, 1000, 3000, 10000]
print(f"{'classes':>8} {'d':>3} {'dense params':>13} {'coded params':>13} {'ratio':>7} {'isolated':>8} {'interf':>7}")
for n in Ns:
    m, _ = build(n)
    ref = {p: m.generate([p] * 2, [0, 1], scale=1) for p in REAL_CLASSES + ["red potion", "", "dragon", "bone"]}
    dense_p = sum(p.numel() for p in m.word_to_class.parameters())
    seq, d = code_layer(m)
    m.word_to_class = seq
    coded_p = sum(p.numel() for p in seq.parameters())
    ok = all(np.array_equal(m.generate([p] * 2, [0, 1], scale=1), ref[p]) for p in ref)
    if n:
        nm = f"syn{n-1:05d}"
        ok = ok and np.array_equal(m.generate([nm + "b"], [0], scale=1)[0], m.generate([nm], [0], scale=1)[0])
    print(f"{len(m.classes):>8} {d:>3} {dense_p:>13,} {coded_p:>13,} {dense_p/coded_p:>6.1f}x {str(ok):>8} {interference(m):>6.0%}")

print("\n(interf = share of 'wordA wordB' prompts over the 21 real classes that resolve to a THIRD class;")
print(" dense layer: 0% by construction, ties go to lower index)")