from __future__ import annotations

import logging
import os
from functools import partial
from typing_extensions import Self, Tuple, Union

import h5py
import jax

jax.config.update("jax_enable_x64", True)

import numpy as np
import jax.numpy as jnp

from .io import BaseShower, SlicedShower
from smiet import units
from jax_radio_tools import hp_spherical_to_cartesian
from jax_radio_tools.atmosphere import Atmosphere
from .utilities.geometry import angle_between
from .utilities.batched_interpolation import batched_interp1d


def amplitude_function(
    params: jax.typing.ArrayLike,
    frequencies: jax.typing.ArrayLike,
    d_noise: float = 0.0,
) -> jax.Array:
    """
    Calculate the amplitude frequency spectrum corresponding to the parameters `params`.

    Parameters
    ----------
    params : jax.typing.ArrayLike
        The spectral parameters. If it is a multidimensional array, the first dimension must contain the parameters.
    frequencies : jax.typing.ArrayLike
        The values of the frequencies to evaluate - remove central frequency beforehand!
    d_noise : float, default=0.0
        The noise floor level

    Returns
    -------
    The evaluated amplitude frequency spectrum with shape VIEW x FREQ x SLICES
    """
    if len(params.shape) == 1:
        params = params.reshape((params.shape[0], 1))
    return (
        params[0, :, None, :]
        * jnp.exp(
            params[1, :, None, :] * frequencies[None, :, None]
            + params[2, :, None, :] * frequencies[None, :, None] ** 2
        )
        + d_noise
    )


# @jax.jit
def get_spectra(
    xmax: float,
    spectral_coeffs: jax.typing.ArrayLike,
    grammages: jax.typing.ArrayLike,
    frequencies: list,
) -> jax.typing.ArrayLike:
    r"""
    Retrieve the amplitude spectra at the specified frequencies, for a given :math:`\Delta X_{max}`.

    Parameters
    ----------
    xmax : float
        The maximum of the atmospheric depth in g/cm^2
    spectral_coeffs : jax.typing.ArrayLike
        the spectral coefficients stored with shape
        {GEO, CE} x VIEW x FREQ
    grammages : jax.typing.ArrayLike
        an array of atmospheric depths in g/cm^2
    frequencies : list
        The list of frequencies at which to evaluate the spectra (after filtering)

    Returns
    -------
    spectrum : jax.typing.ArrayLik
        the amplitude spectrum in shape
        {GEo x CE} x VIEW x FREQ x SLICES
    """
    # expand dimensions for grammage to VIEW x FREQ x SLICES x {GEO, CE}
    gram = jnp.expand_dims(grammages, axis=(0, 1, 3))
    # NB: flip in the spectral parameters is necessary since
    # jnp.polyval evaluates from from the HIGHEST power,
    # whereas np.polynomial.polynomial.polyval evaluates from
    # the LOWEST power.
    # since the amplitude function follows the structure from the
    # numpy code, we need another flip after evaluating the polynomial
    spectral_params = jnp.flip(
        jnp.polyval(jnp.flip(spectral_coeffs.T)[:, :, :, None, :], gram - xmax),
        axis=(0, 1),
    )  # 3 x viewing angles x SLICES x {GEO, CE}

    frequencies = jnp.array(frequencies)

    spectra = spectral_params[0, :, None, :, :] * jnp.exp(
        spectral_params[1, :, None, :, :] * frequencies[None, :, None, None]
        + spectral_params[2, :, None, :, :] * frequencies[None, :, None, None] ** 2
    )
    return jnp.moveaxis(spectra, 3, 0)


