import numpy as np
from matplotlib import gridspec
from matplotlib.lines import Line2D
from charged_shells import functions as fn
import matplotlib.pyplot as plt
import plot_settings


def imp_exp_term(l: int, x: float, abar: float):
    return np.sqrt(2 * l + 1) * fn.sph_bessel_k(l, x) / fn.sph_bessel_k(l+1, x) * abar ** l


def perm_exp_term(l: int, x: float, abar: float):
    return np.sqrt(2 * l + 1) * x ** 2 * fn.sph_bessel_i(l, abar * x) * fn.sph_bessel_k(l, x)


def plot_exp(kappaR: list, abar: float, save_as = None):
    l = np.arange(11)
    fig = plt.figure(figsize=(3, 1.8))
    gs = gridspec.GridSpec(1, 1, figure=fig)
    gs.update(left=0.15, right=0.95, top=0.97, bottom=0.21)
    ax = fig.add_subplot(gs[0, 0])
    for kr in kappaR:
        color = next(plot_settings.COLORS)
        ax.plot(l, imp_exp_term(l, kr, abar) / imp_exp_term(0, kr, abar), label='ICi', marker='o', ls=':', c=color)
        ax.plot(l, perm_exp_term(l, kr, abar) / perm_exp_term(0, kr, abar), label='ICp', marker='^', ls=':', c=color)
    # plt.legend(fontsize=10)
    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=11)
    ax.set_xlabel(r'$l$', fontsize=11)
    ax.set_ylabel('normalized coef', fontsize=11)
    if save_as is not None:
        plt.savefig(save_as, dpi=300)
    plt.show()


def plot_exp2(kappaR: float, abar: list, save_as=None, legend=True):
    l = np.arange(11)

    line1 = Line2D([0], [0], color='black', marker='o', linestyle='None', label='$b_{i,l0}$')
    line2 = Line2D([0], [0], color='black', marker='^', linestyle='None', label='$c_{i,l0}$')

    abar_lines = []

    fig = plt.figure(figsize=(3, 2))
    gs = gridspec.GridSpec(1, 1, figure=fig)
    gs.update(left=0.15, right=0.95, top=0.87, bottom=0.21)
    ax = fig.add_subplot(gs[0, 0])
    for ab in abar:
        color = next(plot_settings.COLORS)
        ax.plot(l, imp_exp_term(l, kappaR, ab) / imp_exp_term(0, kappaR, ab), label='rr', marker='o', ls=':', c=color)
        ax.plot(l, perm_exp_term(l, kappaR, ab) / perm_exp_term(0, kappaR, ab), label='tt', marker='^', ls=':',
                c=color)
        abar_lines.append(Line2D([0], [0], color=color, linewidth=1.2, ls=':', label=rf'$\bar a={ab}$'))

    # main_legend = ax.legend(frameon=False)
    if legend:
        extra_legend = ax.legend(handles=[line1, line2] + abar_lines, loc='upper right', fontsize=11, frameon=False)
    # extra_legend2 = ax.legend(handles=abar_lines, loc='center right', fontsize=11, frameon=False)
    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=11)
    ax.set_xlabel(r'$l$', fontsize=11)
    ax.set_title(f'$\kappa R={kappaR}$', fontsize=11)
    ax.set_ylabel('normalized coef', fontsize=11)
    if save_as is not None:
        plt.savefig(save_as, dpi=300)
    plt.show()


def plot_imp_ex(abar: float):
    kappaR = np.geomspace(0.1, 50, 1000)
    l = np.arange(10)
    fig, ax = plt.subplots()
    ax.plot(kappaR, imp_exp_term(l[None, :], kappaR[:, None], abar))
    ax.set_xscale('log')
    plt.show()


def plot_imp_ex_ratio(abar: float):
    kappaR = np.geomspace(0.1, 50, 1000)
    l = np.arange(10)
    fig, ax = plt.subplots()
    ax.plot(kappaR, imp_exp_term(l[None, :], kappaR[:, None], abar) / imp_exp_term(l[None, :] + 1, kappaR[:, None], abar))
    ax.set_xscale('log')
    plt.show()


def plot_perm_ex(abar: float):
    kappaR = np.geomspace(0.1, 50, 1000)
    l = np.arange(10)
    fig, ax = plt.subplots()
    ax.plot(kappaR, perm_exp_term(l[None, :], kappaR[:, None], abar))
    ax.set_xscale('log')
    plt.show()

def plot_perm_ex_ratio(abar: float):
    kappaR = np.geomspace(0.1, 50, 1000)
    l = np.arange(10)
    fig, ax = plt.subplots()
    ax.plot(kappaR, perm_exp_term(l[None, :], kappaR[:, None], abar) / perm_exp_term(l[None, :] + 1, kappaR[:, None], abar))
    ax.set_xscale('log')
    plt.show()


def plot_kappaR_dep(l: int, abar: float):
    kappaR = np.geomspace(0.1, 50, 1000)
    fig, ax = plt.subplots()
    ax.plot(kappaR, imp_exp_term(l, kappaR, abar))
    ax.plot(kappaR, perm_exp_term(l, kappaR, abar))
    ax.set_xscale('log')
    plt.show()


def plot_kappaR_dep_ratio(l: int, abar: float, l_diff: int = 2):
    kappaR = np.geomspace(0.1, 50, 1000)
    fig, ax = plt.subplots()
    ax.plot(kappaR, imp_exp_term(l + l_diff, kappaR, abar) / imp_exp_term(l, kappaR, abar))
    ax.plot(kappaR, perm_exp_term(l + l_diff, kappaR, abar) / perm_exp_term(l, kappaR, abar))
    ax.set_xscale('log')
    plt.show()


if __name__ == '__main__':

    # plot_exp([1, 10, 50], 0.4,
    #          # save_as=FIGURES_PATH.joinpath('ici_icp_expansion_a08_kr10.png')
    #          )

    plot_exp2(1, [0.5, 0.8], legend=True,
             # save_as=FIGURES_PATH.joinpath('ici_icp_expansion_kr1.png')
             )

    # plot_kappaR_dep_ratio(2, 0.8, l_diff=2)

    # plot_imp_ex_ratio(0.5)
    # plot_perm_ex_ratio(0.5)