ad_trait 0.3.1

use crate::forward_ad::adfn::adfn;
use crate::forward_ad::ForwardADTrait;
#[cfg(feature = "std")]
use crate::reverse_ad::adr::{adr, GlobalComputationGraph};
use crate::AD;
use alloc::rc::Rc;
use alloc::sync::Arc;
use alloc::vec::Vec;
use alloc::{format, vec};
use core::marker::PhantomData;
use nalgebra::{DMatrix, DVector};
#[cfg(feature = "std")]
use rand::distributions::Distribution;
#[cfg(feature = "std")]
use rand::distributions::Uniform;
#[cfg(feature = "std")]
use rand::thread_rng;
#[cfg(feature = "std")]
use rand::Rng;
#[cfg(feature = "std")]
use std::sync::{Mutex, RwLock};

#[cfg(feature = "nightly")]
use crate::simd::f64xn::f64xn;

/// A trait for types that can be "reparameterized" with a different `AD` type.
///
/// This is a critical feature for automatic differentiation, as it allows a function
/// originally defined for `f64` to be converted into a version that uses `adr` or `adfn`
/// for derivative tracking.
pub trait Reparameterize {
    /// The type of the function after reparameterization with type `T2`.
    type SelfType<T2: AD>: DifferentiableFunctionTrait<T2>;
}

impl<R: Reparameterize> Reparameterize for Rc<R> {
    type SelfType<T2: AD> = R::SelfType<T2>;
}
impl<R: Reparameterize> Reparameterize for Arc<R> {
    type SelfType<T2: AD> = R::SelfType<T2>;
}
#[cfg(feature = "std")]
impl<R: Reparameterize> Reparameterize for Mutex<R> {
    type SelfType<T2: AD> = R::SelfType<T2>;
}
#[cfg(feature = "std")]
impl<R: Reparameterize> Reparameterize for RwLock<R> {
    type SelfType<T2: AD> = R::SelfType<T2>;
}

/*
pub trait DifferentiableFunctionClass {
    type FunctionType<T: AD> : DifferentiableFunctionTrait<T>;
}
impl DifferentiableFunctionClass for () {
    type FunctionType<T: AD> = ();
}
*/

/// Defines the interface for a function that can be differentiated.
///
/// Implementors must provide the `call` method to evaluate the function for a given
/// `AD` type `T`, and specify the number of inputs and outputs.
pub trait DifferentiableFunctionTrait<T: AD> {
    // type FunctionClass: DifferentiableFunctionClass;
    /// A human-readable name for the function.
    const NAME: &'static str;

    /// Evaluates the function.
    ///
    /// # Arguments
    /// * `inputs` - A slice of input values of type `T`.
    /// * `freeze` - If true, certain caches or state updates might be skipped (used in optimizations).
    fn call(&self, inputs: &[T], freeze: bool) -> Vec<T>;

    /// The number of input variables the function expects.
    fn num_inputs(&self) -> usize;

    /// The number of output variables the function returns.
    fn num_outputs(&self) -> usize;
}

pub trait ToOtherADType: Reparameterize {
    fn to_other_ad_type<T2: AD>(&self) -> <Self as Reparameterize>::SelfType<T2>;
}

impl<T: AD, F: DifferentiableFunctionTrait<T>> DifferentiableFunctionTrait<T> for Rc<F> {
    const NAME: &'static str = F::NAME;

    fn call(&self, inputs: &[T], freeze: bool) -> Vec<T> {
        (**self).call(inputs, freeze)
    }

    fn num_inputs(&self) -> usize {
        (**self).num_inputs()
    }

    fn num_outputs(&self) -> usize {
        (**self).num_outputs()
    }
}
impl<T: AD, F: DifferentiableFunctionTrait<T>> DifferentiableFunctionTrait<T> for Arc<F> {
    const NAME: &'static str = F::NAME;

    fn call(&self, inputs: &[T], freeze: bool) -> Vec<T> {
        (**self).call(inputs, freeze)
    }

    fn num_inputs(&self) -> usize {
        (**self).num_inputs()
    }

    fn num_outputs(&self) -> usize {
        (**self).num_outputs()
    }
}
#[cfg(feature = "std")]
impl<T: AD, F: DifferentiableFunctionTrait<T>> DifferentiableFunctionTrait<T> for Mutex<F> {
    const NAME: &'static str = F::NAME;

    fn call(&self, inputs: &[T], freeze: bool) -> Vec<T> {
        self.lock().unwrap().call(inputs, freeze)
    }

    fn num_inputs(&self) -> usize {
        self.lock().unwrap().num_inputs()
    }

    fn num_outputs(&self) -> usize {
        self.lock().unwrap().num_outputs()
    }
}
#[cfg(feature = "std")]
impl<T: AD, F: DifferentiableFunctionTrait<T>> DifferentiableFunctionTrait<T> for RwLock<F> {
    const NAME: &'static str = F::NAME;

    fn call(&self, inputs: &[T], freeze: bool) -> Vec<T> {
        self.read().unwrap().call(inputs, freeze)
    }

    fn num_inputs(&self) -> usize {
        self.read().unwrap().num_inputs()
    }

    fn num_outputs(&self) -> usize {
        self.read().unwrap().num_outputs()
    }
}

impl<T: AD> DifferentiableFunctionTrait<T> for () {
    const NAME: &'static str = "()";

    fn call(&self, _inputs: &[T], _freeze: bool) -> Vec<T> {
        vec![]
    }

    fn num_inputs(&self) -> usize {
        0
    }

    fn num_outputs(&self) -> usize {
        0
    }
}
impl Reparameterize for () {
    type SelfType<T2: AD> = ();
}

/*
pub struct DifferentiableFunctionClassZero;
impl DifferentiableFunctionClass for DifferentiableFunctionClassZero {
    type FunctionType<T: AD> = DifferentiableFunctionZero;
}
*/

#[derive(Clone)]
pub struct DifferentiableFunctionZero {
    num_inputs: usize,
    num_outputs: usize,
}
impl DifferentiableFunctionZero {
    pub fn new(num_inputs: usize, num_outputs: usize) -> Self {
        Self {
            num_inputs,
            num_outputs,
        }
    }
}
impl<T: AD> DifferentiableFunctionTrait<T> for DifferentiableFunctionZero {
    const NAME: &'static str = "DifferentiableFunctionZero";

    fn call(&self, _inputs: &[T], _frozen_freeze: bool) -> Vec<T> {
        vec![T::zero(); self.num_outputs]
    }

    fn num_inputs(&self) -> usize {
        self.num_inputs
    }

    fn num_outputs(&self) -> usize {
        self.num_outputs
    }
}

impl Reparameterize for DifferentiableFunctionZero {
    type SelfType<T2: AD> = DifferentiableFunctionZero;
}

pub trait DerivativeMethodClass {
    type DerivativeMethod: DerivativeMethodTrait;
}
impl DerivativeMethodClass for () {
    type DerivativeMethod = ();
}

/// Defines a method for computing the derivative of a `DifferentiableFunctionTrait`.
pub trait DerivativeMethodTrait: Clone {
    /// The `AD` type used by this method (e.g., `f64`, `adr`, `adfn`).
    type T: AD;

