//! B-spline basis function evaluator for KAN.
//!
//! Implements Cox-de Boor recursion on CPU to compute B-spline basis values.
//! Knots are stored as `Vec<f32>` (no gradient needed for the knot vector).
//! The evaluator produces a `(batch, num_bases)` `GpuTensor` for use in KAN layers.
//!
//! When a GPU lookup grid is pre-computed via [`BSplineBasis::precompute_gpu_grid`],
//! evaluation replaces O(n * `num_bases` * 2^(k-1)) recursive CPU calls with a
//! GPU gather + linear interpolation.

use std::sync::Arc;

use cudarc::driver::CudaStream;
use ml_core::MLError;

use crate::gpu_tensor::{
    gpu_add, gpu_clamp, gpu_floor, gpu_gather_dim0, gpu_mul, gpu_scale, gpu_sub,
    gpu_add_scalar, GpuTensor,
};

/// B-spline basis evaluator.
///
/// Stores a uniform knot vector and evaluates all basis functions for a
/// batch of scalar inputs using the Cox-de Boor recursion.
///
/// When `grid_values` is populated (via [`Self::precompute_gpu_grid`]),
/// [`Self::evaluate`] uses GPU gather + lerp instead of CPU recursion.
#[derive(Debug, Clone)]
pub struct BSplineBasis {
    /// Uniform knot vector (length = `grid_size` + 2 * order).
    knots: Vec<f32>,
    /// Spline order (degree + 1).
    order: usize,
    /// Number of basis functions = `grid_size` + order - 1.
    num_bases: usize,
    // -- Pre-computed GPU lookup table --
    /// Shape: (`grid_resolution`, `num_bases`). `None` until `precompute_gpu_grid` is called.
    grid_values: Option<GpuTensor>,
    /// Domain minimum (first knot).
    grid_min: f32,
    /// Domain maximum (last knot).
    grid_max: f32,
    /// Number of uniformly spaced sample points in the lookup table.
    grid_resolution: usize,
}

impl BSplineBasis {
    /// Create a new B-spline basis with uniform knots over `[x_min, x_max]`.
    pub fn new(grid_size: usize, order: usize, x_min: f64, x_max: f64) -> Result<Self, MLError> {
        if grid_size == 0 || order == 0 {
            return Err(MLError::ConfigError("grid_size and order must be > 0".to_owned()));
        }
        if x_max <= x_min {
            return Err(MLError::ConfigError(format!(
                "x_max ({}) must be > x_min ({})",
                x_max, x_min
            )));
        }

        let num_bases = grid_size + order - 1;
        let n_knots = num_bases + order;
        let mut knots = Vec::with_capacity(n_knots);

        let step = (x_max - x_min) / grid_size as f64;
        let pad = (order as f64 - 1.0) * step;
        let start = x_min - pad;

        for i in 0..n_knots {
            knots.push((start + i as f64 * step) as f32);
        }

        Ok(Self {
            knots,
            order,
            num_bases,
            grid_values: None,
            grid_min: 0.0,
            grid_max: 1.0,
            grid_resolution: 0,
        })
    }

    /// Number of basis functions produced.
    pub const fn num_bases(&self) -> usize {
        self.num_bases
    }

    /// Pre-compute basis values on a uniform grid for GPU evaluation.
    pub fn precompute_gpu_grid(
        &mut self,
        resolution: usize,
        stream: &Arc<CudaStream>,
    ) -> Result<(), MLError> {
        if resolution < 2 {
            return Err(MLError::ConfigError(
                "GPU grid resolution must be >= 2".to_owned(),
            ));
        }

        let grid_min = self.knots.first().copied().unwrap_or(0.0);
        let grid_max = self.knots.last().copied().unwrap_or(1.0);
        let range = grid_max - grid_min;
        if range.abs() < 1e-10 {
            return Err(MLError::ConfigError(
                "Knot range is degenerate (all knots equal)".to_owned(),
            ));
        }

        let step = range / (resolution - 1).max(1) as f32;

        let mut grid = vec![0.0_f32; resolution * self.num_bases];
        for i in 0..resolution {
            let x = grid_min + i as f32 * step;
            for j in 0..self.num_bases {
                if let Some(slot) = grid.get_mut(i * self.num_bases + j) {
                    *slot = self.cox_de_boor(x, j, self.order);
                }
            }
        }

        self.grid_values = Some(GpuTensor::from_vec(
            grid,
            &[resolution, self.num_bases],
            stream,
        )?);
        self.grid_min = grid_min;
        self.grid_max = grid_max;
        self.grid_resolution = resolution;
        Ok(())
    }