@partial(jax.jit, static_argnames=["outshape"])
def get_correction_factors(
    spectral_params: jax.typing.ArrayLike,
    ant_v_angles: jax.typing.ArrayLike,
    v_angles_grid: list,
    outshape: tuple,
) -> jax.Array:
    r"""
    Get the correction factors from the spectral parameters.

    Parameters
    ----------
    spectral_params : jax.typing.ArrayLike
        the spectral parameters for the geomagnetic and charge excess emission
        in shape {GEO x CE} x VIEW x FREQ x SLICES
    ant_v_angles : jax.typing.ArrayLike
        the viewing angles for the particular geometry & atmosphere
    v_angles_grid : list, static
        the grid of viewing angles used for interpolation
    outshape : tuple, static
        the shape for the correction factor array
        shape should be {GEO x CE} x ANTS x FREQ x SLICES

    Returns
    -------
    correction_factors : jax.typing.ArrayLike
        the corrections factors in shape
        {GEO x CE} x ANTS x FREQ x SLICES
    """
    # calculate the correction factors
    # shape is {GEO x CE} x ANTS x FREQ x SLICES
    correction_factors = jnp.zeros(outshape)

    for icomp, comp in enumerate(spectral_params):
        # Temporary fix for bad viewing angles in large freq range
        # Here we take log for more accurate interpolation (after taking abs to avoid -log values)
        corr_fact = jnp.log(
            jnp.where(jnp.abs(comp[:-2]) <= 1e-20, 1.0, jnp.abs(1.0 / comp[:-2]))
        )

        # interpolate over viewing angles
        correction_factors = correction_factors.at[icomp, ...].set(
            jnp.exp(
                batched_interp1d(
                    jnp.array(ant_v_angles), jnp.array(v_angles_grid[:-2]), corr_fact
                )
            )
        )

    return correction_factors


