import os
import time
from pathlib import Path

os.environ.setdefault("QLROM_DATA", str(Path.home() / ".cache" / "qlrom"))

import matplotlib.pyplot as plt
import numpy as np
import torch

from qlroms import compute_pod_basis, free_run, metrics
from qlroms.intrusive_qlroms import ks2d
from qlroms.intrusive_qlroms.build_ks import build_fom, get_full_trajectory
# the figure helpers every tutorial shares: chart palette, field panels, chart statistics
from qlroms.utils.plots import plot_evaluation, CHART_COLORS, chart_cmap, chart_stats, close_pdf, field_image, save_figure, snapshot_grid, start_pdf

plt.rcParams["figure.dpi"] = 160

t_start = time.time()
C_TRUTH, C_ROM, C_GLOBAL = "#2a78d6", "#eb6834", "#1baf7a"
SEED = 0

# (K, r) and window sizes are the ones the qlROM-DA study settled on for these cases.
CASES = {
    "travelling": dict(K=5,  r=25, Ntrain=30_000, Ntest=3_000, n_fc=3_000),
    "chaotic_B":  dict(K=10, r=50, Ntrain=60_000, Ntest=5_000, n_fc=2_000),
}

# every figure below also becomes a page of figs/02_local_chart_construction.pdf
start_pdf("02_local_chart_construction")

PosixPath('../outputs/02_local_chart_construction.pdf')

def load_case(name):
    """Trajectory + grid metadata for one KS2D case (cached under QLROM_DATA)."""
    P = CASES[name]
    fom, cfg = build_fom(ks2d, case=name, overrides={"Ntrain": P["Ntrain"], "Ntest": P["Ntest"]})
    traj = get_full_trajectory(Ntot=cfg["i0"] + P["Ntrain"] + P["Ntest"], model=fom, i0=cfg["i0"])
    X = traj[:, :P["Ntrain"] + P["Ntest"]]
    Nx, Ny = int(fom.Nx), int(fom.Ny)
    return dict(name=name, fom=fom, cfg=cfg, P=P, dt=float(fom.dt), Nx=Nx, Ny=Ny,
                shape=(Nx, Ny),          # how a flattened snapshot goes back to a field
                Xtrain=X[:, :P["Ntrain"]], Xtest=X[:, P["Ntrain"]:])


cases = {name: load_case(name) for name in CASES}
for c in cases.values():
    print(f"{c['name']:11s}: {c['Nx']}x{c['Ny']} grid (N_h = {c['Nx'] * c['Ny']}), dt = {c['dt']}, "
          f"train {tuple(c['Xtrain'].shape)}, test {tuple(c['Xtest'].shape)}")

/home/anovoama/.conda/envs/qlrom/lib/python3.11/site-packages/torch/cuda/__init__.py:188: UserWarning: CUDA initialization: The NVIDIA driver on your system is too old (found version 12080). Please update your GPU driver by downloading and installing a new version from the URL: http://www.nvidia.com/Download/index.aspx Alternatively, go to: https://pytorch.org to install a PyTorch version that has been compiled with your version of the CUDA driver. (Triggered internally at /__w/pytorch/pytorch/c10/cuda/CUDAFunctions.cpp:119.)
  return torch._C._cuda_getDeviceCount() > 0

Loading full trajectory from cache: /home/anovoama/.cache/qlrom/ks2d/nu1_0.50_nu2_0.35_Nx32_Ny32/dt0.1000/full_trajectory.pth.

Loading full trajectory from cache: /home/anovoama/.cache/qlrom/ks2d/nu1_0.30_nu2_0.10_Nx64_Ny64/dt0.0100/full_trajectory.pth.

travelling : 32x32 grid (N_h = 1024), dt = 0.1, train (1024, 30000), test (1024, 3000)
chaotic_B  : 64x64 grid (N_h = 4096), dt = 0.01, train (4096, 60000), test (4096, 5000)

