use std::{convert::Infallible, rc::Rc};

use eoa_lib::{
    constraints::LowerThanConstraintFunction, fitness::FitnessFunction,
};
use nalgebra::SVector;

pub struct ArbitraryFitness<const SIZE: usize> {
    fun: Box<dyn Fn(SVector<f64, SIZE>) -> f64>,
}

impl<const SIZE: usize> ArbitraryFitness<SIZE> {
    pub fn new(fun: Box<dyn Fn(SVector<f64, SIZE>) -> f64>) -> Self {
        Self { fun }
    }

    pub fn zero() -> Self {
        Self {
            fun: Box::new(|_| 0.0),
        }
    }
}

impl<const SIZE: usize> FitnessFunction for ArbitraryFitness<SIZE> {
    type In = SVector<f64, SIZE>;
    type Out = f64;
    type Err = Infallible;

    fn fit(&self, inp: &Self::In) -> Result<Self::Out, Self::Err> {
        Ok((self.fun)(*inp))
    }
}

/// A constrained optimization problem with clear field definitions
pub struct ConstrainedProblem<const DIM: usize, const CONSTRAINTS: usize> {
    pub name: String,
    pub objective: ArbitraryFitness<DIM>,
    pub constraints: [LowerThanConstraintFunction<SVector<f64, DIM>, f64>; CONSTRAINTS],
    pub bounds: (SVector<f64, DIM>, SVector<f64, DIM>), // (min, max)
    pub optimal_value: f64,
    pub instantiate_fn: Option<Rc<dyn Fn() -> ConstrainedProblem<DIM, CONSTRAINTS>>>,
}

impl<const DIM: usize, const CONSTRAINTS: usize> ConstrainedProblem<DIM, CONSTRAINTS> {
    pub fn new(instantiate: Rc<dyn Fn() -> ConstrainedProblem<DIM, CONSTRAINTS>>) -> Self {
        let mut problem = instantiate();
        problem.instantiate_fn = Some(instantiate);
        problem
    }

    pub fn clone(&self) -> Self {
        Self::new(self.instantiate_fn.clone().unwrap())
    }
}

pub fn problem_g04() -> ConstrainedProblem<5, 6> {
    ConstrainedProblem::new(Rc::new(|| {
        ConstrainedProblem {
            name: "g04".to_string(),
            objective: ArbitraryFitness::new(Box::new(|vec| {
                5.3578547 * vec[2].powi(2)
                    + 0.8356891 * vec[0] * vec[4]
                    + 37.293239 * vec[0]
                    - 40792.141
            })),
            constraints: [
                LowerThanConstraintFunction::new(Box::new(|vec| {
                    85.334407 + 0.0056858 * vec[1] * vec[4]
                        + 0.0006262 * vec[0] * vec[3]
                        - 0.0022053 * vec[2] * vec[4]
                        - 92.0
                })),
                LowerThanConstraintFunction::new(Box::new(|vec| {
                    -85.334407
                        - 0.0056858 * vec[1] * vec[4]
                        - 0.0006262 * vec[0] * vec[3]
                        + 0.0022053 * vec[2] * vec[4]
                })),
                LowerThanConstraintFunction::new(Box::new(|vec| {
                    80.51249 + 0.0071317 * vec[1] * vec[4]
                        + 0.0029955 * vec[0] * vec[1]
                        + 0.0021813 * vec[2].powi(2)
                        - 110.0
                })),
                LowerThanConstraintFunction::new(Box::new(|vec| {
                    -80.51249
                        - 0.0071317 * vec[1] * vec[4]
                        - 0.0029955 * vec[0] * vec[1]
                        - 0.0021813 * vec[2].powi(2)
                        + 90.0
                })),
                LowerThanConstraintFunction::new(Box::new(|vec| {
                    9.300961 + 0.0047026 * vec[2] * vec[4]
                        + 0.0012547 * vec[0] * vec[2]
                        + 0.0019085 * vec[2] * vec[3]
                        - 25.0
                })),
                LowerThanConstraintFunction::new(Box::new(|vec| {
                    -9.300961
                        - 0.0047026 * vec[2] * vec[4]
                        - 0.0012547 * vec[0] * vec[2]
                        - 0.0019085 * vec[2] * vec[3]
                        + 20.0
                })),
            ],
            bounds: (
                SVector::<f64, 5>::from([78.0, 33.0, 27.0, 27.0, 27.0]), // min bounds
                SVector::<f64, 5>::from([102.0, 45.0, 45.0, 45.0, 45.0]), // max bounds
            ),
            optimal_value: -30665.53867178333,
            instantiate_fn: None,
        }
    }))
}

