import numpy as np
import numba

from despasito.equations_of_state import constants


def calc_Xika(indices, rho, xi, molecular_composition, nk, Fklab, Kklab, gr_assoc):
    r"""
    A wrapper to calculate the fraction of molecules of component i that are not
    bonded at a site of type a on group k. Switched between functions for different
    Kklab

    Parameters
    ----------
    indices : list[list]
        A list of sets of (component, bead, site) to identify the values of the Xika
        matrix that are being fit
    rho : numpy.ndarray
        Number density of system [mol/m^3]
    xi : numpy.ndarray
        Mole fraction of each component, sum(xi) should equal 1.0
    molecular_composition : numpy.array
        :math:`\nu_{i,k}/k_B`, Array of number of components by number of bead types.
        Defines the number of each type of group in each component.
    nk : numpy.ndarray
        For each bead the number of each type of site
    Fklab : numpy.ndarray
        The association strength between a site of type a on a group of type k of
        component i and a site of type b on a group of type l of component j., known
        as the Mayer f-function.
    Kklab : numpy.ndarray
        Bonding volume between each association site
    gr_assoc : numpy.ndarray
        Reference fluid pair correlation function used in calculating association
        sites, (len(rho) x Ncomp x Ncomp)

    Returns
    -------
    Xika : numpy.ndarray
        The fraction of molecules of component i that are not bonded at a site of
        type a on group k. Matrix (len(rho) x Ncomp x Nbeads x len(sitenames))
    err_array : numpy.ndarray
        Of the same length of rho, is a list in the error of the total error Xika for
        each point.
    """

    l_K = len(np.shape(Kklab))

    # Ensure all inputs are numpy arrays
    tmp_array = [rho, xi, molecular_composition, nk, Fklab, Kklab, gr_assoc]
    for i, tmp in enumerate(tmp_array):
        if np.shape(tmp):
            tmp_array[i] = np.array(tmp)
        else:
            tmp_array[i] = np.array([tmp])
    rho, xi, molecular_composition, nk, Fklab, Kklab, gr_assoc = tmp_array

    if l_K == 4:
        Xika_final, err_array = calc_Xika_4(indices, rho, xi, molecular_composition, nk, Fklab, Kklab, gr_assoc)
    if l_K == 6:
        Xika_final, err_array = calc_Xika_6(indices, rho, xi, molecular_composition, nk, Fklab, Kklab, gr_assoc)

    return Xika_final, err_array


# @numba.njit(
#    numba.types.Tuple((numba.f8[:, :], numba.f8[:]))(
#        numba.i8[:, :],
#        numba.f8[:],
#        numba.f8[:],
#        numba.f8[:, :],
#        numba.i8[:, :],
#        numba.f8[:, :, :, :],
#        numba.f8[:, :, :, :],
#        numba.f8[:, :, :],
#    )
# )
@numba.njit()
def calc_Xika_4(
    indices, rho, xi, molecular_composition, nk, Fklab, Kklab, gr_assoc
):  # , maxiter=500, tol=1e-12, damp=.1
    r"""
    Calculate the fraction of molecules of component i that are not bonded at a
    site of type a on group k.

    Parameters
    ----------
    indices : list[list]
        A list of sets of (component, bead, site) to identify the values of the Xika
        matrix that are being fit
    rho : numpy.ndarray
        Number density of system [mol/m^3]
    xi : numpy.ndarray
        Mole fraction of each component, sum(xi) should equal 1.0
    molecular_composition : numpy.array
        :math:`\nu_{i,k}/k_B`, Array of number of components by number of bead types.
        Defines the number of each type of group in each component.
    nk : numpy.ndarray
        For each bead the number of each type of site
    Fklab : numpy.ndarray
        The association strength between a site of type a on a group of type k of
        component i and a site of type b on a group of type l of component j., known
        as the Mayer f-function.
    Kklab : numpy.ndarray
        Bonding volume between each association site
    gr_assoc : numpy.ndarray
        Reference fluid pair correlation function used in calculating association
        sites, (len(rho) x Ncomp x Ncomp)

    Returns
    -------
    Xika : numpy.ndarray
        The fraction of molecules of component i that are not bonded at a site of
        type a on group k. Matrix (len(rho) x Ncomp x Nbeads x len(sitenames))
    err_array : numpy.ndarray
        Of the same length of rho, is a list in the error of the total error Xika
        for each point.
    """

    maxiter = 500
    tol = 1e-12
    damp = 0.1

    nrho = len(rho)
    l_ind = len(indices)

    #    Xika_final = np.ones((nrho,ncomp, nbeads, nsitesmax))
    Xika_final = np.ones((nrho, l_ind))
    err_array = np.zeros(nrho)

    # Parallelize here, wrt rho!
    Xika_elements = 0.5 * np.ones(l_ind)
    for r in range(nrho):
        for knd in range(maxiter):

            Xika_elements_new = np.ones(l_ind)
            ind = 0
            for iind in range(l_ind):
                i, k, a = indices[iind]
                jnd = 0
                for jjnd in range(l_ind):
                    j, l, b = indices[jjnd]
                    delta = Fklab[k, l, a, b] * Kklab[k, l, a, b] * gr_assoc[r, i, j]
                    Xika_elements_new[ind] += (
                        constants.molecule_per_nm3
                        * rho[r]
                        * xi[j]
                        * molecular_composition[j, l]
                        * nk[l, b]
                        * Xika_elements[jnd]
                        * delta
                    )
                    jnd += 1
                ind += 1

            Xika_elements_new = 1.0 / Xika_elements_new
            obj = np.sum(np.abs(Xika_elements_new - Xika_elements))

            if obj < tol:
                break
            else:
                if obj / np.max(Xika_elements) > 1e3:
                    Xika_elements = Xika_elements + damp * (Xika_elements_new - Xika_elements)
                else:
                    Xika_elements = Xika_elements_new

        err_array[r] = obj
        Xika_final[r, :] = Xika_elements

    return Xika_final, err_array


