"""
quant_units.py -- load the int8-quantized datapath units shipped in
`models_int8/`, and (re)verify any unit variant bit-exact over its full domain.

The int8 files are produced by riscv/export_int8.py: the dominant 256x256 hidden
weight is per-channel symmetric int8, everything else is fp16. Dequantizing
reconstructs a normal BitMLP state_dict, so the rest of the stack (build_all /
NeuralUnits / the rv32 core) is untouched -- this is purely an alternate weight
source. All 13 units stay bit-exact, at ~28% of the fp32 size.
"""
import os
import torch
from x86_units import (BitMLP, u_ADC8, u_SBB8, u_logic, u_SHL1, u_SHR1,
                       u_MASK8, u_NOT8, u_FLAGS8, u_prefix, u_modrm, u_sib)

HERE = os.path.dirname(os.path.abspath(__file__))
INT8_DIR = os.path.join(HERE, "models_int8")

# name -> truth-table generator, for exhaustive re-verification of either variant
DOMAINS = {
    "ADC8": u_ADC8, "SBB8": u_SBB8,
    "AND8": lambda: u_logic(lambda a, b: a & b),
    "OR8":  lambda: u_logic(lambda a, b: a | b),
    "XOR8": lambda: u_logic(lambda a, b: a ^ b),
    "SHL1": u_SHL1, "SHR1": u_SHR1, "MASK8": u_MASK8, "NOT8": u_NOT8,
    "FLAGS8": u_FLAGS8, "PREFIX": u_prefix, "MODRM": u_modrm, "SIB": u_sib,
}


def _dequant(blob):
    """int8 blob -> fp32 BitMLP state_dict."""
    sd = {}
    for k, v in blob.items():
        if k.endswith(".q"):
            sd[k[:-2]] = v.float() * blob[k[:-2] + ".s"].float()
        elif k.endswith(".s"):
            continue
        else:
            sd[k] = v.float()
    return sd


def build_int8(int8_dir=INT8_DIR):
    """Like build_all(), but loads the int8 units (dequantized) -- no training."""
    nets = {}
    for name in DOMAINS:
        blob = torch.load(os.path.join(int8_dir, f"{name}.pt"), weights_only=True)
        sd = _dequant(blob)
        nin = sd["net.0.weight"].shape[1]
        nout = sd["net.4.weight"].shape[0]
        net = BitMLP(nin, nout)
        net.load_state_dict(sd)
        net.eval()
        nets[name] = net
    return nets


def verify(nets):
    """Exhaustively re-check every unit over its FULL domain. Returns
    (n_exact, n_total, list_of_failures)."""
    ok = 0; fails = []
    for name, gen in DOMAINS.items():
        X, Y = gen()
        Xt, Yt = torch.tensor(X), torch.tensor(Y)
        good = True
        with torch.no_grad():
            for i in range(0, Xt.shape[0], 65536):
                if not bool(((nets[name](Xt[i:i+65536]) > 0).float()
                             == Yt[i:i+65536]).all()):
                    good = False; break
        ok += int(good)
        if not good:
            fails.append(name)
    return ok, len(DOMAINS), fails


def dir_size(model_dir):
    """Total bytes of the 13 unit files in a model directory."""
    tot = 0
    for name in DOMAINS:
        p = os.path.join(model_dir, f"{name}.pt")
        if os.path.exists(p):
            tot += os.path.getsize(p)
    return tot