use crate::{
    LayoutResult, Layouter,
    misc::{mutate, path_input_from_blocks_positions, score},
    valid_layout::ValidLayout,
};

pub struct GeneticAlgorithm {
    pub mutation_retries: usize,
    pub population_size: usize,
    pub population_keep: usize,
    pub population_new: usize,
    pub generations: usize,
    pub valid_layout: ValidLayout,
}

impl Layouter for GeneticAlgorithm {
    fn layout<R: rand::Rng, P: factorio_pathfinding::Pathfinder>(
        &self,
        input: &crate::LayoutInput,
        pathfinder: &P,
        rng: &mut R,
    ) -> Option<crate::LayoutResult> {
        let mut population = Vec::new();

        for _ in 0..self.population_size {
            loop {
                if let Some(l) = self.valid_layout.layout(input, pathfinder, rng) {
                    let score = score(input, &l);
                    population.push((l, score));
                    break;
                }
            }
        }
        population.sort_by_key(|(_, s)| *s);

        let mut best_result = population[0].clone();

        for g in 0..self.generations {
            for i in 0..(self.population_size - self.population_keep - self.population_new) {
                loop {
                    let parent = &population[i % self.population_keep].0;
                    if let Some(blocks) = mutate(input, parent, rng) {
                        let (connections, map) =
                            path_input_from_blocks_positions(input, parent.size, &blocks);

                        if let Some(paths) =
                            pathfinder.find_paths(factorio_pathfinding::PathInput {
                                connections: &connections,
                                map,
                            })
                        {
                            let r = LayoutResult {
                                positions: blocks,
                                path_result: paths,
                                size: parent.size,
                            };
                            let score = score(input, &r);
                            population.push((r, score));
                            break;
                        }
                    }
                }
            }

            for i in 0..self.population_new {
                loop {
                    if let Some(l) = self.valid_layout.layout(input, pathfinder, rng) {
                        let score = score(input, &l);
                        population[self.population_size - self.population_new + i] = (l, score);
                        break;
                    }
                }
            }

            population.sort_by_key(|(_, s)| *s);
            if population[0].1 < best_result.1 {
                best_result = population[0].clone();
            }
            println!("completed generation {g}\nscore: {}", population[0].1);
        }

        Some(best_result.0)
    }
}