"""A collection of methods to compute overlap integrals for modal scattering matrices.
Essentially this involves solving the following integral

.. math::
    K_{abnm} = \iint^{\infty}_{-\infty} U_{nm}(x,y,q_i) M(x,y) U^*_{ab}(x,y,q_o) \, dy \, dx

:math:`U_nm` is the initial modes in the basis :math:`q_i` we are converting from and
:math:`U_ab` are the target modes in a basis :math:`q_o` we are projecting into.

TODO
----
Should explore if decomposing compute_map_knm_matrix_riemann_optimised into real and imaginary
integrals might be faster. In cases where q_in == q_out then integrals are real, apart from
the map component which can be complex.

Explore use of zgemm3m which is 25% faster than zgemm

Probably look into using CUDA if necessary for more speed.
"""
import cython
import numpy as np
cimport numpy as np
from finesse.cymath cimport complex_t
from scipy.linalg.cython_blas cimport zgemm
from scipy.integrate import newton_cotes
from ..cymath.homs cimport unm_workspace, unm_factor_store, u_n__fast, u_m__fast


@cython.boundscheck(False)
@cython.wraparound(False)
@cython.initializedcheck(False)
cpdef void outer_conj_product(complex_t[:,::1] U, complex_t[:,:,::1] result) except *:
    """Computes U * U^C and returns the output into the result array.
    Result array must be of shape (N,N,M) where U is shape (N,M).
    """
    cdef:
        Py_ssize_t i, j, k
        double a, b, c, d

    A = U.shape[0]
    B = U.shape[1]
    if not (result.shape[0] == result.shape[1] == A) or not (result.shape[2] == B):
        raise RuntimeError(f"result array wrong shape should be ({A,A,B})")

    for i in range(A):
        for j in range(A):
            for k in range(B):
                a = U[i, k].real
                b = U[i, k].imag
                c = U[j, k].real
                d = U[j, k].imag
                result[i, j, k] = a*c + b*d + 1j*(b*c - a*d) # U[i, k] * conj(U[j, k])


@cython.boundscheck(False)
@cython.wraparound(False)
@cython.initializedcheck(False)
cdef void c_outer_conj_product(
        Py_ssize_t Nm,
        Py_ssize_t Ns,
        complex_t *U, # complex[Nm, Ns] C-ordered
        complex_t *result  # complex[Nm, Nm, Ns] C-ordered
    ) noexcept nogil:
    """Computes U * U^C and returns the output into the result array.
    Result array must be of shape (Nm,Nm,Ns) where U is shape (Nm,Ns).
    """
    cdef:
        Py_ssize_t i, j, k
        double a, b, c, d

    for i in range(Nm):
        for j in range(Nm):
            for k in range(Ns):
                a = U[i*Ns + k].real #U[i, k].real
                b = U[i*Ns + k].imag
                c = U[j*Ns + k].real
                d = U[j*Ns + k].imag
                result[i*Nm*Ns + j*Ns + k] = a*c + b*d + 1j*(b*c - a*d) # U[i, k] * conj(U[j, k])


@cython.boundscheck(False)
@cython.wraparound(False)
@cython.initializedcheck(False)
cpdef void outer_conj_product_2(complex_t[:,::1] U1, complex_t[:,::1] U2, complex_t[:,:,::1] result) except *:
    """Computes U1 * U2**C and returns the output into the result array.
    Result array must be of shape (N,N,M) where U is shape (N,M).
    """
    cdef:
        Py_ssize_t i, j, k
        double a, b, c, d

    A = U1.shape[0]
    B = U1.shape[1]
    if not (result.shape[0] == result.shape[1] == A) or not (result.shape[2] == B):
        raise RuntimeError(f"result array wrong shape should be ({A,A,B})")

    for i in range(A):
        for j in range(A):
            for k in range(B):
                a = U1[i, k].real
                b = U1[i, k].imag
                c = U2[j, k].real
                d = U2[j, k].imag
                result[i, j, k] = a*c + b*d + 1j*(b*c - a*d) # U1[i, k] * conj(U2[j, k])


