from charged_shells import expansion, functions as fn, potentials, patch_size
import numpy as np
from charged_shells.parameters import ModelParams
import matplotlib.pyplot as plt
import scipy.special as sps
from pathlib import Path


def point_to_gauss_map(sigma_m, a_bar, lbd, params: ModelParams):
    return (sigma_m * fn.coefficient_Cim(2, params.kappaR) / fn.coefficient_Cpm(2, params.kappaR)
            * np.sinh(lbd) / (lbd * fn.sph_bessel_i(2, lbd)) * a_bar ** 2)


def point_to_cap_map(sigma_m, a_bar, theta0, params: ModelParams):
    return (sigma_m * 10 * fn.coefficient_Cim(2, params.kappaR) / fn.coefficient_Cpm(2, params.kappaR) *
            a_bar ** 2 / (sps.eval_legendre(1, np.cos(theta0)) - sps.eval_legendre(3, np.cos(theta0))))


if __name__ == '__main__':

    target_patch_size = 0.92
    params = ModelParams(R=150, kappaR=3)
    sigma_m = 0.001

    def fn1(x):
        return expansion.MappedExpansionQuad(a_bar=x, kappaR=params.kappaR, sigma_tilde=sigma_m, l_max=30)

    def fn2(x):
        return expansion.GaussianCharges(lambda_k=x, omega_k=np.array([[0, 0], [np.pi, 0]]), sigma1=0.001, l_max=30)

    def fn3(x):
        return expansion.SphericalCap(theta0_k=x, sigma1=0.001, l_max=50, omega_k=np.array([[0, 0], [np.pi, 0]]))

    a_bar = patch_size.inverse_potential_patch_size(target_patch_size, fn1, 0.5, params)
    lbd = patch_size.inverse_potential_patch_size(target_patch_size, fn2, 5, params)
    theta0 = patch_size.inverse_potential_patch_size(target_patch_size, fn3, 0.5, params)

    ex_point = expansion.MappedExpansionQuad(a_bar=a_bar, kappaR=params.kappaR, sigma_tilde=sigma_m, l_max=30)

    gauss_sigma = point_to_gauss_map(sigma_m, a_bar, lbd, params)
    ex_gauss = expansion.GaussianCharges(lambda_k=lbd, omega_k=np.array([[0, 0], [np.pi, 0]]), sigma1=gauss_sigma, l_max=30)

    cap_sigma = point_to_cap_map(sigma_m, a_bar, theta0, params)
    ex_cap = expansion.SphericalCap(theta0_k=theta0, sigma1=cap_sigma, omega_k=np.array([[0, 0], [np.pi, 0]]), l_max=30)

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

    potential_ic = potentials.inverse_patchy_particle_potential(theta, dist, a_bar, -2 * sigma_m, (sigma_m, sigma_m), params, 30)
    potential1 = potentials.charged_shell_potential(theta, phi, dist, ex_point, params)
    potential2 = potentials.charged_shell_potential(theta, phi, dist, ex_gauss, params)
    potential3 = potentials.charged_shell_potential(theta, phi, dist, ex_cap, params)
    # print(potential.shape)
    # print(potential)

    # expansion.plot_theta_profile_multiple([ex_point, ex_gauss, ex_cap], ['IC', 'Gauss', 'cap'], num=1000)

    fig, ax = plt.subplots()
    ax.scatter(theta[::50], 1000 * potential_ic.T[::50], marker='o', label='ICi', facecolors='none', edgecolors='tab:red')
    ax.plot(theta, 1000 * potential1.T, label='CSp - mapped', linewidth=2)
    # ax.plot(theta, potential_ic.T, label='IC', ls=':', linewidth=2, marker='o', markevery=50, mfc='none')
    ax.plot(theta, 1000 * potential2.T, label='CSp - Gauss', linewidth=2, ls='--')
    ax.plot(theta, 1000 * potential3.T, label='CSp - caps', linewidth=2, ls='--')
    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=12)
    ax.set_xlabel(r'$\theta$', fontsize=15)
    ax.set_ylabel(r'$\phi$ [mV]', fontsize=15)
    custom_ticks = [0, np.pi / 4, np.pi / 2, 3 * np.pi / 4, np.pi]
    custom_labels = ['0', r'$\pi/4$', r'$\pi/2$', r'$3\pi/4$', r'$\pi$']
    plt.axhline(y=0, color='black', linestyle=':')
    plt.axvline(x=target_patch_size, color='black', linestyle=':')
    plt.xticks(custom_ticks, custom_labels, fontsize=13)
    plt.legend(fontsize=12)
    plt.tight_layout()
    plt.savefig(Path("/home/andraz/ChargedShells/Figures/potential_shape_comparison.pdf"), dpi=300)
    plt.show()

    # path_comparison = rotational_path.PathExpansionComparison([ex_point, ex_gauss, ex_cap],
    #                                                           path_plot.QuadPath,
    #                                                           dist=2,
    #                                                           params=params)
    # path_comparison.plot(['IC', 'Gauss', 'cap'],
    #                      save_as=Path("/home/andraz/ChargedShells/Figures/energy_shape_comparison_kR1.png"))