#!/usr/bin/env python3
"""D-58 gate: a codebook student exactly as imported, before any training, exported packed. Its GGUF must be the
llama.cpp source file's expert tensors byte for byte (`t27_gguf.py build` then `cb_gate.py compare` on gen).

    python scripts/cb_gate.py export --config configs/p6/CB3a.yaml --out $MT_DATA/runs/p6-gate-CB3a/packed/tok_0
    python scripts/cb_gate.py compare --ours <gguf> --source <llama.cpp gguf>
"""
from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path

import numpy as np
import torch
import yaml

REPO = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(REPO / "src"))
from lowbit.model import student as st  # noqa: E402
from lowbit.quant.ternary import QuantSpec  # noqa: E402


def export(a) -> None:
    cfg = yaml.safe_load(open(a.config))
    prec = st.Precision.load(REPO / cfg["model"]["precision_map"])
    prec.experts = QuantSpec(method=cfg["quant"]["method"], group_size=cfg["quant"]["group_size"])
    model = st.load_student(cfg["model"]["teacher"], prec, quant=True, device=a.device)
    print(json.dumps(st.export_packed_student(model, a.out))[:300])


def compare(a) -> None:
    sys.path.insert(0, os.environ.get("MT_LLAMA", "/path/to/llama.cpp") + "/gguf-py")
    import gguf
    ours = {t.name: t for t in gguf.GGUFReader(a.ours).tensors}
    src = {t.name: t for t in gguf.GGUFReader(a.source).tensors}
    diff = {"experts": [], "other": []}
    for name, t in src.items():
        same = (name in ours and ours[name].tensor_type == t.tensor_type
                and np.asarray(ours[name].data).tobytes() == np.asarray(t.data).tobytes())
        if not same:
            diff["experts" if "_exps" in name else "other"].append(name)
    print(json.dumps({"tensors": len(src), "missing": sorted(set(src) - set(ours)),
                      "expert_mismatches": diff["experts"], "other_mismatches": diff["other"][:20],
                      "n_other_mismatches": len(diff["other"])}))


def main() -> None:
    p = argparse.ArgumentParser()
    sub = p.add_subparsers(dest="cmd", required=True)
    e = sub.add_parser("export")
    e.add_argument("--config", required=True)
    e.add_argument("--out", required=True)
    e.add_argument("--device", default="cuda")
    c = sub.add_parser("compare")
    c.add_argument("--ours", required=True)
    c.add_argument("--source", required=True)
    a = p.parse_args()
    export(a) if a.cmd == "export" else compare(a)


if __name__ == "__main__":
    main()