#!/usr/bin/env python3
"""Đóng băng mọi thứ cần để CHIẾU trường mới lên cơ sở EOF lịch sử.

VÌ SAO CẦN. `calc_eof.py` tính EOF rồi chiếu luôn toàn bộ mẫu, nên nó không cần lưu
lại hai thứ trung gian: (1) khí hậu học theo (hạn dự báo, ngày trong năm) của trường
quy mô rộng, và (2) độ lệch chuẩn dùng để chuẩn hoá từng thành phần chính. Vận hành
thì bắt buộc phải có cả hai: một trường dự báo mới phải trừ ĐÚNG khí hậu học ấy và
chia ĐÚNG độ lệch chuẩn ấy, nếu không 20 thành phần chính trong F3 sẽ nằm trên một
cơ sở khác với lúc huấn luyện — sai lệch mà mô hình nuốt gọn không báo lỗi.

Cơ sở EOF (`patterns`) tuyệt đối không tính lại: nó phải là cơ sở đã phân tích từ
số liệu lịch sử, đúng như bản chất của phương pháp.

Chạy `--verify` để chiếu lại một vài lần khởi tạo lịch sử và so với `pcs` đã lưu;
sai số phải ở mức làm tròn số thực.
"""

from __future__ import annotations

import argparse
import os
import sys

import numpy as np
from joblib import Parallel, delayed

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import ehd_common as C
import build_cube as B
from calc_climatology import bin_by_doy, moments
from calc_eof import WIDE, NMODE, EOF_DIR, _read_year

OUT = os.path.join(EOF_DIR, "eof_frozen_%s.npz")


def wide_clim(name, jobs):
    """Khí hậu học (hạn, ngày trong năm, ô lưới) trên CÁC NĂM HUẤN LUYỆN, và sd."""
    varname, folder, pattern = WIDE[name]
    meta = B.load_meta()
    doy = meta["doy"]
    init_year = meta["init_year"]
    train = np.isin(init_year, C.TRAIN_YEARS)

    parts = Parallel(n_jobs=jobs, verbose=5)(
        delayed(_read_year)(y, varname, folder, pattern) for y in C.YEARS)
    field = np.concatenate(parts, axis=0)
    ninit, nlt, nlat, nlon = field.shape
    ng = nlat * nlon
    flat = field.reshape(ninit, nlt, ng)

    clim = np.empty((nlt, 365, ng), dtype="f4")
    for k in range(nlt):
        s0, s1, s2 = bin_by_doy(flat[train, k, :], doy[train, k], ng)
        clim[k] = moments(s0, s1, s2)[0]
        flat[:, k, :] -= clim[k][doy[:, k] - 1]

    z = np.load(os.path.join(EOF_DIR, "eof_%s.npz" % name), allow_pickle=True)
    lat, lon = z["lat"], z["lon"]
    vecs = z["patterns"].reshape(ng, NMODE)
    w = np.repeat(np.sqrt(np.cos(np.deg2rad(lat))).astype("f4"), nlon)
    flat *= w[None, None, :]
    pcs = (flat.reshape(ninit * nlt, ng) @ vecs)
    sd = pcs[np.repeat(train, nlt)].std(axis=0)
    return clim, sd.astype("f4"), lat, lon, vecs, z


def project(field, doy_row, clim, sd, vecs, lat):
    """Chiếu MỘT trường dự báo (nlt, nlat, nlon), vĩ độ đã tăng dần, lên cơ sở EOF.

    Đây là toàn bộ phép biến đổi mà vận hành cần: trừ khí hậu học đúng theo (hạn,
    ngày trong năm), nhân trọng số sqrt(cos(vĩ độ)) cho đúng diện tích, chiếu lên
    cơ sở lịch sử, rồi chia độ lệch chuẩn của thời kỳ huấn luyện.
    """
    nlt, nlat, nlon = field.shape
    x = field.reshape(nlt, nlat * nlon).astype("f4").copy()
    for k in range(nlt):
        x[k] -= clim[k][doy_row[k] - 1]
    x *= np.repeat(np.sqrt(np.cos(np.deg2rad(lat))).astype("f4"), nlon)[None, :]
    return (x @ vecs) / np.maximum(sd, 1e-9)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--jobs", type=int, default=20)
    ap.add_argument("--vars", nargs="*", default=list(WIDE))
    ap.add_argument("--verify", type=int, default=3,
                    help="số lần khởi tạo lịch sử dùng để kiểm chứng phép chiếu")
    args = ap.parse_args()

    meta = B.load_meta()
    doy = meta["doy"]
    for name in args.vars:
        print("[%s] đọc trường quy mô rộng và dựng khí hậu học ..." % name,
              flush=True)
        clim, sd, lat, lon, vecs, z = wide_clim(name, args.jobs)
        np.savez_compressed(OUT % name, clim=clim, sd=sd, lat=lat, lon=lon,
                            patterns=z["patterns"], nmode=NMODE,
                            train_years=np.array(C.TRAIN_YEARS))
        print("[%s] -> %s (%.0f MB)"
              % (name, OUT % name, os.path.getsize(OUT % name) / 1e6), flush=True)

        if args.verify:
            varname, folder, pattern = WIDE[name]
            worst = 0.0
            for i in np.linspace(0, doy.shape[0] - 1, args.verify).astype(int):
                yr = int(meta["init_year"][i])
                dstr = str(meta["init_dates"][i])
                from netCDF4 import Dataset
                with Dataset(os.path.join(folder, str(yr),
                                          pattern % dstr)) as ds:
                    f = np.ma.filled(ds[varname][:C.NLT], np.nan
                                     ).astype("f4")[:, ::-1, :]
                got = project(f, doy[i], clim, sd, vecs, lat)
                ref = z["pcs"][i]
                d = float(np.nanmax(np.abs(got - ref)))
                worst = max(worst, d)
                print("   kiểm chứng %s: lệch tối đa %.2e" % (dstr, d), flush=True)
            print("[%s] lệch lớn nhất %.2e -> %s"
                  % (name, worst, "ĐẠT" if worst < 1e-3 else "KHÔNG ĐẠT"),
                  flush=True)


if __name__ == "__main__":
    main()