/**
 * PERF-2 P3 (`docs/perf2_campaign.md`): the multigrid V-cycle's maps for a
 * MASKED, VARIABLE-COEFFICIENT five-point operator, batched over K marches.
 * Layout: every per-cell array is [K][n] (march-major, n = nx*ny of the
 * level); the index lists (cells, colours, children) are shared across the
 * batch in this benchmark. blockIdx.y = the march.
 */
extern "C" __global__ void mg_rb_half(
    int n_col, const unsigned int* __restrict__ col, int n,
    const float* __restrict__ ae, const float* __restrict__ aw,
    const float* __restrict__ an, const float* __restrict__ as_,
    const float* __restrict__ ap, const float* __restrict__ b,
    float* __restrict__ x, int nx)
{
    int t = blockIdx.x * blockDim.x + threadIdx.x;
    if (t >= n_col) return;
    int base = blockIdx.y * n;
    int idx = col[t];
    int g = base + idx;
    float s = 0.0f;
    float e = ae[g]; if (e != 0.0f) s += e * x[g + 1];
    float w = aw[g]; if (w != 0.0f) s += w * x[g - 1];
    float nn = an[g]; if (nn != 0.0f) s += nn * x[g + nx];
    float ss = as_[g]; if (ss != 0.0f) s += ss * x[g - nx];
    x[g] = (b[g] + s) / ap[g];
}

extern "C" __global__ void mg_residual(
    int n_cells, const unsigned int* __restrict__ cells, int n,
    const float* __restrict__ ae, const float* __restrict__ aw,
    const float* __restrict__ an, const float* __restrict__ as_,
    const float* __restrict__ ap, const float* __restrict__ b,
    const float* __restrict__ x, float* __restrict__ r, int nx)
{
    int t = blockIdx.x * blockDim.x + threadIdx.x;
    if (t >= n_cells) return;
    int base = blockIdx.y * n;
    int g = base + cells[t];
    float s = 0.0f;
    float e = ae[g]; if (e != 0.0f) s += e * x[g + 1];
    float w = aw[g]; if (w != 0.0f) s += w * x[g - 1];
    float nn = an[g]; if (nn != 0.0f) s += nn * x[g + nx];
    float ss = as_[g]; if (ss != 0.0f) s += ss * x[g - nx];
    r[g] = b[g] - (ap[g] * x[g] - s);
}

/* b_c[c] = sum of r_f over the children of coarse cell c (fixed order). */
extern "C" __global__ void mg_restrict(
    int n_coarse, const unsigned int* __restrict__ coarse_cells,
    const unsigned int* __restrict__ children_ptr, const unsigned int* __restrict__ children_idx,
    int n_f, int n_c, const float* __restrict__ r_f, float* __restrict__ b_c)
{
    int t = blockIdx.x * blockDim.x + threadIdx.x;
    if (t >= n_coarse) return;
    int k = blockIdx.y;
    const float* rf = r_f + (size_t)k * n_f;
    float s = 0.0f;
    for (unsigned int p = children_ptr[t]; p < children_ptr[t + 1]; ++p) s += rf[children_idx[p]];
    b_c[(size_t)k * n_c + coarse_cells[t]] = s;
}

/* x_f += 2 x_c[coarse_of[idx]] */
extern "C" __global__ void mg_prolong(
    int n_cells, const unsigned int* __restrict__ cells, const unsigned int* __restrict__ coarse_of,
    int n_f, int n_c, float* __restrict__ x_f, const float* __restrict__ x_c)
{
    int t = blockIdx.x * blockDim.x + threadIdx.x;
    if (t >= n_cells) return;
    int k = blockIdx.y;
    int idx = cells[t];
    x_f[(size_t)k * n_f + idx] += 2.0f * x_c[(size_t)k * n_c + coarse_of[idx]];
}

extern "C" __global__ void mg_zero(int n_cells, const unsigned int* __restrict__ cells, int n, float* __restrict__ x)
{
    int t = blockIdx.x * blockDim.x + threadIdx.x;
    if (t >= n_cells) return;
    x[(size_t)blockIdx.y * n + cells[t]] = 0.0f;
}

/* The coarsest level: one thread per march, `sweeps` symmetric RED-BLACK
 * sweeps (red, black, black, red — the CPU's ordering, so a nearly singular
 * coarsest system settles on the same constant) from zero. */
extern "C" __global__ void mg_coarsest(
    int K, int n_cells, const unsigned int* __restrict__ cells,
    int n_red, const unsigned int* __restrict__ red,
    int n_black, const unsigned int* __restrict__ black, int n,
    const float* __restrict__ ae, const float* __restrict__ aw,
    const float* __restrict__ an, const float* __restrict__ as_,
    const float* __restrict__ ap, const float* __restrict__ b,
    float* __restrict__ x, int nx, int sweeps)
{
    int k = blockIdx.x * blockDim.x + threadIdx.x;
    if (k >= K) return;
    int base = k * n;
    for (int t = 0; t < n_cells; ++t) x[base + cells[t]] = 0.0f;
    for (int sw = 0; sw < sweeps; ++sw) {
        for (int half = 0; half < 4; ++half) {
            const unsigned int* list = (half == 0 || half == 3) ? red : black;
            int n_list = (half == 0 || half == 3) ? n_red : n_black;
            for (int t = 0; t < n_list; ++t) {
                int g = base + list[t];
                float s = 0.0f;
                float e = ae[g]; if (e != 0.0f) s += e * x[g + 1];
                float w = aw[g]; if (w != 0.0f) s += w * x[g - 1];
                float nn = an[g]; if (nn != 0.0f) s += nn * x[g + nx];
                float ss = as_[g]; if (ss != 0.0f) s += ss * x[g - nx];
                x[g] = (b[g] + s) / ap[g];
            }
        }
    }
}

