"""Reference methods written into the nnlab prediction store, so every new model
can be compared and ensembled with them draw by draw.
    zs_lgbm      zero-shot LightGBM (saved model, raw 393 features)
    zs_xgb_dann  teammate's FINAL before smoothing: (XGBoost + mean of 3 DANN seeds) / 2
    zs_dann      teammate's DANN alone (mean of 3 seeds)
    lgbm_upw     our retrained LightGBM (Madrid + shots x30), one fit per draw
"""
import sys, time
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent))
import lab
import numpy as np
W = lab.W

d = W.load_data()                        # float64, exactly as the reported numbers
ld = lab.load(official60=False)
assert (d["y_ams"] == ld["_ya"]).all() and np.allclose(d["coords_ams"], ld["pa"])
if not lab.have_pred("zs_lgbm", 0, 0):
    lab.save_pred("zs_lgbm", 0, 0, W.predict(W.load_zero_shot(), d["X_ams"]))
z = np.load(lab.PROC / "zeroshot_xgb_dann_ams.npz")
assert (z["pa"] == ld["pa"]).all()
Pd = z["P_dann_seeds"].mean(0)
lab.save_pred("zs_dann", 0, 0, Pd)
lab.save_pred("zs_xgb_dann", 0, 0, (z["P_xgb"] + Pd) / 2)
which = sys.argv[1] if len(sys.argv) > 1 else "all"
if which == "all":
    for n, rep in lab.draws():
        if lab.have_pred("lgbm_upw", n, rep):
            continue
        t0 = time.time()
        shot = lab.shots_for(ld, n, rep)
        m, st = W.train_retrained(d["X_mad"], d["y_mad"], d["X_ams"], d["X_ams"][shot], lab.shot_labels(ld, shot))
        lab.save_pred("lgbm_upw", n, rep, W.predict_retrained(m, st, d["X_ams"]))
        print(f"lgbm_upw n={n} rep={rep} {time.time()-t0:.0f}s", flush=True)
print("done")