import expansion
import functions as fn
import numpy as np
import matplotlib.pyplot as plt


Array = np.ndarray


def charged_shell_potential(theta: Array | float,
                            phi: Array | float,
                            dist: float,
                            ex: expansion.Expansion,
                            kappaR: float,
                            R: float) -> 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 kappaR: screening value
    :param R: particle radius in nm
    :return: potential values in units of 1 / epsilon * epsilon0
    """
    if isinstance(theta, float):
        theta = np.full_like(phi, theta)

    if isinstance(phi, float):
        phi = np.full_like(theta, phi)

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

    l_factors = (fn.coefficient_Cpm(ex.l_array, kappaR) * fn.sph_bessel_k(ex.l_array, kappaR * dist)
                 / fn.sph_bessel_k(ex.l_array, kappaR))

    return np.real(R / kappaR * np.sum(ex.repeat_over_m(l_factors)[None, :] * ex.coeffs
                               * fn.sph_harm(l_array[None, :], m_array[None, :], theta[:, None], phi[:, None]),
                               axis=1))


if __name__ == '__main__':

    kappaR = 3
    R = 150
    ex = expansion.MappedExpansion(1, kappaR, 0.001, max_l=20)

    theta = np.linspace(0, np.pi, 1000)
    phi = 0.
    dist = 1.

    potential = charged_shell_potential(theta, phi, dist, ex, kappaR, R)
    # print(potential)

    plt.plot(potential)
    plt.show()