//! 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 the `interpolate_input` kernel.
//! - Forward passes operate on `GpuTensor` (GPU-in, GPU-out).
//! - Finite-difference gradient approximation perturbs tensors on GPU via the
//!   `perturb_dimension` kernel.
//! - Only the scalar outputs of the forward function and the final attribution
//!   vector 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 std::sync::Arc;

#[cfg(feature = "cuda")]
use cudarc::driver::{CudaFunction, CudaStream, LaunchConfig, PushKernelArg};

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

/// Precompiled cubin containing the IG kernels (`interpolate_input`,
/// `perturb_dimension`). Built by `build.rs` from `src/ig_kernels.cu`.
#[cfg(feature = "cuda")]
static IG_KERNELS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/ig_kernels.cubin"));

/// Number of threads per block for the 1-D IG kernels.
#[cfg(feature = "cuda")]
const IG_BLOCK_DIM: u32 = 256;

/// Epsilon used for central finite-difference gradient estimation.
#[cfg(feature = "cuda")]
const IG_FD_EPS: f32 = 1e-4_f32;

/// Same epsilon for the CPU path.
const IG_FD_EPS_CPU: f32 = 1e-4_f32;

/// Build a 1-D launch configuration covering `n` elements with 256-thread blocks.
#[cfg(feature = "cuda")]
fn elem_cfg(n: usize) -> LaunchConfig {
    let threads = IG_BLOCK_DIM;
    let blocks = ((n as u32) + threads - 1) / threads;
    LaunchConfig {
        grid_dim: (blocks.max(1), 1, 1),
        block_dim: (threads, 1, 1),
        shared_mem_bytes: 0,
    }
}

/// Bundle of compiled IG kernel handles, loaded once per `compute_gpu` call.
#[cfg(feature = "cuda")]
struct IgKernels {
    interpolate_fn: CudaFunction,
    perturb_fn: CudaFunction,
}

#[cfg(feature = "cuda")]
impl IgKernels {
    fn load(stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let module = stream
            .context()
            .load_cubin(IG_KERNELS_CUBIN.to_vec())
            .map_err(|e| MLError::ModelError(format!("IG cubin load: {e}")))?;
        let interpolate_fn = module
            .load_function("interpolate_input")
            .map_err(|e| MLError::ModelError(format!("interpolate_input: {e}")))?;
        let perturb_fn = module
            .load_function("perturb_dimension")
            .map_err(|e| MLError::ModelError(format!("perturb_dimension: {e}")))?;
        Ok(Self {
            interpolate_fn,
            perturb_fn,
        })
    }
}

/// 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
    /// scalar forward outputs plus the final attribution vector cross the
    /// device boundary.
    ///
    /// # 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)] // Raw CUDA kernel launches are inherently unsafe.
    pub fn compute_gpu(
        &self,
        forward_fn: &dyn Fn(&GpuTensor) -> Result<GpuTensor, MLError>,
        input: &[f32],
        baseline: Option<&[f32]>,
        feature_names: &[String],
        stream: &Arc<CudaStream>,
    ) -> 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(),
            });
        }

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

        // Load kernels once per compute call.
        let kernels = IgKernels::load(stream)?;

        // 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)?;

        // Allocate reusable GPU scratch buffers (freed when this fn returns).
        let mut d_interpolated = GpuTensor::zeros(&[num_elements], stream)?;
        let mut d_x_plus = GpuTensor::zeros(&[num_elements], stream)?;
        let mut d_x_minus = GpuTensor::zeros(&[num_elements], stream)?;

        let n_i32 = i32::try_from(num_elements).map_err(|_| {
            MLError::ModelError(format!(
                "IG compute_gpu: num_elements={num_elements} exceeds i32 range"
            ))
        })?;
        let cfg = elem_cfg(num_elements);

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

        let mut accumulated_grads = vec![0.0_f64; num_elements];

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

            // interpolated = baseline + alpha * diff
            // SAFETY: Kernel arguments match `interpolate_input`'s signature:
            //   (const float*, const float*, float*, float, int).
            // All buffers are valid CudaSlice<f32> of `num_elements` elements.
            unsafe {
                stream
                    .launch_builder(&kernels.interpolate_fn)
                    .arg(d_baseline.cuda_data())
                    .arg(d_diff.cuda_data())
                    .arg(d_interpolated.data_mut())
                    .arg(&alpha_f32)
                    .arg(&n_i32)
                    .launch(cfg)
                    .map_err(|e| MLError::ModelError(format!("interpolate_input: {e}")))?;
            }

            for feat_idx in 0..num_elements {
                let dim_i32 = i32::try_from(feat_idx).map_err(|_| {
                    MLError::ModelError(format!(
                        "IG compute_gpu: feature index {feat_idx} exceeds i32 range"
                    ))
                })?;
                let pos_eps = IG_FD_EPS;
                let neg_eps = -IG_FD_EPS;

                // x_plus = interpolated, then x_plus[feat_idx] += eps
                // SAFETY: Kernel arguments match `perturb_dimension`'s signature:
                //   (const float*, float*, int, float, int).
                unsafe {
                    stream
                        .launch_builder(&kernels.perturb_fn)
                        .arg(d_interpolated.cuda_data())
                        .arg(d_x_plus.data_mut())
                        .arg(&dim_i32)
                        .arg(&pos_eps)
                        .arg(&n_i32)
                        .launch(cfg)
                        .map_err(|e| {
                            MLError::ModelError(format!("perturb_dimension(+eps): {e}"))
                        })?;
                }

                // x_minus = interpolated, then x_minus[feat_idx] -= eps
                // SAFETY: same signature as above.
                unsafe {
                    stream
                        .launch_builder(&kernels.perturb_fn)
                        .arg(d_interpolated.cuda_data())
                        .arg(d_x_minus.data_mut())
                        .arg(&dim_i32)
                        .arg(&neg_eps)
                        .arg(&n_i32)
                        .launch(cfg)
                        .map_err(|e| {
                            MLError::ModelError(format!("perturb_dimension(-eps): {e}"))
                        })?;
                }

                // GPU-resident forward passes; returns scalar GpuTensor of shape [1].
                let y_plus = forward_fn(&d_x_plus)?;
                let y_minus = forward_fn(&d_x_minus)?;
                let f_plus = y_plus.to_scalar(stream)?;
                let f_minus = y_minus.to_scalar(stream)?;

                let grad = f64::from(f_plus - f_minus) / (2.0 * f64::from(IG_FD_EPS));
                if let Some(acc) = accumulated_grads.get_mut(feat_idx) {
                    *acc += 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)
    }

    /// 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 = IG_FD_EPS_CPU;
        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)
    }
}

