/**
 * CUDA kernels for diffusion equation
 *
 * Implements various schemes for the diffusion/heat equation:
 * ∂T/∂t = α∇²T
 *
 * Includes explicit, implicit, and Crank-Nicolson schemes.
 */

extern "C" __global__ void diffusion_explicit_1d(
    const float* __restrict__ temp,
    float* __restrict__ temp_new,
    float alpha,
    float dt,
    float dx,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;

    if (i >= n) return;

    float r = alpha * dt / (dx * dx);  // Diffusion number

    // Stability condition: r <= 0.5 for explicit scheme
    if (r > 0.5f) {
        // Issue warning - this should be handled at higher level
        return;
    }

    // Boundary conditions (Dirichlet for now)
    if (i == 0 || i == n - 1) {
        temp_new[i] = temp[i];  // Fixed boundary values
        return;
    }

    // Explicit finite difference: T_i^{n+1} = T_i^n + r(T_{i+1}^n - 2T_i^n + T_{i-1}^n)
    temp_new[i] = temp[i] + r * (temp[i + 1] - 2.0f * temp[i] + temp[i - 1]);
}

extern "C" __global__ void diffusion_explicit_2d(
    const float* __restrict__ temp,
    float* __restrict__ temp_new,
    float alpha,
    float dt,
    float dx,
    float dy,
    int nx,
    int ny
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    int j = blockIdx.y * blockDim.y + threadIdx.y;

    if (i >= nx || j >= ny) return;

    int idx = j * nx + i;

    float rx = alpha * dt / (dx * dx);
    float ry = alpha * dt / (dy * dy);

    // Stability condition: rx + ry <= 0.5 for 2D explicit scheme
    if (rx + ry > 0.5f) {
        return;
    }

    // Boundary conditions
    if (i == 0 || i == nx - 1 || j == 0 || j == ny - 1) {
        temp_new[idx] = temp[idx];  // Fixed boundary values
        return;
    }

    int idx_ip1 = j * nx + (i + 1);
    int idx_im1 = j * nx + (i - 1);
    int idx_jp1 = (j + 1) * nx + i;
    int idx_jm1 = (j - 1) * nx + i;

    // 2D explicit finite difference
    temp_new[idx] = temp[idx] +
                   rx * (temp[idx_ip1] - 2.0f * temp[idx] + temp[idx_im1]) +
                   ry * (temp[idx_jp1] - 2.0f * temp[idx] + temp[idx_jm1]);
}

extern "C" __global__ void diffusion_implicit_1d(
    const float* __restrict__ temp,
    float* __restrict__ temp_new,
    float alpha,
    float dt,
    float dx,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;

    if (i >= n) return;

    // For implicit scheme, this kernel performs one iteration of iterative solver
    // (e.g., Jacobi iteration for the linear system)

    float r = alpha * dt / (dx * dx);

    // Boundary conditions
    if (i == 0 || i == n - 1) {
        temp_new[i] = temp[i];
        return;
    }

    // Jacobi iteration for implicit scheme:
    // (1 + 2r)T_i^{n+1} - r(T_{i+1}^{n+1} + T_{i-1}^{n+1}) = T_i^n
    // T_i^{n+1} = (T_i^n + r(T_{i+1}^{n+1} + T_{i-1}^{n+1})) / (1 + 2r)

    temp_new[i] = (temp[i] + r * (temp_new[i + 1] + temp_new[i - 1])) / (1.0f + 2.0f * r);
}

extern "C" __global__ void diffusion_crank_nicolson_1d(
    const float* __restrict__ temp,
    float* __restrict__ temp_new,
    float alpha,
    float dt,
    float dx,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;

    if (i >= n) return;

    float r = alpha * dt / (dx * dx);

    // Boundary conditions
    if (i == 0 || i == n - 1) {
        temp_new[i] = temp[i];
        return;
    }

    // Crank-Nicolson scheme (θ = 0.5):
    // T_i^{n+1} - 0.5*r*(T_{i+1}^{n+1} - 2*T_i^{n+1} + T_{i-1}^{n+1}) =
    // T_i^n + 0.5*r*(T_{i+1}^n - 2*T_i^n + T_{i-1}^n)

    float rhs = temp[i] + 0.5f * r * (temp[i + 1] - 2.0f * temp[i] + temp[i - 1]);

    // Jacobi iteration for Crank-Nicolson
    temp_new[i] = (rhs + 0.5f * r * (temp_new[i + 1] + temp_new[i - 1])) / (1.0f + r);
}

