#!/usr/bin/env python3
"""Dựng đặc trưng cho MỘT lần khởi tạo mới, không cần khối memmap lịch sử.

`ai_data.Store` chỉ đọc được từ các khối đã dựng sẵn cho 2080 lần khởi tạo lịch sử,
nên vận hành cần một đường riêng đi thẳng từ `EC_all_<ngày>.nc`.  Mọi bước ở đây
sao đúng `build_cube.py` và `build_ai_inputs.py`; chỗ nào lệch là mô hình nhận sai
đầu vào mà không có cách nào biết.

Bốn chỗ dễ sai nhất, ghi ra để sau này còn kiểm:
  * `cp`, `tp` phải qua log1p TRƯỚC khi z-score, vì khí hậu học cũng tính trên
    log1p (build_cube.LOG_VARS);
  * `sv_mod` lấy từ mx2t6 THÔ, không phải từ z-score;
  * ngưỡng Tx của mô hình tra theo (hạn dự báo, ngày trong năm), còn ngưỡng quan
    trắc chỉ theo ngày trong năm;
  * thành phần chính EOF phải CHIẾU lên cơ sở lịch sử đã đóng băng.

Chạy trực tiếp file này với `--check <ngày lịch sử>` để đối chiếu với Store.
"""

from __future__ import annotations

import argparse
import datetime as dt
import os
import sys

import numpy as np
from netCDF4 import Dataset

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import ehd_common as C
import build_cube as B
import calc_climatology as K
import build_ai_inputs as A
from calc_eof import NMODE, EOF_DIR
from eof_freeze import project

STD_FLOOR = A.STD_FLOOR
WIDE_SRC = {"H500": ("gh", "H500", "H500_cf"), "T850": ("t", "T850", "T850_cf")}


class Frozen:
    """Toàn bộ khối tham số khí hậu đã đóng băng, nạp một lần rồi dùng lại."""

    def __init__(self):
        self.grid = C.Grid()
        npt = self.grid.npt
        nv = len(B.EC_VARS)
        self.mean = np.memmap(K.EC_MEAN, "f4", "r", shape=(nv, C.NLT, 365, npt))
        self.std = np.memmap(K.EC_STD, "f4", "r", shape=(nv, C.NLT, 365, npt))
        with Dataset(os.path.join(C.OUT, "Cli", "Qp_tx_obs.nc")) as ds:
            self.thr_obs = self.grid.to_points(
                np.ma.filled(ds["Tx_thres_obs"][:], np.nan).astype("f4"))
        with Dataset(os.path.join(C.OUT, "Cli", "Qp_tx_mod.nc")) as ds:
            self.thr_mod = self.grid.to_points(
                np.ma.filled(ds["Tx_thres_mod"][:], np.nan).astype("f4"))
        self.eof = {}
        for name in WIDE_SRC:
            z = np.load(os.path.join(EOF_DIR, "eof_frozen_%s.npz" % name))
            self.eof[name] = (z["clim"], z["sd"], z["lat"],
                              z["patterns"].reshape(-1, NMODE))


def read_ec_all(path, grid):
    """(46, 40, npt) theo đúng thứ tự biến của build_cube, đã qua log1p."""
    out = np.empty((C.NLT, len(B.EC_VARS), grid.npt), dtype="f4")
    with Dataset(path) as ds:
        for j, v in enumerate(B.EC_VARS):
            a = np.ma.filled(ds[v][:], np.nan).astype("f4")
            a[~np.isfinite(a)] = np.nan
            if v in B.LOG_VARS:
                a = np.log1p(np.maximum(a, 0.0))
            out[:, j, :] = grid.to_points(a)
    return out


def read_wide(root, name, init, year):
    varname, sub, pre = WIDE_SRC[name]
    p = os.path.join(root, sub, str(year), "%s_%s.nc" % (pre, init))
    with Dataset(p) as ds:
        return np.ma.filled(ds[varname][:C.NLT], np.nan).astype("f4")[:, ::-1, :]


