// Test-only arrival-f space-time coupling. No production scheduler changes.
#include <stdint.h>
#include <math.h>
#include "config.h"
#include "core/precision.cuh"
#include "core/descriptors.cuh"
#include "core/layout.cuh"
#include "operators/macro.cuh"
#include "operators/collision_srt.cuh"
#include "streaming/pull_double_buffer.cuh"

#if DIM != 2 || NQ != 9 || COLLISION_MODEL != 0 || STORE_PRECISION != 0 || USE_LES || STREAMING_MODEL != 0
#error "Space-time experiment requires FP32 D2Q9 SRT double buffers without LES"
#endif

__device__ __forceinline__ void collide(float* f, float omega) {
    float rho, ux, uy, feq[NQ];
    compute_rho_u(f, rho, ux, uy);
#if USE_DDF_SHIFTING
    // Same perturbation-preserving SRT branch as production BlockStep.
    const float drho = f[0]+f[1]+f[2]+f[3]+f[4]+f[5]+f[6]+f[7]+f[8];
    const float jx = f[1]-f[2]+f[5]-f[6]+f[7]-f[8];
    const float jy = f[3]-f[4]+f[5]-f[6]-f[7]+f[8];
    compute_feq_from_momentum(drho, jx, jy, ux, uy, feq);
#else
    compute_feq(rho, ux, uy, feq);
#endif
    collide_srt_no_force(f, feq, omega);
}

__device__ __forceinline__ double physical(double value, unsigned int q) {
#if USE_DDF_SHIFTING
    // Match population_storage.py's FP64 weights, not promoted FP32 d_w.
    return value + (q == 0 ? 4.0/9.0 : (q < 5 ? 1.0/9.0 : 1.0/36.0));
#else
    return value;
#endif
}

__device__ __forceinline__ void health(float stored, unsigned int q, unsigned int* counts) {
    const double value = physical((double)stored, q);
    if (!isfinite(value)) atomicAdd(counts + 4, 1u);
    if (value <= 0.0f) atomicAdd(counts + 5, 1u);
}

extern "C" __global__ void CoarseStart(
    const float* coarse, const float* fine, float* out, float omega,
    unsigned int x0, unsigned int y0, unsigned int x1, unsigned int y1,
    unsigned int* counts)
{
    unsigned int x, y; unsigned long k;
    index_from_thread(x, y, k);
    if (x >= NX || y >= NY) return;
    const bool covered = x >= x0 && x < x1 && y >= y0 && y < y1;
    float f[NQ];
    for (unsigned int q = 0; q < NQ; ++q) {
        const unsigned long s = q * 4ul * TOTAL_CELLS + 2ul * y * (2u * NX) + 2u * x;
        f[q] = covered ? 0.25f * (fine[s] + fine[s+1] + fine[s+2u*NX] + fine[s+2u*NX+1])
                       : coarse[index_f(k, q)];
    }
    collide(f, omega);
    for (int q = 0; q < NQ; ++q) if (!covered) health(f[q], q, counts);
    stream_pull_store(k, f, out);
}

extern "C" __global__ void Collide(
    const float* in, float* out, const unsigned char* active, float omega,
    unsigned int* counts)
{
    unsigned int x, y; unsigned long k;
    index_from_thread(x, y, k);
    if (x >= NX || y >= NY) return;
    float f[NQ];
    if (active == nullptr || active[k]) {
        for (int q = 0; q < NQ; ++q) f[q] = in[index_f(k, q)];
        collide(f, omega);
        if (active != nullptr) for (int q = 0; q < NQ; ++q) health(f[q], q, counts);
    } else {
        for (int q = 0; q < NQ; ++q) f[q] = NAN;
    }
    stream_pull_store(k, f, out);
}

extern "C" __global__ void Pull(
    const float* in, float* out, const unsigned char* active, unsigned int* counts)
{
    unsigned int x, y; unsigned long k;
    index_from_thread(x, y, k);
    if (x >= NX || y >= NY) return;
    float f[NQ];
    if (active == nullptr || active[k]) {
        unsigned long j[NQ];
        compute_neighbors(x, y, j);
        stream_pull_load(k, f, in, j);
        if (active != nullptr) for (int q = 0; q < NQ; ++q) health(f[q], q, counts);
    } else {
        for (int q = 0; q < NQ; ++q) f[q] = NAN;
    }
    stream_pull_store(k, f, out);
}

extern "C" __global__ void Interpolate(
    const float* g0, const float* gp, const uint32_t* samples,
    const uint32_t* origins, const int* sizes, const double* weights,
    double* pre, float* h, unsigned int groups, unsigned int frozen,
    unsigned int* counts)
{
    unsigned int group = blockIdx.x * blockDim.x + threadIdx.x;
    if (group >= groups) return;
    const int size = sizes[group];
    const unsigned int q = origins[group] / TOTAL_CELLS;
    for (int child = 0; child < 4; ++child) {
        double value = frozen ? (double)g0[origins[group]] : 0.0;
        if (!frozen) for (int j = 0; j < size; ++j) {
            const float* source = j < size/2 ? g0 : gp;
            value += weights[(4u*group+child)*10u+j] * (double)source[samples[10u*group+j]];
        }
        const float stored = __double2float_rn(value);
        pre[4u*group+child] = value;
        h[4u*group+child] = stored;
        if (!isfinite(physical(value, q))) atomicAdd(counts, 1u);
        if (physical(value, q) <= 0.0) atomicAdd(counts+1, 1u);
        if (!isfinite(physical((double)stored, q))) atomicAdd(counts+2, 1u);
        if (physical((double)stored, q) <= 0.0) atomicAdd(counts+3, 1u);
    }
}

extern "C" __global__ void Inject(
    float* fine, const float* h, const uint32_t* rows, unsigned int count,
    unsigned int omit_first)
{
    unsigned int row = blockIdx.x * blockDim.x + threadIdx.x;
    if (row >= count || (omit_first && row == 0)) return;
    fine[rows[2u*row]] = h[rows[2u*row+1]];
}

extern "C" __global__ void Register(
    const float* predicted, const float* fine, const float* h,
    const uint32_t* targets, const uint32_t* sources, const int* times,
    double* registers, unsigned int count, int substep, int mutation)
{
    unsigned int row = blockIdx.x * blockDim.x + threadIdx.x;
    if (row >= count) return;
    double value = substep == 0 ? -(double)predicted[targets[row]] : registers[row];
    for (int j = 0; j < 4; ++j) {
        const unsigned int slot = 4u*row+j;
        const int when = times[slot];
        if (when != substep && !(substep == 0 && when == -1)) continue;
        if ((mutation == 1 && when == -1) || (mutation == 2 && when == 0)
                || (mutation == 3 && when == 1)) continue;
        value += 0.25 * (double)(when == -1 ? h[sources[slot]] : fine[sources[slot]]);
    }
    registers[row] = value;
}

extern "C" __global__ void ApplyRegisters(
    float* predicted, const double* registers, const uint32_t* targets,
    unsigned int count, unsigned int* counts)
{
    unsigned int row = blockIdx.x * blockDim.x + threadIdx.x;
    if (row >= count) return;
    const unsigned int slot = targets[row];
    predicted[slot] = __double2float_rn((double)predicted[slot] + registers[row]);
    health(predicted[slot], slot / TOTAL_CELLS, counts);
}