fn aoc_07_part1(input: &mut Vec<i32>) -> i32 {
    // Sorting so the median works...
    input.sort();

    // input.len() % 2 will either be 0 or 1, which is what we want...
    let middle = (input.len() + (input.len() % 2)) / 2;
    // Get the median
    let median = input[middle];

    // Sum of the abs of the elements vs the median gives the total consumption
    input.iter().map(|el| i32::abs(el - median)).sum()
}

fn aoc_07_part2(input: &Vec<i32>) -> i32 {
    let mean = input.iter().sum::<i32>() / input.len() as i32;

    // It can be ranging from previous to next value so we need to do a -1, 0, 1
    // and get the min of that range of values...
    (-1..=1).into_iter().map(|epsilon| {
        // Computing the total consumption for that epsilon
        input.iter().map(|el| {
            let val = i32::abs((*el) - mean + epsilon);
            (val * (val + 1)) / 2
        }).sum()
    }).min().unwrap_or(0)

}

fn main() {
    let line = include_str!("../input.txt").lines().next().unwrap_or("");
    let mut s: Vec<i32> = line.split(',').filter_map(|s| s.parse::<i32>().ok()).collect::<Vec<i32>>();
    let result = aoc_07_part1(&mut s);
    println!("Total consumption: {}", result);
    let result = aoc_07_part2(&s);
    println!("New total consumption: {}", result);
}