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

"""
Various tools for the WIP3M project.
"""

import os
import gc
import numpy as np

from wip3m import *
from wip3m.logger import getCustomLogger

logger = getCustomLogger(__name__)


def none_bool_str(value):
    """Convert a string to None, bool, or str.

    Parameters
    ----------
    value : str
        String to convert.

    Returns
    -------
    None, bool, or str
        Converted value.
    """
    if value == "None" or value == None:
        return None
    elif value == "True":
        return True
    elif value == "False":
        return False
    return value


def get_k_max(L, size):
    """
    Compute the maximum wavenumber for a given box size.

    Parameters
    ----------
    L : float
        Size of the box in Mpc/h.
    size : int
        Number of grid cells along each dimension.

    Returns
    -------
    float
        Maximum wavenumber in h/Mpc.
    """
    from numpy import pi, sqrt

    # If kx = ky = kz = k_Nyquist, then |k| = sqrt(3) * k_Nyquist
    return int(1e3 * sqrt(3) * pi * size / L + 1) * 1e-3

def joinstrs(list_of_strs):
    """Join a list of strings into a single string.

    Parameters
    ----------
    list_of_strs : list of str
        List of strings to join.

    Returns
    -------
    str
        Concatenated string.
    """
    return "".join([str(x) for x in list_of_strs if x is not None])


