#!/usr/bin/env python3
# ----------------------------------------------------------------------
# Copyright (C) 2025 Tristan Hoellinger
# Distributed under the GNU General Public License v3.0 (GPLv3).
# See the LICENSE file in the root directory for details.
# SPDX-License-Identifier: GPL-3.0-or-later
# ----------------------------------------------------------------------

__author__ = "Tristan Hoellinger"
__version__ = "0.1.0"
__date__ = "2025"
__license__ = "GPLv3"

"""Plot results of time step convergence tests."""

# pyright: reportWildcardImportFromLibrary=false
from pathlib import Path
import numpy as np

from pysbmy.fft import read_FourierGrid
from pysbmy.field import read_field
from pysbmy.correlations import get_autocorrelation
from pysbmy.timestepping import StandardTimeStepping

from wip3m import *
from wip3m.tools import none_bool_str
from wip3m.plot_utils import *

from wip3m.logger import getCustomLogger, INDENT, UNINDENT

logger = getCustomLogger(__name__)

workdir = ROOT_PATH + "results/"
output_path = OUTPUT_PATH

from argparse import ArgumentParser

parser = ArgumentParser(description="Plot convergence results for WIP3M.")
parser.add_argument(
    "--run_id",
    type=str,
    help="Run ID to use for plotting.",
)
parser.add_argument(
    "--outdir_suffix",
    type=str,
    default="plots",
    help="Suffix to append to the output directory.",
)
parser.add_argument(
    "--ref",
    type=str,
    default="P3M1",
    choices=["P3M1", "PMref"],
    help="Reference field for power spectrum comparison.",
)
parser.add_argument(
    "--lpt",
    type=none_bool_str,
    default=False,
    help="Whether to compute and plot the LPT power spectrum.",
)
parser.add_argument(
    "--cola",
    type=none_bool_str,
    default=True,
    help="Whether to compute and plot the COLA power spectrum.",
)
parser.add_argument(
    "--plot_fields",
    type=none_bool_str,
    default=False,
    help="Whether to plot the fields or not.",
)
args = parser.parse_args()

run_id = args.run_id
outdir_suffix = args.outdir_suffix
ref = args.ref
lpt = args.lpt
cola = args.cola
plot_fields = args.plot_fields

if not run_id:
    raise ValueError("Please provide a run ID using the --run_id argument.")

simdir = output_path + run_id + "/"
outdir = simdir + outdir_suffix + "/"
input_ss_file = simdir + "input_ss_k_grid.h5"
wd = workdir + run_id + "/"

Path(outdir).mkdir(parents=True, exist_ok=True)

