import logging
from typing_extensions import Self, Tuple, Union

import numpy as np
import h5py
import jax

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

import jax.numpy as jnp
from scipy.constants import c as c_vacuum

from jax_radio_tools import units, trace_utils, shower_utils
from jax_radio_tools.atmosphere import Atmosphere
from jax_radio_tools.cstrafo_jax import (
    e_to_geo_ce,
    conversion_fieldstrength_cgs_to_SI,
    hp_get_normalized_angle,
)
from .base_shower import BaseShower

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


class CoreasShower(BaseShower):
    """
    Class to read in showers from CoREAS simulations stored in HDF5 files.

    Parameters
    ----------
    file_path : str
        Path to the HDF5 file containing the CoREAS simulation data.
    bandpass : list, default = [30 * units.MHz, 500 * units.MHz]
        The frequency bandpass in MHz to apply to the traces. Defaults to 30-500 MHz,
        which is the bandwidth in which SMIET is applicable in.
    gdas_file : str, optional
        If provided, the atmosphere will be created using this GDAS file. This will take
        precedence over the atmosphere model specified in the simulation settings.
    """

    def __init__(
        self: Self,
        file_path: str,
        bandpass: list = [30 * units.MHz, 500 * units.MHz],
        gdas_file: str = "",
    ) -> None:
        super().__init__()
        self.__file = file_path  # input file path

        self.__trace_length = None
        self.__delta_t = None
        self.__frequency_range = bandpass
        self.name = file_path.split("/")[-1].split(".")[0]

        self.antenna_array = {"name": None, "position": None, "dis_to_core" : None}

        self.traces = None
        self.trace_times = None

        self.__parse_hdf5(gdas_file=gdas_file)

        # filter traces
        self.traces = trace_utils.filter_trace(
            self.traces,
            f_min=self.__frequency_range[0],
            f_max=self.__frequency_range[1],
            trace_sampling=self.delta_t,
            sample_axis=2,
        )

    def __parse_hdf5(
        self: Self,
        gdas_file: str = "",
    ) -> None:
        file = h5py.File(self.__file)

        # set up longitudinal profile parameters
        self._setup_long_parameters(file)

        # now also setup the corsika parameters
        self._setup_corsika_parameters(file)

        # after this, setup the atmosphere
        # imperative that this is done after setting up the core & zenith angle
        self._setup_atmosphere(file, gdas_file=gdas_file)

        # load the antenna and trace information here
        # this should be called after setting the grammage and long profile
        self._load_antenna_and_traces(file)

        # sortthe antenna array to increasing distance from core & azimuth
        self._sort_antenna_array()

        file.close()

    def _setup_long_parameters(self: Self, file: h5py.File) -> None:
        """
        Setup parameters from the longitudinal profile.
        """
        # longitudinal profile data
        long_data = file["atmosphere"]["NumberOfParticles"][:]

        self.grammages = long_data[:, 0]
        self.long_profile = np.sum(long_data[:, 2:4], axis=1)

        # Gaisser-Hillas parameters
        # in particular used to set nmax and xmax
        # in order: [N, X0, Xmax, p0, p1, p2] (p0,p1,p2 are the polynomial coefficients in denominator for lambda)
        GH_parameters = np.array(
            file["atmosphere"].attrs["Gaisser-Hillas-Fit"]
        )  # GH parameters from fit
        self.nmax = GH_parameters[0]
        self.xmax = GH_parameters[2]

        # convert GH parameter to LR
        lmbda = np.mean(
            GH_parameters[3]
            + GH_parameters[4] * self.grammages
            + GH_parameters[5] * self.grammages**2
        )
        L, R = shower_utils.convert_GH_to_LR(self.xmax, GH_parameters[1], lmbda)

        self.GH_parameters = {
            "x0": GH_parameters[1],
            "lambda": lmbda,
            "L": L,
            "R": R,
            "xmax": self.xmax,
            "nmax": self.nmax,
        }

    def _setup_corsika_parameters(self: Self, file: h5py.File) -> None:
        """
        Setup parameters specific to CORSIKA simulations.
        """
        # get (EM) energy
        self.energy = file["inputs"].attrs["ERANGE"][0] * units.GeV
        self.primary = file["inputs"].attrs["PRMPAR"]  # integer stored for convenience

        self.zenith = np.deg2rad(file["inputs"].attrs["THETAP"][0]) * units.rad
        # Shift the azimuthal angle by 3 pi/2 (180 flip from West -> East plus N-> E due to convention of second coordinate)
        self.azimuth = (
            hp_get_normalized_angle(
                3 * np.pi / 2 + np.deg2rad(file["inputs"].attrs["PHIP"][0])
            )
            * units.rad
        )

        # # 3. define magnetic field vector
        self.magnet = np.array([
            0, file["inputs"].attrs["MAGNET"][0], -1 * file["inputs"].attrs["MAGNET"][1]
        ]) * units.microtesla
        # 4. get simulation core
        self.core = (
            np.array(
                [
                    -1 * file["CoREAS"].attrs["CoreCoordinateWest"],
                    file["CoREAS"].attrs["CoreCoordinateNorth"],
                    file["CoREAS"].attrs["CoreCoordinateVertical"],
                ]
            )
            * units.cm
        )
        self.core[2] = file["inputs"].attrs["OBSLEV"] * units.m

        self.delta_t = float(file["CoREAS"].attrs["TimeResolution"]) * units.s

    def _setup_atmosphere(self: Self, file: h5py.File, gdas_file: str = "") -> None:
        """Setup the atmosphere object, and truncate grammages and longitudinal profile if needed."""
        # if the ATMOD value is -1, then use gdas file as atmosphere
        if file["inputs"].attrs["ATMOD"] == -1:
            self.atm = Atmosphere(
                observation_level=self.core[2] / units.m, gdas_file=gdas_file
            )
        else:
            self.atm = Atmosphere(
                model=file["inputs"].attrs["ATMOD"],
                observation_level=self.core[2] / units.m,
            )

        # filter out grammages that will reach beyond earth
        grammage_lim_mask = self.grammages <= self.atm.get_xmax_from_distance(
            0.0,
            self.zenith / units.rad,
        )

        self.long_profile = self.long_profile[grammage_lim_mask]
        self.grammages = self.grammages[grammage_lim_mask]

    def _load_antenna_and_traces(self: Self, file: h5py.File) -> None:
        """Read in the antenna & trace information from the HDF5 file."""
        # antenna positions at the ground
        self.antenna_array["name"] = list(file["CoREAS"]["observers"].keys())
        self.antenna_array["position"] = np.array(
            [
                file["CoREAS"]["observers"][f"{ant_name}"].attrs["position"] * units.cm
                for ant_name in self.antenna_array["name"]
            ]
        )  # antennas x 3
        self.antenna_array["position"] = self.antenna_array["position"][:, [1, 0, 2]]
        self.antenna_array["position"][:,0] *= -1

        # distance from the shower core
        self.antenna_array["dis_to_core"] = np.linalg.norm(
            (self.get_antenna_position_showerplane()[:, :2]), axis=-1
        )

        # the observer data containing both timing and trace information
        observers = np.array(
            [
                file["CoREAS"]["observers"][f"{ant_name}"][:]
                for ant_name in self.antenna_array["name"]
            ]
        )  # antennas x samples x 4

        # map it to the correct shape of 4 x ANT x SAMPLES
        observers = np.moveaxis(observers, (0, 1, 2), (1, 2, 0))

        # divide into trace times & actual E-field slices
        self.trace_times = observers[0, ...] * units.s  # (Nant, Nsamples)
        self.traces = observers[1:, ...]  # (Npol, Nant, Nsamples)
        self.trace_length = self.traces.shape[2]

        self.traces = self.traces[[1, 0, 2], ...] * conversion_fieldstrength_cgs_to_SI  # convert traces from cgs to SI manually (due to electric charge shenanigans)
        self.traces[0,...] *= -1     

    def _sort_antenna_array(self: Self) -> None:
        """Sort the trace and antenna positions based on increasing distance from the core and azimuth."""
        # 7. finally sort based on increasing distance to core & increasing azimuth.
        # first sort based on increasing distance
        dis_core_mask = np.argsort(self.antenna_array["dis_to_core"])
        self.antenna_array["position"] = self.antenna_array["position"][
            dis_core_mask, :
        ]
        self.antenna_array["name"] = [
            self.antenna_array["name"][i] for i in dis_core_mask
        ]
        self.traces = self.traces[:, dis_core_mask, ...]
        self.trace_times = self.trace_times[dis_core_mask, :]
        self.antenna_array["dis_to_core"] = np.sort(self.antenna_array["dis_to_core"])
        # now sort it based on azimuth for each distance
        # NOTE: key assumption applied where it is sorted based on the fact that there are 8 arms
        # when we introduce interpolator we would not need this anymore
        n_ant_per_arm = len(self.antenna_array["dis_to_core"]) // 8

        # calculate the angles
        angles = np.arctan2(
            self.get_antenna_position_showerplane()[:, 1],
            self.get_antenna_position_showerplane()[:, 0],
        )
        angles = np.around(angles, 15)
        angles[angles < 0] += 2 * np.pi
        # get the indices in increasing order and shift by n_arms
        # then concatenate to get the indices
        ordered_idces = [
            np.argsort(angles[int(8 * iant) : int(8 * (iant + 1))]) + 8 * iant
            for iant in range(n_ant_per_arm)
        ]
        ordered_idces = np.array(np.concatenate(ordered_idces), dtype=int)

        # finally apply to the arrays
        self.antenna_array["position"] = self.antenna_array["position"][
            ordered_idces, :
        ]
        self.antenna_array["name"] = [
            self.antenna_array["name"][i] for i in ordered_idces
        ]
        self.traces = self.traces[:, ordered_idces, ...]
        self.trace_times = self.trace_times[ordered_idces, :]
        self.antenna_array["dis_to_core"] = self.antenna_array["dis_to_core"][ordered_idces]

    @property
    def trace_length(self: Self) -> int:
        """
        Length of the trace.

        Returns
        -------
        trace_length : int
            The number of samples in the trace.
        """
        return self.__trace_length

    @property
    def delta_t(self: Self) -> float:
        """
        Time resolution of the trace.

        Returns
        -------
        delta_t : float
            The time resolution in ns.
        """
        return self.__delta_t

    @trace_length.setter
    def trace_length(self: Self, trace_length: int) -> None:
        self.__trace_length = trace_length

    @delta_t.setter
    def delta_t(self: Self, delta_t: float) -> None:
        self.__delta_t = delta_t

    def get_antenna_position(self: Self) -> np.ndarray:
        """
        Get the antenna positions on the ground.

        Returns
        -------
        ant_position_ground : np.ndarray
            The antenna positions on the ground, transformed to the NRR system.
        """
        return self.antenna_array["position"]

    def get_antenna_position_showerplane(self: Self) -> np.ndarray:
        """
        Get the antenna positions in shower plane.

        Returns
        -------
        ant_position_vB_vvB : jax.typing.ArrayLike
            The antenna positions in the shower plane, transformed to the NRR system.
        """
        return self.get_transformer().transform_to_vxB_vxvxB(
            self.antenna_array["position"], core=self.core
        )

    def get_traces_raw(self: Self) -> np.ndarray:
        """
        Get all traces for all antennas on the ground plane.

        Returns
        -------
        traces_ground : jax.typing.ArrayLike
            the traces at ground in xyz coordinates.
            Shape is (3, Nant, Nsamples), where Nant is the number of antennas and Nsamples is the number of samples.

            When called with SlicedShower, the shape is (3, Nant, Nsamples, Nslices).
        """
        return self.traces

    def get_traces_onsky(self: Self) -> np.ndarray:
        """
        Get the traces from traces at the ground to on-sky components.

        NOTE: the radial component (er) is set to zero manually here.

        Returns
        -------
        traces_onsky : np.ndarray
            The traces in the shower plane, transformed to on-sky components (er, etheta, ephi).
            Shape is (3, Nant, Nsamples), where Nant is the number of antennas and Nsamples is the number of samples.

            When called with SlicedShower, the shape is (3, Nant, Nsamples, Nslices).
        """
        trace_onsky = (
            self.get_transformer()
            .transform_from_ground_to_onsky(
                self.traces.T,
            )
            .T
        )
        trace_onsky[0, ...] = 0.0  # set er component to zero
        return trace_onsky

    def get_traces_vB_vvB(self: Self) -> np.ndarray:
        """
        Get the traces from traces at the ground to vB/vvB components.

        Returns
        -------
        traces_vB_vvB : np.ndarray
            The traces in the shower plane, transformed to vxB and vxvxB components.
            Shape is (3, Nant, Nsamples), where Nant is the number of antennas and Nsamples is the number of samples.

            When called with SlicedShower, the shape is (3, Nant, Nsamples, Nslices).
        """
        trace_vvB = (
            self.get_transformer()
            .transform_to_vxB_vxvxB(self.traces.T, core=self.core)
            .T
        )
        return trace_vvB

    def get_traces_geoce(self: Self) -> np.ndarray:
        """
        Get the traces from traces at the ground to GEO/CE components.

        Returns
        -------
        traces_geo_ce : np.ndarray
            The traces in the shower plane, transformed to GEO and CE components.
            Shape is (2, Nant, Nsamples), where Nant is the number of antennas and Nsamples is the number of samples.

            When called with SlicedShower, the shape is (2, Nant, Nsamples, Nslices).
        """
        traces_vvB = self.get_traces_vB_vvB()
        ant_position_vvB = self.get_antenna_position_showerplane()
        # now convert trace to respective components
        traces_geo, traces_ce = e_to_geo_ce(
            traces_vvB.T, ant_position_vvB[None, :, 0], ant_position_vvB[None, :, 1]
        )
        return np.array([traces_geo.T, traces_ce.T])

    def remove_antennas(self, ant_idx_min: int = 0, ant_idx_max: int = 90) -> None:
        """
        Remove antennas from the data traces.

        Parameter:
        ----------
        ant_idx_min : float
            remove all antennas below this antenna index
        ant_idx_max : float
            remove all antennas above this antenna index
        """
        # d_core_mask = (self.dis_to_core > d_core_min) & (self.dis_to_core < d_core_max)
        d_core_mask = (np.indices(self.dis_to_core.shape)[0] >= ant_idx_min) & (
            np.indices(self.dis_to_core.shape)[0] < ant_idx_max
        )
        self.traces = self.traces[:, d_core_mask, ...]
        self.trace_times = self.trace_times[d_core_mask, :]
        self.antenna_array["position"] = self.antenna_array["position"][d_core_mask, :]
        self.antenna_array["name"] = [
            self.antenna_array["name"][i]
            for i in range(len(self.antenna_array["name"]))
            if d_core_mask[i]
        ]
        self.dis_to_core = self.dis_to_core[d_core_mask]