    /// Computes the function's value and its Jacobian matrix at the given input point.
    ///
    /// # Arguments
    /// * `inputs` - The input values as `f64`.
    /// * `function` - The function to differentiate, which must be reparameterizable.
    fn derivative<D: DifferentiableFunctionTrait<Self::T> + ?Sized>(
        &self,
        inputs: &[f64],
        function: &D,
    ) -> (Vec<f64>, DMatrix<f64>);
}
impl DerivativeMethodTrait for () {
    type T = f64;

    fn derivative<D: DifferentiableFunctionTrait<Self::T> + ?Sized>(
        &self,
        _inputs: &[f64],
        _function: &D,
    ) -> (Vec<f64>, DMatrix<f64>) {
        panic!("derivative should not actually be called on ()");
    }
}

////////////////////////////////////////////////////////////////////////////////////////////////////

pub struct DerivativeMethodClassFiniteDifferencing;
impl DerivativeMethodClass for DerivativeMethodClassFiniteDifferencing {
    type DerivativeMethod = FiniteDifferencing;
}

/// Computes derivatives using the Finite Differencing method.
///
/// This method approximates the Jacobian by evaluating the function at slightly
/// perturbed points. It's safe and works on any function, but can be numerically
/// unstable and slow for a large number of inputs.
#[derive(Clone)]
pub struct FiniteDifferencing {}
impl FiniteDifferencing {
    pub fn new() -> Self {
        Self {}
    }
}
impl DerivativeMethodTrait for FiniteDifferencing {
    type T = f64;

    fn derivative<D: DifferentiableFunctionTrait<Self::T> + ?Sized>(
        &self,
        inputs: &[f64],
        function: &D,
    ) -> (Vec<f64>, DMatrix<f64>) {
        let num_inputs = inputs.len();
        let num_outputs = function.num_outputs();
        let mut out_derivative = DMatrix::zeros(num_outputs, num_inputs);

        let h = 0.0000001;

        let x0 = inputs.to_vec();
        // let f0 = D::call(&x0, args);
        let f0 = function.call(&x0, false);

        for col_idx in 0..num_inputs {
            let mut xh = x0.clone();
            xh[col_idx] += h;
            // let fh = D::call(&xh, args);
            let fh = function.call(&xh, true);
            for row_idx in 0..num_outputs {
                out_derivative[(row_idx, col_idx)] = (fh[row_idx] - f0[row_idx]) / h;
            }
        }

        (f0, out_derivative)
    }
}

#[cfg(feature = "std")]
pub struct DerivativeMethodClassReverseAD;
#[cfg(feature = "std")]
impl DerivativeMethodClass for DerivativeMethodClassReverseAD {
    type DerivativeMethod = ReverseAD;
}

#[cfg(feature = "std")]
/// Computes derivatives using Reverse-mode Automatic Differentiation.
///
/// This method uses a global computation graph to track operations and then
/// performs a backward pass to compute gradients efficiently. It is ideal
/// for functions with few outputs and many inputs.
#[derive(Clone)]
pub struct ReverseAD {}
#[cfg(feature = "std")]
impl ReverseAD {
    pub fn new() -> Self {
        Self {}
    }
}
#[cfg(feature = "std")]
impl DerivativeMethodTrait for ReverseAD {
    type T = adr;

    fn derivative<D: DifferentiableFunctionTrait<Self::T> + ?Sized>(
        &self,
        inputs: &[f64],
        function: &D,
    ) -> (Vec<f64>, DMatrix<f64>) {
        let num_inputs = inputs.len();
        let num_outputs = function.num_outputs();
        let mut out_derivative = DMatrix::zeros(num_outputs, num_inputs);

        GlobalComputationGraph::get().reset();

        let mut inputs_ad = vec![];
        for input in inputs.iter() {
            inputs_ad.push(adr::new_variable(*input, false));
        }

        let f = function.call(&inputs_ad, false);
        assert_eq!(f.len(), num_outputs);
        let out_value = f.iter().map(|x| x.value()).collect();

        for row_idx in 0..num_outputs {
            if f[row_idx].is_constant() {
                for col_idx in 0..num_inputs {
                    out_derivative[(row_idx, col_idx)] = 0.0;
                }
            } else {
                let grad_output = f[row_idx].get_backwards_mode_grad();
                for col_idx in 0..num_inputs {
                    let d = grad_output.wrt(&inputs_ad[col_idx]);
                    out_derivative[(row_idx, col_idx)] = d;
                }
            }
        }

        (out_value, out_derivative)
    }
}

pub struct DerivativeMethodClassForwardAD;
impl DerivativeMethodClass for DerivativeMethodClassForwardAD {
    type DerivativeMethod = ForwardAD;
}

/// Computes derivatives using Forward-mode Automatic Differentiation (Single Tangent).
///
/// This method propagates a single tangent value alongside each computation.
/// To compute a full Jacobian, it evaluates the function once per input dimension.
#[derive(Clone)]
pub struct ForwardAD {}
impl ForwardAD {
    pub fn new() -> Self {
        Self {}
    }
}
impl DerivativeMethodTrait for ForwardAD {
    type T = adfn<1>;

    fn derivative<D: DifferentiableFunctionTrait<Self::T> + ?Sized>(
        &self,
        inputs: &[f64],
        function: &D,
    ) -> (Vec<f64>, DMatrix<f64>) {
        let num_inputs = inputs.len();
        let num_outputs = function.num_outputs();
        let mut out_derivative = DMatrix::zeros(num_outputs, num_inputs);
        let mut out_value = vec![];

        for col_idx in 0..num_inputs {
            let mut inputs_ad = vec![];
            for (i, input) in inputs.iter().enumerate() {
                if i == col_idx {
                    inputs_ad.push(adfn::new(*input, [1.0]))
                } else {
                    inputs_ad.push(adfn::new(*input, [0.0]))
                }
            }

            // let f = D::call(&inputs_ad, args);
            let freeze = if col_idx == 0 { false } else { true };
            let f = function.call(&inputs_ad, freeze);
            assert_eq!(
                f.len(),
                num_outputs,
                "{}",
                format!("does not match {}, {}", f.len(), num_outputs)
            );
            for (row_idx, res) in f.iter().enumerate() {
                if out_value.len() < num_outputs {
                    out_value.push(res.value);
                }
                if res.tangent[0].is_nan() {
                    out_derivative[(row_idx, col_idx)] = res.tangent[0];
                } else {
                    out_derivative[(row_idx, col_idx)] = res.tangent[0];
                }
            }
        }

        (out_value, out_derivative)
    }
}

/// Defines a method for computing the Hessian of a `DifferentiableFunctionTrait`.
///
/// Implementations of this trait are used by `FunctionEngine::hessian` to extract
/// second-order information from recursive AD types.
///
/// # Implementors
/// * `HessianAD<N>`: For Forward-over-Forward Hessians.
/// * `HessianAD_FOR<N>`: For Forward-over-Reverse Hessians.
#[diagnostic::on_unimplemented(
    message = "the derivative method `{Self}` does not support Hessian computation",
    label = "this method does not implement `HessianMethodTrait`",
    note = "Hessian computation requires recursive AD types. Use `HessianAD<N>` (Forward-over-Forward) or `HessianAD_FOR<N>` (Forward-over-Reverse) instead."
)]
pub trait HessianMethodTrait: DerivativeMethodTrait {
    /// Computes the function's value, Jacobian, and Hessian matrices at the given input point.
    fn hessian<D: DifferentiableFunctionTrait<Self::T> + ?Sized>(
        &self,
        inputs: &[f64],
        function: &D,
    ) -> (Vec<f64>, DMatrix<f64>, Vec<DMatrix<f64>>);
}

