"""C3 causal-attribution control: pretrained t5-small (span-corruption
objective, no flan tuning) vs flan-t5-small on the same T-2 sites.
Same seed, same pools, first 150 of each domain's 600 sampled sites.
"""
import sys
from collections import Counter, defaultdict

sys.stdout.reconfigure(encoding="utf-8", errors="replace")
import numpy as np
import torch
from transformers import AutoTokenizer, T5ForConditionalGeneration
import transformers.utils.logging as hlog

hlog.set_verbosity_error()
CODEX = r"E:\mirel\data\bytelex\codex_v1.txt"
N_PER_DOMAIN = 600
N_SUB = 150
rng = np.random.default_rng(20260820)

tkA = AutoTokenizer.from_pretrained("google/flan-t5-small")
tkB = AutoTokenizer.from_pretrained("google/bert_uncased_L-4_H-256_A-4")
tkC = AutoTokenizer.from_pretrained("google-t5/t5-small")
assert tkC.get_vocab() == tkA.get_vocab(), "t5-small vocab != flan vocab"
mA = T5ForConditionalGeneration.from_pretrained("google/flan-t5-small")
mC = T5ForConditionalGeneration.from_pretrained("google-t5/t5-small")
mA.eval()
mC.eval()

raw = open(CODEX, "rb").read().decode("ascii")
lines, story, arith, n = raw.split("\n"), [], [], 0
for ln in lines:
    if not ln:
        n += 1
        continue
    (story if n < 200_000 else arith).append(ln)
    n += len(ln) + 1


def enc(tk, line):
    e = tk(line, add_special_tokens=False, return_offsets_mapping=True)
    return [(i, s, t) for i, (s, t)
            in zip(e["input_ids"], e["offset_mapping"]) if t > s]


def sites(lines_):
    out = []
    for ln in lines_:
        A, B = enc(tkA, ln), enc(tkB, ln)
        bA = {s for _, s, _ in A} | {t for _, _, t in A}
        bB = {s for _, s, _ in B} | {t for _, _, t in B}
        shared = sorted(bA & bB)
        for lo, hi in zip(shared, shared[1:]):
            ca = [i for i, s, t in A if s >= lo and t <= hi]
            cb = [i for i, s, t in B if s >= lo and t <= hi]
            span = ln[lo:hi].lstrip(" ")
            if ca and cb and span:
                shape = "1x1" if (len(ca) == 1 and len(cb) == 1) \
                    else "nm"
                out.append((ln, lo, hi, shape))
    return out


sys.path.insert(0, r"E:\mirel\geolip-bytelex")
from geolip.bytelex.extract import token_byte_table

rows = token_byte_table(tkA)
cbA = np.full(len(tkA), -1, dtype=np.int64)
specA = np.zeros(len(tkA), dtype=bool)
for r in rows:
    if r["is_special"]:
        specA[r["id"]] = True
    if r["is_special"] or not r["hex"]:
        continue
    e = bytes.fromhex(r["hex"]).lstrip(b" ")
    if e:
        cbA[r["id"]] = e[0]

EXTRA0 = tkA.convert_tokens_to_ids("<extra_id_0>")
EX1 = tkA.convert_tokens_to_ids("<extra_id_1>")
EOSA = tkA.eos_token_id
PAD = tkA.pad_token_id


def push(probs, cb):
    overflow = float(probs[len(cb):].sum())
    probs = probs[:len(cb)]
    keep = cb >= 0
    dropped = overflow + float(probs[~keep].sum())
    marg = np.zeros(256)
    np.add.at(marg, cb[keep], probs[keep])
    s = marg.sum()
    return (marg / s if s > 0 else marg), dropped


@torch.no_grad()
def probe(model, ln, lo, hi):
    src = ln[:lo] + "<extra_id_0>" + ln[hi:]
    ids = tkA(src, return_tensors="pt")
    dec = torch.tensor([[PAD, EXTRA0]])
    lg = model(input_ids=ids["input_ids"],
               attention_mask=ids["attention_mask"],
               decoder_input_ids=dec).logits[0, -1]
    p = torch.softmax(lg.float(), -1).numpy()
    (marg, drop) = push(p, cbA)
    return marg, drop, p


res = defaultdict(lambda: defaultdict(float))
cnt = Counter()
top1 = defaultdict(Counter)
for dom, pool in (("story", sites(story)), ("arith", sites(arith))):
    idx = rng.choice(len(pool), size=min(N_PER_DOMAIN, len(pool)),
                     replace=False)
    for k in idx[:N_SUB]:
        ln, lo, hi, shape = pool[int(k)]
        gold = ord(ln[lo:hi].lstrip(" ")[0])
        key = f"{dom}:{shape}"
        cnt[key] += 1
        for tag, model in (("flan", mA), ("t5", mC)):
            marg, drop, p = probe(model, ln, lo, hi)
            res[key][f"{tag}_gold"] += int(marg.argmax() == gold)
            res[key][f"{tag}_drop"] += drop
            res[key][f"{tag}_eos"] += float(p[EOSA])
            res[key][f"{tag}_ex1"] += float(p[EX1])
            top1[f"{key}:{tag}"][
                tkA.convert_ids_to_tokens(int(p.argmax()))
                if int(p.argmax()) < len(tkA)
                else f"PHANTOM_{int(p.argmax())}"] += 1

for key in sorted(cnt):
    n_ = cnt[key]
    r = res[key]
    print(f"[CTRL {key}] n={n_}")
    for tag in ("flan", "t5"):
        print(f"  {tag:4s}: gold={r[f'{tag}_gold']/n_:.4f} "
              f"drop={r[f'{tag}_drop']/n_:.4f} "
              f"eos={r[f'{tag}_eos']/n_:.4f} "
              f"ex1={r[f'{tag}_ex1']/n_:.4f}")
        print(f"        top1: {top1[f'{key}:{tag}'].most_common(8)}")
print("[CTRL] DONE", flush=True)