def generate_sim_params(params_dict, ICs, workdir, outdir, file_ext=None, force=False):
    """Write the parameter file.

    Parameters
    ----------
    params_dict : dict
        Dictionary containing the parameters for the simulation.
    ICs : str
        Path to the initial conditions.
    workdir : str
        Directory where to store the parameter file.
    outdir : str
        Directory where to store the simulation outputs.
    file_ext : str, optional
        Prefix for the output files.
    force : bool, optional
        If True, force regeneration of the parameter file.

    Returns
    -------
    sbmy_path : str
        Path to the parameter file generated.

    """
    from os.path import isfile
    from pysbmy import param_file
    from pysbmy.timestepping import StandardTimeStepping, P3MTimeStepping

    method = params_dict["method"]
    path = workdir + file_ext + "_" if file_ext else workdir
    simpath = outdir + file_ext + "_" if file_ext else outdir
    sbmy_path = joinstrs([path, "example_", method, ".sbmy"])

    # Parameters shared by all methods for this run
    Particles = params_dict["Np"]
    Mesh = params_dict["N"]
    BoxSize = params_dict["L"]
    corner0 = params_dict["corner0"]
    corner1 = params_dict["corner1"]
    corner2 = params_dict["corner2"]
    h = params_dict["h"]
    Omega_m = params_dict["Omega_m"]
    Omega_b = params_dict["Omega_b"]
    n_s = params_dict["n_s"]
    sigma8 = params_dict["sigma8"]

    # Generate the time-stepping distribution
    if method in ["pm", "cola", "spm", "p3m"]:
        TimeStepDistribution = params_dict["TimeStepDistribution"]

        fac_dyn_custom = params_dict.get("fac_dyn_custom", DEFAULT_FAC_DYN_CUSTOM)
        da_max_early_custom = params_dict.get("da_max_early_custom", DEFAULT_DA_MAX_EARLY_CUSTOM)
        fac_H_custom = params_dict.get("fac_H_custom", DEFAULT_FAC_H_CUSTOM)
        fac_bend = params_dict.get("fac_bend", DEFAULT_FAC_BEND)
        sub_bend1 = params_dict.get("sub_bend1", DEFAULT_SUB_BEND1)
        sub_bend2 = params_dict.get("sub_bend2", DEFAULT_SUB_BEND2)
        da_max_late_custom = params_dict.get("da_max_late_custom", DEFAULT_DA_MAX_LATE_CUSTOM)

        ts_filename = path + "ts_" + method + ".h5"
        logger.info("Time-stepping distribution file: %s", ts_filename)
        if not isfile(ts_filename) or force:
            if TimeStepDistribution in [0, 1, 2]:
                ai = params_dict["ai"]
                af = params_dict["af"]
                nsteps = params_dict["nsteps"]
                snapshots = np.full((nsteps), False)
                TS = StandardTimeStepping(ai, af, snapshots, TimeStepDistribution)
            elif TimeStepDistribution == 3:
                ai = params_dict["ai"]
                af = params_dict["af"]
                p3m_fit_coeffs = params_dict.get("p3m_fit_coeffs", None)
                use_p3m_fit = params_dict.get("use_p3m_fit", False)
                fac_p3m_fit = params_dict.get("fac_p3m_fit", DEFAULT_FAC_P3M_FIT)
                TS = P3MTimeStepping(
                    ai=params_dict["ai"],
                    af=params_dict["af"],
                    cosmo=params_dict["cosmo_dict"],
                    fac_dynamical=fac_dyn_custom,
                    delta_a_max_early=da_max_early_custom,
                    fac_hubble=fac_H_custom,
                    fac_bend=fac_bend,
                    sub_bend1=sub_bend1,
                    sub_bend2=sub_bend2,
                    delta_a_max_late=da_max_late_custom,
                    p3m_fit_coeffs=p3m_fit_coeffs,
                    use_p3m_fit=use_p3m_fit,
                    fac_p3m_fit=fac_p3m_fit,
                    forces_every=1,
                )
            else:
                raise ValueError("Invalid TimeStepDistribution: {}".format(TimeStepDistribution))
            TS.write(ts_filename)
            logger.info("TS.ai = %f, TS.af = %f, TS.nsteps = %d", TS.ai, TS.af, TS.nsteps)
        else:
            logger.info("Time-stepping distribution file already exists: %s", ts_filename)
        StandardTimeStepping.read(ts_filename).plot(savepath=path + "ts_" + method + ".png")
    elif method in ["lpt"]:
        pass
    else:
        raise ValueError("Method not supported: {}".format(method))

    # Write the parameter file
    logger.info("Generating parameter file...")
    if params_dict["method"] == "lpt":
        S = param_file(
            OutputRngStateLPT=simpath + "dummy.rng",
            # Basic setup:
            Particles=Particles,
            Mesh=Mesh,
            BoxSize=BoxSize,
            corner0=corner0,
            corner1=corner1,
            corner2=corner2,
            # Initial conditions:
            ICsMode=params_dict["ICsMode"],
            InputWhiteNoise=params_dict["InputWhiteNoise"],  # None or str
            WriteInitialConditions=1,
            OutputInitialConditions=ICs,
            # Power spectrum:
            InputPowerSpectrum=params_dict["InputPowerSpectrum"],
            # Final conditions for LPT:
            OutputLPTSnapshot=simpath + "lpt_particles.gadget3",
            OutputLPTDensity=simpath + "lpt_density.h5",
            #############################
            ## Cosmological parameters ##
            #############################
            h=h,
            Omega_m=Omega_m,
            Omega_b=Omega_b,
            n_s=n_s,
            sigma8=sigma8,
            Omega_q=1.0 - Omega_m,
            Omega_k=0.0,
            w0_fld=-1.0,
            wa_fld=0.0,
        )
    if params_dict["method"] == "pm":
        S = param_file(
            # Basic setup:
            Particles=Particles,
            Mesh=Mesh,
            BoxSize=BoxSize,
            corner0=corner0,
            corner1=corner1,
            corner2=corner2,
            # Initial conditions:
            ICsMode=2,
            InputInitialConditions=ICs,
            # Final conditions for LPT:
            RedshiftLPT=params_dict["RedshiftLPT"],
            WriteLPTSnapshot=0,
            WriteLPTDensity=0,
            ####################
            ## Module PM/COLA ##
            ####################
            ModulePMCOLA=1,
            EvolutionMode=1,
            ParticleMesh=params_dict["Npm"],
            TimeStepDistribution=ts_filename,
            # Final snapshot:
            RedshiftFCs=params_dict["RedshiftFCs"],
            WriteFinalSnapshot=1,
            OutputFinalSnapshot=simpath + "pm_snapshot.gadget3",
            WriteFinalDensity=1,
            OutputFinalDensity=simpath + "final_density_pm.h5",
            RunForceDiagnostic=params_dict["RunForceDiagnostic"],
            nPairsForceDiagnostic=params_dict["nPairsForceDiagnostic"],
            nBinsForceDiagnostic=params_dict["nBinsForceDiagnostic"],
            maxTrialsForceDiagnostic=params_dict["maxTrialsForceDiagnostic"],
            OutputForceDiagnostic=params_dict["OutputForceDiagnostic"],
            #############################
            ## Cosmological parameters ##
            #############################
            h=h,
            Omega_m=Omega_m,
            Omega_b=Omega_b,
            n_s=n_s,
            sigma8=sigma8,
            Omega_q=1.0 - Omega_m,
            Omega_k=0.0,
            w0_fld=-1.0,
            wa_fld=0.0,
        )
    if params_dict["method"] == "cola":
        S = param_file(
            # Basic setup:
            Particles=Particles,
            Mesh=Mesh,
            BoxSize=BoxSize,
            corner0=corner0,
            corner1=corner1,
            corner2=corner2,
            # Initial conditions:
            ICsMode=2,
            InputInitialConditions=ICs,
            # Final conditions for LPT:
            RedshiftLPT=params_dict["RedshiftLPT"],
            WriteLPTSnapshot=0,
            WriteLPTDensity=0,
            ####################
            ## Module PM/COLA ##
            ####################
            ModulePMCOLA=1,
            EvolutionMode=2,
            ParticleMesh=params_dict["Npm"],
            TimeStepDistribution=ts_filename,
            # Final snapshot:
            RedshiftFCs=params_dict["RedshiftFCs"],
            WriteFinalSnapshot=1,
            OutputFinalSnapshot=simpath + "cola_snapshot.gadget3",
            WriteFinalDensity=1,
            OutputFinalDensity=simpath + "final_density_cola.h5",
            RunForceDiagnostic=params_dict["RunForceDiagnostic"],
            nPairsForceDiagnostic=params_dict["nPairsForceDiagnostic"],
            nBinsForceDiagnostic=params_dict["nBinsForceDiagnostic"],
            maxTrialsForceDiagnostic=params_dict["maxTrialsForceDiagnostic"],
            OutputForceDiagnostic=params_dict["OutputForceDiagnostic"],
            #############################
            ## Cosmological parameters ##
            #############################
            h=h,
            Omega_m=Omega_m,
            Omega_b=Omega_b,
            n_s=n_s,
            sigma8=sigma8,
            Omega_q=1.0 - Omega_m,
            Omega_k=0.0,
            w0_fld=-1.0,
            wa_fld=0.0,
        )
    elif params_dict["method"] == "p3m" or params_dict["method"] == "spm":
        if params_dict["EvolutionMode"] == 6 or params_dict["EvolutionMode"] == 7:
            ## Special case for COLA ##
            sub_bend1_diagnostic = params_dict.get("sub_bend1", DEFAULT_SUB_BEND1_COLA)
            sub_bend2_diagnostic = params_dict.get("sub_bend2", DEFAULT_SUB_BEND2_COLA)
        else:
            sub_bend1_diagnostic = params_dict.get("sub_bend1", DEFAULT_SUB_BEND1)
            sub_bend2_diagnostic = params_dict.get("sub_bend2", DEFAULT_SUB_BEND2)
        S = param_file(
            # Basic setup:
            Particles=Particles,
            Mesh=Mesh,
            BoxSize=BoxSize,
            corner0=corner0,
            corner1=corner1,
            corner2=corner2,
            # Initial conditions:
            ICsMode=2,
            InputInitialConditions=ICs,
            # Final conditions for LPT:
            RedshiftLPT=params_dict["RedshiftLPT"],
            WriteLPTSnapshot=0,
            WriteLPTDensity=0,
            ####################
            ## Module PM/COLA ##
            ####################
            ModulePMCOLA=1,
            EvolutionMode=params_dict["EvolutionMode"],
            ParticleMesh=params_dict["Npm"],
            TimeStepDistribution=ts_filename,
            # Final snapshot:
            RedshiftFCs=params_dict["RedshiftFCs"],
            WriteFinalSnapshot=1,
            OutputFinalSnapshot=simpath + "{}_snapshot.gadget3".format(method),
            WriteFinalDensity=1,
            OutputFinalDensity=simpath + "final_density_{}.h5".format(method),
            n_Tiles=params_dict["n_Tiles"],
            RunForceDiagnostic=params_dict.get("RunForceDiagnostic", False),
            nPairsForceDiagnostic=params_dict.get("nPairsForceDiagnostic", 0),
            nBinsForceDiagnostic=params_dict.get("nBinsForceDiagnostic", 0),
            maxTrialsForceDiagnostic=params_dict.get("maxTrialsForceDiagnostic", 0),
            OutputForceDiagnostic=params_dict.get("OutputForceDiagnostic", ""),
            PrintOutputTimestepsLog=params_dict.get("PrintOutputTimestepsLog", False),
            OutputTimestepsLog=params_dict.get("OutputTimestepsLog", ""),
            fac_dyn_custom=params_dict.get("fac_dyn_custom", DEFAULT_FAC_DYN_CUSTOM),
            fac_H_custom=params_dict.get("fac_H_custom", DEFAULT_FAC_H_CUSTOM),
            fac_bend=params_dict.get("fac_bend", DEFAULT_FAC_BEND),
            sub_bend1=params_dict.get("sub_bend1", sub_bend1_diagnostic),
            sub_bend2=params_dict.get("sub_bend2", sub_bend2_diagnostic),
            fac_p3m_fit=params_dict.get("fac_p3m_fit", DEFAULT_FAC_P3M_FIT),
            use_p3m_fit=params_dict.get("use_p3m_fit", False),
            da_max_early_custom=params_dict.get("da_max_early_custom", DEFAULT_DA_MAX_EARLY_CUSTOM),
            da_max_late_custom=params_dict.get("da_max_late_custom", DEFAULT_DA_MAX_LATE_CUSTOM),
            #############################
            ## Cosmological parameters ##
            #############################
            h=h,
            Omega_m=Omega_m,
            Omega_b=Omega_b,
            n_s=n_s,
            sigma8=sigma8,
            Omega_q=1.0 - Omega_m,
            Omega_k=0.0,
            w0_fld=-1.0,
            wa_fld=0.0,
        )

    if not isfile(sbmy_path) or force:
        S.write(sbmy_path)
        logger.info("Parameter file written to %s", sbmy_path)
    else:
        logger.info("Parameter file already exists at %s", sbmy_path)

    return sbmy_path