#[cfg(feature = "hessian")]
use crate::hyper_ad::hyper::HyperAD_ADFN;

#[cfg(feature = "hessian")]
#[derive(Clone)]
pub struct HessianAD<const N: usize> {}
#[cfg(feature = "hessian")]
impl<const N: usize> HessianAD<N> {
    pub fn new() -> Self {
        Self {}
    }
}

#[cfg(feature = "hessian")]
impl<const N: usize> DerivativeMethodTrait for HessianAD<N> {
    type T = HyperAD_ADFN<N>;

    fn derivative<D: DifferentiableFunctionTrait<Self::T> + ?Sized>(
        &self,
        inputs: &[f64],
        function: &D,
    ) -> (Vec<f64>, DMatrix<f64>) {
        // HessianAD can still be used for just derivatives
        let num_inputs = inputs.len();
        let num_outputs = function.num_outputs();
        let mut out_derivative = DMatrix::zeros(num_outputs, num_inputs);
        let mut out_value = vec![];

        let mut inputs_ad = vec![];
        for (i, input) in inputs.iter().enumerate() {
            let mut inner = adfn::<N>::constant(*input);
            if i < N {
                inner.set_tangent_value(i, 1.0);
            }
            let mut outer = HyperAD_ADFN::<N>::new_inner_constant(inner);
            if i < N {
                outer.set_tangent_value(i, 1.0);
            }
            inputs_ad.push(outer);
        }

        let f = function.call(&inputs_ad, false);
        for (row_idx, res) in f.iter().enumerate() {
            out_value.push(res.value());
            let grad = res.tangent_as_vec();
            for (col_idx, g) in grad.iter().enumerate() {
                if col_idx < num_inputs {
                    out_derivative[(row_idx, col_idx)] = *g;
                }
            }
        }

        (out_value, out_derivative)
    }
}

#[cfg(feature = "hessian")]
impl<const N: usize> HessianMethodTrait for HessianAD<N> {
    fn hessian<D: DifferentiableFunctionTrait<Self::T> + ?Sized>(
        &self,
        inputs: &[f64],
        function: &D,
    ) -> (Vec<f64>, DMatrix<f64>, Vec<DMatrix<f64>>) {
        let num_inputs = inputs.len();
        let num_outputs = function.num_outputs();
        let mut out_value = vec![];
        let mut out_jacobian = DMatrix::zeros(num_outputs, num_inputs);
        let mut out_hessians = vec![DMatrix::zeros(num_inputs, num_inputs); num_outputs];

        // Loop over rows and columns in batches of N
        for row_batch_start in (0..num_inputs).step_by(N) {
            for col_batch_start in (0..num_inputs).step_by(N) {
                let mut inputs_ad = vec![];
                for (i, input) in inputs.iter().enumerate() {
                    let mut inner = adfn::<N>::constant(*input);
                    if i >= col_batch_start && i < col_batch_start + N {
                        inner.set_tangent_value(i - col_batch_start, 1.0);
                    }
                    let mut outer = HyperAD_ADFN::<N>::new_inner_constant(inner);
                    if i >= row_batch_start && i < row_batch_start + N {
                        outer.set_tangent_value(i - row_batch_start, 1.0);
                    }
                    inputs_ad.push(outer);
                }

                let f = function.call(&inputs_ad, row_batch_start > 0 || col_batch_start > 0);
                for (row_idx, res) in f.iter().enumerate() {
                    if row_batch_start == 0 && col_batch_start == 0 {
                        out_value.push(res.value());
                    }
                    
                    // Extract Jacobian (only need to do this for one set of column batches per row batch)
                    if row_batch_start == 0 {
                        let grad = res.inner_value().tangent_as_vec();
                        for i in 0..N {
                            if col_batch_start + i < num_inputs {
                                out_jacobian[(row_idx, col_batch_start + i)] = grad[i];
                            }
                        }
                    }

                    // Extract Hessian block
                    for i in 0..N {
                        let r_idx = row_batch_start + i;
                        if r_idx >= num_inputs { break; }
                        
                        let hess_row_chunk = res.tangent[i].tangent_as_vec();
                        for j in 0..N {
                            let c_idx = col_batch_start + j;
                            if c_idx >= num_inputs { break; }
                            out_hessians[row_idx][(r_idx, c_idx)] = hess_row_chunk[j];
                        }
                    }
                }
            }
        }

        (out_value, out_jacobian, out_hessians)
    }
}

#[cfg(all(feature = "hessian", feature = "std"))]
use crate::hyper_ad::hyper_adr::HyperAD_ADR;

#[cfg(all(feature = "hessian", feature = "std"))]
#[derive(Clone)]
#[allow(non_camel_case_types)]
pub struct HessianAD_FOR<const N: usize> {}
#[cfg(all(feature = "hessian", feature = "std"))]
impl<const N: usize> HessianAD_FOR<N> {
    pub fn new() -> Self {
        Self {}
    }
}

#[cfg(all(feature = "hessian", feature = "std"))]
impl<const N: usize> DerivativeMethodTrait for HessianAD_FOR<N> {
    type T = HyperAD_ADR<N>;

    fn derivative<D: DifferentiableFunctionTrait<Self::T> + ?Sized>(
        &self,
        inputs: &[f64],
        function: &D,
    ) -> (Vec<f64>, DMatrix<f64>) {
        let res = self.hessian(inputs, function);
        (res.0, res.1)
    }
}

#[cfg(all(feature = "hessian", feature = "std"))]
impl<const N: usize> HessianMethodTrait for HessianAD_FOR<N> {
    fn hessian<D: DifferentiableFunctionTrait<Self::T> + ?Sized>(
        &self,
        inputs: &[f64],
        function: &D,
    ) -> (Vec<f64>, DMatrix<f64>, Vec<DMatrix<f64>>) {
        let num_inputs = inputs.len();
        let num_outputs = function.num_outputs();
        let mut out_value = vec![];
        let mut out_jacobian = DMatrix::zeros(num_outputs, num_inputs);
        let mut out_hessians = vec![DMatrix::zeros(num_inputs, num_inputs); num_outputs];

        // Loop over rows in batches of N. Each pass recovers full columns via backprop.
        for row_batch_start in (0..num_inputs).step_by(N) {
            let mut inputs_ad = vec![];
            let mut inputs_adr = vec![];
            for (i, input) in inputs.iter().enumerate() {
                // Reset computation graph only on the very first batch
                let adr_var = crate::reverse_ad::adr::adr::new_variable(*input, i == 0 && row_batch_start == 0);
                inputs_adr.push(adr_var);
                let mut outer = HyperAD_ADR::<N>::new_inner_constant(adr_var);
                if i >= row_batch_start && i < row_batch_start + N {
                    outer.set_tangent_value(i - row_batch_start, 1.0);
                }
                inputs_ad.push(outer);
            }

            let f = function.call(&inputs_ad, row_batch_start > 0);
            for (row_idx, res) in f.iter().enumerate() {
                if row_batch_start == 0 {
                    out_value.push(res.value());
                    
                    // Extract full Jacobian from primal ADR
                    let grad = res.value.get_backwards_mode_grad();
                    for (col_idx, adr_var) in inputs_adr.iter().enumerate() {
                        out_jacobian[(row_idx, col_idx)] = grad.wrt(adr_var);
                    }
                }

                // Extract Hessian rows for this batch
                for i in 0..N {
                    let r_idx = row_batch_start + i;
                    if r_idx >= num_inputs { break; }
                    
                    let grad_hess = res.tangent[i].get_backwards_mode_grad();
                    for (c_idx, adr_var) in inputs_adr.iter().enumerate() {
                        out_hessians[row_idx][(r_idx, c_idx)] = grad_hess.wrt(adr_var);
                    }
                }
            }
        }

        (out_value, out_jacobian, out_hessians)
    }
}

