//! Integrated Gradients for feature attribution.
//!
//! Implements the Integrated Gradients method (Sundararajan et al., 2017) for
//! explaining model predictions by attributing the output to each input feature.
//! The method integrates the model's gradients along a straight-line path from
//! a baseline (typically zeros) to the actual input.
//!
//! # Completeness Axiom
//!
//! The sum of attributions equals the difference between the model output at the
//! input and the model output at the baseline:
//!
//!   sum(attributions) = F(input) - F(baseline)
//!
//! This holds exactly in the continuous limit and approximately with finite steps.
//!
//! # GPU Implementation
//!
//! On CUDA builds, all computation stays GPU-resident:
//! - Interpolated inputs are constructed on GPU via a fused CUDA kernel.
//! - Forward passes operate on `GpuTensor` (GPU-in, GPU-out).
//! - Finite-difference gradient approximation perturbs tensors on GPU.
//! - Only the final attribution scores are read back to CPU.
//!
//! On non-CUDA builds, a CPU fallback with `&[f32]` is provided.

use std::collections::HashMap;

use ml_core::MLError;

#[cfg(feature = "cuda")]
use ml_core::cuda_autograd::GpuTensor;

/// Integrated Gradients explainability method.
///
/// Computes per-feature attribution scores by integrating model gradients
/// along a linear interpolation path from a baseline to the input.
#[derive(Debug)]
pub struct IntegratedGradients {
    /// Number of interpolation steps along the path (minimum 2).
    num_steps: usize,
}

impl IntegratedGradients {
    /// Create a new `IntegratedGradients` instance.
    ///
    /// # Arguments
    /// * `num_steps` - Number of interpolation steps. Clamped to a minimum of 2
    ///   (start and end points are always included).
    #[must_use]
    pub fn new(num_steps: usize) -> Self {
        Self {
            num_steps: num_steps.max(2),
        }
    }