@cython.boundscheck(False)
@cython.wraparound(False)
@cython.initializedcheck(False)
cdef void c_outer_conj_product_2(
        Py_ssize_t Nm,
        Py_ssize_t Ns,
        complex_t* U1,
        complex_t* U2,
        complex_t* result
    ) noexcept nogil:
    """Computes U1 * U2**C and returns the output into the result array.
    Result array must be of shape (N,N,M) where U is shape (N,M).
    """
    cdef:
        Py_ssize_t i, j, k
        double a, b, c, d

    for i in range(Nm):
        for j in range(Nm):
            for k in range(Ns):
                a = U1[i*Ns + k].real
                b = U1[i*Ns + k].imag
                c = U2[j*Ns + k].real
                d = U2[j*Ns + k].imag
                result[i*Ns*Nm + j*Ns + k] = a*c + b*d + 1j*(b*c - a*d) # U1[i, k] * conj(U2[j, k])


@cython.boundscheck(False)
@cython.wraparound(False)
@cython.initializedcheck(False)
cdef void update_U_xy_array(
        double *x, Py_ssize_t Nx, double *y, Py_ssize_t Ny,
        complex_t *Un, complex_t *Um, Py_ssize_t Nm,
        unm_workspace *ws, unm_factor_store *store
    ) noexcept nogil:
    """Method to fill Un and Um arrays up to some mode index Nm
    Parameters
    ----------
    x : *double
        Pointer to start of contiguous x array
    Nx : int
        Size of x array
    y : *double
        Pointer to start of contiguous y array
    Ny : int
        Size of y array
    Nm : int
        Max number of modes to calculate up to
    ws
        Pointer to allocated and initialised unm_workspace
    store
        Pointer to allocated and initialised unm_factor_store
    """
    cdef Py_ssize_t i, j
    for i in range(Nm):
        for j in range(Nx):
            Un[i*Nx + j] = u_n__fast(ws, store, i, x[j])

    for i in range(Nm):
        for j in range(Ny):
            Um[i*Ny + j] = u_m__fast(ws, store, i, y[j])


