from __future__ import annotations
import copy
import numpy as np
from scipy.special import eval_legendre
from charged_shells import expansion, functions as fn
import quaternionic

Array = np.ndarray
Quaternion = quaternionic.array
Expansion = expansion.Expansion


def create_expansion24(sigma2: float | Array, sigma4: float | Array, sigma0: float | Array = 0.):
    l_array = np.array([0, 2, 4])
    try:
        broadcasted_arrs = np.broadcast_arrays(np.sqrt(4*np.pi) * sigma0, sigma2, sigma4)
    except ValueError:
        raise ValueError("Given sigma0, sigma2 and sigma4 arrays cannot be broadcast to a common shape.")
    coefs = rot_sym_expansion(l_array, np.stack(broadcasted_arrs, axis=-1))
    return Expansion(l_array, coefs)


def create_mapped_rotsym_expansion(a_bar: Array | float,
                                   kappaR: Array | float,
                                   sigma_tilde: float,
                                   l_array: Array,
                                   sigma0: float | Array = 0) -> Expansion:
    """
    Create Expansion that matches the outside potential of an impermeable particle with
    a rotationally symmetric point charge distribution inside (specifically, either quadrupolar or dipolar).
    :param a_bar: distance between the center and off center charges
    :param kappaR: screening parameter
    :param sigma_tilde: magnitude of off-center charges / 4pi R^2
    :param l_max: maximal ell value for the expansion
    :param sigma0: total (mean) charge density
    """
    if not isinstance(sigma0, Array):
        sigma0 = np.array([sigma0])
    if sigma0.ndim > 1:
        raise ValueError(f'Sigma0 parameter cannot be an array of dimensions larger than 1, got dim={sigma0.ndim}')
    a_bar_bc, kappaR_bc = np.broadcast_arrays(a_bar, kappaR)

    a_bar_bc2, kappaR_bc2, l_array_expanded = np.broadcast_arrays(a_bar_bc[..., None],
                                                                  kappaR_bc[..., None],
                                                                  l_array[None, :])
    coefs = (2 * sigma_tilde * fn.coef_C_diff(l_array_expanded, kappaR_bc2)
             * np.sqrt(4 * np.pi * (2 * l_array_expanded + 1)) * np.power(a_bar_bc2, l_array_expanded))
    coefs = np.squeeze(rot_sym_expansion(l_array, coefs))
    kappaR_bc3, sigma0_bc3 = np.broadcast_arrays(kappaR_bc[..., None], sigma0[None, :])
    coefs = expansion_total_charge(coefs, net_charge_map(sigma0_bc3, kappaR_bc3), l_array)
    return Expansion(l_array, np.squeeze(coefs))


def create_mapped_quad_expansion(a_bar: Array | float,
                                 kappaR: Array | float,
                                 sigma_tilde: float,
                                 l_max: int = 20,
                                 sigma0: float | Array = 0) -> Expansion:
    """Create Expansion mapped from a quadrupolar point charge distribution."""
    l_array = np.array([l for l in range(l_max + 1) if l % 2 == 0])
    return create_mapped_rotsym_expansion(a_bar, kappaR, sigma_tilde, l_array, sigma0)


def create_mapped_dipolar_expansion(a_bar: Array | float,
                                    kappaR: Array | float,
                                    sigma_tilde: float,
                                    l_max: int = 20,
                                    sigma0: float | Array = 0) -> Expansion:
    """Create Expansion mapped from a dipolar point charge distribution."""
    l_array = np.array([l for l in range(l_max + 1) if l % 2 == 1 or l == 0])
    return create_mapped_rotsym_expansion(a_bar, kappaR, sigma_tilde, l_array, sigma0)


def create_gaussian_charge_expansion(omega_k: Array,
                                     lambda_k: Array | float,
                                     sigma1: float,
                                     l_max: int,
                                     sigma0: float | Array = 0,
                                     equal_charges: bool = True)  -> Expansion:
    """
    Create Expansion for a collection of smeared charges on the sphere.
    :param omega_k: array of positions (theta, phi) of all charges
    :param lambda_k: smear parameter for each charge or smear for different cases (if equal_charges = True)
    :param sigma1: scaling
    :param l_max: maximal ell value for the expansion
    :param sigma0: total (mean) charge density
    :param equal_charges: if this is False, length of lambda_k should be N. If True, theta0_k array will be treated
        as different expansion cases
    """
    omega_k = omega_k.reshape(-1, 2)
    if not isinstance(lambda_k, Array):
        lambda_k = np.array([lambda_k])
    if equal_charges:
        if lambda_k.ndim > 1:
            raise ValueError(f'If equal_charges=True, lambda_k should be a 1D array, got shape {lambda_k.shape}')
        lambda_k = np.full((omega_k.shape[0], lambda_k.shape[0]), lambda_k).T
    if lambda_k.shape[-1] != omega_k.shape[0]:
        raise ValueError("Number of charges (length of omega_k) should match the last dimension of lambda_k array.")
    lambda_k = lambda_k.reshape(-1, omega_k.shape[0])
    l_array = np.arange(l_max + 1)
    full_l_array, full_m_array = expansion.full_lm_arrays(l_array)
    theta_k = omega_k[:, 0]
    phi_k = omega_k[:, 1]
    summands = (lambda_k[:, None, :] / np.sinh(lambda_k[:, None, :])
                * fn.sph_bessel_i(full_l_array[None, :, None], lambda_k[:, None, :])
                * np.conj(fn.sph_harm(full_l_array[None, :, None], full_m_array[None, :, None],
                                      theta_k[None, None, :], phi_k[None, None, :])))
    coefs = np.squeeze(4 * np.pi * sigma1 * np.sum(summands, axis=-1))
    coefs = expansion_total_charge(coefs, sigma0, l_array)
    l_array, coefs = purge_unneeded_l(l_array, coefs)
    return Expansion(l_array, coefs)


