import sys

sys.path.insert(0, "/tmp/opencode/tl")
import numpy as np
import torch

import int4hf as H
import mirror3 as M
from export_full import load as load_weights

M.ATT_C, M.PKSH, M.SILU_SH, M.SILU_RND = 66572, 0, 11, False
M.build_tables()

float_m = H.AutoModelForCausalLM.from_pretrained("/tmp/opencode/tl", dtype=torch.float32).eval()
int4_m = H.int4_model()

N = 40


def greedy(model, n):
    ids = torch.tensor([[1]])
    out = []
    with torch.no_grad():
        for _ in range(n):
            lg = model(ids).logits[0, -1]
            t = int(lg.argmax())
            out.append(t)
            ids = torch.cat([ids, torch.tensor([[t]])], 1)
    return out


def mirror_greedy(n):
    st = M.State()
    tok = 1
    out = []
    for pos in range(n):
        lg = M.step(tok, pos, st)
        tok = int(lg.argmax())
        out.append(tok)
    return out


f = greedy(float_m, N)
i4 = greedy(int4_m, N)
mi = mirror_greedy(N)
print("float puro        :", f)
print("int4 pesos float  :", i4)
print("pipeline entero   :", mi)
agree_i4 = next((k for k in range(N) if f[k] != i4[k]), N)
agree_mi = next((k for k in range(N) if f[k] != mi[k]), N)
same = sum(1 for k in range(N) if i4[k] == mi[k])
print(f"float vs int4-float: coinciden {agree_i4}/{N} | float vs pipeline: {agree_mi}/{N}")
print(f"pipeline vs int4-float (fidelidad del kernel): {same}/{N}")