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

"""
Compute time step limiters from baseline (linear / logarithmic) time
step distribution.
"""

# 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

from wip3m import *
from wip3m.tools import 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=128.0, help="Box size in Mpc/h.")
parser.add_argument("--N", type=int, default=256, help="Density grid size per dimension.")
parser.add_argument("--Np", type=int, default=128, help="Number of particles per dimension.")
parser.add_argument("--Npm", type=int, default=256, help="PM mesh resolution per dimension.")
parser.add_argument(
    "--n_Tiles", type=int, default=32, 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=0.0, help="Final redshift.")

# Timestep settings per method
for scheme in ["pmref", "pm1", "pm2", "p3m1", "p3m2", "p3m3"]:
    parser.add_argument(
        f"--nsteps_{scheme}",
        type=int,
        default={
            "pmref": 300,
            "pm1": 200,
            "pm2": 100,
            "p3m1": 500,
            "p3m2": 400,
            "p3m3": 300,
        }[scheme],
        help=f"Number of timesteps for {scheme.upper()}.",
    )
    parser.add_argument(
        f"--tsd_{scheme}",
        type=int,
        default=1,
        help=f"TimeStepDistribution ID for {scheme.upper()}.",
    )
parser.add_argument(
    "--nsteps_p3m4",
    type=int,
    help="Number of timesteps for P3M4.",
)
parser.add_argument(
    "--tsd_p3m4",
    type=int,
    default=1,
    help="TimeStepDistribution ID for P3M4.",
)
parser.add_argument(
    "--nsteps_p3m5",
    type=int,
    help="Number of timesteps for P3M5.",
)
parser.add_argument(
    "--tsd_p3m5",
    type=int,
    default=1,
    help="TimeStepDistribution ID for P3M5.",
)

