"""CUDA kernels (compiled at run time with NVRTC through CuPy) for the
Helmholtz layer potentials on quadrilateral patch surfaces.

Six kernels are handled together, in the order used by fmm3dbie's
helm_rpcomb_imp solver:

    0: S'_k    1: S_kappa    2: S'_kappa    3: D'_k - D'_kappa    4: S_k    5: D_k

with kappa = i|k|, G_k = exp(ikr)/(4 pi r), S' the normal derivative at the
target, D the normal derivative at the source.

Each layer potential at a target x is split by patch. A patch j is *near*
x when |x - c_j| <= rfac r_j (c_j, r_j the patch centroid and radius); this
test is the same device function in every kernel, so the split is
consistent. Far patches are summed with the smooth (oversampled Gauss-
Legendre) rule in dense_apply; near patches use precomputed weights
(near_weights) applied in sparse_apply. Because near patches never enter the
dense sum, the dense sum has no large cancelling terms and can be evaluated
in single precision.

Near weights are computed as Legendre moments I_n = int K P_n |J| du dv,
converted to node weights with the vals-to-coefs matrix U: w_j = sum_n U[n, j] I_n.

  * self patch (target on the patch): the square is split into four
    triangles fanned from the target to the edges. On each triangle we use
    polar-type coordinates (rho along the ray, x the position on the edge)
    and the substitution x = x_c + h_eff sinh(xi) chosen from the first
    fundamental form at the target, which removes the near-singular
    behaviour in the angular direction. Tensor Gauss-Legendre in (rho, xi),
    with the xi range split so that each piece is short relative to the
    singularity and to the wavelength.
  * other near patches: a quadtree on the reference square, refined while the
    target is within rfac_q times the radius of a leaf or the leaf is large
    compared with the wavelength, with a tensor Gauss-Legendre rule on each
    leaf (fmm3dbie's istrat=1 strategy).

The geometry inside the integrals is either the polynomial interpolant of
the node data (gmode 0, as in fmm3dbie) or, for the rounded cube, the exact
superellipsoid parametrization (gmode 1). The exact geometry is much more
accurate at nodes next to patch edges: the double layer kernel is very
sensitive to small inconsistencies between interpolated positions and
normals.
"""

import cupy as cp
import numpy as np

