import os
import sys
import torch




def pack_ternary(w):
    q = torch.empty_like(w, dtype=torch.uint8)
    q[w < 0] = 0
    q[w == 0] = 1
    q[w > 0] = 2

    flat = q.flatten()
    pad = (-len(flat)) % 4
    if pad:
        flat = torch.cat([flat, torch.zeros(pad, dtype=torch.uint8, device=flat.device)])

    flat = flat.view(-1, 4)

    packed = (
        flat[:, 0]
        | (flat[:, 1] << 2)
        | (flat[:, 2] << 4)
        | (flat[:, 3] << 6)
    )

    return packed.cpu(), w.shape


def save_model(model, path="trigram-morph.pt"):
    ternary_weights = {}
    for name, param in model.named_parameters():
        if "weight" in name and param.ndim >= 2 and "embed" not in name:
            T = StickyZoneSTE.apply(param.data, THRESHOLD)
            packed, shape = pack_ternary(T)
            ternary_weights[name] = {"packed": packed, "shape": shape}

    torch.save({
        "model_state_dict": model.state_dict(),
        "config": {
            "vocab": VOCAB,
            "embedding_dim": EMBEDDING_DIM,
            "trigram_dim": HIDDEN_DIM,
            "ffn_hidden": FFN_HIDDEN,
            "ctx": CTX,
            "threshold": THRESHOLD,
        },
        "ternary_packed": ternary_weights,
        "format": "factorized_scaled_ternary",
        "bpw": 1.58,
    }, path)
    total = sum(p.numel() for p in model.parameters())
    print(f"Saved {total:,} params to {path}")


def load_model(path="trigram-morph.pt", device="cpu"):
    checkpoint = torch.load(path, map_location=device, weights_only=False)
    model = ARBModel()
    model.load_state_dict(checkpoint["model_state_dict"], strict=False)
    model.to(device)
    model.eval()
    return model


if __name__ == "__main__":
    from ..trigram import ARBModel
    model = ARBModel()
    total = sum(p.numel() for p in model.parameters())
    ternary = sum(
        p.numel() for n, p in model.named_parameters()
        if "weight" in n and p.ndim >= 2 and "embed" not in n
    )
    fp32 = sum(
        p.numel() for n, p in model.named_parameters()
        if not ("weight" in n and p.ndim >= 2 and "embed" not in n)
    )
    print(f"Total params: {total:,}")
    print(f"Ternary params (1.58 BPW): {ternary:,}")
    print(f"FP32 params: {fp32:,}")
    save_model(model)