from charged_shells import potentials, patch_size, charge_distributions
import numpy as np
from charged_shells.parameters import ModelParams
from matplotlib.lines import Line2D
from config import *
import quadrupole_model_mappings
from plot_settings import *


def ic_cs_comparison():
    kappaR = 15
    params = ModelParams(R=150, kappaR=kappaR)
    sigma_m = 0.001
    sigma0 = 0.0001
    a_bar = 0.3

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

    ex_cp = charge_distributions.create_mapped_quad_expansion(a_bar=a_bar, kappaR=params.kappaR, sigma_tilde=sigma_m, l_max=30,
                                                                    sigma0=sigma0)
    potential_cp = potentials.charged_shell_potential(theta, phi, dist, ex_cp, params)

    potential_ic = potentials.inverse_patchy_particle_potential(theta, dist, a_bar, -2 * sigma_m + sigma0,
                                                                (sigma_m, sigma_m), params, 30)

    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 * potential_cp.T, label='CSp - mapped', linewidth=2)
    # ax.plot(1000 * potential_ic.T - 1000 * potential_cp.T)
    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.xticks(custom_ticks, custom_labels, fontsize=13)
    plt.legend(fontsize=12)
    plt.tight_layout()
    # plt.savefig(Path("/home/andraz/ChargedShells/Figures/potential_shape_comparison_overcharged05.pdf"), dpi=300)
    plt.show()


def models_comparison(save_data=False):
    target_patch_size = 0.92
    kappaR = 3
    params = ModelParams(R=150, kappaR=kappaR)
    sigma_m = 0.001
    sigma0 = 0  # taking this total charge parameter nonzero causes cap model to fail in finding the appropriate theta0

    def fn(x):
        return charge_distributions.create_mapped_quad_expansion(a_bar=x, kappaR=params.kappaR, sigma_tilde=sigma_m, l_max=30, sigma0=sigma0)

    a_bar = patch_size.inverse_potential_patch_size(target_patch_size, fn, 0.5, params)
    ex_point = charge_distributions.create_mapped_quad_expansion(a_bar=a_bar, kappaR=params.kappaR, sigma_tilde=sigma_m, l_max=30, sigma0=sigma0)

    ex_gauss = quadrupole_model_mappings.ic_to_gauss(sigma_m, a_bar, params, l_max=30, sigma0=sigma0)
    ex_cap = quadrupole_model_mappings.ic_to_cap(sigma_m, a_bar, params, l_max=50, sigma0=sigma0)

    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 + sigma0,
                                                                (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)

    if save_data:
        with open(Path("/analysis/config.json")) as config_file:
            config_data = json.load(config_file)
        np.savez(Path(config_data["figure_data"]).joinpath("fig_6_shape_models.npz"),
                 ICi=np.stack((theta, potential_ic)).T,
                 CSp_mapped=np.stack((theta, potential1)).T,
                 CSp_gauss=np.stack((theta, potential2)).T,
                 CSp_caps=np.stack((theta, potential3)).T)

    # 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()


def abar_comparison():
    target_patch_sizes = [0.8, 0.85, 0.92]
    params = ModelParams(R=150, kappaR=3)
    sigma_m = 0.001

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

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

    pots = []
    for ps in target_patch_sizes:
        a_bar = patch_size.inverse_potential_patch_size(ps, fn1, 0.5, params)
        ex_point = charge_distributions.create_mapped_quad_expansion(a_bar=a_bar, kappaR=params.kappaR, sigma_tilde=sigma_m, l_max=30)
        pots.append(potentials.charged_shell_potential(theta, phi, dist, ex_point, params))

    fig, ax = plt.subplots(figsize=(4.125, 3))
    for pot, ps in zip(pots, target_patch_sizes):
        ax.plot(theta, 1000 * pot, label=fr'$\theta_0={180 / np.pi * ps:.1f}^o$', linewidth=2)
    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.xticks(custom_ticks, custom_labels, fontsize=15)
    plt.legend(fontsize=12)
    plt.tight_layout()
    # plt.savefig(Path("/home/andraz/ChargedShells/Figures/potential_amplitude_comparison.pdf"), dpi=300)
    plt.show()


def charge_comparsion():

    charges = np.array([-0.001, 0, 0.001, 0.002])
    a_bar = 0.6
    params = ModelParams(R=150, kappaR=10)
    sigma_m = 0.001

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

    ex_point = charge_distributions.create_mapped_quad_expansion(a_bar=a_bar, kappaR=params.kappaR, sigma_tilde=sigma_m,
                                                                       l_max=30, sigma0=charges)
    pot = potentials.charged_shell_potential(theta, phi, dist, ex_point, params)

    fig, ax = plt.subplots(figsize=(4.125, 3))
    for p, lbl in zip(pot, [fr'$\sigma_0={c}$' for c in charges]):
        ax.plot(theta, 1000 * p, linewidth=2, label=lbl)
    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.xticks(custom_ticks, custom_labels, fontsize=15)
    plt.legend(fontsize=12)
    plt.tight_layout()
    # plt.savefig(FIGURES_PATH.joinpath('potential_charge_comparison.pdf'), dpi=300)
    plt.show()