// Helper kernel for tridiagonal solver (Thomas algorithm)
extern "C" __global__ void thomas_forward_elimination(
    float* __restrict__ a,        // Lower diagonal
    float* __restrict__ b,        // Main diagonal
    float* __restrict__ c,        // Upper diagonal
    float* __restrict__ d,        // Right hand side
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x + 1;  // Start from i=1

    if (i >= n) return;

    // Forward elimination step
    float m = a[i] / b[i - 1];
    b[i] = b[i] - m * c[i - 1];
    d[i] = d[i] - m * d[i - 1];
}

extern "C" __global__ void thomas_backward_substitution(
    const float* __restrict__ b,  // Modified main diagonal
    const float* __restrict__ c,  // Upper diagonal
    float* __restrict__ d,        // Modified RHS / solution
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;

    if (i >= n - 1) return;  // n-2, n-3, ..., 0

    int idx = n - 2 - i;  // Reverse order

    if (idx < n - 1) {
        d[idx] = (d[idx] - c[idx] * d[idx + 1]) / b[idx];
    }
}

// Advanced diffusion kernels with anisotropic diffusion
extern "C" __global__ void anisotropic_diffusion_2d(
    const float* __restrict__ temp,
    float* __restrict__ temp_new,
    const float* __restrict__ kappa_x,  // Diffusivity in x-direction
    const float* __restrict__ kappa_y,  // Diffusivity in y-direction
    float dt,
    float dx,
    float dy,
    int nx,
    int ny
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    int j = blockIdx.y * blockDim.y + threadIdx.y;

    if (i >= nx || j >= ny) return;

    int idx = j * nx + i;

    // Boundary conditions
    if (i == 0 || i == nx - 1 || j == 0 || j == ny - 1) {
        temp_new[idx] = temp[idx];
        return;
    }

    int idx_ip1 = j * nx + (i + 1);
    int idx_im1 = j * nx + (i - 1);
    int idx_jp1 = (j + 1) * nx + i;
    int idx_jm1 = (j - 1) * nx + i;

    // Anisotropic diffusion with spatially varying diffusivity
    float kx = kappa_x[idx];
    float ky = kappa_y[idx];

    float rx = kx * dt / (dx * dx);
    float ry = ky * dt / (dy * dy);

    // Check stability condition
    if (rx + ry > 0.5f) {
        temp_new[idx] = temp[idx];  // Skip update if unstable
        return;
    }

    temp_new[idx] = temp[idx] +
                   rx * (temp[idx_ip1] - 2.0f * temp[idx] + temp[idx_im1]) +
                   ry * (temp[idx_jp1] - 2.0f * temp[idx] + temp[idx_jm1]);
}

// Kernel for nonlinear diffusion (e.g., for turbulent heat transfer)
extern "C" __global__ void nonlinear_diffusion_2d(
    const float* __restrict__ temp,
    float* __restrict__ temp_new,
    float base_alpha,
    float dt,
    float dx,
    float dy,
    int nx,
    int ny
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    int j = blockIdx.y * blockDim.y + threadIdx.y;

    if (i >= nx || j >= ny) return;

    int idx = j * nx + i;

    // Boundary conditions
    if (i == 0 || i == nx - 1 || j == 0 || j == ny - 1) {
        temp_new[idx] = temp[idx];
        return;
    }

    int idx_ip1 = j * nx + (i + 1);
    int idx_im1 = j * nx + (i - 1);
    int idx_jp1 = (j + 1) * nx + i;
    int idx_jm1 = (j - 1) * nx + i;

    // Compute local gradient magnitude
    float grad_x = (temp[idx_ip1] - temp[idx_im1]) / (2.0f * dx);
    float grad_y = (temp[idx_jp1] - temp[idx_jm1]) / (2.0f * dy);
    float grad_mag = sqrtf(grad_x * grad_x + grad_y * grad_y);

    // Nonlinear diffusivity (example: edge-preserving diffusion)
    float alpha = base_alpha / (1.0f + grad_mag * grad_mag);

    float rx = alpha * dt / (dx * dx);
    float ry = alpha * dt / (dy * dy);

    // Check stability
    if (rx + ry > 0.5f) {
        temp_new[idx] = temp[idx];
        return;
    }

    temp_new[idx] = temp[idx] +
                   rx * (temp[idx_ip1] - 2.0f * temp[idx] + temp[idx_im1]) +
                   ry * (temp[idx_jp1] - 2.0f * temp[idx] + temp[idx_jm1]);
}