_SRC = r"""
#include <cupy/complex.cuh>
typedef complex<double> cdouble;

#define NORD1_MAX 13
#define NPOLS_MAX 169
#define NT 128
#define NSUB_MAX 32
#define OVER4PI 0.07957747154594767

__device__ inline void legendre(double x, int norder, double* P) {
    P[0] = 1.0;
    if (norder >= 1) P[1] = x;
    for (int n = 1; n < norder; n++)
        P[n + 1] = ((2 * n + 1) * x * P[n] - n * P[n - 1]) / (n + 1);
}

__device__ inline double ipow(double x, int n) {
    double r = 1.0;
    for (int i = 0; i < n; i++) r *= x;
    return r;
}

// xyz, r_u, r_v at (u, v) of one patch.
// gmode 0: from the Legendre coefficients c[comp * npols + n] of the patch.
// gmode 1: superellipsoid face chart; gp = n(3), e1(3), e2(3), s0, s1, t0, t1.
__device__ inline void geometry(int gmode, const double* coefs, int norder, const double* gpar,
                                double ga, int gpe, int patch, double u, double v, double* g) {
    if (gmode == 1) {
        const double* gp = gpar + 13 * (size_t)patch;
        double hs = 0.5 * (gp[10] - gp[9]), ht = 0.5 * (gp[12] - gp[11]);
        double s = gp[9] + (u + 1) * hs, t = gp[11] + (v + 1) * ht;
        double c[3], S = 0.0;
        for (int q = 0; q < 3; q++) { c[q] = gp[q] + s * gp[3 + q] + t * gp[6 + q]; S += ipow(c[q], gpe); }
        double N = pow(S, 1.0 / gpe);
        double dNu = 0.0, dNv = 0.0;
        for (int q = 0; q < 3; q++) {
            double gN = ipow(c[q] / N, gpe - 1);
            dNu += gN * gp[3 + q]; dNv += gN * gp[6 + q];
        }
        for (int q = 0; q < 3; q++) {
            g[q] = ga * c[q] / N;
            g[3 + q] = ga * hs * (gp[3 + q] / N - c[q] * dNu / (N * N));
            g[6 + q] = ga * ht * (gp[6 + q] / N - c[q] * dNv / (N * N));
        }
        return;
    }
    int n1 = norder + 1, npols = n1 * n1;
    const double* c = coefs + (size_t)patch * 9 * npols;
    double Lu[NORD1_MAX], Lv[NORD1_MAX];
    legendre(u, norder, Lu);
    legendre(v, norder, Lv);
    for (int comp = 0; comp < 9; comp++) {
        double s = 0.0;
        const double* cc = c + comp * npols;
        for (int b = 0; b < n1; b++) {
            double sb = 0.0;
            for (int a = 0; a < n1; a++) sb += cc[b * n1 + a] * Lu[a];
            s += sb * Lv[b];
        }
        g[comp] = s;
    }
}

// The near test, shared by every kernel so that the near/far split is consistent.
__device__ inline bool is_near(const double* x, const double* cen, double r2) {
    double dx = x[0] - cen[0], dy = x[1] - cen[1], dz = x[2] - cen[2];
    return dx * dx + dy * dy + dz * dz <= r2;
}

// The six kernels at one (target, source) pair, times a real weight.
template <typename T>
__device__ inline void kernels6(const T* xt, const T* nt, const T* y, const T* ns, T w,
                                complex<T> zk, T zkabs, int selmask, complex<T>* K) {
    T dx = xt[0] - y[0], dy = xt[1] - y[1], dz = xt[2] - y[2];
    T r2 = dx * dx + dy * dy + dz * dz;
    T r = sqrt(r2);
    T rinv = (T)1 / r;
    T r2inv = rinv * rinv;
    T dnt = dx * nt[0] + dy * nt[1] + dz * nt[2];
    T dns = dx * ns[0] + dy * ns[1] + dz * ns[2];
    complex<T> izkr = complex<T>(0, 1) * zk * r;
    complex<T> Ek(0, 0);
    T Ec = 0;
    // selmask: 2 bits per kernel m at bits 2m, 2m+1
    if (selmask & 0xFC3) Ek = exp(izkr) * ((T)(OVER4PI) * rinv * w);   // kernels 0, 3, 4, 5
    if (selmask & 0x0FC) Ec = exp(-zkabs * r) * ((T)(OVER4PI) * rinv * w);  // kernels 1, 2, 3
    complex<T> fk = (T)1 - izkr;
    T fc = (T)1 + zkabs * r;
    K[4] = Ek;
    K[1] = complex<T>(Ec, 0);
    K[0] = -dnt * fk * Ek * r2inv;
    K[2] = complex<T>(-dnt * fc * Ec * r2inv, 0);
    K[5] = dns * fk * Ek * r2inv;
    if (selmask & 0x0C0) {
        T nn = ns[0] * nt[0] + ns[1] * nt[1] + ns[2] * nt[2];
        complex<T> gk = ((T)3 - (T)3 * izkr - zk * zk * r2);
        T gc = (T)3 + (T)3 * zkabs * r + zkabs * zkabs * r2;
        K[3] = (nn * fk * r2inv - dns * dnt * gk * r2inv * r2inv) * Ek
             - complex<T>((nn * fc * r2inv - dns * dnt * gc * r2inv * r2inv) * Ec, 0);
    } else {
        K[3] = complex<T>(0, 0);
    }
}

// ---------------------------------------------------------------------------
// Near pairs: two passes (counts, then fill using row offsets), one thread per target.
extern "C" __global__ void near_pairs(
    int ntarg, const double* targ_x, int npatches, const double* cen, const double* r2,
    int* counts, const int* row_ptr, int* pair_p)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= ntarg) return;
    double x[3] = {targ_x[3 * i], targ_x[3 * i + 1], targ_x[3 * i + 2]};
    int n = 0;
    for (int j = 0; j < npatches; j++) {
        if (is_near(x, cen + 3 * j, r2[j])) {
            if (pair_p != NULL) pair_p[row_ptr[i] + n] = j;
            n++;
        }
    }
    counts[i] = n;
}

// ---------------------------------------------------------------------------
// Quadtree leaves for the (target, patch) pairs that are not self pairs.
// Two passes: counts only (leaves == NULL), then fill using offsets.
extern "C" __global__ void build_leaves(
    int npairs, const int* pair_t, const int* pair_p, const double* targ_x,
    const int* targ_patch, int gmode, const double* coefs, int norder, const double* gpar,
    double ga, int gpe, double rfac_q, int nlmax, double zkabs, double kR_max,
    int* counts, const int* offsets, double* leaves)
{
    int p = blockIdx.x * blockDim.x + threadIdx.x;
    if (p >= npairs) return;
    int it = pair_t[p], ip = pair_p[p];
    if (targ_patch[it] == ip) { counts[p] = 0; return; }
    double tx = targ_x[3 * it], ty = targ_x[3 * it + 1], tz = targ_x[3 * it + 2];
    double su[64], sv[64], sh[64];
    int sd[64];
    int sp = 1;
    su[0] = 0.0; sv[0] = 0.0; sh[0] = 1.0; sd[0] = 0;
    int n = 0;
    while (sp > 0) {
        sp--;
        double cu = su[sp], cv = sv[sp], h = sh[sp];
        int d = sd[sp];
        bool refine = false;
        if (d < nlmax) {
            // centroid and radius from a 3x3 sample grid (corners alone underestimate
            // the extent of strongly curved patches)
            double xc[9][3], cen[3] = {0, 0, 0}, g[9];
            for (int k = 0; k < 9; k++) {
                geometry(gmode, coefs, norder, gpar, ga, gpe, ip, cu + h * (k % 3 - 1), cv + h * (k / 3 - 1), g);
                for (int q = 0; q < 3; q++) { xc[k][q] = g[q]; cen[q] += g[q] / 9.0; }
            }
            double R2 = 0.0;
            for (int k = 0; k < 9; k++) {
                double s = 0.0;
                for (int q = 0; q < 3; q++) s += (xc[k][q] - cen[q]) * (xc[k][q] - cen[q]);
                R2 = fmax(R2, s);
            }
            double d2 = (tx - cen[0]) * (tx - cen[0]) + (ty - cen[1]) * (ty - cen[1]) + (tz - cen[2]) * (tz - cen[2]);
            // refine while the target is close relative to the leaf, or the leaf is
            // large compared with the wavelength (the leaf rule must resolve exp(ikr))
            refine = (d2 < rfac_q * rfac_q * R2) || (zkabs * zkabs * R2 > kR_max * kR_max);
        }
        if (refine) {
            double hh = 0.5 * h;
            const double cs[4][2] = {{-1, -1}, {1, -1}, {1, 1}, {-1, 1}};
            for (int k = 0; k < 4; k++) {
                su[sp] = cu + hh * cs[k][0]; sv[sp] = cv + hh * cs[k][1]; sh[sp] = hh; sd[sp] = d + 1; sp++;
            }
        } else {
            if (leaves != NULL) {
                size_t o = (size_t)(offsets[p] + n) * 3;
                leaves[o] = cu; leaves[o + 1] = cv; leaves[o + 2] = h;
            }
            n++;
        }
    }
    counts[p] = n;
}

// ---------------------------------------------------------------------------
// Near-field weights: one block per (target, patch) pair.
struct SelfItem { int edge; double xi0, xi1; };

extern "C" __global__ void near_weights(
    int npairs, const int* pair_t, const int* pair_p,
    const double* targ_x, const double* targ_n, const int* targ_patch, const double* targ_uv,
    int gmode, const double* coefs, int norder, const double* gpar, double ga, int gpe,
    const int* leaf_ptr, const double* leaves, int nq, const double* gl_x, const double* gl_w,
    int nqr, const double* glr_x, const double* glr_w, int nqx, const double* glx_x, const double* glx_w,
    double sub_len, double osc_len, const double* umat, cdouble zk, double zkabs, cdouble* W)
{
    __shared__ double sKr[6][NT], sKi[6][NT];
    __shared__ double sLu[NORD1_MAX][NT], sLv[NORD1_MAX][NT];
    __shared__ SelfItem sitems[4 * NSUB_MAX];
    __shared__ double sEdge[4][3];  // per edge: h, x_c, h_eff
    __shared__ int snitems;

    int p = blockIdx.x;
    if (p >= npairs) return;
    int t = threadIdx.x;
    int it = pair_t[p], ip = pair_p[p];
    int n1 = norder + 1, npols = n1 * n1;
    double xt[3] = {targ_x[3 * it], targ_x[3 * it + 1], targ_x[3 * it + 2]};
    double nt[3] = {targ_n[3 * it], targ_n[3 * it + 1], targ_n[3 * it + 2]};
    bool self = (targ_patch[it] == ip);
    double u0 = 0, v0 = 0;

    // per-thread moment accumulators for basis indices n = t, t + NT
    const int NACC = (NPOLS_MAX + NT - 1) / NT;
    double Ir[NACC][6], Ii[NACC][6];
    for (int a = 0; a < NACC; a++) for (int m = 0; m < 6; m++) { Ir[a][m] = 0; Ii[a][m] = 0; }

    int nitems, nodes_per_item;
    if (self) {
        u0 = targ_uv[2 * it]; v0 = targ_uv[2 * it + 1];
        if (t == 0) {
            double g[9];
            geometry(gmode, coefs, norder, gpar, ga, gpe, ip, u0, v0, g);
            double E = g[3] * g[3] + g[4] * g[4] + g[5] * g[5];
            double F = g[3] * g[6] + g[4] * g[7] + g[5] * g[8];
            double G = g[6] * g[6] + g[7] * g[7] + g[8] * g[8];
            double det = sqrt(fmax(E * G - F * F, 1e-300));
            int cnt = 0;
            for (int e = 0; e < 4; e++) {
                double h, x1, x2, xc, heff;
                if (e == 0) { h = 1 - u0; } else if (e == 1) { h = -1 - u0; }
                else if (e == 2) { h = 1 - v0; } else { h = -1 - v0; }
                if (e < 2) { x1 = -1 - v0; x2 = 1 - v0; xc = -h * F / G; heff = fabs(h) * det / G; }
                else       { x1 = -1 - u0; x2 = 1 - u0; xc = -h * F / E; heff = fabs(h) * det / E; }
                sEdge[e][0] = h; sEdge[e][1] = xc; sEdge[e][2] = heff;
                if (fabs(h) < 1e-12) continue;
                double xi0 = asinh((x1 - xc) / heff), xi1 = asinh((x2 - xc) / heff);
                // sub-intervals of length <= sub_len in xi (singular structure) whose
                // extent along the edge is <= dxmax (so that exp(ikr) is resolved)
                double arc = sqrt(e < 2 ? G : E);
                double dxmax = zkabs * arc > 0 ? osc_len / (zkabs * arc) : 1e300;
                double xi = xi0;
                while (xi < xi1 && cnt < 4 * NSUB_MAX) {
                    double xh = xc + heff * sinh(xi);
                    double xn = fmin(xi + sub_len, asinh((xh + dxmax - xc) / heff));
                    if (xn > xi1 || cnt == 4 * NSUB_MAX - 1) xn = xi1;
                    sitems[cnt].edge = e; sitems[cnt].xi0 = xi; sitems[cnt].xi1 = xn; cnt++;
                    xi = xn;
                }
            }
            snitems = cnt;
        }
        __syncthreads();
        nitems = snitems;
        nodes_per_item = nqr * nqx;
    } else {
        nitems = leaf_ptr[p + 1] - leaf_ptr[p];
        nodes_per_item = nq * nq;
    }

    int ntot = nitems * nodes_per_item;
    for (int base = 0; base < ntot; base += NT) {
        int g = base + t;
        double Lu[NORD1_MAX], Lv[NORD1_MAX];
        cdouble K[6];
        bool active = g < ntot;
        if (active) {
            int item = g / nodes_per_item, loc = g - item * nodes_per_item;
            double u, v, wq;
            if (self) {
                int e = sitems[item].edge;
                int ir = loc / nqx, ix = loc - ir * nqx;
                double h = sEdge[e][0], xc = sEdge[e][1], heff = sEdge[e][2];
                double rho = glr_x[ir], wr = glr_w[ir];
                double xi0 = sitems[item].xi0, xi1 = sitems[item].xi1;
                double xi = xi0 + 0.5 * (xi1 - xi0) * (glx_x[ix] + 1.0);
                double wxi = 0.5 * (xi1 - xi0) * glx_w[ix];
                double x = xc + heff * sinh(xi);
                double dx = heff * cosh(xi) * wxi;
                if (e < 2) { u = u0 + rho * h; v = v0 + rho * x; }
                else       { u = u0 + rho * x; v = v0 + rho * h; }
                wq = fabs(h) * rho * wr * dx;
            } else {
                int li = leaf_ptr[p] + item;
                double cu = leaves[3 * li], cv = leaves[3 * li + 1], hf = leaves[3 * li + 2];
                int a = loc / nq, b = loc - a * nq;
                u = cu + hf * gl_x[a]; v = cv + hf * gl_x[b];
                wq = hf * hf * gl_w[a] * gl_w[b];
            }
            double gg[9];
            geometry(gmode, coefs, norder, gpar, ga, gpe, ip, u, v, gg);
            double ns[3] = {gg[4] * gg[8] - gg[5] * gg[7], gg[5] * gg[6] - gg[3] * gg[8], gg[3] * gg[7] - gg[4] * gg[6]};
            double jac = sqrt(ns[0] * ns[0] + ns[1] * ns[1] + ns[2] * ns[2]);
            ns[0] /= jac; ns[1] /= jac; ns[2] /= jac;
            kernels6<double>(xt, nt, gg, ns, wq * jac, zk, zkabs, 0xFFF, K);
            legendre(u, norder, Lu);
            legendre(v, norder, Lv);
        }
        for (int m = 0; m < 6; m++) { sKr[m][t] = active ? K[m].real() : 0.0; sKi[m][t] = active ? K[m].imag() : 0.0; }
        for (int a = 0; a < n1; a++) { sLu[a][t] = active ? Lu[a] : 0.0; sLv[a][t] = active ? Lv[a] : 0.0; }
        __syncthreads();
        for (int acc = 0; acc < NACC; acc++) {
            int n = t + acc * NT;
            if (n < npols) {
                int a = n % n1, b = n / n1;
                for (int i = 0; i < NT; i++) {
                    double P = sLu[a][i] * sLv[b][i];
                    for (int m = 0; m < 6; m++) { Ir[acc][m] += P * sKr[m][i]; Ii[acc][m] += P * sKi[m][i]; }
                }
            }
        }
        __syncthreads();
    }

    // moments -> node weights, w_j = sum_n U[n, j] I_n. Reuse sLu/sLv as I storage.
    double* sIr = &sLu[0][0];  // NORD1_MAX * NT = 1664 >= 6 * NPOLS_MAX = 1014
    double* sIi = &sLv[0][0];
    for (int acc = 0; acc < NACC; acc++) {
        int n = t + acc * NT;
        if (n < npols) for (int m = 0; m < 6; m++) { sIr[m * NPOLS_MAX + n] = Ir[acc][m]; sIi[m * NPOLS_MAX + n] = Ii[acc][m]; }
    }
    __syncthreads();
    for (int j = t; j < npols; j += NT) {
        for (int m = 0; m < 6; m++) {
            double wr = 0, wi = 0;
            for (int n = 0; n < npols; n++) {
                double u = umat[n * npols + j];
                wr += u * sIr[m * NPOLS_MAX + n];
                wi += u * sIi[m * NPOLS_MAX + n];
            }
            W[((size_t)p * 6 + m) * npols + j] = cdouble(wr, wi);
        }
    }
}

// ---------------------------------------------------------------------------
// Far-field pass: out[m][i] = sum over patches j not near x_i of
// sum_l K_m(x_i, y_l) w_l q_l over the oversampled nodes l of patch j, for the
// kernels selected in selmask (4 bits per kernel: 0 skip, 1 density qa, 2 density qb).
// Kernel evaluation in precision T; per-patch partial sums are added into double accumulators.
template <typename T>
__global__ void dense_apply(
    int ntarg, const double* targ_x, const T* targ_n, int npatches, int npols_o,
    const double* cen, const double* r2, const T* src_x, const T* src_n, const T* src_w,
    const complex<T>* qa, const complex<T>* qb, complex<T> zk, T zkabs, int selmask, cdouble* out)
{
    __shared__ T tile[NT][11];
    __shared__ double scen[4];
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    double xd[3] = {0, 0, 0};
    T xt[3] = {0, 0, 0}, nt[3] = {0, 0, 0};
    if (i < ntarg) {
        for (int q = 0; q < 3; q++) { xd[q] = targ_x[3 * i + q]; xt[q] = (T)xd[q]; nt[q] = targ_n[3 * i + q]; }
    }
    int sel[6];
    int kmask = 0;
    for (int m = 0; m < 6; m++) { sel[m] = (selmask >> (4 * m)) & 3; if (sel[m]) kmask |= 3 << (2 * m); }
    cdouble acc[6];
    for (int m = 0; m < 6; m++) acc[m] = cdouble(0, 0);
    for (int j = 0; j < npatches; j++) {
        if (threadIdx.x < 3) scen[threadIdx.x] = cen[3 * j + threadIdx.x];
        if (threadIdx.x == 3) scen[3] = r2[j];
        __syncthreads();
        bool far = (i < ntarg) && !is_near(xd, scen, scen[3]);
        for (int base = 0; base < npols_o; base += NT) {
            int l = base + threadIdx.x;
            if (l < npols_o) {
                size_t s = (size_t)j * npols_o + l;
                for (int q = 0; q < 3; q++) { tile[threadIdx.x][q] = src_x[3 * s + q]; tile[threadIdx.x][3 + q] = src_n[3 * s + q]; }
                tile[threadIdx.x][6] = src_w[s];
                tile[threadIdx.x][7] = qa[s].real(); tile[threadIdx.x][8] = qa[s].imag();
                tile[threadIdx.x][9] = qb[s].real(); tile[threadIdx.x][10] = qb[s].imag();
            }
            __syncthreads();
            if (far) {
                int nl = min(NT, npols_o - base);
                complex<T> part[6];
                for (int m = 0; m < 6; m++) part[m] = complex<T>(0, 0);
                for (int s = 0; s < nl; s++) {
                    complex<T> K[6];
                    kernels6<T>(xt, nt, &tile[s][0], &tile[s][3], tile[s][6], zk, zkabs, kmask, K);
                    complex<T> a(tile[s][7], tile[s][8]), b(tile[s][9], tile[s][10]);
                    for (int m = 0; m < 6; m++) {
                        if (sel[m] == 1) part[m] += K[m] * a;
                        else if (sel[m] == 2) part[m] += K[m] * b;
                    }
                }
                for (int m = 0; m < 6; m++) acc[m] += cdouble((double)part[m].real(), (double)part[m].imag());
            }
            __syncthreads();
        }
    }
    if (i < ntarg) for (int m = 0; m < 6; m++) out[(size_t)m * ntarg + i] = acc[m];
}

// ---------------------------------------------------------------------------
// Near contributions: out[m][i] += sum_{near pairs p of i} sum_l W[p][m][l] q[d_m][patch nodes l]
extern "C" __global__ void sparse_apply(
    int ntarg, const int* row_ptr, const int* pair_p, int npols, int nreq, const int* req_m,
    const int* req_d, const cdouble* W, const cdouble* q, int npts, cdouble* out)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= ntarg) return;
    cdouble acc[6];
    for (int r = 0; r < nreq; r++) acc[r] = cdouble(0, 0);
    for (int p = row_ptr[i]; p < row_ptr[i + 1]; p++) {
        int j = pair_p[p];
        for (int r = 0; r < nreq; r++) {
            int m = req_m[r];
            const cdouble* w = W + ((size_t)p * 6 + m) * npols;
            const cdouble* qq = q + (size_t)req_d[r] * npts + (size_t)j * npols;
            cdouble s(0, 0);
            for (int l = 0; l < npols; l++) s += w[l] * qq[l];
            acc[r] += s;
        }
    }
    for (int r = 0; r < nreq; r++) out[(size_t)req_m[r] * ntarg + i] += acc[r];
}
"""

