use itertools::Itertools;
use ndarray::prelude::*;
use petgraph::algo::all_simple_paths;
use petgraph::prelude::*;
use std::collections::{HashMap, HashSet};
use std::fmt::{Debug, Display, Write};
use std::fs::File;
use std::io::{BufRead, BufReader, Lines};
use std::time::Instant;

// BOILERPLATE
type InputIter = Lines<BufReader<File>>;

fn get_input() -> InputIter {
    let f = File::open("input").unwrap();
    let br = BufReader::new(f);
    br.lines()
}

fn main() {
    let start = Instant::now();
    let ans1 = problem1(get_input());
    let duration = start.elapsed();
    println!("Problem 1 solution: {} [{}s]", ans1, duration.as_secs_f64());

    let start = Instant::now();
    let ans2 = problem2(get_input());
    let duration = start.elapsed();
    println!("Problem 2 solution: {} [{}s]", ans2, duration.as_secs_f64());
}

// PARSE

#[derive(Clone)]
struct Node {
    c: char,
    pos: Position,
}

impl Display for Node {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "({},{})", self.pos.0, self.pos.1)
    }
}

impl Debug for Node {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "({},{})", self.pos.0, self.pos.1)
    }
}

type Position = (usize, usize);

#[derive(Clone)]
struct ForestMap {
    map: Array2<char>,
    indexes: HashMap<Position, NodeIndex>,
    graph: StableDiGraph<Node, u64>,
    start: Position,
    end: Position,
}

const ADJACENCIES: [(isize, isize); 4] = [(-1, 0), (1, 0), (0, -1), (0, 1)];

fn offset_pos(map: &Array2<char>, pos: Position, ofs: (isize, isize)) -> Option<Position> {
    let new_pos = (pos.0 as isize + ofs.0, pos.1 as isize + ofs.1);
    if new_pos.0 >= 0
        && new_pos.0 < map.len_of(Axis(0)) as isize
        && new_pos.1 >= 0
        && new_pos.1 < map.len_of(Axis(1)) as isize
    {
        Some((new_pos.0 as usize, new_pos.1 as usize))
    } else {
        None
    }
}

fn adjacent_to(map: &Array2<char>, pos: Position) -> Vec<Position> {
    ADJACENCIES
        .iter()
        .filter_map(|ofs| offset_pos(map, pos, *ofs))
        .collect()
}

impl<T: BufRead> From<Lines<T>> for ForestMap {
    fn from(lines: Lines<T>) -> Self {
        let rows = lines.map(|line| line.unwrap().chars().collect_vec()).collect_vec();
        let map = Array::from_shape_vec([rows[0].len(), rows.len()], rows.into_iter().flatten().collect_vec())
            .unwrap()
            .reversed_axes();
        let start = (map.slice(s![.., 0]).iter().position(|c| *c == '.').unwrap(), 0);
        let end = (
            map.slice(s![.., map.len_of(Axis(1)) - 1])
                .iter()
                .position(|c| *c == '.')
                .unwrap(),
            map.len_of(Axis(1)) - 1,
        );

        let mut graph = StableGraph::default();
        let mut indexes = HashMap::new();
        for (pos, c) in map.indexed_iter() {
            if *c != '#' {
                indexes.insert(pos, graph.add_node(Node { c: *c, pos }));
            }
        }

        Self {
            map,
            start,
            end,
            graph,
            indexes,
        }
    }
}

impl ForestMap {
    fn build_graph(&mut self) {
        for (pos, c) in self.map.indexed_iter() {
            match c {
                '#' => continue,
                '.' => {
                    adjacent_to(&self.map, pos).iter().for_each(|adj| {
                        if self.indexes.contains_key(&adj) {
                            self.graph.add_edge(self.indexes[&pos], self.indexes[adj], 1);
                        }
                    });
                }
                '^' => {
                    if let Some(adj) = offset_pos(&self.map, pos, (0, -1)) {
                        if self.indexes.contains_key(&adj) {
                            self.graph.add_edge(self.indexes[&pos], self.indexes[&adj], 1);
                        }
                    }
                }
                '>' => {
                    if let Some(adj) = offset_pos(&self.map, pos, (1, 0)) {
                        if self.indexes.contains_key(&adj) {
                            self.graph.add_edge(self.indexes[&pos], self.indexes[&adj], 1);
                        }
                    }
                }
                'v' => {
                    if let Some(adj) = offset_pos(&self.map, pos, (0, 1)) {
                        if self.indexes.contains_key(&adj) {
                            self.graph.add_edge(self.indexes[&pos], self.indexes[&adj], 1);
                        }
                    }
                }
                '<' => {
                    if let Some(adj) = offset_pos(&self.map, pos, (-1, 0)) {
                        if self.indexes.contains_key(&adj) {
                            self.graph.add_edge(self.indexes[&pos], self.indexes[&adj], 1);
                        }
                    }
                }
                c => panic!("invalid map character {}", c),
            }
        }
    }

    fn build_graph2(&mut self) {
        for (pos, c) in self.map.indexed_iter() {
            match c {
                '#' => continue,
                '.' | '^' | '>' | 'v' | '<' => {
                    adjacent_to(&self.map, pos).iter().for_each(|adj| {
                        if self.indexes.contains_key(&adj) {
                            self.graph.add_edge(self.indexes[&pos], self.indexes[adj], 1);
                        }
                    });
                }
                c => panic!("invalid map character {}", c),
            }
        }
    }

