use crate::util::maps::*;
use itertools::Itertools;
use std::collections::HashMap;

pub fn day16() {
    let input: World<char> = World::from_string(&crate::input(16));
    let maze: World<bool> = input.map(|&c| c == '.' || c == 'S' || c == 'E');

    let start: Transform = input
        .enumerate()
        .find(|(_, _, &c)| c == 'S')
        .map(|(x, y, _)| Transform::new(x, y, Direction::East, &maze))
        .unwrap();
    let end: (usize, usize) = input
        .enumerate()
        .find(|(_, _, &c)| c == 'E')
        .map(|(x, y, _)| (x, y))
        .unwrap();

    let (cheapest_cost, paths) = cheapest_paths(start, end);

    let tiles_crossed: usize = paths
        .iter()
        .flatten()
        .sorted_by_key(|p| (p.x, p.y))
        .dedup_by(|a, b| a.x == b.x && a.y == b.y)
        .count();

    println!("Found {} shortest paths.", paths.len());
    println!("Cheapest Path to End: {}", cheapest_cost);
    println!("Tiles Crossed: {}", tiles_crossed);
}

fn cheapest_paths(start: Transform, end: (usize, usize)) -> (u64, Vec<Vec<Transform>>) {
    let paths: Vec<(Vec<Transform>, u64)> = all_paths(start, end)
        .into_iter()
        // path -> (path, cost)
        .map(|path| {
            let cost = path.windows(2).into_iter().fold(0, |acc, w| {
                acc + if w[0].direction == w[1].direction {
                    1
                } else {
                    1000
                }
            });
            (path, cost)
        })
        // keep only the paths tied for cheapest
        .min_set_by_key(|(_, cost)| *cost);
    let cheapest_cost = paths[0].1;
    // (path, cost) -> path
    let paths = paths.into_iter().map(|(path, _)| path).collect();
    (cheapest_cost, paths)
}

fn all_paths(start: Transform, end: (usize, usize)) -> Vec<Vec<Transform>> {
    let costs = maze_costs(start);

    fn collect_paths<'a, 'b>(
        costs: &'b HashMap<Transform<'a>, u64>,
        p: Transform<'a>,
    ) -> Vec<Vec<Transform<'a>>> {
        let cost_p = costs[&p];

        // cost 0 means this is the start
        if cost_p == 0 {
            return vec![vec![p]];
        };
        // otherwise we look at all ways that we might get to p
        p.comes_from()
            .into_iter()
            // no cost stored means that x is a silly place to ever come from (a dead end)
            .filter(|x| costs.contains_key(x))
            // if cost_x < cost_p, x most be closer to the start than p is
            // otherwise, we don't want to check x (that would be anti-progress)
            .filter(|x| costs[x] < cost_p)
            .flat_map(|x| {
                // now we recur, getting all paths from start to x
                let mut paths = collect_paths(costs, x);
                // and add ourself to the paths, then send the paths up the chain
                paths.iter_mut().for_each(|path| path.push(p));
                paths
            })
            .collect()
    }

    // every facing direction at the end may have a different cost
    let endpoints = costs.keys().filter(|t| t.x == end.0 && t.y == end.1);
    endpoints
        .into_iter()
        .flat_map(|&endpoint| collect_paths(&costs, endpoint))
        .collect()
}

fn maze_costs(start: Transform) -> HashMap<Transform, u64> {
    fn calculate_costs<'a, 'b>(costs: &'b mut HashMap<Transform<'a>, u64>, p: Transform<'a>) {
        // cost to any x adjacent to p is the cost to get to p + the cost to travel p -> x
        // and if that's cheaper than any previous way we've found, repeat for each x as the new p
        for (travel_cost, adjacent) in p.goes_to() {
            let total_cost = costs[&p] + travel_cost;
            let existing_cost = costs.get(&adjacent).copied().unwrap_or(u64::MAX);
            if total_cost < existing_cost {
                costs.insert(adjacent, total_cost);
                calculate_costs(costs, adjacent);
            }
        }
    }

    let mut costs = HashMap::new();
    costs.insert(start, 0);
    calculate_costs(&mut costs, start);
    costs
}

#[derive(Copy, Clone)]
struct Transform<'a> {
    x: usize,
    y: usize,
    direction: Direction,
    map: &'a World<bool>,
}

impl<'a> Transform<'a> {
    fn new(x: usize, y: usize, direction: Direction, map: &'a World<bool>) -> Self {
        Self {
            x,
            y,
            direction,
            map,
        }
    }

    fn goes_to(self) -> Vec<(u64, Transform<'a>)> {
        let mut ret: Vec<(u64, Transform)> = vec![
            self.direction.clockwise(),
            self.direction.counterclockwise(),
        ]
        .into_iter()
        .filter(|d| self.map.travel_get(self.x, self.y, *d) == Some(true))
        .map(|d| (1000, Self::new(self.x, self.y, d, self.map)))
        .collect();
        if let Some((nx, ny, true)) = self.map.travel_and_get(self.x, self.y, self.direction) {
            ret.push((1, Self::new(nx, ny, self.direction, self.map)));
        }
        ret
    }

    fn comes_from(self) -> Vec<Transform<'a>> {
        let mut ret: Vec<Transform<'a>> = vec![
            self.direction.clockwise(),
            self.direction.counterclockwise(),
        ]
        .into_iter()
        .map(|d| Self::new(self.x, self.y, d, self.map))
        .collect();
        let backward = self
            .map
            .travel_and_get(self.x, self.y, self.direction.opposite());
        if let Some((nx, ny, true)) = backward {
            ret.push(Self::new(nx, ny, self.direction, self.map));
        }
        ret
    }
}

impl<'a> PartialEq for Transform<'a> {
    fn eq(&self, other: &Self) -> bool {
        std::ptr::eq(self.map, other.map)
            && self.x == other.x
            && self.y == other.y
            && self.direction == other.direction
    }
}

impl<'a> Eq for Transform<'a> {}

impl<'a> std::hash::Hash for Transform<'a> {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.x.hash(state);
        self.y.hash(state);
        self.direction.hash(state);
    }
}