/* ---------------------------------------------------------------------------
 * PERF-2 P3-iii (`MarchPlane`): the same maps over K LANES that each own
 * their operator (their own mask, coefficients, index lists and work
 * buffers). A LaneLevel is one lane's level: device pointers plus counts;
 * the kernels take a [K] table of them per level and blockIdx.y = the lane.
 * The per-cell arithmetic and its order are the K = 1 kernels' above, so a
 * lane's correction does not depend on K or on its neighbours. The layout
 * (14 pointers, then 6 ints) is mirrored by `LaneLevel` in device.rs.
 * ------------------------------------------------------------------------- */
struct LaneLevel {
    const float* ae; const float* aw; const float* an; const float* as_; const float* ap;
    float* b; float* x; float* r;
    const unsigned int* cells; const unsigned int* red; const unsigned int* black;
    const unsigned int* coarse_of; const unsigned int* children_ptr; const unsigned int* children_idx;
    int n; int nx; int n_cells; int n_red; int n_black; int pad;
};

extern "C" __global__ void ml_rb_half(const LaneLevel* __restrict__ L, int colour)
{
    const LaneLevel& l = L[blockIdx.y];
    int t = blockIdx.x * blockDim.x + threadIdx.x;
    const unsigned int* list = colour ? l.black : l.red;
    int n_list = colour ? l.n_black : l.n_red;
    if (t >= n_list) return;
    int g = list[t];
    int nx = l.nx;
    const float* x = l.x;
    float s = 0.0f;
    float e = l.ae[g]; if (e != 0.0f) s += e * x[g + 1];
    float w = l.aw[g]; if (w != 0.0f) s += w * x[g - 1];
    float nn = l.an[g]; if (nn != 0.0f) s += nn * x[g + nx];
    float ss = l.as_[g]; if (ss != 0.0f) s += ss * x[g - nx];
    l.x[g] = (l.b[g] + s) / l.ap[g];
}

extern "C" __global__ void ml_residual(const LaneLevel* __restrict__ L)
{
    const LaneLevel& l = L[blockIdx.y];
    int t = blockIdx.x * blockDim.x + threadIdx.x;
    if (t >= l.n_cells) return;
    int g = l.cells[t];
    int nx = l.nx;
    const float* x = l.x;
    float s = 0.0f;
    float e = l.ae[g]; if (e != 0.0f) s += e * x[g + 1];
    float w = l.aw[g]; if (w != 0.0f) s += w * x[g - 1];
    float nn = l.an[g]; if (nn != 0.0f) s += nn * x[g + nx];
    float ss = l.as_[g]; if (ss != 0.0f) s += ss * x[g - nx];
    l.r[g] = l.b[g] - (l.ap[g] * x[g] - s);
}

/* coarse.b[c] = sum of fine.r over the children of coarse cell c (fixed order). */
extern "C" __global__ void ml_restrict(const LaneLevel* __restrict__ F, const LaneLevel* __restrict__ C)
{
    const LaneLevel& f = F[blockIdx.y];
    const LaneLevel& c = C[blockIdx.y];
    int t = blockIdx.x * blockDim.x + threadIdx.x;
    if (t >= c.n_cells) return;
    float s = 0.0f;
    for (unsigned int p = f.children_ptr[t]; p < f.children_ptr[t + 1]; ++p) s += f.r[f.children_idx[p]];
    c.b[c.cells[t]] = s;
}

/* fine.x += 2 coarse.x[coarse_of[idx]] */
extern "C" __global__ void ml_prolong(const LaneLevel* __restrict__ F, const LaneLevel* __restrict__ C)
{
    const LaneLevel& f = F[blockIdx.y];
    const LaneLevel& c = C[blockIdx.y];
    int t = blockIdx.x * blockDim.x + threadIdx.x;
    if (t >= f.n_cells) return;
    int idx = f.cells[t];
    f.x[idx] += 2.0f * c.x[f.coarse_of[idx]];
}

extern "C" __global__ void ml_zero(const LaneLevel* __restrict__ L)
{
    const LaneLevel& l = L[blockIdx.y];
    int t = blockIdx.x * blockDim.x + threadIdx.x;
    if (t >= l.n_cells) return;
    l.x[l.cells[t]] = 0.0f;
}

/* The coarsest level: one thread per lane, the CPU's red, black, black, red
 * ordering from zero (see mg_coarsest). */
extern "C" __global__ void ml_coarsest(const LaneLevel* __restrict__ L, int K, int sweeps)
{
    int k = blockIdx.x * blockDim.x + threadIdx.x;
    if (k >= K) return;
    const LaneLevel& l = L[k];
    int nx = l.nx;
    float* x = l.x;
    for (int t = 0; t < l.n_cells; ++t) x[l.cells[t]] = 0.0f;
    for (int sw = 0; sw < sweeps; ++sw) {
        for (int half = 0; half < 4; ++half) {
            const unsigned int* list = (half == 0 || half == 3) ? l.red : l.black;
            int n_list = (half == 0 || half == 3) ? l.n_red : l.n_black;
            for (int t = 0; t < n_list; ++t) {
                int g = list[t];
                float s = 0.0f;
                float e = l.ae[g]; if (e != 0.0f) s += e * x[g + 1];
                float w = l.aw[g]; if (w != 0.0f) s += w * x[g - 1];
                float nn = l.an[g]; if (nn != 0.0f) s += nn * x[g + nx];
                float ss = l.as_[g]; if (ss != 0.0f) s += ss * x[g - nx];
                x[g] = (l.b[g] + s) / l.ap[g];
            }
        }
    }
}