use std::convert::Infallible;
#[cfg(test)]
use crate::test_infra::load_test_file;

use super::FitnessFunction;

pub struct Rosenbrock;

impl Rosenbrock {
    pub fn new() -> Self {
        Rosenbrock
    }
}

impl FitnessFunction for Rosenbrock {
    type In = Vec<f64>;
    type Out = f64;
    type Err = Infallible;

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

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

    for test in data {
        assert_eq!(
            Rosenbrock::new().fit(&test.inp).unwrap(),
            test.out
        )
    }
}