#!/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"

"""
Perform time step convergence tests with custom time step distributions,
using COLA with different force evaluations (PM, sPM, P3M).
"""

# pyright: reportWildcardImportFromLibrary=false
from os.path import isfile
from pathlib import Path
import numpy as np

from pysbmy.power import PowerSpectrum
from pysbmy.fft import FourierGrid
from pysbmy.timestepping import StandardTimeStepping, P3MTimeStepping

from wip3m import *
from wip3m.tools import none_bool_str, get_k_max, generate_sim_params, generate_white_noise_Field, run_simulation
from wip3m.params import (
    params_planck_kmax_missing,
    cosmo_small_to_full_dict,
    z2a,
    BASELINE_SEEDPHASE,
)
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="Run convergence tests towards implementing P3M gravity.")

parser.add_argument("--run_id", type=str, help="Simulation run identifier.")
parser.add_argument("--L", type=float, default=32.0, help="Box size in Mpc/h.")
parser.add_argument("--N", type=int, default=128, help="Density grid size per dimension.")
parser.add_argument("--Np", type=int, default=32, help="Number of particles per dimension.")
parser.add_argument("--Npm", type=int, default=64, help="PM mesh resolution per dimension.")
parser.add_argument(
    "--n_Tiles", type=int, default=8, help="Number of tiles per dimension for short-range PP."
)
parser.add_argument("--z_i", type=float, default=199.0, help="Initial redshift.")
parser.add_argument("--z_f", type=float, default=19.0, help="Final redshift.")
parser.add_argument(
    "--plot_fields",
    type=none_bool_str,
    default=False,
    help="Whether to plot the fields or not.",
)
parser.add_argument(
    "--lpt",
    type=none_bool_str,
    default=False,
    help="Whether to compute and plot the LPT power spectrum.",
)
parser.add_argument(
    "--scale_limiter",
    type=none_bool_str,
    default=False,
    help="Which limiter should be scaled by the scaling arguments? ",
)
parser.add_argument(
    "--only_plot",
    type=none_bool_str,
    default=False,
    help="Do not re-run simulations. Warning: it is up to the user to ensure that the required files are present.",
)

# Timestep settings per method
for scheme in ["pmref", "pm1", "pm2", "spm", "p3m1", "p3m2", "p3m3"]:
    parser.add_argument(
        f"--scaling_{scheme}",
        type=float,
        default={
            "pmref": 1.,
            "pm1": 0.95,
            "pm2": 0.9,
            "spm": 1.,
            "p3m1": 1.,
            "p3m2": 0.95,
            "p3m3": 0.9,
        }[scheme],
        help=f"Relax the step limiter by this factor for {scheme.upper()}.",
    )

args = parser.parse_args()

