"""Utility module to take care of transformations to shower plane etc."""

import jax

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

import jax.numpy as jnp
from jax_radio_tools import geo_ce_to_vB_vvB_v


def transform_traces_on_vxB(
    traces_geoce: jax.typing.ArrayLike,
    x: jax.typing.ArrayLike,
    y: jax.typing.ArrayLike,
    vxB_axis_mode: str = "replace",
) -> jax.Array:
    """
    Adjust the (synthesised) traces on the vxB axis.

    This function is meant to alter the traces on the vxB axis, as in the current
    version of the SMIET software the vxB axis cannot be synthesised. There are several
    strategies to handle this, and applying them should be facilitated by this function.

    Here it is assumed that the antennas are ordered with increasing phi, where phi=0 is oriented West.

    Parameters
    ----------
    traces_geoce : jax.typing.ArrayLike
        the synthesised traces in geomagnetic and charge-excess polarisation.
        Shape must be of NPOL x NANTS x NSAMPLES
    x : jax.typing.ArrayLike
        the antenna positions along the vxB axis, shaped as (NANTS,)
    y : jax.typing.ArrayLike
        the antenna positions along the vxvxB axis, shaped as (NANTS,)
    vxB_axis_mode : {"replace", "zero", "average"}
        Operation to perform on the vxB axis. If "replace" (the default), the traces from the
        vxvxB component will be copied over to the vxB component. If "zero", the vxB component
        will be set to zero.

    Returns
    -------
    traces_vB_vvB : jax.Array
        the synthesised traces in vB x vvB axis (shower plane)
    """
    # here the threshold cut is rather high, I am not sure why such a high cut is necessary
    # but it is somehow necessary for the antenna set I am using
    ant_at_vvB_axis = jnp.abs(x) < 5
    ant_at_vB_axis = jnp.abs(y) < 5

    if vxB_axis_mode == "replace":
        traces_geoce = traces_geoce.at[:, ant_at_vB_axis, :].set(
            traces_geoce[:, ant_at_vvB_axis, :]
        )
    elif vxB_axis_mode == "zero":
        # for the vxB component, set to zero
        traces_geoce = traces_geoce.at[:, ant_at_vB_axis, :].set(0.0)
    elif vxB_axis_mode == "average":
        # for the vxB component, average the vxvxB and vxB components
        arm_idces = jnp.arange(0, len(x), 8, dtype=int)
        traces_geoce = traces_geoce.at[:, arm_idces, :].set(
            0.5
            * (traces_geoce[:, arm_idces + 1, :] + traces_geoce[:, arm_idces + 7, :])
        )
        traces_geoce = traces_geoce.at[:, arm_idces + 4, :].set(
            0.5
            * (traces_geoce[:, arm_idces + 3, :] + traces_geoce[:, arm_idces + 5, :])
        )

    else:
        raise ValueError(
            f"Unknown vxB_axis_mode: {vxB_axis_mode}. Must be one of 'replace' or 'zero'."
        )

    return traces_geoce


def reconstruct_starshape_from_vxvxB(
    geo_traces, ce_traces, antenna_positions_vvB, number_of_arms=8
):
    """
    Reconstruct traces on all arms from vxvxB-arm simulation (where geo and ce are orthogonal).

    This is used for generating star-shapes from simulations that only have antennas on the
    vxvxB arm. The traces on the other arms are simply copied from the vxvxB arm, and the
    antenna positions are rotated accordingly.

    Traces synthesised for all slices can also be passed. In that case, the output traces will
    have the same number of slices. If only at the ground, the n_slices = 1.

    Parameters
    ----------
    geo_traces : np.ndarray
        Geomagnetic traces on the vxvxB arm, shape (n_antennas, n_samples, n_slices)
    ce_traces : np.ndarray
        Charge-excess traces on the vxvxB arm, shape (n_antennas, n_samples, n_slices)
    antenna_positions_vvB : np.ndarray
        Positions of antennas on vxvxB arm (distance in shower plane), shape (n_antennas,)
    number_of_arms : int, default=8
        Total number of star-shape arms to reconstruct

    Returns
    -------
    geo_traces_starshape : np.ndarray
        Geomagnetic traces on all arms, shape (n_antennas*number_of_arms, n_samples, n_slices)
    ce_traces_starshape : np.ndarray
        Charge-excess traces on all arms, shape (n_antennas*number_of_arms, n_samples, n_slices)
    antenna_positions_starshape : np.ndarray
        Antenna positions on all arms, shape (n_antennas*number_of_arms, 3)
    """
    n_ant_per_arm = geo_traces.shape[0]

    # get the total number of antennas
    n_ant = n_ant_per_arm * number_of_arms

    # initialise arrays
    antenna_positions_starshape = jnp.zeros((n_ant, 3))
    geo_traces_starshape = jnp.zeros((n_ant, *geo_traces.shape[1:]))
    ce_traces_starshape = jnp.zeros((n_ant, *geo_traces.shape[1:]))
    antenna_labels = []

    phi_angles = jnp.linspace(0, 360, number_of_arms, endpoint=False)

    for ip, phi in enumerate(phi_angles):
        phi_rad = jnp.deg2rad(phi)
        geo_traces_starshape = geo_traces_starshape.at[
            ip * n_ant_per_arm : (ip + 1) * n_ant_per_arm, ...
        ].set(geo_traces)
        ce_traces_starshape = ce_traces_starshape.at[
            ip * n_ant_per_arm : (ip + 1) * n_ant_per_arm, ...
        ].set(ce_traces)
        antenna_positions_starshape = antenna_positions_starshape.at[
            ip * n_ant_per_arm : (ip + 1) * n_ant_per_arm, 0
        ].set(antenna_positions_vvB * jnp.cos(phi_rad))
        antenna_positions_starshape = antenna_positions_starshape.at[
            ip * n_ant_per_arm : (ip + 1) * n_ant_per_arm, 1
        ].set(antenna_positions_vvB * jnp.sin(phi_rad))

        antenna_labels += [f"ant_r{antenna_positions_vvB[i]:.0f}_phi{phi:.0f}" for i in range(n_ant_per_arm)]

    return geo_traces_starshape, ce_traces_starshape, antenna_positions_starshape, antenna_labels

def geo_ce_to_e(traces_geo: jax.typing.ArrayLike, traces_ce : jax.typing.ArrayLike, x, y) -> jax.Array:
    """
    Wrapper function for jax_radio_tools.trace_utils.geo_ce_to_vxB_vxvxB
    """
    return geo_ce_to_vB_vvB_v(
        jnp.array([traces_geo, traces_ce]),
        x,
        y
    )
