From 42ead11aa07aa560606ee5ab6f66a2de426102da Mon Sep 17 00:00:00 2001
From: Johannes <freelon@gmx.net>
Date: Sun, 12 Nov 2023 19:15:31 +0100
Subject: [PATCH] Day 16 part 2.

---
 src/main.rs        |  2 +-
 src/tasks/day16.rs | 78 ++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 76 insertions(+), 4 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 09be118..49e35d9 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,3 +1,3 @@
 fn main() {
-    aoc_2018::tasks::day16::task1();
+    aoc_2018::tasks::day16::task2();
 }
diff --git a/src/tasks/day16.rs b/src/tasks/day16.rs
index 40f28d2..cb18fdc 100644
--- a/src/tasks/day16.rs
+++ b/src/tasks/day16.rs
@@ -1,3 +1,4 @@
+use std::collections::{HashMap, HashSet};
 use std::ops::{BitAnd, BitOr};
 
 use itertools::Itertools;
@@ -15,6 +16,65 @@ pub fn task1() {
     println!("Day 16 part1: {count}");
 }
 
+pub fn task2() {
+    let (samples, program) = include_str!("../../input/day16.txt")
+        .split_once("\n\n\n\n")
+        .unwrap();
+
+    let code_to_op_candidates = samples
+        .split("\n\n")
+        .map(|sample_input| parse(sample_input))
+        .map(|sample| (sample.1.opcode, matches_ops(sample)))
+        .into_group_map_by(|x| x.0);
+
+    // let code_to_op: HashMap<usize, HashSet<Op>> =
+    let mut code_to_op_candidates: HashMap<_, _> = code_to_op_candidates
+        .into_iter()
+        .map(|(code, ops)| {
+            let mut set: HashSet<_> = HashSet::from_iter(OPS.iter().map(ToOwned::to_owned));
+            for (_, candidate_set) in ops {
+                set = set
+                    .intersection(&candidate_set)
+                    .map(ToOwned::to_owned)
+                    .collect();
+            }
+
+            (code, set)
+        })
+        .collect();
+
+    let mut codes_to_op: HashMap<usize, Op> = HashMap::new();
+    loop {
+        let single = code_to_op_candidates.iter().find(|(_, ops)| ops.len() == 1);
+        let Some(single) = single else {
+            break;
+        };
+
+        let code = *single.0;
+        let op = *single.1.iter().next().unwrap();
+
+        codes_to_op.insert(code, op);
+
+        code_to_op_candidates.iter_mut().for_each(|(_, ops)| {
+            ops.remove(&op);
+        });
+    }
+
+    if codes_to_op.len() != OPS.len() {
+        panic!("Bad final opcode assignments: {codes_to_op:?}");
+    }
+
+    let program: Vec<Instruction> = program.lines().map(parse_instruction).collect();
+
+    let mut registers = Registers([0, 0, 0, 0]);
+    for instruction in program {
+        let op: Op = codes_to_op[&instruction.opcode];
+        registers = op.process(registers, instruction);
+    }
+
+    println!("Day 16 part 2: {}", registers.0[0]);
+}
+
 fn parse(sample: &str) -> Sample {
     let (before, instruction, after) = sample.lines().collect_tuple().unwrap();
     let before: Vec<_> = before
@@ -29,6 +89,12 @@ fn parse(sample: &str) -> Sample {
         .split(", ")
         .collect();
     let after = Registers::from(after);
+    let instruction = parse_instruction(instruction);
+
+    Sample(before, instruction, after)
+}
+
+fn parse_instruction(instruction: &str) -> Instruction {
     let instruction: Vec<usize> = instruction.split(" ").map(|x| x.parse().unwrap()).collect();
     let instruction = Instruction {
         opcode: instruction[0],
@@ -36,8 +102,7 @@ fn parse(sample: &str) -> Sample {
         b: instruction[2],
         c: instruction[3],
     };
-
-    Sample(before, instruction, after)
+    instruction
 }
 
 fn matches(sample: Sample, op: Op) -> bool {
@@ -49,6 +114,13 @@ fn matches_count(sample: Sample) -> usize {
     vec.len()
 }
 
+fn matches_ops(sample: Sample) -> HashSet<Op> {
+    OPS.iter()
+        .filter(|op| matches(sample, **op))
+        .map(|op| op.to_owned())
+        .collect()
+}
+
 const OPS: &[Op] = &[
     Op::AddR,
     Op::AddI,
@@ -68,7 +140,7 @@ const OPS: &[Op] = &[
     Op::EqRR,
 ];
 
-#[derive(Copy, Clone, Debug)]
+#[derive(Copy, Clone, Debug, Hash, Eq, PartialEq)]
 enum Op {
     AddR,
     AddI,