use factorio_core::{
    pathfield::PathField,
    prelude::*,
    visualize::{Color, Symbol, Visualization, Visualize, image_grid},
};
use factorio_graph::priority_queue::binary_heap::FastBinaryHeap;
use factorio_pathfinding::belt_finding::{self, conflict_avoidance::ConflictAvoidance};
use rand::{Rng, seq::IndexedRandom};
use serde::{Deserialize, Serialize};
use std::{sync::atomic::AtomicU32, time::Instant};

static OUTFILEINDEX: AtomicU32 = AtomicU32::new(0);

pub struct GeneticAlgorithm<'a> {
    problem: &'a Problem,
    population: Vec<PathLayout<'a>>,
    population_size: usize,
    population_keep: usize,
    population_new: usize,
}

impl<'a> GeneticAlgorithm<'a> {
    pub fn new<R: Rng + ?Sized>(
        problem: &'a Problem,
        population_size: usize,
        population_keep: usize,
        population_new: usize,
        rng: &mut R,
    ) -> GeneticAlgorithm<'a> {
        let mut population = Vec::new();

        let start = Instant::now();

        let mut count: usize = 0;

        while population.len() < population_size {
            count += 1;
            if let Some(p) = PathLayout::new(Layout::new(problem, rng)) {
                population.push(p);
            }
        }

        println!("Layouts accepted: {}/{}", population_size, count);

        population.sort_by_key(|p| p.score());

        println!(
            "Best score: {}. Time: {:.2}s",
            population[0].score(),
            start.elapsed().as_secs_f32()
        );
        population[0].print_visualization();

        GeneticAlgorithm {
            problem,
            population,
            population_size,
            population_keep,
            population_new,
        }
    }

    pub fn generation<R: Rng + ?Sized>(&mut self, rng: &mut R) {
        let start_new = Instant::now();
        for i in self.population_keep..(self.population_keep + self.population_new) {
            loop {
                if let Some(p) = PathLayout::new(Layout::new(self.problem, rng)) {
                    self.population[i] = p;
                    break;
                }
            }
        }

        let duration_new = start_new.elapsed();
        let start_mutate = Instant::now();

        for i in (self.population_keep + self.population_new)..self.population_size {
            let j = i - (self.population_keep + self.population_new);
            loop {
                if let Some(p) =
                    PathLayout::new(self.population[j % self.population_keep].layout.mutate(rng))
                {
                    self.population[i] = p;
                    break;
                }
            }
        }
        let duration_mutate = start_mutate.elapsed();

        self.population.sort_by_key(|p| p.score());
        println!(
            "Best score: {}. Time new: {:.2}s. Time mutate: {:.2}s",
            self.population[0].score(),
            duration_new.as_secs_f32(),
            duration_mutate.as_secs_f32()
        );
        self.population[0].print_visualization();
        let v: Vec<_> = self.population.iter().map(|p| p.visualize()).collect();
        let img = image_grid(&v, v[0].size().x as u32, v[0].size().y as u32, 5);
        let mut file = std::fs::File::create("generation.png").unwrap();
        img.write_to(&mut file, image::ImageFormat::Png).unwrap();
    }

    pub fn output_population(&self) {
        println!("Population:");
        for (i, p) in self.population.iter().enumerate() {
            println!("{i:3}: {}", p.score());
            p.print_visualization();
        }
    }
}

pub fn genetic_algorithm2<'a, R: Rng + ?Sized>(
    problem: &'a Problem,
    new_layouts: usize,
    mutation_timeout: usize,
    max_mutations: usize,
    rng: &'_ mut R,
) -> PathLayout<'a> {
    let mut m = (0..new_layouts)
        .map(|_| valid_path_layout(problem, rng))
        .min_by_key(|p| p.score())
        .unwrap();

    // m.print_visualization();

    let mut last_improvement = 0;
    let mut count = 0;

    while last_improvement < mutation_timeout && count < max_mutations {
        last_improvement += 1;
        count += 1;
        if let Some(p) = PathLayout::new(m.layout.mutate(rng)) {
            if p.score() < m.score() {
                m = p;
                // println!("Step: {count}");
                // m.print_visualization();
                last_improvement = 0;
            }
        }
    }
    m
}