save_dirs = {}
for c in cases.values():
    P = c["P"]
    save_dirs[c["name"]] = ks2d.get_simulation_path(model=c["fom"], Ntrain=P["Ntrain"],
                                                    Ntest=P["Ntest"], dt=c["dt"])
    t0 = time.time()
    c["rom"] = ks2d.build_local_model(c["Xtrain"], c["fom"], r=P["r"], K=P["K"],
                                      save_dir=save_dirs[c["name"]], method="galerkin",
                                      clustering_kwargs={"random_state": SEED})
    # affiliation of every training snapshot, and the phase-space view of it
    c["labels"] = torch.cdist(c["Xtrain"].T, 
                              c["rom"].centroids.cpu()).argmin(dim=1).numpy()
    qbar = c["Xtrain"].mean(dim=1, keepdim=True)
    c["Phi_gl"] = compute_pod_basis(c["Xtrain"][:, ::10] - qbar, 3, method="randomized")
    c["Z"] = (c["Phi_gl"].T @ (c["Xtrain"] - qbar)).numpy()                   # (3, Ntrain)
    c["Zc"] = (c["Phi_gl"].T @ (c["rom"].centroids.cpu().T - qbar)).numpy()   # (3, K)
    print(f"{c['name']:11s}: ql-Galerkin with K={c['rom'].K} charts, r={c['rom'].r} modes "
          f"({time.time() - t0:.0f} s)")

travelling : ql-Galerkin with K=5 charts, r=25 modes (0 s)

chaotic_B  : ql-Galerkin with K=10 charts, r=50 modes (2 s)

c = cases["travelling"]
rom = c["rom"]
print(f"per-chart operator shapes: b {tuple(rom.b_all[0].shape)}, A {tuple(rom.A_all[0].shape)}, "
      f"B {tuple(rom.B_all[0].shape)}  ->  {rom.B_all[0].numel():,} quadratic coefficients per chart")

fig, axs = plt.subplots(1, c["P"]["K"], figsize=(2.1 * c["P"]["K"], 2.4), layout="constrained")
vmax = float(rom.A_all.abs().max()) * 0.4
for k in range(c["P"]["K"]):
    im = field_image(axs[k], rom.A_all[k], vmax=vmax, origin="upper")   # a matrix, not a field
    axs[k].set_title(f"$A_{k}$", fontsize=10)
fig.colorbar(im, ax=axs, shrink=0.85, pad=0.01)
fig.suptitle("travelling: each chart gets its OWN linear operator (same modes index, different dynamics)",
             fontsize=10)
plt.show()

per-chart operator shapes: b (25,), A (25, 25), B (25, 25, 25)  ->  15,625 quadratic coefficients per chart

for c in cases.values():
    n_fc = c["P"]["n_fc"]
    t0 = time.time()
    X_rom, path = free_run(c["rom"], c["Xtest"][:, 0], n_steps=n_fc)
    X_rom = X_rom.cpu()
    # a quadratic ROM can leave its training regime and blow up: cut the run where it
    # stops being a forecast at all (non-finite, or twice the size of the truth)
    finite = torch.isfinite(X_rom).all(dim=0)
    too_big = X_rom.norm(dim=0) > 2 * c["Xtest"][:, :n_fc].norm(dim=0)
    bad = (~finite) | too_big
    n_valid = int(bad.to(torch.int).argmax()) if bad.any() else n_fc
    c["n_valid"], c["diverged"] = n_valid, n_valid < n_fc
    c["X_rom"], c["path"] = X_rom[:, :n_valid], path[:n_valid]
    c["X_ref"] = c["Xtest"][:, :n_valid]
    c["err"] = metrics.relative_error_series(c["X_ref"], c["X_rom"])
    c["horizon"] = metrics.prediction_horizon(c["err"], c["dt"], threshold=0.5)
    c["t"] = np.arange(n_valid) * c["dt"]
    note = f"BLEW UP after {n_valid} steps" if c["diverged"] else "finite throughout"
    print(f"{c['name']:11s}: {n_fc} closed-loop steps in {time.time() - t0:.1f} s | {note} "
          f"({n_valid * c['dt']:.1f} time units) | mean relative error {c['err'].mean():.3f} | "
          f"T_ph(0.5) = {c['horizon']:.2f} t.u.")

travelling : 3000 closed-loop steps in 3.5 s | finite throughout (300.0 time units) | mean relative error 0.058 | T_ph(0.5) = 300.00 t.u.

chaotic_B  : 2000 closed-loop steps in 7.7 s | BLEW UP after 1940 steps (19.4 time units) | mean relative error 1.412 | T_ph(0.5) = 2.67 t.u.

