from .qpms_cdefs cimport *
from .qpms_c cimport ScatteringSystem, _ScatteringSystemAtOmega, _ScatteringSystemAtOmegaK
from .cycommon import EwaldPart, BesselType
from libc.stdlib cimport malloc, free, calloc
import numpy as np
from cython.parallel cimport prange, parallel
from cython cimport boundscheck, wraparound

cdef extern from "ewald.h":
    void ewald3_2_sigma_long_Delta(cdouble *target, double *err,  int maxn, cdouble x, int xbranch, cdouble z)
    void ewald3_2_sigma_long_Delta_series(cdouble *target, double *err,  int maxn, cdouble x, int xbranch, cdouble z)
    void ewald3_2_sigma_long_Delta_recurrent(cdouble *target, double *err,  int maxn, cdouble x, int xbranch, cdouble z, bint bigimz)
    int complex_gamma_inc_e(double a, cdouble x, int xbranch, qpms_csf_result *result)
    extern int ewald_factor_ipow_l
    extern int ewald_factor_ipow_m

cdef extern from "translations_dbg.h":
    int qpms_trans_calculator_test_sswf(const qpms_trans_calculator *c,
        cdouble *dest, const csph_t kdlj, const  qpms_bessel_t J) nogil
    int qpms_trans_calculator_test_sswf_lc3p(const qpms_trans_calculator *c,
        cdouble *dest, const cdouble k, const cart3_t r, const qpms_bessel_t J) nogil
    int qpms_trans_calculator_test_sswf_e32(const qpms_trans_calculator *c,
        cdouble * const sigmas_total, double * serr_total,
        const double eta, const cdouble k,
        const cart2_t b1, const cart2_t b2,
        const cart2_t beta,
        const cart3_t particle_shift,
        double maxR, double maxK,
        const qpms_ewald_part parts) nogil

cdef extern from "scatsystem_dbg.h":
    int qpms_scatsyswk_test_sswf_basis_e(cdouble *target, const qpms_scatsys_at_omega_k_t *sswk,
            qpms_bessel_t typ, cart3_t evalpoint, qpms_ewald_part parts) nogil

def e32_Delta(int maxn, cdouble x, cdouble z, int xbranch = 0, get_err=True, method='auto'):
    cdef np.ndarray[double, ndim=1] err_np
    cdef double[::1] err_view
    cdef np.ndarray[np.complex_t, ndim=1] target_np = np.empty((maxn+1,), dtype=complex, order='C')
    cdef cdouble[::1] target_view = target_np
    if get_err:
        err_np = np.empty((maxn+1,), order='C')
        err_view = err_np
    if method == 'recurrent':
        ewald3_2_sigma_long_Delta_recurrent(&target_view[0], &err_view[0] if get_err else NULL, maxn, x, xbranch, z, False)
    if method == 'recurrent_bigimz':
        ewald3_2_sigma_long_Delta_recurrent(&target_view[0], &err_view[0] if get_err else NULL, maxn, x, xbranch, z, True)
    elif method == 'series':
        ewald3_2_sigma_long_Delta_series(&target_view[0], &err_view[0] if get_err else NULL, maxn, x, xbranch, z)
    else:
        ewald3_2_sigma_long_Delta(&target_view[0], &err_view[0] if get_err else NULL, maxn, x, xbranch, z)
    if get_err:
        return target_np, err_np
    else:
        return target_np

