import logging
from typing_extensions import Self

import jax

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

import jax.numpy as jnp
import numpy as np

from smiet import units
from jax_radio_tools import hp_spherical_to_cartesian, cstrafo
from ..utilities.geometry import angle_between
from ..utilities.jax_utils import ModelMeta

logger = logging.getLogger("smiet.jax.io")


class BaseShower(metaclass=ModelMeta):
    """
    Base class to contain the shower parameters.

    The Shower class is used to hold the geometric information for a shower, like the zenith/azimuth,
    as well as the longitudinal profile. It can be used to specify the target parameters for a shower
    in template synthesis.
    """

    __parameter_names = ("xmax", "nmax", "zenith", "azimuth", "magnetic_field")

    def __init__(self: Self) -> None:
        self.logger = logging.getLogger("smiet.jax.io.Shower")

        # shower parameters
        self.__xmax = None
        self.__nmax = None
        self.__grammages = None
        self.__long_profile = None
        self.__GH_parameters = None
        self.__transformer = None

        # observational parameters
        self.__zenith = None
        self.__azimuth = None
        self.__geometry = None  # not used, there for backwards compatibility
        self.__magnetic_field = None
        self.__simulation_core = jnp.array([0, 0, 0])  # Currently fixed

    @property
    def xmax(self: Self) -> float:
        """The $X_{max}$ of the shower, from fitting a GH to the longitudinal profile."""
        if self.__xmax is not None:
            return self.__xmax
        else:
            self.logger.error("Xmax has not been set")

    @property
    def nmax(self: Self) -> float:
        """The $N_{max}$ of the shower, from fitting a GH to the longitudinal profile."""
        if self.__nmax is not None:
            return self.__nmax
        else:
            self.logger.error("Nmax has not been set")

    @property
    def grammages(self: Self) -> jax.Array:
        """Array of atmospheric slices in g/cm^2."""
        if self.__grammages is not None:
            return self.__grammages
        else:
            self.logger.error("The grammage has not been set")

    @property
    def slice_grammage(self: Self) -> float:
        """The spacing between each grammage in g/cm^2."""
        if self.__grammages is not None:
            return self.__grammages[1] - self.__grammages[0]
        else:
            self.logger.error(
                "Longitudinal profile has not been set yet, cannot calculate slicing thickness"
            )

    @property
    def long_profile(self: Self) -> float:
        """The longitudinal profile of electrons and positrons in each atmospheric slice."""
        if self.__long_profile is not None:
            # ensure that the length of grammages and ___long_profile is the same
            assert len(self.__long_profile) == len(self.__grammages), (
                "Length not same as grammages!"
            )
            return self.__long_profile
        else:
            self.logger.error("Set the longitudinal parameters first.")

    @property
    def GH_parameters(self: Self) -> dict:
        """Get the Gaisser-Hillas parameters as a dictionary."""
        if self.__GH_parameters is not None:
            return self.__GH_parameters
        else:
            self.logger.error("GH parameters have not been set")

    @property
    def nr_slices(self: Self) -> int:
        """The number of slices in the array."""
        if self.__long_profile is not None or self.__grammages is not None:
            # ensure that the length of grammages and __long_profile is the same
            assert len(self.__long_profile) == len(self.__grammages), (
                "Length not same as grammages!"
            )
            return len(self.__long_profile)
        else:
            self.logger.error("Set the longitudinal parameters first.")

    @property
    def geometry(self: Self) -> jax.Array:
        """Store the zenith and azimuth. These must be provided in the internal units system."""
        if self.__geometry is not None:
            return self.__geometry
        else:
            self.logger.error("Geometry has not been set")

    @property
    def zenith(self: Self) -> float:
        """The zenith angle in radians."""
        if self.__zenith is not None:
            return self.__zenith
        else:
            self.logger.error("Geometry has not been set")

    @property
    def azimuth(self: Self) -> float:
        """The azimuthal angle in radians."""
        if self.__azimuth is not None:
            return self.__azimuth
        else:
            self.logger.error("Geometry has not been set")

    @property
    def core(self: Self) -> jax.Array:
        """The core (x, y, z in NRR CS) where the EAS hit in the simulation."""
        if self.__simulation_core is not None:
            return self.__simulation_core
        else:
            self.logger.error("The simulation core is not known")

    @property
    def magnet(self: Self) -> float:
        """Magnetic field vector in the NRR coordinate system."""
        if self.__magnetic_field is not None:
            return self.__magnetic_field
        else:
            self.logger.error("The magnetic field vector has not been set")

    @property
    def geomagnetic_angle(self: Self) -> float:
        """The angle between the magnetic field vector and the shower axis."""
        shower_axis = hp_spherical_to_cartesian(
            self.__zenith / units.rad, self.__azimuth / units.rad
        )

        return angle_between(self.__magnetic_field, shower_axis)

    @xmax.setter
    def xmax(self: Self, xmax: float) -> None:
        self.__xmax = xmax

    @nmax.setter
    def nmax(self: Self, nmax: float) -> None:
        self.__nmax = nmax

    @zenith.setter
    def zenith(self: Self, zenith: float) -> None:
        self.__zenith = zenith

    @azimuth.setter
    def azimuth(self: Self, azimuth: float) -> None:
        self.__azimuth = azimuth

    @geometry.setter
    def geometry(self: Self, geo: jax.typing.ArrayLike) -> None:
        assert len(geo) == 2, (
            "Please provide zenith and azimuth components in internal units"
        )

        self.__geometry = jnp.array(geo)

    @grammages.setter
    def grammages(self: Self, grammages: jax.typing.ArrayLike) -> None:
        self.__grammages = grammages

    @long_profile.setter
    def long_profile(self: Self, long_profile: jax.typing.ArrayLike) -> None:
        self.__long_profile = long_profile

    @GH_parameters.setter
    def GH_parameters(self: Self, params: dict) -> None:
        self.__GH_parameters = params

    @magnet.setter
    def magnet(self: Self, magnet_field_vector: jax.typing.ArrayLike) -> None:
        assert len(magnet_field_vector) == 3, (
            "B-field vector must contain three components"
        )

        self.__magnetic_field = magnet_field_vector

    @core.setter
    def core(self: Self, core: jax.typing.ArrayLike) -> None:
        self.__simulation_core = core

    def get_transformer(self: Self):
        """
        Get the transformer from jax_radio_tools that contains the coordinate transformations from shower to ground frame, as well to on-sky coordinates.

        Returns
        -------
        transformer : cstrafo
            The coordinate transformer from jax_radio_tools.
        """
        if self.__transformer is None:
            self.__transformer = cstrafo(
                zenith=self.zenith / units.rad,
                azimuth=self.azimuth / units.rad,
                magnetic_field_vector=self.magnet,
            )

        return self.__transformer

    def set_parameters(
        self: Self,
        grammages: jax.typing.ArrayLike,
        long_profile: jax.typing.ArrayLike,
        xmax : jax.typing.ArrayLike,
        nmax : jax.typing.ArrayLike,
        zenith : jax.typing.ArrayLike,
        azimuth : jax.typing.ArrayLike,
        magnetic_field : jax.typing.ArrayLike,
        core : jax.typing.ArrayLike,
    ) -> None:
        """
        Set the parameters of the shower model from a dictionary of parameters.

        This is a convenience function to set all parameters at once.

        Parameters
        ----------
        grammages : jax.typing.ArrayLike
            an array of atmospheric depth in g/cm^2
        long_profile : jax.typing.ArrayLike
            an array of the longitudinal profile of the shower
        xmax : jax.typing.ArrayLike
            depth of shower maximum in g/cm^2
        nmax : jax.typing.ArrayLike
            number of particles at shower maximum
        zenith : jax.typing.ArrayLike
            zenith angle in radians
        azimuth : jax.typing.ArrayLike
            azimuthal angle in radians
        magnetic_field : jax.typing.ArrayLike
            magnetic field vector in Gauss. Must be a 3-vector
        core : jax.typing.ArrayLike
            core position in meters. Must be a 3-vector
        """
        self.grammages = grammages
        self.long_profile = long_profile

        self.xmax = xmax
        self.nmax = nmax

        self.zenith = zenith
        self.azimuth = azimuth
        self.magnet = magnetic_field
        self.core = core

    def transform_profile_to_origin(
        self: Self, origin_grammages: jax.typing.ArrayLike
    ) -> None:
        """
        Transform the longitudinal profile to match those of the origin shower.

        This is needed such that the grammages of the origin and target match.

        Parameters
        ----------
        origin_grammages : jax.typing.ArrayLike
            The grammages of the origin shower.
        """
        self.long_profile = jnp.interp(
            x=origin_grammages, xp=self.grammages, fp=self.long_profile
        )
        self.grammages = origin_grammages