from pysmt.shortcuts import Symbol, Solver, Int, Equals, NotEquals
from pysmt.typing import INT


# Given a row of symbols (e.g. [S E N D]), make a term representing the number
# they represent (e.g. S * 10^3 + E * 10^2 + N * 10^1 + D * 10^0)
def make_sum(row):
    fac = 1
    res = Int(0)

    for s in reversed(row):
        res += s * Int(fac)
        fac *= 10

    return res


def solve_sum(first_row, second_row, result_row):
    def symbolize(row):
        return [Symbol(c, INT) for c in row.upper()]

    first = symbolize(first_row)
    second = symbolize(second_row)
    result = symbolize(result_row)

    symbols = set(first + second + result)

    with Solver() as solver:
        # All symbols represent digits
        for s in symbols:
            solver.add_assertion(0 <= s)
            solver.add_assertion(s <= 9)

        # Different symbols are different digits
        for s0 in symbols:
            for s1 in symbols:
                if s0.symbol_name() != s1.symbol_name():
                    solver.add_assertion(NotEquals(s0, s1))

        # The leading digits are nonzero
        solver.add_assertion(NotEquals(first[0], Int(0)))
        solver.add_assertion(NotEquals(second[0], Int(0)))
        solver.add_assertion(NotEquals(result[0], Int(0)))

        # The arithmetic works
        solver.add_assertion(Equals(make_sum(first) + make_sum(second), make_sum(result)))

        if solver.solve():
            print("Solution:")
            print(solver.get_model())
        else:
            print("No solution")


#   9 5 6 7
#   1 0 8 5
# 1 0 6 5 2
solve_sum("send", "more", "money")

#   2 1
#   8 1
# 1 0 2
solve_sum("to", "go", "out")

# No solution
solve_sum("abc", "def", "gh")

# No solution
solve_sum("a", "a", "a")