import numpy as np
from charged_shells import expansion, interactions, mapping
from charged_shells.parameters import ModelParams
from functools import partial

Expansion = expansion.Expansion


def sEE_minimum(ex: Expansion, params: ModelParams, accuracy=1e-3, dist=2., match_expansion_axis_to_params=None):
    ex2 = ex.clone()
    zero_to_pi_half = np.linspace(0, np.pi / 2, int(1 / accuracy), endpoint=True)

    ex.rotate_euler(alpha=0, beta=zero_to_pi_half, gamma=0)
    ex2.rotate_euler(alpha=0, beta=zero_to_pi_half, gamma=0)

    energy_fn = mapping.parameter_map_two_expansions(partial(interactions.charged_shell_energy, dist=dist),
                                                     match_expansion_axis_to_params)
    energy = energy_fn(ex, ex2, params)

    min_idx = np.argmin(energy, axis=0)
    min_energy = np.min(energy, axis=0)

    return min_energy, zero_to_pi_half[min_idx]


def main():

    kappaR = np.array([1, 3, 10])
    params = ModelParams(R=150, kappaR=kappaR)
    ex = expansion.MappedExpansionQuad(0.4, kappaR, 0.001)
    min_energy, min_angle = sEE_minimum(ex, params, match_expansion_axis_to_params=1)
    print(min_energy, min_angle + np.pi/2)



if __name__ == '__main__':

    main()