//! Loss functions with gradient output as CUDA kernels.
//!
//! Each loss function computes both the scalar loss value and the gradient
//! `dL/dpred` in a single kernel launch, avoiding a separate backward pass.

use std::sync::Arc;

use cudarc::driver::{CudaFunction, CudaStream, LaunchConfig, PushKernelArg};
use cudarc::nvrtc::Ptx;

use crate::MLError;
use super::gpu_tensor::GpuTensor;

/// Compiled loss kernels.
#[allow(missing_debug_implementations)]
pub struct LossKernels {
    mse_kernel: CudaFunction,
    huber_kernel: CudaFunction,
}

/// Result of a loss computation.
pub struct LossResult {
    /// Scalar loss value (mean over batch).
    pub loss: GpuTensor,
    /// Gradient of loss w.r.t. predictions: shape matches predictions.
    pub grad: GpuTensor,
}

impl std::fmt::Debug for LossResult {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "LossResult {{ loss: {:?}, grad: {:?} }}", self.loss, self.grad)
    }
}

impl LossKernels {
    /// Compile all loss kernels.
    pub fn new(stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let context = stream.context();
        let ptx: Ptx = crate::cuda_compile::compile_ptx_for_device(LOSS_CUDA_SRC, &context)
            .map_err(|e| MLError::ModelError(format!("loss kernel compilation: {e}")))?;
        let module = context.load_module(ptx).map_err(|e| {
            MLError::ModelError(format!("loss module load: {e}"))
        })?;

        Ok(Self {
            mse_kernel: module.load_function("mse_loss_with_grad").map_err(|e| {
                MLError::ModelError(format!("mse_loss_with_grad load: {e}"))
            })?,
            huber_kernel: module.load_function("huber_loss_with_grad").map_err(|e| {
                MLError::ModelError(format!("huber_loss_with_grad load: {e}"))
            })?,
        })
    }

    /// MSE loss: `L = mean((pred - target)^2)`, `dL/dpred = 2*(pred - target)/n`.
    ///
    /// Both `pred` and `target` must be 1-D with the same number of elements.
    /// Returns a scalar loss `[1]` and gradient with same shape as `pred`.
    pub fn mse(
        &self,
        pred: &GpuTensor,
        target: &GpuTensor,
        stream: &Arc<CudaStream>,
    ) -> Result<LossResult, MLError> {
        let n = pred.numel();
        if n != target.numel() {
            return Err(MLError::DimensionMismatch {
                expected: n,
                actual: target.numel(),
            });
        }

        let grad = GpuTensor::zeros(&pred.shape, stream)?;
        let mut loss_buf = GpuTensor::zeros(&[1], stream)?;
        let n_i32 = n as i32;

        // Kernel: per-element computes (pred-target)^2 → atomicAdd to loss, writes grad
        let threads = 256_u32;
        let blocks = ((n as u32) + threads - 1) / threads;
        let cfg = LaunchConfig {
            grid_dim: (blocks, 1, 1),
            block_dim: (threads, 1, 1),
            shared_mem_bytes: 0,
        };

        // Zero loss accumulator
        stream.memset_zeros(&mut loss_buf.data).map_err(|e| {
            MLError::ModelError(format!("mse zero loss: {e}"))
        })?;

        unsafe {
            stream
                .launch_builder(&self.mse_kernel)
                .arg(&pred.data)
                .arg(&target.data)
                .arg(&grad.data)
                .arg(&loss_buf.data)
                .arg(&n_i32)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("mse_loss_with_grad: {e}")))?;
        }

        Ok(LossResult {
            loss: loss_buf,
            grad,
        })
    }

    /// Huber loss with gradient.
    ///
    /// `delta` is the threshold between L2 and L1 behavior.
    /// When `|pred - target| < delta`: L2 (MSE-like).
    /// When `|pred - target| >= delta`: L1 (linear) with continuous gradient.
    pub fn huber(
        &self,
        pred: &GpuTensor,
        target: &GpuTensor,
        delta: f32,
        stream: &Arc<CudaStream>,
    ) -> Result<LossResult, MLError> {
        let n = pred.numel();
        if n != target.numel() {
            return Err(MLError::DimensionMismatch {
                expected: n,
                actual: target.numel(),
            });
        }

        let grad = GpuTensor::zeros(&pred.shape, stream)?;
        let mut loss_buf = GpuTensor::zeros(&[1], stream)?;
        let n_i32 = n as i32;

        let threads = 256_u32;
        let blocks = ((n as u32) + threads - 1) / threads;
        let cfg = LaunchConfig {
            grid_dim: (blocks, 1, 1),
            block_dim: (threads, 1, 1),
            shared_mem_bytes: 0,
        };

        stream.memset_zeros(&mut loss_buf.data).map_err(|e| {
            MLError::ModelError(format!("huber zero loss: {e}"))
        })?;

        unsafe {
            stream
                .launch_builder(&self.huber_kernel)
                .arg(&pred.data)
                .arg(&target.data)
                .arg(&grad.data)
                .arg(&loss_buf.data)
                .arg(&delta)
                .arg(&n_i32)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("huber_loss_with_grad: {e}")))?;
        }

        Ok(LossResult {
            loss: loss_buf,
            grad,
        })
    }
}

