use tracing::{Level, field::Empty, info, span, trace, warn};

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> {
        assert!(self.population_new + self.population_keep <= self.population_size);
        let _complete_span = span!(Level::TRACE, "genetic_algorithm_v1").entered();

        let mut population = Vec::new();

        {
            let _initial_generation_span = span!(Level::TRACE, "initial generation").entered();
            for i in 0..self.population_size {
                let _layout_span = span!(Level::TRACE, "layout", i).entered();
                loop {
                    if let Some(l) = self.valid_layout.layout(input, pathfinder, rng) {
                        let score = score(input, &l);
                        population.push((l, score));
                        break;
                    }
                    warn!("Unable to generate valid layout");
                }
            }
        }
        population.sort_by_key(|(_, s)| *s);

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

        for g in 0..self.generations {
            let generation_span =
                span!(Level::TRACE, "generation", generation = g, score = Empty).entered();

            {
                let _mutate_span = span!(Level::TRACE, "mutate").entered();
                for i in 0..(self.population_size - self.population_keep - self.population_new) {
                    let _layout_span = span!(Level::TRACE, "layout", i).entered();
                    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;
                            }
                        }
                        trace!("unsuccesfull mutation");
                    }
                }
            }

            {
                let _new_span = span!(Level::TRACE, "new").entered();
                for i in 0..self.population_new {
                    let _layout_span = span!(Level::TRACE, "layout", i).entered();
                    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;
                        }
                        warn!("Unable to generate valid layout");
                    }
                }
            }

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

        Some(best_result.0)
    }
}