for c in cases.values():
    n_v = c["n_valid"]
    snaps = [0, n_v // 4, n_v // 2, n_v - 1]
    snapshot_grid([("truth", c["X_ref"]), ("ql-Galerkin", c["X_rom"]),
                   ("error", c["X_ref"] - c["X_rom"])],
                  snaps, c["shape"], figsize=(2.4 * len(snaps), 5.6),
                  col_titles=[f"$t={n * c['dt']:.1f}$ (chart {c['path'][n]})" for n in snaps],
                  suptitle=f"{c['name']}: closed-loop forecast")
    plt.show()

fig, axs = plt.subplots(2, 1, figsize=(10, 5), layout="constrained")
for ax, c in zip(axs, cases.values(), strict=True):
    cmap, norm = chart_cmap(c["P"]["K"])
    ax.semilogy(c["t"], c["err"], color="0.5", lw=1.0, zorder=1)
    ax.scatter(c["t"], c["err"], c=c["path"], cmap=cmap, norm=norm, s=4, zorder=2)
    ax.axhline(0.5, color="0.3", ls="--", lw=0.9)
    if c["horizon"] < c["t"][-1]:
        ax.axvline(c["horizon"], color="0.3", ls=":", lw=1.2)
        ax.text(c["horizon"], 1.1e-3, f"  $T_{{ph}}={c['horizon']:.1f}$", fontsize=9)
    ax.set(xlabel="$t$", ylabel="relative error", title=f"{c['name']}: error, colored by active chart")
    ax.grid(alpha=0.25, lw=0.5)
plt.show()

def to_atlas(model, X_phys):
    """Physical snapshots (N, Nt) -> exact global-atlas coordinates (r_g, Nt)."""
    apod = model.project_state(torch.as_tensor(X_phys, dtype=model.rdtype, device=model.device))
    a, ids = apod[:-1], apod[-1].long()
    return torch.einsum("mgr,rm->gm", model.Tgk[ids], a) + model.dgk[ids].T


for c in cases.values():
    ids_true = c["rom"].project_state(
        torch.as_tensor(c["X_ref"], dtype=c["rom"].rdtype,
                        device=c["rom"].device))[-1].long().cpu().numpy()
    fig, mt = plot_evaluation(
        to_atlas(c["rom"], c["X_ref"]).cpu().numpy(),
        {"qlROM closed loop": dict(z=to_atlas(c["rom"], c["X_rom"]).cpu().numpy(),
                                   err=c["err"], ids=c["path"])},
        dt=c["dt"], ids_true=ids_true, threshold=0.5,
        title=f"{c['name']}: K = {c['P']['K']} charts, r = {c['P']['r']}")
    save_figure(fig)
    plt.show()
    m = mt["qlROM closed loop"]
    horizon = "> window" if m["T_ph"] is None else f"{m['T_ph']:.1f}"
    print(f"{c['name']:11s} T_ph = {horizon} | mean rel err {m['err_mean']:.4f} | "
          f"charts matched {m['chart_match']:.1%}")

travelling  T_ph = > window | mean rel err 0.0577 | charts matched 96.5%

chaotic_B   T_ph = 2.7 | mean rel err 0.2197 | charts matched 26.6%

for c in cases.values():
    K, dt = c["P"]["K"], c["dt"]
    s = chart_stats(c["path"], K=K)         # occupancy, dwell segments, transition matrix
    dwell = s["seg_len"] * dt

    fig, axs = plt.subplots(1, 3, figsize=(11, 2.6), layout="constrained")
    axs[0].bar(np.arange(K), s["occupancy"], color=CHART_COLORS[:K])
    axs[0].set(xlabel="chart", ylabel="fraction of time", title="occupancy")
    axs[1].hist(dwell, bins=20, color=C_ROM)
    axs[1].set(xlabel="dwell time", ylabel="count", title=f"dwell (mean {dwell.mean():.2f} t.u.)")
    im = axs[2].imshow(s["T"], cmap="magma", vmin=0, vmax=1)
    axs[2].set(xlabel="to chart", ylabel="from chart", title=r"$P(j \mid i)$")
    fig.colorbar(im, ax=axs[2], pad=0.02)
    fig.suptitle(f"{c['name']}: {len(dwell) - 1} chart switches along the forecast", fontsize=10)
    plt.show()

c = cases["chaotic_B"]
Nx, Ny = c["Nx"], c["Ny"]
K_STAT = 40      # the chart count that keeps a chaotic_B free run finite (tutorial 1, sec 4)
rom_stat = ks2d.build_local_model(c["Xtrain"], c["fom"], r=c["P"]["r"], K=K_STAT,
                                  save_dir=save_dirs["chaotic_B"], method="galerkin",
                                  clustering_kwargs={"random_state": SEED})
X_run, _ = free_run(rom_stat, c["Xtest"][:, 0], n_steps=c["P"]["n_fc"])
X_run = X_run.cpu()
finite = torch.isfinite(X_run).all(dim=0)
X_stat = X_run[:, :int(finite.to(torch.int).argmin()) if (~finite).any() else X_run.shape[1]]
X_stat_ref = c["Xtest"][:, :X_stat.shape[1]]
print(f"statistics over {X_stat.shape[1]} steps ({X_stat.shape[1] * c['dt']:.0f} t.u.) "
      f"of the K={K_STAT} free run")


def radial_spectrum(X):
    """Time-averaged radial energy spectrum of flattened 2-D fields (N_h, T)."""
    f = np.asarray(X.detach().cpu()).T.reshape(-1, Nx, Ny)
    P = np.abs(np.fft.fft2(f)) ** 2 / (Nx * Ny) ** 2
    kx = np.fft.fftfreq(Nx) * Nx
    kk = np.rint(np.sqrt(kx[:, None] ** 2 + kx[None, :] ** 2)).astype(int)
    nb = kk.max() + 1
    return np.array([P[:, kk == q].sum(axis=1).mean() for q in range(nb)])


E_ref, E_rom = radial_spectrum(X_stat_ref), radial_spectrum(X_stat)
fig, axs = plt.subplots(1, 2, figsize=(10, 3.2), layout="constrained")
axs[0].hist(np.asarray(X_stat_ref).ravel(), bins=80, density=True, histtype="step",
            color=C_TRUTH, lw=1.6, label="truth")
axs[0].hist(np.asarray(X_stat).ravel(), bins=80, density=True, histtype="step",
            color=C_ROM, lw=1.6, label=f"ql-Galerkin free run (K={K_STAT})")
axs[0].set(xlabel="$u$", ylabel="pdf", title="(a) Pointwise distribution", yscale="log")
axs[1].loglog(np.arange(1, len(E_ref)), E_ref[1:], color=C_TRUTH, lw=1.6, label="truth")
axs[1].loglog(np.arange(1, len(E_rom)), E_rom[1:], color=C_ROM, lw=1.6, label="ql-Galerkin")
axs[1].set(xlabel=r"wavenumber $\kappa$", ylabel=r"$E(\kappa)$", ylim=(1e-14, None),
           title="(b) Radial energy spectrum")   # below ~1e-12 is round-off, not physics
for ax in axs:
    ax.legend(frameon=False)
    ax.grid(alpha=0.25, lw=0.5)
fig.suptitle("chaotic_B: after the trajectories part, do the statistics still agree?", fontsize=10)
plt.show()

statistics over 2000 steps (20 t.u.) of the K=40 free run

print("SUMMARY")
for c in cases.values():
    print(f"  {c['name']:11s} K={c['P']['K']:2d} r={c['P']['r']:2d} | free run {c['P']['n_fc']} "
          f"steps: mean rel err {c['err'].mean():.3f}, T_ph(0.5) = {c['horizon']:.1f} t.u.")
print(f"  runtime {time.time() - t_start:.0f}s")

SUMMARY
  travelling  K= 5 r=25 | free run 3000 steps: mean rel err 0.058, T_ph(0.5) = 300.0 t.u.
  chaotic_B   K=10 r=50 | free run 2000 steps: mean rel err 1.412, T_ph(0.5) = 2.7 t.u.
  runtime 37s

close_pdf()

9 figures written to ../outputs/02_local_chart_construction.pdf

PosixPath('../outputs/02_local_chart_construction.pdf')

object	what it is	built by
$\bm{c}_k,\ \beta$	the partition and the affiliation rule	k-means (section 5 of tutorial 1)
$\bar{\bm{q}}_k,\ \bm{\Phi}_k$	one affine chart per regime	per-cluster POD
$\bm{b}_k,\mathbf{A}_k,\mathsf{B}_k$	the reduced dynamics of that regime	intrusive projection (here)
$\mathbf{T}_{ji},\bm{d}_{ji}$	how coordinates move between charts	transition maps (tutorial 4)

Tutorial 2: building a local-chart ROM¶

1. The qlROM¶

1.1 The intrusive step: projecting the equations onto a chart¶

2. Running the model: one step, one chart¶

3. The chart machinery, watched from outside¶

4. Chaos: statistics instead of trajectories¶

5. What you just built¶