use std::{convert::Infallible, f64::consts::PI, marker::PhantomData};

use nalgebra::{allocator::Allocator, DefaultAllocator, Dim, OVector, U1};

use super::FitnessFunction;

pub struct Linear<D>
where
    D: Dim,
    DefaultAllocator: Allocator<D>
{
    a0: f64,
    a: OVector<f64, D>,
}

impl<D> Linear<D>
where
    D: Dim,
    DefaultAllocator: Allocator<D>
{
    pub fn new(a0: f64, a: OVector<f64, D>) -> Self {
        Self {
            a0,
            a
        }
    }
}

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

    fn fit(self: &Self, inp: &Self::In) -> Result<Self::Out, Self::Err> {
        Ok(self.a0 + (self.a.transpose() * inp).x)
    }
}

pub struct Step<D>
where
    D: Dim,
    DefaultAllocator: Allocator<D>,
{
    a0: f64,
    a: OVector<f64, D>,
}

impl<D> Step<D>
where
    D: Dim,
    DefaultAllocator: Allocator<D>,
{
    pub fn new(a0: f64, a: OVector<f64, D>) -> Self {
        Self {
            a0,
            a
        }
    }
}

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

    fn fit(self: &Self, inp: &Self::In) -> Result<Self::Out, Self::Err> {
        Ok(self.a0 + inp.component_mul(&self.a).map(|x| x.floor()).sum())
    }
}

pub struct Rastrigin<D: Dim> {
    _phantom: PhantomData<D>
}

impl<D: Dim> Rastrigin<D> {
    pub fn new() -> Self {
        Self {
            _phantom: PhantomData
        }
    }
}

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

    fn fit(self: &Self, inp: &Self::In) -> Result<Self::Out, Self::Err> {
        let dim = inp.len() as f64;
        Ok(10.0 * dim + inp.iter()
           .map(|x| x.powi(2) - 10.0 * (2.0 * PI * x).cos())
           .sum::<f64>())
    }
}

pub struct Griewank<D: Dim> {
    _phantom: PhantomData<D>
}

impl<TDim: Dim> Griewank<TDim> {
    pub fn new() -> Self {
        Self {
            _phantom: PhantomData
        }
    }
}

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

    fn fit(self: &Self, inp: &Self::In) -> Result<Self::Out, Self::Err> {
        Ok(1.0 + inp.map(|x| x.powi(2)).sum() / 4000.0 - inp
           .iter()
           .enumerate()
           .map(|(i, x)| (x / ((i + 1) as f64).sqrt()).cos())
           .product::<f64>())
    }
}

pub struct Schwefel<D: Dim> {
    _phantom: PhantomData<D>
}

impl<D: Dim> Schwefel<D> {
    pub fn new() -> Self {
        Self {
            _phantom: PhantomData
        }
    }
}

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

    fn fit(self: &Self, inp: &Self::In) -> Result<Self::Out, Self::Err> {
        Ok(-inp
            .map(|x| x * x.abs().sqrt().sin())
            .sum())
    }
}

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

    use crate::{fitness::{real::{Linear, Rastrigin, Schwefel, Step}, FitnessFunction}, test_infra::load_test_file};

    use super::Griewank;

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

        for test in data {
            let offset = OVector::<f64, Dyn>::from_element(test.inp.len(), 1.0);
            let inp = OVector::<f64, Dyn>::from_vec(test.inp);

            let linear = Linear::new(1.0, offset);

            assert_eq!(
                linear.fit(&inp).unwrap(),
                test.out
            );
        }
    }

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

        for test in data {
            let offset = OVector::<f64, Dyn>::from_iterator(test.inp.len(), (2..=test.inp.len()+1).map(|x| x as f64));
            let linear = Linear::new(1.0, offset);
            let inp = OVector::<f64, Dyn>::from_vec(test.inp);

            assert_eq!(
                linear.fit(&inp).unwrap(),
                test.out
            );
        }
    }

    #[test]
    fn test_step_1() {
        const MAX_LEN: usize = 10;

        let data = load_test_file::<f64, f64>("tests/step_1.txt");
        let offset = SVector::repeat(1.0);
        let linear = Step::new(1.0, offset);

        for test in data {
            let filled = test.inp.iter()
                .chain(vec![0f64; MAX_LEN - test.inp.len()].iter())
                .map(|x| *x)
                .collect::<Vec<_>>();
            let inp = SVector::<f64, MAX_LEN>::from_vec(filled);

            assert_eq!(
                linear.fit(&inp).unwrap(),
                test.out
            );
        }
    }

    #[test]
    fn test_step_2() {
        const MAX_LEN: usize = 10;

        let data = load_test_file::<f64, f64>("tests/step_2.txt");
        let offset = SVector::from_vec((2..=11).map(|x| x as f64).collect());
        let linear = Step::new(1.0, offset);

        for test in data {
            let filled = test.inp.iter()
                .chain(vec![0f64; MAX_LEN - test.inp.len()].iter())
                .map(|x| *x)
                .collect::<Vec<_>>();
            let inp = SVector::<f64, MAX_LEN>::from_vec(filled);

            assert_eq!(
                linear.fit(&inp).unwrap(),
                test.out
            );
        }
    }

    #[test]
    fn test_rastrigin() {
        let data = load_test_file::<f64, f64>("tests/rastrigin.txt");
        let linear = Rastrigin::<Dyn>::new();

        for test in data {
            let inp = OVector::<f64, Dyn>::from_vec(test.inp);

            assert_eq!(
                linear.fit(&inp).unwrap(),
                test.out
            );
        }
    }

    #[test]
    fn test_griewank() {
        let data = load_test_file::<f64, f64>("tests/griewank.txt");
        let griewank = Griewank::<Dyn>::new();

        for test in data {
            let inp = OVector::<f64, Dyn>::from_vec(test.inp);

            assert_eq!(
                griewank.fit(&inp).unwrap(),
                test.out
            );
        }
    }

    #[test]
    fn test_schwefel() {
        const MAX_LEN: usize = 10;

        let data = load_test_file::<f64, f64>("tests/schwefel.txt");
        let schwefel = Schwefel::new();

        for test in data {
            let filled = test.inp.iter()
                .chain(vec![0f64; MAX_LEN - test.inp.len()].iter())
                .map(|x| *x)
                .collect::<Vec<_>>();
            let inp = SVector::<f64, MAX_LEN>::from_vec(filled);

            assert_eq!(
                schwefel.fit(&inp).unwrap(),
                test.out
            );
        }
    }
}