from analysis.peak_heigth import *

Array = np.ndarray
Expansion = expansion.Expansion
RotConfig = Literal['ep', 'pp', 'sep']


def get_kappaR_higher_multipole_dicts(peak: Peak, params: ModelParams, emanuele_data: Array, abar: float,
                            kappaR: Array, dist_cs: list, max_kappaR: float = 50, min_kappa: float = 0.01,
                                      only_loaded: bool = False):

    correct_abar_indices, = np.nonzero(np.isclose(emanuele_data[:, 1], abar))
    emanuele_data = emanuele_data[correct_abar_indices]

    distances, inverse = np.unique(emanuele_data[:, 0], return_inverse=True)

    energy_fn = mapping.parameter_map_two_expansions(interactions.charged_shell_energy,
                                                     peak.kappaR_axis_in_expansion)

    energy = []
    for d in dist_cs:
        if not only_loaded:
            energy.append(energy_fn(peak.ex1, peak.ex2, params, dist=d))
        else:
            energy.append(np.full((len(kappaR),), 1.))

    data_dict_ic = {}
    data_dict_cs = {}
    k = 0
    for i in range(len(distances)):
        d = np.around(distances[i], 5)
        if d in np.around(dist_cs, 5):
            idx, = np.nonzero(inverse == i)
            relevant_kR_indices = (emanuele_data[idx, 2] <= max_kappaR) * (min_kappa <= emanuele_data[idx, 2])
            kR = emanuele_data[idx, 2][relevant_kR_indices]
            en = emanuele_data[idx, peak.ici_data_column][relevant_kR_indices]
            sort = np.argsort(kR)
            if peak.log_y:
                data_dict_ic[d] = np.stack((kR[sort], np.abs(en)[sort])).T
                data_dict_cs[d] = np.stack((kappaR, np.abs(energy[k]))).T
            else:
                data_dict_ic[d] = np.stack((kR[sort], en[sort])).T
                data_dict_cs[d] = np.stack((kappaR, energy[k])).T
            k += 1

    return data_dict_ic, data_dict_cs


def IC_peak_energy_plot(dist: list,
                        which: RotConfig,
                        abar=0.8,
                        min_kappaR: float = 0.01,
                        max_kappaR: float = 50.,
                        R: float = 150,
                        save_as: Path = None,
                        save_data: bool = False,
                        quad_correction: bool = False,
                        log_y: bool = True):

    # em_data_path = (ICI_DATA_PATH.joinpath("FIG_11"))
    em_data_path = (ICI_DATA_PATH.joinpath("TEST_HIGHER_MULTIPOLES"))
    em_data = np.load(em_data_path.joinpath("higher_multiplot_energy_2.npz"))
    data2 = em_data['L2']
    data4 = em_data['L4']
    data6 = em_data['L6']


    num = 100 if which == 'sep' else 30
    kappaR = np.geomspace(min_kappaR, max_kappaR, num)
    params = ModelParams(R=R, kappaR=kappaR)

    ex = charge_distributions.create_mapped_quad_expansion(abar, params.kappaR, 0.00099, l_max=20)

    if which == 'ep':
        peak = PeakEP(ex, log_y, kappaR_axis_in_expansion=0)
    elif which == 'pp':
        peak = PeakPP(ex, log_y, kappaR_axis_in_expansion=0)
    elif which == 'sep':
        peak = PeakSEP(ex, log_y, kappaR_axis_in_expansion=0)
    else:
        raise ValueError

    data_dict_ic2, data_dict_cs = get_kappaR_higher_multipole_dicts(peak, params, data2,abar, kappaR, dist, max_kappaR, min_kappaR)
    data_dict_ic4, _ = get_kappaR_higher_multipole_dicts(peak, params, data4, abar, kappaR, dist, max_kappaR, min_kappaR)
    data_dict_ic6, _ = get_kappaR_higher_multipole_dicts(peak, params, data6, abar, kappaR, dist, max_kappaR, min_kappaR)

    colors = cycle(plt.rcParams['axes.prop_cycle'].by_key()['color'])

    fig = plt.figure(figsize=(2, 2))
    gs = gridspec.GridSpec(1, 1, figure=fig)
    gs.update(left=0.25, right=0.95, top=0.9, bottom=0.21)
    ax = fig.add_subplot(gs[0, 0])
    for (key, data2), data_cs, data4, data6 in zip(data_dict_ic2.items(), data_dict_cs.values(), data_dict_ic4.values(), data_dict_ic6.values()):
        current_color = next(colors)
        ax.plot(data2[:, 0], data2[:, 1], label=r'$l=2$', c=current_color)
        ax.plot(data_cs[:, 0], data_cs[:, 1], ls='--', c='k')
        ax.plot(data4[:, 0], data4[:, 1], ls='-', c=next(colors), label=r'$l=4$')
        ax.plot(data6[:, 0], data6[:, 1], ls='-', c=next(colors), label=r'$l=6$')
    ax.legend(fontsize=9, ncol=1, frameon=False, handlelength=0.7, loc='upper right',
              bbox_to_anchor=(0.5, 0.5),
              # bbox_to_anchor=(0.42, 0.9)
              )
    ax.set_title(f'$\\bar a = {abar:.1f}$, $\\rho = {key:.1f}$', loc='center',
                 # x=0.95, y=0.8,
                 y=0.96,
                 # fontsize=9
                 )
    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=11)
    ax.set_xlabel(r'$\kappa R$', fontsize=11)

    ax.set_ylabel(peak.y_label, fontsize=11)
    if log_y:
        ax.set_yscale('log')
    ax.set_xscale('log')
    # plt.tight_layout()
    if save_as is not None:
        plt.savefig(save_as, dpi=300)
    plt.show()


if __name__ == '__main__':

    dist = [3.]
    IC_peak_energy_plot(abar=0.5, dist=dist, which='pp', min_kappaR=0.1, max_kappaR=50,
                        save_as=FIGURES_PATH.joinpath("ICi_data").joinpath("peak_pp_kappaR_higher_correction_abar05_rho30.png"),
                        save_data=False,
                        quad_correction=False,
                        log_y=False
                        )