_module = None


def module():
    global _module
    if _module is None:
        _module = cp.RawModule(code=_SRC, options=("-std=c++17",),
                               name_expressions=["dense_apply<double>", "dense_apply<float>"])
        _module.compile()  # compile (or load from CuPy's cache) now rather than on first use
    return _module


def _i32(a):
    return cp.ascontiguousarray(cp.asarray(a, dtype=cp.int32))


def _f64(a):
    return cp.ascontiguousarray(cp.asarray(a, dtype=cp.float64))


class Geometry:
    """Device-side description of the patch geometry used inside the integrals."""

    def __init__(self, surf):
        npols = surf.npols
        self.norder = surf.norder
        self.coefs = cp.asarray(np.ascontiguousarray(
            surf.srccoefs.reshape(9, surf.npatches, npols).transpose(1, 0, 2)))
        if getattr(surf, "exact", None) is not None:
            self.gmode = 1
            self.gpar = _f64(surf.exact["params"])
            self.ga = float(surf.exact["a"])
            self.gpe = int(surf.exact["p"])
        else:
            self.gmode, self.gpar, self.ga, self.gpe = 0, cp.zeros(13), 0.0, 0

    def args(self):
        return (np.int32(self.gmode), self.coefs, np.int32(self.norder), self.gpar,
                np.float64(self.ga), np.int32(self.gpe))


