import expansion
import functions as fn
import numpy as np
import matplotlib.pyplot as plt


Array = np.ndarray
Expansion = expansion.Expansion


def prefactor(R: float, kappaR: float, c0: float):
    return R * kappaR * 1e4 / (12.04 * c0)


def c0(R: float, kappaR: float):
    return 10 * kappaR ** 2 / (0.329 ** 2 * R ** 2)


def expansions_to_common_l(ex1: Expansion, ex2: expansion) -> (Expansion, Expansion):
    common_l_array = np.union1d(ex1.l_array, ex2.l_array)
    missing_l1 = np.setdiff1d(common_l_array, ex1.l_array, assume_unique=True)
    missing_l2 = np.setdiff1d(common_l_array, ex2.l_array, assume_unique=True)

    fill_1 = np.zeros(np.sum(2 * missing_l1 + 1))
    fill_2 = np.zeros(np.sum(2 * missing_l2 + 1))

    full_l_array1, _ = ex1.lm_arrays
    full_l_array2, _ = ex2.lm_arrays

    # we search for where to place missing coeffs with the help of a boolean array and argmax function
    bool1 = (full_l_array1[:, None] - missing_l1[None, :]) > 0
    bool2 = (full_l_array2[:, None] - missing_l2[None, :]) > 0

    # we set last element to True so that argmax returns last idx if all missing l > max_l
    bool1[-1, :] = True
    bool2[-1, :] = True

    indices1 = np.argmax(bool1, axis=0)
    indices2 = np.argmax(bool2, axis=0)

    new_coeffs1 = np.insert(ex1.coeffs, np.repeat(indices1, 2 * missing_l1 + 1), fill_1)
    new_coeffs2 = np.insert(ex2.coeffs, np.repeat(indices2, 2 * missing_l2 + 1), fill_2)

    assert len(new_coeffs1) == len(new_coeffs2)

    return Expansion(common_l_array, new_coeffs1), Expansion(common_l_array, new_coeffs2)


def charged_shell_energy(ex1: Expansion, ex2: Expansion, dist: float, kappaR: float, R: float):

    ex1, ex2 = expansions_to_common_l(ex1, ex2)

    full_l_array, full_m_array = ex1.lm_arrays

    coefficient_C = fn.interaction_coeff_C(ex1.l_array[:, None], ex2.l_array[None, :], kappaR)
    full_coefficient_C = ex1.repeat_over_m(ex2.repeat_over_m(coefficient_C, axis=1), axis=0)

    indices, _ = np.nonzero(full_m_array[:, None] == full_m_array[None, :])

    flat_l = full_l_array[indices]
    flat_m = full_m_array[indices]

    flat_C = full_coefficient_C[indices1, indices2]

    flat_sigma1 = ex1.coeffs[indices1]
    flat_sigma2 = ex2.coeffs[indices2]

    # charge_factor = -1 ** (flat_l + flat_m) * np.real(flat_sigma1 * np.conj(flat_sigma2) + (-1) ** (flat_l + flat_p) * )

    return



if __name__ == '__main__':

    kappaR = 3
    R = 150
    ex1 = expansion.MappedExpansion(1, kappaR, 0.001, max_l=10)
    ex2 = expansion.MappedExpansion(1, kappaR, 0.001, max_l=5)

    dist = 2.

    ex1, ex2 = expansions_to_common_l(ex1, ex2)
    print(ex1.coeffs)
    print(ex2.coeffs)

    # energy = charged_shell_energy(ex1, ex2, dist, kappaR, R)
    # print(potential)

    # plt.plot(energy)
    # plt.show()