"""
modelCreate for the round-2 pick (FINAL-2): two models and one fixed rule.

    1. zero-shot anchor   teammate's XGBoost + DANN, trained on Madrid labels and
                          UNLABELLED Amsterdam features (cached probabilities,
                          data/processed/zeroshot_xgb_dann_ams.npz)
    2. lgbm_feda          LightGBM on Madrid + the labelled shots (weight 30, classes
                          balanced, each city robust-z scaled with its own median/IQR)
                          with Daume's feature augmentation: the top 120 features get
                          a Madrid-only copy and an Amsterdam-only copy
    3. ensemble           log p = w0 * log p_zero_shot + (1 - w0) * log p_feda,  w0 = 10 / (n + 10)
    4. decision           neighbour smoothing (r = 6 px, alpha 0.5, shots excluded), then
                          tau-0.5 logit adjustment + macro-F1 plug-in rule
                          (spatial=False skips the smoothing: no coordinates needed)

    from model_create2 import model_create
    m = model_create(shot_idx, shot_labels)        # rows of the Amsterdam feature table, labels 1..4
    labels = m.predict()                           # labels 1..4 for m.keep (every non-shot row)

Needs at fit time: the Madrid table, ALL unlabelled Amsterdam features (per-city
scaling, the DANN inside the anchor), pixel coordinates for the smoothing. The
decision is batch-level: it labels the whole city at once.

CLI (scores one protocol draw; the only place non-shot labels are read):
    python tools/nnlab/model_create2.py --n 25 --rep 0 [--no-spatial]
"""
from __future__ import annotations
import argparse, sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent))
import lab
import numpy as np
import pandas as pd
W = lab.W
K = 10.0
_S = {}


def _state():
    if not _S:
        import lgbm_var as V
        d = lab.load(official60=False)
        feats = W.feature_names()
        Xm = pd.read_parquet(W.DATA / "madrid_features.parquet", columns=feats)[feats].to_numpy(float)
        Xa = pd.read_parquet(W.DATA / "amsterdam_features.parquet", columns=feats)[feats].to_numpy(float)
        z = np.load(lab.PROC / "zeroshot_xgb_dann_ams.npz")
        assert (z["pa"] == d["pa"]).all()
        _S.update(V=V, d=d, Zm=W.robustz(Xm, W.robustz_fit(Xm)), Za=W.robustz(Xa, W.robustz_fit(Xa)),
                  ctx={"top": V.top_gain_idx()},
                  P_zs=((z["P_xgb"] + z["P_dann_seeds"].mean(0)) / 2).astype(np.float32).astype(np.float64),
                  params={**W.LGBM_PARAMS, **W.RETRAIN_OVERRIDES, "n_jobs": lab.THREADS})
    return _S


class Model:
    def __init__(self, proba, keep, spatial, members):
        self.proba, self.keep, self.spatial, self.members = proba, keep, spatial, members

    def predict(self):
        S = _state()
        P = self.proba[self.keep]
        if self.spatial:
            P = lab.smooth(P, S["d"]["pa"], keep=self.keep)
        return W.decide(P, S["d"]["prior"])


def model_create(shot_idx, shot_labels, n=None, spatial=True):
    S = _state()
    shot = np.asarray(shot_idx)
    ys = np.asarray(shot_labels)
    n = len(shot) / 4 if n is None else n
    P_feda = S["V"].fit_predict("lgbm_feda", S["Zm"], S["d"]["ym"], S["Za"], shot, ys, S["ctx"], S["params"])
    P_feda = P_feda.astype(np.float32).astype(np.float64)              # as stored by the lab
    w0 = K / (n + K)
    L = w0 * np.log(np.clip(S["P_zs"], 1e-6, 1)) + (1 - w0) * np.log(np.clip(P_feda, 1e-6, 1))
    keep = np.setdiff1d(np.arange(len(S["Za"])), shot)
    return Model(lab.softmax(L), keep, spatial, {"zs_xgb_dann": S["P_zs"], "lgbm_feda": P_feda})


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--n", type=int, default=25)
    ap.add_argument("--rep", type=int, default=0)
    ap.add_argument("--no-spatial", action="store_true")
    a = ap.parse_args()
    S = _state()
    d = S["d"]
    shot = lab.shots_for(d, a.n, a.rep)
    m = model_create(shot, lab.shot_labels(d, shot), n=a.n, spatial=not a.no_spatial)
    f1 = W.macro_f1(d["_ya"][m.keep], m.predict())                      # scoring only
    print(f"FINAL-2 n={a.n} rep={a.rep} spatial={m.spatial}: macro F1 {f1:.6f}")
    ref = lab.pred_path("ens_e12", a.n, a.rep)
    if ref.exists():
        print(f"max |proba - stored ens_e12| = {np.abs(m.proba - np.load(ref)).max():.2e}")