def scatsyswk_sswf_basis(_ScatteringSystemAtOmegaK sswk, evalpos, btyp=QPMS_HANKEL_PLUS, ewaldparts=EwaldPart.FULL):
    """Evaluate Ewald-summed scalar SWFs that are used to compose the tranlation operators.

    Parameters
    ----------
    evalpos: array_like of floats and shape (..., 3)
        Evaluation points in cartesian coordinates.

    Returns
    -------
    ndarray of complex, with the shape `evalpos.shape[:-1] + (particle_count, pnelem)`
    """        
    if(btyp != QPMS_HANKEL_PLUS):
        raise NotImplementedError("Only first kind Bessel function-based fields are supported")
    cdef qpms_bessel_t btyp_c = BesselType(btyp)
    cdef qpms_ewald_part ewaldparts_c = EwaldPart(ewaldparts)
    evalpos = np.array(evalpos, dtype=float, copy=False)
    if evalpos.shape[-1] != 3:
        raise ValueError("Last dimension of evalpos has to be 3")

    cdef ScatteringSystem ss = sswk.ssw_pyref.ss_pyref
    cdef qpms_scatsys_t *ss_c = ss.s
    cdef qpms_scatsys_at_omega_k_t *sswk_c = sswk.rawpointer()
    cdef const qpms_trans_calculator *c = ss_c[0].c
    cdef qpms_l_t sswf_lMax = 2*c[0].lMax+1
    cdef qpms_y_t pnelem = qpms_lMax2nelem_sc(sswf_lMax)
    cdef qpms_ss_pi_t particle_count = sswk.ssw_pyref.ss_pyref.particle_count

    cdef np.ndarray[double,ndim=2] evalpos_a = evalpos.reshape(-1,3)
    cdef np.ndarray[complex, ndim=3] results = np.empty((evalpos_a.shape[0], particle_count, pnelem), dtype=complex)
    cdef cdouble *res
    cdef cart3_t pos
    cdef Py_ssize_t i, pi, j
    with nogil, wraparound(False), parallel():
        res = <cdouble *> malloc(particle_count * pnelem * sizeof(cdouble))
        for i in prange(evalpos_a.shape[0]):
            pos.x = evalpos_a[i,0]
            pos.y = evalpos_a[i,1]
            pos.z = evalpos_a[i,2]
            qpms_scatsyswk_test_sswf_basis_e(res, sswk_c, btyp_c, pos, ewaldparts_c)
            for pi in range(particle_count):
                for j in range(pnelem):
                    results[i,pi,j] = res[pi * pnelem + j]
        free(res)
    return results.reshape(evalpos.shape[:-1] + (particle_count, pnelem))

def scatsysw_eval_sswf(_ScatteringSystemAtOmega ssw, evalpos, btyp=QPMS_HANKEL_PLUS):
    """Evaluate sswfs by internal qpms_trans_calculator at given locations.

    Parameters
    ----------
    ssw: _ScatteringSystemAtOmega
        We take the wavenumber and qpms_trans_calculator from this, nothing else.
        (In a way, this is an overkill, but there is no up-to-date standalone cython
        wrapper for qpms_trans_calculator.)
    evalpos: array_like of floats and shape (..., 3)
        Evaluation points in cartesian coordinates.
    btyp: BesselType
        Kind of Bessel functions to be used

    Returns
    -------
    ndarray of complex, with the shape `evalpos.shape[:-1] + (pnelem,)`
    where `pnelem` is TODO
    """
    cdef qpms_bessel_t btyp_c = BesselType(btyp)
    evalpos = np.array(evalpos, dtype=float, copy=False)
    if evalpos.shape[-1] != 3:
        raise ValueError("Last dimension of evalpos has to be 3")

    cdef ScatteringSystem ss = ssw.ss_pyref
    cdef qpms_scatsys_t *ss_c = ss.s
    cdef const qpms_trans_calculator *c = ss_c[0].c
    cdef qpms_l_t sswf_lMax = 2*c[0].lMax+1
    cdef qpms_y_t pnelem = qpms_lMax2nelem_sc(sswf_lMax)

    cdef cdouble k = ssw.wavenumber

    cdef np.ndarray[double,ndim=2] evalpos_a = evalpos.reshape(-1,3)
    cdef np.ndarray[complex,ndim=2] results = np.empty((evalpos_a.shape[0], pnelem), dtype=complex)
    cdef cdouble *res
    cdef cart3_t pos
    cdef csph_t kdlj
    cdef Py_ssize_t i, j
    with nogil, wraparound(False), parallel():
        res = <cdouble *> malloc(pnelem * sizeof(cdouble))
        for i in prange(evalpos_a.shape[0]):
            pos.x = evalpos_a[i,0]
            pos.y = evalpos_a[i,1]
            pos.z = evalpos_a[i,2]
            kdlj = cart2csph(pos)
            kdlj.r *= k
            qpms_trans_calculator_test_sswf(c, res, kdlj, btyp_c)
            for j in range(pnelem):
                results[i,j] = res[j]
        free(res)
    return results.reshape(evalpos.shape[:-1] + (pnelem,))

def gamma_inc(double a, cdouble x, int xbranch=0):
    cdef qpms_csf_result res
    complex_gamma_inc_e(a, x, xbranch, &res)
    return res.val

def ipow_l_factor(n = None):
    global ewald_factor_ipow_l
    if n is not None:
        ewald_factor_ipow_l = n
    return ewald_factor_ipow_l

def ipow_m_factor(n = None):
    global ewald_factor_ipow_m
    if n is not None:
        ewald_factor_ipow_m = n
    return ewald_factor_ipow_m