def find_near_pairs(targ_x, cen, r2):
    """Pairs (target, patch) with |x_t - c_j|^2 <= r2_j, grouped by target.

    Returns (pair_t, pair_p, row_ptr) as int32 device arrays.
    """
    ntarg = int(targ_x.shape[0])
    npatches = int(cen.shape[0])
    ker = module().get_function("near_pairs")
    counts = cp.zeros(ntarg, dtype=cp.int32)
    nb = (ntarg + 127) // 128
    ker((nb,), (128,), (np.int32(ntarg), targ_x, np.int32(npatches), cen, r2, counts,
                        cp.zeros(1, cp.int32), np.uintp(0)))
    row_ptr = cp.zeros(ntarg + 1, dtype=cp.int32)
    row_ptr[1:] = cp.cumsum(counts).astype(cp.int32)
    npairs = int(row_ptr[-1])
    pair_p = cp.zeros(max(npairs, 1), dtype=cp.int32)
    ker((nb,), (128,), (np.int32(ntarg), targ_x, np.int32(npatches), cen, r2, counts, row_ptr, pair_p))
    pair_p = pair_p[:npairs]
    pair_t = cp.repeat(cp.arange(ntarg, dtype=cp.int32), counts.tolist()) if npairs else cp.zeros(0, cp.int32)
    return pair_t.astype(cp.int32), pair_p, row_ptr