cpdef compute_map_scattering_coeffs_riemann_optimised(
            double dA,
            complex_t[:,::1] Z,
            complex_t[:,:,::1] Unn_,
            complex_t[:,:,::1] Umm_,
            complex_t[:,:,::1] tmp,
            complex_t[:,:,:,::1] result,
        ) :
    """Calculates a mode scattering matrix using a Riemann sum. This method
    uses an computationally optimised approach making use of fast BLAS
    functions. This requires the input modes to be specified in specific
    formats and memory layouts. What this functions computes is the following
    via a Riemann
    sum:

    .. math::
        K_{abnm} = \int^{\infty}_{-\infty}  u_{m}(y,q^y_i) u^*_{b}(y,q^y_o) \Bigg[ \int^{\infty}_{-\infty} Z(x,y) u_{n}(x,q^x_i) u^*_{a}(x,q^x_o) \, dx \Bigg] \, dy

    This integral is not actually performed to infinity, it is bound by the
    dimensions of the discretised map :math:`Z`. The map bound and uniform
    discretisation must be chosen to efficiently sample the size of any of
    the beams and maximum mode order being used.

    Due to the optimised calculation method, the result indexing is not
    [a,b,n,m] but [m,a,n,b]. To convert it back to a more usable
    indexing use:

    >>> result = np.transpose(result, (2,1,0,3))

    Notes
    -----
    Nx - number of x samples
    Ny - number of y samples
    Nm - number of modes (n, m) being calculated

    Parameters
    ----------
    dA : double
        Area of discrete integral, dx * dy

    Z : array[complex]
        2D map of size [Ny, Nx]

    Unn : array[complex]
        3D array of size [Nm, Nm, Nx]. This should contain the
        Un(x) * Un'(x)**C products

    Umm : array[complex]
        3D array of size [Nm, Nm, Ny]. This should contain the
        Um(x) * Um'(x)**C products

    tmp : array[complex]
        Temporary storage that can be used to compute the dot
        products between Z and Unn.
        Should be of size (Nm, Nm, Ny)

    result : array[complex]
        Resulting Knmnm output of size [Nm, Nm, Nm, Nm].
        IMPORTANT: Note output indexing of result K[m,a,n,b]
    """
    if not (Z.ndim == 2):
        raise RuntimeError("Z.ndim != 2")
    if not (result.ndim == 4):
        raise RuntimeError("result.ndim != 2")
    cdef:
        int Nx = Z.shape[1]
        int Ny = Z.shape[0]
        int Nm = result.shape[0]
    if not (result.shape[0] == Nm and result.shape[1] == Nm and result.shape[2] == Nm and result.shape[3] == Nm):
        raise RuntimeError("result.shape != (Nm, Nm, Nm, Nm)")
    if not (Unn_.shape[0] == Nm and Unn_.shape[1] == Nm and Unn_.shape[2] == Nx):
        raise RuntimeError("Unn_.shape != (Nm, Nm, Nx)")
    if not (Umm_.shape[0] == Nm and Umm_.shape[1] == Nm and Umm_.shape[2] == Ny):
        raise RuntimeError("Umm_.shape != (Nm, Nm, Ny)")
    if not (tmp.shape[0] == Nm and tmp.shape[1] == Nm and tmp.shape[2] == Ny):
        raise RuntimeError("tmp.shape != (Nm, Nm, Ny)")
    c_riemann_optimised(
        Nx, Ny, Nm, dA,
        &Z[0,0],
        &Unn_[0,0,0],
        &Umm_[0,0,0],
        &tmp[0,0,0],
        &result[0,0,0,0],
    )


@cython.boundscheck(False)
@cython.wraparound(False)
@cython.initializedcheck(False)
cdef void c_riemann_optimised(
            int Nx,
            int Ny,
            int Nm,
            double dA,
            complex_t* Z, # [Ny, Nx]
            complex_t* Unn_, # [Nm, Nm, Nx]
            complex_t* Umm_, # [Nm, Nm, Ny]
            complex_t* tmp,  # [Nm, Nm, Ny]
            complex_t* result, # [Nm, Nm, Nm, Nm]
        ) noexcept nogil:
    """Raw interface that can be accessed via a Python interface using
    `compute_map_scattering_coeffs_riemann_optimised`, see that function
    for more details.

    Parameters
    ----------
    Nx : int
        Number of x samples

    Ny : int
        Number of y samples

    Nm : int
        Number of 1D modes

    dA : double
        Area of discrete integral, dx * dy

    Z : array[complex]
        2D map of size [Ny, Nx]

    Unn_ : array[complex]
        3D array of size [Nm, Nm, Nx]. This should contain the
        Un(x) * Un'(x)**C products

    Umm_ : array[complex]
        3D array of size [Nm, Nm, Ny]. This should contain the
        Um(x) * Um'(x)**C products

    tmp : array[complex]
        Temporary storage that can be used to compute the dot
        products between Z and Unn_.
        Should be of size (Nm, Nm, Ny)

    result : array[complex]
        Resulting Knmnm output of size [Nm, Nm, Nm, Nm].
        IMPORTANT: Note output indexing of result K[m,a,n,b]
    """
    cdef:
        Py_ssize_t i, k
        int lda, ldb, ldc, M, N, K
        complex_t alpha = dA
        complex_t beta = 0
        complex_t *A
        complex_t *B
        complex_t *C

    # setup blas zgemm variables for map product with x-modes
    # store this in a temporary memory and then apply the
    M = Ny
    N = Nm
    K = Nx
    lda = K
    ldb = K
    ldc = M
    A = Z

    for i in range(Nm):
        B = Unn_ + i*Nm*Nx # &Unn_[i,0,0]
        C = tmp + i*Nm*Ny #&tmp[i,0,0]
        zgemm("T", "N",
            &M, &N, &K,
            &alpha, A, &lda, B, &ldb,
            &beta, C, &ldc
        )

    # Now we need to perform the product of the y-modes
    # with all the map and x-mode products
    M = Nm
    N = Nm
    K = Ny
    lda = Ny
    ldb = Ny
    ldc = Nm
    alpha = 1 # reset as already applied dA
    for i in range(Nm):
        for k in range(Nm):
            A = Umm_ + i*Nm*Ny # &Umm_[i,0,0]
            B = tmp + k*Nm*Ny # &tmp[k,0,0]
            C = result + i*Nm*Nm*Nm + k*Nm*Nm # &result[i,k,0,0]
            zgemm("T", "N",
                &M, &N, &K,
                &alpha, A, &lda, B, &ldb,
                &beta, C, &ldc
            )


