import expansion
import parameters
import functions as fn
import numpy as np
from py3nj import wigner3j
import time
import matplotlib.pyplot as plt


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


# 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

    # if all elements of bool are false, i.e. missing_l > all existing_l, additional l values come to the end
    indices1 = np.where(np.any(bool1, axis=0), np.argmax(bool1, axis=0), full_l_array1.shape[0])
    indices2 = np.where(np.any(bool2, axis=0), np.argmax(bool2, axis=0), full_l_array2.shape[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)

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


def charged_shell_energy(ex1: Expansion, ex2: Expansion, dist: float, params: ModelParams):

    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, :], params.kappaR)
    full_coefficient_C = ex1.repeat_over_m(ex2.repeat_over_m(coefficient_C, axis=1), axis=0)

    charge_factor = np.real(ex1.coeffs[:, None] * np.conj(ex2.coeffs[None, :]) +
                            (-1) ** (full_l_array[:, None] + full_l_array[None, :])
                            * ex1.coeffs[None, :] * ex1.coeffs[:, None])

    indices_l, indices_p = np.nonzero(full_m_array[:, None] == full_m_array[None, :])
    flat_l = full_l_array[indices_l]
    flat_p = full_l_array[indices_p]
    flat_m = full_m_array[indices_l]  # the same as full_m_array[indices_p]

    all_s_array = np.arange(2 * ex1.max_l + 1)
    bessels = fn.sph_bessel_k(all_s_array, params.kappa * dist)

    s_bool1 = np.abs(flat_l[:, None] - all_s_array[None, :]) <= flat_p[:, None]
    s_bool2 = flat_p[:, None] <= (flat_l[:, None] + all_s_array[None, :])
    indices_lpm, indices_s = np.nonzero(s_bool1 * s_bool2)

    l_vals = flat_l[indices_lpm]
    p_vals = flat_p[indices_lpm]
    m_vals = flat_m[indices_lpm]
    C_vals = full_coefficient_C[indices_l, indices_p][indices_lpm]
    charge_vals = charge_factor[indices_l, indices_p][indices_lpm]
    s_vals = all_s_array[indices_s]
    bessel_vals = bessels[indices_s]

    lps_terms = (2 * s_vals + 1) * np.sqrt((2 * l_vals + 1) * (2 * p_vals + 1))

    wigner1 = wigner3j(l_vals, s_vals, p_vals,
                       0, 0, 0)
    wigner2 = wigner3j(l_vals, s_vals, p_vals,
                       -m_vals, 0, m_vals)

    return (params.R ** 2 / (params.kappa * params.epsilon * params.epsilon0)
            * np.sum(C_vals * (-1) ** (l_vals + m_vals) * charge_vals * lps_terms * bessel_vals * wigner1 * wigner2))


if __name__ == '__main__':

    params = ModelParams(1, 3, 1, 1)
    ex1 = expansion.MappedExpansion(1, params.kappaR, 0.001, max_l=5)
    ex2 = expansion.MappedExpansion(1, params.kappaR, 0.001, max_l=5)

    dist = 2.

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

    t0 = time.perf_counter()
    energy = charged_shell_energy(ex1, ex2, dist, params)
    t1 = time.perf_counter()

    print('energy: ', energy)
    print('time: ', t1 - t0)

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