from charged_shells import expansion, parameters
from charged_shells import functions as fn
from charged_shells import units_and_constants as uc
import numpy as np
from scipy.special import eval_legendre, kv


Array = np.ndarray
ModelParams = parameters.ModelParams
Expansion = expansion.Expansion


def charged_shell_potential(theta: Array | float,
                            phi: Array | float,
                            dist: float,
                            ex: Expansion,
                            params: ModelParams,
                            print_idx: int = None) -> Array:
    """
    Electrostatic potential around a charged shell with patches given by expansion over spherical harmonics.

    :param theta: array of azimuthal angles
    :param phi: array of polar angles
    :param dist: distance between the particles in units of radius R
    :param ex: Expansion object detailing patch distribution
    :param params: ModelParams object specifying parameter values for the model
    :param print_idx: if not None, print a single term for debugging purposes
    """
    theta, phi = np.broadcast_arrays(theta, phi)
    angles_shape = theta.shape
    theta = theta.reshape(-1)  # ensures that arrays are 1D
    phi = phi.reshape(-1)

    if not theta.shape == phi.shape:
        raise ValueError('theta and phi arrays should have the same shape.')
    l_array, m_array = ex.lm_arrays

    dist = dist * params.R
    l_factors = (fn.coefficient_Cpm(ex.l_array, params.kappaR) * fn.sph_bessel_k(ex.l_array, params.kappa * dist)
                 / fn.sph_bessel_k(ex.l_array, params.kappaR))
    l_factors = ex.repeat_over_m(l_factors)

    factors = l_factors[:, None] * ex.coefs[..., None] * fn.sph_harm(l_array[:, None], m_array[:, None], theta[None, :], phi[None, :])
    if print_idx is not None:
        print(l_array[print_idx], m_array[print_idx], np.real(factors[print_idx]))

    pot = (1 / (params.kappa * params.epsilon * uc.CONSTANTS.epsilon0) * np.real(np.sum(factors, axis=-2)))

    return pot.reshape(ex.shape + angles_shape)


def inverse_patchy_particle_potential(theta: Array | float,
                                      dist: float,
                                      abar: float,
                                      Zc: float,
                                      Zp: tuple[float, float],
                                      params: ModelParams,
                                      lmax: int = 20):
    dist = dist * params.R

    out0 = ((Zc + Zp[0] + Zp[1]) * (np.exp(params.kappaR) / (1. + params.kappaR)) *
            (np.exp(-params.kappa * dist) / dist) / (params.epsilon * uc.CONSTANTS.epsilon0)) * params.R ** 2

    for l in range(2, lmax, 2):
        out1 = (Zp[0] * np.power(abar, l) * eval_legendre(l, np.cos(theta))
                + Zp[1] * np.power(abar, l) * eval_legendre(l, np.cos(np.pi - theta)))
        out0 += (((2 * l + 1.) * (kv(l + 0.5, params.kappa * dist) / kv(l + 1.5, params.kappaR))) * out1
                 / (params.kappa * params.epsilon * uc.CONSTANTS.epsilon0))

    return out0