import argparse
import os
import sys
import numpy as np

sys.path.insert(0, os.path.dirname(__file__))
from modeling import load_model, predict


def make_cls(seed=1234, ns=48, nq=4, nf=6, k=3):
    rng = np.random.default_rng(seed)
    W = rng.standard_normal((nf, k)).astype(np.float32)
    X = rng.standard_normal((ns + nq, nf)).astype(np.float32) * 1.6 + 0.2
    y = (X @ W).argmax(1).astype(np.int64)
    return (X[:ns], y[:ns], X[ns:], k)


def make_reg(seed=5678, ns=48, nq=4, nf=6):
    rng = np.random.default_rng(seed)
    w = rng.standard_normal(nf).astype(np.float32)
    X = rng.standard_normal((ns + nq, nf)).astype(np.float32) * 2.0 - 0.3
    y = (X @ w + 0.1 * rng.standard_normal(ns + nq)).astype(np.float32)
    return (X[:ns], y[:ns], X[ns:])


REF_CLS_PROBS = [
    [0.67778, 0.29555, 0.02667],
    [0.43088, 0.44604, 0.12308],
    [0.8469, 0.11382, 0.03928],
    [0.34364, 0.60266, 0.0537],
]
REF_REG_MEAN = [-3.22947, -1.90098, 0.24437, -3.92093]
TOL_PROB = 0.06
TOL_REG = 0.15


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--weights", default=None, help="explicit .safetensors path")
    ap.add_argument("--tier", default="small", choices=["small", "big"])
    args = ap.parse_args()
    src = args.weights or "."
    model = load_model(src, tier=args.tier)
    Xc_tr, yc_tr, Xc_te, k = make_cls()
    Xr_tr, yr_tr, Xr_te = make_reg()
    probs = predict(model, Xc_tr, yc_tr, Xc_te, "classification", n_classes=k)
    mean, std = predict(model, Xr_tr, yr_tr, Xr_te, "regression")
    ok = True
    if probs.shape != (4, 3):
        print(f"FAILURE: cls output shape {probs.shape} != (4, 3)")
        sys.exit(1)
    if mean.shape != (4,) or std.shape != (4,):
        print(f"FAILURE: reg shapes mean {mean.shape} std {std.shape} != (4,)")
        sys.exit(1)
    print(f"Shapes OK: cls {probs.shape}, reg mean {mean.shape} std {std.shape}")
    if not np.all(std > 0):
        print(f"FAILURE: regression std not all positive: {std.tolist()}")
        sys.exit(1)
    if args.tier == "small":
        d_cls = float(np.abs(probs - np.array(REF_CLS_PROBS)).max())
        scale = float(np.std(REF_REG_MEAN)) + 1e-06
        d_reg = float(np.abs(mean - np.array(REF_REG_MEAN)).max()) / scale
        print(f"cls max abs prob diff = {d_cls:.4f} (tol {TOL_PROB})")
        print(f"reg max rel mean diff = {d_reg:.4f} (tol {TOL_REG})")
        if d_cls > TOL_PROB:
            print("FAILURE: classification probabilities off reference")
            ok = False
        if d_reg > TOL_REG:
            print("FAILURE: regression means off reference")
            ok = False
    else:
        print(
            "(reference stored for the small tier only; big tier ran a shape + sanity check)"
        )
    print("example cls probs[0]:", np.round(probs[0], 4).tolist())
    print("example reg mean/std[0]:", round(float(mean[0]), 4), round(float(std[0]), 4))
    print("VERIFY:", "PASS" if ok else "FAIL")
    sys.exit(0 if ok else 1)


if __name__ == "__main__":
    main()