use std::{convert::Infallible, marker::PhantomData};
use nalgebra::{allocator::Allocator, DefaultAllocator, Dim, OVector};

use super::FitnessFunction;

pub struct Rosenbrock<D> {
    _phantom: PhantomData<D>
}

impl<D> Rosenbrock<D>
where
    D: Dim,
    DefaultAllocator: Allocator<D>
{
    pub fn new() -> Self {
        Rosenbrock {
            _phantom: PhantomData
        }
    }
}

impl<D> FitnessFunction for Rosenbrock<D>
where
    D: Dim,
    DefaultAllocator: Allocator<D>
{
    type In = OVector<f64, D>;
    type Out = f64;
    type Err = Infallible;

    fn fit(self: &Self, inp: &OVector<f64, D>) -> Result<f64, Infallible> {
        Ok(inp.as_slice().windows(2)
           .map(|xs| 100.0 * (xs[1] - xs[0].powi(2)).powi(2) + (1.0 - xs[0]).powi(2))
           .sum())
    }
}

#[cfg(test)]
pub mod tests {
    use nalgebra::{Dyn, OVector};

    use crate::{fitness::{rosenbrock::Rosenbrock, FitnessFunction}, test_infra::load_test_file};

    #[test]
    fn test_rosenbrock() {
        let data = load_test_file::<f64, f64>("tests/rosenbrock.txt");

        for test in data {
            assert_eq!(
                Rosenbrock::<Dyn>::new().fit(&OVector::<_, Dyn>::from_vec(test.inp)).unwrap(),
                test.out
            )
        }
    }
}