    /// Evaluate all basis functions for a batch of scalars.
    ///
    /// # Arguments
    /// * `x` - `GpuTensor` of shape `(n,)` containing scalar inputs
    ///
    /// # Returns
    /// `GpuTensor` of shape `(n, num_bases)` with basis values.
    ///
    /// Requires a pre-computed GPU grid (via [`Self::precompute_gpu_grid`]).
    /// Falls back to CPU Cox-de Boor if no grid is present (test-only path).
    pub fn evaluate(&self, x: &GpuTensor) -> Result<GpuTensor, MLError> {
        // Fast path: GPU lookup table when pre-computed
        if let Some(ref grid) = self.grid_values {
            return self.evaluate_gpu(x, grid);
        }

        // Slow path: CPU Cox-de Boor (test-only, grid not pre-computed)
        self.evaluate_cpu(x)
    }

    /// CPU evaluation via Cox-de Boor recursion (test-only fallback).
    #[cfg(test)]
    fn evaluate_cpu(&self, x: &GpuTensor) -> Result<GpuTensor, MLError> {
        let x_vec = x.to_vec()?; // test-only CPU fallback readback
        let n = x_vec.len();

        let mut result = vec![0.0_f32; n * self.num_bases];
        for (i, &xi) in x_vec.iter().enumerate() {
            for j in 0..self.num_bases {
                let val = self.cox_de_boor(xi, j, self.order);
                if let Some(slot) = result.get_mut(i * self.num_bases + j) {
                    *slot = val;
                }
            }
        }

        GpuTensor::from_vec(result, &[n, self.num_bases], &x.stream)
    }

    /// CPU evaluation via Cox-de Boor recursion (non-test: error).
    #[cfg(not(test))]
    fn evaluate_cpu(&self, _x: &GpuTensor) -> Result<GpuTensor, MLError> {
        Err(MLError::ConfigError(
            "BSplineBasis::evaluate requires precompute_gpu_grid() in production".to_owned(),
        ))
    }

    /// GPU-native evaluation via linear interpolation on pre-computed grid.
    ///
    /// Uses GPU ops for normalization, floor/ceil, gather, and lerp.
    /// No host roundtrip in the production forward pass.
    fn evaluate_gpu(&self, x: &GpuTensor, grid: &GpuTensor) -> Result<GpuTensor, MLError> {
        let n = x.dim(0)?;

        if n == 0 {
            return GpuTensor::zeros(&[0, self.num_bases], &x.stream);
        }

        let range = self.grid_max - self.grid_min;
        if range.abs() < 1e-10 {
            return GpuTensor::zeros(&[n, self.num_bases], &x.stream);
        }

        let res_max = (self.grid_resolution - 1) as f32;
        let inv_range_scaled = res_max / range;

        // Normalize x to [0, grid_resolution - 2] on GPU:
        //   normalized = (x - grid_min) * inv_range_scaled, clamped to [0, res_max - 1]
        let shifted = gpu_add_scalar(x, -self.grid_min)?;
        let scaled = gpu_scale(&shifted, inv_range_scaled)?;
        let normalized = gpu_clamp(&scaled, 0.0, res_max - 1.0)?;

        // Floor indices: (n,)
        let floor_idx = gpu_floor(&normalized)?;

        // Ceil indices: min(floor + 1, res_max)
        let ceil_idx = gpu_clamp(&gpu_add_scalar(&floor_idx, 1.0)?, 0.0, res_max)?;

        // Fractional part: (n,)
        let frac = gpu_sub(&normalized, &floor_idx)?;

        // Broadcast floor/ceil indices to (n, num_bases) for gather
        let nb = self.num_bases;
        let floor_2d = crate::gpu_tensor::gpu_unsqueeze_broadcast(&floor_idx, nb)?;
        let ceil_2d = crate::gpu_tensor::gpu_unsqueeze_broadcast(&ceil_idx, nb)?;

        // Gather rows from grid: (n, num_bases) each
        let floor_vals = gpu_gather_dim0(grid, &floor_2d)?;
        let ceil_vals = gpu_gather_dim0(grid, &ceil_2d)?;

        // Lerp: result = floor_vals + frac * (ceil_vals - floor_vals)
        let diff = gpu_sub(&ceil_vals, &floor_vals)?;
        let frac_2d = crate::gpu_tensor::gpu_unsqueeze_broadcast(&frac, nb)?;
        let interp = gpu_mul(&frac_2d, &diff)?;
        gpu_add(&floor_vals, &interp)
    }

