"""Container to generate origin showers and also copy relevant files."""

import os
import shutil
import numpy as np
from typing_extensions import Self

from . import write_simulation_to_file, generate_simulation
from smiet import units

import logging

import logging

logger = logging.getLogger("smiet.corsika.generate_origin_library")


class OriginLibraryGenerator:
    corsika_exec_hadr_label : dict = {
        "sibyll": "SIBYLL",
        "qgsjetII": "QGSII",
    }
    """Class to generate the origin shower library using CORSIKA simulations."""

    def __init__(
        self: Self,
        sim_primary: str = "proton",
        sim_energy: float = 10**8 * units.GeV,
        atm_model: int = 17,
        site: str = "lofar",
    ) -> None:
        """
        Initialize the generator.

        Parameters
        ----------
        sim_primary : str, default="proton"
            Primary particle type (e.g., "proton", "iron").

            Currently only "proton" and "iron" are supported.
        sim_energy : float, default=10**8 * units.GeV
            Energy of the primary particle in GeV.
        atm_model : int
            Atmospheric model to use (default is 17).
        site : str
            Site location (default is "lofar") for the magnetic field location
            and the observation level.
        """
        self.sim_primary = sim_primary
        self.sim_energy = sim_energy
        self.atm_model = atm_model
        self.site = site
        self.obs_lev = self.__get_obs_lvl(site)

        self.corsika_exec_name = None
        self.corsika_runner_path = None
        self.run_dir = None

    def __get_obs_lvl(self, site: str) -> float:
        """Set the observation level based on the site."""
        site_obs_levels = {
            "lofar": 0.0 * units.m,
            "auger": 1400.0 * units.m,
        }
        return site_obs_levels[site]

    def setup_paths(
        self: Self,
        corsika_runner_path: str,
        run_dir: str,
        cluster_type: str = "horeka",
        hadr_model: str = "sibyll",
    ) -> None:
        """
        Configure the paths for CORSIKA runner and run directory.

        Parameters
        ----------
        corsika_runner_path : str
            Path to the CORSIKA runner executable.
        run_dir : str
            Directory where the CORSIKA simulations will be run and saved.
        cluster_type : str, default="horeka"
            Type of cluster that the jobs are submitted into.

            The submit scripts should be placed in the submit_templates directory,
            with the naming scheme submit_{cluster_type}.sh.

            Currently only "horeka" is supported.
        hadr_model : str, default="sibyll"
            Hadronic interaction model used for the fitting functions.
            Currently only "sibyll" and "qgsjetII" are supported.
        """
        self.corsika_runner_path = corsika_runner_path
        self.run_dir = run_dir

        # ensure that the directories exist
        if not os.path.exists(self.corsika_runner_path):
            raise FileNotFoundError(
                f"CORSIKA runner path {self.corsika_runner_path} not found."
            )
        if not os.path.exists(self.run_dir):
            os.makedirs(self.run_dir, exist_ok=True)

        self.hadr_model = hadr_model
        # name of coreas runner executable
        self.corsika_exec_name = f"mpi_corsika77550Linux_{self.corsika_exec_hadr_label[hadr_model]}_urqmd_thin_coreas_parallel_runner"

        # ensure that the corsika executable exists
        if not os.path.exists(
            os.path.join(corsika_runner_path, self.corsika_exec_name)
        ):
            print(
                f"CORSIKA executable {self.corsika_exec_name} not found in {corsika_runner_path}. Exiting."
            )
            raise FileNotFoundError()

        path_to_submit_tmpl = os.path.join(
            os.path.abspath(os.path.dirname(__file__)),
            "submit_templates",
            f"submit_{cluster_type}.sh",
        )

        # make sure it exists
        if not os.path.exists(path_to_submit_tmpl):
            raise FileNotFoundError(
                f"Submit template for cluster type {cluster_type} not found at {path_to_submit_tmpl}."
            )

        with open(path_to_submit_tmpl, "r") as f:
            self.submit_file_templ = f.read()

    def generate_submit_files(
        self: Self,
        zenith_angles: tuple = (0, 63, 3),
        xmax_values: tuple = (600, 1300, 100),
        n_sims: int = 3,
        n_nodes: int = 2,
        n_cores_per_node: int = 76,
        vxvxB: bool = True,
    ) -> list:
        """
        Generate steering files for CORSIKA simulations.

        Parameters
        ----------
        zenith_angles : tuple
            Tuple specifying (min, max, step) for zenith angles in degrees.
            NOte that max should be less than 60 degrees, and include the step.
        xmax_values : tuple
            Tuple specifying (min, max, step) for Xmax values in g/cm^2.
            NOte that max should include the step.
        n_sims : int, default=3
            The number of times to repeat each simulation (default is 3).
            This is to ensure that we have enough statistics.
        n_nodes : int
            Number of nodes to use for each simulation (default is 2).
        n_cores_per_node : int
            Number of cores per node (default is 76).
        vxvxB : bool
            Whether to only generate showers on the vxvxB axis or not.
            Default is True. Otherwise, a rotated starshape is used.

        Returns
        -------
        list
            List of directories where simulations are set up.
        """
        if vxvxB:
            nr_of_arms = 1
            rotate = False
        else:
            nr_of_arms = 8
            rotate = True
        sim_cores = int(n_nodes * n_cores_per_node)

        zeniths = np.arange(*zenith_angles)
        xmaxs = np.arange(*xmax_values)

        sim_rundirs = []

        # check if the run directory has been set up
        if self.run_dir is None:
            raise ValueError(
                "Run directory not set up. Please run setup_paths() first."
            )

        # for bookkeeping, generate a file that takes in the simulation number, zenith, and xmax
        sim_labels = ["SIM\tZENITH(deg)\tXMAX(g/cm2)\tITER"]

        for izen, ixmax, isim in np.ndindex(len(zeniths), len(xmaxs), n_sims):

            sim_zenith = zeniths[izen] * units.deg
            sim_xmax = xmaxs[ixmax]
            # compute the simulation number based on a basic criteria
            sim_nr = int(f"00{izen:02d}{isim:01d}{ixmax:01d}")

            inp_file, lst_file, reas_file = generate_simulation(
                sim_nr=sim_nr,
                sim_primary=self.sim_primary,
                sim_cores=sim_cores,
                sim_zenith=sim_zenith,
                sim_energy=self.sim_energy,
                sim_xmax=sim_xmax,
                sim_azimuth=0.0 * units.deg,  # fixed here
                slice_gram=5.0,  # fixed here
                magnetic_field=self.site,
                atmosphere=self.atm_model,
                nr_of_arms=nr_of_arms,
                rotate=rotate,
                core=np.array([0, 0, self.obs_lev]),
            )

            # create the necessary directories
            sim_dir = os.path.join(
                self.run_dir, f"zen{sim_zenith / units.deg:.0f}", f"xmax{sim_xmax:.0f}"
            )
            sim_rundir = os.path.join(sim_dir, f"SIM{sim_nr:06d}")
            os.makedirs(sim_dir, exist_ok=True)

            # check if a simulation has already been ran, if so then skip
            if os.path.exists(os.path.join(sim_rundir, f"DAT{sim_nr:06d}.long")):
                logger.warning(f"SIM{sim_nr:06d} already ran. Skipping.")
                continue
            # manually modify the input file to add the direct and datdir inputs
            inp_file = self.modify_input_file(inp_file, sim_rundir)

            # write to directory
            write_simulation_to_file(
                inp_file, lst_file, reas_file, sim_directory=sim_dir, reset=True
            )

            # create a logs directory
            os.makedirs(os.path.join(sim_dir, "logs"), exist_ok=True)

            # copy the coreas_to_hdf5_origin.py script to the directory
            shutil.copy(
                os.path.join(os.path.dirname(__file__), "coreas_to_hdf5_origin.py"),
                os.path.join(self.run_dir, "coreas_to_hdf5_origin.py"),
            )

            # also create a hdf5 directory to store the hdf5 files
            os.makedirs(os.path.join(self.run_dir, "hdf5_sims"), exist_ok=True)

            # also create the submit file
            submit_file = self.submit_file_templ.format(
                sim_xmax,
                sim_zenith / units.deg,
                sim_dir,
                n_nodes,
                n_cores_per_node,
                self.corsika_runner_path,
                self.corsika_exec_name,
                sim_nr,
                self.run_dir,
            )
            with open(os.path.join(sim_rundir, "submit.sh"), "w") as f:
                f.write(submit_file)

            sim_labels.append(
                f"{sim_nr:06d}\t{sim_zenith / units.deg:.1f}\t{sim_xmax:.1f}\t{isim+1}"
            )
            sim_rundirs.append(sim_rundir)

        logger.info(
            f"Generated {len(sim_rundirs)} simulation directories in {self.run_dir}"
        )

        # write the sim_labels to a file
        with open(os.path.join(self.run_dir, "sim_labels.txt"), "w") as f:
            f.write("\n".join(sim_labels))

        return sim_rundirs

    def submit_jobs(self: Self, sim_rundirs: list) -> None:
        """
        Submit the generated jobs to the SLURM scheduler.

        Parameters
        ----------
        sim_rundirs : list
            List of directories where simulations are set up.
        """
        for sim_rundir in sim_rundirs:
            submit_cmd = f"sbatch {os.path.join(sim_rundir, 'submit.sh')}"
            logger.info(f"Submitting job with command: {submit_cmd}")
            os.system(submit_cmd)

    def modify_input_file(self: Self, inp_file: list, direct: str) -> list:
        """
        Modify the input file to include direct and datdir commands.

        Parameters
        ----------
        inp_file : list
            List of strings representing the lines of the input file.

        Returns
        -------
        list
            Modified list of strings for the input file.
        """
        modified_inp = []
        for line in inp_file:
            if line.startswith("DIRECT"):
                modified_inp.append(f"DATDIR {self.corsika_runner_path}\n")
                modified_inp.append(f"DIRECT {direct}/\n")
            else:
                modified_inp.append(line)
        return modified_inp

    def generate_single_shower(
        self: Self,
        sim_zenith: float,
        sim_xmax: float,
        n_nodes: int = 2,
        n_cores_per_node: int = 76,
        vxvxB: bool = True,
        sim_dir="./",
        sim_nr: int = 999999,
    ) -> str:
        """
        Generate a single origin shower simulation. This is useful for testing and debugging purposes.

        Parameters
        ----------
        sim_zenith : float
            Zenith angle in degrees.
        sim_xmax : float
            Xmax value in g/cm^2.
         n_nodes : int
            Number of nodes to use for each simulation (default is 2).
        n_cores_per_node : int
            Number of cores per node (default is 76).
        vxvxB : bool
            Whether to only generate showers on the vxvxB axis or not.
            Default is True. Otherwise, a rotate starshape is used.
        sim_dir : str
            Directory where the simulation will be set up.
        sim_nr : int, default=999999
            Simulation number to use for the shower.
        """
        if vxvxB:
            nr_of_arms = 1
            rotate = False
        else:
            nr_of_arms = 8
            rotate = True
        inp_file, lst_file, reas_file = generate_simulation(
            sim_nr=sim_nr,
            sim_primary=self.sim_primary,
            sim_cores=int(n_nodes * n_cores_per_node),
            sim_zenith=sim_zenith * units.deg,
            sim_energy=self.sim_energy,
            sim_xmax=sim_xmax,
            slice_gram=5.0,
            magnetic_field=self.site,
            atmosphere=self.atm_model,
            nr_of_arms=nr_of_arms,
            rotate=rotate,
            core=np.array([0, 0, self.obs_lev]),
        )

        # manually modify the input file to add the direct and datdir inputs
        inp_file = self.modify_input_file(inp_file, sim_dir)

        write_simulation_to_file(
            inp_file, lst_file, reas_file, sim_directory=sim_dir, reset=True
        )

        return inp_file, lst_file, reas_file
