use std::collections::HashSet;

fn distance<T>(a: T, b: T) -> T
where
    T: Ord + std::ops::Sub<Output = T>,
{
    if a < b {
        return b - a;
    } else {
        return a - b;
    }
}

fn check_line_safety_1(line: &Vec<u32>) -> bool {
    let len = line.len();
    let mut prev = line[0];
    let up_or_down = line[0] < line[1];

    for i in 1..len {
        if prev == line[i] {
            return false;
        }

        let dist = distance(prev, line[i]);
        if dist < 1 || dist > 3 {
            return false;
        }

        if up_or_down && prev > line[i] {
            return false;
        } else if !up_or_down && prev < line[i] {
            return false;
        }

        prev = line[i];
    }

    return true;
}

fn solve_1(input: &str) -> u32 {
    let data: Vec<Vec<u32>> = input
        .lines()
        .map(|line| {
            line.split_whitespace()
                .map(|number| number.parse::<u32>().unwrap())
                .collect()
        })
        .collect();

    return data.into_iter().filter(check_line_safety_1).count() as u32;
}

fn check_line_safety_2(line: &Vec<u32>) -> bool {
    let len = line.len();
    let mut prev = line[0];
    let mut ups: HashSet<usize> = HashSet::new();
    let mut downs: HashSet<usize> = HashSet::new();
    let mut distance_violations: HashSet<usize> = HashSet::new();

    for i in 1..len {
        let dist = distance(prev, line[i]);

        if dist < 1 || dist > 3 {
            distance_violations.insert(i);
        }

        if prev < line[i] {
            ups.insert(i);
        } else if prev > line[i] {
            downs.insert(i);
        }

        prev = line[i];
    }

    if ups.len() > 2 && downs.len() > 2 {
        return false;
    }
    if distance_violations.len() > 2 {
        return false;
    } else if distance_violations.len() == 0 && (ups.len() == 0 || downs.len() == 0) {
        return true;
    }

    let minor = if ups.len() > downs.len() {
        downs.clone()
    } else {
        ups.clone()
    };

    let union: HashSet<_> = minor.union(&distance_violations).to_owned().collect();
    if union.len() == 1 {
        let mut line_copy_1 = line.clone();
        let mut line_copy_2 = line.clone();
        let delete_spot = *union.into_iter().next().unwrap();

        line_copy_1.remove(delete_spot);
        line_copy_2.remove(delete_spot - 1);

        return check_line_safety_1(&line_copy_1) || check_line_safety_1(&line_copy_2);
    } else {
        let mut line_copy_1 = line.clone();
        let mut line_copy_2 = line.clone();
        let mut line_iter = union.into_iter();

        line_copy_1.remove(*line_iter.next().unwrap());
        line_copy_2.remove(*line_iter.next().unwrap());

        return check_line_safety_1(&line_copy_1) || check_line_safety_1(&line_copy_2);
    }
}

fn solve_2(input: &str) -> u32 {
    let data: Vec<Vec<u32>> = input
        .lines()
        .map(|line| {
            line.split_whitespace()
                .map(|number| number.parse::<u32>().unwrap())
                .collect()
        })
        .collect();

    return data.into_iter().filter(check_line_safety_2).count() as u32;
}

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

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

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

    let result_2 = solve_2(input);
    println!("The number of safe reports (version 2) 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), 2);
    }

    #[test]
    fn test_2() {
        let test_input = include_str!("../test.txt");

        assert_eq!(solve_2(test_input), 4);
    }
}