    // Cull nodes that don't change the topology of the graph and combine their cost
    fn simplify_graph(&mut self) {
        let mut idxs: Vec<_> = self
            .graph
            .neighbors(self.indexes[&self.start])
            .map(|idx| (self.indexes[&self.start], idx))
            .collect();
        let mut visited = HashSet::from([self.indexes[&self.start]]);
        while let Some((last_idx, cur_idx)) = idxs.pop() {
            if !visited.insert(cur_idx) {
                continue;
            }
            let our_neighbors = self.graph.neighbors(cur_idx).collect_vec();

            // if we have exactly 2 neighbours, then one is where we came from, and we can shortcut this node with a
            // pair of new edges A <-> C and break the existing 4 edges between them
            if our_neighbors.len() == 2 {
                let next_idx = our_neighbors.iter().find(|n| **n != last_idx).unwrap();

                // remove the 4 existing edges
                // careful of order of operations, as removing edges invalidates edge indexes
                let forward_cost = self
                    .graph
                    .remove_edge(self.graph.find_edge(cur_idx, *next_idx).unwrap())
                    .unwrap();
                let last_forward_cost = self
                    .graph
                    .remove_edge(self.graph.find_edge(last_idx, cur_idx).unwrap())
                    .unwrap();
                let backward_cost = self
                    .graph
                    .remove_edge(self.graph.find_edge(cur_idx, last_idx).unwrap())
                    .unwrap();
                let next_backward_cost = self
                    .graph
                    .remove_edge(self.graph.find_edge(*next_idx, cur_idx).unwrap())
                    .unwrap();
                let new_forward_cost = forward_cost + last_forward_cost;
                let new_backward_cost = backward_cost + next_backward_cost;

                // add edge from last to next
                self.graph.add_edge(last_idx, *next_idx, new_forward_cost);
                self.graph.add_edge(*next_idx, last_idx, new_backward_cost);

                self.graph.remove_node(cur_idx);
                // push the next node
                idxs.push((last_idx, *next_idx));
            } else {
                // don't do anything about nodes with > 2 edges, just push them onto the stack, if there are some
                idxs.append(
                    &mut self
                        .graph
                        .neighbors(cur_idx)
                        .into_iter()
                        .map(|next_idx| (cur_idx, next_idx))
                        .collect(),
                );
            }
        }
    }
}

impl Debug for ForestMap {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        for y in 0..self.map.len_of(Axis(1)) {
            for c in self.map.index_axis(Axis(1), y) {
                f.write_char(*c)?;
            }
            writeln!(f)?;
        }
        Ok(())
    }
}

// PROBLEM 1 solution

fn problem1<T: BufRead>(input: Lines<T>) -> u64 {
    let mut map = ForestMap::from(input);
    map.build_graph();
    // println!("{:?}", map);
    let paths = all_simple_paths::<Vec<_>, _>(&map.graph, map.indexes[&map.start], map.indexes[&map.end], 0, None)
        .collect_vec();
    let longest = paths.iter().max_by_key(|path| path.len()).unwrap();

    longest.len() as u64 - 1
}

fn calc_path_length(map: &ForestMap, path: &Vec<NodeIndex>) -> u64 {
    path.iter().tuple_windows().fold(0, |accum, (prev, next)| {
        accum + map.graph[map.graph.find_edge(*prev, *next).unwrap()]
    })
}

// PROBLEM 2 solution
fn problem2<T: BufRead>(input: Lines<T>) -> u64 {
    let mut map = ForestMap::from(input);
    map.build_graph2();
    map.simplify_graph();

    let paths = all_simple_paths::<Vec<_>, _>(&map.graph, map.indexes[&map.start], map.indexes[&map.end], 0, None)
        .collect_vec();
    let longest = paths.iter().max_by_key(|path| calc_path_length(&map, &path)).unwrap();
    longest.iter().tuple_windows().fold(0, |accum, (prev, next)| {
        accum + map.graph[map.graph.find_edge(*prev, *next).unwrap()]
    })
}

#[cfg(test)]
mod tests {
    use crate::*;
    use std::io::Cursor;

    const EXAMPLE: &str = &"#.#####################
#.......#########...###
#######.#########.#.###
###.....#.>.>.###.#.###
###v#####.#v#.###.#.###
###.>...#.#.#.....#...#
###v###.#.#.#########.#
###...#.#.#.......#...#
#####.#.#.#######.#.###
#.....#.#.#.......#...#
#.#####.#.#.#########v#
#.#...#...#...###...>.#
#.#.#v#######v###.###v#
#...#.>.#...>.>.#.###.#
#####v#.#.###v#.#.###.#
#.....#...#...#.#.#...#
#.#########.###.#.#.###
#...###...#...#...#.###
###.###.#.###v#####v###
#...#...#.#.>.>.#.>.###
#.###.###.#.###.#.#v###
#.....###...###...#...#
#####################.#";

    #[test]
    fn problem1_example() {
        let c = Cursor::new(EXAMPLE);
        assert_eq!(problem1(c.lines()), 94);
    }

    #[test]
    fn problem2_example() {
        let c = Cursor::new(EXAMPLE);
        assert_eq!(problem2(c.lines()), 154);
    }
}