const INPUT: &str = include_str!("../../data/day14");

use std::collections::HashMap;

#[derive(Clone, Copy, Debug, Default)]
pub struct Mask {
    fixed: u64,
    select_value: u64,
}

impl Mask {
    pub fn parse(data: &str) -> Self {
        let (fixed, select_value) = data.trim().chars().fold((0, 0), |(fixed, select), c| {
            match c {
                '0' => (2*fixed + 0, 2*select + 0),
                '1' => (2*fixed + 1, 2*select + 0),
                'X' => (2*fixed + 0, 2*select + 1),
                _ => panic!("invalid mask character: {:?}", c),
            }
        });
        assert_eq!(0, select_value & fixed);
        Self { fixed, select_value }
    }

    pub fn apply(self, value: u64) -> u64 {
        self.fixed | (self.select_value & value)
    }
}

#[derive(Clone, Debug, Default)]
pub struct Fluctuations {
    clear: u64,
    flucts: Vec<u64>,
}

impl Fluctuations {
    pub fn apply(&self, address: u64) -> impl Iterator<Item = u64> + '_ {
        let address = address & !self.clear;
        self.flucts.iter().map(move |&mask| mask | address)
    }
}

impl From<Mask> for Fluctuations {
    fn from(mask: Mask) -> Self {
        let mut flucts = vec![mask.fixed];
        let mut fluct = mask.select_value;
        let mut current = 0;
        // println!("fluctuate         : {:036b}", fluct);
        while 0 != fluct {
            let shift = fluct.trailing_zeros();
            current += shift;
            let bit = 1 << current;
            // println!("extending with bit: {:036b}", bit);
            let new_flucts = flucts.iter().map(|&m| m | bit).collect::<Vec<_>>();
            flucts.extend(new_flucts);
            current += 1;
            fluct = fluct >> (shift + 1);
        }
        Self {
            clear: mask.select_value,
            flucts,
        }
    }
}

#[derive(Clone, Copy, Debug)]
pub enum Instruction {
    SetMask { mask: Mask },
    SetMemory { address: u64, value: u64 },
}

impl Instruction {
    pub fn parse(line: &str) -> Self {
        if let Some(rem) = line.strip_prefix("mask = ") {
            let mask = Mask::parse(rem);
            Self::SetMask { mask }
        } else if let Some(rem) = line.strip_prefix("mem[") {
            let p = rem.find("] = ").unwrap();
            let address = &rem[..p];
            let address = address.parse().unwrap_or_else(|e| panic!("invalid address: {:?}: {}", address, e));
            let value = &rem[p+4..];
            let value = value.parse().unwrap_or_else(|e| panic!("invalid value: {:?}: {}", value, e));
            Self::SetMemory { address, value }
        } else {
            panic!("invalid input: {:?}", line);
        }
    }

    pub fn parse_list(input: &str) -> Vec<Self> {
        input.lines().map(Self::parse).collect()
    }
}

fn main() {
    let instructions = Instruction::parse_list(INPUT);
    {
        let mut memory = HashMap::new();
        let mut current_mask = Mask::default();
        for instr in &instructions {
            match *instr {
                Instruction::SetMask { mask } => {
                    current_mask = mask;
                },
                Instruction::SetMemory { address, value } => {
                    memory.insert(address, current_mask.apply(value));
                },
            }
        }
        println!("Sum of values: {}", memory.values().sum::<u64>());
    }
    {
        let mut memory = HashMap::new();
        let mut flucts = Fluctuations::default();
        for instr in &instructions {
            match *instr {
                Instruction::SetMask { mask } => {
                    flucts = Fluctuations::from(mask);
                },
                Instruction::SetMemory { address, value } => {
                    for address in flucts.apply(address) {
                        memory.insert(address, value);
                    }
                },
            }
        }
        println!("Sum of values: {}", memory.values().sum::<u64>());
    }
}