From 8cd9c9b7a02b683067a1f7c800a40cfdf556e2b6 Mon Sep 17 00:00:00 2001
From: Juris <jurisk@users.noreply.github.com>
Date: Sat, 16 Dec 2023 06:30:01 +0200
Subject: [PATCH] Optimizer solution for 2022-21

---
 .../jurisk/adventofcode/y2018/Advent23.scala  |   2 +-
 .../jurisk/adventofcode/y2022/Advent21.scala  | 124 +++++++++++++++---
 2 files changed, 109 insertions(+), 17 deletions(-)

diff --git a/scala2/src/main/scala/jurisk/adventofcode/y2018/Advent23.scala b/scala2/src/main/scala/jurisk/adventofcode/y2018/Advent23.scala
index 3cfc02e4..a5addac1 100644
--- a/scala2/src/main/scala/jurisk/adventofcode/y2018/Advent23.scala
+++ b/scala2/src/main/scala/jurisk/adventofcode/y2018/Advent23.scala
@@ -47,7 +47,7 @@ object Advent23 {
     val ctx = new Context
     import ctx._
 
-    val o = mkOptimize
+    val o = mkOptimize()
 
     val List(x, y, z) = List("x", "y", "z").map(mkIntConst)
 
diff --git a/scala2/src/main/scala/jurisk/adventofcode/y2022/Advent21.scala b/scala2/src/main/scala/jurisk/adventofcode/y2022/Advent21.scala
index 6c1d4d73..77d734a8 100644
--- a/scala2/src/main/scala/jurisk/adventofcode/y2022/Advent21.scala
+++ b/scala2/src/main/scala/jurisk/adventofcode/y2022/Advent21.scala
@@ -1,6 +1,7 @@
 package jurisk.adventofcode.y2022
 
 import cats.implicits._
+import com.microsoft.z3.{Context, IntNum}
 import jurisk.adventofcode.y2022.Advent21.Expression._
 import jurisk.adventofcode.y2022.Advent21.Operation._
 import jurisk.utils.FileInput._
@@ -144,29 +145,100 @@ object Advent21 {
     translate(target)
   }
 
-  def part1(data: Parsed): Value = {
-    val commands = data.toMap
-    val root     = convertCommandsToExpression(Root, commands, none)
-    root.rearrange match {
-      case Literal(value) => value
-      case other          => s"Expected Literal here but got $other".fail
+  private def solveOptimizer(
+    commands: Map[Name, Monkey],
+    calculate: Name,
+    extraEquality: Option[(Name, Name)],
+  ): Value = {
+    val ctx = new Context
+    import ctx._
+
+    val o = mkOptimize()
+
+    extraEquality foreach { case (a, b) =>
+      o.Add(mkEq(mkIntConst(a), mkIntConst(b)))
+    }
+
+    commands foreach { case (name, monkey) =>
+      val n = mkIntConst(name)
+      monkey match {
+        case Monkey.BinaryMonkey(aName, operation, bName) =>
+          val a = mkIntConst(aName)
+          val b = mkIntConst(bName)
+
+          val clauses = operation match {
+            case Operation.Plus     => mkEq(n, mkAdd(a, b)) :: Nil
+            case Operation.Minus    => mkEq(n, mkSub(a, b)) :: Nil
+            case Operation.Multiply => mkEq(n, mkMul(a, b)) :: Nil
+            case Operation.Divide   =>
+              mkEq(n, mkDiv(a, b)) ::
+                mkEq(mkRem(a, b), mkInt(0)) ::
+                Nil
+          }
+
+          clauses foreach { clause =>
+            o.Add(clause)
+          }
+        case Monkey.Literal(value)                        => o.Add(mkEq(n, mkInt(value)))
+      }
+    }
+
+    println(o.Check())
+
+    val model = o.getModel
+
+    val result = model.evaluate(mkIntConst(calculate), true)
+
+    result match {
+      case intNum: IntNum => intNum.getInt64
+      case _              => s"Expected IntNum: $result".fail
     }
   }
 
-  def part2(data: Parsed): Value = {
-    val Human = "humn"
+  private def solvePart1Optimizer(
+    commands: Map[Name, Monkey],
+    calculate: Name,
+  ): Value = solveOptimizer(commands, calculate, none)
 
-    val allCommands    = data.toMap - Human
-    val rootCommand    = allCommands(Root)
-    val (rootA, rootB) = rootCommand match {
-      case c: Monkey.BinaryMonkey => (c.a, c.b)
-      case _                      => s"Unexpectedly $rootCommand is not a binary command".fail
+  private def solvePart2Optimizer(
+    commands: Map[Name, Monkey],
+    leftSideName: Name,
+    rightSideName: Name,
+    calculate: Name,
+  ): Value =
+    solveOptimizer(commands, calculate, (leftSideName, rightSideName).some)
+
+  private def solvePart1Algebraic(
+    commands: Map[Name, Monkey],
+    calculate: Name,
+  ): Value = {
+    val expression = convertCommandsToExpression(calculate, commands, none)
+    println(s"Part 1 expression for $calculate: $expression")
+
+    expression.rearrange match {
+      case Literal(value) => value
+      case other          => s"Expected Literal here but got $other".fail
     }
+  }
 
-    val commands: Map[Name, Monkey] = allCommands - Root
+  def part1(data: Parsed): Value = {
+    val commands = data.toMap
+    val (a, b)   =
+      (solvePart1Algebraic(commands, Root), solvePart1Optimizer(commands, Root))
+    assert(a == b, s"Got different results: $a and $b")
+    a
+  }
 
-    val leftSide  = convertCommandsToExpression(rootA, commands, Human.some)
-    val rightSide = convertCommandsToExpression(rootB, commands, Human.some)
+  private def solvePart2Algebraic(
+    commands: Map[Name, Monkey],
+    leftSideName: Name,
+    rightSideName: Name,
+    calculate: Name,
+  ): Value = {
+    val leftSide  =
+      convertCommandsToExpression(leftSideName, commands, calculate.some)
+    val rightSide =
+      convertCommandsToExpression(rightSideName, commands, calculate.some)
 
     println(s"${leftSide.display} = ${leftSide.display}")
 
@@ -211,6 +283,26 @@ object Advent21 {
     result
   }
 
+  def part2(data: Parsed): Value = {
+    val Human = "humn"
+
+    val allCommands    = data.toMap - Human
+    val rootCommand    = allCommands(Root)
+    val (rootA, rootB) = rootCommand match {
+      case c: Monkey.BinaryMonkey => (c.a, c.b)
+      case _                      => s"Unexpectedly $rootCommand is not a binary command".fail
+    }
+
+    val commands: Map[Name, Monkey] = allCommands - Root
+
+    val (s1, s2) = (
+      solvePart2Algebraic(commands, rootA, rootB, Human),
+      solvePart2Optimizer(commands, rootA, rootB, Human),
+    )
+    assert(s1 == s1, s"Got different results: $s1 and $s2")
+    s1
+  }
+
   def main(args: Array[String]): Unit = {
     val testData = readFileText("2022/21-test.txt")
     val realData = readFileText("2022/21.txt")