#!/usr/bin/env python3
"""GĐ3: huấn luyện và ĐÓNG BĂNG bộ tham số vận hành.

Khác với mọi script trước trong dự án -- vốn chỉ tính điểm số rồi vứt mô hình đi --
script này LƯU LẠI mọi thứ cần để chạy nghiệp vụ:

  * 8 họ hồi quy của tầng 1 và 6 họ phân loại của tầng 2, khớp GỘP TOÀN QUỐC;
  * bảng hiệu chỉnh độ tán (edges, sd) của từng họ tầng 1 -- không có nó thì không
    tái tạo được thành viên phân vị, tức không tính được p = M/N;
  * ngưỡng cổng p1 RIÊNG TỪNG Ô LƯỚI;
  * ngưỡng quyết định tầng 2 RIÊNG TỪNG Ô LƯỚI, suy từ r chung của cửa sổ;
  * base rate khí hậu của từng ô, cần để suy ngưỡng khi đổi r.

RÀNG BUỘC VẬN HÀNH (nhắc lại để người dựng hệ nghiệp vụ không phạm phải):
  * Dữ liệu vào PHẢI là control forecast của ECMWF, KHÔNG được dùng trung bình tổ
    hợp: toàn bộ huấn luyện dựa trên control, trung bình tổ hợp là đối tượng thống
    kê khác (trơn hơn, phương sai nhỏ hơn) nên z-score sẽ co về 0 và hệ sẽ dự báo
    thiếu một cách hệ thống.
  * KHÔNG dùng số liệu quan trắc tại thời điểm khởi tạo; mọi đặc trưng lấy từ chính
    lần chạy EC.
  * Khối khí hậu (clim_ec_mean/std, EOF, Tx_thres) phải đóng băng theo model cycle
    và đi kèm bộ mô hình này; đổi cycle thì huấn luyện lại.
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time

import numpy as np
from joblib import Parallel, delayed, dump

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import ehd_common as C
import ai_data as D
import exp3_offset as OFF
from exp3_multi import conditional_spread, apply_spread
from phase1_arch import (WINDOWS, NQ, ZQ, NFOLD, build_window, zoo_s1, zoo_s2,
                         p_of)
from exp3_week_s1 import pick_p1

PACK = os.path.join(C.OUT, "operational")
# r đã chốt ở GĐ2, trên trường ngoài mẫu của 16 năm huấn luyện, kèm biên an toàn
# r đã chốt trên trường ngoài mẫu.  W1 là điểm ĐẠT TRỌN BỘ chỉ tiêu (mục 6).
# W2-W6 không có điểm nào đạt ở bất kỳ r nào, nên lấy r = 1 cho trường nhị phân
# KHÔNG THIÊN LỆCH (FBI ~ 1): điểm max-SEDI chỉ hơn 0.02-0.04 SEDI nhưng đẩy FBI
# lên 1.2-1.57, tức dự báo thừa tới hơn 50%.  Với W2-W6, sản phẩm chính là TRƯỜNG
# XÁC SUẤT, trường nhị phân chỉ là tiện ích.
R_BY_WINDOW = {"W1": 0.84, "W2": 1.00, "W3": 1.00, "W4": 1.00, "W5": 1.00,
               "W6": 1.00, "W1-2": 0.78, "W1-3": 0.72, "W1-6": 0.65}


def fit_s1_family(mk, Xtr, ztr):
    """Khớp một họ tầng 1 và trả về CẢ bảng hiệu chỉnh độ tán, để tái tạo phân vị."""
    m = mk(); m.fit(Xtr, ztr)
    mu = np.asarray(m.predict(Xtr)).ravel()
    edges, sd = conditional_spread(mu, ztr - mu)
    return m, edges, sd


def votes_from(models, X, thr):
    """p = M/N từ các họ đã khớp; thr là ngưỡng bỏ phiếu theo từng hàng."""
    v = np.zeros((len(models), X.shape[0]), dtype="i1")
    for i, (m, edges, sd) in enumerate(models):
        mu = np.asarray(m.predict(X)).ravel()
        sg = apply_spread(mu, edges, sd)
        v[i] = ((mu[:, None] + sg[:, None] * ZQ[None, :])
                >= np.asarray(thr, "f8")[:, None]).sum(axis=1)
    return p_of(v)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--window", required=True, choices=list(WINDOWS))
    ap.add_argument("--fset", default="F3")
    ap.add_argument("--jobs", type=int, default=20)
    ap.add_argument("--pack-dir", default=PACK,
                    help="thư mục đóng gói; mặc định là gói nghiệp vụ hiện hành")
    ap.add_argument("--s2-weights", default="",
                    help="trọng số họ tầng 2, dạng 'Ten=so:Ten=so'; rỗng = đều "
                         "nhau. Dùng để trộn hai phương án thành một ensemble có "
                         "trọng số (Ope1b).")
    ap.add_argument("--s2-families", default="",
                    help="danh sách họ tầng 2 cách nhau bằng dấu phẩy; rỗng = cả 6 "
                         "họ như PA1/OPE1")
    ap.add_argument("--years", default="train", choices=["train", "all"],
                    help="'train' = 16 năm huấn luyện, để số liệu khớp với báo cáo "
                         "kiểm định; 'all' = cả 20 năm, dùng khi phát hành nghiệp vụ")
    args = ap.parse_args()
    PACKDIR = args.pack_dir
    os.makedirs(PACKDIR, exist_ok=True)
    t0 = time.time()

    store = D.Store()
    pts = np.arange(store.grid.reg.size)
    lts = np.concatenate([C.WEEKS[i][1] for i in WINDOWS[args.window]])
    years = C.TRAIN_YEARS if args.years == "train" else C.YEARS
    A, S, Cnt, Yr, P, names = build_window(store, pts, years, args.fset, lts)
    ok = np.isfinite(A).all(axis=1) & np.isfinite(S) & np.isfinite(Cnt)
    A, S, Cnt, Yr, P = A[ok], S[ok], Cnt[ok], Yr[ok], P[ok]
    y = (Cnt >= 1).astype("f8")
    thr = np.zeros(A.shape[0])
    print("[op] %s: %d hàng, %d đặc trưng, %d năm, base rate %.3f"
          % (args.window, A.shape[0], A.shape[1], len(years), y.mean()), flush=True)

    z1, z2 = zoo_s1(A.shape[0]), zoo_s2(A.shape[0])
    if args.s2_families:
        want = [s.strip() for s in args.s2_families.split(",") if s.strip()]
        have = [n for n, _ in z2]
        miss = [w for w in want if w not in have]
        if miss:
            sys.exit("Không có họ tầng 2: %s (có: %s)" % (miss, have))
        z2 = [(n, mk) for n, mk in z2 if n in want]
        print("[op] tầng 2 chỉ dùng: %s" % ", ".join(n for n, _ in z2), flush=True)
    w2 = None
    if args.s2_weights:
        wmap = {}
        for it in args.s2_weights.replace(":", ",").split(","):
            if not it.strip():
                continue
            k, v = it.split("=")
            wmap[k.strip()] = float(v)
        have = [n for n, _ in z2]
        bad = [k for k in wmap if k not in have]
        if bad:
            sys.exit("Không có họ tầng 2: %s (có: %s)" % (bad, have))
        w2 = np.array([wmap.get(n, 0.0) for n, _ in z2], dtype="f8")
        if w2.sum() <= 0:
            sys.exit("Tổng trọng số tầng 2 phải > 0")
        w2 /= w2.sum()
        print("[op] trọng số tầng 2: %s"
              % ", ".join("%s %.4f" % (n, x) for (n, _), x in zip(z2, w2)),
              flush=True)
    par = Parallel(n_jobs=args.jobs)

    # ---- tầng 1: khớp đầy đủ (đem đi vận hành) và khớp ngoài mẫu (để chốt ngưỡng)
    fitted1 = par(delayed(fit_s1_family)(mk, A, S.astype("f8")) for _, mk in z1)
    folds = [f for f in np.array_split(
        np.random.default_rng(0).permutation(np.unique(Yr)), NFOLD) if f.size]
    masks = [np.isin(Yr, f) for f in folds]
    masks = [m for m in masks if m.any() and not m.all()]

    def oof_one(j, i):
        m, e, sd = fit_s1_family(z1[i][1], A[~masks[j]], S[~masks[j]].astype("f8"))
        mu = np.asarray(m.predict(A[masks[j]])).ravel()
        sg = apply_spread(mu, e, sd)
        return j, i, ((mu[:, None] + sg[:, None] * ZQ[None, :])
                      >= thr[masks[j]][:, None]).sum(axis=1).astype("i1")

    v = np.zeros((len(z1), A.shape[0]), dtype="i1")
    for j, i, o in par(delayed(oof_one)(j, i)
                       for j in range(len(masks)) for i in range(len(z1))):
        v[i, masks[j]] = o
    q_oof = p_of(v)
    print("[op] tầng 1 xong, %.0f phút" % ((time.time() - t0) / 60), flush=True)

    # ---- cổng: p1 riêng từng ô lưới
    p1 = {}
    keep = np.zeros(A.shape[0], bool)
    for p in pts:
        m = P == p
        if not m.any():
            continue
        p1[int(p)] = float(pick_p1(q_oof[m], y[m]))
        keep[m] = q_oof[m] >= p1[int(p)]

    # ---- tầng 2: khớp đầy đủ để đem đi vận hành
    def fit_s2_one(nm, mk, X, yy):
        if mk == "ann":
            m = OFF.OffsetNet(hidden=(16,), act="relu", epochs=30, batch=2048,
                              seed=0)
            m.fit(X, yy, np.zeros(X.shape[0]))
        else:
            m = mk(); m.fit(X, yy.astype(int))
        return nm, m

    def prob_s2(models, X):
        ps = []
        for nm, m in models:
            p = (m.predict_proba(X, np.zeros(X.shape[0])) if nm == "ANN_16"
                 else m.predict_proba(X)[:, 1])
            ps.append(np.clip(np.asarray(p, "f8"), 0, 1))
        return np.average(ps, axis=0, weights=w2)

    fitted2 = par(delayed(fit_s2_one)(nm, mk, A[keep], y[keep]) for nm, mk in z2)
    print("[op] tầng 2 xong trên %d mẫu, %.0f phút"
          % (int(keep.sum()), (time.time() - t0) / 60), flush=True)

    # ---- xác suất NGOÀI MẪU của tầng 2, để hiệu chỉnh ngưỡng.
    # Dùng mô hình đã khớp trên chính hàng đó thì ngưỡng sẽ quá tự tin và hệ sẽ dự
    # báo thiếu khi vận hành, nên phải khớp lại theo fold.
    s_oof = np.full(A.shape[0], -1.0)
    for j in range(len(masks)):
        trn = (~masks[j]) & keep
        tst = masks[j] & keep
        if not tst.any() or trn.sum() < 50 or np.unique(y[trn]).size < 2:
            continue
        mods = par(delayed(fit_s2_one)(nm, mk, A[trn], y[trn]) for nm, mk in z2)
        s_oof[tst] = prob_s2(mods, A[tst])
    print("[op] xác suất ngoài mẫu của tầng 2 xong, %.0f phút"
          % ((time.time() - t0) / 60), flush=True)

    # ---- ngưỡng tầng 2 riêng từng ô, suy từ r chung của cửa sổ
    r = R_BY_WINDOW[args.window]
    th2, base = {}, {}
    for p in pts:
        m = P == p
        if not m.any():
            continue
        b = float(y[m].mean()); base[int(p)] = b
        rate = float(np.clip(r * b, 1e-4, 0.999))
        th2[int(p)] = float(np.quantile(s_oof[m], 1.0 - rate))

    out = os.path.join(PACKDIR, "model_%s.joblib" % args.window)
    dump({"window": args.window, "lts": lts.tolist(), "fset": args.fset,
          "feature_names": names, "years": years,
          "stage1": fitted1, "stage1_names": [n for n, _ in z1],
          "stage2": fitted2, "s2_weights": (None if w2 is None else w2.tolist()),
          "vote_threshold": 0.0, "n_quantile": NQ,
          "zq": ZQ}, out, compress=3)
    meta = {"window": args.window, "r": r, "fset": args.fset,
            "n_train_rows": int(A.shape[0]), "years": [int(x) for x in years],
            "n_family_s1": len(z1), "n_family_s2": len(z2),
            "gate_p1": p1, "stage2_threshold": th2, "base_rate": base,
            "constraints": {
                "input": "ECMWF control forecast only, never the ensemble mean",
                "no_obs_at_init": True,
                "frozen_blocks": ["clim_ec_mean", "clim_ec_std", "EOF", "Tx_thres"]}}
    with open(os.path.join(PACKDIR, "meta_%s.json" % args.window), "w") as fh:
        json.dump(meta, fh, indent=1)
    print("[op] -> %s (%.0f MB) + meta, tổng %.0f phút"
          % (out, os.path.getsize(out) / 1e6, (time.time() - t0) / 60), flush=True)


if __name__ == "__main__":
    main()