def build_leaves(pair_t, pair_p, targ_x, targ_patch, geo, zk=0.0, rfac_q=1.5, nlmax=14, kR_max=3.0):
    npairs = int(pair_t.size)
    if npairs == 0:
        return cp.zeros(1, cp.int32), cp.zeros((1, 3), cp.float64)
    ker = module().get_function("build_leaves")
    counts = cp.zeros(npairs, dtype=cp.int32)
    nb = (npairs + 127) // 128
    args = (np.int32(npairs), pair_t, pair_p, targ_x, targ_patch) + geo.args() + (
        np.float64(rfac_q), np.int32(nlmax), np.float64(abs(zk)), np.float64(kR_max))
    ker((nb,), (128,), args + (counts, cp.zeros(1, cp.int32), np.uintp(0)))
    leaf_ptr = cp.zeros(npairs + 1, dtype=cp.int32)
    leaf_ptr[1:] = cp.cumsum(counts).astype(cp.int32)
    nleaves = int(leaf_ptr[-1])
    leaves = cp.zeros((max(nleaves, 1), 3), dtype=cp.float64)
    ker((nb,), (128,), args + (counts, leaf_ptr, leaves))
    return leaf_ptr, leaves


def near_weights(pair_t, pair_p, targ_x, targ_n, targ_patch, targ_uv, geo, leaf_ptr, leaves,
                 umat, zk, nq=8, nqr=10, nqx=8, sub_len=2.0, osc_len=2.0):
    """Near weights W (npairs, 6, npols), complex128 on the device."""
    npairs = int(pair_t.size)
    npols = (geo.norder + 1) ** 2
    W = cp.zeros((max(npairs, 1), 6, npols), dtype=cp.complex128)
    if npairs == 0:
        return W
    gx, gw = np.polynomial.legendre.leggauss(nq)
    rx, rw = np.polynomial.legendre.leggauss(nqr)
    rx, rw = 0.5 * (rx + 1), 0.5 * rw  # on [0, 1]
    xx, xw = np.polynomial.legendre.leggauss(nqx)
    ker = module().get_function("near_weights")
    ker((npairs,), (128,), (
        np.int32(npairs), pair_t, pair_p, targ_x, targ_n, targ_patch, targ_uv) + geo.args() + (
        leaf_ptr, leaves, np.int32(nq), _f64(gx), _f64(gw),
        np.int32(nqr), _f64(rx), _f64(rw), np.int32(nqx), _f64(xx), _f64(xw), np.float64(sub_len),
        np.float64(osc_len), umat, np.complex128(zk), np.float64(abs(zk)), W))
    return W


