use super::util;
use std::collections::HashMap;

pub fn solve() {

    let lines = util::read_file("input/day14.txt");

    let part1 = solve_part1(&lines);
    println!("Day X Part 1: {}", part1);


    let part2 = solve_part2(&lines);
    println!("Day X Part 2: {}", part2);
}

fn solve_part1(lines: &Vec<String>) -> u64 {

    let mut mem = vec!{0u64; 1 << 16};
    let mut mask = BitMask { ones_mask: 0, zeros_mask: !0 };
    for line in lines {

        if line.starts_with("mask = ") {
            mask = BitMask::parse(&String::from(&line[("mask = ".len())..]))
        } else {
            let instr = Instruction::parse(line);
            mem[instr.pos] = mask.apply(instr.n);
        }
    }
    return mem.iter().sum();
}

fn solve_part2(lines: &Vec<String>) -> u64 {

    /*
    Idea for performance optimalization: use a Vec containing the "root" memory position, and
    the mask of the floats used to write to that position.
    Whenever something is written, check (using bit operations) if a collision between existing
    memory positions is possible. If that is the case, modify the old (existing) entry by
     removing all overlapping entries in the float mask.
     */

    let mut mem: HashMap<u64, u64> = HashMap::new();
    let mut mask = MemoryBitMask { ones_mask: 0, float_masks: Vec::new() };
    for line in lines {

        if line.starts_with("mask = ") {
            mask = MemoryBitMask::parse(&String::from(&line[("mask = ".len())..]))
        } else {
            let instr = Instruction::parse(line);
            mask.set_all(instr.pos as u64, instr.n, &mut mem);
        }
    }
    return mem.values().sum();
}

struct BitMask {
    /// One (1) everywhere except where the mask should be zero (0)
    zeros_mask: u64,
    /// Zero everywhere except where the mask is 1
    ones_mask: u64,
}

impl BitMask {

    fn parse(input: &String) -> BitMask {

        let mut zeros = u64::max_value();
        let mut ones = 0u64;

        let mut i = 0usize;
        for c in input.chars().rev() {

            if c == '0'  {
                zeros &= !(1 << i);
            } else if c == '1' {
                ones |= 1 << i;
            }
            i += 1;
        }
        return BitMask {
            zeros_mask: zeros,
            ones_mask: ones,
        }

    }


    fn apply(&self, x: u64) -> u64 {
        return (x & self.zeros_mask) | self.ones_mask;
    }
}

struct MemoryBitMask {

    ones_mask: u64,
    ///Positions of the floating numbers
    float_masks: Vec<u64>,

}
impl MemoryBitMask {

    fn parse(input: &String) -> MemoryBitMask {

        let mut ones = 0u64;
        let mut float_masks: Vec<u64> = Vec::new();

        let mut i = 0usize;
        for c in input.chars().rev() {

            if c == 'X'  {
                float_masks.push(1 << i);
            } else if c == '1' {
                ones |= 1 << i;
            }
            i += 1;
        }
        return MemoryBitMask {
            float_masks,
            ones_mask: ones,
        }

    }

    fn set_all(&self, pos: u64, n: u64, mem: &mut HashMap<u64, u64>) {

        let base_pos = pos | self.ones_mask;
        for float_opt in 0..(1 << self.float_masks.len()) {

            let mut real_pos = base_pos;
            for i in 0..self.float_masks.len() {

                let float_mask = self.float_masks[i];
                // Clear the bit in the original pos
                real_pos &= !float_mask;
                let bit_mask = 1 << i;
                if float_opt & bit_mask != 0 {
                    real_pos |= float_mask;
                }
            }
            mem.insert(real_pos, n);
        }
    }
}

struct Instruction {
    pos: usize,
    n: u64,
}

impl Instruction {
    fn parse(input: &String) -> Instruction {
        let pos_start = input.find("[").unwrap() + 1;
        let pos_end = input.find("]").unwrap();
        let pos = input[pos_start..pos_end].parse::<usize>().unwrap();

        let n_start = input.find(" = ").unwrap() + 3;
        let n = input[n_start..].parse::<u64>().unwrap();

        return Instruction { pos, n };
    }
}