// CelerisLab – boundary/ibm_kernels.cuh
// Immersed Boundary Method (IBM) kernels: Euler↔Lagrangian coupling.
// Migrated from ref/main.cu: Eul2Lag, Lag2Eul, Update_Lagrangian.
//
// Uses Peskin's 4-point delta function for interpolation/spreading.
//
// Data layout (per rigid body):
//   RigidBodyState2D: {x, y, theta, vx, vy, omega}
//   Lagrangian points: motionL[NT_LAG * 4]
//     motionL[4*k + 0] = target vx at point k
//     motionL[4*k + 1] = target vy at point k
//     motionL[4*k + 2] = XL (x-position)
//     motionL[4*k + 3] = YL (y-position)
//   forceL[NT_LAG * 2] = Lagrangian force (fx, fy) per point
// ============================================================================

#ifndef CELERIS_BOUNDARY_IBM_KERNELS_CUH
#define CELERIS_BOUNDARY_IBM_KERNELS_CUH

#include "../core/params.cuh"  // RigidBodyState2D, RigidBodyControl2D

// IBM is currently 2D only. Guard all kernels.
#if DIM == 2

// ---------------------------------------------------------------------------
// Peskin 4-point delta function  (1D component)
// ---------------------------------------------------------------------------
__device__ __forceinline__ float peskin_delta_1d(float r) {
    float ar = fabsf(r);
    if (ar < 1.0f)
        return 0.125f * (3.0f - 2.0f * ar + sqrtf(1.0f + 4.0f * ar - 4.0f * ar * ar));
    else if (ar < 2.0f)
        return 0.125f * (5.0f - 2.0f * ar - sqrtf(-7.0f + 12.0f * ar - 4.0f * ar * ar));
    else
        return 0.0f;
}

// ---------------------------------------------------------------------------
// Peskin 4-point delta function  (2D = product of 1D components)
// ---------------------------------------------------------------------------
__device__ __forceinline__ float peskin_delta_2d(float dx, float dy) {
    return peskin_delta_1d(dx) * peskin_delta_1d(dy);
}

// ---------------------------------------------------------------------------
// update_lagrangian_control:
//   Update Lagrangian point positions using RigidBodyState2D.
//   Supports circular body (points on circumference).
//
//   Thread k ∈ [0, NT_LAG)  — one thread per Lagrangian point.
// ---------------------------------------------------------------------------
__global__ void update_lagrangian_control(
    float* __restrict__       motionL,   // [NT_LAG * 4]
    const RigidBodyState2D*   state,     // single rigid body state
    float                     radius,    // body radius
    int                       NT_LAG)    // number of Lagrangian points
{
    int k = threadIdx.x + blockIdx.x * blockDim.x;
    if (k >= NT_LAG) return;

    float angle = state->theta + (float)k * 2.0f * 3.14159265358979323846f / (float)NT_LAG;

    // Target velocity at this Lagrangian point (rigid-body kinematics)
    float r_sin = radius * sinf(angle);
    float r_cos = radius * cosf(angle);

    motionL[4 * k + 0] = state->vx - state->omega * r_sin;  // target vx
    motionL[4 * k + 1] = state->vy + state->omega * r_cos;  // target vy
    motionL[4 * k + 2] = state->x + r_cos;                  // XL
    motionL[4 * k + 3] = state->y + r_sin;                  // YL
}

