use std::fs;
use std::path::Path;

/* AOC21 Day 9: https://adventofcode.com/2021/day/9 */
fn main() {
    let input = Path::new("resources").join("input.txt");
    let content = fs::read_to_string(input).expect("Unable to read input file");
    let width = content.lines().next().expect("Empty input").len();
    let heightmap: Vec<u8> = content.lines().flat_map(|l| l.bytes().map(|c| c-b'0')).collect();
    println!("Ex1: The risk level is {}", risk_level(&heightmap, width));
    println!("Ex2: The result is {}", find_basins(&heightmap, width));
}

fn risk_level(hm: &[u8], w: usize) -> u32 {
    (0..hm.len()).map(|i|
        if ( i < w || hm[i-w] > hm[i] ) &&
           ( i+w >= hm.len() || hm[i+w] > hm[i] ) &&
           ( i%w == 0 || hm[i-1] > hm[i] ) &&
           ( (i+1)%w == 0 || hm[i+1] > hm[i] ) {
               1 + hm[i] as u32
        } else { 0 }
    ).sum()
}

fn find_basins(hm: &[u8], w: usize) -> u32 {
    let mut hmask = vec![false;hm.len()];
    let mut basins: Vec<u32> = vec![];
    for i in 0..hm.len() {
        let size = propagate_basin(hm, &mut hmask, i, w);
        if size > 0 { basins.push(size); }
    }
    basins.sort_by(|a,b| b.cmp(a));
    basins.iter().take(3).product()
}

fn propagate_basin(hm: &[u8], hmask: &mut [bool], i: usize, w: usize) -> u32 {
    if !hmask[i] && hm[i] != 9 {
        hmask[i] = true;
        1 + if i >= w { propagate_basin(hm, hmask, i-w, w) } else { 0 } +
            if i+w < hm.len() { propagate_basin(hm, hmask, i+w, w) } else { 0 } +
            if i%w != 0 { propagate_basin(hm, hmask, i-1, w) } else { 0 } +
            if (i+1)%w != 0 { propagate_basin(hm, hmask, i+1, w) } else { 0 }
    } else {
        0
    }
}