import numpy as np
from scipy.optimize import bisect
import expansion
from matplotlib import pyplot as plt
import parameters
import potentials
from pathlib import Path

Expansion = expansion.Expansion
Array = np.ndarray
ModelParams = parameters.ModelParams


@expansion.map_over_expansion
def charge_patch_size(ex: Expansion, phi: float = 0, theta0: Array | float = 0, theta1: Array | float = np.pi / 2):
    return bisect(lambda theta: ex.charge_value(theta, phi), theta0, theta1)


def potential_patch_size(ex: Expansion, params: ModelParams,
                         phi: float = 0, theta0: Array | float = 0, theta1: Array | float = np.pi / 2,
                         match_expansion_axis_to_params: int = None):

    # this is more complicate to map over leading axes of the expansion as potential also depends on model parameters,
    # with some, such as kappaR, also being the parameters of the expansion in the first place. When mapping,
    # we must therefore provide the expansion axis that should match the collection of parameters in params.

    meap = match_expansion_axis_to_params  # just a shorter variable name

    @expansion.map_over_expansion
    def potential_zero(exp: Expansion, prms: ModelParams):
        return bisect(lambda theta: potentials.charged_shell_potential(theta, phi, 1, exp, prms), theta0, theta1)

    params_list = params.unravel()
    if meap is not None:
        expansion_list = [Expansion(ex.l_array, np.take(ex.coefs, i, axis=meap)) for i in range(ex.shape[meap])]
    else:
        expansion_list = [ex for _ in params_list]

    if not len(expansion_list) == len(params_list):
        raise ValueError(f'Axis of expansion that is supposed to match params does not have the same length, got '
                         f'len(params.unravel()) = {len(params_list)} and '
                         f'expansion.shape[{meap}] = {len(expansion_list)}')

    results = []
    for exp, prms in zip(expansion_list, params_list):
        results.append(potential_zero(exp, prms))

    if meap is not None:
        return np.array(results).swapaxes(0, meap)  # return the params-matched axis to where it belongs
    return np.squeeze(np.array(results))


if __name__ == '__main__':

    a_bar = np.linspace(0.2, 0.8, 100)
    kappaR = np.array([0.26, 1, 3, 10, 26])
    params = ModelParams(R=150, kappaR=kappaR)
    ex = expansion.MappedExpansionQuad(a_bar=a_bar[:, None], sigma_m=0.001, l_max=30, kappaR=kappaR[None, :])

    # patch_size = charge_patch_size(ex)
    patch_size = potential_patch_size(ex, params, match_expansion_axis_to_params=1)

    fig, ax = plt.subplots()
    for patch, kR in zip(patch_size.T, kappaR):
        ax.plot(a_bar, patch * 180 / np.pi, label=rf'$\kappa$ = {kR}')
    ax.tick_params(which='both', direction='in', top=True, right=True, labelsize=12)
    ax.set_xlabel(r'$\bar a$', fontsize=13)
    ax.set_ylabel(r'$\theta$', fontsize=13)
    plt.legend(fontsize=12)
    plt.savefig(Path("/home/andraz/ChargedShells/Figures/patch_size_potential.png"), dpi=600)
    plt.show()