#!/usr/bin/env python3
"""GĐ3: chạy dự báo nghiệp vụ từ bộ tham số đã đóng băng.

Đầu vào là MỘT lần khởi tạo của ECMWF (control forecast).  Đầu ra là, cho từng cửa
sổ hạn dự báo và từng ô lưới trong 541 ô: xác suất xảy ra EHDE và quyết định nhị
phân đã áp ngưỡng riêng của ô đó.

Quy trình đúng ba bước, giống hệt lúc huấn luyện:
  1. gộp các ngày trong cửa sổ thành một véc-tơ đặc trưng;
  2. tầng 1 -- N thành viên bỏ phiếu, p = M/N; p < p1 của ô đó thì kết luận nEHDE
     và DỪNG (xuất ra "không", chứ không phải "không phát báo");
  3. tầng 2 -- trung bình xác suất của các họ phân loại, so với ngưỡng của ô đó.

CẢNH BÁO: phải nạp đúng control forecast.  Nếu nạp trung bình tổ hợp, z-score sẽ co
về 0 và hệ dự báo thiếu một cách hệ thống mà không báo lỗi gì cả.
"""

from __future__ import annotations

import argparse
import datetime as dt
import json
import os
import sys

import numpy as np
from joblib import load
from netCDF4 import Dataset

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import ehd_common as C
import ehd_metrics as M
import ai_data as D
import ope1_features as OF
from exp3_multi import apply_spread
from exp3_week import MAX_VARS
from phase1_arch import p_of

PACK = os.path.join(C.OUT, "operational")


def aggregate(X, names, ni, npt, lts):
    """Gộp cửa sổ từ ma trận (ni*46*npt, nf) -> (ni*npt, 81)."""
    imax = [names.index(v) for v in MAX_VARS if v in names]
    isv = names.index("sv_mod")
    nf = X.shape[1]
    X = X.reshape(ni, C.NLT, npt, nf)[:, lts, :, :]
    return np.concatenate(
        [np.nanmean(X, axis=1), np.nanmax(X[:, :, :, imax], axis=1),
         (X[:, :, :, isv] >= 0.0).sum(axis=1, dtype="f4")[:, :, None]],
        axis=2).reshape(ni * npt, -1)


def accum_sev(X, names, ni, npt, lts):
    """Tổng Sev>=0 tích luỹ trong cửa sổ, (ni*npt,), đơn vị °C·ngày.

    Đây là đại lượng "tổng Sev>=0 dự báo trong tuần": một ngày vượt ngưỡng thì
    nhỏ (trừ khi ngày đó rất nóng), nhiều ngày vượt thì lớn, không ngày nào vượt
    thì đúng bằng 0.  Lấy thẳng từ sv_mod của chính lần chạy EC nên KHÔNG phải
    sản phẩm của mô hình AI.
    """
    isv = names.index("sv_mod")
    # lấy cột sv_mod TRƯỚC rồi mới cắt cửa sổ: nếu viết [:, lts, :, isv] thì mảng
    # chỉ số và số nguyên bị cách nhau bởi slice, NumPy đẩy chiều advanced lên đầu
    # và ra (len(lts), ni, npt) chứ không phải (ni, len(lts), npt).
    sv = X.reshape(ni, C.NLT, npt, X.shape[1])[..., isv][:, lts, :]
    return np.nansum(np.maximum(sv, 0.0), axis=1).reshape(ni * npt)


def features_hist(store, points, isel, fset, lts):
    """Đường LỊCH SỬ: lấy đặc trưng từ khối memmap đã dựng sẵn."""
    X, _, _ = store.build_isel(points, isel, fset)
    nm = store.feature_names(fset)
    return (aggregate(X, nm, isel.size, points.size, lts),
            accum_sev(X, nm, isel.size, points.size, lts))


