from charged_shells import patch_size, parameters, charge_distributions, expansion
from charged_shells import functions as fn
from scipy.special import eval_legendre
import numpy as np


ModelParams = parameters.ModelParams
Array = np.ndarray


def point_to_gauss_magnitude(sigma_tilde: Array, a_bar: Array, lbd: Array, kappaR: Array):
    return (sigma_tilde * fn.coefficient_Cim(2, kappaR) / fn.coefficient_Cpm(2, kappaR)
            * np.sinh(lbd) / (lbd * fn.sph_bessel_i(2, lbd)) * a_bar ** 2)


def point_to_cap_magnitude(sigma_tilde: Array, a_bar: Array, theta0: Array, kappaR: Array):
    return (sigma_tilde * 10 * fn.coefficient_Cim(2, kappaR) / fn.coefficient_Cpm(2, kappaR) *
            a_bar ** 2 / (eval_legendre(1, np.cos(theta0)) - eval_legendre(3, np.cos(theta0))))


def ic_to_gauss(sigma_tilde, a_bar, params: ModelParams, l_max: int = 30,
                sigma0: float = 0) -> expansion.Expansion:

    ex_mapped = charge_distributions.create_mapped_quad_expansion(a_bar=a_bar, kappaR=params.kappaR,
                                                                        sigma_tilde=sigma_tilde, l_max=30, sigma0=sigma0)
    target_patch_size = patch_size.potential_patch_size(ex_mapped, params)
    sigma0_mapped = charge_distributions.net_charge_map(sigma0, params.kappaR)

    def fn_gauss(x):
        return charge_distributions.create_gaussian_charge_expansion(lambda_k=x, omega_k=np.array([[0, 0], [np.pi, 0]]),
                                                                   sigma1=sigma_tilde, l_max=l_max, sigma0=sigma0_mapped)

    lbd = patch_size.inverse_potential_patch_size(target_patch_size, fn_gauss, 5, params)
    gauss_sigma = point_to_gauss_magnitude(sigma_tilde, a_bar, lbd, params.kappaR)
    return charge_distributions.create_gaussian_charge_expansion(lambda_k=lbd, omega_k=np.array([[0, 0], [np.pi, 0]]),
                                                               sigma1=gauss_sigma, l_max=l_max, sigma0=sigma0_mapped)


def ic_to_cap(sigma_tilde, a_bar, params: ModelParams, l_max: int = 30, sigma0: float = 0) -> expansion.Expansion:

    ex_mapped = charge_distributions.create_mapped_quad_expansion(a_bar=a_bar, kappaR=params.kappaR,
                                                                        sigma_tilde=sigma_tilde, l_max=30, sigma0=sigma0)
    target_patch_size = patch_size.potential_patch_size(ex_mapped, params)
    sigma0_mapped = charge_distributions.net_charge_map(sigma0, params.kappaR)

    def fn_cap(x):
        return charge_distributions.create_spherical_cap_expansion(theta0_k=x, sigma1=sigma_tilde, l_max=l_max,
                                                                omega_k=np.array([[0, 0], [np.pi, 0]]), sigma0=sigma0_mapped)

    theta0 = patch_size.inverse_potential_patch_size(target_patch_size, fn_cap, 0.5, params)
    cap_sigma = point_to_cap_magnitude(sigma_tilde, a_bar, theta0, params.kappaR)
    return charge_distributions.create_spherical_cap_expansion(theta0_k=theta0, sigma1=cap_sigma,
                                                            omega_k=np.array([[0, 0], [np.pi, 0]]), l_max=l_max, sigma0=0)