pub struct DerivativeMethodClassForwardADMulti<A: AD + ForwardADTrait>(PhantomData<A>);
impl<A: AD + ForwardADTrait> DerivativeMethodClass for DerivativeMethodClassForwardADMulti<A> {
    type DerivativeMethod = ForwardADMulti<A>;
}

/// Computes derivatives using Forward-mode Automatic Differentiation (Multi-Tangent).
///
/// This method allows for computing multiple columns of the Jacobian in a single pass
/// by propagating a vector of tangents. This can significantly speed up computation
/// by taking advantage of SIMD and reducing function overhead.
#[derive(Clone)]
pub struct ForwardADMulti<A: AD + ForwardADTrait> {
    phantom_data: PhantomData<A>,
}
impl<A: AD + ForwardADTrait> ForwardADMulti<A> {
    pub fn new() -> Self {
        Self {
            phantom_data: PhantomData::default(),
        }
    }
}
impl<A: AD + ForwardADTrait> DerivativeMethodTrait for ForwardADMulti<A> {
    type T = A;

    fn derivative<D: DifferentiableFunctionTrait<Self::T> + ?Sized>(
        &self,
        inputs: &[f64],
        function: &D,
    ) -> (Vec<f64>, DMatrix<f64>) {
        let num_inputs = inputs.len();
        let num_outputs = function.num_outputs();
        let mut out_derivative = DMatrix::zeros(num_outputs, num_inputs);
        let mut out_value = vec![];

        let mut curr_idx = 0;

        let mut freeze = false;
        let k = Self::T::tangent_size();
        'l1: loop {
            let mut inputs_ad = vec![];
            for input in inputs.iter() {
                // inputs_ad.push(adf::new(*input, [0.0; K]))
                inputs_ad.push(Self::T::constant(*input));
            }

            'l2: for i in 0..k {
                if curr_idx + i >= num_inputs {
                    break 'l2;
                }
                // inputs_ad[curr_idx+i].tangent[i] = 1.0;
                inputs_ad[curr_idx + i].set_tangent_value(i, 1.0);
            }

            let f = function.call(&inputs_ad, freeze);
            freeze = true;
            assert_eq!(f.len(), num_outputs);

            for (row_idx, res) in f.iter().enumerate() {
                if out_value.len() < num_outputs {
                    out_value.push(res.value());
                }
                let curr_tangent = res.tangent_as_vec();
                'l3: for i in 0..k {
                    if curr_idx + i >= num_inputs {
                        break 'l3;
                    }
                    // out_derivative[(row_idx, curr_idx+i)] = res.tangent[i];
                    if curr_tangent[i].is_nan() {
                        out_derivative[(row_idx, curr_idx + i)] = curr_tangent[i];
                    } else {
                        out_derivative[(row_idx, curr_idx + i)] = curr_tangent[i];
                    }
                }
            }

            curr_idx += k;
            if curr_idx >= num_inputs {
                break 'l1;
            }
        }

        return (out_value, out_derivative);
    }
}

#[cfg(feature = "nightly")]
pub struct DerivativeMethodClassFiniteDifferencingMulti<const K: usize>;
#[cfg(feature = "nightly")]
impl<const K: usize> DerivativeMethodClass for DerivativeMethodClassFiniteDifferencingMulti<K> {
    type DerivativeMethod = FiniteDifferencingMulti2<K>;
}

#[cfg(feature = "nightly")]
#[derive(Clone)]
pub struct FiniteDifferencingMulti2<const K: usize>;
#[cfg(feature = "nightly")]
impl<const K: usize> FiniteDifferencingMulti2<K> {
    pub fn new() -> Self {
        Self {}
    }
}
#[cfg(feature = "nightly")]
impl<const K: usize> DerivativeMethodTrait for FiniteDifferencingMulti2<K> {
    type T = f64xn<K>;

    fn derivative<D: DifferentiableFunctionTrait<Self::T> + ?Sized>(
        &self,
        inputs: &[f64],
        function: &D,
    ) -> (Vec<f64>, DMatrix<f64>) {
        let num_inputs = inputs.len();
        let num_outputs = function.num_outputs();
        let mut out_derivative = DMatrix::zeros(num_outputs, num_inputs);
        let mut out_value = vec![];

        let h = 0.0000001;

        let mut curr_idx = 0;
        let mut first_loop = true;

        'l1: loop {
            let mut inputs_ad = vec![];
            for input in inputs.iter() {
                inputs_ad.push(f64xn::<K>::splat(*input));
            }

            if first_loop {
                'l2: for i in 0..K {
                    if curr_idx + i >= num_inputs {
                        break 'l2;
                    }
                    if i + 1 >= K {
                        break 'l2;
                    }
                    inputs_ad[curr_idx + i].value[i + 1] += h;
                }
            } else {
                'l2: for i in 0..K {
                    if curr_idx + i >= num_inputs {
                        break 'l2;
                    }
                    if i >= K {
                        break 'l2;
                    }
                    inputs_ad[curr_idx + i].value[i] += h;
                }
            }

            // let f = D::call(&inputs_ad, args);
            let f = function.call(&inputs_ad, false);
            assert_eq!(f.len(), num_outputs);

            if first_loop {
                for res in f.iter() {
                    out_value.push(res.value[0]);
                }
            }

            for (row_idx, res) in f.iter().enumerate() {
                if first_loop {
                    'l3: for i in 0..K {
                        if curr_idx + i >= num_inputs {
                            break 'l3;
                        }
                        if i + 1 >= K {
                            break 'l3;
                        }
                        out_derivative[(row_idx, curr_idx + i)] =
                            (res.value[i + 1] - out_value[row_idx]) / h;
                    }
                } else {
                    'l3: for i in 0..K {
                        if curr_idx + i >= num_inputs {
                            break 'l3;
                        }
                        if i >= K {
                            break 'l3;
                        }
                        out_derivative[(row_idx, curr_idx + i)] =
                            (res.value[i] - out_value[row_idx]) / h;
                    }
                }
            }

            if first_loop {
                first_loop = false;
                curr_idx += K - 1;
            } else {
                curr_idx += K;
            }

            if curr_idx >= num_inputs {
                break 'l1;
            }
        }

        return (out_value, out_derivative);
    }
}

#[cfg(feature = "std")]
#[derive(Clone)]
pub struct WASP {
    cache: Arc<RwLock<WASPCache>>,
    num_f_calls: Arc<RwLock<usize>>,
    d_theta: f64,
    d_ell: f64,
}
#[cfg(feature = "std")]
impl WASP {
    pub fn new(n: usize, m: usize, orthonormal_delta_x: bool, d_theta: f64, d_ell: f64) -> Self {
        Self {
            cache: Arc::new(RwLock::new(WASPCache::new(n, m, orthonormal_delta_x))),
            num_f_calls: Arc::new(RwLock::new(0)),
            d_theta,
            d_ell,
        }
    }
    pub fn reset_cache(&self) {
        self.cache.write().unwrap().reset();
    }
    pub fn new_default(n: usize, m: usize) -> Self {
        Self::new(n, m, true, 0.3, 0.3)
    }
    pub fn num_f_calls(&self) -> usize {
        return self.num_f_calls.read().unwrap().clone();
    }
}
#[cfg(feature = "std")]
impl DerivativeMethodTrait for WASP {
    type T = f64;

