use std::collections::HashMap;

static DAY: u8 = 4;

fn main() {
    let input = advent::read_lines(DAY);
    println!("{DAY}a: {}", count_xmas(&input));
    println!("{DAY}b: {}", count_x_mas(&input));
}

fn get_letter_grid(input: &[String]) -> HashMap<(isize,isize), char> {
    let mut grid = HashMap::new();
    for (y, line) in input.iter().enumerate() {
        for (x, c) in line.chars().enumerate() {
            grid.insert((x as isize, y as isize), c);
        }
    }
    grid
}

fn word_search(grid: &HashMap<(isize,isize), char>, x: isize, y: isize, diffx: isize, diffy: isize, word: &str) -> bool {
    for (i, c) in word.chars().enumerate() {
        let i = i as isize;
        if let Some(letter) = grid.get(&(x + i*diffx, y + i*diffy)) {
            if *letter != c {
                return false;
            }
        } else {
            return false;
        }
    }
    true
}

fn word_up(grid: &HashMap<(isize,isize), char>, x: isize, y: isize, word: &str) -> bool {
    word_search(grid, x, y, 0, -1, word)
}

fn word_down(grid: &HashMap<(isize,isize), char>, x: isize, y: isize, word: &str) -> bool {
    word_search(grid, x, y, 0, 1, word)
}

fn word_left(grid: &HashMap<(isize,isize), char>, x: isize, y: isize, word: &str) -> bool {
    word_search(grid, x, y, -1, 0, word)
}

fn word_right(grid: &HashMap<(isize,isize), char>, x: isize, y: isize, word: &str) -> bool {
    word_search(grid, x, y, 1, 0, word)
}

fn word_ne(grid: &HashMap<(isize,isize), char>, x: isize, y: isize, word: &str) -> bool {
    word_search(grid, x, y, 1, -1, word)
}

fn word_se(grid: &HashMap<(isize,isize), char>, x: isize, y: isize, word: &str) -> bool {
    word_search(grid, x, y, 1, 1, word)
}

fn word_sw(grid: &HashMap<(isize,isize), char>, x: isize, y: isize, word: &str) -> bool {
    word_search(grid, x, y, -1, 1, word)
}

fn word_nw(grid: &HashMap<(isize,isize), char>, x: isize, y: isize, word: &str) -> bool {
    word_search(grid, x, y, -1, -1, word)
}

fn count_xmas(input: &[String]) -> usize {
    let grid = get_letter_grid(input);
    let max_y = input.len() as isize;
    let max_x = input[0].len() as isize;
    let mut xmas_count = 0;

    for x in 0..max_x {
        for y in 0..max_y {
            if word_up(&grid, x, y, "XMAS") { xmas_count += 1; }
            if word_down(&grid, x, y, "XMAS") { xmas_count += 1; }
            if word_left(&grid, x, y, "XMAS") { xmas_count += 1; }
            if word_right(&grid, x, y, "XMAS") { xmas_count += 1; }
            if word_ne(&grid, x, y, "XMAS") { xmas_count += 1; }
            if word_se(&grid, x, y, "XMAS") { xmas_count += 1; }
            if word_sw(&grid, x, y, "XMAS") { xmas_count += 1; }
            if word_nw(&grid, x, y, "XMAS") { xmas_count += 1; }
        }
    }

    xmas_count
}

fn count_x_mas(input: &[String]) -> usize {
    let grid = get_letter_grid(input);
    let max_y = input.len() as isize;
    let max_x = input[0].len() as isize;
    let mut xmas_count = 0;

    for x in 0..max_x {
        for y in 0..max_y {
            for word1 in ["MAS", "SAM"] {
                for word2 in ["MAS", "SAM"] {
                    if word_se(&grid, x, y, word1) && word_sw(&grid, x+2, y, word2) { xmas_count += 1; }
                }
            }
        }
    }

    xmas_count
}

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

    #[test]
    fn test() {
        let input = [
            "MMMSXXMASM",
            "MSAMXMSMSA",
            "AMXSXMAAMM",
            "MSAMASMSMX",
            "XMASAMXAMM",
            "XXAMMXXAMA",
            "SMSMSASXSS",
            "SAXAMASAAA",
            "MAMMMXMMMM",
            "MXMXAXMASX",
        ].iter().map(|&x| String::from(x)).collect::<Vec<_>>();
        assert_eq!(count_xmas(&input), 18);
        assert_eq!(count_x_mas(&input), 9);
    }
}