# @numba.njit(
#    numba.types.Tuple((numba.f8[:, :], numba.f8[:]))(
#        numba.i8[:, :],
#        numba.f8[:],
#        numba.f8[:],
#        numba.f8[:, :],
#        numba.i8[:, :],
#        numba.f8[:, :, :, :],
#        numba.f8[:, :, :, :, :, :],
#        numba.f8[:, :, :],
#    )
# )
@numba.njit()
def calc_Xika_6(
    indices, rho, xi, molecular_composition, nk, Fklab, Kklab, gr_assoc
):  # , maxiter=500, tol=1e-12, damp=.1
    r"""
    Calculate the fraction of molecules of component i that are not bonded at a site
    of type a on group k.

    Parameters
    ----------
    indices : list[list]
        A list of sets of (component, bead, site) to identify the values of the Xika
        matrix that are being fit
    rho : numpy.ndarray
        Number density of system [mol/m^3]
    xi : numpy.ndarray
        Mole fraction of each component, sum(xi) should equal 1.0
    molecular_composition : numpy.array
        :math:`\nu_{i,k}/k_B`, Array of number of components by number of bead types.
        Defines the number of each type of group in each component.
    nk : numpy.ndarray
        For each bead the number of each type of site
    Fklab : numpy.ndarray
        The association strength between a site of type a on a group of type k of
        component i and a site of type b on a group of type l of component j., known
        as the Mayer f-function.
    Kklab : numpy.ndarray
        Bonding volume between each association site
    gr_assoc : numpy.ndarray
        Reference fluid pair correlation function used in calculating association
        sites, (len(rho) x Ncomp x Ncomp)

    Returns
    -------
    Xika : numpy.ndarray
        The fraction of molecules of component i that are not bonded at a site of type
        a on group k. Matrix (len(rho) x Ncomp x Nbeads x len(sitenames))
    err_array : numpy.ndarray
        Of the same length of rho, is a list in the error of the total error Xika for
        each point.
    """

    maxiter = 500
    tol = 1e-12
    damp = 0.1

    nrho = len(rho)
    l_ind = len(indices)

    #    Xika_final = np.ones((nrho,ncomp, nbeads, nsitesmax))
    Xika_final = np.ones((nrho, l_ind))
    err_array = np.zeros(nrho)

    # Parallelize here, wrt rho!
    Xika_elements = 0.5 * np.ones(l_ind)
    for r in range(nrho):
        for knd in range(maxiter):

            Xika_elements_new = np.ones(l_ind)
            ind = 0
            for iind in range(l_ind):
                i, k, a = indices[iind]
                jnd = 0
                for jjnd in range(l_ind):
                    j, l, b = indices[jjnd]
                    delta = Fklab[k, l, a, b] * Kklab[i, j, k, l, a, b] * gr_assoc[r, i, j]
                    Xika_elements_new[ind] += (
                        constants.molecule_per_nm3
                        * rho[r]
                        * xi[j]
                        * molecular_composition[j, l]
                        * nk[l, b]
                        * Xika_elements[jnd]
                        * delta
                    )
                    jnd += 1
                ind += 1

            Xika_elements_new = 1.0 / Xika_elements_new
            obj = np.sum(np.abs(Xika_elements_new - Xika_elements))

            if obj < tol:
                break
            else:
                if obj / np.max(Xika_elements) > 1e3:
                    Xika_elements = Xika_elements + damp * (Xika_elements_new - Xika_elements)
                else:
                    Xika_elements = Xika_elements_new

        err_array[r] = obj
        Xika_final[r, :] = Xika_elements

    return Xika_final, err_array
