use aoc_runner_derive::{aoc, aoc_generator};

use bitvec::prelude::*;

type Mbv = BitVec<Msb0>;

#[aoc_generator(day3)]
pub fn parse_diagnostics(input: &str) -> Vec<Mbv> {
    input
        .lines()
        .map(|line| {
            let length = line.len();
            let mut element = usize::from_str_radix(line, 2).unwrap();
            element <<= usize::BITS - length as u32;
            let mut vec: Mbv = Mbv::from_element(element);
            vec.truncate(length);
            vec
        })
        .collect()
}

#[aoc(day3, part1)]
pub fn calculate_power_consumption(input: &[Mbv]) -> usize {
    let gamma_factor = most_common_bit_vec(input);
    let epsilon_factor = invert_bit_vec(&gamma_factor);
    let gamma_factor = bit_vec_to_usize(&gamma_factor);
    let epsilon_factor = bit_vec_to_usize(&epsilon_factor);
    gamma_factor * epsilon_factor
}

pub fn most_common_bit(input: &[Mbv], pos: usize) -> bool {
    let count = input.iter().filter(|vec| *vec.get(pos).unwrap()).count();
    count * 2 >= input.len()
}

pub fn most_common_bit_vec(input: &[Mbv]) -> Mbv {
    let mut ret_val = Mbv::new();
    for i in 0..input[0].len() {
        ret_val.push(most_common_bit(input, i))
    }
    ret_val
}

pub fn invert_bit_vec(val: &Mbv) -> Mbv {
    val.iter().map(|val| !*val).collect()
}

pub fn bit_vec_to_usize(val: &Mbv) -> usize {
    let mut ret_val = 0;
    for bit in val {
        ret_val <<= 1;
        if *bit {
            ret_val += 1;
        }
    }
    ret_val
}

pub fn get_oxygen_generator_rating(input: &[Mbv]) -> usize {
    bit_vec_to_usize(&get_recursive_prefix_match(input, false))
}

pub fn get_co2_scrubber_rating(input: &[Mbv]) -> usize {
    bit_vec_to_usize(&get_recursive_prefix_match(input, true))
}

pub fn get_recursive_prefix_match(input: &[Mbv], invert: bool) -> Mbv {
    let mut input: Vec<Mbv> = input.to_vec();
    for i in 0..input[0].len() {
        if input.len() > 1 {
            let mut target = most_common_bit(&input, i);
            if invert {
                target = !target;
            }
            input.retain(|val| val.get(i).unwrap() == target);
        } else {
            break;
        }
    }
    assert_eq!(input.len(), 1);
    input[0].clone()
}

#[aoc(day3, part2)]
pub fn calculate_life_support_rating(input: &[Mbv]) -> usize {
    get_oxygen_generator_rating(input) * get_co2_scrubber_rating(input)
}

#[cfg(test)]
mod test {
    const EXAMPLE_INPUT: &str = "00100
11110
10110
10111
10101
01111
00111
11100
10000
11001
00010
01010";

    #[test]
    fn test_power_consumption_example() {
        let result = super::calculate_power_consumption(&super::parse_diagnostics(EXAMPLE_INPUT));
        assert_eq!(result, 198);
    }

    #[test]
    fn test_life_support_example() {
        let result = super::calculate_life_support_rating(&super::parse_diagnostics(EXAMPLE_INPUT));
        assert_eq!(result, 230);
    }
}