// ---------------------------------------------------------------------------
// euler_to_lagrangian (Eul2Lag):
//   Interpolate Eulerian velocity to Lagrangian points.
//   Compute Lagrangian force = penalty * (target_vel - interpolated_vel) * dL.
//
//   Thread k ∈ [0, NT_LAG)
// ---------------------------------------------------------------------------
__global__ void euler_to_lagrangian(
    float* __restrict__       forceL,     // [NT_LAG * 2] output forces
    const float* __restrict__ motionL,    // [NT_LAG * 4]
    const float* __restrict__ u_field,    // [DIM * N] SoA velocity field
    int                       NT_LAG)
{
    int k = threadIdx.x + blockIdx.x * blockDim.x;
    if (k >= NT_LAG) return;

    float XL = motionL[4 * k + 2];
    float YL = motionL[4 * k + 3];
    int ix = (int)XL;
    int iy = (int)YL;

    // Arc-length element dL (central difference)
    float dL;
    if (k > 0 && k < NT_LAG - 1) {
        float dx1 = motionL[4*k+2] - motionL[4*(k-1)+2];
        float dy1 = motionL[4*k+3] - motionL[4*(k-1)+3];
        float dx2 = motionL[4*(k+1)+2] - motionL[4*k+2];
        float dy2 = motionL[4*(k+1)+3] - motionL[4*k+3];
        dL = 0.5f * (sqrtf(dx1*dx1 + dy1*dy1) + sqrtf(dx2*dx2 + dy2*dy2));
    } else if (k == 0) {
        float dx1 = motionL[4*1+2] - motionL[4*0+2];
        float dy1 = motionL[4*1+3] - motionL[4*0+3];
        dL = sqrtf(dx1*dx1 + dy1*dy1);
    } else {
        float dx1 = motionL[4*k+2] - motionL[4*(k-1)+2];
        float dy1 = motionL[4*k+3] - motionL[4*(k-1)+3];
        dL = sqrtf(dx1*dx1 + dy1*dy1);
    }

    // Interpolate velocity from Euler grid using Peskin 4-point delta
    float UL = 0.0f, VL = 0.0f;
    for (int dy = -2; dy <= 2; dy++) {
        for (int dx = -2; dx <= 2; dx++) {
            int ex = ix + dx;
            int ey = iy + dy;
            if (ex < 0 || ex >= NX || ey < 0 || ey >= NY) continue;

            float ww = peskin_delta_2d((float)ex - XL, (float)ey - YL);
            unsigned long en = (unsigned long)ey * NX + (unsigned long)ex;
            UL += u_field[en] * ww;                       // ux
            VL += u_field[TOTAL_CELLS + en] * ww;         // uy
        }
    }

    // Force = (target_vel - interpolated_vel) * dL
    forceL[2 * k + 0] = 2.0f * (motionL[4*k+0] - UL) * dL;   // fx
    forceL[2 * k + 1] = 2.0f * (motionL[4*k+1] - VL) * dL;   // fy
}

// ---------------------------------------------------------------------------
// lagrangian_to_euler (Lag2Eul):
//   Spread Lagrangian forces back to the Eulerian grid.
//
//   Thread (x, y) — one thread per Euler node.
//   Only processes nodes near the body (bounding box check).
// ---------------------------------------------------------------------------
__global__ void lagrangian_to_euler(
    float* __restrict__       forceE,     // [DIM * N] output Euler force field
    const float* __restrict__ forceL,     // [NT_LAG * 2]
    const float* __restrict__ motionL,    // [NT_LAG * 4]
    float  xc, float yc, float L0,       // body center and extent
    int    NT_LAG)
{
    unsigned int x, y;
    unsigned long k;
#if DIM == 2
    index_from_thread(x, y, k);
#endif

    if (x >= (unsigned int)NX || y >= (unsigned int)NY) return;

    // Bounding box check
    if ((float)x < xc - L0 || (float)x > xc + L0 ||
        (float)y < yc - L0 || (float)y > yc + L0) return;

    float forcex = 0.0f, forcey = 0.0f;

    for (int kl = 0; kl < NT_LAG; kl++) {
        float XL = motionL[4 * kl + 2];
        float YL = motionL[4 * kl + 3];
        float ww = peskin_delta_2d((float)x - XL, (float)y - YL);
        forcex += forceL[2 * kl + 0] * ww;
        forcey += forceL[2 * kl + 1] * ww;
    }

    forceE[k]               = forcex;
    forceE[TOTAL_CELLS + k] = forcey;
}

#endif // DIM == 2

#endif // CELERIS_BOUNDARY_IBM_KERNELS_CUH
