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

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


def inverse_potential_patch_size(target_patch_size: float,
                                 ex_generator: Callable[[float], Expansion],
                                 x0: float,
                                 params: ModelParams, **ps_kwargs):

    def patch_size_dif(x):
        ex = ex_generator(x)
        return potential_patch_size(ex, params, **ps_kwargs) - target_patch_size

    root_result = root_scalar(patch_size_dif, x0=x0)

    if not root_result.converged:
        raise ValueError('No convergence. Patches of desired size might not be achievable in the given model. '
                         'Conversely, a common mistake might be target patch size input in degrees.')

    return root_result.root


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=10)
    ex = expansion.MappedExpansionQuad(a_bar=a_bar[:, None], sigma_m=0.001, l_max=30, kappaR=kappaR[None, :])

    # fn = lambda x: expansion.SphericalCap(theta0_k=x, sigma1=0.001, l_max=30, omega_k=np.array([[0, 0], [np.pi, 0]]))
    # print(inverse_potential_patch_size(48.8 * np.pi / 180, fn, 0.4, params))

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