def composite_newton_cotes_weights(N, order):
    """
    Constructs the weights for a composite Newton-Cotes rule for integration
    along 1-dimensional line with equally spaced steps. Newton-Cotes are a
    generalisation of a various discrete integration methods that are
    approximating an integrated by some polynomial. Common methods are:

        N = 0 : Riemann sum
        N = 1 : Trapezoid rule
        N = 2 : Simpsons rule
        N = 4 : Simpsons 3/8 rule

    Approximating a large bound with a high order polynomial can be numerically
    problematic. Thus a composite rule is generated by subdividing a larger area
    into multiple chunks and applying each rule along it.

    If the order Newton-Cotes order specified does not produce a rule that fits
    into N, the order is decremented and that is used to fill gaps that do
    no fit.

    See https://mathworld.wolfram.com/Newton-CotesFormulas.html

    Parameters
    ----------
    N : integer >= 0
        1D size of array being ingrated over

    order : integer >= 0
        Order of Newton-Cotes rule to use

    Returns
    -------
    weights : array
        Array of weights to multiply data with before summing
    """
    # ensure we have a 1xN array
    if order == 0:
        return np.ones((N,))
    W = newton_cotes(order, 1)[0]
    wx = np.zeros(N, dtype=np.float64)
    i = 0
    while i < N - 1:
        try:
            wx[i:(i+len(W))] += W
            i += len(W)-1
        except ValueError as ex:
            # if we can't fit a newton-cotes rule in to the space
            # needed then pick a slightly smaller one
            if len(wx[i:(i+len(W))]) < len(W):
                order -= 1
                W = newton_cotes(order, 1)[0]
    return wx


@cython.boundscheck(False)
@cython.wraparound(False)
@cython.initializedcheck(False)
def map_coupling_matrix_riemann(complex_t[:,::1] Y, double dx, double dy, complex_t[:,::1] Un, complex_t[:,::1] Um, long[:,::1] index_map):
    cdef:
        Py_ssize_t i, j, k, l, n, m, nn, mm, Nm
        complex_t a, b
        complex_t[:, ::1] K
        double D = dx * dy
        complex_t sum = 0
        complex_t [:, ::1] Un_
        complex_t [:, ::1] Um_

    Nm = index_map.shape[0]
    Nx = Un.shape[1]
    Ny = Um.shape[1]
    K = np.zeros((Nm, Nm), dtype=np.complex128)
    Un_ = np.conj(Un)
    Um_ = np.conj(Um)

    for i in range(Nm):
        n = index_map[i, 0]
        m = index_map[i, 1]
        for j in range(Nm):
            nn = index_map[j, 0]
            mm = index_map[j, 1]
            sum = 0
            for k in range(Nx):
                for l in range(Ny):
                    a = Un[n, k] * Un_[nn, k]
                    b = Um[m, l] * Um_[mm, l]
                    sum += a * b * Y[k, l]

            K[i, j] = K[i, j]  + sum * D

    return np.asarray(K)