#[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);
    }

    // ── GPU smoke test ────────────────────────────────────────────────────
    //
    // Exercises the full `compute_gpu` pipeline on a tiny linear model whose
    // forward pass is expressible with GpuTensor primitives (element-wise
    // multiply + mean reduction gives a scalar output). Verifies:
    //   1. Cubins load and both kernels launch without error.
    //   2. Completeness axiom holds within 1% (same tolerance as CPU).
    //   3. GPU attributions match the CPU path within 5%.
    //
    // `#[ignore]` because it requires a CUDA runtime; run with
    //   `cargo test -p ml-explainability --all-features -- --ignored --nocapture`.
    #[cfg(feature = "cuda")]
    #[test]
    #[ignore = "requires CUDA runtime"]
    fn test_ig_compute_gpu_linear_model() {
        use ml_core::cuda_autograd::GpuTensor;

        let stream = cudarc::driver::CudaContext::new(0)
            .expect("CUDA context")
            .new_stream()
            .expect("fork stream");

        // Linear weights on-device: f(x) = sum_i (w_i * x_i) = numel * mean(w * x).
        // For numel=4 and weights w, mean(w*x) * 4 == w . x.
        // We construct f as a closure that returns a [1]-tensor whose value is
        // numel * mean_all(w * x). mean_all returns an f32 scalar.
        let weights_host = [2.0_f32, -1.0, 0.5, 3.0];
        let num_features = weights_host.len();
        let d_weights = GpuTensor::from_host(
            &weights_host,
            vec![num_features],
            &stream,
        )
        .expect("upload weights");

        let stream_ref = stream.clone();
        let weights_ref = d_weights.clone();
        let num_features_f32 = num_features as f32;

        let gpu_forward = move |x: &GpuTensor| -> Result<GpuTensor, MLError> {
            let wx = weights_ref.mul(x, &stream_ref)?;
            let mean = wx.mean_all(&stream_ref)?;
            let dot = mean * num_features_f32;
            // Return a [1] GpuTensor holding the scalar.
            GpuTensor::from_host(&[dot], vec![1], &stream_ref)
        };

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

        let ig = IntegratedGradients::new(50);
        let gpu_attributions = ig
            .compute_gpu(&gpu_forward, &input, None, &feature_names, &stream)
            .expect("GPU IG should succeed");

        assert_eq!(gpu_attributions.len(), num_features);

        // Completeness: sum == F(input) - F(baseline).
        let f_input: f32 = weights_host
            .iter()
            .zip(input.iter())
            .map(|(w, x)| w * x)
            .sum();
        let expected_diff = f64::from(f_input);
        let sum_attr: f64 = gpu_attributions.values().sum();
        let rel_err = if expected_diff.abs() > 1e-10 {
            (sum_attr - expected_diff).abs() / expected_diff.abs()
        } else {
            (sum_attr - expected_diff).abs()
        };
        assert!(
            rel_err < 0.01,
            "GPU completeness violated: sum={:.6}, F(x)-F(b)={:.6}, rel_err={:.4}",
            sum_attr, expected_diff, rel_err,
        );

        // Cross-check against CPU path.
        let cpu_forward = |x: &[f32]| -> Result<f32, MLError> {
            Ok(weights_host.iter().zip(x.iter()).map(|(w, v)| w * v).sum())
        };
        let cpu_attributions = ig
            .compute(&cpu_forward, &input, None, &feature_names)
            .expect("CPU IG should succeed");

        for name in &feature_names {
            let gpu_v = gpu_attributions.get(name).copied().unwrap_or(0.0);
            let cpu_v = cpu_attributions.get(name).copied().unwrap_or(0.0);
            let denom = cpu_v.abs().max(1e-6);
            let rel = (gpu_v - cpu_v).abs() / denom;
            assert!(
                rel < 0.05,
                "{name}: gpu={gpu_v:.6} vs cpu={cpu_v:.6} rel_err={rel:.4}",
            );
        }
    }
}
