#!/usr/bin/env python3
"""Đánh giá lại cổng tầng 1 cho các TỔ HỢP CON của 28 họ, từ phiếu đã lưu.

Không khớp lại mô hình nào: exp3_week_s1.py đã ghi số phiếu của từng họ ra .npz,
nên mọi tổ hợp con đều tính lại được trong vài giây.  Đây là cách trả lời câu hỏi
"tăng số thành viên có làm cổng tốt lên không" bằng số chứ không bằng lập luận.
"""

from __future__ import annotations

import json
import os
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import ehd_common as C
import ehd_metrics as M
from exp3_week_s1 import (OUT_JSON, OUT_NPZ, NQ, P1_GRID, MAX_LEAK,
                          gate_scores, pick_p1)


def main():
    z = np.load(OUT_NPZ, allow_pickle=True)
    fam = [str(x) for x in z["families"]]
    with open(OUT_JSON) as fh:
        res = json.load(fh)
    keys = sorted({k.rsplit("|", 1)[0] for k in z.files if k.endswith("|val")})
    rank = [n for n, _ in sorted(res["family"].items(),
                                 key=lambda kv: -kv[1]["AUC_val"])]

    def idx(names):
        return np.array([fam.index(n) for n in names])

    subs = {
        "7 họ gốc": np.arange(7),
        "28 họ — tất cả": np.arange(len(fam)),
        "5 họ tuyến tính tốt nhất": idx(["BayesRidge", "ElasticNet", "Huber",
                                         "Ridge", "PLS8"]),
        "10 họ AUC cao nhất": idx(rank[:10]),
        "chỉ 7 ANN torch": np.array([i for i, n in enumerate(fam)
                                     if n.startswith("ANN_")]),
        "chỉ 3 CNN 1-D": np.array([i for i, n in enumerate(fam)
                                   if n.startswith("CNN_")]),
    }

    def p_of(v, sel):
        sub = v[sel]
        good = sub >= 0
        return np.where(good.any(axis=0),
                        np.where(good, sub, 0).sum(axis=0)
                        / np.maximum(good.sum(axis=0) * NQ, 1), np.nan)

    out = {}
    for name, sel in subs.items():
        blk = {"n_family": int(sel.size), "n_member": int(sel.size) * NQ,
               "by_week": {}}
        aucs = []
        for w in C.WEEK_NAMES:
            ks = [k for k in keys if k.split("|")[1] == w]
            if not ks:
                continue
            po = np.concatenate([p_of(z[k + "|oof"], sel) for k in ks])
            pv = np.concatenate([p_of(z[k + "|val"], sel) for k in ks])
            oo = np.concatenate([z[k + "|obs_oof"] for k in ks])
            ov = np.concatenate([z[k + "|obs_val"] for k in ks])
            p1 = pick_p1(po, oo)
            g = gate_scores(pv, ov, p1)
            g["AUC"] = float(M.auc_score(pv, ov))
            blk["by_week"][w] = g
            aucs.append(g["AUC"])
        blk["AUC_mean"] = float(np.mean(aucs)) if aucs else float("nan")
        out[name] = blk
        print("%-26s AUC TB %.3f  |  %s" % (name, blk["AUC_mean"],
              "  ".join("%s %.3f" % (w, blk["by_week"][w]["AUC"])
                        for w in C.WEEK_NAMES if w in blk["by_week"])), flush=True)

    res["subsets"] = out
    with open(OUT_JSON, "w") as fh:
        json.dump(res, fh, indent=1, default=float)
    print("\n[done] -> %s" % OUT_JSON, flush=True)


if __name__ == "__main__":
    main()