# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: 'Python 3.10.0 64-bit (''adventofcode-mOkh6lsX'': pipenv)'
#     language: python
#     name: python3
# ---

# # Day 24, symbolic evaluation
#
# * <https://adventofcode.com/2021/day/24>
#
# This puzzle is different from most AoC problems in that the description and tests are not actually all that much use. You need to study the puzzle input too, as it is the specific mathematical expressions created from the input that'll determine when, given the 14 different inputs (each between 1 and 9), you'll get a zero a the output.
#
# ## Puzzle input patterns
#
# The input consists of 14 repeated sections like this:
#
# <table>
#   <thead>
#     <tr>
#       <th align="right">#</th>
#       <th >opcode</th>
#       <th align="right">op1</th>
#       <th align="right">op2</th>
#       <th align="left" style="text-align: left">interpretation</th>
#     </tr>
#   </thead>
#   <tbody>
#     <tr>
#       <td align="left" style="text-align: left">1</td>
#       <td align="left" style="text-align: left">inp</td>
#       <td align="right">w</td>
#       <td align="right"></td>
#       <td align="left" style="text-align: left"><tt>w = <i>input_digit</i></tt></td>
#     </tr>
#     <tr>
#       <td align="left" style="text-align: left">2</td>
#       <td align="left" style="text-align: left">mul</td>
#       <td align="right">x</td>
#       <td align="right">0</td>
#       <td align="left" rowspan="3" style="text-align: left; vertical-align: top">
#         <tt>x = z % 26</tt><br/>
#         Here, <tt>z</tt> is the output of the previous section.
#       </td>
#     </tr>
#     <tr>
#       <td align="left" style="text-align: left">3</td>
#       <td align="left" style="text-align: left">add</td>
#       <td align="right">x</td>
#       <td align="right">z</td>
#     </tr>
#     <tr>
#       <td align="left" style="text-align: left">4</td>
#       <td align="left" style="text-align: left">mod</td>
#       <td align="right">x</td>
#       <td align="right">26</td>
#     </tr>
#     <tr>
#       <td align="left" style="text-align: left">5</td>
#       <td align="left" style="text-align: left">div</td>
#       <td align="right">z</td>
#       <td align="right"><i>A</i></td>
#       <td align="left" style="text-align: left">
#         <tt>z = z / <i>A</i></tt><br/>
#         <i>A</i> is either 1 or 26, depending on <i>B</i>
#       </td>
#     </tr>
#     <tr>
#       <td align="left" style="text-align: left">6</td>
#       <td align="left" style="text-align: left">add</td>
#       <td align="right">x</td>
#       <td align="right"><i>B</i></td>
#       <td align="left" style="text-align: left">
#         <tt>x = x + <i>B</i><tt><br/>
#         <i>B</i> is a number between -15 and +15.
#       </td>
#     </tr>
#     <tr>
#       <td align="left" style="text-align: left">7</td>
#       <td align="left" style="text-align: left">eql</td>
#       <td align="right">x</td>
#       <td align="right">w</td>
#       <td align="left" rowspan="2" style="text-align: left; vertical-align: top">
#         <tt>x = 0 if x == w else 1</tt><br/>
#       </td>
#     </tr>
#     <tr>
#       <td align="left" style="text-align: left">8</td>
#       <td align="left" style="text-align: left">eql</td>
#       <td align="right">x</td>
#       <td align="right">0</td>
#     </tr>
#     <tr>
#       <td align="left" style="text-align: left">9</td>
#       <td align="left" style="text-align: left">mul</td>
#       <td align="right">y</td>
#       <td align="right">0</td>
#       <td align="left" rowspan="4" style="text-align: left; vertical-align: top">
#         <tt>y = 25 * x + 1</tt><br/>
#         <tt>x</tt> is either 0 or 1, so <tt>y</tt> is now either 1 or 26.
#       </td>
#     </tr>
#     <tr>
#       <td align="left" style="text-align: left">10</td>
#       <td align="left" style="text-align: left">add</td>
#       <td align="right">y</td>
#       <td align="right">x</td>
#     </tr>
#     <tr>
#       <td align="left" style="text-align: left">11</td>
#       <td align="left" style="text-align: left">mul</td>
#       <td align="right">y</td>
#       <td align="right">25</td>
#     </tr>
#     <tr>
#       <td align="left" style="text-align: left">12</td>
#       <td align="left" style="text-align: left">add</td>
#       <td align="right">y</td>
#       <td align="right">1</td>
#     </tr>
#     <tr>
#       <td align="left" style="text-align: left">13</td>
#       <td align="left" style="text-align: left">mul</td>
#       <td align="right">z</td>
#       <td align="right">y</td>
#       <td align="left" style="text-align: left"><tt>z = z * y</tt></td>
#     </tr>
#     <tr>
#       <td align="left" style="text-align: left">14</td>
#       <td align="left" style="text-align: left">mul</td>
#       <td align="right">y</td>
#       <td align="right">0</td>
#       <td align="left" rowspan="4" style="text-align: left; vertical-align: top">
#         <tt>y = (w + <i>C</i>) * x</tt>
#         <br/><i>C</i> is a positive, non-zero integer. <tt>x</tt> is either 0 or 1.
#       </td>
#     </tr>
#     <tr>
#       <td align="left" style="text-align: left">15</td>
#       <td align="left" style="text-align: left">add</td>
#       <td align="right">y</td>
#       <td align="right">w</td>
#     </tr>
#     <tr>
#       <td align="left" style="text-align: left">16</td>
#       <td align="left" style="text-align: left">add</td>
#       <td align="right">y</td>
#       <td align="right"><i>C</i></td>
#     </tr>
#     <tr>
#       <td align="left" style="text-align: left">17</td>
#       <td align="left" style="text-align: left">mul</td>
#       <td align="right">y</td>
#       <td align="right">x</td>
#     </tr>
#     <tr>
#       <td align="left" style="text-align: left">18</td>
#       <td align="left" style="text-align: left">add</td>
#       <td align="right">z</td>
#       <td align="right">y</td>
#       <td align="left" style="text-align: left"><tt>z = z + y</tt></td>
#     </tr>
# </table>
#
# The values for <i>A</i>, <i>B</i> and <i>C</i> are the only values that vary between the parts, and, in fact, between puzzle inputs for everyone participating in AoC. Moreover, <i>A</i> depends on <i>B</i>; it is 26 only if <i>B</i> is a positive number (zero or greater).
#
# So, expressed as Python, the sections come down to:
#
# ```python
# def section(input, z, B, C):
#     x = z % 26 + B
#     if B >= 0:
#         z //= 26
#     if input != x:
#       z = z * 26 + input + C
#     return z
# ```
#
# From this, you can see that `z` will never be negative, and can only be 0 if, by the time we reach the last block, it is smaller than 26 (as `z //= 26` is the only point where `z` decreases, and only for values smaller than 26 would floor division give 0 there).
#
# The other conclusion we can make is that the outcome _branches_, based on the values of the input digits; at least, for those blocks where `B` is not larger than 9, as that would _guarantee_ that `input` is not equal to `x`. *One* of those branches will end up being zero, for a given set of conditions. Our job will be to find that set of conditions, because from that we can deduce the permissible range of each input variable.
#
# Finally, I note that only the _condition_ has to rely on modulo operations. If we play our cards right, then each variant of the expression being processed is going to be a [linear polynomial](https://en.wikipedia.org/wiki/Polynomial#linear_polynomial) with all positive [coefficients](https://en.wikipedia.org/wiki/Coefficient). Put differently, it'll be a rather simple $ai_0 + bi_1 + ci_2 + ... + zi_n$ expression, something we can make use of when trying to simplify expressions or prune branches.
#
# ## Using Sympy to track branching
#
# I decided to solve this problem by using [sympy](https://www.sympy.org/) to parse and solve the equation, as it'll let us combine the parts into a single equation and track branching. Braching is tracked via [`Piecewise` objects](https://docs.sympy.org/latest/modules/functions/elementary.html#sympy.functions.elementary.piecewise.Piecewise), and Sympy will automatically eliminate branches if it recognises the condition always applies or can never be met. Sympy can do this because keeps track of various properties of the symbols (variables) involved, such as the fact that all our inputs are going to be non-zero positive integers.
#
# However, there are a few challenges to overcome:
#
# - The ALU division operation needs to floor the outcome (if the signs of the operands are the same. truncate towards zero. We don't have to worry about negative numbers however, as the only division that takes place is either by 1 or by 26. We can't just use `floor()` here, because then Sympy generally won't be able to simplify the expression further.
# - The expresion rapidly grows to a size where manipulating it gets _very_ slow, so we need to find strategies to simplify it further than the standard Sympy simplifcation methods can achieve.
#
# ### Recasting division to floor the result
#
# The first problem can be solved by redefining the operation in terms that Sympy can process and even simplify. Floor division can de defined by first subtracting the remainder from the dividend before dividing:
#
# $$
# \lfloor \frac a b \rfloor = \frac {a - (a \mod b)} {b}
# $$
#
# Sympy knows how to handle modulo operations, so that's what we'll use to translate the `div` operator.
#
# We don't have to worry about rounding towards negative infinity, as for this puzzle, neither operand is ever smaller than zero. However, should the need arise, you can expand on this by testing for either $a$ or $b$ being negative:
#
# $$
# \begin{cases}
# \frac {a + (-a \mod b)} {b} & \text{if } a < 0 \land b > 0 \\
# \frac {a + (a \mod -b)} {b} & \text{if } a > 0 \land b < 0 \\
# \frac {a - (a \mod b)} {b} & \text{if } ab >= 0
# \end{cases}
# $$
#
# In Sympy, you can then model those multiple cases in a `Piecewise()` object. I didn't bother with this however, as the first two cases would simply be dropped instantly, anyway.
#
# ### Eliminating modulo operations
#
# Next, we can assist Sympy by eliminating modulo operations if we know the left-hand $a$ value is always going to be lower than the right-hand value $b$, which in our case is always going to be 26 (either from the `mod x 26` operation form line 4, or one of the `div z 26` operations on line 5).
#
# One way we could do this is to try and test the expression $a < b$ for each free symbol (input variable) in $a$ using the [`solveset()` function](https://docs.sympy.org/latest/modules/solvers/solveset.html#sympy.solvers.solveset.solveset) and a [`Range()` set](https://docs.sympy.org/latest/modules/sets.html#sympy.sets.fancysets.Range) as the domain. If this produces the same range of values again, we know that for all possible values for that input, the modulo operation will not have any effect and can be eliminated.
#
# However, because the left-hand-side expression in our modulo operations are always linear polynomials with positive coefficients (only `+` and `*` operations), you can instead substitute all input symbols with $9$ to determine the highest possible value. If the result is then lower than $b$, we know the modulo can be removed.
#
# ### Collapsing equality tests
#
# We can do something similar for equality tests, but this time we'll have to stick with `solveset()`, as the alternative would have to be testing each possible combination of the inputs involved.
#
# For each free $symbol$ in the $expression$ (each an input variable), test what `solveset(expression, symbol, Range(1, 10))` returns. This will give us a new set, the set of all values for that that symbol for which the outcome will be true. There are three possible outcomes:
#
# * The empty set: the equality test is _always false_, regardless of what the value is for that input.
# * The `Range(1, 10)` set: the equality test is _always true_, for all possible inputs.
# * Some other set, which is always a subset of the input domain.
#
# For the first two outcomes, the equality can be replaced by a boolean constant.
#
# ### Eliminating branches
#
# From the above analysis we know that $z$ can only ever be zero if, by the time we reach the very last section, $z$ is a value between 0 and 25 inclusive, and the only way $z$ is going to get there is by division by 26. If you count the number times $z$ is divided by 26, you can test any given branch by substituting all inputs with 1 and seeing if the result is equal to or greater than 26 raised to the power of the number of divisions that are still left.
#
# However, because we also eliminate branches that can never be taken (by collapsing equality tests), we can't know how many divisions will remain until we've parsed all the sections. So instead, we start with merging expressions that have already been simplified into a single branch into the current expression. The moment the merged expression still has two branches, we start a new set of expressions to merge.
#
# Once we then have a series of branched expressions, we count how many of these divide by 26, so we know the limit beyond which a given expression will no longer reach zero. Each branch will have one or more inequality conditions, in the form of `inputA - inputB != number`; the remaining branches need to be updated with the _inverse_ of those conditions, because these conditions are what limit the input values to a smaller range. If you end up with a _single_ branch, you've found a path to `z == 0`, so we need to keep trace of those conditions.
#
# ### Finding the minimum and maximum possible version numbers
#
# Merging all branches this way, results in a single `0` expression, and a single condition, a conjunction of equality tests. Each of those equality tests can be turned into a maximum value for the smaller of the two inputs. E.g., the expression `inputA - inputB = 5` can only be true if `inputB` is smaller than `inputA`, and can, at most, be `4`. If it was `5`, then the condition would have matched one of the already eliminated branches, ones that don't reach zero!
#
# To determine the maximum version number, then, start with a list of all `9` digits, and adjust those for inputs that must be smaller to meet the conditions they are involved in. For part two, do the same with a list of `1` digits, adjusted upwards to keep the other input large enough for the condition to apply.
#