pub fn problem_g05() -> ConstrainedProblem<4, 5> {
    ConstrainedProblem::new(Rc::new(|| {
        ConstrainedProblem {
            name: "g05".to_string(),
            objective: ArbitraryFitness::new(Box::new(|vec| {
                3.0 * vec[0]
                    + 0.000001 * vec[0].powi(3)
                    + 2.0 * vec[1]
                    + (0.000002 / 3.0) * vec[1].powi(3)
            })),
            constraints: [
                LowerThanConstraintFunction::new(Box::new(|vec| -vec[3] + vec[2] - 0.55)),
                LowerThanConstraintFunction::new(Box::new(|vec| -vec[2] + vec[3] - 0.55)),
                LowerThanConstraintFunction::new(Box::new(|vec| {
                    1000.0 * (-vec[2] - 0.25).sin() + 1000.0 * (-vec[3] - 0.25).sin() - vec[0] + 894.8
                })),
                LowerThanConstraintFunction::new(Box::new(|vec| {
                    1000.0 * (vec[2] - 0.25).sin()
                        + 1000.0 * (vec[2] - vec[3] - 0.25).sin()
                        - vec[1]
                        + 894.8
                })),
                LowerThanConstraintFunction::new(Box::new(|vec| {
                    1000.0 * (vec[3] - 0.25).sin() + 1000.0 * (vec[3] - vec[2] - 0.25).sin() + 1294.8
                })),
            ],
            bounds: (
                SVector::<f64, 4>::from([0.0, 0.0, -0.55, -0.55]), // min bounds
                SVector::<f64, 4>::from([1200.0, 1200.0, 0.55, 0.55]), // max bounds
            ),
            optimal_value: 5126.4967140071,
            instantiate_fn: None,
        }
    }))
}

pub fn problem_g06() -> ConstrainedProblem<2, 2> {
    ConstrainedProblem::new(Rc::new(|| {
        ConstrainedProblem {
            name: "g06".to_string(),
            objective: ArbitraryFitness::new(Box::new(|vec| {
                (vec[0] - 10.0).powi(3) + (vec[1] - 20.0).powi(3)
            })),
            constraints: [
                LowerThanConstraintFunction::new(
                    Box::new(|vec| -(vec[0] - 5.0).powi(2) - (vec[1] - 5.0).powi(2) + 100.0),
                ),
                LowerThanConstraintFunction::new(
                    Box::new(|vec| (vec[0] - 6.0).powi(2) + (vec[1] - 5.0).powi(2) - 82.81),
                ),
            ],
            bounds: (
                SVector::<f64, 2>::new(0.0, 0.0),   // min bounds
                SVector::<f64, 2>::new(50.0, 50.0), // max bounds
            ),
            optimal_value: -6961.8137558015,
            instantiate_fn: None,
        }
    }))
}

pub fn problem_g08() -> ConstrainedProblem<2, 2> {
    ConstrainedProblem::new(Rc::new(|| ConstrainedProblem {
        name: "g08".to_string(),
        objective: ArbitraryFitness::new(Box::new(|vec| {
            let num = (2.0 * std::f64::consts::PI * vec[0]).sin().powi(3)
                * (2.0 * std::f64::consts::PI * vec[1]).sin();
            let den = vec[0].powi(3) * (vec[0] + vec[1]);
            -num / den
        })),
        constraints: [
            LowerThanConstraintFunction::new(Box::new(move |vec| {
                let x1 = vec[0];
                let x2 = vec[1];
                x1.powi(2) - x2 + 1.0
            })),
            LowerThanConstraintFunction::new(Box::new(move |vec| {
                let x1 = vec[0];
                let x2 = vec[1];
                1.0 - x1 + (x2 - 4.0).powi(2)
            })),
        ],
        bounds: (
            SVector::<f64, 2>::new(0.0, 0.0),   // min bounds
            SVector::<f64, 2>::new(10.0, 10.0), // max bounds
        ),
        optimal_value: -0.0958250414180359,
        instantiate_fn: None,
    }))
}

