"""v0.2 extraction: first+last subtoken states at MULTIPLE layers,
both models, all 55,192 sites. One forward per line per model yields
every layer. Student t5-enc L2..L8; teacher bert L4,6,8,10,12."""
import json
import sys

sys.path.insert(0, r"E:\mirel\geolip-bytelex")
sys.stdout.reconfigure(encoding="utf-8", errors="replace")
import numpy as np
import torch
from transformers import AutoTokenizer, AutoModel, T5EncoderModel
import transformers.utils.logging as hlog

hlog.set_verbosity_error()
D = r"E:\mirel\data\bytelex\proto_frame"
DEV = "cuda"
LB = [4, 6, 8, 10, 12]
LT = [2, 3, 4, 5, 6, 7, 8]
BATCH = 128

dump = np.load(rf"{D}\frame_dump_v2.npz")
sid = dump["sid"]
N = len(sid)
lines = open(r"E:\mirel\data\bytelex\codex_v1.txt",
             "rb").read().decode("ascii").split("\n")
tkB = AutoTokenizer.from_pretrained("bert-base-uncased")
tkA = AutoTokenizer.from_pretrained("google/flan-t5-small")
mB = AutoModel.from_pretrained("bert-base-uncased").to(DEV).eval()
mT = T5EncoderModel.from_pretrained("google/flan-t5-small").to(DEV).eval()

by_line = {}
for k in range(N):
    by_line.setdefault(int(dump["line"][k]), []).append(k)
line_ids = sorted(by_line)

SB = np.zeros((N, len(LB), 2, 768), dtype=np.float16)
ST = np.zeros((N, len(LT), 2, 512), dtype=np.float32)
with torch.no_grad():
    for bs in range(0, len(line_ids), BATCH):
        chunk = line_ids[bs:bs + BATCH]
        texts = [lines[li] for li in chunk]
        eb = tkB(texts, return_offsets_mapping=True, padding=True,
                 return_tensors="pt")
        et = tkA(texts, return_offsets_mapping=True, padding=True,
                 return_tensors="pt")
        hb = mB(input_ids=eb["input_ids"].to(DEV),
                attention_mask=eb["attention_mask"].to(DEV),
                output_hidden_states=True).hidden_states
        ht = mT(input_ids=et["input_ids"].to(DEV),
                attention_mask=et["attention_mask"].to(DEV),
                output_hidden_states=True).hidden_states
        hb = [hb[l].cpu() for l in LB]
        ht = [ht[l].cpu() for l in LT]
        for r, li in enumerate(chunk):
            offB = eb["offset_mapping"][r].tolist()
            offT = et["offset_mapping"][r].tolist()
            for k in by_line[li]:
                lo, hi = int(dump["lo"][k]), int(dump["hi"][k])
                ixB = [i for i, (s, t) in enumerate(offB)
                       if t > s and s < hi and t > lo]
                ixT = [i for i, (s, t) in enumerate(offT)
                       if t > s and s < hi and t > lo]
                if not ixB or not ixT:
                    continue
                for j in range(len(LB)):
                    SB[k, j, 0] = hb[j][r, ixB[0]].numpy()
                    SB[k, j, 1] = hb[j][r, ixB[-1]].numpy()
                for j in range(len(LT)):
                    ST[k, j, 0] = ht[j][r, ixT[0]].numpy()
                    ST[k, j, 1] = ht[j][r, ixT[-1]].numpy()
        if (bs // BATCH) % 10 == 0:
            print(f"[v02x] {bs}/{len(line_ids)}", flush=True)

np.savez(rf"{D}\sweep_dump_v02.npz", SB=SB, ST=ST,
         LB=np.array(LB), LT=np.array(LT))
print(f"[v02x] COMPLETE, vram peak "
      f"{torch.cuda.max_memory_allocated()/2**30:.1f}GB", flush=True)