def build(init, ec_all, wide_root, fr, fset="F3"):
    """Ma trận đặc trưng (46*npt, nfeat) theo đúng thứ tự cột của Store.build_isel."""
    grid = fr.grid
    npt = grid.npt
    doy = np.array([C.doy_noleap(d) for d in C.valid_days(init)], dtype="i4")
    raw = read_ec_all(ec_all, grid)

    z = np.empty_like(raw)
    for k in range(C.NLT):
        d = doy[k] - 1
        z[k] = (raw[k] - fr.mean[:, k, d, :]) / np.maximum(fr.std[:, k, d, :],
                                                           STD_FLOOR)
    jm = B.EC_VARS.index("mx2t6")
    thm = np.stack([fr.thr_mod[k][doy[k] - 1] for k in range(C.NLT)])   # (46,npt)
    tho = fr.thr_obs[doy - 1]                                           # (46,npt)
    sv = raw[:, jm, :] - thm
    ev = (sv >= 0).astype("f4")

    pcs = {}
    for name in WIDE_SRC:
        clim, sd, lat, vecs = fr.eof[name]
        f = read_wide(wide_root, name, init.isoformat(), init.year)
        pcs[name] = project(f, doy, clim, sd, vecs, lat)                # (46,NMODE)

    lead = (np.arange(C.NLT, dtype="f4") / 45.0)
    cols = [sv[:, :, None], ev[:, :, None],
            np.broadcast_to(lead[:, None, None], (C.NLT, npt, 1)),
            np.broadcast_to(np.cos(2 * np.pi * doy / 365.0).astype("f4")[:, None, None],
                            (C.NLT, npt, 1)),
            np.broadcast_to(np.sin(2 * np.pi * doy / 365.0).astype("f4")[:, None, None],
                            (C.NLT, npt, 1)),
            z.transpose(0, 2, 1)]
    for name in ("H500", "T850"):
        cols.append(np.broadcast_to(pcs[name][:, None, :], (C.NLT, npt, NMODE)))
    cols += [np.broadcast_to(grid.lat[grid.iy].astype("f4")[None, :, None],
                             (C.NLT, npt, 1)),
             np.broadcast_to(grid.lon[grid.ix].astype("f4")[None, :, None],
                             (C.NLT, npt, 1)),
             np.broadcast_to(grid.reg.astype("f4")[None, :, None],
                             (C.NLT, npt, 1)),
             tho[:, :, None], thm[:, :, None]]
    X = np.concatenate([c.astype("f4") for c in cols],
                       axis=2).reshape(C.NLT * npt, -1)
    return X, doy


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--check", required=True,
                    help="ngày khởi tạo LỊCH SỬ để đối chiếu với ai_data.Store")
    ap.add_argument("--ec-root", default=os.path.join(C.DATA, "EC"))
    ap.add_argument("--combined", default=os.path.join(C.DATA, "EC_Combined"))
    args = ap.parse_args()

    init = dt.date.fromisoformat(args.check)
    fr = Frozen()
    ec_all = os.path.join(args.combined, str(init.year),
                          "EC_all_%s.nc" % init.isoformat())
    X, doy = build(init, ec_all, args.ec_root, fr)

    import ai_data as D
    store = D.Store()
    dates = np.array([str(x) for x in store.st["init_dates"]])
    i = int(np.nonzero(dates == init.isoformat())[0][0])
    pts = np.arange(store.grid.reg.size)
    Xs, _, _ = store.build_isel(pts, np.array([i]), "F3")
    names = store.feature_names("F3")

    print("[check] khởi tạo %s | đường vận hành %s | Store %s"
          % (init, X.shape, Xs.shape))
    assert X.shape == Xs.shape, (X.shape, Xs.shape)
    d = np.abs(X - Xs)
    scale = np.maximum(np.abs(Xs), 1.0)
    rel = np.nanmax(d / scale, axis=0)
    print("%-16s %12s %12s" % ("cột", "lệch tuyệt đối", "lệch tương đối"))
    worst = np.argsort(-rel)[:8]
    for j in worst:
        print("%-16s %12.3e %12.3e" % (names[j], np.nanmax(d[:, j]), rel[j]))
    ok = float(np.nanmax(rel))
    print("\n=> lệch tương đối lớn nhất trên %d cột: %.3e -> %s"
          % (X.shape[1], ok, "ĐẠT" if ok < 1e-4 else "KHÔNG ĐẠT"))


if __name__ == "__main__":
    main()