import argparse
import os
import sys
import numpy as np

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


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument(
        "--weights",
        default=os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
        help="repo dir with model_<tier>[_int8].safetensors + config, or a .safetensors path",
    )
    ap.add_argument("--tier", default="small", choices=["small", "big"])
    ap.add_argument(
        "--dataset", default="breast_cancer", choices=["breast_cancer", "wine", "iris"]
    )
    ap.add_argument(
        "--bag", type=int, default=1, help="inference-bagging passes (classification)"
    )
    args = ap.parse_args()
    from sklearn.model_selection import train_test_split
    from sklearn.datasets import load_breast_cancer, load_wine, load_iris

    loader = {
        "breast_cancer": load_breast_cancer,
        "wine": load_wine,
        "iris": load_iris,
    }[args.dataset]
    d = loader()
    X, y = (d.data.astype(np.float32), d.target.astype(np.int64))
    names = list(getattr(d, "target_names", [str(i) for i in range(int(y.max()) + 1)]))
    Xtr, Xte, ytr, yte = train_test_split(
        X, y, test_size=0.3, random_state=0, stratify=y
    )
    model = load_model(args.weights, tier=args.tier)
    ncls = int(y.max() + 1)
    print(
        f"dataset={args.dataset}  support rows={len(Xtr)}  query rows={len(Xte)}  features={X.shape[1]}  classes={ncls}  tier={args.tier}"
    )
    if args.bag > 1:
        probs = predict_bagged(
            model, Xtr, ytr, Xte, "classification", n_classes=ncls, n_bag=args.bag
        )
    else:
        probs = predict(model, Xtr, ytr, Xte, "classification", n_classes=ncls)
    acc = float((probs.argmax(1) == yte).mean())
    print(f"zero-shot accuracy: {acc:.3f}  (in-context, no training)")
    i = 0
    p = probs[i]
    print(
        f"example query row 0 -> predicted '{names[p.argmax()]}' ({p.max() * 100:.1f}% confidence); true '{names[yte[i]]}'"
    )
    print(
        "  full class probabilities: "
        + ", ".join((f"{names[c]}={p[c] * 100:.1f}%" for c in range(ncls)))
    )
    Xr = np.delete(X, 0, axis=1).astype(np.float32)
    yr = X[:, 0].astype(np.float32)
    Xr_tr, Xr_te, yr_tr, yr_te = train_test_split(Xr, yr, test_size=0.3, random_state=0)
    mean, std = predict(model, Xr_tr, yr_tr, Xr_te, "regression")
    print(
        f"regression demo (predict feature 0): example -> mean {mean[0]:.3f} +/- {std[0]:.3f} (true {yr_te[0]:.3f}); this is the honest Gaussian error bar"
    )


if __name__ == "__main__":
    main()