class TemplateSynthesis:
    """
    This class is the main interface for synthesising pulses.

    The main workflow consists of creating an instance and reading in the spectral parameters
    from a particular file.
    Then one can pass a ``SlicedShower`` as an origin to the :meth:`TemplateSynthesis.make_template`
    method, which construct all the necessary arrays to perform the synthesis.
    To perform the synthesis, you then call the :meth:`TemplateSynthesis.map_template` method with a
    ``Shower`` instance, which should be equipped with a longitudinal profile and the parameters
    from a Gaisser-Hillas fit.

    Parameters
    ----------
    freq_ar : list, default=[30, 500, 100]
        The frequency range (min, max, central) in MHz for which to load the spectral parameters.
        Currently available options are:
            - [30, 500, 100] : can be used for most general purposes and experiments
            - [30, 80, 50] : a smaller sub-band which can be used for experiments such as LOFAR, AERA, etc
        The default value is [30, 500, 100].

        NOTE: The frequency range provided should not be multiplied by units.MHz, as this is done internally!


    Attributes
    ----------
    has_spectral_coefficients : bool
        Whether the spectral coefficients have been loaded
    atm : radiotools.atmosphere.models.Atmosphere
        The atmosphere model from the origin shower, used to make the template
    spectral_params : jax.typing.ArrayLike
        The spectral coefficients with shape {GEO, CE} x VIEW x FREQ
    viewing_angles : jax.typing.ArrayLike
        The viewing angles (in units of Cherenkov angle) for which we have the spectral coefficients
    amplitudes : jax.typing.ArrayLike
        The amplitude spectrum for each antenna, frequency, slice, and emission component
    phases : jax.typing.ArrayLike
        The phase spectrum for each antenna, frequency, slice, and emission component
    frequencies : np.ndarray
        The frequencies corresponding to the Fourier transform of the origin
    frequency_range : tuple of float
        The minimum, maximum and central frequency for which the spectral parameters where fitted.
        This is read from the HDF5 file containing the spectral parameters.
    freq_range_mask : np.ndarray
        A boolean mask to select only the frequencies within the frequency range
    truncated_frequencies : np.ndarray
        The frequencies after truncating all frequencies outside the frequency range
    n_samples : int
        The number of time samples in the traces
    delta_t : float
        The time step between each time sample in the traces
    antenna_information : np.ndarray
        The antenna information, containing the time axis, position, position in the shower plane, and name of each antenna
    n_antennas : int
        The number of antennas in the template
    template_information : dict
        The information about the template, such as the name, geometry, magnetisation, :math:`X_{max}`, core and creation time
    slices : list
        The list of slices in the template, each as a :obj:`SliceSynthesis` object
    grammages : list
        The grammage at the bottom of each slice

    Notes
    -----
    It is implicitly assumed that the origin shower's longitudinal profile is sampled with the same step
    size as the antennas are configured to observe (i.e. if the antennas are set up to observer slices
    with a thickness of 5 g/cm2, the longitudinal profile should also be sampled with a step size of
    5 g/cm2 - which is set using the LONGI keyword in the CORSIKA input file).
    """

    def __init__(
        self: Self,
        freq_ar: list = [30, 500, 100],
        ce_linear : bool = True
    ) -> None:
        self.logger = logging.getLogger("smiet.jax.synthesis.TemplateSynthesis")

        self.atm = None

        # spectral parameters
        self.has_spectral_coefficients = None
        self.spectral_params = None  # {GEO, CE} x VIEW x FREQ
        self.viewing_angles = None  # in units of Cherenkov angle, shape VIEW
        self.amplitudes = None  # {GEO, CE} x ANT x FREQ x SLICES
        self.phases = None  # {GEO, CE} x ANT x FREQ x SLICES

        # frequency properties
        self.frequency_range = None
        self.frequencies = None
        self.freq_range_mask = None
        self.truncated_frequencies = None
        self.n_samples = None  # time bins, not frequency bins
        self.delta_t = None

        # antenna properties, now all contained in a
        # single dictionary
        self.antenna_information = {
            "name": None,
            "position": None,
            "position_showerplane": None,
            "time_axis": None,
        }
        self.n_antennas = None

        # slice properties
        self.template_information = {
            "zenith": None,
            "azimuth": None,
            "magnet": None,
            "xmax": None,
            "nmax": None,
            "core": None,
            "long_profile": None,
        }
        self.grammages = None
        self.nr_slices = None
        

        # read spectral parameters from file
        self.logger.info(
            f"Loading spectral coefficients with frequency range {freq_ar} MHz"
        )
        spectral_filename = f"spectral_parameters_{int(freq_ar[0])}_{int(freq_ar[1])}_{int(freq_ar[2])}.hdf5"
        self.read_spectral_file(spectral_filename, ce_linear=ce_linear)

    def read_spectral_file(self: Self, filename: str, ce_linear : bool = True) -> None:
        """
        Read spectral parameters from a file with `filename` in the spectral_parameters/ directory.

        Parameters
        ----------
        filename : str
           The name of the spectral parameters file
        """
        path_to_file = os.path.join(
            os.path.dirname(os.path.dirname(__file__)), "spectral_parameters", filename
        )

        if not os.path.exists(path_to_file):
            raise FileNotFoundError(
                f"Filename {filename} does not exist in the spectral_parameters/ directory."
                "Did you provide the correct frequency range?"
            )

        # read spectral parmaeters from hdf5 file
        with h5py.File(path_to_file) as spectral_file:
            self.frequency_range = tuple(
                spectral_file["/Metadata/Frequency_MHz"][:] * units.MHz
            )
            self.viewing_angles = np.array(spectral_file["/Metadata/ViewingAngles"][:])

            geo = np.array(spectral_file["SpectralFitParams/GEO"][:])
            ce = np.array(spectral_file["SpectralFitParams/CE"][:])
            if ce_linear:
                ce = np.array(spectral_file["SpectralFitParams/CE_LIN"][:])
            # self.spectral_params[2, ...] = np.array(
            #     spectral_file["SpectralFitParams/CE_LIN"][:]
            # )

            self.spectral_coeffs = np.array([geo, ce])

            self.has_spectral_coefficients = True

        self.logger.debug(f"Loaded in the spectral coefficients from {filename}")

    def _initialise_shower_properties(self: Self, shower: SlicedShower) -> None:
        """
        Initialise the parameters from the origin shower.

        This is just a convenience function to make the make_template function look better.

        Parameters
        ----------
        shower : SlicedShower
            The sliced shower object conatining antennas information
        """
        self.antenna_information["name"] = list(shower.antenna_array["name"])
        self.antenna_information["position"] = jnp.asarray(shower.antenna_array["position"])
        self.antenna_information["position_showerplane"] = jnp.asarray(
            shower.get_antenna_position_showerplane()
        )
        self.antenna_information["time_axis"] = shower.trace_times
        self.n_antennas = len(shower.antenna_array["name"])

        self.atm = shower.atm

        self.atm.obs_lvl = (
            shower.core[2] / units.m
        )  # (need to convert from internal units to m)

        # Calculate frequencies from shower
        self.frequencies = jnp.fft.rfftfreq(shower.trace_length, d=shower.delta_t)

        # finally store the number of samples that we have
        self.delta_t = shower.delta_t
        self.n_samples = shower.trace_length

        # gather shower parameters to compute the relevant shower properties
        self.grammages = jnp.asarray(shower.grammages)
        self.nr_slices = shower.nr_slices

        shower.long_profile = jnp.asarray(shower.long_profile)

    def make_template(self: Self, origin: SlicedShower) -> None:
        """
        Process a ``SlicedShower`` into a template.

        Parameters
        ----------
        origin : smiet.jax.io.sliced_shower.SlicedShower
            The origin shower
        """
        if not self.has_spectral_coefficients:
            raise RuntimeError(
                "Please make sure the spectral coefficients are loaded before making a template"
            )

        # initialise all antenna parameters
        self._initialise_shower_properties(origin)

        # initialise parameters
        self.freq_range_mask = np.logical_and(
            self.frequency_range[0] <= self.frequencies,
            self.frequencies <= self.frequency_range[1],
        )
        self.truncated_frequencies = (
            self.frequencies[self.freq_range_mask] - self.frequency_range[2]
        )

        # computing shower properties (distance, correction factors, normalisation factors)
        # we compute these first since they are static and need not be computed repeatedly
        # in future implementations
        ant_viewing_angles, ant_distances_origin = (
            self._calculate_distance_viewing_angles(
                origin.zenith, origin.azimuth, origin.core
            )
        )
        norm_factors_origin = self._calculate_normalisation_factor(origin)
        # remove slices with few particles
        norm_factors_origin = jnp.where(
            origin.long_profile < 0.001 * origin.nmax, 0.0, norm_factors_origin
        )

        spectral_params_origin = self._get_spectra(origin.xmax)
        correction_factors_origin = self._get_correction_factors(
            spectral_params_origin, ant_viewing_angles
        )

        # set inverse of normalisation factor to zero where particles are too small so as to not process those slices
        norm_factors_origin_inv = jnp.where(
            norm_factors_origin == 0, 0.0, 1.0 / norm_factors_origin
        )

        # get traces from the origin shower
        # {GEO, CE} x ANT x SAMPLES x SLICES
        origin_traces = jnp.array(origin.get_traces_geoce())

        # RFFT traces to frequency domain, half of traces (+1) due to taking only the real part of FT
        # {GEO, CE} x ANT x FREQ x SLICES
        origin_traces_fft = jnp.fft.rfft(origin_traces, norm="ortho", axis=2)
        origin_traces_fft *= (
            ant_distances_origin[None, :, None, :]
            * norm_factors_origin_inv[:, None, None, :]
        )

        # amplitude and phase arrays
        # {GEO, CE} x ANTENNAS x FREQUENCIES x SLICES
        self.amplitudes = jnp.abs(origin_traces_fft) * correction_factors_origin
        self.phases = jnp.angle(origin_traces_fft)

        # set the vxB arm to zero
        ant_on_vxB = self.antenna_information["position_showerplane"][:, 1] < 1e-3
        self.amplitudes = self.amplitudes.at[:, ant_on_vxB, :, :].set(0.0)
        self.phases = self.phases.at[:, ant_on_vxB, :, :].set(0.0)

        # store the shower properties in template_information
        self.template_information = {
            "name" : origin.name,
            "zenith": origin.zenith,  # in radians
            "azimuth": origin.azimuth,  # in radians
            "magnet": origin.magnet,  # in Gauss
            "xmax": origin.xmax,
            "nmax": origin.nmax,
            "core": origin.core,
            "long_profile": origin.long_profile,
        }

    def _get_spectra(self: Self, xmax: float) -> jax.Array:
        """
        Wrap the JAX jitted function for the computation of the spectra.

        Parameters
        ----------
        xmax : float
            the atmospheric depth at shower maximum
        """
        return jax.jit(get_spectra)(
            xmax,
            self.spectral_coeffs,
            self.grammages,
            list(self.truncated_frequencies),
        )

    def _get_correction_factors(
        self: Self,
        spectral_params: jax.typing.ArrayLike,
        ant_viewing_angles: jax.typing.ArrayLike,
    ) -> jax.Array:
        """
        Wrap the JAX jitted function to get the correction factors.

        Parameters
        ----------
        spectral_params : jax.typing.ArrayLike
            the computed spectral parameters
        ant_viewing_angles : jax.typing.ArrayLike
            the viewing angles for this particular shower

        Returns
        -------
        the correction factors
        """
        corr_fact_shape = (
            2,
            self.n_antennas,
            len(self.truncated_frequencies),
            self.nr_slices,
        )
        correction_factors = jnp.zeros(
            (
                2,
                self.n_antennas,
                len(self.frequencies),
                self.nr_slices,
            )
        )
        correction_factors = correction_factors.at[..., self.freq_range_mask, :].set(
            get_correction_factors(
                spectral_params,
                ant_viewing_angles,
                self.viewing_angles,
                corr_fact_shape,
            )
        )
        return correction_factors
    
    def map_template_to_slices(self: Self, target: BaseShower) -> jax.Array:
        """
        Map the template to a target profile, represented in a target BaseShower.

        Calculates the trace for every antenna present in the template at all slices.

        Parameters
        ----------
        target :  smiet.jax.io.base_shower.BaseShower
            The target BaseShower object, containing the longitudinal profile,
            zenith, azimuth, geomagnetic angle, xmax and nmax

        Returns
        -------
        total_synth : jax.Array
            The synthesised geomagnetic & charge-excess trace for all antennas.
            Shape is {GEO, CE} x ANT x SAMPLES
        """
        # computing shower properties (distance, correction factors, normalisation factors)
        # we compute these first since they are static and need not be computed repeatedly
        # in future implementations
        ant_viewing_angles, ant_distances_target = (
            self._calculate_distance_viewing_angles(
                target.zenith, target.azimuth, target.core
            )
        )
        norm_factors_target = self._calculate_normalisation_factor(target)
        # remove slices with few particles
        norm_factors_target = jnp.where(
            target.long_profile < 0.001 * target.nmax, 0, norm_factors_target
        )

        spectral_params_target = self._get_spectra(target.xmax)
        correction_factors_target = self._get_correction_factors(
            spectral_params_target, ant_viewing_angles
        )

        # same shape as amplitudes, {GEO, CE} x ANT x FREQ x SLICES
        target_amplitudes = (
            self.amplitudes
            * norm_factors_target[:, None, None, :]
            / ant_distances_target[None, :, None, :]
        )

        # take into account corresction factors
        target_amplitudes /= jnp.where(
            correction_factors_target == 0, 1.0, correction_factors_target
        )

        synthesised_traces = jnp.fft.irfft(
            target_amplitudes * jnp.exp(1j * self.phases),
            norm="ortho",
            axis=2,  # inverse FFT on frequency axis
        )  # shape same as sliced traces, {GEO, CE} x ANT x SAMPLES x SLICES

        return synthesised_traces

    def map_template(self: Self, target: BaseShower) -> jax.Array:
        """
        Map the template to a target profile, represented in a target BaseShower.

        Calculates the trace for every antenna present in the template.

        Parameters
        ----------
        target :  smiet.jax.io.base_shower.BaseShower
            The target BaseShower object, containing the longitudinal profile,
            zenith, azimuth, geomagnetic angle, xmax and nmax

        Returns
        -------
        total_synth : jax.Array
            The synthesised geomagnetic & charge-excess trace for all antennas.
            Shape is {GEO, CE} x ANT x SAMPLES x 1
        """
        synthesised_traces = self.map_template_to_slices(target)

        # total synthesised trace is the sum over all slices
        total_synth = jnp.sum(
            synthesised_traces,
            axis=-1
        )

        return total_synth

    def _calculate_distance_viewing_angles(
        self: Self,
        zenith: float,
        azimuth: float,
        core: np.ndarray,
    ) -> tuple[np.ndarray, np.ndarray]:
        """
        Calculate the viewing angles of the saved antenna's with respect to the slice, as well as their distance to this slice.

        Parameters
        ----------
        zenith : float
            the zenith angle in radians
        azimuth : float
            the azumithal angle in radians.

        Return
        ------
        vangles, distances : np.ndarray
            tuple of viewing angles (units in Cherenkov angle) and distances (units of m)
            for all antennas
        """
        # set the observation level based on the last entry ofg the core
        self.atm.obs_lvl = (
            core[2] / units.cm
        )  # (need to convert from internal units to m)
        # geometric distance from each slice, shape of ANT x SLICES x 3
        dis_from_slices = jnp.expand_dims(
            self.atm.get_geometric_distance_grammage(self.grammages, zenith / units.rad)
            * units.m,
            axis=(0, 2),
        )

        # shower axis as a unit vector, also with shape of ANT x SLICES x 3
        unit_shower_axis = jnp.expand_dims(
            hp_spherical_to_cartesian(zenith / units.rad, azimuth / units.rad), axis=(0, 1)
        )

        # slices as vectors, shape with ANT x SLICES x 3
        slice_vectors = unit_shower_axis * dis_from_slices

        # shape of ANT x SLICES x 3, need to expand antenna positions since
        # only defined as ANT x 3
        slice_to_ant_vectors = self.antenna_information["position"][:, None, :] - (
            slice_vectors + jnp.expand_dims(core, axis=(0, 1))
        )

        # viewing angles and distances returns ANTs x SLICES
        vangles = (
            angle_between(slice_to_ant_vectors, -1 * slice_vectors)
            / self.atm.get_cherenkov_angle(self.grammages, zenith / units.rad)[None, :]
            * units.rad
        )
        distances = jnp.linalg.norm(slice_to_ant_vectors, axis=-1)

        return vangles, distances

    def _calculate_normalisation_factor(self: Self, shower: BaseShower) -> np.ndarray:
        """
        Calculate the normalisation factor based on the atmospheric properties.

        Parameters
        ----------
        shower :  smiet.jax.io.base_shower.BaseShower
            the shower object containing the geometry (zenith, azimuth), geomagnetic angle
            and the longitudinal profile (nparts_per_slice)
        """
        # normalisation factors
        # first entry for normlisation of geomagnetic emission using geomagnetic angle / density
        # second entry same for C-E emission, which is just sin(cherenkov angle)
        # also expand dimensions of nparts per slice to get shape of {GEO,CE} x SLICES
        norm_factors = jnp.array(
            [
                jnp.sin(shower.geomagnetic_angle)
                / self.atm.get_density(self.grammages, shower.zenith / units.rad),
                jnp.sin(
                    self.atm.get_cherenkov_angle(
                        self.grammages, shower.zenith / units.rad
                    )
                    * units.rad
                ),
            ]
        ) * jnp.expand_dims(shower.long_profile, axis=0)

        self.logger.debug("Computed normalisation factors")

        return norm_factors

    def save_template(
        self: Self,
        save_dir: str = os.path.join(
            os.path.dirname(__file__), "..", "templates", "jax"
        ),
        template_file: Union[None, str] = None,
    ) -> None:
        """
        Save the internal state of the synthesis class to disk.

        Parameters
        ----------
        template_file : str, default='default_template.h5'
            the file to save the template into
        save_dir : str, default='smiet/templates'
            the directory to save the template into
        """
        if template_file is None:
            template_file = f"template_{self.template_information['name']}.hdf5"
        with h5py.File(os.path.join(save_dir, template_file), "w") as f:
            prop_grp = f.create_group("shower")
            prop_grp.create_dataset("name", data=self.template_information["name"])
            prop_grp.create_dataset("zenith", data=self.template_information["zenith"])
            prop_grp.create_dataset(
                "azimuth", data=self.template_information["azimuth"]
            )
            prop_grp.create_dataset("magnet", data=self.template_information["magnet"])
            prop_grp.create_dataset("xmax", data=self.template_information["xmax"])
            prop_grp.create_dataset("nmax", data=self.template_information["nmax"])
            prop_grp.create_dataset("core", data=self.template_information["core"])
            prop_grp.create_dataset(
                "long_profile", data=self.template_information["long_profile"], compression="gzip"
            )

            ant_grp = f.create_group("antennas")
            ant_grp.create_dataset("name", data=self.antenna_information["name"])
            ant_grp.create_dataset(
                "position", data=self.antenna_information["position"], compression="gzip"
            )
            ant_grp.create_dataset(
                "position_showerplane",
                data=self.antenna_information["position_showerplane"], compression="gzip"
            )
            ant_grp.create_dataset(
                "time_axis", data=self.antenna_information["time_axis"], compression="gzip"
            )
            ant_grp.create_dataset("n_antennas", data=self.n_antennas)

            freq_grp = f.create_group("frequencies")
            freq_grp.create_dataset("frequency_range", data=self.frequency_range)
            freq_grp.create_dataset("frequencies", data=self.frequencies)
            freq_grp.create_dataset("frequency_mask", data=self.freq_range_mask)
            freq_grp.create_dataset("trunc_frequency", data=self.truncated_frequencies)
            freq_grp.create_dataset("delta_t", data=self.delta_t)
            freq_grp.create_dataset("n_samples", data=self.n_samples)

            spect_grp = f.create_group("spect_params")
            spect_grp.create_dataset("spectral_coeffs", data=self.spectral_coeffs, compression="gzip")
            spect_grp.create_dataset("viewing_angles", data=self.viewing_angles, compression="gzip")

            templ_grp = f.create_group("template")
            templ_grp.create_dataset("amplitudes", data=self.amplitudes, compression="gzip")
            templ_grp.create_dataset("phases", data=self.phases, compression="gzip")
            templ_grp.create_dataset("grammages", data=self.grammages, compression="gzip")
            templ_grp.create_dataset("nr_slices", data=self.nr_slices)

            atm_grp = f.create_group("atmosphere")
            atm_grp.create_dataset("model", data=self.atm.model)
            atm_grp.create_dataset("n0", data=self.atm.n0)

    def load_template(
        self: Self,
        file_path,
        gdas_file : Union[str, None] = None,
    ) -> None:
        """
        Load the template from a saved state, as done by save_template().

        Parameters
        ----------
        file_path : str, default=None
            the file path to load the template from.
        gdas_file : str, default=None
            the GDAS file to use for the atmosphere model.
            If None, then the model number from the loaded template will be used as the 
            atmosphere object. Otherwise the provided GDAS file will be used.
            THe model number must be the same as that from CORSIKA.
        """
        # verify that the file exists
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"Template file {file_path} does not exist.")

        with h5py.File(file_path, "r") as f:
            self.antenna_information["name"] = f["antennas/name"][()]
            self.antenna_information["position"] = f["antennas/position"][()]
            self.antenna_information["position_showerplane"] = f[
                "antennas/position_showerplane"
            ][()]
            self.antenna_information["time_axis"] = f["antennas/time_axis"][()]
            self.n_antennas = len(self.antenna_information["name"])

            self.frequency_range = f["frequencies/frequency_range"][()]
            self.frequencies = f["frequencies/frequencies"][()]
            self.freq_range_mask = f["frequencies/frequency_mask"][()]
            self.truncated_frequencies = f["frequencies/trunc_frequency"][()]
            self.n_samples = f["frequencies/n_samples"][()]
            self.delta_t = f["frequencies/delta_t"][()]

            self.spectral_coeffs = f["spect_params/spectral_coeffs"][()]
            self.viewing_angles = f["spect_params/viewing_angles"][()]

            self.amplitudes = f["template/amplitudes"][()]
            self.phases = f["template/phases"][()]
            self.grammages = f["template/grammages"][()]
            self.nr_slices = f["template/nr_slices"][()]

            atm_model = int(f["atmosphere/model"][()])
            atm_n0 = f["atmosphere/n0"][()]

            self.template_information = {
                "name": f["shower/name"][()],
                "zenith": f["shower/zenith"][()],
                "azimuth": f["shower/azimuth"][()],
                "magnet": f["shower/magnet"][()],
                "xmax": f["shower/xmax"][()],
                "nmax": f["shower/nmax"][()],
                "core": f["shower/core"][()],
                "long_profile": f["shower/long_profile"][()],
            }

        # make the atmosphere object separately
        self.atm = Atmosphere(model=atm_model, n0=atm_n0, observation_level=self.template_information["core"][2], gdas_file=gdas_file)

    def get_time_axis(self: Self) -> jax.Array:
        """
        Get the time axis for all antennas.

        Returns
        -------
        time_axis : np.ndarray
            The time axis for each antenna, shaped as (# antennas, # time samples)

        """
        return self.antenna_information["time_axis"]

    def get_antenna_names(self: Self) -> list[str]:
        """
        Get the names of all internal antennas.

        Returns
        -------
        ant_names : list of str

        """
        return self.antenna_information["name"]

    def truncate_atmosphere(self: Self, starting_grammage: float = 200) -> None:
        """
        Truncate the starting point of the atmosphere grid, and subsequent arrays used for
        the synthesis process.

        Parameters
        ----------
        starting_grammage : float, default=200
            the grammage in which we want to start the atmospheric grid,
            i.e. where we want to truncate from.

        Returns
        -------
        None, but all objects related to the slices will be truncated.
        """
        truncated_gram_idces = jnp.argwhere(self.grammages > starting_grammage)
        self.amplitudes = jnp.squeeze(
            self.amplitudes[..., truncated_gram_idces], axis=-1
        )
        self.phases = jnp.squeeze(self.phases[..., truncated_gram_idces], axis=-1)
        self.grammages = jnp.squeeze(self.grammages[truncated_gram_idces], axis=-1)
        self.nr_slices = len(self.grammages)