def features_rt(init, ec_all, wide_root, fr, names, lts, npt):
    """Đường THỜI GIAN THỰC: đi thẳng từ EC_all_<ngày>.nc, không cần khối lịch sử."""
    X, _ = OF.build(init, ec_all, wide_root, fr)
    return aggregate(X, names, 1, npt, lts), accum_sev(X, names, 1, npt, lts)


def forecast(pack, meta, A, points, ni):
    """Hai tầng -> xác suất và quyết định cho từng (init, ô lưới)."""
    thr = np.full(A.shape[0], pack["vote_threshold"], dtype="f8")
    zq = pack["zq"]
    v = np.zeros((len(pack["stage1"]), A.shape[0]), dtype="i1")
    sev_sum = np.zeros(A.shape[0], dtype="f8")      # cộng dồn Sev>=0 của mọi thành viên
    for i, (m, edges, sd) in enumerate(pack["stage1"]):
        mu = np.asarray(m.predict(A)).ravel()
        sg = apply_spread(mu, edges, sd)
        mem = mu[:, None] + sg[:, None] * zq[None, :]      # (nrow, NQ) thành viên
        v[i] = (mem >= thr[:, None]).sum(axis=1)
        sev_sum += np.maximum(mem, 0.0).sum(axis=1)
    p_gate = p_of(v)
    n_member = len(pack["stage1"]) * zq.size
    sev_ai = sev_sum / n_member                     # trung bình trên N thành viên

    ps = []
    for nm, m in pack["stage2"]:
        p = (m.predict_proba(A, np.zeros(A.shape[0])) if nm == "ANN_16"
             else m.predict_proba(A)[:, 1])
        ps.append(np.clip(np.asarray(p, "f8"), 0, 1))
    sw = pack.get("s2_weights")
    p2 = np.mean(ps, axis=0) if sw is None else np.average(
        ps, axis=0, weights=np.asarray(sw, "f8"))

    pid = np.broadcast_to(np.asarray(points)[None, :],
                          (ni, len(points))).reshape(-1)
    gate_p1 = np.array([meta["gate_p1"][str(int(p))] for p in pid])
    th2 = np.array([meta["stage2_threshold"][str(int(p))] for p in pid])
    open_gate = p_gate >= gate_p1
    prob = np.where(open_gate, p2, 0.0)          # cổng đóng => khẳng định KHÔNG
    pred = open_gate & (p2 > th2)
    return (prob.reshape(ni, -1), pred.reshape(ni, -1),
            p_gate.reshape(ni, -1), sev_ai.reshape(ni, -1))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--init", required=True,
                    help="ngày khởi tạo YYYY-MM-DD; 'last' = lần khởi tạo cuối "
                         "cùng; 'valid' = toàn bộ khởi tạo của 4 năm kiểm định, "
                         "dùng để KIỂM CHỨNG rằng bộ tham số đã đóng băng tái tạo "
                         "đúng kỹ năng đã báo cáo")
    ap.add_argument("--windows", nargs="*", default=["W1", "W1-2", "W1-3"])
    ap.add_argument("--out", default=os.path.join(C.OUT, "operational",
                                                  "forecast.nc"))
    ap.add_argument("--from-ecall", default=None,
                    help="đường dẫn EC_all_<ngày>.nc; có thì chạy đường THỜI GIAN "
                         "THỰC, không cần khối memmap lịch sử")
    ap.add_argument("--wide-root", default=os.path.join(C.DATA, "EC"),
                    help="thư mục chứa H500/ và T850/ để chiếu EOF")
    ap.add_argument("--verify", action="store_true",
                    help="chấm điểm ngay kết quả vừa dự báo (chỉ có nghĩa với "
                         "--init valid, vì cần quan trắc)")
    ap.add_argument("--pack", default=PACK,
                    help="thư mục chứa model_<W>.joblib và meta_<W>.json")
    args = ap.parse_args()
    globals()["PACK"] = args.pack

    store = D.Store()
    points = np.arange(store.grid.reg.size)
    frozen = OF.Frozen() if args.from_ecall else None
    dates = np.array([str(x) for x in store.st["init_dates"]])
    if args.from_ecall:
        # Đường thời gian thực không tra khối lịch sử: ngày khởi tạo có thể chưa
        # từng tồn tại trong bộ 2080 mẫu.  isel chỉ còn dùng để đếm số lần khởi tạo.
        isel = np.array([0])
    elif args.init == "last":
        isel = np.array([store.init_year.size - 1])
    elif args.init == "valid":
        isel = np.nonzero(np.isin(store.init_year, C.VALID_YEARS))[0]
    else:
        hit = np.nonzero(np.char.find(dates, args.init) >= 0)[0]
        isel = hit if hit.size else np.array([int(args.init)])
    if isel.size == 0:
        sys.exit("Không tìm thấy initialisation %s" % args.init)
    if args.from_ecall:
        print("[fcs] khởi tạo %s (đường thời gian thực, từ %s)"
              % (args.init, os.path.basename(args.from_ecall)), flush=True)
    else:
        print("[fcs] %d lần khởi tạo, từ %s tới %s"
              % (isel.size, dates[isel[0]], dates[isel[-1]]), flush=True)

    with Dataset(args.out, "w") as ds:
        ds.createDimension("point", points.size)
        ds.createDimension("init", isel.size)
        ds.source = ("EHD S2S, hai tầng, mô hình gộp toàn quốc + ngưỡng riêng từng "
                     "ô lưới. Đầu vào PHẢI là control forecast của ECMWF.")
        for w in args.windows:
            pf = os.path.join(PACK, "model_%s.joblib" % w)
            mf = os.path.join(PACK, "meta_%s.json" % w)
            if not (os.path.exists(pf) and os.path.exists(mf)):
                print("[bỏ qua] chưa có bộ tham số cho %s" % w, flush=True)
                continue
            pack = load(pf)
            with open(mf) as fh:
                meta = json.load(fh)
            lts = np.asarray(pack["lts"])
            if args.from_ecall:
                A, acc = features_rt(dt.date.fromisoformat(args.init),
                                     args.from_ecall, args.wide_root, frozen,
                                     store.feature_names(pack["fset"]), lts,
                                     points.size)
            else:
                A, acc = features_hist(store, points, isel, pack["fset"], lts)
            prob, pred, pg, sev_ai = forecast(pack, meta, A, points, isel.size)
            acc = acc.reshape(isel.size, -1)
            for nm, arr, typ in (("prob", prob, "f4"),
                                 ("pred", pred.astype("i1"), "i1"),
                                 ("p_gate", pg, "f4"),
                                 ("sev_ai", sev_ai, "f4"),
                                 ("sev_acc", acc, "f4")):
                v = ds.createVariable("%s_%s" % (nm, w.replace("-", "_")), typ,
                                      ("init", "point"), zlib=True)
                v[:] = arr
            msg = "[fcs] %s: %d/%d ô lưới báo có EHDE" % (w, int(pred.sum()),
                                                           pred.size)
            if args.verify:
                # quan trắc chỉ dùng để CHẤM ĐIỂM, không bao giờ vào đặc trưng
                _, yv, _ = store.build_isel(points, isel, pack["fset"])
                evt = yv["evt"].reshape(isel.size, C.NLT, points.size)[:, lts, :]
                obs = (np.nansum(evt, axis=1) >= 1.0)
                ok = np.isfinite(prob) & np.isfinite(obs.astype("f4"))
                g = M.scores(*M.contingency(pred[ok], obs[ok]))
                msg += ("  |  POD %.3f POFD %.3f FAR %.3f FBI %.3f CSI %.3f "
                        "SEDI %.3f (n=%d)" % (g["POD"], g["POFD"], g["FAR"],
                                              g["FBI"], g["CSI"], g["SEDI"],
                                              g["n"]))
            print(msg, flush=True)
    print("[done] -> %s" % args.out, flush=True)


if __name__ == "__main__":
    main()