pub fn valid_path_layout<'a, R: Rng + ?Sized>(
    problem: &'a Problem,
    rng: &'_ mut R,
) -> PathLayout<'a> {
    loop {
        if let Some(p) = PathLayout::new(Layout::new(problem, rng)) {
            return p;
        }
    }
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub(crate) struct MacroBlock {
    pub(crate) size: Position,
    pub(crate) input: Vec<Interface>,
    pub(crate) output: Vec<Interface>,
}

#[derive(Debug, Serialize, Deserialize, Clone, Copy)]
pub(crate) struct Interface {
    pub(crate) offset: Position,
    pub(crate) dir: Direction,
}

#[derive(Debug, Serialize, Deserialize, Clone, Copy)]
pub(crate) struct Connection {
    pub(crate) startblock: usize,
    pub(crate) startpoint: usize,
    pub(crate) endblock: usize,
    pub(crate) endpoint: usize,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct Problem {
    pub(crate) size: Position,
    pub(crate) blocks: Vec<MacroBlock>,
    pub(crate) connections: Vec<Connection>,
}

// #[derive(Debug, Clone, Copy)]
// pub struct BlockHandle(usize);

#[derive(Debug, Clone)]
pub struct Layout<'a> {
    pub(crate) problem: &'a Problem,
    pub(crate) blocks: Vec<Block>,
}

pub struct PathLayout<'a> {
    layout: Layout<'a>,
    paths: Vec<Vec<PathField>>,
    score: usize,
}

pub fn beltfinding_problem_from_layout(l: &Layout) -> belt_finding::Problem {
    let mut p = belt_finding::Problem::new(l.problem.size.x as usize, l.problem.size.y as usize);

    for b in &l.blocks {
        let aabb = b.get_aabb();

        p.set_blocked_range(
            aabb.min().x as usize,
            aabb.min().y as usize,
            aabb.max().x as usize,
            aabb.max().y as usize,
            true,
        );
    }

    for c in &l.problem.connections {
        let start_transform = l.blocks[c.startblock].block_to_world();
        let startpos = l.problem.blocks[c.startblock].output[c.startpoint]
            .offset
            .transform(start_transform);
        let startdir = l.problem.blocks[c.startblock].output[c.startpoint]
            .dir
            .transform(start_transform);
        let end_transform = l.blocks[c.endblock].block_to_world();
        let endpos = l.problem.blocks[c.endblock].input[c.endpoint]
            .offset
            .transform(end_transform);
        let enddir = l.problem.blocks[c.endblock].input[c.endpoint]
            .dir
            .transform(end_transform);
        p.add_connection(
            (startpos, startdir),
            (endpos.in_direction(&enddir, -1), enddir),
        );
    }

    p
}

impl<'a> PathLayout<'a> {
    pub fn new(layout: Layout<'a>) -> Option<PathLayout<'a>> {
        let mut p = beltfinding_problem_from_layout(&layout);

        if !p.find_path::<FastBinaryHeap<_>>() {
            return None;
        }

        let mut c = ConflictAvoidance::new(&p);

        let start = std::time::Instant::now();

        if !c.remove_all_conflicts(Some(std::time::Duration::from_secs(2))) {
            if start.elapsed().as_secs_f32() > 0.5 {
                // println!("Conflict avoidance: {:.2}", start.elapsed().as_secs_f32());
                // c.print_visualization();
                let file = std::fs::File::create(format!(
                    "out/{}.json",
                    OUTFILEINDEX.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
                ))
                .unwrap();
                serde_json::to_writer(file, &p).unwrap();
                // println!("Saved slow solve.");
            }
            return None;
        }

        let paths = c.get_paths().to_vec();

        let score = paths
            .iter()
            .map(|path| path.iter().skip(1).map(|p| p.cost()).sum::<usize>())
            .sum();

        Some(PathLayout {
            layout,
            paths,
            score,
        })
    }

    pub fn score(&self) -> usize {
        self.score
    }
}

impl Visualize for PathLayout<'_> {
    fn visualize(&self) -> factorio_core::visualize::Visualization {
        let mut v = self.layout.visualize();
        let offset = self.layout.blocks.len();

        for (i, path) in self.paths.iter().enumerate() {
            for p in &path[1..] {
                match p {
                    PathField::Belt { pos, dir } => {
                        v.add_symbol(
                            *pos,
                            Symbol::Arrow(*dir),
                            Some(Color::index(i + offset)),
                            None,
                        );
                    }
                    PathField::Underground { pos, dir, len } => {
                        v.add_symbol(
                            *pos,
                            Symbol::ArrowEnter(*dir),
                            Some(Color::index(i + offset)),
                            None,
                        );
                        v.add_symbol(
                            pos.in_direction(dir, *len as i32),
                            Symbol::ArrowEnter(*dir),
                            Some(Color::index(i + offset)),
                            None,
                        );
                    }
                }
            }
        }

        v
    }
}

