"""Fixed-rule ensembles of methods already in the prediction store.

Rule, fixed before any ensemble was scored (from the literature shortlist, item 6):
    log p = w0 * mean(log p of zero-shot members) + (1 - w0) * mean(log p of shot-using members)
    w0 = K / (n + K),  K = 10      (the zero-shot anchor fades as the labelled sample grows)
No weights are fitted, so no Amsterdam label is used beyond what the members used.

    python tools/nnlab/ensemble.py NAME --zs zs_xgb_dann --shot lgbm_upw,tab_dann_rz_joint [--k 10]
writes method ens_NAME.
"""
import argparse, sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent))
import lab
import numpy as np

ap = argparse.ArgumentParser()
ap.add_argument("name")
ap.add_argument("--zs", default="")
ap.add_argument("--shot", default="")
ap.add_argument("--k", type=float, default=10.0)
a = ap.parse_args()
zs = [m for m in a.zs.split(",") if m]
sh = [m for m in a.shot.split(",") if m]
assert zs or sh
name = "ens_" + a.name
logm = lambda ms, n, r: np.mean([np.log(np.clip(lab.load_pred(m, n, r), 1e-6, 1)) for m in ms], 0)
done = 0
for n, rep in lab.draws():
    if any(not lab.have_pred(m, n, rep) for m in sh) or any(not (lab.have_pred(m, n, rep) or lab.have_pred(m, 0, 0)) for m in zs):
        continue
    w0 = (a.k / (n + a.k) if sh else 1.0) if zs else 0.0
    L = (w0 * logm(zs, n, rep) if zs else 0) + ((1 - w0) * logm(sh, n, rep) if sh else 0)
    lab.save_pred(name, n, rep, lab.softmax(L), meta=dict(zs=zs, shot=sh, k=a.k, rule="w0=k/(n+k), mean of log-probs"))
    done += 1
print(f"{name}: {done} draws written")