#!/usr/bin/env python3
"""Chấm điểm cổng tầng 1 cho phương án (b) từ phiếu đã lưu, không khớp lại mô hình.

exp3_week_s1b.py ghi số phiếu của từng họ ra .npz trước khi tính chỉ số, nên toàn bộ
bảng so sánh dựng lại được trong vài giây.  Biến thể A (không neo) lấy từ .npz của
lần chạy trước, KÈM quan trắc của chính lần đó: hai lần chạy lọc NaN khác nhau (lần
trước lọc thêm chuỗi lead day cho CNN) nên số hàng lệch vài đơn vị.
"""

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
import exp3_zoo as Z
import exp3_offset as OFF
from exp3_week_s1 import gate_scores, pick_p1

RES = os.path.join(C.OUT, "AI_Results")
NQ = 9


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


def main():
    z = np.load(os.path.join(RES, "exp3_week_s1b_votes.npz"), allow_pickle=True)
    fam_b1 = [str(x) for x in z["fam_b1"]]
    fam_b2 = [str(x) for x in z["fam_b2"]]
    keys = sorted({k.rsplit("|", 1)[0] for k in z.files if k.endswith("|praw_val")})
    pts = sorted({int(k.split("|")[0]) for k in keys})

    za = fa = None
    prev = os.path.join(RES, "exp3_week_s1_votes.npz")
    if os.path.exists(prev):
        za = np.load(prev, allow_pickle=True)
        fl = [str(x) for x in za["families"]]
        fa = np.array([fl.index(n) for n in fam_b1 if n in fl])

    res = {"config": {"points": pts, "weeks": C.WEEK_NAMES, "fset": "F3",
                      "n_quantile": NQ, "n_family_b1": len(fam_b1),
                      "n_family_b2": len(fam_b2)},
           "models_b1": [{"name": n, "desc": d}
                         for n, _, d, _ in Z.zoo_full() if n in fam_b1],
           "models_b2": [{"name": n, "desc": d} for n, d, _ in OFF.offset_zoo()],
           "by_week": {}, "family_b2": {}}

    for w in C.WEEK_NAMES:
        ks = [k for k in keys if k.split("|")[1] == w]
        if not ks:
            continue
        oo = np.concatenate([z[k + "|obs_oof"] for k in ks])
        ov = np.concatenate([z[k + "|obs_val"] for k in ks])
        cands = {
            "p_raw (EC gốc)": (np.concatenate([z[k + "|praw_oof"] for k in ks]),
                               np.concatenate([z[k + "|praw_val"] for k in ks]),
                               oo, ov),
            "B1 neo bằng Sev": (np.concatenate([p_of(z[k + "|b1_oof"]) for k in ks]),
                                np.concatenate([p_of(z[k + "|b1_val"]) for k in ks]),
                                oo, ov),
            "B2 neo bằng xác suất": (
                np.concatenate([p_of(z[k + "|b2_oof"]) for k in ks]),
                np.concatenate([p_of(z[k + "|b2_val"]) for k in ks]), oo, ov),
        }
        if za is not None:
            cands["A không neo"] = (
                np.concatenate([p_of(za[k + "|oof"][fa]) for k in ks]),
                np.concatenate([p_of(za[k + "|val"][fa]) for k in ks]),
                np.concatenate([za[k + "|obs_oof"] for k in ks]),
                np.concatenate([za[k + "|obs_val"] for k in ks]))
        blk = {}
        for nm, (po, pv, o1, o2) in cands.items():
            p1 = pick_p1(po, o1)
            g = gate_scores(pv, o2, p1)
            g["AUC"] = float(M.auc_score(pv, o2))
            g["AUC_oof"] = float(M.auc_score(po, o1))
            blk[nm] = g
        res["by_week"][w] = blk
        print("[%s] " % w + " | ".join(
            "%s AUC %.3f đóng %.0f%% NPV %.3f rò %.1f%%"
            % (nm, g["AUC"], 100 * g["closure"], g["NPV"], 100 * g["leakage"])
            for nm, g in blk.items()), flush=True)

    desc = {n: d for n, d, _ in OFF.offset_zoo()}
    for i, nm in enumerate(fam_b2):
        pv = np.concatenate([z[k + "|b2_val"][i].astype("f8") for k in keys])
        ov = np.concatenate([z[k + "|obs_val"] for k in keys])
        good = pv >= 0
        res["family_b2"][nm] = {
            "desc": desc.get(nm, ""),
            "AUC_val": float(M.auc_score(pv[good], ov[good])) if good.any() else np.nan,
            "fail_frac": float(1.0 - good.mean())}

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


if __name__ == "__main__":
    main()