impl Problem {
    pub fn new(size: Position) -> Self {
        Self {
            size,
            blocks: Vec::new(),
            connections: Vec::new(),
        }
    }

    // pub fn add_block(&mut self, size: Position) -> BlockHandle {
    //     self.blocks.push(Block {
    //         size,
    //         input: Vec::new(),
    //         output: Vec::new(),
    //     });

    //     BlockHandle(self.blocks.len() - 1)
    // }

    // pub fn add_connection(
    //     &mut self,
    //     starthandle: BlockHandle,
    //     startoffset: Position,
    //     endhandle: BlockHandle,
    //     endoffset: Position,
    // ) {
    //     let startinterface = self.blocks[starthandle.0].output.len();
    //     let endinterface = self.blocks[endhandle.0].input.len();

    //     self.blocks[starthandle.0].output.push(Interface {
    //         offset: startoffset,
    //         target: (endhandle.0, endinterface),
    //     });

    //     self.blocks[endhandle.0].input.push(Interface {
    //         offset: endoffset,
    //         target: (starthandle.0, startinterface),
    //     })
    // }
}

impl Layout<'_> {
    /// Create a new valid layout
    pub fn new<'a, R: Rng + ?Sized>(problem: &'a Problem, rng: &'_ mut R) -> Layout<'a> {
        let mut blocks = Vec::new();

        assert!(Self::place_block(problem, &mut blocks, rng));

        let mut l = Layout { problem, blocks };
        l.center();
        l
    }

    fn place_block<R: Rng + ?Sized>(
        problem: &'_ Problem,
        blocks: &mut Vec<Block>,
        rng: &'_ mut R,
    ) -> bool {
        if problem.blocks.len() == blocks.len() {
            return true;
        }

        let b = &problem.blocks[blocks.len()];
        for _ in 0..1000 {
            let dir = rng.random::<Direction>();

            let pos = match dir {
                Direction::Up => Position::new(
                    rng.random_range(0..=(problem.size.x - b.size.x)),
                    rng.random_range(0..=(problem.size.y - b.size.y)),
                ),
                Direction::Right => Position::new(
                    rng.random_range((b.size.y - 1)..problem.size.x),
                    rng.random_range(0..=(problem.size.y - b.size.x)),
                ),
                Direction::Down => Position::new(
                    rng.random_range((b.size.x - 1)..problem.size.x),
                    rng.random_range((b.size.y - 1)..problem.size.y),
                ),
                Direction::Left => Position::new(
                    rng.random_range(0..=(problem.size.x - b.size.y)),
                    rng.random_range((b.size.x - 1)..problem.size.y),
                ),
            };

            let current = Block::new(pos, dir, problem.blocks[blocks.len()].size);

            let current_aabb = current.get_aabb();

            if blocks
                .iter()
                .all(|b| !AABB::collision(b.get_aabb(), current_aabb))
            {
                blocks.push(current);

                if Self::place_block(problem, blocks, rng) {
                    return true;
                }

                blocks.pop();
            }
        }

        false
    }

    pub fn get_aabb(&self) -> AABB {
        self.blocks
            .iter()
            .map(|b| b.get_aabb())
            .reduce(AABB::combine)
            .expect("At least one block is required.")
    }

    pub fn center(&mut self) {
        let aabb = self.get_aabb();

        let rest = self.problem.size - aabb.size();

        let new_min = Position::new(rest.x / 2, rest.y / 2);

        let t = Transformation::new(Direction::Up, new_min - aabb.min());

        for b in &mut self.blocks {
            *b = b.transform(t);
        }
    }

    /// Mutate existing layout, creating a valid layout
    pub fn mutate<R: Rng + ?Sized>(&self, rng: &mut R) -> Self {
        let mut s = self.clone();

        #[allow(clippy::type_complexity)]
        let r: &[(&dyn Fn(&mut Layout, &mut R) -> bool, _)] = &[
            (&Self::mutate_replace::<R>, 30),
            (&Self::mutate_flip::<R>, 50),
            (&Self::mutate_jiggle::<R>, 160),
        ];

        loop {
            let p = r.choose_weighted(rng, |i| i.1).unwrap();

            if p.0(&mut s, rng) && rng.random_bool(0.5) {
                break;
            }
        }

        s.center();
        s
    }

    fn mutate_replace<R: Rng + ?Sized>(layout: &mut Layout, rng: &mut R) -> bool {
        let i = rng.random_range(0..layout.blocks.len());

        let dir = rng.random::<Direction>();

        let b = &layout.problem.blocks[i];

        let pos = match dir {
            Direction::Up => Position::new(
                rng.random_range(0..=(layout.problem.size.x - b.size.x)),
                rng.random_range(0..=(layout.problem.size.y - b.size.y)),
            ),
            Direction::Right => Position::new(
                rng.random_range((b.size.y - 1)..layout.problem.size.x),
                rng.random_range(0..=(layout.problem.size.y - b.size.x)),
            ),
            Direction::Down => Position::new(
                rng.random_range((b.size.x - 1)..layout.problem.size.x),
                rng.random_range((b.size.y - 1)..layout.problem.size.y),
            ),
            Direction::Left => Position::new(
                rng.random_range(0..=(layout.problem.size.x - b.size.y)),
                rng.random_range((b.size.x - 1)..layout.problem.size.y),
            ),
        };

        let current = Block::new(pos, dir, b.size);
        let current_aabb = current.get_aabb();

        if layout
            .blocks
            .iter()
            .enumerate()
            .all(|(j, b)| j == i || !AABB::collision(b.get_aabb(), current_aabb))
        {
            layout.blocks[i] = current;
            true
        } else {
            false
        }
    }

    fn mutate_flip<R: Rng + ?Sized>(layout: &mut Layout, rng: &mut R) -> bool {
        let i = rng.random_range(0..layout.blocks.len());
        let b = &mut layout.blocks[i];
        let block = &layout.problem.blocks[i];

        let new_pos = match b.dir() {
            Direction::Up => b.pos() + block.size - Position::new(1, 1),
            Direction::Right => b.pos() + Position::new(1 - block.size.y, block.size.x - 1),
            Direction::Down => b.pos() - block.size + Position::new(1, 1),
            Direction::Left => b.pos() + Position::new(block.size.y - 1, 1 - block.size.x),
        };
        let new_dir = b.dir().reverse();

        *b = Block::new(new_pos, new_dir, b.size());

        true
    }

    fn mutate_jiggle<R: Rng + ?Sized>(layout: &mut Layout, rng: &mut R) -> bool {
        let i = rng.random_range(0..layout.blocks.len());
        let dir = rng.random::<Direction>();
        let step = [(1, 10), (2, 5), (3, 5)]
            .choose_weighted(rng, |i| i.1)
            .unwrap()
            .0;
        // let step = 1;

        let b = &layout.blocks[i];

        let current = Block::new(b.pos().in_direction(&dir, step), b.dir(), b.size());
        let current_aabb = current.get_aabb();

        if current_aabb.min().x < 0
            || current_aabb.min().y < 0
            || current_aabb.max().x >= layout.problem.size.x
            || current_aabb.max().y >= layout.problem.size.y
        {
            return false;
        }

        if layout
            .blocks
            .iter()
            .enumerate()
            .all(|(j, b)| j == i || !AABB::collision(b.get_aabb(), current_aabb))
        {
            layout.blocks[i] = current;
            true
        } else {
            false
        }
    }
}

impl Visualize for Layout<'_> {
    fn visualize(&self) -> Visualization {
        let mut v = Visualization::new(self.problem.size);

        for (i, (b, mb)) in self
            .blocks
            .iter()
            .zip(self.problem.blocks.iter())
            .enumerate()
        {
            let c = Color::index(i);

            let aabb = b.get_aabb();

            for x in aabb.min().x..=aabb.max().x {
                for y in aabb.min().y..=aabb.max().y {
                    v.add_symbol(Position::new(x, y), Symbol::Block, Some(c), None);
                }
            }

            v.add_symbol(b.pos(), Symbol::Char('X'), Some(c), None);

            let transform = b.block_to_world();

            for input in &mb.input {
                v.add_symbol(
                    input.offset.transform(transform),
                    Symbol::Char('i'),
                    Some(c),
                    None,
                );
            }

            for output in &mb.output {
                v.add_symbol(
                    output.offset.transform(transform),
                    Symbol::Char('o'),
                    Some(c),
                    None,
                );
            }
        }

        v
    }
}