    fn derivative<D: DifferentiableFunctionTrait<Self::T> + ?Sized>(
        &self,
        inputs: &[f64],
        function: &D,
    ) -> (Vec<f64>, DMatrix<f64>) {
        let mut num_f_calls = 0;
        let f_k = function.call(inputs, false);
        let f_k_dv = DVector::from_column_slice(&f_k);
        num_f_calls += 1;
        let epsilon = 0.000001;

        let mut cache = self.cache.write().unwrap();
        let n = inputs.len();

        let x = DVector::<f64>::from_column_slice(inputs);

        loop {
            let i = cache.i.clone();

            let delta_x_i = cache.delta_x.column(i);

            let x_k_plus_delta_x_i: DVector<f64> = &x + epsilon * &delta_x_i;
            let f_k_plus_delta_x_i = DVector::<f64>::from_column_slice(
                &function.call(x_k_plus_delta_x_i.as_slice(), true),
            );
            num_f_calls += 1;
            let delta_f_i = (&f_k_plus_delta_x_i - &f_k_dv) / epsilon;
            let delta_f_i_hat = cache.delta_f_t.row(i);
            let delta_f_i_hat = DVector::from_column_slice(delta_f_i_hat.transpose().as_slice());
            let return_result = close_enough(&delta_f_i, &delta_f_i_hat, self.d_theta, self.d_ell);

            cache.delta_f_t.set_row(i, &delta_f_i.transpose());
            let c_1_mat = &cache.c_1[i];
            let c_2_mat = &cache.c_2[i];
            let delta_f_t = &cache.delta_f_t;

            let d_t_star = c_1_mat * delta_f_t + c_2_mat * delta_f_i.transpose();
            let d_star = d_t_star.transpose();

            let tmp = &d_star * &cache.delta_x;
            cache.delta_f_t = tmp.transpose();

            let mut new_i = i + 1;
            if new_i >= n {
                new_i = 0;
            }
            cache.i = new_i;

            if return_result {
                *self.num_f_calls.write().unwrap() = num_f_calls;
                return (f_k, d_star);
            }
        }
    }
}

#[cfg(feature = "std")]
#[derive(Clone, Debug)]
pub struct WASPCache {
    pub n: usize,
    pub m: usize,
    pub i: usize,
    pub delta_f_t: DMatrix<f64>,
    pub delta_x: DMatrix<f64>,
    pub c_1: Vec<DMatrix<f64>>,
    pub c_2: Vec<DVector<f64>>,
}
#[cfg(feature = "std")]
impl WASPCache {
    pub fn new(n: usize, m: usize, orthonormal_delta_x: bool) -> Self {
        let delta_f_t = DMatrix::<f64>::identity(n, m);
        let delta_x = get_tangent_matrix(n, orthonormal_delta_x);
        let mut c_1 = vec![];
        let mut c_2 = vec![];

        let a_mat: DMatrix<f64> = 2.0 * &delta_x * &delta_x.transpose();
        let a_inv_mat = a_mat.try_inverse().unwrap();

        for i in 0..n {
            let delta_x_i = DVector::<f64>::from_column_slice(delta_x.column(i).as_slice());
            let s_i = (delta_x_i.transpose() * &a_inv_mat * &delta_x_i)[(0, 0)];
            let s_i_inv = 1.0 / s_i;
            let c_1_mat = &a_inv_mat
                * (DMatrix::<f64>::identity(n, n)
                    - s_i_inv * &delta_x_i * delta_x_i.transpose() * &a_inv_mat)
                * 2.0
                * &delta_x;
            let c_2_mat = s_i_inv * &a_inv_mat * delta_x_i;
            c_1.push(c_1_mat);
            c_2.push(c_2_mat);
        }

        return Self {
            n,
            m,
            i: 0,
            delta_f_t,
            delta_x,
            c_1,
            c_2,
        };
    }
    pub fn reset(&mut self) {
        self.delta_f_t = DMatrix::<f64>::identity(self.n, self.m);
        self.i = 0;
    }
}

/*
#[derive(Clone)]
pub struct WASP2 {
    cache: Arc<RwLock<WASPCache2>>,
    num_f_calls: Arc<RwLock<usize>>,
    d_theta: f64,
    d_ell: f64
}
impl WASP2 {
    pub fn new(n: usize, m: usize, alpha:f64, orthonormal_delta_x: bool, d_theta: f64, d_ell: f64) -> Self {
        Self {
            cache: Arc::new(RwLock::new(WASPCache2::new(n, m, alpha, orthonormal_delta_x))),
            num_f_calls: Arc::new(RwLock::new(0)),
            d_theta,
            d_ell,
        }
    }
    pub fn new_default(n: usize, m: usize) -> Self {
        Self::new(n, m, 0.98, true, 0.3, 0.3)
    }
    pub fn num_f_calls(&self) -> usize {
        return self.num_f_calls.read().unwrap().clone()
    }
}
impl DerivativeMethodTrait for WASP2 {
    type T = f64;

    fn derivative<D: DifferentiableFunctionTrait<Self::T> + ?Sized>(&self, inputs: &[f64], function: &D) -> (Vec<f64>, DMatrix<f64>) {
        let mut num_f_calls = 0;
        let f_k = function.call(inputs, false);
        let f_k_dv = DVector::from_column_slice(&f_k);
        num_f_calls += 1;
        let epsilon = 0.000001;

        let mut cache = self.cache.write().unwrap();
        let n = inputs.len();

        let x = DVector::<f64>::from_column_slice(inputs);

        loop {
            let i = cache.i.clone();

            let delta_x_i = cache.delta_x.column(i);

            let x_k_plus_delta_x_i: DVector<f64> = &x + epsilon*&delta_x_i;
            let f_k_plus_delta_x_i = DVector::<f64>::from_column_slice(&function.call(x_k_plus_delta_x_i.as_slice(), true));
            num_f_calls += 1;
            let delta_f_i = (&f_k_plus_delta_x_i - &f_k_dv) / epsilon;
            let delta_f_i_hat = &cache.curr_d * &delta_x_i;
            // let delta_f_i_hat = cache.delta_f_t.row(i);
            // let delta_f_i_hat = DVector::from_column_slice(delta_f_i_hat.transpose().as_slice());
            let return_result = close_enough(&delta_f_i, &delta_f_i_hat, self.d_theta, self.d_ell);

            cache.delta_f_t.set_row(i, &delta_f_i.transpose());
            let c_1_mat = &cache.c_1[i];
            let c_2_mat = &cache.c_2[i];
            let delta_f_t = &cache.delta_f_t;

            let d_t_star = c_1_mat*delta_f_t + c_2_mat*delta_f_i.transpose();
            let d_star = d_t_star.transpose();
            cache.curr_d = d_star.clone();

            let mut new_i = i + 1;
            if new_i >= n { new_i = 0; }
            cache.i = new_i;

            if return_result {
                *self.num_f_calls.write().unwrap() = num_f_calls;
                return (f_k, d_star);
            }
        }
    }
}

pub struct WASPCache2 {
    pub i: usize,
    pub curr_d: DMatrix<f64>,
    pub delta_f_t: DMatrix<f64>,
    pub delta_x: DMatrix<f64>,
    pub c_1: Vec<DMatrix<f64>>,
    pub c_2: Vec<DVector<f64>>
}
impl WASPCache2 {
    pub fn new(n: usize, m: usize, alpha: f64, orthonormal_delta_x: bool) -> Self {
        assert!(alpha > 0.0 && alpha < 1.0);

        let curr_d = DMatrix::<f64>::identity(m, n);
        let delta_f_t = DMatrix::<f64>::identity(n, m);
        let delta_x = get_tangent_matrix(n, orthonormal_delta_x);
        let mut c_1 = vec![];
        let mut c_2 = vec![];

        for i in 0..n {
            let delta_x_i = DVector::<f64>::from_column_slice(delta_x.column(i).as_slice());
            let mut w_i = DMatrix::<f64>::zeros(n, n);
            for j in 0..n {
                let exponent = math_mod(i as i32 - j as i32, n as i32) as f64 / (n as i32 - 1) as f64;
                w_i[(j, j)] = alpha * (1.0 - alpha).powf(exponent);
            }
            let w_i_2 = &w_i * &w_i;

            let a_i = 2.0 * &delta_x * &w_i_2 * &delta_x.transpose();
            let a_i_inv = a_i.clone().try_inverse().unwrap();

            let s_i = (delta_x_i.transpose() * &a_i_inv * &delta_x_i)[(0,0)];
            let s_i_inv = 1.0 / s_i;
            let c_1_mat = &a_i_inv * (DMatrix::<f64>::identity(n, n) - s_i_inv * &delta_x_i * delta_x_i.transpose() * &a_i_inv) * 2.0 * &delta_x * &w_i_2;
            let c_2_mat = s_i_inv * &a_i_inv * delta_x_i;
            c_1.push(c_1_mat);
            c_2.push(c_2_mat);
        }

        return Self {
            i: 0,
            curr_d,
            delta_f_t,
            delta_x,
            c_1,
            c_2,
        }
    }
}

pub fn math_mod(a: i32, b: i32) -> i32 {
    return ((a % b) + b) % b;
}
*/

