from lib import *

EXAMPLE = """1=-0-2
12111
2=0=
21
2=01
111
20012
112
1=-1=
1-12
12
1=
122
"""

def snafu_to_dec(n: str) -> int:
    r = 0
    for i, c in enumerate(reversed(n)):
        c = "=-012".index(c) - 2
        r += c * 5**i
    return r

def dec_to_snafu(n: int) -> str:
    r = ""
    while n != 0:
        rem = n % 5
        n //= 5
        if rem <= 2:
            r = str(rem) + r
        else:
            r = "   =-"[rem] + r
            n += 1
    return r

def dec_to_snafu_cs(n: int) -> str:
    import cpmpy as cp
    import numpy as np
    size = 20
    xs = cp.intvar(-2, 2, shape=size, name="xs")
    ns = np.array([5**i for i in range(size)])
    m = cp.Model(cp.sum(ns * xs) == n)
    m.solve()

    value = xs.value()
    assert value is not None
    # print(value)

    r = ""
    for v in value:
        if v >= 0:
            r += str(v)
        elif v == -1:
            r += "-"
        elif v == -2:
            r += "="
        else:
            assert False
    r = "".join(reversed(r))
    r = r.lstrip("0")
    return r

def solve(input: Input):
    res = 0
    for s in input.lines():
        d = snafu_to_dec(s)
        res += d
    try:
        assert dec_to_snafu_cs(res) == dec_to_snafu(res)
    except ModuleNotFoundError:
        print("Install 'cpmpy' to solve with constraint solver.")
    return dec_to_snafu(res)

def main():
    DAY_INPUT = "d25.txt"
    print("Example  1:", solve(Input(EXAMPLE)))
    print("Solution 1:", solve(Input(DAY_INPUT)))


if __name__ == "__main__":
    main()