use std::collections::BTreeMap;
use std::fs::read_to_string;
use std::time::Instant;

fn find_path(layout: &[&[u8]], min_consecutive: u8, max_consecutive: u8,
             out_path: &mut Vec<(i32, i32)>, out_history: &mut Vec<(i32, i32)>) -> i32 {
    #[derive(Copy, Clone, PartialEq)]
    #[repr(u8)]
    enum Direction {
        Left = 0,
        Right = 1,
        Up = 2,
        Down = 3,
    }
    const ALL_DIRECTIONS: [Direction; 4] = [Direction::Left, Direction::Right, Direction::Up, Direction::Down];
    const OPPOSITE_DIRECTION: [Direction; 4] = [Direction::Right, Direction::Left, Direction::Down, Direction::Up];
    const DIRECTION_OFFSET: [(i32, i32); 4] = [(-1, 0), (1, 0), (0, -1), (0, 1)];
    struct Node {
        x: i32,
        y: i32,
        cost: i32,
        dir: Direction,
        consecutive: u8,
    }
    let width = layout[0].len();
    let height = layout.len();
    let index_multiplier = max_consecutive as usize * 4;
    // Evaluation score for visited nodes, -1 is unvisited
    // Node is treated as a different one unless direction and consecutive counter are the same
    let mut scores: Vec<i32> = vec![-1; width * height * index_multiplier];
    let mut previous: Vec<i32> = vec![-1; width * height * index_multiplier];
    let mut queue: BTreeMap<i32, Node> = BTreeMap::new();
    let get_node_index = |node: &Node| -> usize {
        (node.y as usize * width + node.x as usize) * index_multiplier
            + node.dir as usize * max_consecutive as usize + node.consecutive as usize - 1
    };
    let target_x = width as i32 - 1;
    let target_y = height as i32 - 1;
    let get_heuristic = |node: &Node| -> i32 {
        // Heuristic is manhattan distance with addition of correction due to forced turns
        // This correction seems to give minor but nonzero improvement of performance
        let dx = target_x - node.x;
        let dy = target_y - node.y;
        let x_correction = dx / max_consecutive as i32 * min_consecutive as i32;
        let y_correction = dy / max_consecutive as i32 * min_consecutive as i32;
        dx + x_correction + dy + y_correction
    };
    let mut current_node = Node { x: 0, y: 0, cost: 0, dir: Direction::Right, consecutive: 0 };
    let mut current_index: i32 = -1;
    'outer: loop {
        out_history.push((current_node.x, current_node.y));
        for dir in ALL_DIRECTIONS {
            // Can't turn 180
            if OPPOSITE_DIRECTION[current_node.dir as usize] == dir {
                continue;
            }
            // Can't go in one direction less than min_consecutive tiles, 0 is only for start
            if current_node.consecutive > 0 && current_node.consecutive < min_consecutive
                && current_node.dir != dir {
                continue;
            }
            // Compute new node consecutive counter
            let mut new_node = Node {
                x: current_node.x,
                y: current_node.y,
                cost: current_node.cost,
                dir: dir,
                consecutive: if current_node.dir == dir { current_node.consecutive + 1 } else { 1 },
            };
            // Can't go in one direction more than max_consecutive tiles
            if new_node.consecutive > max_consecutive as u8 {
                continue;
            }
            // Calculating new node position and bounds checking
            let (offset_x, offset_y) = DIRECTION_OFFSET[dir as usize];
            new_node.x += offset_x;
            new_node.y += offset_y;
            if new_node.x < 0 || new_node.y < 0 || new_node.x as usize >= width || new_node.y as usize >= height {
                continue;
            }
            // If the node hasn't been visited yet
            let node_index = get_node_index(&new_node);
            if scores[node_index] < 0 {
                let new_tile_cost = (layout[new_node.y as usize][new_node.x as usize] - b'0') as i32;
                new_node.cost += new_tile_cost;
                let new_score = new_node.cost + get_heuristic(&new_node);
                scores[node_index] = new_score;
                previous[node_index] = current_index;
                if new_node.x == target_x && new_node.y == target_y && new_node.consecutive >= min_consecutive {
                    current_node = new_node;
                    break 'outer;
                }
                // Combines node score with node index to provide unique keys
                queue.insert((new_score as usize * scores.len() + node_index) as i32, new_node);
            }
        }
        match queue.pop_first() {
            Some((_, node)) => {
                current_node = node;
                current_index = get_node_index(&current_node) as i32;
            }
            None => { return -1; }
        }
    }
    let mut temp_x = current_node.x;
    let mut temp_y = current_node.y;
    let mut temp_index = get_node_index(&current_node) as i32;
    while temp_index >= 0 {
        out_path.push((temp_x, temp_y));
        temp_index = previous[temp_index as usize];
        temp_x = temp_index / index_multiplier as i32;
        temp_y = temp_x / width as i32;
        temp_x -= temp_y * width as i32;
    }
    current_node.cost
}

fn visualize_path(layout: &[&[u8]], path: &Vec<(i32, i32)>) {
    let mut layout_copy = layout.iter().map(|&bytes| Vec::from(bytes)).collect::<Vec<_>>();
    for &(x, y) in path {
        layout_copy[y as usize][x as usize] = b'.';
    }
    for line in layout_copy.iter() {
        println!("{}", std::str::from_utf8(&line).unwrap());
    }
}

fn visualize_history(layout: &[&[u8]], history: &Vec<(i32, i32)>) {
    let grayscale = " .:-=+*#%@%#*+=-:. ".as_bytes();
    let mut layout_history = layout.iter().map(|&bytes| Vec::from(bytes)).collect::<Vec<_>>();
    for (i, &(x, y)) in history.iter().enumerate() {
        layout_history[y as usize][x as usize] = grayscale[i * grayscale.len() / history.len()];
    }
    for line in layout_history.iter() {
        println!("{}", std::str::from_utf8(&line).unwrap());
    }

}

fn main() {
    let time_start = Instant::now();
    let input_str = read_to_string("input.txt").unwrap();
    let time_start_no_io = Instant::now();
    let layout = input_str.lines().map(|str| str.as_bytes()).collect::<Vec<_>>();
    let mut path: Vec<(i32, i32)> = vec![];
    let mut history: Vec<(i32, i32)> = vec![];
    let mut path2: Vec<(i32, i32)> = vec![];
    let mut history2: Vec<(i32, i32)> = vec![];
    let loss1 = find_path(&layout, 0, 3, &mut path, &mut history);
    let loss2 = find_path(&layout, 4, 10, &mut path2, &mut history2);
    let elapsed = time_start.elapsed().as_micros();
    let elapsed_no_io = time_start_no_io.elapsed().as_micros();
    println!("Part 1 path:");
    visualize_path(&layout, &path);
    println!("\nPart 1 search order:");
    visualize_history(&layout, &history);
    println!("\nPart 2 path:");
    visualize_path(&layout, &path2);
    println!("\nPart 2 search order:");
    visualize_history(&layout, &history2);
    println!("Time: {}us", elapsed);
    println!("Time without file i/o: {}us", elapsed_no_io);
    println!("Loss1: {}", loss1);
    println!("Loss2: {}", loss2);
}