import expansion
import functions as fn
import numpy as np
import parameters
import units_and_constants as uc
import matplotlib.pyplot as plt


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) -> 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
    """
    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)

    pot = (1 / (params.kappa * params.epsilon * uc.CONSTANTS.epsilon0)
            * np.real(np.sum(l_factors[:, None] * ex.coefs[..., None]
                             * fn.sph_harm(l_array[:, None], m_array[:, None], theta[None, :], phi[None, :]), axis=-2)))

    return pot.reshape(ex.shape + angles_shape)


if __name__ == '__main__':

    params = ModelParams(R=150, kappaR=3)
    ex = expansion.MappedExpansionQuad(np.array([0.44, 0.5]), params.kappaR, 0.001, l_max=10)

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

    potential = charged_shell_potential(theta, phi, dist, ex, params)
    print(potential.shape)
    # print(potential)

    plt.plot(theta, potential.T)
    plt.show()