use std::error::Error;
use std::fmt::Debug;
use crate::binary_string::{BinaryString, BinaryStringConversionError};
use crate::evolutionary_strategy::{EvolutionaryStrategy, IdentityStrategy};
use crate::fitness::FitnessFunction;
use crate::terminating::TerminatingCondition;
use crate::perturbation::PerturbationOperator;
use crate::comparison::BetterThanOperator;
use full_palette::BROWN;
use nalgebra::allocator::Allocator;
use nalgebra::{DefaultAllocator, Dim, SVector, U2};
use plotters::prelude::*;

// Functions
#[derive(Debug, Clone, PartialEq)]
pub struct LocalSearchCandidate<TInput, TResult>
{
    pub fit: TResult,
    pub pos: TInput,
    pub cycle: usize
}

impl<TInput, TResult> LocalSearchCandidate<TInput, TResult> {
    pub fn new(fit: TResult, pos: TInput, cycle: usize) -> Self {
        Self {
            fit,
            pos,
            cycle
        }
    }
}

#[derive(Debug, Clone, PartialEq)]
pub struct LocalSearchStats<TInput, TResult> {
    stats: Vec<LocalSearchCandidate<TInput, TResult>>
}

impl<D, TResult> LocalSearchStats<BinaryString<D>, TResult>
where
    D: Dim,
    DefaultAllocator: Allocator<D>
{
    pub fn to_real(self, min: &SVector<f64, 2>, max: &SVector<f64, 2>) -> Result<LocalSearchStats<SVector<f64, 2>, TResult>, BinaryStringConversionError> {
        Ok(LocalSearchStats::<SVector<f64, 2>, TResult>::from_vec(
            self.stats
                .into_iter()
                .map(|candidate|
                     Ok(LocalSearchCandidate {
                         fit: candidate.fit,
                         pos: candidate.pos.to_real::<U2>(&min, &max)?,
                         cycle: candidate.cycle,
                     }))
                .collect::<Result<Vec<_>, _>>()?))
    }
}

impl<TInput, TResult> LocalSearchStats<TInput, TResult> {
    fn from_vec(stats: Vec<LocalSearchCandidate<TInput, TResult>>)  -> Self {
        Self {
            stats
        }
    }

    pub fn new() -> Self {
        Self {
            stats: vec![]
        }
    }

    pub fn candidates(&self) -> &Vec<LocalSearchCandidate<TInput, TResult>> {
        &self.stats
    }

    pub fn cycles(&self) -> usize {
        self.best().cycle
    }

    pub fn best(&self) -> &LocalSearchCandidate<TInput, TResult> {
        self.stats.last().expect("No candidates...")
    }

    pub fn append(&mut self, candidate: LocalSearchCandidate<TInput, TResult>) {
        self.stats.push(candidate)
    }
}

#[derive(Debug, Clone, PartialEq)]
pub struct LocalSearchResult<TInput, TResult>
    where TResult: Clone
{
    pub best_candidate: LocalSearchCandidate<TInput, TResult>,

    // How many cycles there were
    pub cycles: usize,

    // statistics
    pub stats: LocalSearchStats<TInput, TResult>
}

pub fn local_search_first_improving<
        TInput, TResult, TFit, TTerminatingCondition, TPerturbationOperator, TBetterThanOperator>(
    fit: &TFit,
    terminating_condition: &mut TTerminatingCondition,
    perturbation_operator: &mut TPerturbationOperator,
    better_than_operator: &TBetterThanOperator,
    initial: &TInput
) -> Result<LocalSearchResult<TInput, TResult>, Box<dyn Error>>
where
    TResult: Clone,
    TInput: Clone,
    TFit: FitnessFunction<In = TInput, Out = TResult>,
    TTerminatingCondition: TerminatingCondition<TInput, TResult>,
    TPerturbationOperator: PerturbationOperator<Chromosome = TInput>,
    TBetterThanOperator: BetterThanOperator<TResult>,
{
    local_search_first_improving_evolving(
        fit,
        terminating_condition,
        perturbation_operator,
        better_than_operator,
        &mut IdentityStrategy,
        initial
    )
}

