import logging
from typing_extensions import Self, Tuple, Union

import numpy as np
import h5py
import jax

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

from smiet import units
from jax_radio_tools import (
    e_to_geo_ce,
    conversion_fieldstrength_cgs_to_SI,
)
from .coreas_shower import CoreasShower

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


class SlicedShower(CoreasShower):
    """
    Class to read in showers from each slice from CoREAS simulations, inherited from CoreasShower.

    Parameters
    ----------
    file_path : str
        The filepath to the simulation to read.
    slicing_grammage : int, default=5
        The width between atmospheric slices in g/cm^2
    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,
        slicing_grammage: int = 5,  # not used since its read from the file itself
        bandpass: list = [30 * units.MHz, 500 * units.MHz],
        gdas_file : str = "",
    ) -> None:
        
        self.logger = logging.getLogger("smiet.jax.io.SlicedShower")
        self.__slice_gram = slicing_grammage  # g/cm2, TODO: change how to read this since its fixed from filepath  # Not used
        super().__init__(file_path, bandpass=bandpass, gdas_file=gdas_file)



    def _load_antenna_and_traces(self : Self, file : h5py.File) -> None:
        """Read in the antenna information from the HDF5 file."""
        # observer properties
        self.antenna_array["name"] = list(
            sorted(
                set([key.split("x")[0] for key in file["CoREAS"]["observers"].keys()])
            )
        )

        # antenna positions at the ground
        self.antenna_array["position"] = np.array(
            [
                file["CoREAS"]["observers"][f"{ant_name}x{self.__slice_gram}"].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 for each slice, containing both timing and trace information
        observer_slices = np.array(
            [
                [
                    file["CoREAS"]["observers"][f"{ant_name}x{int(gram):d}"][:]
                    for ant_name in self.antenna_array["name"]
                ]
                for gram in self.grammages
            ]
        )  # slices x antennas x samples x 4

        # map it to the correct shape of
        # 4 x ANT x SAMPLES x SLICE as according to numpy version
        observer_slices = np.moveaxis(observer_slices, (0, 1, 2, 3), (3, 1, 2, 0))

        # divide into trace times & actual E-field slices
        # we also do not care about the time for each slice, since they arrive at the same time
        self.trace_times = (
            observer_slices[0, ..., -1] * units.s
        )  # shape of (Nant, Nsamples), not needed but keep incase
        self.traces = observer_slices[
            1:, ...
        ]  # shape of (Npol, Nant, Nsamples, Nslices)

        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 get_coreas_settings(self: Self) -> dict:
        """Get specific configurations from the CoREAS simulation that is useful for the synthesis.

        Returns
        -------
        coreas_settings : dict
            A dictionary containing important information about the configuration of the CoREAS simulation
        """
        file = h5py.File(self.__file)

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

        return {"time_resolution": time_resolution}
    
    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, None, :, 0], ant_position_vvB[None, None, :, 1]
        )
        return np.array([traces_geo.T, traces_ce.T])
