#!/usr/bin/env python3
"""Settle the UTC-to-local day alignment of CFS Tx against observed Tx.

CFS TMAX is a 6-hour maximum stamped in UTC; VnGAD Tx is a local-day (UTC+7)
maximum. Two candidate groupings of the four 6-hour buckets into a "day" differ by
one calendar day, and choosing wrong costs a full day of lead skill without ever
raising an error. See cfs_common for the definitions of A and B.

The test is deliberately blunt: pair each candidate daily Tx with the observed Tx
of day init+k+delta for delta in -1, 0, +1 and see which (alignment, delta) has the
highest correlation at short lead. The right alignment must peak at delta = 0.
Nearest-neighbour is used for the spatial mapping -- this is a question about time,
and the interpolation scheme cannot change its answer.
"""

from __future__ import annotations

import argparse
import datetime as dt
import os
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import cfs_common as C
import build_obs as B


def nearest_map():
    """For each of the 541 land points, the (j, i) of the closest Gaussian cell."""
    g = C.flux_grid()
    z = np.load(B.TX_NPZ)
    olat, olon = z["lat"][z["jj"]], z["lon"][z["ii"]]
    jj = np.abs(olat[:, None] - g["wlat"][None, :]).argmin(axis=1)
    ii = np.abs(olon[:, None] - g["wlon"][None, :]).argmin(axis=1)
    return jj, ii


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--max-lead", type=int, default=10)
    ap.add_argument("--max-init", type=int, default=400)
    a = ap.parse_args()

    z = np.load(B.TX_NPZ)
    obs, oord = z["tx"], z["ordinal"]
    row = {int(o): r for r, o in enumerate(oord)}          # ordinal -> obs row
    pj, pi = nearest_map()
    ktmax = C.FLUX_NAMES.index("tmax")
    nlead = a.max_lead

    inits = [s for s in C.both_inits() if os.path.exists(C.flux6h_path(s))]
    if a.max_init:
        inits = inits[:: max(1, len(inits) // a.max_init)][:a.max_init]
    print("using %d inits, leads 0..%d, %d points" % (len(inits), nlead - 1, len(pj)))

    deltas = (-1, 0, 1)
    mod = {"A": [], "B": []}
    ob = {d: [] for d in deltas}
    for s in inits:
        arr = np.load(C.flux6h_path(s))[:, ktmax][:, pj, pi] - 273.15   # (184, npt)
        A = arr.reshape(C.NLT, 4, -1).max(axis=1)[:nlead]
        Bv = np.empty_like(A)
        Bv[0] = arr[0:3].max(axis=0)
        for k in range(1, nlead):
            Bv[k] = arr[4 * k - 1:4 * k + 3].max(axis=0)
        base = C.to_date(s).toordinal()
        ok = True
        cols = {}
        for d in deltas:
            rows = [row.get(base + k + d, -1) for k in range(nlead)]
            if min(rows) < 0:
                ok = False
                break
            cols[d] = obs[rows]
        if not ok:
            continue
        mod["A"].append(A); mod["B"].append(Bv)
        for d in deltas:
            ob[d].append(cols[d])

    for k in mod:
        mod[k] = np.concatenate(mod[k])
    for d in deltas:
        ob[d] = np.concatenate(ob[d])
    print("paired samples: %d" % mod["A"].size)

    print("\ncorrelation of daily Tx, model vs observed shifted by delta days")
    print("%-6s %9s %9s %9s" % ("align", "d=-1", "d=0", "d=+1"))
    best = None
    for al in ("A", "B"):
        rs = []
        for d in deltas:
            m, o = mod[al].ravel(), ob[d].ravel()
            g = np.isfinite(m) & np.isfinite(o)
            rs.append(float(np.corrcoef(m[g], o[g])[0, 1]))
        print("%-6s %9.4f %9.4f %9.4f" % (al, rs[0], rs[1], rs[2]))
        for d, r in zip(deltas, rs):
            if best is None or r > best[2]:
                best = (al, d, r)
    print("\nhighest: alignment %s at delta=%+d (r=%.4f)" % best)

    print("\nmean error (model - obs) at delta=0, degC")
    for al in ("A", "B"):
        m, o = mod[al].ravel(), ob[0].ravel()
        g = np.isfinite(m) & np.isfinite(o)
        print("  %s  ME=%+.3f  MAE=%.3f  RMSE=%.3f"
              % (al, float((m[g] - o[g]).mean()),
                 float(np.abs(m[g] - o[g]).mean()),
                 float(np.sqrt(((m[g] - o[g]) ** 2).mean()))))

    print("\nper-lead correlation at delta=0")
    npt = mod["A"].shape[1]
    ni = mod["A"].shape[0] // nlead
    print("%-5s %8s %8s" % ("lead", "A", "B"))
    for k in range(nlead):
        out = []
        for al in ("A", "B"):
            m = mod[al].reshape(ni, nlead, npt)[:, k].ravel()
            o = ob[0].reshape(ni, nlead, npt)[:, k].ravel()
            g = np.isfinite(m) & np.isfinite(o)
            out.append(float(np.corrcoef(m[g], o[g])[0, 1]))
        print("%-5d %8.4f %8.4f" % (k, out[0], out[1]))


if __name__ == "__main__":
    main()