# coding: utf-8 import os, sys
from z3 import *

# get the playground information
if len(sys.argv) != 2:
    print("Usage: python3 colours.py <map-file>")
    sys.exit(0)

with open(sys.argv[1]) as f:
    playground = f.read()
playground = playground.strip().split("\n")

# get the playground size
size_y = len(playground)
assert(size_y != 0)
size_x = len(playground[0])
assert(size_x != 0)
assert(all([len(p) == size_x for p in playground]))
# get the number of regions
regions = set("".join(playground))

################################### Colours ####################################

# create the solver
solver = Solver()

# todo: create a datatype Colour
# ...
# todo: declare four valid colours
# ...
# hint: don't forget to call "create" for the datatype to get a Sort
# ...

# todo: create a colour variable for each playground cell
# hint: you can do this with Const("var_name", YourSort)
# hint: use something like the coordinates as part of the variable name
p_colours = [[None for _j in range(size_x)] for _j in range(size_y)]
for i in range(size_y):
    for j in range(size_x):
        # p_colours[i][j] = ... replace this with your code
        pass

print("All regions found in this map: '{}'.".format(",".join(regions)))
for r in regions:
    # todo: find all cells that belong to the same region
    region_vars = []
    for i in range(size_y):
        for j in range(size_x):
            pass # ... add the variable to region_vars here

    # todo: make all colours in the same region identical
    # hint: use the transitivitiy of equality x1 = x2 & x2 = x3 -> x1 = x3
    for c in range(len(region_vars) - 1):
        pass # ... replace this with your code ...


def get_neighbours(i, j):
    regions = []
    colours = []
    if i > 0:
        colours.append( p_colours[i-1][j])
        regions.append(playground[i-1][j])
    if i < size_y - 1:
        colours.append( p_colours[i+1][j])
        regions.append(playground[i+1][j])
    if j > 0:
        colours.append( p_colours[i][j-1])
        regions.append(playground[i][j-1])
    if j < size_x - 1:
        colours.append( p_colours[i][j+1])
        regions.append(playground[i][j+1])
    return zip(regions, colours)


# todo:  if two neighboring cells are from different regions
#        then they must have different colours
for i in range(size_y):
    for j in range(size_x):
        cell_r = playground[i][j]
        cell_c = p_colours[i][j]
        for r, c in get_neighbours(i, j):
            pass # ... replace this with your code ...

# call the solver and check satisfiability
res = solver.check()
if res != sat:
    print("unsat")
    sys.exit(1)

# todo make a translation to strings "1", "2", "3", "4"
def translate(s):
    # hint: do something analogouns to the following line here
    # if s.sexpr() == "C1":     return "1"
    # ...
    return "0"

################################################################################

from colorama import Fore, Back, Style
cols = [Back.BLACK, Back.RED, Back.GREEN, Back.BLUE, Back.MAGENTA]
cols = [Style.BRIGHT + Fore.WHITE + c for c in cols]

# print the model as translated numbers
m = solver.model()
names = [x.name() for x in m.decls()]
for i in range(size_y):
    row = ""
    for j in range(size_x):
        pc = p_colours[i][j]
        invalid = (pc is None or pc.decl() not in m.decls())
        t = "0" if invalid else translate(m[pc])
        if sys.stdout.isatty(): t = cols[int(t)] + t + Style.RESET_ALL
        row += t
    print(row)