if __name__ == "__main__":
    logger.info(f"Plotting convergence results for run ID: {run_id}")
    logger.info(f"Output directory: {outdir}")

    logger.info(f"Reading simulation parameters from {wd}sim_params.txt...")
    with open(wd + "sim_params.txt", "r") as f:
        sim_params = eval(f.read())
    L = sim_params["L"]  # Box size in Mpc/h
    N = sim_params["N"]  # Density grid size
    Np = sim_params["Np"]  # Number of dark matter particles per spatial dimension
    Npm = sim_params["Npm"]  # PM grid size
    nsteps_pmref = sim_params["nsteps_pmref"]
    nsteps_pm1 = sim_params["nsteps_pm1"]
    nsteps_pm2 = sim_params["nsteps_pm2"]
    nsteps_cola = sim_params["nsteps_cola"]
    nsteps_spm = sim_params["nsteps_spm"]
    nsteps_p3m1 = sim_params["nsteps_p3m1"]
    nsteps_p3m2 = sim_params["nsteps_p3m2"]
    nsteps_p3m3 = sim_params["nsteps_p3m3"]
    tsd_spm = sim_params["tsd_spm"]
    tsd_p3m1 = sim_params["tsd_p3m1"]
    tsd_p3m2 = sim_params["tsd_p3m2"]
    tsd_p3m3 = sim_params["tsd_p3m3"]
    logger.info(f"Reading simulation parameters from {wd}sim_params.txt done.")

    if plot_fields:
        logger.info(f"Loading fields from {simdir}...")
        INDENT()
        slice_ijk = (N // 2, slice(None), slice(None))
        DELTA_LPT = read_field(simdir + "lpt_density.h5").data[slice_ijk]
        DELTA_PMREF = read_field(simdir + f"pmref_nsteps{nsteps_pmref}_final_density_pm.h5").data[
            slice_ijk
        ]
        DELTA_PM1 = read_field(simdir + f"pm1_nsteps{nsteps_pm1}_final_density_pm.h5").data[
            slice_ijk
        ]
        DELTA_PM2 = read_field(simdir + f"pm2_nsteps{nsteps_pm2}_final_density_pm.h5").data[
            slice_ijk
        ]
        DELTA_COLA = read_field(simdir + f"cola_nsteps{nsteps_cola}_final_density_cola.h5").data[
            slice_ijk
        ]
        DELTA_SPM = read_field(simdir + f"spm_nsteps{nsteps_spm}_final_density_spm.h5").data[
            slice_ijk
        ]
        DELTA_P3M1 = read_field(simdir + f"p3m1_nsteps{nsteps_p3m1}_final_density_p3m.h5").data[
            slice_ijk
        ]
        DELTA_P3M2 = read_field(simdir + f"p3m2_nsteps{nsteps_p3m2}_final_density_p3m.h5").data[
            slice_ijk
        ]
        DELTA_P3M3 = read_field(simdir + f"p3m3_nsteps{nsteps_p3m3}_final_density_p3m.h5").data[
            slice_ijk
        ]
        UNINDENT()
        logger.info(f"Loading fields from {simdir} done.")

        diff_pm1_pmref = DELTA_PM1 - DELTA_PMREF
        diff_pm2_pmref = DELTA_PM2 - DELTA_PMREF
        diff_p3m1_pmref = DELTA_P3M1 - DELTA_PMREF
        diff_p3m2_pm1 = DELTA_P3M2 - DELTA_PM1
        diff_p3m3_pm2 = DELTA_P3M3 - DELTA_PM2
        diff_p3m1_spm = DELTA_P3M1 - DELTA_SPM

        fields = [
            "pmref",
            "pm1",
            "pm2",
            "p3m1",
            "p3m2",
            "p3m3",
            "diff_p3m1_spm",
            "diff_pm1_pmref",
            "diff_pm2_pmref",
            "diff_p3m1_pmref",
            "diff_p3m2_pm1",
            "diff_p3m3_pm2",
        ]  # fields to plot
        ncols = 6  # Max panels per row

        figname = "_".join(fields)
        slices_dict = {
            "lpt": DELTA_LPT,
            "cola": DELTA_COLA,
            "pmref": DELTA_PMREF,
            "pm1": DELTA_PM1,
            "pm2": DELTA_PM2,
            "spm": DELTA_SPM,
            "p3m1": DELTA_P3M1,
            "p3m2": DELTA_P3M2,
            "p3m3": DELTA_P3M3,
            "diff_pm1_pmref": diff_pm1_pmref,
            "diff_pm2_pmref": diff_pm2_pmref,
            "diff_p3m1_pmref": diff_p3m1_pmref,
            "diff_p3m2_pm1": diff_p3m2_pm1,
            "diff_p3m3_pm2": diff_p3m3_pm2,
            "diff_p3m1_spm": diff_p3m1_spm,
        }
        titles_dict = {
            "lpt": "LPT",
            "pmref": f"PMref $n_\\mathrm{{steps}}={nsteps_pmref}$",
            "pm1": f"PM1 $n_\\mathrm{{steps}}={nsteps_pm1}$",
            "pm2": f"PM2 $n_\\mathrm{{steps}}={nsteps_pm2}$",
            "cola": f"COLA $n_\\mathrm{{steps}}={nsteps_cola}$",
            "spm": f"sPM $n_\\mathrm{{steps}}={nsteps_spm}$",
            "p3m1": f"P3M1 $n_\\mathrm{{steps}}={nsteps_p3m1}$",
            "p3m2": f"P3M2 $n_\\mathrm{{steps}}={nsteps_p3m2}$",
            "p3m3": f"P3M3 $n_\\mathrm{{steps}}={nsteps_p3m3}$",
            "diff_pm1_pmref": r"$\delta_{\rm PM1}-\delta_{\rm PMref}$",
            "diff_pm2_pmref": r"$\delta_{\rm PM2}-\delta_{\rm PMref}$",
            "diff_p3m1_pmref": r"$\delta_{\rm P3M1}-\delta_{\rm PMref}$",
            "diff_p3m2_pm1": r"$\delta_{\rm P3M2}-\delta_{\rm PM1}$",
            "diff_p3m3_pm2": r"$\delta_{\rm P3M3}-\delta_{\rm PM2}$",
            "diff_p3m1_spm": r"$\delta_{\rm P3M1}-\delta_{\rm sPM}$",
        }

        logger.info(f"Plotting fields: {fields}...")
        npanels = len(fields)
        nrows = np.ceil(npanels / ncols).astype(int)

        fig, axs = plt.subplots(
            nrows=nrows, ncols=ncols, figsize=(3 * ncols, 4.6 * nrows), sharey=True
        )

        axs = axs.flatten()
        ims = []

        for i, key in enumerate(fields):
            ax = axs[i]
            data = slices_dict[key]
            title = titles_dict[key]

            if key.startswith("diff"):
                norm = TwoSlopeNorm(
                    vmin=-np.log(1 + np.abs(np.min(data))),
                    vcenter=0,
                    vmax=np.log10(1 + np.abs(np.max(data))),
                )
                im = ax.imshow(np.sign(data) * np.log(1 + np.abs(data)), cmap="RdBu_r", norm=norm)
            else:
                im = ax.imshow(np.log10(2 + data), cmap=cmap)

            ims.append((im, key))
            ax.set_title(title, fontsize=fs_titles)
            for spine in ax.spines.values():
                spine.set_visible(False)

        # Hide unused axes
        for i in range(npanels, len(axs)):
            axs[i].axis("off")

        # Shared axes labels
        for i, ax in enumerate(axs[:npanels]):
            if i % ncols == 0:
                ax.set_yticks([0, N // 2, N])
                ax.set_yticklabels([f"{-L/2:.0f}", "0", f"{L/2:.0f}"], fontsize=fs)
                ax.set_ylabel(r"Mpc/$h$", size=GLOBAL_FS_SMALL)

            ax.set_xticks([0, N // 2, N])
            ax.set_xticklabels([f"{-L/2:.0f}", "0", f"{L/2:.0f}"], fontsize=fs)
            ax.set_xlabel(r"Mpc/$h$", size=GLOBAL_FS_SMALL)

        # Colorbars
        for ax, (im, key) in zip(axs[:npanels], ims):
            divider = make_axes_locatable(ax)
            cax = divider.append_axes("bottom", size="5%", pad=0.6)
            cb = fig.colorbar(im, cax=cax, orientation="horizontal")
            if key.startswith("diff"):
                cb.set_label(
                    r"$\textrm{sgn}\left(\Delta\delta\right)\log_{10}(1 + |\Delta\delta|)$",
                    fontsize=fs,
                )
            else:
                cb.set_label(r"$\log_{10}(2 + \delta)$", fontsize=fs)
            cb.ax.tick_params(labelsize=fs)
            cax.xaxis.set_ticks_position("bottom")
            cax.xaxis.set_label_position("bottom")

        fig.savefig(outdir + f"{figname}.png", bbox_inches="tight", dpi=300, transparent=True)
        fig.savefig(outdir + f"{figname}.pdf", bbox_inches="tight", dpi=300)
        plt.close(fig)
        logger.info(f"Plotting fields done. Saved to {outdir}{figname}.png and {figname}.pdf")

    logger.info("Computing and plotting power spectra...")
    INDENT()
    logger.info("Loading Fourier grid...")
    G = read_FourierGrid(input_ss_file)
    logger.info("Loading Fourier grid done.")

    k = G.k_modes
    AliasingCorr = False

    logger.info("Computing power spectra...")
    DELTA = read_field(simdir + "initial_density.h5")
    Pk_INI, _ = get_autocorrelation(DELTA, G, AliasingCorr)

    if lpt:
        DELTA = read_field(simdir + "lpt_density.h5")
        Pk_LPT, _ = get_autocorrelation(DELTA, G, AliasingCorr)

    DELTA = read_field(simdir + f"pmref_nsteps{nsteps_pmref}_final_density_pm.h5")
    Pk_PMref, _ = get_autocorrelation(DELTA, G, AliasingCorr)

    DELTA = read_field(simdir + f"pm1_nsteps{nsteps_pm1}_final_density_pm.h5")
    Pk_PM1, _ = get_autocorrelation(DELTA, G, AliasingCorr)

    DELTA = read_field(simdir + f"pm2_nsteps{nsteps_pm2}_final_density_pm.h5")
    Pk_PM2, _ = get_autocorrelation(DELTA, G, AliasingCorr)

    if cola:
        DELTA = read_field(simdir + f"cola_nsteps{nsteps_cola}_final_density_cola.h5")
        Pk_COLA, _ = get_autocorrelation(DELTA, G, AliasingCorr)

    DELTA = read_field(simdir + f"spm_nsteps{nsteps_spm}_final_density_spm.h5")
    Pk_sPM, _ = get_autocorrelation(DELTA, G, AliasingCorr)

    DELTA = read_field(simdir + f"p3m1_nsteps{nsteps_p3m1}_final_density_p3m.h5")
    Pk_P3M1, _ = get_autocorrelation(DELTA, G, AliasingCorr)

    DELTA = read_field(simdir + f"p3m2_nsteps{nsteps_p3m2}_final_density_p3m.h5")
    Pk_P3M2, _ = get_autocorrelation(DELTA, G, AliasingCorr)

    DELTA = read_field(simdir + f"p3m3_nsteps{nsteps_p3m3}_final_density_p3m.h5")
    Pk_P3M3, _ = get_autocorrelation(DELTA, G, AliasingCorr)
    logger.info("Computing power spectra done.")

    if ref == "P3M1":
        nsteps_reference = nsteps_p3m1
        Pk_ref = Pk_P3M1
        label_ref = f"Ref (P3M, $n_\\mathrm{{steps}}={nsteps_reference}$)"
    elif ref == "PMref":
        nsteps_reference = nsteps_pm1
        Pk_ref = Pk_PMref
        label_ref = f"Ref (PM, $n_\\mathrm{{steps}}={nsteps_reference}$)"
    else:
        raise ValueError("Invalid reference field. Choose 'P3M1' or 'PMref'.")

    logger.info("Plotting time step diagnostics...")
    INDENT()

    file_ext = f"spm_nsteps{nsteps_spm}"
    aiDrift = StandardTimeStepping.read(wd + file_ext + f"_ts_spm.h5").aiDrift
    plot_timestepping_diagnostics(
        log_path=simdir + "spm_timestep_log.txt",
        aiDrift=aiDrift,
        TimeStepDistribution=tsd_spm,
        nsteps=nsteps_spm,
        save_path=outdir + "time_step_diagnostics_spm.pdf",
    )

    file_ext = f"p3m1_nsteps{nsteps_p3m1}"
    aiDrift = StandardTimeStepping.read(wd + file_ext + f"_ts_p3m.h5").aiDrift
    plot_timestepping_diagnostics(
        log_path=simdir + "p3m1_timestep_log.txt",
        aiDrift=aiDrift,
        TimeStepDistribution=tsd_p3m1,
        nsteps=nsteps_p3m1,
        save_path=outdir + "time_step_diagnostics_p3m1.pdf",
    )

    file_ext = f"p3m2_nsteps{nsteps_p3m2}"
    aiDrift = StandardTimeStepping.read(wd + file_ext + f"_ts_p3m.h5").aiDrift
    plot_timestepping_diagnostics(
        log_path=simdir + "p3m2_timestep_log.txt",
        aiDrift=aiDrift,
        TimeStepDistribution=tsd_p3m2,
        nsteps=nsteps_p3m2,
        save_path=outdir + "time_step_diagnostics_p3m2.pdf",
    )

    file_ext = f"p3m3_nsteps{nsteps_p3m3}"
    aiDrift = StandardTimeStepping.read(wd + file_ext + f"_ts_p3m.h5").aiDrift
    plot_timestepping_diagnostics(
        log_path=simdir + "p3m3_timestep_log.txt",
        aiDrift=aiDrift,
        TimeStepDistribution=tsd_p3m3,
        nsteps=nsteps_p3m3,
        save_path=outdir + "time_step_diagnostics_p3m3.pdf",
    )

    logger.info(f"Plotting power spectra...")
    INDENT()
    fig, ax = plt.subplots(figsize=(7, 4))

    ax.set_xscale("log")
    k = G.k_modes
    kmin, kmax = k.min(), k.max()
    logger.info(f"kmin = {kmin:.2e}, kmax = {kmax:.2e}")
    log_pad = 0.02
    log_k_min = np.log10(kmin)
    log_k_max = np.log10(kmax)
    log_range = log_k_max - log_k_min
    xlim_min = 10 ** (log_k_min - log_pad * log_range)
    xlim_max = 10 ** (log_k_max + log_pad * log_range)

    plt.xlim(xlim_min, xlim_max)
    # ax.set_ylim([0.2, 1.8])
    dark_grey_bnd = 0.01
    medium_grey_bnd = 0.05
    light_grey_bnd = 0.1

    ax.plot([kmin, kmax], [1, 1], color="black", linestyle="-", label=label_ref)
    if lpt:
        ax.plot(k, Pk_LPT / Pk_ref, label="2LPT", linestyle="--")
    fields_to_plot = [
        ("PMref", Pk_PMref),
        ("PM1", Pk_PM1),
        ("PM2", Pk_PM2),
        ("sPM", Pk_sPM),
        ("P3M1", Pk_P3M1),
        ("P3M2", Pk_P3M2),
        ("P3M3", Pk_P3M3),
    ]
    if cola:
        fields_to_plot.append(("COLA", Pk_COLA))
    # Remove the field corresponding to the reference:
    fields_to_plot = [(field_name, Pk) for field_name, Pk in fields_to_plot if field_name != ref]
    for field_name, Pk in fields_to_plot:
        label = f"{field_name}, $n_\\mathrm{{steps}}={eval(f'nsteps_{field_name.lower()}')}$"
        if field_name in ["PMref", "P3M1"]:
            linestyle = "-"
            zorder = 1
        else:
            linestyle = "--"
            zorder = 2
        ax.plot(k, Pk / Pk_ref, label=label, linestyle=linestyle)

    ax.axhspan(1 - dark_grey_bnd, 1 + dark_grey_bnd, color="grey", alpha=0.3)
    ax.axhspan(1 - medium_grey_bnd, 1 + medium_grey_bnd, color="grey", alpha=0.2)
    ax.axhspan(1 - light_grey_bnd, 1 + light_grey_bnd, color="grey", alpha=0.1)

    for i in range(1, len(k)):
        ax.axvline(k[i], color="black", linestyle=":", linewidth=1, alpha=0.1, zorder=0)
    ax.yaxis.set_major_locator(plt.MaxNLocator(6))
    ax.yaxis.get_major_ticks()[0].label1.set_visible(False)
    ax.set_xlabel("$k$ [$h/\\mathrm{Mpc}$]", fontsize=fs)
    ax.set_ylabel("$P(k)/P_\\mathrm{ref}(k)$", fontsize=fs)
    ax.tick_params(which="both", direction="in")
    ax.tick_params(axis="both", which="major", labelsize=fs)
    ax.tick_params(axis="both", which="minor", labelsize=fs)

    # Characteristic vertical reference scales
    nyquist = np.pi * N / L
    nyquist_PM = np.pi * Npm / L
    epsilon = 0.03 * L / Np
    particle_length = 2 * epsilon
    xs = 1.25 * L / Npm
    xr = 4.5 * xs
    particle_wavenumber = 2 * np.pi / particle_length  # Too large to be shown
    xs_inv = 2 * np.pi / xs
    xr_inv = 2 * np.pi / xr
    if nyquist <= xlim_max:
        line1 = ax.axvline(
            x=nyquist, color="black", linestyle="--", lw=2, label="Nyquist (density grid)", zorder=0
        )
    if nyquist_PM <= xlim_max:
        line2 = ax.axvline(
            x=nyquist_PM, color="black", linestyle="-", lw=1, label="Nyquist (PM grid)", zorder=0
        )
    if xs_inv <= xlim_max:
        line3 = ax.axvline(
            x=xs_inv, color="gray", linestyle="-.", lw=2, label=r"Split wavenumber $x_s$", zorder=0
        )
    if xr_inv <= xlim_max:
        line4 = ax.axvline(
            x=xr_inv, color="gray", linestyle=":", lw=2, label=r"Short-range reach $x_r$", zorder=0
        )

    logger.info(f"Nyquist (density grid): {nyquist:.2f} h/Mpc")
    logger.info(f"Nyquist (PM grid): {nyquist_PM:.2f} h/Mpc")
    logger.info(f"Particle wavenumber: {particle_wavenumber:.2f} h/Mpc")
    logger.info(f"Split wavenumber: {xs_inv:.2f} h/Mpc")
    logger.info(f"Short-range reach: {xr_inv:.2f} h/Mpc")

    handles, labels = plt.gca().get_legend_handles_labels()
    # empty_patch = mpatches.Patch(color="none", label="")
    # handles = [empty_patch, *handles]
    # labels = ["", *labels]
    plt.legend(
        handles,
        labels,
        loc="upper center",
        ncol=2,
        bbox_to_anchor=(0.5, -0.2),
        fontsize=fs,
        frameon=False,
    )
    fig.savefig(
        outdir + "power_spectra.png",
        bbox_inches="tight",
        dpi=300,
        transparent=True,
    )
    fig.savefig(
        outdir + "power_spectra.pdf",
        bbox_inches="tight",
        dpi=300,
    )
    plt.close(fig)
    UNINDENT()
    logger.info(
        f"Power spectra plotted and saved to {outdir}power_spectra.png and power_spectra.pdf"
    )
    UNINDENT()
    logger.info("Computing and plotting power spectra done.")
    logger.info("All plots completed successfully.")
