from pathlib import Path
import os
import sys


SCRIPT_DIR = Path(__file__).resolve().parent
ROOT = SCRIPT_DIR.parent
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))
MPLCONFIGDIR = SCRIPT_DIR.parent / ".mplconfig_pic_full_hovmoller"
MPLCONFIGDIR.mkdir(exist_ok=True)
for lock_file in MPLCONFIGDIR.glob("*.matplotlib-lock"):
    try:
        lock_file.unlink()
    except OSError:
        pass
os.environ.setdefault("MPLCONFIGDIR", str(MPLCONFIGDIR))

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.colors import LinearSegmentedColormap

from lda.ks import KSParams, integrate_steps, random_initial_state, saved_trajectory


OUTDIR = SCRIPT_DIR
PNG_PATH = OUTDIR / "Fig01_full_truth_hovmoller_T1600.png"
PDF_PATH = OUTDIR / "Fig01_full_truth_hovmoller_T1600.pdf"
FORECAST_TRUTH_PATH = OUTDIR / "Fig01_forecast_truth_from_DA_start_T0_to_T600_dt005.npz"
RGB_CMAP_PATH = Path(r"amwg_blueyellowred.rgb")
USE_AMWG_CMAP = False


def load_rgb_colormap(rgb_path):
    colors = []
    with Path(rgb_path).open("r", encoding="utf-8") as f:
        for raw_line in f:
            line = raw_line.strip()
            if not line or line.startswith(";") or line.startswith("#") or line.lower().startswith("ncolors"):
                continue
            parts = line.split()
            if len(parts) < 3:
                continue
            colors.append([float(parts[0]) / 255.0, float(parts[1]) / 255.0, float(parts[2]) / 255.0])
    if not colors:
        raise ValueError(f"No RGB colors found in {rgb_path}")
    return LinearSegmentedColormap.from_list(Path(rgb_path).stem, colors, N=256)


def main():
    params = KSParams(
        n=2048,
        length=256.0,
        linear_growth=1.0,
        third_order_dispersion=0.1,
        fifth_order_dispersion=0.1,
        fourth_order_damping=1.0,
        dealias=True,
        max_abs_state=1.0e6,
    )
    dt = 0.05
    truth_spinup_steps = 10000
    ensemble_spinup_steps = 10000
    da_steps = 8000
    forecast_extension_steps = 4000
    da_plus_forecast_steps = da_steps + forecast_extension_steps
    total_steps = truth_spinup_steps + ensemble_spinup_steps + da_plus_forecast_steps
    total_time = total_steps * dt
    save_interval = 0.25

    rng = np.random.default_rng(76001)
    field0 = random_initial_state(
        params,
        rng,
        std=0.5,
        spectral_slope=4.0,
        max_initial_mode=32,
    )
    truth = saved_trajectory(
        field0,
        total_time=total_time,
        save_interval=save_interval,
        dt=dt,
        params=params,
        include_initial=True,
    ).astype(np.float64)

    da_start_state = integrate_steps(
        field0,
        truth_spinup_steps + ensemble_spinup_steps,
        dt,
        params,
    )
    forecast_truth = saved_trajectory(
        da_start_state,
        total_time=da_plus_forecast_steps * dt,
        save_interval=dt,
        dt=dt,
        params=params,
        include_initial=True,
    )
    np.savez_compressed(
        FORECAST_TRUTH_PATH,
        truth=forecast_truth.astype(np.float32),
        dt=np.asarray(dt, dtype=np.float64),
        da_steps=np.asarray(da_steps, dtype=np.int64),
        forecast_extension_steps=np.asarray(forecast_extension_steps, dtype=np.int64),
        truth_spinup_steps=np.asarray(truth_spinup_steps, dtype=np.int64),
        ensemble_spinup_steps=np.asarray(ensemble_spinup_steps, dtype=np.int64),
    )

    n_save, n_grid = truth.shape
    length = params.length
    t = np.arange(n_save, dtype=np.float64) * save_interval
    x = np.arange(n_grid, dtype=np.float64)

    q01, q99 = np.percentile(truth, [1.0, 99.0])
    vmax = max(abs(q01), abs(q99))

    plt.rcParams.update(
        {
            "font.family": "DejaVu Sans",
            "font.size": 15,
            "mathtext.fontset": "dejavusans",
            "axes.labelsize": 19,
            "xtick.labelsize": 16,
            "ytick.labelsize": 16,
        }
    )
    cmap = load_rgb_colormap(RGB_CMAP_PATH) if USE_AMWG_CMAP and RGB_CMAP_PATH.exists() else "RdBu_r"
    fig, ax = plt.subplots(figsize=(8.2, 9.2), facecolor="white")
    im = ax.imshow(
        truth,
        origin="lower",
        aspect="auto",
        interpolation="nearest",
        extent=[x[0], x[-1], t[0], t[-1]],
        cmap=cmap,
        vmin=-vmax,
        vmax=vmax,
    )

    t_truth_spinup = truth_spinup_steps * dt
    t_da_start = (truth_spinup_steps + ensemble_spinup_steps) * dt
    line1 = ax.axhline(t_truth_spinup, color="0.05", lw=3.0, ls="--", alpha=0.95)
    line2 = ax.axhline(t_da_start, color="0.05", lw=3.0, ls=":", alpha=0.95)

    ax.set_xlabel("Physical grid index")
    ax.set_ylabel("Time")
    ax.set_xlim(0, n_grid - 1)
    ax.set_ylim(0, total_time)
    ax.set_xticks([0, 512, 1024, 1536, 2047])

    ax.legend(
        [line1, line2],
        [r"$t=500$: truth spin-up ends", r"$t=1000$: DA period begins"],
        loc="upper left",
        frameon=True,
        facecolor="white",
        edgecolor="none",
        framealpha=0.82,
        fontsize=10.5,
        handlelength=2.7,
        borderpad=0.45,
        labelspacing=0.35,
    )

    cbar = fig.colorbar(im, ax=ax, pad=0.018)
    cbar.set_label(r"State $u(x,t)$", fontsize=17)
    cbar.ax.tick_params(labelsize=15)

    fig.tight_layout()
    fig.savefig(PNG_PATH, dpi=300, bbox_inches="tight")
    fig.savefig(PDF_PATH, bbox_inches="tight")
    print(f"saved {PNG_PATH}")
    print(f"saved {PDF_PATH}")
    print(f"saved {FORECAST_TRUTH_PATH}")


if __name__ == "__main__":
    main()
