//! 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 crate::MLError;
use super::gpu_tensor::GpuTensor;

/// Compiled loss kernels.
#[allow(missing_debug_implementations)]
pub struct LossKernels {
    mse_kernel: CudaFunction,
    huber_kernel: CudaFunction,
    bce_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 {
    /// Load loss kernels from precompiled cubin (build.rs).
    pub fn new(stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        static CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/loss_kernels.cubin"));
        let module = super::cubin_loader::load_cubin(CUBIN, stream)?;
        let f = |name: &str| super::cubin_loader::get_fn(&module, name);
        Ok(Self {
            mse_kernel: f("mse_loss_with_grad")?,
            huber_kernel: f("huber_loss_with_grad")?,
            bce_kernel: f("bce_with_logits_loss_with_grad")?,
        })
    }

    /// 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,
        })
    }

    /// Binary cross-entropy with logits.
    ///
    /// Numerically stable form:
    ///   `L = mean( max(z,0) - z*y + log(1 + exp(-|z|)) )`
    /// Gradient:
    ///   `dL/dz = (sigmoid(z) - y) / n`
    ///
    /// `logits` (the raw pre-sigmoid output) and `target` (in {0.0, 1.0}, soft
    /// targets in `[0, 1]` are also valid) must have the same `numel`. Returns
    /// scalar mean loss `[1]` and elementwise gradient with same shape as
    /// `logits`. Unlike sigmoid+MSE, the gradient never vanishes for confidently
    /// wrong predictions — `|dL/dz|` is bounded in `[0, 1/n)` regardless of `|z|`,
    /// which is exactly the property we want for asymmetric class problems
    /// (Phase 1a direction; Phase 3 meta-labeling).
    pub fn bce_with_logits(
        &self,
        logits: &GpuTensor,
        target: &GpuTensor,
        stream: &Arc<CudaStream>,
    ) -> Result<LossResult, MLError> {
        let n = logits.numel();
        if n != target.numel() {
            return Err(MLError::DimensionMismatch {
                expected: n,
                actual: target.numel(),
            });
        }

        let grad = GpuTensor::zeros(&logits.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!("bce zero loss: {e}"))
        })?;

        unsafe {
            stream
                .launch_builder(&self.bce_kernel)
                .arg(&logits.data)
                .arg(&target.data)
                .arg(&grad.data)
                .arg(&loss_buf.data)
                .arg(&n_i32)
                .launch(cfg)
                .map_err(|e| MLError::ModelError(format!("bce_with_logits_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() < 0.1, "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() < 0.1, "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() < 0.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() < 0.5, "grad[0]={}", grad_host[0]);
        assert!((grad_host[1] - 2.0).abs() < 0.5, "grad[1]={}", grad_host[1]);
    }

    #[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() < 0.1,
            "Huber loss should be ~0.125, got {}",
            loss_host[0]
        );
    }

    /// At z = 0 (uninformative logits), BCE-with-logits MUST equal `ln 2 ≈
    /// 0.693` regardless of the label distribution. This is the canonical
    /// sanity check from every reference impl and pins the formula's constant
    /// term — if the kernel's `log1p(exp(-|z|))` branch is implemented
    /// incorrectly (e.g. drops the constant), this test catches it.
    #[test]
    fn test_bce_at_zero_logits_equals_ln2() {
        let stream = make_stream();
        let kernels = LossKernels::new(&stream).unwrap();

        // 4 zero logits with mixed labels — answer is ln(2) regardless.
        let logits = GpuTensor::from_host(&[0.0, 0.0, 0.0, 0.0], vec![4], &stream).unwrap();
        let target = GpuTensor::from_host(&[0.0, 1.0, 0.0, 1.0], vec![4], &stream).unwrap();

        let result = kernels.bce_with_logits(&logits, &target, &stream).unwrap();
        let loss_host = result.loss.to_host(&stream).unwrap();
        let ln2 = std::f32::consts::LN_2;
        assert!(
            (loss_host[0] - ln2).abs() < 1e-4,
            "BCE at z=0 should equal ln(2) = {ln2}, got {}",
            loss_host[0]
        );

        // Gradients at z=0 are (sigmoid(0) - y)/n = (0.5 - y)/4.
        // → [+0.125, -0.125, +0.125, -0.125]
        let grad_host = result.grad.to_host(&stream).unwrap();
        for (i, &g) in grad_host.iter().enumerate() {
            let expected = if i % 2 == 0 { 0.125 } else { -0.125 };
            assert!(
                (g - expected).abs() < 1e-5,
                "BCE grad[{i}] expected {expected}, got {g}"
            );
        }
    }

    /// BCE gradient must be bounded and have the correct sign in all four
    /// (label, confidence) quadrants. This is the property that makes BCE
    /// preferable to MSE-on-±1 for binary classification — confidently-wrong
    /// predictions get the strongest pull-back signal (|grad| → 1/n), and
    /// confidently-correct ones get near-zero gradient (no wasted updates).
    #[test]
    fn test_bce_gradient_invariants() {
        let stream = make_stream();
        let kernels = LossKernels::new(&stream).unwrap();

        // Layout:
        //   i=0: z=+5,  y=1 (confident correct  → grad ≈ 0, slight neg)
        //   i=1: z=+5,  y=0 (confident wrong    → grad ≈ +1/4, strong pos)
        //   i=2: z=-5,  y=0 (confident correct  → grad ≈ 0, slight pos)
        //   i=3: z=-5,  y=1 (confident wrong    → grad ≈ -1/4, strong neg)
        let logits = GpuTensor::from_host(&[5.0, 5.0, -5.0, -5.0], vec![4], &stream).unwrap();
        let target = GpuTensor::from_host(&[1.0, 0.0, 0.0, 1.0], vec![4], &stream).unwrap();

        let result = kernels.bce_with_logits(&logits, &target, &stream).unwrap();
        let grad_host = result.grad.to_host(&stream).unwrap();

        // |grad_i| ≤ 1/n = 0.25 strictly. This is the bounded-gradient property.
        for (i, &g) in grad_host.iter().enumerate() {
            assert!(g.abs() <= 0.25 + 1e-6, "grad[{i}] = {g} violates |grad| ≤ 1/n = 0.25");
        }

        // Confident correct → near zero.
        assert!(grad_host[0].abs() < 0.01, "confident-correct grad[0] should ≈ 0, got {}", grad_host[0]);
        assert!(grad_host[2].abs() < 0.01, "confident-correct grad[2] should ≈ 0, got {}", grad_host[2]);

        // Confident wrong → close to ±0.25 with correct sign.
        assert!(grad_host[1] > 0.24, "confident-wrong-positive grad[1] should ≈ +0.25, got {}", grad_host[1]);
        assert!(grad_host[3] < -0.24, "confident-wrong-negative grad[3] should ≈ -0.25, got {}", grad_host[3]);
    }
}