# +
from __future__ import annotations

from functools import cached_property, reduce, singledispatchmethod
from operator import add, mod, mul
from typing import Callable, Final

import sympy as sy
from sympy import piecewise_fold, simplify_logic, solveset


OPCODES: Final[dict[str, Callable[[sy.Basic, sy.Basic], sy.Basic]]] = {
    "add": add,
    "mul": mul,
    "div": lambda a, b: (a - a % b) / b,  # we can assume a * b >= 0, always.
    "mod": mod,
    "eql": lambda a, b: sy.Piecewise((1, sy.Eq(a, b)), (0, True)),
}
Z: Final[sy.Symbol] = sy.Symbol("z", integer=True, negative=False)



class MONAD:
    _condition: sy.Boolean = sy.S.true
    _limit: int = 0
    _min: int | None = None
    _max: int | None = None

    def __init__(self, instructions: str) -> None:
        self._parse(instructions)

    def _parse(self, instructions: str) -> None:
        reg: dict[str, sy.Basic] = dict.fromkeys("xyz", sy.S.Zero)
        ws: list[sy.Symbol] = []
        branches: list[sy.Basic] = [sy.S.Zero]

        for block in instructions.split("inp w\n")[1:]:
            w = sy.Symbol(f"w{len(ws)}", integer=True, positive=True, nonzero=True)
            ws.append(w)
            reg |= {"w": w, "z": Z}
            for line in block.splitlines():
                instr, target, *args = line.split()
                args = [reg[p] if p in reg else sy.Integer(p) for p in args]
                reg[target] = OPCODES[instr](reg[target], *args)
            if not branches[-1].is_Piecewise:
                reg["z"] = reg["z"].subs({Z: branches.pop()})
            expr = piecewise_fold(reg["z"]).replace(*self._replace_args)
            branches.append(expr)

        # combine all branched expressions into a single expression, while
        # removing branches that are never going to reach zero.
        expr = sy.S.Zero
        self._limit = 26 ** sum(1 for br in branches if br.has(sy.S.One / 26))
        for branch in branches:
            self._limit //= 26 if branch.has(sy.S.One / 26) else 1
            expr = piecewise_fold(branch.subs({Z: expr})).replace(*self._replace_args)
    
    def _find_extrema(self):
        """Turn the final 0 condition into boundaries for the 14 digits"""
        ws = sorted(self._condition.free_symbols, key=sy.default_sort_key)
        mins, maxs = [1] * len(ws), [9] * len(ws)
        for cond in self._condition.args:
            # each condition is an inequality between two inputs. It is always
            # in the form inputA - inputB == C so we only need to know the value
            # of C and the indexes of the input variables involved.
            w1, w2, diff = cond.lhs.args[0], -cond.lhs.args[1], cond.rhs.p
            if diff < 0:
                w1, w2, diff = w2, w1, -diff
            wi1, wi2 = ws.index(w1), ws.index(w2)
            mins[wi1], maxs[wi2] = max(mins[wi1], 1 + diff), min(maxs[wi2], 9 - diff)

        self._min = reduce(lambda a, b: a * 10 + b, mins)
        self._max = reduce(lambda a, b: a * 10 + b, maxs)
    
    @property
    def minimum(self) -> int:
        if self._min is None:
            self._find_extrema()
        return self._min
    
    @property
    def maximum(self) -> int:
        if self._max is None:
            self._find_extrema()
        return self._max

    @singledispatchmethod
    def _simplify(self, _: sy.Basic) -> sy.Basic | None:
        """Handler for simplification handlers via single dispatch
        
        Individual methods below are registered to simplify a specific Sympy
        object type.

        """
        return None

    @cached_property
    def _replace_args(
        self,
    ) -> tuple[Callable[[sy.Basic], bool], Callable[[sy.Basic], sy.Basic | None]]:
        """Argument pair for Expr.replace(), dispatching to the _simplify() method
        
        For each expression element for which the first callable returns True,
        sympy calls the second method, which in turn will call the registered
        hook method for the specific type of object.

        """
        # this is way harder than it should be, singledispatchmethod should
        # really add registry on the generated method directly. Access _simplify
        # via the class namespace so the descriptor protocol doesn't kick in,
        # so we can then access the dispatcher registry.
        dispatch_registry = vars(type(self))["_simplify"].dispatcher.registry 
        types = tuple(dispatch_registry.keys() - {object})
        return ((lambda a: isinstance(a, types)), self._simplify)

    @_simplify.register
    def _simplify_mod(self, mod: sy.Mod) -> sy.Basic | None:
        """Unwrap a modulo operation if a is always smaller than b"""
        (a, b), subs = mod.args, dict.fromkeys(mod.free_symbols, 9)
        if not mod.has(Z) and b.is_number and a.subs(subs) < b:
            return a
        return None

    @_simplify.register
    def _simplify_eq(self, eq: sy.Eq) -> sy.Basic | None:
        """Simplify an equality expression if it's always true or false"""
        for sym in eq.free_symbols - {Z}:
            match solveset(eq, sym, sy.Range(1, 10)):
                case sy.EmptySet:
                    return sy.S.false
                case sy.Range(1, 10):
                    return sy.S.true
        return None

    @_simplify.register
    def _simplify_ne(self, ne: sy.Ne) -> sy.Basic | None:
        """Simplify an inequality expression if it's always true or false"""
        if (result := self._simplify_eq(~ne)) is not None:
            return ~result
        return None

    @_simplify.register
    def _simplify_piecewise(self, pw: sy.Piecewise) -> sy.Basic | None:
        """Eliminate branches that will exceed the limit"""
        limit = self._limit
        if not limit:
            return None
        elim, new_pairs, subs = sy.S.true, [], dict.fromkeys(pw.free_symbols, 1)
        for br, cond in pw.args:
            if br.subs(subs) >= limit:
                elim &= ~cond
                continue
            new_pairs.append((br, cond))
        new_pairs = [(e, simplify_logic(c & elim)) for e, c in new_pairs]
        if len(new_pairs) == 1:
            # all other branches eliminated; update the condition that applies
            # to this single branch.
            (expr, cond), = new_pairs
            self._condition &= cond
            return expr
        return pw.func(*new_pairs)



# +
import aocd

alu_instructions = aocd.get_data(day=24, year=2021)
expr = MONAD(alu_instructions)
print("Part 1:", expr.maximum)
print("Part 2:", expr.minimum)


StarcoderdataJupyterScriptsDedupFiltered