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


if __name__ == '__main__':

    params = ModelParams(R=150, kappaR=3)
    ex1 = expansion.MappedExpansionQuad(0.44, params.kappaR, 0.001, l_max=30)
    ex2 = expansion.SphericalCap(np.array([[0, 0], [np.pi, 0]]), 0.5, 0.003, l_max=70)

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

    potential1 = charged_shell_potential(theta, phi, dist, ex1, params)
    potential2 = charged_shell_potential(theta, phi, dist, ex2, params)
    # print(potential.shape)
    # print(potential)

    plt.plot(theta, potential1.T)
    plt.plot(theta, potential2.T)
    plt.show()