pub fn problem_g09() -> ConstrainedProblem<7, 4> {
    ConstrainedProblem::new(Rc::new(|| {
        ConstrainedProblem {
            name: "g09".to_string(),
            objective: ArbitraryFitness::new(Box::new(|vec| {
                (vec[0] - 10.0).powi(2)
                    + 5.0 * (vec[1] - 12.0).powi(2)
                    + vec[2].powi(4)
                    + 3.0 * (vec[3] - 11.0).powi(2)
                    + 10.0 * vec[4].powi(6)
                    + 7.0 * vec[5].powi(2)
                    + vec[6].powi(4)
                    - 4.0 * vec[5] * vec[6]
                    - 10.0 * vec[5]
                    - 8.0 * vec[6]
            })),
            constraints: [
                LowerThanConstraintFunction::new(Box::new(|vec| {
                    -127.0
                        + 2.0 * vec[0].powi(2)
                        + 3.0 * vec[1].powi(4)
                        + vec[2]
                        + 4.0 * vec[3].powi(2)
                        + 5.0 * vec[4]
                })),
                LowerThanConstraintFunction::new(Box::new(|vec| {
                    -282.0 + 7.0 * vec[0] + 3.0 * vec[1] + 10.0 * vec[2].powi(2) + vec[3] - vec[4]
                })),
                LowerThanConstraintFunction::new(Box::new(|vec| {
                    -196.0 + 23.0 * vec[0] + vec[1].powi(2) + 6.0 * vec[5].powi(2) - 8.0 * vec[6]
                })),
                LowerThanConstraintFunction::new(Box::new(|vec| {
                    4.0 * vec[0].powi(2) + vec[1].powi(2) - 3.0 * vec[0] * vec[1]
                        + 2.0 * vec[2].powi(2)
                        + 5.0 * vec[5]
                        - 11.0 * vec[6]
                })),
            ],
            bounds: (
                SVector::<f64, 7>::from([-10.0; 7]), // min bounds (all -10)
                SVector::<f64, 7>::from([10.0; 7]),  // max bounds (all 10)
            ),
            optimal_value: 680.6300573745,
            instantiate_fn: None,
        }
    }))
}

pub fn problem_g11(eps: f64) -> ConstrainedProblem<2, 1> {
    let problem = ConstrainedProblem::new(Rc::new(move || {
        ConstrainedProblem {
            name: "g11".to_string(),
            objective: ArbitraryFitness::new(Box::new(|vec| {
                // Minimize f(x) = x1^2 + (x2 - 1)^2
                vec[0].powi(2) + (vec[1] - 1.0).powi(2)
            })),
            constraints: [
                // Equality h(x) = x2 - x1^2 = 0
                // |h| - eps >= 0.0
                LowerThanConstraintFunction::new(Box::new(move |vec| {
                    let h = vec[1] - vec[0].powi(2);
                    h.abs() - eps
                })),
            ],
            bounds: (
                SVector::<f64, 2>::new(-50.0, -50.0), // min bounds
                SVector::<f64, 2>::new(50.0, 50.0),   // max bounds
            ),
            optimal_value: 0.7499, // Best known optimum
            instantiate_fn: None,
        }
    }));

    problem
}