#[cfg(feature = "std")]
pub(crate) fn get_tangent_matrix(n: usize, orthogonal: bool) -> DMatrix<f64> {
    let mut rng = thread_rng();
    let uniform = Uniform::new(-1.0, 1.0);

    let t = DMatrix::<f64>::from_fn(n, n, |_, _| uniform.sample(&mut rng));

    return if orthogonal {
        let svd = t.svd(true, true);
        let delta_x = svd.u.as_ref().unwrap() * svd.v_t.as_ref().unwrap();
        delta_x
    } else {
        t
    };
}

pub(crate) fn close_enough(a: &DVector<f64>, b: &DVector<f64>, d_theta: f64, d_ell: f64) -> bool {
    let a_n = a.norm();
    let b_n = b.norm();

    let tmp = ((a.dot(&b) / (a_n * b_n)) - 1.0).abs();
    if tmp > d_theta {
        return false;
    }

    let tmp1 = if b_n != 0.0 {
        ((a_n / b_n) - 1.0).abs()
    } else {
        f64::MAX
    };
    let tmp2 = if a_n != 0.0 {
        ((b_n / a_n) - 1.0).abs()
    } else {
        f64::MAX
    };

    if f64::min(tmp1, tmp2) > d_ell {
        return false;
    }

    return true;
}