    /// Compute feature attributions on GPU.
    ///
    /// The forward function operates entirely on GPU tensors (`GpuTensor` in,
    /// `GpuTensor` out). Interpolated inputs are built on GPU, gradient
    /// approximation (central finite differences) runs on GPU, and only the
    /// final attribution vector is read back to CPU.
    ///
    /// # Arguments
    /// * `forward_fn` - GPU-resident forward pass: takes a 1-D `GpuTensor` of
    ///   shape `[num_features]` and returns a 1-D `GpuTensor` of shape `[1]`
    ///   (scalar output on GPU).
    /// * `input` - Host-side input feature vector to explain.
    /// * `baseline` - Optional host-side baseline. Defaults to zeros.
    /// * `feature_names` - Names for each input feature. Length must match `input`.
    /// * `stream` - CUDA stream for all GPU operations.
    ///
    /// # Errors
    /// Returns `MLError` if dimensions mismatch or if any GPU operation fails.
    #[cfg(feature = "cuda")]
    #[allow(unsafe_code)] // CUDA kernel launches require unsafe FFI.
    pub fn compute_gpu(
        &self,
        forward_fn: &dyn Fn(&GpuTensor) -> Result<GpuTensor, MLError>,
        input: &[f32],
        baseline: Option<&[f32]>,
        feature_names: &[String],
        stream: &std::sync::Arc<cudarc::driver::CudaStream>,
    ) -> Result<HashMap<String, f64>, MLError> {
        use cudarc::driver::{LaunchConfig, PushKernelArg};

        let num_elements = input.len();

        if feature_names.len() != num_elements {
            return Err(MLError::DimensionMismatch {
                expected: num_elements,
                actual: feature_names.len(),
            });
        }

        let zeros = vec![0.0_f32; num_elements];
        let baseline_data = baseline.unwrap_or(&zeros);

        if baseline_data.len() != num_elements {
            return Err(MLError::DimensionMismatch {
                expected: num_elements,
                actual: baseline_data.len(),
            });
        }

        // diff = input - baseline (computed on host, uploaded once)
        let diff: Vec<f32> = input
            .iter()
            .zip(baseline_data.iter())
            .map(|(a, b)| a - b)
            .collect();

        // Upload baseline and diff to GPU once
        let d_baseline = GpuTensor::from_host(baseline_data, vec![num_elements], stream)?;
        let d_diff = GpuTensor::from_host(&diff, vec![num_elements], stream)?;

        // Compile the interpolation + perturbation kernels
        let context = stream.context();
        let ptx = ml_core::cuda_compile::compile_ptx_for_device(
            IG_CUDA_SRC,
            context,
        ).map_err(|e| MLError::ModelError(format!("IG kernel compile: {e}")))?;
        let module = context.load_module(ptx).map_err(|e| {
            MLError::ModelError(format!("IG module load: {e}"))
        })?;
        let interpolate_fn = module.load_function("interpolate_input").map_err(|e| {
            MLError::ModelError(format!("interpolate_input load: {e}"))
        })?;
        let perturb_fn = module.load_function("perturb_dimension").map_err(|e| {
            MLError::ModelError(format!("perturb_dimension load: {e}"))
        })?;

        let threads = 256_u32;
        #[allow(clippy::integer_division)]
        let blocks = (num_elements as u32).div_ceil(threads);
        let elem_cfg = LaunchConfig {
            grid_dim: (blocks, 1, 1),
            block_dim: (threads, 1, 1),
            shared_mem_bytes: 0,
        };

        let eps = 1e-4_f32;
        let n_i32 = num_elements as i32;

        let divisor = if self.num_steps > 1 {
            (self.num_steps - 1) as f64
        } else {
            1.0
        };

        // Accumulated gradients: one f64 per feature
        let mut accumulated_grads = vec![0.0_f64; num_elements];

        for step in 0..self.num_steps {
            let alpha = (step as f64 / divisor) as f32;

            // 1) Build interpolated input on GPU: baseline + alpha * diff
            let d_interp = GpuTensor::zeros(&[num_elements], stream)?;
            // SAFETY: kernel args match (baseline, diff, output, alpha, n).
            // All buffers are GPU-allocated with `num_elements` size.
            unsafe {
                stream
                    .launch_builder(&interpolate_fn)
                    .arg(d_baseline.data())
                    .arg(d_diff.data())
                    .arg(d_interp.data())
                    .arg(&alpha)
                    .arg(&n_i32)
                    .launch(elem_cfg)
                    .map_err(|e| MLError::ModelError(format!("interpolate launch: {e}")))?;
            }

            // 2) For each feature dimension, compute central finite difference
            //    entirely on GPU: perturb +eps and -eps, run forward, subtract.
            for (feat_idx, acc_grad) in accumulated_grads.iter_mut().enumerate() {
                let feat_i32 = feat_idx as i32;

                // x_plus = interp; x_plus[feat_idx] += eps
                let d_x_plus = GpuTensor::zeros(&[num_elements], stream)?;
                // SAFETY: kernel args match (src, dst, dim, delta, n).
                // All buffers are GPU-allocated with `num_elements` size.
                unsafe {
                    stream
                        .launch_builder(&perturb_fn)
                        .arg(d_interp.data())
                        .arg(d_x_plus.data())
                        .arg(&feat_i32)
                        .arg(&eps)
                        .arg(&n_i32)
                        .launch(elem_cfg)
                        .map_err(|e| {
                            MLError::ModelError(format!("perturb +eps launch: {e}"))
                        })?;
                }

                // x_minus = interp; x_minus[feat_idx] -= eps
                let neg_eps = -eps;
                let d_x_minus = GpuTensor::zeros(&[num_elements], stream)?;
                // SAFETY: same as above — perturb kernel with negative delta.
                unsafe {
                    stream
                        .launch_builder(&perturb_fn)
                        .arg(d_interp.data())
                        .arg(d_x_minus.data())
                        .arg(&feat_i32)
                        .arg(&neg_eps)
                        .arg(&n_i32)
                        .launch(elem_cfg)
                        .map_err(|e| {
                            MLError::ModelError(format!("perturb -eps launch: {e}"))
                        })?;
                }

                // Forward passes on GPU
                let f_plus = forward_fn(&d_x_plus)?;
                let f_minus = forward_fn(&d_x_minus)?;

                // Read back the scalar outputs (1 f32 each -- minimal DtoH)
                let f_plus_host = f_plus.to_host(stream)?;
                let f_minus_host = f_minus.to_host(stream)?;

                let fp = f_plus_host.first().copied().unwrap_or(0.0_f32);
                let fm = f_minus_host.first().copied().unwrap_or(0.0_f32);
                let grad = f64::from(fp - fm) / (2.0 * f64::from(eps));
                *acc_grad += grad;
            }
        }

        // Average the accumulated gradients and multiply by diff
        let n_steps = self.num_steps as f64;
        let mut result = HashMap::with_capacity(feature_names.len());
        for (i, name) in feature_names.iter().enumerate() {
            let avg_grad = accumulated_grads.get(i).copied().unwrap_or(0.0) / n_steps;
            let d = diff.get(i).copied().unwrap_or(0.0_f32);
            let attribution = avg_grad * f64::from(d);
            result.insert(name.clone(), attribution);
        }

        Ok(result)
    }

