use std::ops::AddAssign;

#[derive(Clone)]
struct Oneness(isize);

fn part1(input: &[String])
{
    let scale = calculate_scale(input);

    let mut gamma = 0;
    let mut epsilon = 0;
    for (i, oneness) in scale.iter().rev().enumerate() {
        gamma += (oneness.is_one() || oneness.is_balanced()) as usize * 2_usize.pow(i as u32);
        epsilon += (oneness.is_zero() || oneness.is_balanced()) as usize * 2_usize.pow(i as u32);
    }

    println!("Solution to part 1: {}", gamma * epsilon);
}

fn part2(input: Vec<String>)
{
    let mut oxygen = input.clone();
    let mut pos: usize = 0;
    while oxygen.len() != 1 {
        let scale = calculate_scale(&oxygen);
        let looking_for_one = scale[pos].is_one() || scale[pos].is_balanced();
        oxygen = oxygen
            .into_iter()
            .filter(|s| {
                let char_is_one = s.as_bytes()[pos] == '1' as u8;
                looking_for_one == char_is_one
            })
            .collect();

        pos += 1;
    }

    let mut co2 = input;
    pos = 0;
    while co2.len() != 1 {
        let scale = calculate_scale(&co2);
        let looking_for_one = scale[pos].is_zero();
        co2 = co2
            .into_iter()
            .filter(|s| {
                let char_is_one = s.as_bytes()[pos] == '1' as u8;
                looking_for_one == char_is_one
            })
            .collect();

        pos += 1;
    }

    let oxygen = usize::from_str_radix(&oxygen[0], 2).unwrap();
    let co2 = usize::from_str_radix(&co2[0], 2).unwrap();

    println!("Solution to part 2: {}", oxygen * co2);
}

fn calculate_scale(input: &[String]) -> Vec<Oneness>
{
    let mut one_scale = vec![Oneness::balanced(); input[0].len()];
    for s in input.iter() {
        for (i, c) in s.char_indices() {
            one_scale[i] += match c {
                '1' => true,
                '0' => false,
                o => panic!("Invalid input: {}", o),
            };
        }
    }

    one_scale
}

pub fn run(input: Vec<String>)
{
    // Determine the number of bits by taking a look at the first string.
    let num_bits = input[0].len();
    let mut one_scale = vec![Oneness::balanced(); num_bits];

    part1(&input);
    part2(input);
}

impl Oneness
{
    pub const fn balanced() -> Self { Self(0) }

    pub fn is_one(&self) -> bool { self.0 > 0 }

    pub fn is_zero(&self) -> bool { self.0 < 0 }

    pub fn is_balanced(&self) -> bool { self.0 == 0 }
}

impl AddAssign<bool> for Oneness
{
    fn add_assign(&mut self, rhs: bool)
    {
        match rhs {
            true => self.0 += 1,
            false => self.0 -= 1,
        }
    }
}