pub fn local_search_first_improving_evolving<
        TInput, TResult, TFit, TTerminatingCondition, TPerturbationOperator, TBetterThanOperator, TEvolutionaryStrategy>(
    fit: &TFit,
    terminating_condition: &mut TTerminatingCondition,
    perturbation_operator: &mut TPerturbationOperator,
    better_than_operator: &TBetterThanOperator,
    evolutionary_strategy: &mut TEvolutionaryStrategy,
    initial: &TInput
) -> Result<LocalSearchResult<TInput, TResult>, Box<dyn Error>>
where
    TResult: Clone,
    TInput: Clone,
    TFit: FitnessFunction<In = TInput, Out = TResult>,
    TTerminatingCondition: TerminatingCondition<TInput, TResult>,
    TPerturbationOperator: PerturbationOperator<Chromosome = TInput>,
    TEvolutionaryStrategy: EvolutionaryStrategy<TResult, TPerturbationOperator>,
    TBetterThanOperator: BetterThanOperator<TResult>
{
    let mut best_candidate = LocalSearchCandidate {
        pos: initial.clone(),
        fit: fit.fit(&initial)?,
        cycle: 0
    };

    let mut stats = LocalSearchStats::new();
    let mut cycle: usize = 0;

    while !terminating_condition.should_terminate(&best_candidate, &stats, cycle) {
        let perturbed = perturbation_operator.perturb(&best_candidate.pos);
        let perturbed_fit = fit.fit(&perturbed)?;

        // Minimize
        let better = if better_than_operator.better_than(&perturbed_fit, &best_candidate.fit) {
            best_candidate = LocalSearchCandidate {
                pos: perturbed.clone(),
                fit: perturbed_fit,
                cycle
            };

            stats.append(best_candidate.clone());

            true
        } else {
            false
        };

        evolutionary_strategy.step(
            perturbation_operator,
            better,
            &stats)?;

        cycle += 1;
    }

    Ok(LocalSearchResult {
        best_candidate,
        stats,
        cycles: cycle
    })
}

