use std::cmp::{max, min};

fn parse(input: &str) -> Vec<Vec<char>> {
    input.lines().map(|s| s.chars().collect()).collect()
}

fn check_xmas(chars: &Vec<Vec<char>>, start: &(i32, i32), direction: (i32, i32)) -> bool {
    let h = chars.len() as i32;
    let w = chars[0].len() as i32;

    if start.0 + 3 * direction.0 < 0
        || start.0 + 3 * direction.0 >= w
        || start.1 + 3 * direction.1 < 0
        || start.1 + 3 * direction.1 >= h
    {
        return false;
    }

    let xmas = vec!['X', 'M', 'A', 'S'];

    for x in 0..4 {
        let cur_pos = (start.0 + x * direction.0, start.1 + x * direction.1);

        if chars[cur_pos.1 as usize][cur_pos.0 as usize] != xmas[x as usize] {
            return false;
        }
    }

    true
}

fn solve_1(input: &str) -> usize {
    let input = parse(input);

    let h = input.len() as i32;
    let w = input[0].len() as i32;

    let mut result = 0;

    for j in 0..h {
        for i in 0..w {
            if input[j as usize][i as usize] == 'X' {
                for ii in max(0, i - 1)..=min(i + 1, w - 1) {
                    for jj in max(0, j - 1)..=min(j + 1, h - 1) {
                        if input[jj as usize][ii as usize] == 'M' {
                            let dir = (ii - i, jj - j);

                            if check_xmas(&input, &(i, j), dir) {
                                result += 1;
                            }
                        }
                    }
                }
            }
        }
    }

    result
}

fn check_x_mas(chars: &Vec<Vec<char>>, candidate: (i32, i32)) -> usize {
    let h = chars.len() as i32;
    let w = chars[0].len() as i32;
    let (x, y) = candidate;

    if x == 0 || x == w - 1 || y == 0 || y == h - 1 {
        return 0;
    }

    let mut result = 0;

    let mut diagonal_ms = vec![];
    let mut diagonal_ss = vec![];
    for i in 0..2 {
        for j in 0..2 {
            let coord = (x as usize - 1 + 2 * i, y as usize - 1 + 2 * j);
            match chars[coord.1][coord.0] {
                'M' => diagonal_ms.push(coord),
                'S' => diagonal_ss.push(coord),
                _ => (),
            }
        }
    }
    if diagonal_ms.len() >= 2 && diagonal_ss.len() >= 2 {
        if diagonal_ms
            .into_iter()
            .filter(|&(cx, cy)| {
                let x_diff = cx as i32 - candidate.0;
                let y_diff = cy as i32 - candidate.1;

                diagonal_ss.contains(&(
                    (candidate.0 - x_diff) as usize,
                    (candidate.1 - y_diff) as usize,
                ))
            })
            .count()
            == 2
        {
            result += 1;
        }
    }

    result
}

fn solve_2(input: &str) -> usize {
    let input = parse(input);

    let h = input.len() as i32;
    let w = input[0].len() as i32;

    let mut result = 0;

    for j in 0..h {
        for i in 0..w {
            if input[j as usize][i as usize] == 'A' {
                result += check_x_mas(&input, (i, j));
            }
        }
    }

    result
}

fn main() {
    println!("Hello, this is Patrick!");

    let input = include_str!("../input.txt");

    let result_1 = solve_1(input);
    println!("The number of XMASs is {}", result_1);

    let result_2 = solve_2(input);
    println!("The number of X-MASs is {}", result_2);
}

#[cfg(test)]

mod tests {
    use super::*;

    #[test]
    fn test_1() {
        let test_input = include_str!("../test.txt");
        assert_eq!(solve_1(test_input), 18);
    }

    #[test]
    fn test_2() {
        let test_input = include_str!("../test.txt");
        assert_eq!(solve_2(test_input), 9);
    }
}