import sys

sys.path.insert(0, "/tmp/opencode/tl")
import importlib

import numpy as np

import export_full as E

GREEDY_FLOAT = [4612, 2619, 260, 647, 31844, 635, 393, 260, 1496, 2789, 3557, 21075, 31843, 1100,
                4838, 289, 826, 351, 600, 17360, 1086, 31843, 1591, 1124, 31844, 700, 2815, 260,
                1421, 3476, 288, 600, 2103, 31843, 1100, 393, 900, 6055, 291, 2598]


def q_grid(w, grid=13):
    w = w.astype(np.float32)
    best = None
    for c in np.linspace(0.5, 1.0, grid):
        s = np.abs(w).max(-1, keepdims=True) * c / 7.0
        s[s == 0] = 1e-12
        q = np.clip(np.round(w / s), -8, 7)
        e = ((q * s - w) ** 2).sum(-1, keepdims=True)
        if best is None:
            best = (q, s, e)
        else:
            m = e < best[2]
            best = (np.where(m, q, best[0]), np.where(m, s, best[1]), np.minimum(e, best[2]))
    return best[0], best[1][..., 0]


def q_diffuse(w, grid=13):
    q, s = q_grid(w)
    q = q.copy().astype(np.float32)
    w = w.astype(np.float32)
    for i in range(w.shape[0]):
        order = np.argsort(-np.abs(w[i]))
        buf = w[i].copy()
        for j in order:
            v = np.clip(np.round(buf[j] / s[i]), -8, 7)
            q[i, j] = v
            err = buf[j] - v * s[i]
            if j + 1 < len(buf):
                buf = buf + 0.0
        q[i] = np.clip(np.round(w[i] / s[i]), -8, 7)
    return q.astype(np.int8), s


def q_gptq(w, grid=13, damp=0.01):
    """column-wise error feedback with an identity Hessian approximation (no calibration needed)"""
    w = w.astype(np.float64)
    q, s = q_grid(w)
    sr = s.astype(np.float64)
    H = np.eye(w.shape[1])
    for i in range(w.shape[0]):
        W = w[i].copy()
        for j in range(w.shape[1]):
            qq = np.clip(np.round(W[j] / sr[i]), -8, 7)
            q[i, j] = qq
            err = (W[j] - qq * sr[i]) / H[j, j]
            W = W - err * H[j]
    return q.astype(np.int8), s


def run(tag, quant):
    E.quantize = quant
    for m in list(sys.modules):
        if m in ("mirror3",):
            del sys.modules[m]
    import mirror3 as M

    M.ATT_C, M.PKSH, M.SILU_SH, M.SILU_RND = 66572, 0, 11, False
    M.build_tables()
    st = M.State()
    tok = 1
    out = []
    for pos in range(40):
        lg = M.step(tok, pos, st)
        tok = int(lg.argmax())
        out.append(tok)
    n = 0
    for a, b in zip(out, GREEDY_FLOAT):
        if a != b:
            break
        n += 1
    print(f"{tag:22s} coinciden {n}/40   tokens: {out[:4]}...{out[36:]}", flush=True)


run("grid (actual)", q_grid)
run("difusion de error", q_diffuse)
run("GPTQ identidad", q_gptq)