import numpy as np
from charged_shells import patch_size, charge_distributions
from charged_shells.parameters import ModelParams
import time
from config import *
from plot_settings import *


def plot_abar_dependence(*, save=False, save_data=False):
    a_bar = np.linspace(0.2, 0.8, 101)
    kappaR = np.array([1, 3, 10])
    params = ModelParams(R=150, kappaR=kappaR)
    ex = charge_distributions.create_mapped_quad_expansion(a_bar=a_bar[:, None], sigma_tilde=0.001, l_max=30, kappaR=kappaR[None, :])

    ps = patch_size.potential_patch_size(ex, params, match_expansion_axis_to_params=1)

    marked_point_ex = charge_distributions.create_mapped_quad_expansion(a_bar=0.5, sigma_tilde=0.001, l_max=30, kappaR=3)
    mark_ps = patch_size.potential_patch_size(marked_point_ex, ModelParams(R=150, kappaR=3),
                                              match_expansion_axis_to_params=None)

    markers = ['o', 's', '^', 'D', 'v', '<', '>', 'p', '*', 'H', '+', 'x']

    if save_data:
        data_dict = {}
        for key, patch in zip(["kR=01", "kR=1", "kR=3", "kR=10", "kR=30"], ps.T):
            data_dict[key] = np.stack((a_bar, patch)).T
        np.savez(FIGURES_DATA_PATH.joinpath("fig_6.npz"), **data_dict)

    # fig, ax = plt.subplots(figsize=plt.figaspect(0.5)/1.5)
    fig, ax = plt.subplots(figsize=(2, 1.7))
    for patch, kR, marker in zip(ps.T, kappaR, markers):
        ax.plot(a_bar, patch, label=rf'$\kappa R$ = {kR}', linewidth=1.5,
                # marker=marker, markerfacecolor='none', markersize=8, markevery=10
                )
    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=11)
    ax.set_xlabel(r'$\bar a$', fontsize=11)
    ax.set_ylabel(r'$\theta_0$', fontsize=11)
    plt.legend(handlelength=0.6, fontsize=9, frameon=False, loc='upper right', bbox_to_anchor=(0.55, 0.63))
    # plt.tight_layout()
    plt.subplots_adjust(left=0.22, right=0.97, top=0.95, bottom=0.22)
    plt.axhline(y=mark_ps, color='black', linestyle=':')
    plt.axvline(x=0.5, color='black', linestyle=':')
    ax.xaxis.set_label_coords(0.5, -0.17)
    if save:
        plt.savefig(FIGURES_PATH.joinpath('final_figures').joinpath('patch_size_potential.png'), dpi=300)
    plt.show()


def plot_abar_charge_dependence(*, save=False):
    a_bar = np.linspace(0.2, 0.8, 101)
    kappaR = np.array([0.1, 1, 3, 10, 30])
    ex = charge_distributions.create_mapped_quad_expansion(a_bar=a_bar[:, None], sigma_tilde=0.001, l_max=30, kappaR=kappaR[None, :])

    ps = patch_size.charge_patch_size(ex)

    markers = ['o', 's', '^', 'D', 'v', '<', '>', 'p', '*', 'H', '+', 'x']

    fig, ax = plt.subplots(figsize=plt.figaspect(0.5)/1.5)
    for patch, kR, marker in zip(ps.T, kappaR, markers):
        ax.plot(a_bar, patch, label=rf'$\kappa R$ = {kR}', marker=marker, markerfacecolor='none', markevery=5)
    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=12)
    ax.set_xlabel(r'$\bar a$', fontsize=15)
    ax.set_ylabel(r'$\theta_0$', fontsize=15)
    plt.legend(fontsize=14)
    plt.tight_layout()
    if save:
        plt.savefig(FIGURES_PATH.joinpath('patch_size_charge.pdf'), dpi=300)
    plt.show()


def plot_kappaR_charge_dependence(*, normalized=False, save=False):
    a_bar = np.array([0.2, 0.4, 0.6, 0.8])
    kappaR = np.linspace(0.01, 30, 100)
    ex = charge_distributions.create_mapped_quad_expansion(a_bar=a_bar[None, :], sigma_tilde=0.001, l_max=30, kappaR=kappaR[:, None])

    ps = patch_size.charge_patch_size(ex)
    if normalized:
        ps = ps / ps[0][None, :]

    markers = ['o', 's', '^', 'D', 'v', '<', '>', 'p', '*', 'H', '+', 'x']

    fig, ax = plt.subplots(figsize=plt.figaspect(0.5)/1.5)
    for patch, abar, marker in zip(ps.T, a_bar, markers):
        ax.plot(kappaR, patch, label=rf'$\bar a$ = {abar}', marker=marker, markerfacecolor='none', markevery=5)
    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=12)
    ax.set_xlabel(r'$\kappa R$', fontsize=15)
    ax.set_ylabel(r'$\theta_0$', fontsize=15)
    plt.legend(fontsize=14)
    plt.tight_layout()
    if save:
        plt.savefig(FIGURES_PATH.joinpath('patch_size_charge_nonmonotonicity.pdf'), dpi=300)
    plt.show()


