"""Load a paper checkpoint for evaluation; for arm C, round the ternary weights ONCE so inference runs at float speed.
(train.TLinear re-rounds weights on every forward pass, which is right for training and 2x slower for evaluation.)"""
import numpy as np, torch, torch.nn.functional as F
import train


class FrozenTernary(torch.nn.Module):
    def __init__(self, lin):
        super().__init__()
        w = lin.weight.detach()
        sw = 1.0 / w.abs().mean().clamp_min(1e-5)
        self.register_buffer("wq", ((w * sw).round().clamp(-1, 1) / sw))
        self.bias = lin.bias

    def forward(self, x):
        xf = x.float()
        xn = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + 1e-6)
        sx = 127.0 / xn.abs().amax(-1, keepdim=True).clamp_min(1e-5)
        xq = (xn * sx).round().clamp(-128, 127) / sx
        return F.linear(xq.to(x.dtype), self.wq.to(x.dtype), self.bias)


def load(arm, ckpt, counts, learned_bias=None):
    if learned_bias is None:
        learned_bias = arm == "C"
    m = train.Model(arm, counts, learned_bias=learned_bias).cuda().eval()
    st = torch.load(ckpt, map_location="cuda", weights_only=False)
    m.load_state_dict(st["model"])
    if arm == "C":
        for mod in list(m.body.layers.modules()):
            for name, child in list(mod.named_children()):
                if isinstance(child, train.TLinear):
                    setattr(mod, name, FrozenTernary(child))
    return m, st.get("tag", arm)