/*

pub fn math_modulus(a: i64, b: i64) -> usize {
    (((a % b) + b) % b) as usize
}

pub fn get_tangent_matrix(n: usize, orthonormalize: bool) -> DMatrix<f64> {
    let mut out = DMatrix::zeros(n, n);
    let mut rng = rand::rng();

    for i in 0..n {
        for j in 0..n {
            out[(i, j)] = rng.random_range(-1.0..=1.0);
        }
    }

    if orthonormalize {
        let svd = out.svd(true, true);
        out = svd.u.as_ref().unwrap()*svd.v_t.as_ref().unwrap();
    }

    return out;
}

pub fn wasp_projection<D: DifferentiableFunctionTrait<f64> + ?Sized>(f: &D, f_x_k: &DVector<f64>, x_k: &[f64], cache: &WASPCache) -> DMatrix<f64> {
    let epsilon = 0.00001;
    let x_k = DVector::from_column_slice(x_k);
    let i = cache.i.lock().unwrap();
    let c_1_mat = &cache.c_1_mats[*i];
    let c_2_mat = &cache.c_2_mats[*i];
    let delta_x_i = DVector::from_column_slice(cache.delta_x_mat.column(*i).as_slice());
    let f_x_k_delta = DVector::from_column_slice(&f.call((&x_k + epsilon*&delta_x_i).as_slice(), true));
    let delta_f_i = (f_x_k_delta - f_x_k) / epsilon;
    let mut delta_f_hat_t = cache.delta_f_mat_t.lock().unwrap();
    delta_f_hat_t.set_row(*i, &delta_f_i.transpose());
    return c_1_mat*&*delta_f_hat_t + c_2_mat*&delta_f_i.transpose();
}

pub fn wasp_projection2<D: DifferentiableFunctionTrait<f64> + ?Sized>(f: &D, f_x_k: &DVector<f64>, x_k: &[f64], cache: &WASPCache2) -> DMatrix<f64> {
    let epsilon = 0.00001;
    let x_k = DVector::from_column_slice(x_k);
    let i = cache.i.lock().unwrap();
    let c_1_mat = &cache.c_1_mats[*i];
    let c_2_mat = &cache.c_2_mats[*i];
    let delta_x_i = DVector::from_column_slice(cache.delta_x_mat.column(*i).as_slice());
    let f_x_k_delta = DVector::from_column_slice(&f.call((&x_k + epsilon*&delta_x_i).as_slice(), true));
    let delta_f_i = (f_x_k_delta - f_x_k) / epsilon;
    let mut delta_f_hat_t = cache.delta_f_mat_t.lock().unwrap();
    delta_f_hat_t.set_row(*i, &delta_f_i.transpose());
    return c_1_mat*&*delta_f_hat_t + c_2_mat*&delta_f_i.transpose();
}

pub fn close_enough(d_a_t_mat: &DMatrix<f64>, d_b_t_mat: &DMatrix<f64>, l: usize, m: usize, d_theta: f64) -> bool {
    let mut numbers: Vec<usize> = (0..m).collect();

    let mut rng = rng();
    numbers.shuffle(&mut rng);

    let js: Vec<usize> = numbers.into_iter().take(l).collect();

    // println!("{}", d_a_t_mat);
    // println!("{}", d_b_t_mat);
    // println!("---");

    for j in js {
        let d_a = DVector::from_column_slice(d_a_t_mat.column(j).as_slice());
        let d_b = DVector::from_column_slice(d_b_t_mat.column(j).as_slice());

        let d_a_n = d_a.norm();
        let d_b_n = d_b.norm();

        let dot = d_a.dot(&d_b);
        let angle = (dot / (d_a_n * d_b_n)).acos();
        // println!("{:?}", angle);

        if angle > d_theta { return false; }
    }

    return true;
}

#[inline(always)]
pub fn close_enough2(a: &DVector<f64>, b: &DVector<f64>, d_theta: f64, d_l: f64) -> bool {
    let an = a.norm();
    let bn = b.norm();
    let d = a.dot(b);

    if (d / (an * bn) - 1.0).abs() > d_theta { return false; }
    if (an / bn - 1.0).abs() > d_l { return false; }

    return true;
}

pub fn derivative_angular_distance(d_a_t_mat: &DMatrix<f64>, d_b_t_mat: &DMatrix<f64>) -> f64 {
    let m = d_a_t_mat.ncols();

    let mut max_angle = f64::MIN;

    for j in 0..m {
        let d_a = DVector::from_column_slice(d_a_t_mat.column(j).as_slice());
        let d_b = DVector::from_column_slice(d_b_t_mat.column(j).as_slice());

        let d_a_n = d_a.norm();
        let d_b_n = d_b.norm();

        let dot = d_a.dot(&d_b);
        let angle = (dot / (d_a_n * d_b_n)).acos();
        if angle > max_angle { max_angle = angle; }
    }

    return max_angle;
}

#[derive(Clone)]
pub struct WASPCache {
    pub delta_f_mat_t: Arc<Mutex<DMatrix<f64>>>,
    pub delta_x_mat: DMatrix<f64>,
    pub c_1_mats: Vec<DMatrix<f64>>,
    pub c_2_mats: Vec<DVector<f64>>,
    pub i: Arc<Mutex<usize>>
}
impl WASPCache {
    pub fn new(n: usize, m: usize, alpha: f64, orthonormalize: bool) -> Self {
        let delta_x_mat = get_tangent_matrix(n, orthonormalize);
        let mut c_1_mats = vec![];
        let mut c_2_mats = vec![];

        for i in 0..n {
            let delta_x_i = DVector::from_column_slice(delta_x_mat.column(i).as_slice());
            let mut w_i_mat = DMatrix::zeros(n, n);
            for j in 0..n {
                let exp = math_modulus(i as i64 - j as i64, n as i64) as f64 / ((n - 1) as f64);
                w_i_mat[(j,j)] = alpha*(1.0 - alpha).pow(  exp );
            }
            let w_i_mat_2 = &w_i_mat * & w_i_mat;
            let a_i_mat = 2.0*(&delta_x_mat * &w_i_mat_2 * &delta_x_mat.transpose());
            let a_i_mat_inv = a_i_mat.clone().try_inverse().unwrap();
            let s_i = (&delta_x_i.transpose() * &a_i_mat_inv * &delta_x_i)[0];
            let s_i_inv = 1.0 / s_i;
            let c_1_mat = &a_i_mat_inv*(DMatrix::identity(n, n) - s_i_inv*&delta_x_i*&delta_x_i.transpose()*&a_i_mat_inv)*2.0*&delta_x_mat*&w_i_mat_2;
            let c_2_mat = s_i_inv*&a_i_mat_inv*&delta_x_i;
            c_1_mats.push(c_1_mat);
            c_2_mats.push(c_2_mat);
        }

        Self {
            delta_f_mat_t: Arc::new(Mutex::new(DMatrix::zeros(n, m))),
            delta_x_mat,
            c_1_mats,
            c_2_mats,
            i: Arc::new(Mutex::new(0)),
        }
    }
}

#[derive(Clone)]
pub struct WASPCache2 {
    pub delta_f_mat_t: Arc<Mutex<DMatrix<f64>>>,
    pub delta_x_mat: DMatrix<f64>,
    pub c_1_mats: Vec<DMatrix<f64>>,
    pub c_2_mats: Vec<DVector<f64>>,
    pub i: Arc<Mutex<usize>>
}
impl WASPCache2 {
    pub fn new(n: usize, m: usize, orthonormalize: bool) -> Self {
        let delta_x_mat = get_tangent_matrix(n, orthonormalize);
        let mut c_1_mats = vec![];
        let mut c_2_mats = vec![];

        for i in 0..n {
            let delta_x_i = DVector::from_column_slice(delta_x_mat.column(i).as_slice());
            let a_i_mat = 2.0*(&delta_x_mat * &delta_x_mat.transpose());
            let a_i_mat_inv = a_i_mat.clone().try_inverse().unwrap();
            let s_i = (&delta_x_i.transpose() * &a_i_mat_inv * &delta_x_i)[0];
            let s_i_inv = 1.0 / s_i;
            let c_1_mat = &a_i_mat_inv*(DMatrix::identity(n, n) - s_i_inv*&delta_x_i*&delta_x_i.transpose()*&a_i_mat_inv)*2.0*&delta_x_mat;
            let c_2_mat = s_i_inv*&a_i_mat_inv*&delta_x_i;
            c_1_mats.push(c_1_mat);
            c_2_mats.push(c_2_mat);
        }

        Self {
            delta_f_mat_t: Arc::new(Mutex::new(DMatrix::zeros(n, m))),
            delta_x_mat,
            c_1_mats,
            c_2_mats,
            i: Arc::new(Mutex::new(0)),
        }
    }
}

pub struct DerivativeMethodClassWASP;
impl DerivativeMethodClass for DerivativeMethodClassWASP {
    type DerivativeMethod = WASP;
}

#[derive(Clone)]
pub struct WASP {
    pub cache: WASPCache2,
    pub d_theta: f64,
    pub d_l: f64,
    pub num_f_calls: Arc<Mutex<usize>>
}
impl WASP {
    pub fn new(n: usize, m: usize, d_theta: f64, d_l: f64, orthonormalize: bool) -> Self {
        Self {
            cache: WASPCache2::new(n, m, orthonormalize),
            d_theta,
            d_l,
            num_f_calls: Arc::new(Mutex::new(0)),
        }
    }

    pub fn get_num_f_calls(&self) -> usize {
        self.num_f_calls.lock().unwrap().clone()
    }
}
impl DerivativeMethodTrait for WASP {
    type T = f64;

    fn derivative<D: DifferentiableFunctionTrait<Self::T> + ?Sized>(&self, inputs: &[f64], function: &D) -> (Vec<f64>, DMatrix<f64>) {
        let mut num_f_calls = self.num_f_calls.lock().unwrap();
        *num_f_calls = 0;
        let f_x_k_vec = function.call(inputs, false);
        let f_x_k = DVector::from_column_slice(&f_x_k_vec);
        let x_k = DVector::from_column_slice(inputs);
        *num_f_calls += 1;

        let mut return_result;

        loop {
            let mut i = self.cache.i.lock().unwrap();

            let epsilon = 0.00001;
            let c_1_mat = &self.cache.c_1_mats[*i];
            let c_2_mat = &self.cache.c_2_mats[*i];
            let delta_x_i = DVector::from_column_slice(self.cache.delta_x_mat.column(*i).as_slice());
            let f_x_k_delta = DVector::from_column_slice(&function.call((&x_k + epsilon * &delta_x_i).as_slice(), true));
            *num_f_calls += 1;
            let delta_f_i = (f_x_k_delta - &f_x_k) / epsilon;
            let mut delta_f_hat_t = self.cache.delta_f_mat_t.lock().unwrap();
            let delta_f_i_hat = DVector::from_column_slice(delta_f_hat_t.row(*i).transpose().as_slice());

            return_result = close_enough2(&delta_f_i, &delta_f_i_hat, self.d_theta, self.d_l);

            delta_f_hat_t.set_row(*i, &delta_f_i.transpose());
            let d_t = c_1_mat * &*delta_f_hat_t + c_2_mat * &delta_f_i.transpose();
            *delta_f_hat_t = &self.cache.delta_x_mat.transpose() * &d_t;

            *i = (*i + 1) % inputs.len();

            if return_result {
                return (f_x_k_vec, d_t.transpose());
            }
        }
    }
}

pub struct DerivativeMethodClassWASPNec;
impl DerivativeMethodClass for DerivativeMethodClassWASPNec {
    type DerivativeMethod = WASPNec;
}
#[derive(Clone)]
pub struct WASPNec {
    pub cache: WASPCache,
    pub first_call: Arc<Mutex<bool>>,
    pub num_f_calls: Arc<Mutex<usize>>
}
impl WASPNec {
    pub fn new(n: usize, m: usize, alpha: f64, orthonormalize: bool) -> Self {
        Self {
            cache: WASPCache::new(n, m, alpha, orthonormalize),
            first_call: Arc::new(Mutex::new(true)),
            num_f_calls: Arc::new(Mutex::new(0)),
        }
    }

    pub fn get_num_f_calls(&self) -> usize {
        self.num_f_calls.lock().unwrap().clone()
    }
}
impl DerivativeMethodTrait for WASPNec {
    type T = f64;

    fn derivative<D: DifferentiableFunctionTrait<Self::T> + ?Sized>(&self, inputs: &[f64], function: &D) -> (Vec<f64>, DMatrix<f64>) {
        let mut num_f_calls = self.num_f_calls.lock().unwrap();
        *num_f_calls = 0;
        let f_x_k_vec = function.call(inputs, false);
        let f_x_k = DVector::from_column_slice(&f_x_k_vec);
        *num_f_calls += 1;

        let mut first_call = self.first_call.lock().unwrap();
        if *first_call {
            let epsilon = 0.00001;
            let x_k = DVector::from_column_slice(inputs);
            let n = inputs.len();
            for i in 0..n {
                let delta_x_i = DVector::from_column_slice(self.cache.delta_x_mat.column(i).as_slice());
                let f_x_k_delta = DVector::from_column_slice(&function.call((&x_k + epsilon*&delta_x_i).as_slice(), true));
                let delta_f_i = (f_x_k_delta - &f_x_k) / epsilon;
                let mut delta_f_hat_t = self.cache.delta_f_mat_t.lock().unwrap();
                delta_f_hat_t.set_row(i, &delta_f_i.transpose());
            }
            *first_call = false;
        }

        let d_t = wasp_projection(function, &f_x_k, inputs, &self.cache);
        *num_f_calls += 1;
        let mut i = self.cache.i.lock().unwrap();
        *i = (*i + 1) % inputs.len();

        return (f_x_k_vec, d_t.transpose());
    }
}

pub struct DerivativeMethodClassWASPEc;
impl DerivativeMethodClass for DerivativeMethodClassWASPEc {
    type DerivativeMethod = WASPEc;
}

#[derive(Clone)]
pub struct WASPEc {
    pub cache_a: WASPCache,
    pub cache_b: WASPCache,
    pub l: usize,
    pub d_theta: f64,
    pub num_f_calls: Arc<Mutex<usize>>
}
impl WASPEc {
    pub fn new(n: usize, m: usize, alpha: f64, orthonormalize: bool, l: usize, d_theta: f64) -> Self {
        assert!(l <= m);

        Self {
            cache_a: WASPCache::new(n, m, alpha, orthonormalize),
            cache_b: WASPCache::new(n, m, alpha, orthonormalize),
            l,
            d_theta,
            num_f_calls: Arc::new(Mutex::new(0)),
        }
    }

    pub fn get_num_f_calls(&self) -> usize {
        self.num_f_calls.lock().unwrap().clone()
    }
}
impl DerivativeMethodTrait for WASPEc {
    type T = f64;

    fn derivative<D: DifferentiableFunctionTrait<Self::T> + ?Sized>(&self, inputs: &[f64], function: &D) -> (Vec<f64>, DMatrix<f64>) {
        let mut num_f_calls = self.num_f_calls.lock().unwrap();
        *num_f_calls = 0;
        let f_x_k_vec = function.call(inputs, false);
        let f_x_k = DVector::from_column_slice(&f_x_k_vec);
        *num_f_calls += 1;

        loop {
            let d_a_t = wasp_projection(function, &f_x_k, inputs, &self.cache_a);
            let d_b_t = wasp_projection(function, &f_x_k, inputs, &self.cache_b);
            *num_f_calls += 2;
            let mut i_a = self.cache_a.i.lock().unwrap();
            let mut i_b = self.cache_b.i.lock().unwrap();
            *i_a = (*i_a + 1) % inputs.len();
            *i_b = (*i_b + 1) % inputs.len();

            if close_enough(&d_a_t, &d_b_t, self.l, function.num_outputs(), self.d_theta) {
                return (f_x_k_vec, (d_a_t.transpose() + d_b_t.transpose()) * 0.5);
            }
        }
    }
}

*/