# TODO: comparison of patch shape at different kappaR and a neutral particle, with fixed patch size


def ic_cs_comparison2(save_data=False):
    # target_patch_size = 0.92
    kappaR = 3
    params = ModelParams(R=150, kappaR=kappaR)
    sigma_m = 0.001
    sigma0_array = np.array([-0.0002, 0, 0.0002])

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

    def fn1(x):
        return charge_distributions.create_mapped_quad_expansion(a_bar=x, kappaR=params.kappaR, sigma_tilde=sigma_m, l_max=30, sigma0=0)
    # a_bar is determined only for a neutral particle (patch size only well-defined in this case)
    # a_bar = patch_size.inverse_potential_patch_size(target_patch_size, fn1, 0.5, params)
    # print(a_bar)

    a_bar = 0.5
    target_patch_size = patch_size.potential_patch_size(
        charge_distributions.create_mapped_quad_expansion(a_bar=a_bar,
                                                                kappaR=params.kappaR,
                                                                sigma_tilde=sigma_m,
                                                                l_max=30,
                                                                sigma0=0),
        params)

    potential_ic = []
    potential_cs = []
    for sigma0 in sigma0_array:

        ex_point = charge_distributions.create_mapped_quad_expansion(a_bar=a_bar, kappaR=params.kappaR,
                                                                           sigma_tilde=sigma_m, l_max=30, sigma0=sigma0)

        potential_ic.append(potentials.inverse_patchy_particle_potential(theta, dist, a_bar, -2 * sigma_m + sigma0,
                                                                    (sigma_m, sigma_m), params, 30))
        potential_cs.append(potentials.charged_shell_potential(theta, phi, dist, ex_point, params))

    if save_data:
        with open(Path("/analysis/config.json")) as config_file:
            config_data = json.load(config_file)
        np.savez(Path(config_data["figure_data"]).joinpath("fig_6a.npz"),
                 ICi_eta_neg=np.stack((theta, potential_ic[0])).T,
                 ICi_eta_0=np.stack((theta, potential_ic[1])).T,
                 ICi_eta_pos=np.stack((theta, potential_ic[2])).T,
                 CSp_eta_neg=np.stack((theta, potential_cs[0])).T,
                 CSp_eta_0=np.stack((theta, potential_cs[1])).T,
                 CSp_eta_pos=np.stack((theta, potential_cs[2])).T)

    fig, ax = plt.subplots(figsize=(2, 1.7))
    lines = []
    # lines.append(plt.scatter([], [], marker='x', c='k', label='CSp'))
    # lines.append(plt.scatter([], [], marker='o', facecolor='none', edgecolor='k', label='ICi'))
    for p_ic, p_cs, charge in zip(potential_ic, potential_cs, sigma0_array):
        l, = ax.plot(theta, 1000 * p_cs.T, label=rf'$\eta = {charge/sigma_m}$', linewidth=1.5, zorder=0)
        l2, = ax.plot(theta, 1000 * p_cs.T, linewidth=1.5, zorder=0, c='k', ls='--', dashes=(5, 5))
        lines.append(l)
        # ax.scatter(theta[::100], 1000 * p_ic.T[::100], marker='o', facecolors='none', edgecolors='k', s=50)
        # ax.scatter(theta[::100], 1000 * p_cs.T[::100], marker='x', c='k', s=50)
    # ax.plot(theta, potential_ic.T, label='IC', ls=':', linewidth=2, marker='o', markevery=50, mfc='none')
    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=11)
    ax.set_xlabel(r'$\theta$', fontsize=11)
    ax.set_ylabel(r'$\psi [mV]$', fontsize=11)
    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=11)

    # lines.append(Line2D([0], [0], color='k', label='CSp'))
    # lines.append(plt.scatter([], [], marker='o', facecolor='none', edgecolor='red', label='ICi'))
    plt.legend(handles=lines, fontsize=9, handlelength=0.8, frameon=False,
               bbox_to_anchor=(0.57, 1.0), loc='upper center'
               )

    # plt.tight_layout()
    plt.subplots_adjust(left=0.22, right=0.97, top=0.95, bottom=0.22)
    ax.xaxis.set_label_coords(0.5, -0.17)
    plt.savefig(FIGURES_PATH.joinpath('final_figures').joinpath('potential_ic_cs_comparison.png'), dpi=300)
    plt.show()


def main():
    # models_comparison(save_data=False)
    # ic_cs_comparison()
    ic_cs_comparison2()

    # abar_comparison()
    # charge_comparsion()


if __name__ == '__main__':

    main()