import os,sys,json
os.environ["NNLAB_THREADS"]="2"; os.environ["OMP_NUM_THREADS"]="2"
import numpy as np, pandas as pd
sys.path.insert(0,"./tools/nnlab")
import lab
S="./reports/nnlab/scores/"
steps=[("lgbm_upw","rule"),("zs_xgb_dann","rule"),("zs_xgb_dann","rule_smooth"),("ens_e2","rule"),("ens_e4","rule"),("ens_e4","rule_smooth"),("ens_e4__lp","rule"),("ens_e4__lp","rule_smooth")]
tab={}
for m,c in steps:
    s=json.load(open(S+m+".json"))
    tab[f"{m}:{c}"]=[r[c] for r in sorted(s["draws"],key=lambda r:(r["n"],r["rep"])) if r["n"]==5]
df=pd.DataFrame(tab).T; df.columns=[f"r{i}" for i in range(5)]; df["mean"]=df.mean(1); df["d_prev"]=df["mean"].diff()
pd.set_option("display.width",200); print(df.round(4))
# own anchor recomputation at n=5
ams=pd.read_parquet(lab.PROC/"amsterdam_features.parquet",columns=["age_class"]); ya=ams.age_class.to_numpy().astype(int)
ym=pd.read_parquet(lab.PROC/"madrid_features.parquet",columns=["age_class"]).age_class.to_numpy().astype(int)
prior=np.array([(ym==c).mean() for c in (1,2,3,4)])
def shots(n,rep):
    rng=np.random.default_rng(42+1000*n+rep)
    return np.concatenate([rng.choice(np.flatnonzero(ya==c),n,replace=False) for c in (1,2,3,4)])
from sklearn.metrics import f1_score
for m in ("zs_lgbm","lgbm_upw","zs_xgb_dann"):
    v=[]
    for r in range(5):
        P=lab.load_pred(m,5,r); s=shots(5,r); keep=np.setdiff1d(np.arange(len(ya)),s)
        v.append(f1_score(ya[keep],lab.W.decide(P[keep],prior),labels=[1,2,3,4],average="macro"))
    print(m,"own rule n5:",np.round(v,4),"mean",round(float(np.mean(v)),4))
# member agreement / zero-shot weight at n=5
w0=10/15; print("w0 at n=5 =",round(w0,3),"; each shot member weight =",round((1-w0)/4,3))