    /// CPU fallback for feature attribution computation.
    ///
    /// Used on non-CUDA builds or when no GPU is available. The forward function
    /// operates on host `&[f32]` slices.
    pub fn compute(
        &self,
        forward_fn: &dyn Fn(&[f32]) -> Result<f32, MLError>,
        input: &[f32],
        baseline: Option<&[f32]>,
        feature_names: &[String],
    ) -> Result<HashMap<String, f64>, MLError> {
        let num_elements = input.len();

        if feature_names.len() != num_elements {
            return Err(MLError::DimensionMismatch {
                expected: num_elements,
                actual: feature_names.len(),
            });
        }

        let zeros = vec![0.0_f32; num_elements];
        let baseline_data = baseline.unwrap_or(&zeros);

        if baseline_data.len() != num_elements {
            return Err(MLError::DimensionMismatch {
                expected: num_elements,
                actual: baseline_data.len(),
            });
        }

        let diff: Vec<f32> = input
            .iter()
            .zip(baseline_data.iter())
            .map(|(a, b)| a - b)
            .collect();

        let divisor = if self.num_steps > 1 {
            (self.num_steps - 1) as f64
        } else {
            1.0
        };

        let eps = 1e-4_f32;
        let mut accumulated_grads = vec![0.0_f64; num_elements];

        for step in 0..self.num_steps {
            let alpha = step as f64 / divisor;
            let alpha_f32 = alpha as f32;

            let interpolated: Vec<f32> = baseline_data
                .iter()
                .zip(diff.iter())
                .map(|(b, d)| b + alpha_f32 * d)
                .collect();

            for (feat_idx, acc_grad) in accumulated_grads.iter_mut().enumerate() {
                let mut x_plus = interpolated.clone();
                let mut x_minus = interpolated.clone();
                if let Some(v) = x_plus.get_mut(feat_idx) {
                    *v += eps;
                }
                if let Some(v) = x_minus.get_mut(feat_idx) {
                    *v -= eps;
                }

                let f_plus = forward_fn(&x_plus)?;
                let f_minus = forward_fn(&x_minus)?;

                let grad = f64::from(f_plus - f_minus) / (2.0 * f64::from(eps));
                *acc_grad += grad;
            }
        }

        let n_steps = self.num_steps as f64;
        let mut result = HashMap::with_capacity(feature_names.len());
        for (i, name) in feature_names.iter().enumerate() {
            let avg_grad = accumulated_grads.get(i).copied().unwrap_or(0.0) / n_steps;
            let d = diff.get(i).copied().unwrap_or(0.0_f32);
            let attribution = avg_grad * f64::from(d);
            result.insert(name.clone(), attribution);
        }

        Ok(result)
    }
}

// ── CUDA kernels for GPU-resident integrated gradients ──────────────────
#[cfg(feature = "cuda")]
const IG_CUDA_SRC: &str = r#"
// Compute interpolated = baseline + alpha * diff (element-wise)
extern "C" __global__
void interpolate_input(
    const float* __restrict__ baseline,
    const float* __restrict__ diff,
    float*       __restrict__ output,
    float alpha,
    int n)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        output[i] = baseline[i] + alpha * diff[i];
    }
}