    /// Cox-de Boor recursion for basis function `j` of given `order` at point `x`.
    fn cox_de_boor(&self, x: f32, j: usize, k: usize) -> f32 {
        if k == 1 {
            let t_j = self.knots.get(j).copied().unwrap_or(0.0);
            let t_j1 = self.knots.get(j + 1).copied().unwrap_or(0.0);
            return if x >= t_j && x < t_j1 { 1.0 } else { 0.0 };
        }

        let t_j = self.knots.get(j).copied().unwrap_or(0.0);
        let t_jk = self.knots.get(j + k).copied().unwrap_or(0.0);
        let t_jk1 = self.knots.get(j + k - 1).copied().unwrap_or(0.0);
        let t_j1 = self.knots.get(j + 1).copied().unwrap_or(0.0);

        let denom_left = t_jk1 - t_j;
        let left = if denom_left.abs() > 1e-10 {
            ((x - t_j) / denom_left) * self.cox_de_boor(x, j, k - 1)
        } else {
            0.0
        };

        let denom_right = t_jk - t_j1;
        let right = if denom_right.abs() > 1e-10 {
            ((t_jk - x) / denom_right) * self.cox_de_boor(x, j + 1, k - 1)
        } else {
            0.0
        };

        left + right
    }
}

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

    fn test_stream() -> Arc<CudaStream> {
        let ctx = CudaContext::new(0).expect("CUDA required");
        ctx.new_stream().expect("Failed to create stream")
    }

    #[test]
    fn test_bspline_basis_shape() {
        let stream = test_stream();
        let basis = BSplineBasis::new(5, 4, -1.0, 1.0).unwrap();
        assert_eq!(basis.num_bases(), 8); // 5 + 4 - 1

        let x = GpuTensor::from_vec(vec![0.0_f32, 0.5, -0.5, 0.1], &[4], &stream).unwrap();
        let out = basis.evaluate(&x).unwrap();
        assert_eq!(out.shape, vec![4, 8]);
    }

    #[test]
    fn test_bspline_partition_of_unity() {
        let stream = test_stream();
        let basis = BSplineBasis::new(5, 4, -1.0, 1.0).unwrap();
        let x = GpuTensor::from_vec(vec![0.0_f32, 0.3, -0.3], &[3], &stream).unwrap();
        let out = basis.evaluate(&x).unwrap();
        let out_host = out.to_vec().unwrap(); // test-only readback
        for i in 0..3 {
            let sum: f32 = (0..8)
                .map(|j| out_host.get(i * 8 + j).copied().unwrap_or(0.0))
                .sum();
            assert!(
                (sum - 1.0).abs() < 0.15,
                "Basis sum at index {} = {} (expected ~1.0)",
                i,
                sum
            );
        }
    }

    #[test]
    fn test_bspline_non_negative() {
        let stream = test_stream();
        let basis = BSplineBasis::new(5, 4, -1.0, 1.0).unwrap();
        let x = GpuTensor::from_vec(vec![0.0_f32, 0.5, -0.5, 0.9], &[4], &stream).unwrap();
        let out = basis.evaluate(&x).unwrap();
        let flat = out.to_vec().unwrap(); // test-only readback
        for &v in &flat {
            assert!(v >= -1e-6, "Basis value {} is negative (< -epsilon)", v);
        }
    }

    #[test]
    fn test_bspline_gpu_grid_shape() {
        let stream = test_stream();
        let mut basis = BSplineBasis::new(5, 4, -1.0, 1.0).unwrap();
        basis.precompute_gpu_grid(512, &stream).unwrap();

        let x = GpuTensor::from_vec(vec![0.0_f32, 0.5, -0.5, 0.1], &[4], &stream).unwrap();
        let out = basis.evaluate(&x).unwrap();
        assert_eq!(out.shape, vec![4, 8]);
    }

    #[test]
    fn test_bspline_gpu_vs_cpu_match() {
        let stream = test_stream();
        let cpu_basis = BSplineBasis::new(5, 4, -1.0, 1.0).unwrap();
        let mut gpu_basis = cpu_basis.clone();
        gpu_basis.precompute_gpu_grid(1024, &stream).unwrap();

        let x = GpuTensor::from_vec(
            vec![0.0_f32, 0.3, -0.3, 0.7, -0.7, 0.1, -0.1, 0.5],
            &[8],
            &stream,
        )
        .unwrap();

        let cpu_out = cpu_basis.evaluate(&x).unwrap().to_vec().unwrap(); // test-only readback
        let gpu_out = gpu_basis.evaluate(&x).unwrap().to_vec().unwrap(); // test-only readback

        assert_eq!(cpu_out.len(), gpu_out.len());
        for (i, (c, g)) in cpu_out.iter().zip(gpu_out.iter()).enumerate() {
            assert!(
                (c - g).abs() < 0.02,
                "Mismatch at index {}: cpu={} gpu={} (diff={})",
                i, c, g, (c - g).abs()
            );
        }
    }

    #[test]
    fn test_bspline_gpu_partition_of_unity() {
        let stream = test_stream();
        let mut basis = BSplineBasis::new(5, 4, -1.0, 1.0).unwrap();
        basis.precompute_gpu_grid(1024, &stream).unwrap();

        let x = GpuTensor::from_vec(vec![0.0_f32, 0.3, -0.3], &[3], &stream).unwrap();
        let out = basis.evaluate(&x).unwrap();
        let out_host = out.to_vec().unwrap(); // test-only readback
        for i in 0..3 {
            let sum: f32 = (0..8)
                .map(|j| out_host.get(i * 8 + j).copied().unwrap_or(0.0))
                .sum();
            assert!(
                (sum - 1.0).abs() < 0.15,
                "GPU basis sum at index {} = {} (expected ~1.0)",
                i,
                sum
            );
        }
    }

    #[test]
    fn test_bspline_gpu_empty_input() {
        let stream = test_stream();
        let mut basis = BSplineBasis::new(5, 4, -1.0, 1.0).unwrap();
        basis.precompute_gpu_grid(256, &stream).unwrap();

        let x = GpuTensor::zeros(&[0], &stream).unwrap();
        let out = basis.evaluate(&x).unwrap();
        assert_eq!(out.shape, vec![0, 8]);
    }

    #[test]
    fn test_bspline_gpu_grid_resolution_minimum() {
        let stream = test_stream();
        let mut basis = BSplineBasis::new(5, 4, -1.0, 1.0).unwrap();
        let err = basis.precompute_gpu_grid(1, &stream);
        assert!(err.is_err());
        basis.precompute_gpu_grid(2, &stream).unwrap();
    }

    #[test]
    fn test_bspline_gpu_large_batch() {
        let stream = test_stream();
        let mut basis = BSplineBasis::new(5, 4, -1.0, 1.0).unwrap();
        basis.precompute_gpu_grid(1024, &stream).unwrap();

        let x = GpuTensor::randn(&[512], 1.0, &stream).unwrap();
        let out = basis.evaluate(&x).unwrap();
        assert_eq!(out.shape, vec![512, 8]);
    }
}