if __name__ == "__main__":
    try:
        go_beyond_Nyquist_ss = True  # for the summary statistics
        force = False
        force_hard = False
        only_plot = args.only_plot

        run_id = args.run_id
        L = args.L
        N = args.N
        Np = args.Np
        Npm = args.Npm
        n_Tiles = args.n_Tiles
        scaling_pmref = args.scaling_pmref
        scaling_pm1 = args.scaling_pm1
        scaling_pm2 = args.scaling_pm2
        scaling_spm = args.scaling_spm
        scaling_p3m1 = args.scaling_p3m1
        scaling_p3m2 = args.scaling_p3m2
        scaling_p3m3 = args.scaling_p3m3
        name_scaled = args.scale_limiter
        if name_scaled is not False and name_scaled not in [
            "fac_dyn_custom",
            "da_max_early_custom",
            "fac_H_custom",
            "fac_bend",
            "sub_bend1",
            "sub_bend2",
            "da_max_late",
        ]:
            raise ValueError(
                f"Invalid name_scaled: {name_scaled}. Must be one of the limiter names or False."
            )
        RedshiftLPT = args.z_i
        RedshiftFCs = args.z_f
        logger.info("Running convergence tests...")
        logger.info(f"Run ID: {run_id}")
        logger.info(f"Box size: {L} Mpc/h")
        logger.info(f"Density grid size: {N}^3")
        logger.info(f"Particles per dimension: {Np}^3")
        logger.info(f"PM grid size: {Npm}^3")
        logger.info(f"Number of tiles: {n_Tiles}^3")
        logger.info(f"Limiter to scale: {name_scaled if name_scaled else 'None'}")
        if name_scaled:
            logger.info(f"Limiter rescaling for PMREF: {scaling_pmref}")
            logger.info(f"Limiter rescaling for PM1: {scaling_pm1}")
            logger.info(f"Limiter rescaling for PM2: {scaling_pm2}")
            logger.info(f"Limiter rescaling for sPM: {scaling_spm}")
            logger.info(f"Limiter rescaling for P3M1: {scaling_p3m1}")
            logger.info(f"Limiter rescaling for P3M2: {scaling_p3m2}")
            logger.info(f"Limiter rescaling for P3M3: {scaling_p3m3}")
        INDENT()

        lpt = args.lpt
        plot_fields = args.plot_fields

        corner = -L / 2.0
        ai = z2a(RedshiftLPT)
        af = z2a(RedshiftFCs)
        k_max = get_k_max(L, N)  # k_max in h/Mpc
        cosmo = params_planck_kmax_missing.copy()
        cosmo["k_max"] = k_max

        wd = workdir + run_id + "/"
        simdir = output_path + run_id + "/"
        logdir = simdir + "logs/"
        outdir = simdir + "plots/"
        if force_hard:
            import shutil

            if Path(simdir).exists():
                shutil.rmtree(simdir)
            if Path(wd).exists():
                shutil.rmtree(wd)
        Path(wd).mkdir(parents=True, exist_ok=True)
        Path(logdir).mkdir(parents=True, exist_ok=True)
        Path(outdir).mkdir(parents=True, exist_ok=True)

        input_white_noise_file = simdir + "input_white_noise.h5"
        input_seed_phase_file = simdir + "seed"
        ICs_path = simdir + "initial_density.h5"
        simpath = simdir

        # Path to the input matter power spectrum (generated later)
        input_power_file = simdir + "input_power.h5"

        sim_params = {
            "L": L,
            "N": N,
            "Np": Np,
            "Npm": Npm,
            "n_Tiles": n_Tiles,
            "RedshiftLPT": RedshiftLPT,
            "RedshiftFCs": RedshiftFCs,
            "name_scaled": name_scaled,
            "scaling_pmref": scaling_pmref,
            "scaling_pm1": scaling_pm1,
            "scaling_pm2": scaling_pm2,
            "scaling_spm": scaling_spm,
            "scaling_p3m1": scaling_p3m1,
            "scaling_p3m2": scaling_p3m2,
            "scaling_p3m3": scaling_p3m3,
        }
        with open(wd + "sim_params.txt", "w") as f:
            f.write(f"{sim_params}\n")

        logger.info("Setting up simulation parameters...")

        common_params = {
            "Np": Np,
            "N": N,
            "L": L,
            "corner0": corner,
            "corner1": corner,
            "corner2": corner,
            "h": cosmo["h"],
            "Omega_m": cosmo["Omega_m"],
            "Omega_b": cosmo["Omega_b"],
            "n_s": cosmo["n_s"],
            "sigma8": cosmo["sigma8"],
        }

        lpt_params = common_params.copy()
        lpt_params["method"] = "lpt"
        lpt_params["InputPowerSpectrum"] = input_power_file
        lpt_params["ICsMode"] = 1
        lpt_params["InputWhiteNoise"] = input_white_noise_file

        fac_dyn_custom = DEFAULT_FAC_DYN_CUSTOM
        da_early = DEFAULT_DA_MAX_EARLY_CUSTOM
        fac_H_custom = DEFAULT_FAC_H_CUSTOM
        fac_bend = DEFAULT_FAC_BEND
        sub_bend1 = DEFAULT_SUB_BEND1_COLA
        sub_bend2 = DEFAULT_SUB_BEND2_COLA
        da_late = DEFAULT_DA_MAX_LATE_CUSTOM

        limiters = {
            "fac_dyn_custom": fac_dyn_custom,
            "da_max_early_custom": da_early,
            "fac_H_custom": fac_H_custom,
            "fac_bend": fac_bend,
            "sub_bend1": sub_bend1,
            "sub_bend2": sub_bend2,
            "da_max_late": da_late,
        }

        common_params_num = common_params.copy()
        common_params_num["ai"] = ai
        common_params_num["af"] = af
        common_params_num["RedshiftLPT"] = RedshiftLPT
        common_params_num["RedshiftFCs"] = RedshiftFCs
        common_params_num["Npm"] = Npm
        common_params_num["RunForceDiagnostic"] = False
        common_params_num["nBinsForceDiagnostic"] = 20
        common_params_num["nPairsForceDiagnostic"] = 3
        common_params_num["maxTrialsForceDiagnostic"] = int(1e8)
        common_params_num["OutputForceDiagnostic"] = simdir + "force_diagnostic.txt"
        common_params_num["TimeStepDistribution"] = 3  # Custom time step distribution
        common_params_num["cosmo_dict"] = cosmo
        common_params_num["fac_dyn_custom"] = limiters["fac_dyn_custom"]
        common_params_num["da_max_early_custom"] = limiters["da_max_early_custom"]
        common_params_num["fac_H_custom"] = limiters["fac_H_custom"]
        common_params_num["fac_bend"] = limiters["fac_bend"]
        common_params_num["sub_bend1"] = limiters["sub_bend1"]
        common_params_num["sub_bend2"] = limiters["sub_bend2"]
        common_params_num["da_max_late_custom"] = limiters["da_max_late"]

        pmref_params = common_params_num.copy()
        pmref_params["method"] = "cola"  # COLA+PM
        pmref_params["relax"] = scaling_pmref  # plot unscaled values during the simulation
        if name_scaled:
            pmref_params[name_scaled] /= scaling_pmref

        pm1_params = common_params_num.copy()
        pm1_params["method"] = "cola"  # COLA+PM
        pm1_params["relax"] = scaling_pm1
        if name_scaled:
            pm1_params[name_scaled] /= scaling_pm1

        pm2_params = common_params_num.copy()
        pm2_params["method"] = "cola"  # COLA+PM
        pm2_params["relax"] = scaling_pm2
        if name_scaled:
            pm2_params[name_scaled] /= scaling_pm2

        spm_params = common_params_num.copy()
        spm_params["method"] = "spm"
        spm_params["EvolutionMode"] = 6  # COLA+sPM
        spm_params["n_Tiles"] = n_Tiles
        spm_params["PrintOutputTimestepsLog"] = True
        spm_params["OutputTimestepsLog"] = simdir + "colaspm_timestep_log.txt"
        spm_params["relax"] = scaling_spm
        if name_scaled:
            spm_params[name_scaled] /= scaling_spm

        p3m1_params = common_params_num.copy()
        p3m1_params["method"] = "p3m"
        p3m1_params["EvolutionMode"] = 7  # COLA+P3M
        p3m1_params["n_Tiles"] = n_Tiles
        p3m1_params["PrintOutputTimestepsLog"] = True
        p3m1_params["OutputTimestepsLog"] = simdir + "colap3m1_timestep_log.txt"
        p3m1_params["relax"] = scaling_p3m1
        if name_scaled:
            p3m1_params[name_scaled] /= scaling_p3m1

        p3m2_params = common_params_num.copy()
        p3m2_params["method"] = "p3m"
        p3m2_params["EvolutionMode"] = 7  # COLA+P3M
        p3m2_params["n_Tiles"] = n_Tiles
        p3m2_params["PrintOutputTimestepsLog"] = True
        p3m2_params["OutputTimestepsLog"] = simdir + "colap3m2_timestep_log.txt"
        p3m2_params["relax"] = scaling_p3m2
        if name_scaled:
            p3m2_params[name_scaled] /= scaling_p3m2

        p3m3_params = common_params_num.copy()
        p3m3_params["method"] = "p3m"
        p3m3_params["EvolutionMode"] = 7  # COLA+P3M
        p3m3_params["n_Tiles"] = n_Tiles
        p3m3_params["PrintOutputTimestepsLog"] = True
        p3m3_params["OutputTimestepsLog"] = simdir + "colap3m3_timestep_log.txt"
        p3m3_params["relax"] = scaling_p3m3
        if name_scaled:
            p3m3_params[name_scaled] /= scaling_p3m3

        logger.info("Setting up simulation parameters done.")

        if not only_plot:
            logger.info("Generating simulation parameters...")
            INDENT()
            reset_plotting()  # Default style for Simbelmynë
            generate_sim_params(lpt_params, ICs_path, wd, simdir, None, force)

            all_sim_params = [
                ("pmref", pmref_params),
                ("pm1", pm1_params),
                ("pm2", pm2_params),
                ("spm", spm_params),
                ("p3m1", p3m1_params),
                ("p3m2", p3m2_params),
                ("p3m3", p3m3_params),
            ]
            for name, parameters in all_sim_params:
                logger.info(
                    f"Generating parameters for {name.upper()} with relax = {parameters['relax']}..."
                )
                file_ext = f"{name}_relax{str(parameters['relax']).replace('.', '_')}"
                generate_sim_params(parameters, ICs_path, wd, simdir, file_ext, force)
                logger.info(f"Generating parameters for {name.upper()} done.")

            UNINDENT()
            logger.info("Generating simulation parameters done.")

            setup_plotting()  # Reset plotting style for this project

            logger.info("Generating white noise field...")
            generate_white_noise_Field(
                L=L,
                size=N,
                corner=corner,
                seedphase=BASELINE_SEEDPHASE,
                fname_whitenoise=input_white_noise_file,
                seedname_whitenoise=input_seed_phase_file,
                force_phase=force,
            )
            logger.info("Generating white noise field done.")

            # If cosmo["WhichSpectrum"] == "class", then classy is required.
            if not isfile(input_power_file) or force:
                logger.info("Generating input power spectrum...")
                Pk = PowerSpectrum(L, L, L, N, N, N, cosmo_small_to_full_dict(cosmo))
                Pk.write(input_power_file)
                logger.info("Generating input power spectrum done.")

            logger.info("Generating Fourier grid...")
            # k grid used to compute the final overdensity power spectrum
            Pinit = 100
            trim_threshold = 100  # Merge bins until this minimum number of modes per bin is reached
            log_kmin = np.log10(2 * np.pi / (np.sqrt(3) * L))  # Minimum non-zero k in h/Mpc
            if go_beyond_Nyquist_ss:
                k_max_ss = get_k_max(L, N)
            else:
                k_max_ss = get_k_max(L, N) / np.sqrt(3)  # 1D Nyquist frequency
            Pbins_left_bnds = np.logspace(log_kmin, np.log10(k_max_ss), Pinit + 1, dtype=np.float32)
            Pbins_left_bnds = Pbins_left_bnds[:-1]
            input_ss_file = simdir + "input_ss_k_grid.h5"
            Gk = FourierGrid(
                L,
                L,
                L,
                N,
                N,
                N,
                k_modes=Pbins_left_bnds,
                kmax=k_max_ss,
                trim_bins=True,
                trim_threshold=trim_threshold,
            )
            Gk.write(input_ss_file)
            logger.info("Generating Fourier grid done.")

            logger.info("Running simulations...")
            INDENT()
            logger.info("Running LPT simulation...")
            run_simulation("lpt", lpt_params, wd, logdir)
            logger.info("Running LPT simulation done.")

            for name, parameters in all_sim_params:
                logger.info(f"Running {name.upper()} simulation...")
                run_simulation(name, parameters, wd, logdir)
                logger.info(f"Running {name.upper()} simulation done.")
            UNINDENT()
            logger.info("All simulations done.")

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

        nsteps_dict = {}
        for name, parameters in all_sim_params:
            TimeStepDistribution = parameters["TimeStepDistribution"]
            file_ext = f"{name}_relax{str(parameters['relax']).replace('.', '_')}"
            method = parameters["method"]
            TSpath = wd + file_ext + f"_ts_{method}.h5" if file_ext else wd + f"ts_{method}.h5"
            if TimeStepDistribution in [0, 1, 2]:
                TS = StandardTimeStepping.read(TSpath)
                aiDrift = TS.aiDrift
                nsteps = TS.nsteps
            elif TimeStepDistribution == 3:
                TS = P3MTimeStepping.read(TSpath)
                aiDrift = TS.aiDrift
                nsteps = TS.nsteps
            else:
                raise ValueError(f"Invalid TimeStepDistribution: {TimeStepDistribution}")
            nsteps_dict[name] = nsteps
            if name.startswith("p3m") or name.startswith("spm"):
                plot_timestepping_diagnostics(
                    log_path=parameters["OutputTimestepsLog"],
                    aiDrift=aiDrift,
                    TimeStepDistribution=TimeStepDistribution,
                    nsteps=nsteps,
                    ymin=2e-3,
                    exact=False,
                    save_path=outdir + f"time_step_diagnostics_{name}.pdf",
                )
                plot_custom_timestepping_diagnostics(
                    log_path=parameters["OutputTimestepsLog"],
                    aiDrift=aiDrift,
                    TimeStepDistribution=TimeStepDistribution,
                    nsteps=nsteps,
                    ymin=2e-3,
                    fac_hubble=fac_H_custom,
                    fac_bend=fac_bend,
                    da_max_early=da_early,
                    da_max_late=DEFAULT_DA_MAX_LATE_CUSTOM,
                    save_path=outdir + f"time_step_diagnostics_custom_{name}.pdf",
                )

        UNINDENT()
        logger.info("Plotting time step diagnostics done.")

        logger.info("Computing and plotting power spectra...")
        INDENT()

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

        logger.info("Loading Fourier grid...")
        G = read_FourierGrid(input_ss_file)
        logger.info("Loading Fourier grid done.")

        k = G.k_modes
        AliasingCorr = False
        slice_ijk = (N // 2, slice(None), slice(None))  # for plot_fields

        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_relax{str(pmref_params['relax']).replace('.', '_')}_final_density_cola.h5")
        if plot_fields:
            delta_pmref = DELTA.data[slice_ijk]
        Pk_PMref, _ = get_autocorrelation(DELTA, G, AliasingCorr)

        DELTA = read_field(simdir + f"pm1_relax{str(pm1_params['relax']).replace('.', '_')}_final_density_cola.h5")
        if plot_fields:
            delta_pm1 = DELTA.data[slice_ijk]
        Pk_PM1, _ = get_autocorrelation(DELTA, G, AliasingCorr)

        DELTA = read_field(simdir + f"pm2_relax{str(pm2_params['relax']).replace('.', '_')}_final_density_cola.h5")
        if plot_fields:
            delta_pm2 = DELTA.data[slice_ijk]
        Pk_PM2, _ = get_autocorrelation(DELTA, G, AliasingCorr)

        DELTA = read_field(simdir + f"spm_relax{str(spm_params['relax']).replace('.', '_')}_final_density_spm.h5")
        if plot_fields:
            delta_spm = DELTA.data[slice_ijk]
        Pk_sPM, _ = get_autocorrelation(DELTA, G, AliasingCorr)

        DELTA = read_field(simdir + f"p3m1_relax{str(p3m1_params['relax']).replace('.', '_')}_final_density_p3m.h5")
        if plot_fields:
            delta_p3m1 = DELTA.data[slice_ijk]
        Pk_P3M1, _ = get_autocorrelation(DELTA, G, AliasingCorr)

        DELTA = read_field(simdir + f"p3m2_relax{str(p3m2_params['relax']).replace('.', '_')}_final_density_p3m.h5")
        if plot_fields:
            delta_p3m2 = DELTA.data[slice_ijk]
        Pk_P3M2, _ = get_autocorrelation(DELTA, G, AliasingCorr)

        DELTA = read_field(simdir + f"p3m3_relax{str(p3m3_params['relax']).replace('.', '_')}_final_density_p3m.h5")
        if plot_fields:
            delta_p3m3 = DELTA.data[slice_ijk]
        Pk_P3M3, _ = get_autocorrelation(DELTA, G, AliasingCorr)
        logger.info("Computing power spectra done.")

        ref = "P3M1"

        scaling_reference = scaling_p3m1
        Pk_ref = Pk_P3M1
        label_ref = f"P3M, scaling={scaling_reference}, ns={nsteps_dict['p3m1']}"

        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),
        ]
        # 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}, scaling={eval(f'scaling_{field_name.lower()}')}, ns={nsteps_dict[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)
        ax.axhspan(1 - dark_grey_bnd, 1 + dark_grey_bnd, color="grey", alpha=0.3)
        ax.axhline(1 - medium_grey_bnd, color="grey", linestyle="-", linewidth=1)
        ax.axhline(1 + medium_grey_bnd, color="grey", linestyle="-", linewidth=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("Plotting power spectra done.")

        logger.info(f"Plotting P3M 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)
        fields_to_plot = [
            ("P3M1", Pk_P3M1),
            ("P3M2", Pk_P3M2),
            ("P3M3", Pk_P3M3),
        ]
        # 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}, scaling={eval(f'scaling_{field_name.lower()}')}, ns={nsteps_dict[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.axhline(1 - medium_grey_bnd, color="grey", linestyle="-", linewidth=1)
        ax.axhline(1 + medium_grey_bnd, color="grey", linestyle="-", linewidth=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()
        plt.legend(
            handles,
            labels,
            loc="upper center",
            ncol=2,
            bbox_to_anchor=(0.5, -0.2),
            fontsize=fs,
            frameon=False,
        )
        fig.savefig(
            outdir + "p3m_power_spectra.png",
            bbox_inches="tight",
            dpi=300,
            transparent=True,
        )
        fig.savefig(
            outdir + "p3m_power_spectra.pdf",
            bbox_inches="tight",
            dpi=300,
        )
        plt.close(fig)
        UNINDENT()
        logger.info(
            f"P3M power spectra plotted and saved to {outdir}p3m_power_spectra.png and power_spectra.pdf"
        )
        logger.info("Pplotting P3M power spectra done.")
        UNINDENT()

        logger.info("Computing and plotting power spectra done.")

        if plot_fields:
            logger.info("Reading lpt field...")
            delta_lpt = read_field(simdir + "lpt_density.h5").data[slice_ijk]
            logger.info("Reading lpt field done.")

            diff_pm1_pmref = delta_pm1 - delta_pmref
            diff_pm2_pmref = delta_pm2 - delta_pmref
            diff_p3m1_pmref = delta_p3m1 - delta_pmref
            diff_p3m2_p3m1 = delta_p3m2 - delta_p3m1
            diff_p3m3_p3m1 = delta_p3m3 - delta_p3m1
            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_p3m1",
                "diff_p3m3_p3m1",
            ]  # fields to plot
            ncols = 6  # Max panels per row

            figname = "_".join(fields)
            slices_dict = {
                "lpt": delta_lpt,
                "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_p3m1": diff_p3m2_p3m1,
                "diff_p3m3_p3m1": diff_p3m3_p3m1,
                "diff_p3m1_spm": diff_p3m1_spm,
            }
            titles_dict = {
                "lpt": "LPT",
                "pmref": f"PMref relax={scaling_pmref}",
                "pm1": f"PM1 relax={scaling_pm1}",
                "pm2": f"PM2 relax={scaling_pm2}",
                "spm": f"sPM relax={scaling_spm}",
                "p3m1": f"P3M1 relax={scaling_p3m1}",
                "p3m2": f"P3M2 relax={scaling_p3m2}",
                "p3m3": f"P3M3 relax={scaling_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_p3m1": r"$\delta_{\rm P3M2}-\delta_{\rm P3M1}$",
                "diff_p3m3_p3m1": r"$\delta_{\rm P3M3}-\delta_{\rm P3M1}$",
                "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"):
                    vmin = -np.log(1 + np.abs(np.min(data)))
                    vmax = np.log10(1 + np.abs(np.max(data)))
                    if vmax > 0 > vmin:
                        norm = TwoSlopeNorm(vmin=vmin, vcenter=0, vmax=vmax)
                        im = ax.imshow(
                            np.sign(data) * np.log(1 + np.abs(data)), cmap="RdBu_r", norm=norm
                        )
                    else:
                        im = ax.imshow(np.log10(1 + np.abs(data)), cmap="RdBu_r")
                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("All plots completed successfully.")

    except OSError as e:
        logger.error("Directory or file access error: %s", str(e))
        raise
    except Exception as e:
        logger.critical("An unexpected error occurred: %s", str(e))
        raise RuntimeError("Failed.") from e
    finally:
        UNINDENT()
        logger.info("Running convergence tests done.")
