factorio_blueprint/factorio-pathfinding/src/belt_finding/mod.rs

use std::marker::PhantomData;

use crate::Connection;
use crate::Map as _;
use crate::SinglePathInput;
use crate::SinglePathfinder;
use crate::examples::HashMapMap;
use factorio_core::misc::Map;
use factorio_core::{prelude::*, visualize::Visualize};
use factorio_graph::priority_queue::PriorityQueueByKey;
use factorio_graph::wheighted_graph::WheightedGraph;
use factorio_graph::wheighted_graph::shortest_path::QueueObject;
use factorio_graph::wheighted_graph::shortest_path::a_star;
use serde::{Deserialize, Serialize};
use tracing::Level;
use tracing::span;

pub mod brute_force;

pub mod conflict_avoidance;

pub struct ConflictAvoidance<P> {
    pub timeout: Option<std::time::Duration>,
    pub priority_queue: PhantomData<P>,
}

impl<P> SinglePathfinder for ConflictAvoidance<P>
where
    P: PriorityQueueByKey<QueueObject<(Position, Direction)>>,
{
    fn find_paths<M: crate::Map>(
        &self,
        input: SinglePathInput<M>,
    ) -> Option<Vec<Vec<factorio_core::pathfield::PathField>>> {
        let _span = span!(Level::TRACE, "find_paths").entered();

        let (pos, size) = input.map.area();
        let mut p = Problem::new(size.x as usize, size.y as usize);

        for x in 0..size.x {
            for y in 0..size.y {
                if !input.map.get_position(Position::new(x - pos.x, y - pos.y)) {
                    p.set_blocked(x as usize, y as usize, true);
                }
            }
        }

        for i in input.connections {
            p.add_connection(
                (i.start_pos, i.start_dir),
                (i.end_pos.in_direction(&i.end_dir, -1), i.end_dir),
            );
        }

        // p.print_visualization();

        if p.find_path::<P>() {
            let mut c = conflict_avoidance::ConflictAvoidance::new(&p);
            // c.print_visualization();

            if c.remove_all_conflicts(self.timeout) {
                // c.print_visualization();
                Some(c.get_paths().to_vec())
            } else {
                None
            }
        } else {
            None
        }
    }
}

#[derive(Debug, Default, Serialize, Deserialize, Clone, Copy)]
pub struct Field {
    pub blocked: bool,
    weight: i64,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct Problem {
    map: Map<Field>,
    start: Vec<(Position, Direction)>,
    end: Vec<(Position, Direction)>,
    path: Vec<Vec<(Position, Direction)>>,
}

impl Problem {
    pub fn new(width: usize, height: usize) -> Self {
        Self {
            map: Map::new(width, height),
            start: Vec::new(),
            end: Vec::new(),
            path: Vec::new(),
        }
    }

    pub fn add_connection(&mut self, start: (Position, Direction), end: (Position, Direction)) {
        self.start.push(start);
        self.end.push(end);
        self.path.push(Vec::new());
    }

    pub fn set_blocked(&mut self, x: usize, y: usize, blocked: bool) {
        self.map.get_mut(x, y).blocked = blocked;
    }

    pub fn set_blocked_range(&mut self, x1: usize, y1: usize, x2: usize, y2: usize, blocked: bool) {
        for x in x1..=x2 {
            for y in y1..=y2 {
                self.set_blocked(x, y, blocked);
            }
        }
    }