// ── CUDA source ──────────────────────────────────────────────────────────

const LOSS_CUDA_SRC: &str = r#"
// ═══════════════════════════════════════════════════════════════════════
// Loss kernels with fused gradient computation.
// Each kernel writes per-element gradients AND atomicAdd into a scalar loss.
// ═══════════════════════════════════════════════════════════════════════

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 diff = pred[i] - target[i];
        float sq = diff * diff;
        // Mean loss: each element contributes sq / n
        atomicAdd(loss_out, sq / (float)n);
        // Gradient: d/dpred (pred - target)^2 / n = 2*(pred - target)/n
        grad[i] = 2.0f * diff / (float)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,
                          int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        float diff = pred[i] - target[i];
        float abs_diff = fabsf(diff);
        float inv_n = 1.0f / (float)n;

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

#[cfg(test)]
mod tests {
    use super::*;

    fn make_stream() -> Arc<CudaStream> {
        cudarc::driver::CudaContext::new(0)
                .expect("CUDA required")
                .new_stream()
                .expect("fork stream")
    }

    #[test]
    fn test_mse_zero_loss() {
        let stream = make_stream();
        let kernels = LossKernels::new(&stream).unwrap();

        let pred = GpuTensor::from_host(&[1.0, 2.0, 3.0], vec![3], &stream).unwrap();
        let target = GpuTensor::from_host(&[1.0, 2.0, 3.0], vec![3], &stream).unwrap();

        let result = kernels.mse(&pred, &target, &stream).unwrap();
        let loss_host = result.loss.to_host(&stream).unwrap();
        assert!(loss_host[0].abs() < 1e-6, "MSE should be 0 for identical tensors, got {}", loss_host[0]);

        let grad_host = result.grad.to_host(&stream).unwrap();
        for g in &grad_host {
            assert!(g.abs() < 1e-6, "Gradient should be 0, got {g}");
        }
    }

    #[test]
    fn test_mse_known_value() {
        let stream = make_stream();
        let kernels = LossKernels::new(&stream).unwrap();

        // pred = [1, 2], target = [0, 0]
        // MSE = mean((1-0)^2 + (2-0)^2) = (1 + 4) / 2 = 2.5
        let pred = GpuTensor::from_host(&[1.0, 2.0], vec![2], &stream).unwrap();
        let target = GpuTensor::from_host(&[0.0, 0.0], vec![2], &stream).unwrap();

        let result = kernels.mse(&pred, &target, &stream).unwrap();
        let loss_host = result.loss.to_host(&stream).unwrap();
        assert!(
            (loss_host[0] - 2.5).abs() < 1e-5,
            "MSE should be 2.5, got {}",
            loss_host[0]
        );

        // Gradients: 2*(pred-target)/n = [2*1/2, 2*2/2] = [1, 2]
        let grad_host = result.grad.to_host(&stream).unwrap();
        assert!((grad_host[0] - 1.0).abs() < 1e-5);
        assert!((grad_host[1] - 2.0).abs() < 1e-5);
    }

    #[test]
    fn test_huber_small_diff() {
        let stream = make_stream();
        let kernels = LossKernels::new(&stream).unwrap();

        // Small difference (within delta) should behave like MSE
        let pred = GpuTensor::from_host(&[0.5], vec![1], &stream).unwrap();
        let target = GpuTensor::from_host(&[0.0], vec![1], &stream).unwrap();

        let result = kernels.huber(&pred, &target, 1.0, &stream).unwrap();
        let loss_host = result.loss.to_host(&stream).unwrap();
        // Huber with delta=1.0, diff=0.5: 0.5 * 0.25 / 1 = 0.125
        assert!(
            (loss_host[0] - 0.125).abs() < 1e-5,
            "Huber loss should be 0.125, got {}",
            loss_host[0]
        );
    }
}