def create_spherical_cap_expansion(omega_k: Array,
                                   theta0_k: Array | float,
                                   sigma1: float,
                                   l_max: int,
                                   sigma0: float | Array = 0,
                                   equal_sizes: bool = True) -> Expansion:
    """
    Create Expansion for a collection of spherical caps.
    :param omega_k: array of positions (theta, phi) of all spherical caps
    :param theta0_k: sizes of each spherical caps or cap sizes for different cases (if equal_sizes = True)
    :param sigma1: charge magnitude for the single cap, currently assumes that this is equal for all caps
    :param l_max: maximal ell value for the expansion
    :param sigma0: total (mean) charge density
    :param equal_sizes: if this is False, length of theta0_k should be N. If True, theta0_k array will be treated as
        different expansion cases
    """
    omega_k = omega_k.reshape(-1, 2)
    if not isinstance(theta0_k, Array):
        theta0_k = np.array(theta0_k)
    if equal_sizes:
        if theta0_k.ndim == 0:
            theta0_k = np.full(omega_k.shape[0], theta0_k)
        elif theta0_k.ndim == 1:
            theta0_k = np.full((omega_k.shape[0], theta0_k.shape[0]), theta0_k)
        else:
            raise ValueError(f'If equal_charges=True, theta0_k should be a 1D array, got shape {theta0_k.shape}')
    if theta0_k.shape[0] != omega_k.shape[0]:
        raise ValueError("Number of charges (length of omega_k) should match the last dimension of theta0_k array.")
    rotations = Quaternion(np.stack((np.cos(omega_k[..., 0] / 2),
                                     np.sin(omega_k[..., 1]) * np.sin(omega_k[..., 0] / 2),
                                     np.cos(omega_k[..., 1]) * np.sin(omega_k[..., 0] / 2),
                                     np.zeros_like(omega_k[..., 0]))).T)
    l_array = np.arange(l_max + 1)
    coefs_l0 = -sigma1 * (np.sqrt(np.pi / (2 * l_array[None, :] + 1)) *
                          (eval_legendre(l_array[None, :] + 1, np.cos(theta0_k[..., None]))
                           - eval_legendre(l_array[None, :] - 1, np.cos(theta0_k[..., None]))))
    coefs = rot_sym_expansion(l_array, coefs_l0)
    coefs_all_single_caps = expansion.expansion_rotation(rotations, coefs, l_array)
    # Rotating is implemented in such a way that it rotates every patch to every position,
    # hence the redundancy of out of diagonal elements.
    coefs_all = np.sum(np.diagonal(coefs_all_single_caps), axis=-1)
    coefs_all = expansion_total_charge(coefs_all, sigma0, l_array)
    return Expansion(l_array, np.squeeze(coefs_all))


def net_charge_map(sigma0: float | Array, kappaR: float | Array):
    return sigma0 * np.exp(kappaR) / np.sinh(kappaR) * kappaR / (1 + kappaR)


def rot_sym_expansion(l_array: Array, coefs: Array) -> Array:
    """Create full expansion array for rotationally symmetric distributions with only m=0 terms different form 0."""
    full_coefs = np.zeros(coefs.shape[:-1] + (np.sum(2 * l_array + 1),))
    full_coefs[..., np.cumsum(2 * l_array + 1) - l_array - 1] = coefs
    return full_coefs


def expansion_total_charge(coefs: Array, sigma0: float | Array | None, l_vals: Array):
    """Adds a new axis to the expansion coefficients that modifies expansion based on given net charge density."""
    if l_vals[0] != 0:
        raise NotImplementedError("To modify total charge, the first expansion term should be l=0. "
                                  "Adding l=0 term not yet implemented (TODO).")
    if sigma0 is None:
        return coefs
    if not isinstance(sigma0, Array):
        x = copy.deepcopy(coefs)
        x[..., 0] = sigma0 * np.sqrt(4 * np.pi)
        return x

    # insert new axis in the 2nd-to-last place and repeat the expansion data over this new axis
    x = np.repeat(np.expand_dims(coefs, -2), sigma0.shape[-1], axis=-2)
    x[..., 0] = sigma0 * np.sqrt(4 * np.pi)
    return x


def purge_unneeded_l(l_array: Array, coefs: Array) -> (Array, Array):
    """Remove ell values from expansion for which all (ell, m) coefficients are zero."""
    def delete_zero_entries(l, l_arr, cfs):
        l_idx = np.where(l_arr == l)[0][0]
        m_indices = expansion.m_indices_at_l(l_arr, l_idx)
        if np.all(cfs[..., m_indices] == 0):
            return np.delete(l_arr, l_idx), np.delete(cfs, m_indices, axis=-1)
        return l_arr, cfs
    for l in l_array:
        l_array, coefs = delete_zero_entries(l, l_array, coefs)
    return l_array, coefs