    fn calculate_wheights(&mut self, without: usize) {
        for x in 0..self.map.width {
            for y in 0..self.map.height {
                let mut weight = 0;

                if self.map.get(x, y).blocked {
                    weight += 10000;
                }

                self.map.get_mut(x, y).weight = weight;
            }
        }

        for (i, path) in self.path.iter().enumerate() {
            if i != without && !path.is_empty() {
                for p in &path[1..] {
                    let weight = 1000;
                    let x = p.0.x as usize;
                    let y = p.0.y as usize;

                    self.map.get_mut(x, y).weight += weight;
                }
            }
        }

        // for p in &self.start {
        //     self.map.get_mut(p.0.x as usize, p.0.y as usize).weight = f64::INFINITY;
        // }

        // for p in &self.end {
        //     self.map.get_mut(p.0.x as usize, p.0.y as usize).weight = f64::INFINITY;
        // }
    }
}

impl From<(Vec<Connection>, HashMapMap)> for Problem {
    fn from((input, map): (Vec<Connection>, HashMapMap)) -> Self {
        let (pos, size) = map.area();
        let mut p = Problem::new(size.x as usize, size.y as usize);

        for x in 0..size.x {
            for y in 0..size.y {
                if !map.get_position(Position::new(x - pos.x, y - pos.y)) {
                    p.set_blocked(x as usize, y as usize, true);
                }
            }
        }

        for i in input {
            p.add_connection((i.start_pos, i.start_dir), (i.end_pos, i.end_dir));
        }

        p
    }
}

impl Visualize for Problem {
    fn visualize(&self) -> factorio_core::visualize::Visualization {
        let mut v = factorio_core::visualize::Visualization::new(Position::new(
            self.map.width as i32,
            self.map.height as i32,
        ));

        for x in 0..self.map.width {
            for y in 0..self.map.height {
                if self.map.get(x, y).blocked {
                    v.add_symbol(
                        Position::new(x as i32, y as i32),
                        factorio_core::visualize::Symbol::Block,
                        Some(factorio_core::visualize::Color::white()),
                        None,
                    )
                }
            }
        }

        for (i, (p, _d)) in self.start.iter().enumerate() {
            v.add_symbol(
                *p,
                factorio_core::visualize::Symbol::Char('S'),
                Some(factorio_core::visualize::Color::index(i)),
                None,
            )
        }

        for (i, (p, _d)) in self.end.iter().enumerate() {
            v.add_symbol(
                *p,
                factorio_core::visualize::Symbol::Char('T'),
                Some(factorio_core::visualize::Color::index(i)),
                None,
            )
        }

        for (i, p) in self.path.iter().enumerate() {
            for (pos, dir) in p {
                v.add_symbol(
                    *pos,
                    factorio_core::visualize::Symbol::Arrow(*dir),
                    Some(factorio_core::visualize::Color::index(i)),
                    None,
                )
            }
        }

        v
    }
}

struct MapInternal<'a> {
    map: &'a Map<Field>,
    end: (Position, Direction),
}

impl WheightedGraph for MapInternal<'_> {
    type Node = (Position, Direction);

    fn edge(
        &self,
        node: &(Position, Direction),
        num: usize,
    ) -> Option<((Position, Direction), i64)> {
        let next = node.0.in_direction(&node.1, 1);
        next.in_range(
            &Position::new(0, 0),
            &Position::new(
                self.map.width as PositionType,
                self.map.height as PositionType,
            ),
        )?;
        if self.map.get(next.x as usize, next.y as usize).blocked && self.end != (next, node.1) {
            return None;
        }

        let penalty = self.map.get(next.x as usize, next.y as usize).weight;
        match num {
            0 => Some(((next, node.1), 3 + penalty)),
            1 => Some(((next, node.1.counter_clockwise()), 4 + penalty)),
            2 => Some(((next, node.1.clockwise()), 4 + penalty)),
            _ => {
                let mut count = 2;
                for l in 2..=6 {
                    let n = node.0.in_direction(&node.1, l);
                    n.in_range(
                        &Position::new(0, 0),
                        &Position::new(
                            self.map.width as PositionType,
                            self.map.height as PositionType,
                        ),
                    )?;
                    if !self.map.get(n.x as usize, n.y as usize).blocked {
                        count += 1;
                        if count == num {
                            let penalty = penalty + self.map.get(n.x as usize, n.y as usize).weight;
                            return Some(((n, node.1), 35 + penalty));
                        }
                    }
                }
                None
            }
        }
    }
}

impl Problem {
    pub fn find_path<P>(&mut self) -> bool
    where
        P: PriorityQueueByKey<QueueObject<(Position, Direction)>>,
    {
        let _span = span!(Level::TRACE, "find_path").entered();
        for i in 0..self.start.len() {
            self.calculate_wheights(i);
            let m = MapInternal {
                map: &self.map,
                end: self.end[i],
            };
            let p = {
                // dijkstra::<MapInternal, FastBinaryHeap<_>>(&m, self.start[i], self.end[i])
                a_star::<MapInternal, P, _, _>(
                    &m,
                    self.start[i],
                    |&n| n == self.end[i],
                    |&(p, _)| {
                        3 * (PositionType::abs_diff(p.x, self.end[i].0.x)
                            + PositionType::abs_diff(p.y, self.end[i].0.y))
                            as i64
                    },
                )
            };
            if let Some(p) = p {
                self.path[i] = p;
            } else {
                return false;
            }
        }
        true
    }
}