#[cfg(feature = "std")]
#[derive(Clone)]
pub struct SPSA;
#[cfg(feature = "std")]
impl SPSA {
    pub fn new() -> Self {
        Self {}
    }
}
#[cfg(feature = "std")]
impl DerivativeMethodTrait for SPSA {
    type T = f64;

    fn derivative<D: DifferentiableFunctionTrait<Self::T> + ?Sized>(
        &self,
        inputs: &[f64],
        function: &D,
    ) -> (Vec<f64>, DMatrix<f64>) {
        let f0 = function.call(inputs, false);

        let mut rng = rand::thread_rng();

        let epsilon = 0.00000001;

        let r: Vec<f64> = (0..inputs.len())
            .into_iter()
            .map(|_x| rng.gen_range(-1.0..=1.0))
            .collect();
        let x = DVector::from_column_slice(inputs);
        let delta_k = DVector::from_column_slice(&r);
        let xpos: DVector<f64> = &x + epsilon * &delta_k;
        let xneg: DVector<f64> = &x - epsilon * &delta_k;
        let fpos = DVector::from_column_slice(&function.call(xpos.as_slice(), false));
        let fneg = DVector::from_column_slice(&function.call(xneg.as_slice(), false));
        let v = (&fpos - &fneg) / (2.0 * epsilon);
        let delta_k_inverse =
            DVector::from_column_slice(&delta_k.iter().map(|x| 1.0 / *x).collect::<Vec<f64>>());
        let out = &v * &delta_k_inverse.transpose();

        (f0, out)
    }
}

pub struct DerivativeMethodClassAlwaysZero;
impl DerivativeMethodClass for DerivativeMethodClassAlwaysZero {
    type DerivativeMethod = DerivativeAlwaysZero;
}

#[derive(Clone)]
pub struct DerivativeAlwaysZero;
impl DerivativeAlwaysZero {
    pub fn new() -> Self {
        Self {}
    }
}
impl DerivativeMethodTrait for DerivativeAlwaysZero {
    type T = f64;

    fn derivative<D: DifferentiableFunctionTrait<Self::T> + ?Sized>(
        &self,
        _inputs: &[f64],
        function: &D,
    ) -> (Vec<f64>, DMatrix<f64>) {
        let num_outputs = function.num_outputs();
        let num_inputs = function.num_inputs();
        (
            vec![0.0; num_outputs],
            DMatrix::from_vec(num_outputs, num_inputs, vec![0.0; num_outputs * num_inputs]),
        )
    }
}