# WIP integration function
# needs to separate the real H(x) functions for integrand and the complex
# scaling factors still.

# @cython.boundscheck(False)
# @cython.wraparound(False)
# @cython.initializedcheck(False)
# cpdef void compute_map_knm_scatterins_riemann_optimised_mode_matched(
#     double dA,
#     double[:,::1] Z_real,
#     double[:,::1] Z_imag,
#     double[:,:,::1] Unn_,
#     double[:,:,::1] Umm_,
#     double[:,:,::1] tmp_real,
#     double[:,:,::1] tmp_imag,
#     double[:,:,:,::1] result_real,
#     double[:,:,:,::1] result_imag,
#     complex_t[:,:,:,::1] result,
# ) except *:
#     cdef:
#         Py_ssize_t i, j, k, l,
#         int Nx, Ny, Nm, inc=1, lda, ldb, ldc, M, N, K
#         double alpha = dA
#         double beta = 0
#         double *A
#         double *B
#         double *C
#         double *X
#         double *Y

#     Nx = Z_real.shape[1]
#     Ny = Z_real.shape[0]
#     Nm = result_real.shape[0] # computing (Nm, Nm, Nm, Nm) output
#     assert(result_real.shape[0] == result_real.shape[1] == result_real.shape[2] == result_real.shape[3])
#     assert(result_imag.shape[0] == result_imag.shape[1] == result_imag.shape[2] == result_imag.shape[3])
#     assert(Unn_.shape[0] == Unn_.shape[1] == result_imag.shape[0] == result_real.shape[0])
#     assert(Unn_.shape[0] == Unn_.shape[1] == result_imag.shape[0] == result_real.shape[0])
#     assert(Unn_.shape[2] == Nx)
#     assert(Umm_.shape[2] == Ny)

#     # setup blas zgemm variables for map product with x-modes
#     # store this in a temporary memory and then apply the
#     M = Ny
#     N = Nm
#     K = Nx
#     lda = K
#     ldb = K
#     ldc = M
#     for i in range(Nm):
#         B = &Unn_[i,0,0]

#         A = &Z_real[0, 0]
#         C = &tmp_real[i,0,0]
#         dgemm("T", "N",
#             &M, &N, &K,
#             &alpha, A, &lda, B, &ldb,
#             &beta, C, &ldc
#         )
#         A = &Z_imag[0, 0]
#         C = &tmp_imag[i,0,0]
#         dgemm("T", "N",
#             &M, &N, &K,
#             &alpha, A, &lda, B, &ldb,
#             &beta, C, &ldc
#         )

#     # Now we need to perform the product of the y-modes
#     # with all the map and x-mode products
#     M = Nm
#     N = Nm
#     K = Ny
#     lda = Ny
#     ldb = Ny
#     ldc = Nm
#     alpha = 1 # reset as already applied dA
#     for i in range(Nm):
#         for k in range(Nm):
#             A = &Umm_[i,0,0]
#             B = &tmp_real[k,0,0]
#             C = &result_real[i,k,0,0]
#             dgemm("T", "N",
#                 &M, &N, &K,
#                 &alpha, A, &lda, B, &ldb,
#                 &beta, C, &ldc
#             )
#             B = &tmp_imag[k,0,0]
#             C = &result_imag[i,k,0,0]
#             dgemm("T", "N",
#                 &M, &N, &K,
#                 &alpha, A, &lda, B, &ldb,
#                 &beta, C, &ldc
#             )

#     for i in range(Nm):
#         for j in range(Nm):
#             for k in range(Nm):
#                 for l in range(Nm):
#                     result[i,j,k,l].real = result_real[i,j,k,l]

#     for i in range(Nm):
#         for j in range(Nm):
#             for k in range(Nm):
#                 for l in range(Nm):
#                     result[i,j,k,l].imag = result_imag[i,j,k,l]
