"""
Main class to run SMIET on
"""

import os
import logging
import numpy as np
from typing_extensions import Self, Union

import smiet.units as units
import smiet.numpy as smiet_np

try:
    import jax
    import jax.numpy as jnp
    import smiet.jax as smiet_jax
except ImportError:
    jax = None
    jnp = None
    smiet_jax = None


class SMIETRun:
    """
    Main class to run SMIET on. The idea of this class is so that all users can easily apply SMIET to their own analyses.

    Main Features:
    --------------
    - generating templates from origin showers, in particular from generated origin shower libraries
    - loading templates from existing template libraries
    - loading target showers from existing Coreas simulations
    - generating target showers from user-provided longitudinal profiles
    - synthesizing target showers from origin shower templates + target showers, with interpolation (by default)
    - retrieving time traces, frequency spectra, and fluences for synthesized showers

    The user is expected to give the site in which SMIET will apply on, and the backend (jax or numpy).
    """

    __site_parameters = {
        "base": {
            "magnetic_field_vector": np.array(
                [0.004675, 0.186270, -0.456412]
            ),  # in Gauss
            "frequency_range": [30, 500, 100],  # in MHz
        },
        "lofar": {
            "magnetic_field_vector": np.array(
                [0.004675, 0.186270, -0.456412]
            ),  # in Gauss
            "frequency_range": [30, 80, 50],  # in MHz
        },
        "ska": {
            "magnetic_field_vector": np.array([0.0, 2.760, 4.827]),  # in Gauss
            "frequency_range": [30, 500, 100],  # in MHz
        },
        "auger": {
            "magnetic_field_vector": np.array(
                [0.00871198, 0.19693423, 0.1413841]
            ),  # in Gauss
            "frequency_range": [30, 80, 50],  # in MHz
        },
    }

    def __init__(
        self: Self,
        site: str = "base",
        backend: str = "numpy",
    ) -> None:
        """
        Initialize the SMIET runner.

        This will set up default parameters and configurations.

        Parameters:
        ----------
        site : str, default = "base"
            The site to run the simulation for. Default is "base", which uses the LOFAR magnetic field with maximal frequency range ([30, 500] MHz), which are the values which SMIET is tested on.

            Current available options are:
            - "base": LOFAR magnetic field with maximal frequency range ([30, 500] MHz)
            - "lofar": LOFAR site parameters (magnetic field and frequency range [30, 80] MHz)
            - "ska": SKA site parameters (magnetic field and frequency range [30, 500] MHz)
            - "auger": Pierre Auger Observatory (AERA) site parameters (magnetic field and frequency range [30, 80] MHz)
        backend : str
            The backend to use for computations. Options are "numpy", "jax_cpu", or "jax_cuda". Default is "numpy".

            If "jax_cpu" or "jax_cuda" is selected but JAX is not installed, it will fall back to "numpy".
            In this case, a warning to install SMIET with the JAX backend will be printed.

            Here, "jax_cpu" will use JAX with CPU as the device, while "jax_cuda" will use JAX with GPU (CUDA) as the device, if available.
            If a GPU is not available, it will fall back to CPU.

        Returns:
        --------
        None
        """
        self.site = site
        self.__get_site_parameters(site)
        self.logger = logging.getLogger("smiet.core.SMIETRun")

        if backend == "numpy":
            self.backend = "numpy"
            self.smiet = smiet_np
        elif backend.find("jax") != -1:
            raise NotImplementedError(
                "JAX backend is not yet implemented in this version of SMIET. Please use the numpy backend for now."
            )
            if jax is None or smiet_jax is None:
                self.logger.warning(
                    "JAX backend selected but JAX is not installed. Falling back to numpy backend. To use JAX, please install SMIET with the JAX backend."
                )
                self.backend = "numpy"
                self.smiet = smiet_np
            else:
                self.backend = "jax"
                if backend == "jax_cpu":
                    jax.config.update("jax_default_device", jax.devices("cpu")[0])
                self.smiet = smiet_jax

        # some flags
        self.__loaded_templates = False
        self.__loaded_targets = False
        self.__synthesized_showers = False

        # containers for origins
        self.origin_xmaxs = None
        self.origin_zeniths = None
        self.syntheses = []

        # containers for targets
        self.target_xmaxs = None
        self.target_zeniths = None
        self.targets = []

        # containers for synthesis
        self.synth_traces = []
        self.synth_times = []

        # additional properties that are useful for later
        self.antenna_positions = []
        self.antenna_labels = None
        self.delta_t = None

    def __get_site_parameters(self: Self, site: str) -> None:
        """
        Retrieve the parameters related to each site.

        This includes the magnetic field and frequency range.

        Parameters:
        ----------
        site : str
            The site to retrieve parameters for.
        """
        self.freq_ar = self.__site_parameters[site]["frequency_range"]
        self.magnet = self.__site_parameters[site]["magnetic_field_vector"]

    def generate_templates(
        self: Self,
        origin_shower_repo: str,
        template_repo: str,
    ) -> None:
        """
        Generate templates from a list of origin showers generated from the origin shower library.

        This is meant to be ran once to generate the templates used for the origin shower library.

        Parameters:
        ----------
        origin_shower_repo : str
            The path to the origin shower repository containing all the origin showers to generate templates from.
        template_repo : str
            The path to the repository to save the generated templates.
        """
        # first check if the template repo exists
        if not os.path.exists(template_repo):
            raise FileNotFoundError(
                f"Template repository {template_repo} does not exist."
            )

        # get the path to all the templates
        path_to_templates = os.path.join(template_repo, self.site, self.backend)

        # read the sim_labels.txt file and extract the info for the configurations
        # this loads the following columns:
        # SIM_ID  XMAX  ZENITH  ITERATION
        config_data = np.genfromtxt(
            os.path.join(origin_shower_repo, "sim_labels.txt"),
            delimiter="\t",
            skip_header=1,
        )

        # get an estimate of how long it will take
        n_showers = len(
            [f for f in os.listdir(origin_shower_repo) if f.endswith(".hdf5")]
        )
        self.logger.warning(
            f"Generating {n_showers} templates from origin showers. This will take at most {n_showers * 10 // 60} minutes."
        )

        # for each origin shower, make a template and save it
        config_txt = "xmax,zenith\n"
        for origin_shower_file in os.listdir(origin_shower_repo):
            if not origin_shower_file.endswith(".hdf5"):
                continue

            # add sim id identification
            # find the corresponding row in the config file
            origin_shower_sim_id = origin_shower_file.replace("SIM", "").replace(
                ".hdf5", ""
            )

            sim_row_idx = np.argwhere(config_data[:, 0] == int(origin_shower_sim_id))[
                0
            ][0]
            # get the properties
            zenith = int(config_data[sim_row_idx, 1])
            xmax = int(config_data[sim_row_idx, 2])
            iter_num = int(config_data[sim_row_idx, 3])

            config_txt += f"{xmax},{zenith}\n"

            origin_shower = self.smiet.SlicedShower(
                os.path.join(origin_shower_repo, origin_shower_file)
            )

            # get a label for the directory based on xmax and zenith
            save_path = os.path.join(path_to_templates, f"xmax{xmax}_zen{zenith}")
            template_name = f"template_sim{iter_num}.hdf5"

            # check if it exists already
            if os.path.exists(os.path.join(save_path, template_name)):
                self.logger.info(
                    f"Template for origin shower {origin_shower_file} already exists. Skipping."
                )
                continue

            # make if it doesnt exist already
            if not os.path.exists(save_path):
                os.makedirs(save_path, exist_ok=True)

            synthesis = self.smiet.TemplateSynthesis(freq_ar=self.freq_ar)

            synthesis.make_template(origin_shower)

            # now save the template based on the iteration number
            synthesis.save_template(
                save_dir=save_path, template_file=f"template_sim{iter_num}.hdf5"
            )

            if self.backend == "jax":
                # clear cache everytime to avoid memory issues
                jax.clear_caches()

        # finally for each configuration, get the xmax value and set the closest one to
        # the true value to be sim1
        for xmax_zenith in os.listdir(path_to_templates):
            config_path = os.path.join(path_to_templates, xmax_zenith)
            if not os.path.isdir(config_path):
                continue
            xmax_diff = []
            for config_file in os.listdir(config_path):
                if not config_file.endswith(".hdf5"):
                    continue
                template_path = os.path.join(config_path, config_file)
                synthesis = self.smiet.TemplateSynthesis(freq_ar=self.freq_ar)
                synthesis.load_template(template_path)

                # check if this is the closest to the true xmax
                true_xmax = float(xmax_zenith.split("_")[0].replace("xmax", ""))
                synth_xmax = synthesis.template_information["xmax"]
                xmax_diff.append(np.abs(true_xmax - synth_xmax))
            # get the index of the closest one
            closest_index = np.argmin(xmax_diff)
            # rename the file to sim1
            closest_file = os.listdir(config_path)[closest_index]
            if closest_file != "template_sim1.hdf5":
                # rename the the original sim1 file to something temporary first
                os.rename(
                    os.path.join(config_path, "template_sim1.hdf5"),
                    os.path.join(config_path, "temp.hdf5"),
                )
                # then rename the closest file to sim 1
                os.rename(
                    os.path.join(config_path, closest_file),
                    os.path.join(config_path, "template_sim1.hdf5"),
                )
                # finally replace temp file to the closest file name
                os.rename(
                    os.path.join(config_path, "temp.hdf5"),
                    os.path.join(config_path, closest_file),
                )

        self.logger.info(
            f"Generated templates from origin showers in {origin_shower_repo} and saved to {template_repo}."
        )

        # make a configuration file to record all the configurations
        with open(os.path.join(path_to_templates, "config.txt"), "w") as config_file:
            config_file.write(config_txt)

    def load_templates(
        self: Self,
        template_repo: str,
        xmax_range: Union[tuple, None] = None,
        zenith_range: Union[tuple, None] = None,
        randomize: bool = False,
    ) -> None:
        """
        Load the templates for synthesis. This will load the necessary templates which will be used for the run.

        This will automatically load the templates if they are available, otherwise new templates will be generated from origin showers.

        The origin shower library will be used by default. If you want to perform synthesis with a single origin shower (e.g. for testing purposes), use the `synthesize_with_single_origin` function instead.

        If you only want to load a subset of xmax / zenith values for the synthesis, you can do so by specifying the ranges. If no ranges are provided, all available templates will be loaded.

        Parameters:
        ----------
        template_repo : str
            The path to the repository containing the origin shower library (templates).
        xmax_range : Union[tuple, None], default=None
            The range of xmax values to load templates for. If None, all xmax values will be loaded.
        zenith_range : Union[tuple, None], default=None
            The range of zenith angles to load templates for. If None, all zenith angles will be loaded.
        randomize : bool, default=False
            Whether to randomly select one template per configuration. Default is False.

        Returns:
        --------
        None
        """
        # first check if the template repo exists
        if not os.path.exists(template_repo):
            raise FileNotFoundError(
                f"Template repository {template_repo} does not exist."
            )

        # get the path to all the templates
        path_to_templates = os.path.join(template_repo, self.site, self.backend)

        # get all configurations of xmax and zenith
        config_data = np.genfromtxt(
            os.path.join(path_to_templates, "config.txt"), delimiter=",", skip_header=1
        )
        all_xmaxs = config_data[:, 0]
        all_zeniths = config_data[:, 1]

        if xmax_range is None or zenith_range is None:
            self.logger.warning(
                "No xmax_range or zenith_range provided. Loading all available templates. This may be memory intensive (> 10 GB)."
            )

        if xmax_range is None:
            xmax_range = (min(all_xmaxs), max(all_xmaxs))
        if zenith_range is None:
            zenith_range = (min(all_zeniths), max(all_zeniths))

        all_template_paths = self.__get_all_templates(
            path_to_templates, xmax_range, zenith_range, randomize
        )
        self.origin_xmaxs = np.zeros(len(all_template_paths))
        self.origin_zeniths = np.zeros(len(all_template_paths))

        for itemp, template_path in enumerate(all_template_paths):
            synthesis = self.smiet.TemplateSynthesis(freq_ar=self.freq_ar)
            synthesis.load_template(template_path)

            self.syntheses.append(synthesis)
            self.origin_xmaxs[itemp] = synthesis.template_information["xmax"]
            self.origin_zeniths[itemp] = np.rad2deg(
                synthesis.template_information["zenith"]
            )

        # sort in increasing xmax
        sort_indices = np.argsort(self.origin_xmaxs)
        self.origin_xmaxs = self.origin_xmaxs[sort_indices]
        self.origin_zeniths = self.origin_zeniths[sort_indices]
        self.syntheses = [self.syntheses[i] for i in sort_indices]

        # also set the delta_t here, it will be the same for all origins
        self.delta_t = self.syntheses[0].delta_t

        # give a log message
        self.logger.info(
            f"Loaded {len(self.syntheses)} templates for synthesis from {template_repo} within xmax range {xmax_range} and zenith range {zenith_range}."
        )

        self.__loaded_templates = True

    def __get_all_templates(
        self: Self,
        path_to_templates: str,
        xmax_range: tuple,
        zenith_range: tuple,
        randomize: bool = False,
    ) -> list[str]:
        """
        Get all templates names that are within the range of xmax and zenith angles.

        Parameters:
        ----------
        path_to_templates : str
            The path to the templates.
        xmax_range : tuple
            The range of xmax values to filter templates.
        zenith_range : tuple
            The range of zenith angles to filter templates.
        randomize : bool, default=False
            Whether to randomly select one template per configuration. Default is False.
        """
        all_templates = []

        rng = np.random.default_rng()

        for xmax_zenith in os.listdir(path_to_templates):
            if not os.path.isdir(os.path.join(path_to_templates, xmax_zenith)):
                continue
            config_path = os.path.join(path_to_templates, xmax_zenith)

            # get xmax and zenith from naming structure
            xmax_str, zenith_str = xmax_zenith.split("_")
            xmax = float(xmax_str.replace("xmax", ""))
            zenith = float(zenith_str.replace("zen", ""))

            n_sims_per_config = len(os.listdir(config_path))

            if (
                xmax >= xmax_range[0]
                and xmax <= xmax_range[1]
                and zenith >= zenith_range[0]
                and zenith <= zenith_range[1]
            ):
                # take a random number, since there are some number of simulations per directory
                rand_sim = (
                    rng.integers(low=1, high=n_sims_per_config + 1) if randomize else 1
                )
                template_path = os.path.join(
                    path_to_templates, xmax_zenith, f"template_sim{rand_sim}.hdf5"
                )

                all_templates.append(template_path)

        # make sure that its not empty
        if len(all_templates) == 0:
            raise ValueError(
                f"No templates found in {template_path} for the given xmax range {xmax_range} and zenith range {zenith_range}."
            )

        return all_templates

    def load_coreas_showers(
        self: Self,
        coreas_shower_paths: list[str],
        gdas_atmosphere_paths: Union[list[str], None] = None,
    ):
        """
        Load target showers in the form of existing coreas simulations.

        Parameters:
        -----------
        coreas_shower_paths : list[str]
            A list of paths to each coreas simulation in HDF5 format. Each path corresponds to a single target shower.
        gdas_atmosphere_paths : Union[list[str], None], default=None
            The path to the GDAS atmosphere for each simulation. Default is None, where the model set globally is used. If provided, the entries and length of this list must match the length of coreas_shower_paths.
        """
        self.target_xmaxs = np.zeros(len(coreas_shower_paths))
        self.target_zeniths = np.zeros(len(coreas_shower_paths))

        for ishower, coreas_shower_path in enumerate(coreas_shower_paths):
            # make sure paths exist
            if not os.path.exists(coreas_shower_path):
                raise FileNotFoundError(
                    f"Coreas shower path {coreas_shower_path} does not exist."
                )

            gdas_atmosphere_path = (
                gdas_atmosphere_paths[ishower]
                if gdas_atmosphere_paths is not None
                else None
            )

            target_shower = self.smiet.CoreasShower(
                coreas_shower_path, gdas_file=gdas_atmosphere_path
            )

            self.target_xmaxs[ishower] = target_shower.xmax
            self.target_zeniths[ishower] = np.rad2deg(target_shower.zenith)
            self.targets.append(target_shower)

        # raise error if target xmaxs or zeniths are outside the range of origin showers
        if not self.__check_xmax_zenith_range():
            raise ValueError(
                "Target shower xmax or zenith values are outside the range of origin shower values. Load templates with a wider xmax / zenith range."
            )

        self.logger.info(
            f"Loaded {len(self.targets)} target showers from Coreas simulations."
        )
        self.__loaded_targets = True

    def __check_xmax_zenith_range(self: Self) -> bool:
        """
        Check whether the xmax and zenith of the target showers are within the range of the origin showers.
        """
        within_xmax = min(self.target_xmaxs) > min(self.origin_xmaxs) and max(
            self.target_xmaxs
        ) < max(self.origin_xmaxs)
        within_zenith = min(self.target_zeniths) > min(self.origin_zeniths) and max(
            self.target_zeniths
        ) < max(self.origin_zeniths)

        return within_xmax and within_zenith

    def generate_target_showers(
        self: Self,
        target_grammage: np.ndarray,
        target_longs: list[np.ndarray],
        target_zeniths: np.ndarray,
        target_azimuths: Union[np.ndarray, None] = None,
        target_cores: Union[np.ndarray, None] = None,
    ) -> None:
        """
        Generate and set target showers based on user-provided longitudinal profiles and geometries.

        Parameters:
        -----------
        target_grammage : np.ndarray
            The grammage slices corresponding to the provided longitudinal profiles. Should be a numpy array of shape (N,) where N is the number of grammage slices.
        target_longs : list[np.ndarray]
            A list of longitudinal profiles for each target shower. Each profile should be a numpy array of shape (N,) where N is the number of grammage slices. The number of grammage slices need not match those in the templates, as interpolation will be performed internally.
        target_zeniths : np.ndarray
            The zenith angles for each target shower. Should be a numpy array of shape (M,) where M is the number of target showers.
        target_azimuths : Union[np.ndarray, None], default=None
            The azimuth angles for each target shower. Should be a numpy array of shape (M,) where M is the number of target showers. If None, azimuths will be set to 0 degrees by default.
        target_cores : Union[np.ndarray, None], default=None
            The core positions for each target shower. Should be a numpy array of shape (M, 2) where M is the number of target showers and each entry contains the (x, y) core position in meters. If None, cores will be set to (0, 0) by default.
        """
        target_azimuths = (
            np.zeros(len(target_longs)) if target_azimuths is None else target_azimuths
        )
        target_cores = (
            np.zeros((len(target_longs), 3)) if target_cores is None else target_cores
        )

        if self.backend == "jax":
            target_zeniths = jnp.array(target_zeniths)
            target_azimuths = jnp.array(target_azimuths)
            target_cores = jnp.array(target_cores)
            target_grammage = jnp.array(target_grammage)

        self.target_xmaxs = np.zeros(len(target_longs))
        for ishower, long in enumerate(target_longs):
            # should try to fit GH to get xmax and nmax
            # since argmax and max are limited by the
            # grid size.
            # for now its fine, we will adjust this later
            # in any case, for the np version the GH parameters will be fit
            xmax = target_grammage[long.argmax()]
            nmax = long.max()

            target_shower = self.smiet.Shower()

            target_shower.set_parameters(
                grammages=target_grammage,
                long_profile=long,
                xmax=xmax,
                nmax=nmax,
                zenith=target_zeniths[ishower],
                azimuth=target_azimuths[ishower],
                core=target_cores[ishower],
                magnetic_field=self.magnet,
            )

            self.target_xmaxs[ishower] = target_shower.xmax
            self.targets.append(target_shower)

        self.target_zeniths = np.rad2deg(target_zeniths)

        # raise error if target xmaxs or zeniths are outside the range of origin showers
        if not self.__check_xmax_zenith_range():
            raise ValueError(
                "Target shower xmax or zenith values are outside the range of origin shower values. Load templates with a wider xmax / zenith range."
            )

        self.logger.info(
            f"Generated {len(self.targets)} target showers from user-provided longitudinal profiles."
        )
        self.__loaded_targets = True

    def synthesize_showers(
        self: Self,
        interpolation: bool = True,
        slices: bool = False,
    ) -> None:
        """
        Run the synthesis for all loaded target showers.

        This will perform the interpolation synthesis method by default, where the combination of two synthesis, weighted by some function, will be used. It is most adviced to use this method, so only disable the interpolation if you need to.

        Parameters:
        -----------
        interpolation : bool, default=True
            Whether to use interpolation synthesis method. If False, the nearest neighbor synthesis object will be used.
        slices : bool, default=False
            Whether to synthesise traces for all slices. If False, only synthesises at ground level.
        """
        # first make sure all flags are checked
        if not self.__loaded_templates:
            raise RuntimeError(
                "Templates have not been loaded. Please load templates before running synthesis."
            )
        if not self.__loaded_targets:
            raise RuntimeError(
                "Target showers have not been loaded. Please load target showers before running synthesis."
            )
        if slices:
            self.logger.warning(
                "Synthesis with slices=True is enabled. Setting interpolation=False since it does not work with interpolated synthesis."
            )
            interpolation = False

        for ishower, target_shower in enumerate(self.targets):
            # find the closest origin shower, the smaller one
            # for zenith, synthesis only works if theta_origin > theta_target
            # otherwise synthesis will artificially generate showers
            diff_xmaxs = self.origin_xmaxs - target_shower.xmax
            diff_zenith = self.origin_zeniths - target_shower.zenith
            lower_indices = np.where(np.logical_and(diff_xmaxs <= 0, diff_zenith > 0))[
                0
            ]
            upper_indices = np.where(np.logical_and(diff_xmaxs > 0, diff_zenith > 0))[0]

            if len(lower_indices) == 0 or len(upper_indices) == 0:
                raise ValueError(
                    f"Target shower xmax {target_shower.xmax} is outside the range of origin shower xmax values."
                )

            lower_index = lower_indices[-1]
            upper_index = upper_indices[0]

            lower_synth = self.syntheses[lower_index]
            upper_synth = self.syntheses[upper_index]

            # here we get the synthesised traces on a star shape
            # from a vxvxB synthesis
            target_shower.transform_profile_to_origin(lower_synth.grammages)
            synth_traces_geo_vvB, synth_traces_ce_vvB = (
                lower_synth.map_template_to_slices(target_shower)
                if slices
                else lower_synth.map_template(target_shower)
            )
            synth_traces_geo, synth_traces_ce, ant_pos_starshape, ant_labels = (
                self.smiet.utilities.reconstruct_starshape_from_vxvxB(
                    synth_traces_geo_vvB,
                    synth_traces_ce_vvB,
                    lower_synth.antenna_information["position_showerplane"][:, 1],
                )
            )

            synth_traces = np.array([synth_traces_geo, synth_traces_ce])

            if interpolation:
                target_shower.transform_profile_to_origin(upper_synth.grammages)
                synth_traces_geo_upper_vvB, synth_traces_ce_upper_vvB = (
                    upper_synth.map_template(target_shower)
                )
                synth_traces_geo_upper, synth_traces_ce_upper, _, _ = (
                    self.smiet.utilities.reconstruct_starshape_from_vxvxB(
                        synth_traces_geo_upper_vvB,
                        synth_traces_ce_upper_vvB,
                        upper_synth.antenna_information["position_showerplane"][:, 1],
                    )
                )

                synth_traces_upper = np.array(
                    [synth_traces_geo_upper, synth_traces_ce_upper]
                )

                # calculate weight factor for interpolation
                # TODO: vectorize this!
                w_xmax = (target_shower.xmax - self.origin_xmaxs[lower_index]) / (
                    self.origin_xmaxs[upper_index] - self.origin_xmaxs[lower_index]
                )

                synth_traces *= 1 - w_xmax
                synth_traces += w_xmax * synth_traces_upper

            if not slices:
                synth_traces = synth_traces[..., np.newaxis]  # add slice axis

            self.synth_traces.append(synth_traces)
            self.synth_times.append(
                np.repeat(lower_synth.get_time_axis(), repeats=8, axis=0)
            )  # 8 is number of antennas in star shape, hardcoded
            self.antenna_positions.append(ant_pos_starshape)

        # antenna labels only need to be assigned once
        self.antenna_labels = ant_labels

        self.logger.info(f"Synthesized {len(self.synth_traces)} target showers.")

        self.__synthesized_showers = True

    def __find_target_idx(self: Self, xmax: float, zenith: float) -> int:
        """
        Find the index corresponding to the target shower, given a xmax and zenith value

        Parameters:
        -----------
        xmax : float
            The xmax of the target shower to find.
        zenith : float
            The zenith angle of the target shower to find.
        """
        # find the target shower
        target_index = None
        for ishower, target_shower in enumerate(self.targets):
            if np.isclose(target_shower.xmax, xmax) and np.isclose(
                np.rad2deg(target_shower.zenith), zenith
            ):
                target_index = ishower
                break

        if target_index is None:
            raise ValueError(
                f"Target shower with xmax {xmax} and zenith {zenith} not found."
            )
        return target_index
    
    def __get_sample_axis(
        self: Self,
        ant_name: Union[float, None] = None,
        slice_grammage: Union[float, None] = None,
    ):
        """Get the sample axis based on the provided antenna name and slice grammage. Could be implemented more efficiently!"""
        if ant_name is not None:
            return 1
        elif slice_grammage is not None:
            return 2
        elif (ant_name is not None) and (slice_grammage is None):
            return 1
        elif (ant_name is None) and (slice_grammage is not None):
            return 1
        elif (ant_name is None) and (slice_grammage is None):
            return 2
        else:
            return 2
        

    def get_traces(
        self: Self,
        xmax: float,
        zenith: float,
        ant_name: Union[str, None] = None,
        geoce: bool = True,
        slice_grammage: Union[float, None] = None,
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        Get the time traces for a specific target shower.

        Parameters:
        -----------
        xmax : float
            The xmax of the target shower to retrieve traces for.
        zenith : float
            The zenith angle of the target shower to retrieve traces for.
        ant_name : Union[str, None], default=None
            The name of the antenna to retrieve traces for. If None, all antennas will be returned.

            The nomenclature for antenna names is given as 'ant_r{r}_phi{phi}`, where `phi` is the azimuthal angle in degrees and `r` is the radial distance in meters.

            The avaliable antenna names can be found in the `antenna_labels` attribute after synthesis.
        geoce : bool, default=True
            Whether to return traces in the geomagnetic + charge-excess emissions. If False, returns in the shower plane coordinate system.
        slice_grammage : Union[float, None], default=None
            The grammage slice to retrieve traces for. If None, retrieves trace at the ground. If traces are only synthesised at the ground, the argument will default to None.

        Returns:
        --------
        tuple[np.ndarray, np.ndarray, np.ndarray]
            A tuple containing:
            - traces: np.ndarray of shape (N_pol, N_antennas, N_timebins, n_slices)
                The traces either in geo/ce components (if geoce is True) or in vxB/vxvxB components (if geoce is False).
            - times: np.ndarray of shape (N_antennas, N_timebins, n_slices)
                The time axes corresponding to each antenna trace.
            - antenna_positions: np.ndarray of shape (N_antennas, 3)
                The positions of each antenna in meters in the shower plane.

            If ant_name is not None, N_antennas will be 1 and the corresponding antenna trace & time will be returned.

        """
        if not self.__synthesized_showers:
            raise RuntimeError(
                "Showers have not been synthesized yet. Please run synthesis before retrieving traces."
            )

        target_index = self.__find_target_idx(xmax, zenith)

        synth_traces = self.synth_traces[target_index]
        trace_times = self.synth_times[target_index]
        ant_positions = self.antenna_positions[target_index]

        if not geoce:
            # convert to vxB / vxvxB
            synth_traces = self.smiet.utilities.geo_ce_to_e(
                synth_traces[0],
                synth_traces[1],
                ant_positions[:, 0],
                ant_positions[:, 1],
            )

            # match the shape of {Npol, Nant, Ntime}
            synth_traces = np.swapaxes(synth_traces.T, 3, 1)

        if ant_name is not None:
            if ant_name not in self.antenna_labels:
                raise ValueError(
                    f"Antenna name {ant_name} not found in available antenna labels."
                )
            ant_index = self.antenna_labels.index(ant_name)

            synth_traces = np.take(synth_traces, indices=ant_index, axis=1)
            trace_times = np.take(trace_times, indices=ant_index, axis=0)
            ant_positions = np.take(ant_positions, indices=ant_index, axis=0)

        # force slice grammage to None if only synthesised at ground
        if synth_traces.shape[-1] == 1:
            self.logger.warning(
                "Traces were only synthesised at ground level. Ignoring provided slice_grammage and returning ground level traces."
            )
            slice_grammage = None

        if slice_grammage is not None:
            # find the closest slice
            grammages = self.targets[target_index].grammages
            slice_index = np.argmin(np.abs(grammages - slice_grammage))
            synth_traces = np.take(synth_traces, indices=slice_index, axis=-1)
        else:
            synth_traces = np.sum(synth_traces, axis=-1)

        return synth_traces, trace_times, ant_positions

    def get_spectrum(
        self: Self,
        xmax: float,
        zenith: float,
        ant_name: Union[str, None] = None,
        geoce: bool = True,
        slice_grammage: Union[float, None] = None,
    ) -> tuple[np.ndarray, np.ndarray]:
        """
        Get the frequency spectrum for a given synthesised shower and antenna name (optional).

        Parameters:
        -----------
        xmax : float
            The xmax of the target shower to retrieve spectrum for.
        zenith : float
            The zenith angle of the target shower to retrieve spectrum for.
        ant_name : Union[str, None], default=None
            The name of the antenna to retrieve spectrum for. If None, all antennas will be returned.
        geoce : bool, default=True
            Whether to return traces in the geomagnetic + charge-excess emissions. If False, returns in the shower plane coordinate system.
        slice_grammage : Union[float, None], default=None
            The grammage slice to retrieve traces for. If None, retrieves traces at all slices. If traces are only synthesised at the ground, then that will be returned instead.

        Returns:
        --------
        tuple[np.ndarray, np.ndarray]
            A tuple containing:
            - spectrum: np.ndarray of shape (N_pol, N_antennas, N_freqbins, Nslices)
                The frequency spectrum either in geo/ce components (if geoce is True) or in vxB/vxvxB components (if geoce is False).
            - freq_grid: np.ndarray of shape (N_freqbins,)
                The frequency axis corresponding to the spectrum.

        If ant_name is not None, N_antennas will be 1 and the corresponding antenna spectrum will be returned.
        """
        # basically call the get_traces function to get the relevant traces
        synth_traces, trace_times, _ = self.get_traces(
            xmax, zenith, ant_name=ant_name, geoce=geoce, slice_grammage=slice_grammage
        )
        sample_idx = self.__get_sample_axis(ant_name, slice_grammage)
        # then perform an FFT to get the frequency spectrum
        synth_spect = np.fft.rfft(synth_traces, axis=sample_idx, norm="ortho", n=synth_traces.shape[sample_idx])
        synth_freq_grid = np.fft.rfftfreq(synth_traces.shape[sample_idx], d=self.delta_t)

        return synth_spect, synth_freq_grid

    def get_fluence(
        self: Self, xmax: float, zenith: float, geoce: bool = True, slice_grammage: Union[float, None] = None,
    ) -> tuple[np.ndarray, np.ndarray]:
        """
        Get the fluence for a given synthesised shower.

        Parameters:
        ----------
        xmax : float
            The xmax of the target shower to retrieve fluence for.
        zenith : float
            The zenith angle of the target shower to retrieve fluence for.
        geoce : bool, default=True
            Whether to return traces in the geomagnetic + charge-excess emissions. If False, returns in the shower plane coordinate system.
        slice_grammage : Union[float, None], default=None
            The grammage slice to retrieve traces for. If None, retrieves traces at all slices. If traces are only synthesised at the ground, then that will be returned instead.
        """
        # again call the get_traces function to get the relevant traces
        # but here call without antenna name to get all antennas
        synth_traces, _, ant_positions = self.get_traces(
            xmax, zenith, ant_name=None, geoce=geoce, slice_grammage=slice_grammage
        )
        sample_idx = self.__get_sample_axis(ant_name=None, slice_grammage=slice_grammage)
        # then perform a fluence calculation
        conversion_factor_integrated_signal = (
            2.65441729e-3 * 6.24150934e18
        )  # to convert to eV/m^2 from (V/m)^2 s
        synth_fluence = (
            np.sum(synth_traces**2, axis=(0, sample_idx))
            * (self.delta_t / units.s)
            * conversion_factor_integrated_signal
        )
        return synth_fluence, ant_positions

    def get_truths(
        self: Self,
        xmax: float,
        zenith: float,
        mode: str = "fluence",
        geoce: bool = True,
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        Return the truth values for a given synthesised shower.

        Note that this is only available if the target showers were loaded from Coreas simulations. Otherwise, an error will be raised.

        Parameters:
        ----------
        xmax : float
            The xmax of the target shower to retrieve truths for.
        zenith : float
            The zenith angle of the target shower to retrieve truths for.
        mode : str, default="fluence"
            The mode of truth values to retrieve. Options are "fluence", "spectrum", or "traces".
        geoce : bool, default=True
            Whether to return traces in the geomagnetic + charge-excess emissions. If False, returns in the shower plane coordinate system.
        """
        # first check if target showers are coreas showers
        if not isinstance(self.targets[0], self.smiet.CoreasShower):
            raise RuntimeError(
                "Target showers are not Coreas showers. Truth values are only available for Coreas showers."
            )
        target_index = self.__find_target_idx(xmax, zenith)
        ant_names = self.targets[target_index].antenna_names
        ant_pos_showerplane = self.targets[
            target_index
        ].get_antenna_position_showerplane(ant_names)

        if self.backend == "numpy":
            if geoce:
                traces_truth_geo, traces_truth_ce, traces_times = self.targets[
                    target_index
                ].get_traces_geoce(
                    ant_names=ant_names,
                    return_times=True,
                )
                traces_truth = np.array([traces_truth_geo, traces_truth_ce])
            else:
                traces_truth_vB, traces_truth_vvB, traces_times = self.targets[
                    target_index
                ].get_traces_vB_vvB(
                    ant_names=ant_names,
                    return_times=True,
                )
                traces_truth = np.array([traces_truth_vB, traces_truth_vvB])
        else:
            raise NotImplementedError("will implement it later.")

        if mode == "traces":
            return traces_truth, traces_times, ant_pos_showerplane
        elif mode == "spectrum":
            # then perform an FFT to get the frequency spectrum
            spect_truth = np.fft.rfft(traces_truth, axis=-1, norm="ortho")
            freq_grid = np.fft.rfftfreq(traces_times.shape[-1], d=self.targets[target_index].delta_t)
            return spect_truth, freq_grid, ant_pos_showerplane
        elif mode == "fluence":
            # then perform a fluence calculation
            conversion_factor_integrated_signal = 2.65441729e-3 * 6.24150934e18
            fluence_truth = (
                np.sum(traces_truth**2, axis=(0, -1))
                * (self.targets[target_index].delta_t / units.s)
                * conversion_factor_integrated_signal
            )
            return fluence_truth, ant_pos_showerplane
        else:
            raise ValueError(
                f"Mode {mode} not recognized. Options are 'fluence', 'spectrum', or 'traces'."
            )

    def synthesize_with_single_origin(
        self: Self,
        origin_shower_path: str,
        target_long: np.ndarray,
        target_grammage: np.ndarray,
        target_zenith: float,
        target_azimuth: float = 0.0,
        target_core: np.ndarray = np.array([0.0, 0.0, 0.0]),
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        Perform synthesis with a single origin shower.

        This function is mostly for testing purposes, but is available if users want to do this.

        Parameters:
        ----------
        origin_shower_path : str
            The path to the origin shower to use for synthesis.
        target_long : np.ndarray
            The longitudinal profile of the target shower.
        target_grammage : np.ndarray
            The grammage slices corresponding to the provided longitudinal profile.
        target_zenith : float
            The zenith angle of the target shower.
        target_azimuth : float, default=0.0
            The azimuth angle of the target shower. Default is 0.0 degrees.
        target_core : np.ndarray, default=np.array([0.0, 0.0, 0.0])
            The core position of the target shower in meters. Default is (0.0, 0.0, 0.0).

        Returns:
        --------
        np.ndarray
            The synthesised time traces for the target shower and the time axes.
        """
        # load an origin shower
        if not os.path.exists(origin_shower_path):
            raise FileNotFoundError(
                f"Origin shower path {origin_shower_path} does not exist."
            )
        origin_shower = self.smiet.SlicedShower(origin_shower_path)

        # make a template
        synthesis = self.smiet.TemplateSynthesis(freq_ar=self.freq_ar)
        synthesis.make_template(origin_shower)

        # generate a target shower
        target_shower = self.smiet.Shower()
        target_shower.set_parameters(
            grammages=target_grammage,
            long_profile=target_long,
            xmax=target_grammage[target_long.argmax()],
            nmax=target_long.max(),
            zenith=np.deg2rad(target_zenith),
            azimuth=np.deg2rad(target_azimuth),
            core=target_core,
            magnetic_field=self.magnet,
        )

        target_shower.transform_profile_to_origin(origin_shower.grammages)

        # map the template
        synth_traces_geo, synth_traces_ce = synthesis.map_template(target_shower)

        # return the synthesised traces on the vxvxB arm (or starshape?)
        return synth_traces_geo, synth_traces_ce, synthesis.get_time_axis()

    def target_properties(self: Self) -> dict:
        """
        Return the properties of each target shower
        """
        return {
            "xmaxs": self.target_xmaxs,
            "zeniths": self.target_zeniths,
            "long": [self.targets[i].long_profile for i in range(len(self.targets))],
            "grammage": [self.targets[i].grammages for i in range(len(self.targets))],
        }

    def origin_properties(self: Self) -> dict:
        """
        Return the properties of each origin shower
        """
        return {
            "xmaxs": self.origin_xmaxs,
            "zeniths": self.origin_zeniths,
            "long": [
                self.syntheses[i].template_information["long_profile"]
                for i in range(len(self.syntheses))
            ],
            "grammage": [
                self.syntheses[i].grammages for i in range(len(self.syntheses))
            ],
        }
