use std::convert::Infallible;

use eoa_lib::{
    bounded::BoundedOVector, comparison::MinimizingOperator, constraints::{stochastic_ranking_evolution_algorithm, ConstraintFunction, LowerThanConstraintFunction}, crossover::{ArithmeticCrossover, BoundedCrossover, BoundedCrossoverStrategy}, evolution::EvolutionResult, fitness::FitnessFunction, initializer::{Initializer, RandomInitializer}, pairing::AdjacentPairing, perturbation::{BoundedPerturbation, BoundedPerturbationStrategy, MutationPerturbation, RandomDistributionPerturbation}, population::Population, replacement::{BestReplacement, GenerationalReplacement}, selection::{BestSelection, TournamentSelection}
};
use nalgebra::SVector;
use rand::RngCore;
use rand_distr::Normal;

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
        }
    }
}

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 CONSTRAINTS: usize> {
    pub objective: ArbitraryFitness<2>,
    pub constraints: [LowerThanConstraintFunction<SVector<f64, 2>, f64>; CONSTRAINTS],
    pub bounds: (SVector<f64, 2>, SVector<f64, 2>), // (min, max)
    pub optimal_value: f64,
}

fn problem_g06() -> ConstrainedProblem<2> {
    ConstrainedProblem {
        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(100.0, 100.0)  // max bounds
        ),
        optimal_value: -6961.8137558015,
    }
}

fn problem_g08() -> ConstrainedProblem<2> {
    ConstrainedProblem {
        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,
    }
}

pub fn problem_g11(eps: f64) -> ConstrainedProblem<2> {
    ConstrainedProblem {
        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
            // Transformed 1: h - eps >= 0  =>  -(h - eps) <= 0
            LowerThanConstraintFunction::new(
                Box::new(move |vec| {
                    let h = vec[1] - vec[0].powi(2);
                    h - eps
                })
            ),
            // Transformed 2: eps - h >= 0  =>  -(eps - h) <= 0
            LowerThanConstraintFunction::new(
                Box::new(move |vec| {
                    let h = vec[1] - vec[0].powi(2);
                    eps - h
                })
            ),
        ],
        bounds: (
            SVector::<f64, 2>::new(-1.0, -1.0),  // min bounds
            SVector::<f64, 2>::new(1.0, 1.0)     // max bounds
        ),
        optimal_value: 0.7499, // Best known optimum
    }
}

pub fn problem_g24() -> ConstrainedProblem<2> {
    ConstrainedProblem {
        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
    }
}

/// Solve a constrained optimization problem using stochastic ranking
///
/// Returns the evolution result with feasible fractions for each iteration
pub fn solve_with_stochastic_ranking<const CONSTRAINTS: usize>(
    mut problem: ConstrainedProblem<CONSTRAINTS>,
    population_size: usize,
    parents_count: usize,
    iterations: usize,
    stochastic_params: (usize, f64), // (N, p) for stochastic ranking
    mutation_std_dev: f64,
    rng: &mut dyn RngCore,
) -> Result<(ConstrainedProblem<CONSTRAINTS>, EvolutionResult<SVector<f64, 2>, f64>, Vec<f64>), Box<dyn std::error::Error>> {
    // Create initial population
    let initializer = RandomInitializer::new(
        Box::new(BoundedOVector::new(problem.bounds.0, problem.bounds.1))
    );
    let initial_population = Population::from_vec(
        initializer.initialize(nalgebra::Const::<2>, population_size, rng)
    );

    // Setup components as specified
    // let mut selection = TournamentSelection::new(5, 0.9);
    let mut selection = TournamentSelection::new(5, 0.95);
    let mut replacement = GenerationalReplacement;
    let mut pairing = AdjacentPairing::new();
    let mut crossover = BoundedCrossover::<nalgebra::Const<2>, 2, _>::new(
        ArithmeticCrossover::new(),
        problem.bounds.0,
        problem.bounds.1,
        BoundedCrossoverStrategy::Trim
    );

    // Setup bounded random distribution perturbation with Normal distribution
    let normal_perturbation = RandomDistributionPerturbation::<2, Normal<f64>>::normal(mutation_std_dev)?;
    let perturbation = BoundedPerturbation::new(
        normal_perturbation,
        problem.bounds.0,
        problem.bounds.1,
        BoundedPerturbationStrategy::Trim
    );
    let mut mutation = MutationPerturbation::new(Box::new(perturbation), 0.1);

    // Constraint weights (initially set to 1.0)
    let constraint_weights = [1.0; CONSTRAINTS];

    let (N, p) = stochastic_params;
    let better_than = MinimizingOperator::new();

    // Convert constraint array references
    let constraint_refs = problem.constraints.iter().collect::<Vec<_>>().try_into()
        .map_err(|_| "Failed to convert constraint references")?;

    stochastic_ranking_evolution_algorithm(
        initial_population,
        parents_count,
        N,
        p,
        &mut problem.objective,
        constraint_refs,
        constraint_weights,
        &mut pairing,
        &mut selection,
        &mut crossover,
        &mut mutation,
        &mut replacement,
        &better_than,
        iterations,
        rng,
    ).map(|(x, y)| (problem, x, y))
}

fn main() {
    let mut rng = rand::rng();

    // Example usage with problem G06
    let problem = problem_g11(0.01);

    let population_size = 500;
    let result = solve_with_stochastic_ranking(
        problem,
        population_size,      // population_size
        population_size / 2,      // parents_count
        1000,     // iterations
        (population_size * 2, 0.2), // (N, p) for stochastic ranking
        0.001,     // mutation_std_dev
        &mut rng,
    );

    match result {
        Ok((problem, evolution_result, feasible_fractions)) => {
            if let Some(best_candidate) = evolution_result.best_candidate {
                println!("{:?}", feasible_fractions);
                println!("Best solution found:");
                println!("  Chromosome: {:?}", best_candidate.chromosome);
                println!("  Fitness: {}", best_candidate.evaluation);
                println!("  Iterations: {}", evolution_result.iterations);
                println!("  Evaluations: {}", evolution_result.evaluations);
                println!("  Final feasible fraction: {:.2}%", feasible_fractions.last().unwrap_or(&0.0) * 100.0);

                // Sanity check: verify the best candidate is feasible
                let best_chromosome = &best_candidate.chromosome;

                println!("\nFeasibility check for best solution:");
                for (i, constraint) in problem.constraints.iter().enumerate() {
                    match constraint.evaluate(best_chromosome) {
                        Ok(value) => {
                            let is_feasible = value <= 0.0;
                            println!("  Constraint {}: {} ({})", i+1, value, if is_feasible { "FEASIBLE" } else { "INFEASIBLE" });
                        }
                        Err(e) => {
                            println!("  Constraint {}: Error evaluating - {}", i+1, e);
                        }
                    }
                }
            } else {
                println!("Could not find any feasible solution!")
            }
        }
        Err(e) => {
            println!("Error: {}", e);
        }
    }
}