// Copy src to dst, then perturb dst[dim] += delta
extern "C" __global__
void perturb_dimension(
    const float* __restrict__ src,
    float*       __restrict__ dst,
    int dim,
    float delta,
    int n)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        float val = src[i];
        if (i == dim) {
            val += delta;
        }
        dst[i] = val;
    }
}
"#;

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

    #[test]
    fn test_integrated_gradients_basic() {
        // Simple linear model: f(x) = 2*x[0] - 1*x[1] + 0.5*x[2] + 3*x[3]
        let forward = |x: &[f32]| -> Result<f32, MLError> {
            let x0 = x.first().copied().unwrap_or(0.0);
            let x1 = x.get(1).copied().unwrap_or(0.0);
            let x2 = x.get(2).copied().unwrap_or(0.0);
            let x3 = x.get(3).copied().unwrap_or(0.0);
            Ok(2.0 * x0 - 1.0 * x1 + 0.5 * x2 + 3.0 * x3)
        };

        let input = [1.0_f32, -0.5, 0.3, 2.0];
        let feature_names: Vec<String> = (0..4).map(|i| format!("feature_{}", i)).collect();

        let ig = IntegratedGradients::new(50);
        let attributions = ig
            .compute(&forward, &input, None, &feature_names)
            .expect("IG computation should succeed");

        // Verify we got attributions for all features
        assert_eq!(attributions.len(), 4);
        for name in &feature_names {
            let score = attributions.get(name).expect("feature should have a score");
            assert!(score.is_finite(), "score for {} should be finite: {}", name, score);
        }
    }

    #[test]
    fn test_ig_completeness_axiom() {
        // Linear model: f(x) = 2*x[0] - 1*x[1] + 0.5*x[2] + 3*x[3]
        let forward = |x: &[f32]| -> Result<f32, MLError> {
            let x0 = x.first().copied().unwrap_or(0.0);
            let x1 = x.get(1).copied().unwrap_or(0.0);
            let x2 = x.get(2).copied().unwrap_or(0.0);
            let x3 = x.get(3).copied().unwrap_or(0.0);
            Ok(2.0 * x0 - 1.0 * x1 + 0.5 * x2 + 3.0 * x3)
        };

        let input = [1.0_f32, -0.5, 0.3, 2.0];
        let baseline = [0.0_f32; 4];

        // F(input) - F(baseline)
        let f_input = forward(&input).expect("forward should succeed");
        let f_baseline = forward(&baseline).expect("forward should succeed");
        let expected_diff = f64::from(f_input) - f64::from(f_baseline);

        let feature_names: Vec<String> = (0..4).map(|i| format!("feature_{}", i)).collect();
        // For linear networks, even 10 steps gives exact results.
        let ig = IntegratedGradients::new(50);
        let attributions = ig
            .compute(&forward, &input, Some(&baseline), &feature_names)
            .expect("IG computation should succeed");

        let sum_attributions: f64 = attributions.values().sum();

        // Completeness axiom: for linear functions, IG is exact.
        // Allow 1% tolerance for f32 accumulation rounding.
        let relative_error = if expected_diff.abs() > 1e-10 {
            (sum_attributions - expected_diff).abs() / expected_diff.abs()
        } else {
            (sum_attributions - expected_diff).abs()
        };

        assert!(
            relative_error < 0.01,
            "Completeness axiom violated: sum(attributions)={:.6}, F(x)-F(b)={:.6}, \
             relative_error={:.4}",
            sum_attributions,
            expected_diff,
            relative_error,
        );
    }

    #[test]
    fn test_ig_dimension_mismatch() {
        let forward = |x: &[f32]| -> Result<f32, MLError> {
            Ok(x.iter().sum())
        };

        let input = [1.0_f32, -0.5, 0.3, 2.0];
        // Wrong number of feature names (3 instead of 4)
        let feature_names: Vec<String> = (0..3).map(|i| format!("feature_{}", i)).collect();

        let ig = IntegratedGradients::new(10);
        let result = ig.compute(&forward, &input, None, &feature_names);
        assert!(result.is_err(), "Should fail with dimension mismatch");
    }

    #[test]
    fn test_ig_min_steps_clamped() {
        let ig = IntegratedGradients::new(0);
        assert_eq!(ig.num_steps, 2);

        let ig = IntegratedGradients::new(1);
        assert_eq!(ig.num_steps, 2);

        let ig = IntegratedGradients::new(50);
        assert_eq!(ig.num_steps, 50);
    }
}