class FarSources:
    """Oversampled nodes of all patches, patch-major, in the precision of the far field."""

    def __init__(self, x, n, w, npols_o, cen, r2, dtype):
        self.real = np.dtype(dtype)
        self.cplx = np.complex128 if self.real == np.float64 else np.complex64
        self.x = cp.ascontiguousarray(cp.asarray(x).astype(self.real))
        self.n = cp.ascontiguousarray(cp.asarray(n).astype(self.real))
        self.w = cp.ascontiguousarray(cp.asarray(w).astype(self.real))
        self.npols_o = npols_o
        self.cen, self.r2 = cen, r2
        self.npatches = int(cen.shape[0])


def dense_apply(targ_x, targ_n, far, qa, qb, zk, sel):
    """out (6, ntarg) complex128: far-patch sums for the kernels in sel (6 entries in {0,1,2})."""
    ntarg = int(targ_x.shape[0])
    name = "dense_apply<double>" if far.real == np.float64 else "dense_apply<float>"
    ker = module().get_function(name)
    ccast = lambda a: cp.ascontiguousarray(cp.asarray(a).astype(far.cplx, copy=False))
    out = cp.zeros((6, ntarg), dtype=cp.complex128)
    selmask = sum(int(s) << (4 * m) for m, s in enumerate(sel))
    nb = (ntarg + 127) // 128
    ker((nb,), (128,), (
        np.int32(ntarg), targ_x, cp.ascontiguousarray(targ_n.astype(far.real)), np.int32(far.npatches),
        np.int32(far.npols_o), far.cen, far.r2, far.x, far.n, far.w, ccast(qa), ccast(qb),
        far.cplx(zk), far.real.type(abs(zk)), np.int32(selmask), out))
    return out


def sparse_apply(row_ptr, pair_p, npols, W, q, reqs, out):
    """out (6, ntarg) += near contributions; reqs = list of (kernel m, density index d); q (ndens, npts)."""
    ntarg = int(row_ptr.size) - 1
    if len(reqs) == 0 or int(pair_p.size) == 0:
        return out
    ker = module().get_function("sparse_apply")
    req_m = _i32([m for m, _ in reqs])
    req_d = _i32([d for _, d in reqs])
    q = cp.ascontiguousarray(q, dtype=cp.complex128)
    nb = (ntarg + 127) // 128
    ker((nb,), (128,), (np.int32(ntarg), row_ptr, pair_p, np.int32(npols), np.int32(len(reqs)),
                        req_m, req_d, W, q, np.int32(q.shape[1]), out))
    return out