pub fn plot_fitness_evolution(file: &str, stats: LocalSearchStats<SVector<f64, 2>, f64>) -> Result<(), Box<dyn std::error::Error>> {
    let root = BitMapBackend::new(file, (1920, 768)).into_drawing_area();
    root.fill(&WHITE)?;

    let (left, right) = root.split_horizontally(root.relative_to_width(0.50));

    let mut chart = ChartBuilder::on(&left)
        .caption("Best candidate evaluation", ("sans-serif", 25).into_font())
        .margin(5)
        .x_label_area_size(30)
        .y_label_area_size(30)
        .build_cartesian_2d(0..stats.cycles(), stats.best().fit..stats.candidates().first().unwrap().fit)?;

    chart.configure_mesh()
        .x_desc("Cycle")
        .y_desc("Evaluation")
        .draw()?;

    chart
        .draw_series(LineSeries::new(
            stats.candidates().iter().map(|candidate| (candidate.cycle, candidate.fit)),
            &RED,
        ))?
        .label("Linear fitness")
        .legend(|(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], &RED));

    chart
        .configure_series_labels()
        .background_style(&WHITE.mix(0.8))
        .border_style(&BLACK)
        .draw()?;

    let mut chart = ChartBuilder::on(&right)
        .caption("Locations of best candidates so far", ("sans-serif", 25).into_font())
        .margin(5)
        .x_label_area_size(50)
        .y_label_area_size(50)
        // TODO
        .build_cartesian_2d(-10.0..10.0f64, -10.0..10.0f64)?;

    chart.configure_mesh()
        .x_desc("X coordinate")
        .y_desc("Y coordinate")
        .draw()?;

    let mut points = stats.candidates().iter()
        .map(|candidate| (candidate.pos.x, candidate.pos.y));

    chart
        .draw_series(LineSeries::new(
            points.clone(),
            &BLACK,
        ))?;

    chart
        .draw_series(PointSeries::of_element(
            points.clone(),
            3,
            &BLACK,
            &| c, s, st| {
                return EmptyElement::at(c)
                    + Circle::new((0,0),s,st.filled());
            },
        ))?
        .label("Candidates")
        .legend(|(x, y)| Circle::new((x, y), 3, BLACK.filled()));

    // TODO:
    let start = points.next().unwrap();
    let best = points.last().unwrap();

    chart
        .draw_series(std::iter::once(Circle::new(
            start, 4, BROWN.filled())
        ))?
        .label("Start")
        .legend(|(x, y)| Circle::new((x, y), 3, BROWN.filled()));

    chart
        .draw_series(std::iter::once(Circle::new(
            best, 4, GREEN.filled())
        ))?
        .label("Best")
        .legend(|(x, y)| Circle::new((x, y), 3, GREEN.filled()));

    chart
        .draw_series(std::iter::once(Text::new(
            "Start",
            start,
            ("sans-serif", 20).into_font(),
        )))?;

    chart
        .configure_series_labels()
        .background_style(&WHITE.mix(0.8))
        .border_style(&BLACK)
        .draw()?;

    root.present()?;
    Ok(())
}

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

    use crate::{binary_string::BinaryString, bounded::BoundedOVector, comparison::MinimizingOperator, evolutionary_strategy::{IdentityStrategy, OneToFiveStrategy}, fitness::{one_max::OneMax, real::Linear, rosenbrock::Rosenbrock, sphere::Sphere, BinaryFitnessWrapper}, initializer::{Initializer, RandomInitializer}, local_search::{local_search_first_improving, local_search_first_improving_evolving, plot_fitness_evolution}, perturbation::{BinaryStringBitPerturbation, BoundedPerturbation, BoundedPerturbationStrategy, PatternPerturbation, RandomDistributionPerturbation}, terminating::{AndTerminatingConditions, CloserThanTerminatingCondition, EqualTerminatingCondition, NoBetterForCyclesTerminatingCondition}};

    #[test]
    fn test_local_search_sphere_binary() {
        let optimum = BinaryString::<U10>::new(vec![0, 0, 1, 0, 0,
                                             0, 0, 1, 0, 0]);
        let min = SVector::<f64, 2>::from_element(0.0);
        let max = SVector::<f64, 2>::from_element(31.0);
        let optimum_real = optimum.to_real(&min, &max).unwrap();
        let sphere = Sphere::new(optimum_real);
        let sphere_wrapped = BinaryFitnessWrapper::new(sphere, min.clone(), max.clone());

        let result = local_search_first_improving(
            &sphere_wrapped,
            &mut
                AndTerminatingConditions::new(
                    vec![
                        &mut EqualTerminatingCondition::new_remembered(optimum.clone()),
                        &mut NoBetterForCyclesTerminatingCondition::new(100)
                    ]
                ),
            &mut BinaryStringBitPerturbation::new(0.3),
            &MinimizingOperator::new(),
            &BinaryString::new(vec![1; 10]),
        ).unwrap();

        println!("{:?}", result);

        assert_eq!(
            result.best_candidate.fit,
            0.0
        );

        assert_eq!(
            result.best_candidate.pos,
            optimum
        );

        plot_fitness_evolution("test_results/sphere_binary.png", result.stats.to_real(&min, &max).unwrap()).unwrap();
    }

    #[test]
    fn test_local_search_sphere() {
        let optimum = SVector::<f64, 2>::repeat(4.0);
        let sphere = Sphere::new(optimum);

        let result = local_search_first_improving_evolving(
            &sphere,
            &mut
                AndTerminatingConditions::new(
                    vec![
                        &mut CloserThanTerminatingCondition::new_remembered(optimum.clone(), 0.1),
                        &mut NoBetterForCyclesTerminatingCondition::new(100)
                    ]
                ),
            &mut RandomDistributionPerturbation::normal(0.3).unwrap(),
            &MinimizingOperator::new(),
            &mut IdentityStrategy,
            &SVector::<f64, 2>::repeat(-5.0),
        ).unwrap();

        println!("{:?}", result);

        assert!(
            result.best_candidate.fit <= 0.1
        );

        assert!(
            (result.best_candidate.pos - optimum).magnitude() <= 0.1
        );

        plot_fitness_evolution("test_results/sphere.png", result.stats).unwrap();
    }

    #[test]
    fn test_local_search_one_max() {
        let one_max = OneMax::<U10>::new();
        let optimum = BinaryString::<U10>::new(vec![0; 10]);

        let result = local_search_first_improving(
            &one_max,
            &mut
                AndTerminatingConditions::new(
                    vec![
                        &mut EqualTerminatingCondition::new_remembered(optimum.clone()),
                        &mut NoBetterForCyclesTerminatingCondition::new(100)
                    ]
                ),
            &mut BinaryStringBitPerturbation::new(0.3),
            &MinimizingOperator::new(),
            &BinaryString::<U10>::new(vec![1; 10]),
        ).unwrap();

        println!("{:?}", result);

        assert_eq!(
            result.best_candidate.fit,
            0
        );

        assert_eq!(
            result.best_candidate.pos,
            optimum
        );
    }

    // #[test]
    // fn test_local_search_labs() {
    //     let labs = LABS::new();
    //     let optimum = BinaryString::new(vec![0; 20]);

    //     let result = local_search_first_improving(
    //         &labs,
    //         &mut
    //             AndTerminatingConditions::new(
    //                 vec![
    //                     // &mut EqualTerminatingCondition::new_remembered(optimum.clone()),
    //                     &mut NoBetterForCyclesTerminatingCondition::new(1000)
    //                 ]
    //             ),
    //         &mut BinaryStringBitPerturbation::new(0.1),
    //         &MinimizingOperator::new(),
    //         &BinaryString::new(vec![0; 20]),
    //     ).unwrap();

    //     println!("{:?}", result);

    //     assert_eq!(
    //         result.best_candidate.fit,
    //         0
    //     );

    //     assert_eq!(
    //         result.best_candidate.pos,
    //         optimum
    //     );
    // }

    #[test]
    fn test_local_search_rosenbrock() {
        let rosenbrock = Rosenbrock::new();
        let optimum = BinaryString::<U10>::new(vec![1, 0, 0, 0, 1, 1, 0, 0, 0, 1]);

        let min = SVector::<f64, 2>::from_element(-16.0);
        let max = SVector::<f64, 2>::from_element(15.0);
        let rosenbrock_wrapped = BinaryFitnessWrapper::new(rosenbrock, min, max);

        let result = local_search_first_improving(
            &rosenbrock_wrapped,
            &mut
                AndTerminatingConditions::new(
                    vec![
                        &mut EqualTerminatingCondition::new_remembered(optimum.clone()),
                        &mut NoBetterForCyclesTerminatingCondition::new(100)
                    ]
                ),
            &mut BinaryStringBitPerturbation::new(0.1),
            &MinimizingOperator::new(),
            &BinaryString::new(vec![0; 10]),
        ).unwrap();

        println!("{:?}", result);

        assert_eq!(
            result.best_candidate.fit,
            0.0
        );

        assert_eq!(
            result.best_candidate.pos,
            optimum
        );

        plot_fitness_evolution("test_results/rosenbrock.png", result.stats.to_real(&min, &max).unwrap()).unwrap();
    }

    #[test]
    fn test_local_search_linear() {
        let optimum = SVector::<f64, 2>::from_vec(vec![-10.0, 10.0]);
        let max = SVector::<f64, 2>::from_vec(vec![10.0, 10.0]);
        let min = -SVector::<f64, 2>::from_vec(vec![10.0, 10.0]);

        let mut initializer =
            RandomInitializer::<U2, OVector<f64, U2>>::new(Box::new(BoundedOVector::<U2>::new(min, max)));

        let linear = Linear::new(7.0, SVector::<f64, 2>::from_vec(vec![0.2, -0.5]));

        let result = local_search_first_improving(
            &linear,
            &mut
                AndTerminatingConditions::new(
                    vec![
                        &mut EqualTerminatingCondition::new_remembered(optimum.clone()),
                        &mut NoBetterForCyclesTerminatingCondition::new(100)
                    ]
                ),
            &mut BoundedPerturbation::new(
                RandomDistributionPerturbation::normal(0.5).unwrap(),
                min,
                max,
                BoundedPerturbationStrategy::Retry(10)),
            &MinimizingOperator::new(),
            &initializer.initialize_single(U2),
        ).unwrap();

        println!("{:?}", result);

        assert_eq!(
            result.best_candidate.fit,
            -0.0
        );

        assert_eq!(
            result.best_candidate.pos,
            optimum
        );

        plot_fitness_evolution("test_results/linear.png", result.stats).unwrap();
    }

    #[test]
    fn test_local_search_linear_pattern() {
        let optimum = SVector::<f64, 2>::from_vec(vec![-10.0, 10.0]);
        let max = SVector::<f64, 2>::from_vec(vec![10.0, 10.0]);
        let min = -SVector::<f64, 2>::from_vec(vec![10.0, 10.0]);

        let linear = Linear::new(7.0, SVector::<f64, 2>::from_vec(vec![0.2, -0.5]));

        let result = local_search_first_improving(
            &linear,
            &mut
                AndTerminatingConditions::new(
                    vec![
                        &mut EqualTerminatingCondition::new_remembered(optimum.clone()),
                        &mut NoBetterForCyclesTerminatingCondition::new(100)
                    ]
                ),
            &mut BoundedPerturbation::new(
                PatternPerturbation::new(0.5),
                min,
                max,
                BoundedPerturbationStrategy::Retry(10)),
            &MinimizingOperator::new(),
            &SVector::<f64, 2>::zeros(),
        ).unwrap();

        println!("{:?}", result);

        assert_eq!(
            result.best_candidate.fit,
            -0.0
        );

        assert_eq!(
            result.best_candidate.pos,
            optimum
        );

        plot_fitness_evolution("test_results/linear_pattern.png", result.stats).unwrap();
    }

    #[test]
    fn test_local_search_linear_onetofive() {
        let optimum = SVector::<f64, 2>::from_vec(vec![-10.0, 10.0]);
        let max = SVector::<f64, 2>::from_vec(vec![10.0, 10.0]);
        let min = -SVector::<f64, 2>::from_vec(vec![10.0, 10.0]);

        let linear = Linear::new(7.0, SVector::<f64, 2>::from_vec(vec![0.2, -0.5]));

        let result = local_search_first_improving_evolving(
            &linear,
            &mut
                AndTerminatingConditions::new(
                    vec![
                        &mut EqualTerminatingCondition::new_remembered(optimum.clone()),
                        &mut NoBetterForCyclesTerminatingCondition::new(100)
                    ]
                ),
            &mut BoundedPerturbation::new(
                RandomDistributionPerturbation::normal(0.5).unwrap(),
                min,
                max,
                BoundedPerturbationStrategy::Retry(10)),
            &MinimizingOperator::new(),
            &mut OneToFiveStrategy,
            &SVector::<f64, 2>::zeros(),
        ).unwrap();

        println!("{:?}", result);

        assert_eq!(
            result.best_candidate.fit,
            -0.0
        );

        assert_eq!(
            result.best_candidate.pos,
            optimum
        );

        plot_fitness_evolution("test_results/linear_onetofive.png", result.stats).unwrap();
    }
}