def read_field(*args):
    """
    Read a field from a file.

    Parameters
    ----------
    args : tuple
        Arguments to pass to the read_field function from pysbmy.

    Returns
    -------
    field : str
        The field read from the file.
    """
    from io import BytesIO
    from wip3m.low_level import stdout_redirector
    from pysbmy.field import read_field as _read_field

    with BytesIO() as f:
        with stdout_redirector(f):
            return _read_field(*args)


def generate_white_noise_Field(
    L,
    size,
    corner,
    seedphase,
    fname_whitenoise,
    seedname_whitenoise,
    force_phase=False,
):
    """
    Generate a white noise realisation in physical space and write it to
    disk.

    Parameters
    ----------
    L : float
        Size of the simulation box (in Mpc/h).
    size : int
        Number of grid points along each axis.
    corner : float
        Position of the corner (in Mpc/h).
    seedphase : int or list of int
        User-provided seed to generate the initial white noise.
    fname_whitenoise : str
        File path to write the white noise realisation.
    seedname_whitenoise : str
        File path to write the seed state of the RNG.
    force_phase : bool, optional
        If True, forces regeneration of the random phases. Default is
        False.

    Raises
    ------
    OSError
        If file writing fails or directory paths are invalid.
    RuntimeError
        For unexpected issues.
    """
    if not os.path.exists(fname_whitenoise) or force_phase:
        from pysbmy.field import BaseField

        try:
            logger.debug("Generating white noise for L=%.2f, size=%d", L, size)
            rng = np.random.default_rng(seedphase)

            logger.debug("Saving RNG state to %s", seedname_whitenoise)
            np.save(seedname_whitenoise, rng.bit_generator.state)
            with open(seedname_whitenoise + ".txt", "w") as f:
                f.write(str(rng.bit_generator.state))

            data = rng.standard_normal(size=size**3)
            wn = BaseField(L, L, L, corner, corner, corner, 1, size, size, size, data)
            del data

            wn.write(fname_whitenoise)
            logger.debug("White noise field written to %s", fname_whitenoise)
            del wn
        except OSError as e:
            logger.error("Writing white noise failed at '%s': %s", fname_whitenoise, str(e))
            raise
        except Exception as e:
            logger.critical("Unexpected error in generate_white_noise_Field: %s", str(e))
            raise RuntimeError("generate_white_noise_Field failed.") from e
        finally:
            gc.collect()

def run_simulation(name, params, wd, logdir):
    from io import BytesIO
    from wip3m.low_level import stderr_redirector
    from pysbmy import pySbmy

    if params.get("nsteps") is not None:
        file_ext = f"{name}_nsteps{params['nsteps']}"
    elif params.get("relax") is not None:
        file_ext = f"{name}_relax{str(params['relax']).replace('.', '_')}"
    elif params.get("file_ext") is not None:
        file_ext = params["file_ext"]
    else:
        file_ext = None
    method = params["method"]
    path = wd + file_ext + "_" if file_ext else wd
    sbmy_path = joinstrs([path, "example_", method, ".sbmy"])
    log_path = joinstrs([logdir, file_ext, method, ".txt"])

    f = BytesIO()
    with stderr_redirector(f):
        pySbmy(sbmy_path, log_path)
    f.close()