args = parser.parse_args()

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

        run_id = args.run_id
        L = args.L
        N = args.N
        Np = args.Np
        Npm = args.Npm
        n_Tiles = args.n_Tiles
        nsteps_pmref = args.nsteps_pmref
        nsteps_pm1 = args.nsteps_pm1
        nsteps_pm2 = args.nsteps_pm2
        nsteps_p3m1 = args.nsteps_p3m1
        nsteps_p3m2 = args.nsteps_p3m2
        nsteps_p3m3 = args.nsteps_p3m3
        tsd_pmref = args.tsd_pmref
        tsd_pm1 = args.tsd_pm1
        tsd_pm2 = args.tsd_pm2
        tsd_p3m1 = args.tsd_p3m1
        tsd_p3m2 = args.tsd_p3m2
        tsd_p3m3 = args.tsd_p3m3
        if args.nsteps_p3m4 is None:
            p3m4 = False
        else:
            p3m4 = True
            nsteps_p3m4 = args.nsteps_p3m4
            tsd_p3m4 = args.tsd_p3m4
        if args.nsteps_p3m5 is None:
            p3m5 = False
        else:
            p3m5 = True
            nsteps_p3m5 = args.nsteps_p3m5
            tsd_p3m5 = args.tsd_p3m5
        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"Number of timesteps for PMREF: {nsteps_pmref}")
        logger.info(f"Number of timesteps for PM1: {nsteps_pm1}")
        logger.info(f"Number of timesteps for PM2: {nsteps_pm2}")
        logger.info(f"Number of timesteps for P3M1: {nsteps_p3m1}")
        logger.info(f"Number of timesteps for P3M2: {nsteps_p3m2}")
        logger.info(f"Number of timesteps for P3M3: {nsteps_p3m3}")
        if p3m4:
            logger.info(f"Number of timesteps for P3M4: {nsteps_p3m4}")
        if p3m5:
            logger.info(f"Number of timesteps for P3M5: {nsteps_p3m5}")
        INDENT()

        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/"
        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)

        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,
            "nsteps_pmref": nsteps_pmref,
            "nsteps_pm1": nsteps_pm1,
            "nsteps_pm2": nsteps_pm2,
            "nsteps_p3m1": nsteps_p3m1,
            "nsteps_p3m2": nsteps_p3m2,
            "nsteps_p3m3": nsteps_p3m3,
            "ai": ai,
            "af": af,
            "cosmo": cosmo,
            "tsd_pmref": tsd_pmref,
            "tsd_pm1": tsd_pm1,
            "tsd_pm2": tsd_pm2,
            "tsd_p3m1": tsd_p3m1,
            "tsd_p3m2": tsd_p3m2,
            "tsd_p3m3": tsd_p3m3,
        }
        if p3m4:
            sim_params["nsteps_p3m4"] = nsteps_p3m4
            sim_params["tsd_p3m4"] = tsd_p3m4
        if p3m5:
            sim_params["nsteps_p3m5"] = nsteps_p3m5
            sim_params["tsd_p3m5"] = tsd_p3m5
        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

        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"

        pmref_params = common_params_num.copy()
        pmref_params["method"] = "pm"
        pmref_params["TimeStepDistribution"] = tsd_pmref
        pmref_params["nsteps"] = nsteps_pmref

        pm1_params = common_params_num.copy()
        pm1_params["method"] = "pm"
        pm1_params["TimeStepDistribution"] = tsd_pm1
        pm1_params["nsteps"] = nsteps_pm1

        pm2_params = common_params_num.copy()
        pm2_params["method"] = "pm"
        pm2_params["TimeStepDistribution"] = tsd_pm2
        pm2_params["nsteps"] = nsteps_pm2

        p3m1_params = common_params_num.copy()
        p3m1_params["method"] = "p3m"
        p3m1_params["EvolutionMode"] = 4
        p3m1_params["TimeStepDistribution"] = tsd_p3m1
        p3m1_params["nsteps"] = nsteps_p3m1
        p3m1_params["n_Tiles"] = n_Tiles
        p3m1_params["PrintOutputTimestepsLog"] = True
        p3m1_params["OutputTimestepsLog"] = simdir + "p3m1_timestep_log.txt"

        p3m2_params = common_params_num.copy()
        p3m2_params["method"] = "p3m"
        p3m2_params["EvolutionMode"] = 4
        p3m2_params["TimeStepDistribution"] = tsd_p3m2
        p3m2_params["nsteps"] = nsteps_p3m2
        p3m2_params["n_Tiles"] = n_Tiles
        p3m2_params["PrintOutputTimestepsLog"] = True
        p3m2_params["OutputTimestepsLog"] = simdir + "p3m2_timestep_log.txt"

        p3m3_params = common_params_num.copy()
        p3m3_params["method"] = "p3m"
        p3m3_params["EvolutionMode"] = 4
        p3m3_params["TimeStepDistribution"] = tsd_p3m3
        p3m3_params["nsteps"] = nsteps_p3m3
        p3m3_params["n_Tiles"] = n_Tiles
        p3m3_params["PrintOutputTimestepsLog"] = True
        p3m3_params["OutputTimestepsLog"] = simdir + "p3m3_timestep_log.txt"

        if p3m4:
            p3m4_params = common_params_num.copy()
            p3m4_params["method"] = "p3m"
            p3m4_params["EvolutionMode"] = 4
            p3m4_params["TimeStepDistribution"] = tsd_p3m4
            p3m4_params["nsteps"] = nsteps_p3m4
            p3m4_params["n_Tiles"] = n_Tiles
            p3m4_params["PrintOutputTimestepsLog"] = True
            p3m4_params["OutputTimestepsLog"] = simdir + "p3m4_timestep_log.txt"

        if p3m5:
            p3m5_params = common_params_num.copy()
            p3m5_params["method"] = "p3m"
            p3m5_params["EvolutionMode"] = 4
            p3m5_params["TimeStepDistribution"] = tsd_p3m5
            p3m5_params["nsteps"] = nsteps_p3m5
            p3m5_params["n_Tiles"] = n_Tiles
            p3m5_params["PrintOutputTimestepsLog"] = True
            p3m5_params["OutputTimestepsLog"] = simdir + "p3m5_timestep_log.txt"

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

        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),
            ("p3m1", p3m1_params),
            ("p3m2", p3m2_params),
            ("p3m3", p3m3_params),
        ]
        if p3m4:
            all_sim_params.append(("p3m4", p3m4_params))
        if p3m5:
            all_sim_params.append(("p3m5", p3m5_params))
        for name, parameters in all_sim_params:
            logger.info(
                f"Generating parameters for {name.upper()} with nsteps = {parameters['nsteps']}..."
            )
            file_ext = f"{name}_nsteps{parameters['nsteps']}"
            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()

        for name, parameters in all_sim_params:
            if name.startswith("p3m") or name.startswith("spm"):
                file_ext = f"{name}_nsteps{parameters['nsteps']}"
                method = parameters["method"]
                aiDrift = StandardTimeStepping.read(wd + file_ext + f"_ts_{method}.h5").aiDrift
                plot_timestepping_diagnostics(
                    log_path=parameters["OutputTimestepsLog"],
                    aiDrift=aiDrift,
                    TimeStepDistribution=parameters["TimeStepDistribution"],
                    nsteps=parameters["nsteps"],
                    save_path=wd+"time_step_diagnostics.pdf"
                )
        UNINDENT()
        logger.info("Plotting time step diagnostics done.")

    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.")