def plot_sigma0_dependence(*, save=False):

    # kappaR = np.array([0.1, 1, 3, 10, 30])
    kappaR = np.linspace(0.1, 15, 201)
    # sigma0 = np.linspace(-0.001, 0.0015, 6)
    sigma0 = np.array([-0.0002, 0, 0.0002])
    params = ModelParams(R=150, kappaR=kappaR)
    ex = charge_distributions.create_mapped_quad_expansion(a_bar=0.3, sigma_tilde=0.001, l_max=30, kappaR=kappaR, sigma0=sigma0)

    t0 = time.perf_counter()
    ps = patch_size.potential_patch_size(ex, params, match_expansion_axis_to_params=0, skip_nozero_cases=True)
    t1 = time.perf_counter()
    print(f'Time: {t1 - t0}')

    markers = ['o', 's', '^', 'D', 'p', 'v', '<', '>', 'x', '*', 'H', '+']

    fig, ax = plt.subplots(figsize=plt.figaspect(0.5)/1.5)
    for patch, sgm, marker in zip(ps.T, sigma0, markers):
        nonzero_ps = np.nonzero(patch)
        ax.plot(kappaR[nonzero_ps], patch[nonzero_ps], label=rf'$\tilde \sigma_0$ = {sgm}', marker=marker, markerfacecolor='none', markevery=10)
    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=12)
    ax.set_xlabel(r'$\kappa R$', fontsize=15)
    ax.set_ylabel(r'$\theta_0$', fontsize=15)
    plt.legend(fontsize=14, loc='upper right')
    plt.tight_layout()
    if save:
        plt.savefig(FIGURES_PATH.joinpath('patch_size_potential_charge_abar03.pdf'), dpi=300)
    plt.show()


def plot_sigma0_dependence_relative(*, save=False):

    # kappaR = np.array([0.1, 1, 3, 10, 30])
    kappaR = np.linspace(0.1, 15, 201)
    # sigma0 = np.linspace(-0.001, 0.0015, 6)
    sigma0 = np.array([-0.0002, 0, 0.0002])
    params = ModelParams(R=150, kappaR=kappaR)
    ex = charge_distributions.create_mapped_quad_expansion(a_bar=0.8, sigma_tilde=0.001, l_max=30, kappaR=kappaR, sigma0=sigma0)
    ex_neutral = charge_distributions.create_mapped_quad_expansion(a_bar=0.3, sigma_tilde=0.001, l_max=30, kappaR=kappaR, sigma0=0)

    ps = patch_size.potential_patch_size(ex, params, match_expansion_axis_to_params=0, skip_nozero_cases=True)
    ps_neutral = patch_size.potential_patch_size(ex_neutral, params, match_expansion_axis_to_params=0)

    markers = ['o', 's', '^', 'D', 'p', 'v', '<', '>', 'x', '*', 'H', '+']

    fig, ax = plt.subplots(figsize=plt.figaspect(0.5)/1.5)
    for patch, sgm, marker in zip(ps.T, sigma0, markers):
        nonzero_ps = np.nonzero(patch)
        ax.plot(kappaR[nonzero_ps], patch[nonzero_ps] / ps_neutral[nonzero_ps],
                label=rf'$\tilde \sigma_0$ = {sgm}', marker=marker, markerfacecolor='none', markevery=10)
    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=12)
    ax.set_xlabel(r'$\kappa R$', fontsize=15)
    ax.set_ylabel(r'$\theta_0$', fontsize=15)
    plt.legend(fontsize=14)
    plt.tight_layout()
    if save:
        plt.savefig(FIGURES_PATH.joinpath('patch_size_potential_charge_abar03_relative.pdf'), dpi=300)
    plt.show()


def main():

    plot_abar_dependence(save=True)

    # plot_sigma0_dependence(save=True)

    # plot_sigma0_dependence_relative(save=True)

    # plot_abar_charge_dependence(save=True)
    # plot_kappaR_charge_dependence(normalized=True, save=True)


if __name__ == '__main__':

    main()