// ===================================================================
// Loss kernels with fused gradient computation -- f32-native.
//
// Common header (common_device_functions.cuh) is prepended by build.rs
// providing: float, atomicAdd(), etc.
//
// Each kernel writes per-element gradients AND atomically accumulates
// into a scalar loss buffer.
// ===================================================================

extern "C" __global__
void mse_loss_with_grad(const float* __restrict__ pred,
                        const float* __restrict__ target,
                        float* __restrict__ grad,
                        float* __restrict__ loss_out,
                        int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        float p = pred[i];
        float t = target[i];
        float diff = p - t;
        float sq = diff * diff;
        float inv_n = 1.0f / (float)n;
        // Atomic add to loss scalar
        atomicAdd(loss_out, sq * inv_n);
        // Gradient: 2*(pred - target)/n
        float two = 2.0f;
        grad[i] = two * diff * inv_n;
    }
}

extern "C" __global__
void huber_loss_with_grad(const float* __restrict__ pred,
                          const float* __restrict__ target,
                          float* __restrict__ grad,
                          float* __restrict__ loss_out,
                          float delta_f32,
                          int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        float p = pred[i];
        float t = target[i];
        float diff = p - t;
        float abs_diff = fabsf(diff);
        float inv_n = 1.0f / (float)n;
        float delta = delta_f32;
        float half = 0.5f;
        float zero = 0.0f;

        if (abs_diff < delta) {
            // Quadratic region: 0.5 * diff^2
            atomicAdd(loss_out, half * diff * diff * inv_n);
            grad[i] = diff * inv_n;
        } else {
            // Linear region: delta * (|diff| - 0.5 * delta)
            atomicAdd(loss_out, (delta * abs_diff - half * delta * delta) * inv_n);
            float one = 1.0f;
            float neg_one = zero - one;
            grad[i] = delta * ((diff > zero) ? one : neg_one) * inv_n;
        }
    }
}
