use std::collections::VecDeque;
use std::fs;

use itertools::Itertools;

fn to_grid(lines: Vec<&str>) -> Vec<Vec<char>> {
    let mut grid: Vec<Vec<char>> = Vec::new();

    for line in lines {
        let mut row = Vec::new();
        for ch in line.chars() {
            row.push(ch);
        }
        grid.push(row);
    }

    grid
}

fn find_goals(grid: &Vec<Vec<char>>) -> Vec<(usize, usize)> {
    let mut result = Vec::new();

    // The goal locations are represented by digits in the grid
    for (y, row) in grid.iter().enumerate() {
        for (x, ch) in row.iter().enumerate() {
            if ch.is_ascii_digit() {
                result.push((x, y));
            }
        }
    }

    result
}

fn goal_distances(grid: &Vec<Vec<char>>, goals: &Vec<(usize, usize)>) -> Vec<Vec<usize>> {
    let mut dists = Vec::new();
    for _i in 0..goals.len() {
        dists.push([0].repeat(goals.len()));
    }

    // For each goal location, find the distance to each other goal
    for (i, &g) in goals.iter().enumerate() {
        let mut visited = vec![g];
        let mut queue = VecDeque::from([(g, 0)]);

        // Do a breadth-first search to find the shortest paths
        while queue.len() > 0 {
            let (loc, depth) = queue.pop_front().unwrap();

            // Have we reached another goal? If so, which one?
            let oj = goals.iter().position(|&xy| xy == loc);
            if oj.is_some() {
                let j = oj.unwrap();
                dists[i][j] = depth;
            }

            // Which adjacent locations are accessible?
            let mut adjs = Vec::new();
            let u = (loc.0, loc.1 - 1); // up
            let d = (loc.0, loc.1 + 1); // down
            let l = (loc.0 - 1, loc.1); // left
            let r = (loc.0 + 1, loc.1); // right
            if grid[u.1][u.0] != '#' {
                adjs.push(u);
            }
            if grid[d.1][d.0] != '#' {
                adjs.push(d);
            }
            if grid[l.1][l.0] != '#' {
                adjs.push(l);
            }
            if grid[r.1][r.0] != '#' {
                adjs.push(r);
            }

            // Standard breadth-first search stuff
            for adj in adjs {
                if !visited.contains(&adj) {
                    visited.push(adj);
                    queue.push_back((adj, depth + 1));
                }
            }
        }
    }

    dists
}

fn solve_puzzle(grid: &Vec<Vec<char>>, go_back: bool) -> usize {
    let goals = find_goals(&grid);
    let dists = goal_distances(&grid, &goals);

    // This is the travelling salesman problem, so we must try all paths.
    // We just brute-force the total distance for every order of the goals,
    // and only keep the result if the path started at "0" as it should.
    let mut shortest = 13371337;
    for p in (0..goals.len()).permutations(goals.len()) {
        let mut total = 0;
        for i in 1..p.len() {
            total += dists[p[i - 1]][p[i]];
        }
        // Part 2: go back to the start too
        if go_back {
            total += dists[p[p.len() - 1]][p[0]];
        }

        // Is this the shortest path we've seen that start at "0"?
        let first = goals[p[0]];
        if grid[first.1][first.0] == '0' && total < shortest {
            shortest = total;
        }
    }

    shortest
}

fn main() {
    // Read layout from input text file
    let input = fs::read_to_string("input.txt").unwrap();
    let lines = input.lines().collect();
    let grid = to_grid(lines);

    // Part 1 gives 430 for me
    println!("Part 1 solution: {}", solve_puzzle(&grid, false));

    // Part 2 gives 700 for me
    println!("Part 2 solution: {}", solve_puzzle(&grid, true));
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn part1_example1() {
        let lines = vec![
            "###########",
            "#0.1.....2#",
            "#.#######.#",
            "#4.......3#",
            "###########",
        ];
        let grid = to_grid(lines);
        assert_eq!(solve_puzzle(&grid, false), 14);
    }
}