pub fn problem_g21() -> ConstrainedProblem<7, 6> {
    ConstrainedProblem::new(Rc::new(|| ConstrainedProblem {
        name: "g21".to_string(),
        objective: ArbitraryFitness::new(Box::new(|vec| vec[0])),
        constraints: [
            LowerThanConstraintFunction::new(Box::new(|vec| {
                -vec[0] + 35.0 * vec[1].powf(0.6) + 35.0 * vec[2].powf(0.6)
            })),
            LowerThanConstraintFunction::new(Box::new(|vec| {
                -300.0 * vec[2]
                    + 7500.0 * vec[4]
                    - 7500.0 * vec[5]
                    - 25.0 * vec[3] * vec[4]
                    + 25.0 * vec[3] * vec[5]
                    + vec[2] * vec[3]
            })),
            LowerThanConstraintFunction::new(Box::new(|vec| {
                100.0 * vec[1]
                    + 155.365 * vec[3]
                    + 2500.0 * vec[6]
                    - vec[1] * vec[3]
                    - 25.0 * vec[3] * vec[6]
                    - 15536.5
            })),
            LowerThanConstraintFunction::new(Box::new(|vec| -vec[4] + (-vec[3] + 900.0).ln())),
            LowerThanConstraintFunction::new(Box::new(|vec| -vec[5] + (vec[3] + 300.0).ln())),
            LowerThanConstraintFunction::new(
                Box::new(|vec| -vec[6] + (-2.0 * vec[3] + 700.0).ln()),
            ),
        ],
        bounds: (
            // x1..x7 Min
            SVector::<f64, 7>::from([0.0, 0.0, 0.0, 100.0, 6.3, 5.9, 4.5]),
            // x1..x7 Max
            SVector::<f64, 7>::from([1000.0, 40.0, 40.0, 300.0, 6.7, 6.4, 6.25]),
        ),
        optimal_value: 193.724510070035,
        instantiate_fn: None,
    }))
}

pub fn problem_g24() -> ConstrainedProblem<2, 2> {
    ConstrainedProblem::new(Rc::new(|| ConstrainedProblem {
        name: "g24".to_string(),
        objective: ArbitraryFitness::new(Box::new(|vec| {
            // Minimize f(x) = -x1 - x2
            -vec[0] - vec[1]
        })),
        constraints: [
            // g1(x) = -2x1^4 + 8x1^3 - 8x1^2 + x2 - 2 <= 0
            LowerThanConstraintFunction::new(Box::new(|vec| {
                -2.0 * vec[0].powi(4)
                    + 8.0 * vec[0].powi(3)
                    - 8.0 * vec[0].powi(2)
                    + vec[1]
                    - 2.0
            })),
            // g2(x) = -4x1^4 + 32x1^3 - 88x1^2 + 96x1 + x2 - 36 <= 0
            LowerThanConstraintFunction::new(Box::new(|vec| {
                -4.0 * vec[0].powi(4)
                    + 32.0 * vec[0].powi(3)
                    - 88.0 * vec[0].powi(2)
                    + 96.0 * vec[0]
                    + vec[1]
                    - 36.0
            })),
        ],
        bounds: (
            SVector::<f64, 2>::new(0.0, 0.0), // min bounds
            SVector::<f64, 2>::new(3.0, 4.0), // max bounds
        ),
        optimal_value: -5.50801327159536, // Best known optimum
        instantiate_fn: None,
    }))
}

pub struct StochasticRankingConfig {
    pub population_size: usize,
    pub parents_count: usize,
    pub iterations: usize,
    pub n_param: usize,
    pub p_param: f64,
    pub mutation_std_dev: f64,
}

pub struct NsgaConfig {
    pub population_size: usize,
    pub parents_count: usize,
    pub iterations: usize,
    pub mutation_std_dev: f64,
}

#[cfg(test)]
mod tests {
    use super::*;
    use nalgebra::SVector;
    use eoa_lib::constraints::ConstraintFunction;

    #[test]
    fn test_g05_optimal_feasibility() {
        let problem = problem_g05();
        let optimal_solution = SVector::<f64, 4>::from([
            679.945148297028709,
            1026.06697600004691,
            0.118876369094410433,
            -0.39623348521517826,
        ]);

        let epsilon = 1e-4; // Tolerance for floating point comparisons

        for (i, constraint) in problem.constraints.iter().enumerate() {
            let evaluation = constraint.evaluate(&optimal_solution).unwrap();
            // For inequality constraints g(x) <= 0, we check if evaluation <= epsilon
            // For equality constraints h(x) = 0, which are transformed to |h(x)| - eps <= 0,
            // we check if evaluation <= epsilon
            assert!(
                evaluation <= epsilon,
                "Constraint {} (g05) failed: Expected evaluation <= {}, got {}",
                i + 1,
                epsilon,
                evaluation
            );
        }
    }
}