#!/usr/bin/env python3

# what's the most sugar canes that can be planted with a single water source?
# restriction: all plants must be at the same height.
# restriction: the water source must be at that same height.
# this version is suitable for use with dispensers,
# and requires no additional preparation.

from z0 import Goal, SolverFor, Unary, BitVec, Bool, \
               If, IfUnary, Implies, And, Or, Not, Totalizer, sat, unsat, \
               binreduce

If, IfBitVec = IfUnary, If

import sys

# configure terminal colors

def esc(x):
    return "\x1B[" + str(x) + "m"

color_none = esc(0) + esc(90)
color_plant = esc(102) + esc(30)
color_level_lo = esc(44) + esc(30)
color_level_hi = esc(104) + esc(30)
color_level_zero_lo = esc(100) + esc(34)
color_level_zero_hi = esc(100) + esc(92)
color_reset = esc(0)

def get(el):
    if el is None:
        return None
    elif isinstance(el, Bool):
        return model[el.name]
    elif isinstance(el, BitVec):
        return model[el.name]
    elif isinstance(el, Unary):
        return model[el.name].bit_length()
    else:
        assert False, el
    return 0

def render(level, plant, block):
    if plant is None:
        return color_none + ":"
    if level is None:
        return color_none + ";"  # erroneous
    if plant == 1:
        return color_plant + "X"
    if block == 1:  # upper
        if level == 0:
            return color_level_zero_hi + "O"
        else:
            return color_level_hi + str(level)
    else:  # lower
        if level == 0:
            return color_level_zero_lo + "0"
        else:
            return color_level_lo + str(level)

def render_short(short):
    if short is None:
        return color_none + ":"
    if short == 0:
        return color_level_zero_lo + "0"
    else:
        return color_level_lo + str(short)

###

start_level = 8
size = 1 + 2 * (start_level - 2 +  # upper
                start_level + 1)   # lower
halfsize = size // 2

if 1:
    symmetry = 4  # rotational
    cross_heuristic = True
    no_preparation = True
    previous = 200  # best: 240 without preparation, 260 with.
    level_max = start_level
    max_upper = start_level - 1

else:
    symmetry = 4
    cross_heuristic = True
    no_preparation = False
    previous = 0
    level_max = start_level #+ 1
    max_upper = 1

def mandist(x, y):
    return abs(x - halfsize) + abs(y - halfsize)

def make2d(f, w=size, h=size):
    return [[f(x, y) for x in range(w)] for y in range(h)]

# generate the unobstructed diamond shape by manhattan distances.
usable = make2d(lambda x, y: mandist(x, y) <= halfsize)

usables = sum(1 if t else 0 for row in usable for t in row)
plant_len = usables.bit_length()

def makeuse(f, usable=usable):
    return make2d(lambda x, y: f(x, y) if usable[y][x] else None)

# since the water level is strictly decreasing,
# we can skip generating equations for blocks where usable[y][x] == False.
xy_fmt = "{:02X}{:02X}" if symmetry == 4 else "x{:02}y{:02}"
xyi_fmt = "{:02X}{:02X}u{:01X}" if symmetry == 4 else "x{:02}y{:02}u{:01X}"
levels = makeuse(lambda x, y: Unary("L" + xy_fmt.format(x, y), level_max))
blocks = makeuse(lambda x, y: Bool("B" + xy_fmt.format(x, y)))
plants = makeuse(lambda x, y: Bool("P" + xy_fmt.format(x, y)))

# find the shortest path down to the lower level by flowing backwards.
# this needs only a smaller section than the previous variables.
short_usable = makeuse(lambda x, y: (no_preparation and
                                     mandist(x, y) <= start_level))
shorts = makeuse(lambda x, y: Unary("S" + xy_fmt.format(x, y), level_max),
                 short_usable)

arrays = (levels, blocks, plants, shorts)

if symmetry == 1:
    pass

elif symmetry == 2:
    for array in arrays:
        for y_a in range(halfsize + 1):
            y_b = size - y_a - 1
            for x_a in range(size):
                x_b = size - x_a - 1
                array[y_b][x_b] = array[y_a][x_a]

elif symmetry == 4:
    middle = halfsize + 1
    for array in arrays:
        for y in range(middle):
            for x in range(middle):
                # bottom right corner:
                # equivalent to mirroring both axes.
                array[size - y - 1][size - x - 1] = array[y][x]
                # bottom left corner:
                # as x increases, y decreases.
                # as y increases, x increases.
                array[size - x - 1][y] = array[y][x]
                # top right corner:
                # as x increases, y increases.
                # as y increases, x decreases.
                array[x][size - y - 1] = array[y][x]

else:
    raise Exception(f"unsupported value of symmetry: {symmetry}")

def check(x, y):
    return x >= 0 and y >= 0 and x < size and y < size and usable[y][x]

def maybe(x, y):
    return (levels[y][x], blocks[y][x], shorts[y][x]) if check(x, y) else None

def xyiter(width, height):
    for y in range(height):
        for x in range(width):
            yield x, y

###

g = Goal()

if symmetry == 1:
    sym_width, sym_height = size, size
elif symmetry == 2:
    sym_width, sym_height = size, (size + 1) // 2
elif symmetry == 4:
    sym_width, sym_height = (size + 1) // 2, (size + 1) // 2

plant_pairs = []

for x, y in xyiter(sym_width, sym_height):
    if not usable[y][x]:
        continue

    level_var = levels[y][x]
    block_var = blocks[y][x]
    plant_var = plants[y][x]
    short_var = shorts[y][x]

    for var in (level_var, short_var):
        if var is None:
            continue

        # bound the water levels.
        g.add(var <= start_level)

    distance = mandist(x, y)
    # the center block is the water source block.
    centered = distance == 0

    # plants must be on blocks.
    g.add(Implies(plant_var, block_var))

    # plants cannot be placed in water!
    g.add(Implies(plant_var, level_var == 0))

    if cross_heuristic and (x == halfsize or y == halfsize):
        # 8765432 876543210
        # 0123456 789ABCDEF
        step_kink = max_upper
        if distance < step_kink:
            g.add(level_var == start_level - distance)
            g.add(block_var == True)
            g.add(plant_var == False)
        elif distance == start_level + step_kink:
            g.add(level_var == 0)
            g.add(block_var == True)
            g.add(plant_var == True)
        else:
            g.add(level_var == max(start_level + step_kink - distance, 0))
            g.add(block_var == False)
            g.add(plant_var == False)
    elif centered:
        # configure the water source.
        g.add(level_var == start_level)
        g.add(block_var == True)
        g.add(plant_var == False)

    # take account for symmetry.
    plant_worth = 1
    if symmetry == 2:
        plant_worth = 2 if y != halfsize else 1
    elif symmetry == 4:
        if x != halfsize:
            plant_worth *= 2
        if y != halfsize:
            plant_worth *= 2

    plant_pairs.append((plant_var, plant_worth))

    # gather information on cardinal directions.
    neighbors = [maybe(x - 1, y),  # west, "left"
                 maybe(x + 1, y),  # east, "right"
                 maybe(x, y - 1),  # north, "up"
                 maybe(x, y + 1)]  # south, "down"
    neighbors = [n for n in neighbors if n is not None]

    # gather information about neighboring cells.
    max_level_hi = 0
    max_level_lo = 0
    max_short = 0
    for n_level, n_block, n_short in neighbors:
        max_level_hi = Unary.max(max_level_hi, If(n_block, n_level, 0))
        max_level_lo = Unary.max(max_level_lo, If(n_block, 0, n_level))
        if n_short is not None:
            max_short = Unary.max(max_short, n_short)

    # some convenience...
    upper = block_var
    lower = Not(block_var)
    any_level = level_var != 0

    #assert max_upper == start_level - 1, "unimplemented (FIXME)"
    if max_upper < start_level:
        g.add(Implies(And(upper, any_level), level_var > start_level - max_upper))

    # this isn't necessary, but might assist the solver:
    # NOTE: should probably include max_level_lo > 0 as well.
    # NOTE: this turns out to be necessary for strictly correct no_preparation.
    g.add(Implies(And(upper, Not(any_level)), plant_var))

    # plants must be placed next to water alongside and below its block.
    g.add(Implies(plant_var, max_level_lo != 0))

    if not no_preparation:
        short = True
    elif short_var is None:
        short = False
    else:
        # find the shortest path down from the upper level.
        # annoyingly, this is how water flows in minecraft.

        # the usual flow, but backwards.
        g.add(Implies(And(upper, max_short != 0),
                      short_var == max_short - 1))

        g.add(Implies(lower, short_var == start_level))

        source_short_var = shorts[halfsize][halfsize]
        short = short_var == source_short_var + distance

    # trying some alternative definitions:
    if not centered:
        g.add(Implies(upper,
                      level_var == If(And(short, max_level_hi != 0),
                                      max_level_hi - 1, 0)))
    g.add(Implies(lower,
                  level_var == If(max_level_hi != 0, start_level,
                                  If(max_level_lo != 0,
                                     max_level_lo - 1, 0))))

if 1:
    plant_sum = BitVec("Plants", plant_len, 0)
    plant_scores = [IfBitVec(var, worth, 0, plant_len) for var, worth in plant_pairs]
    #plant_sum = sum(plant_scores)
    plant_sum = binreduce(plant_scores, lambda a, b: a + b)
else:
    plant_scores = sum(([var] * worth for var, worth in plant_pairs), [])
    plant_sum = Totalizer(plant_scores)

# TODO: optimal(?) sum reduction (assuming inputs are 1-bit) would be:
# A = 1 + 1 + 1 = 3 (2 bits)
# B = A + A + 1 = 3 + 3 + 1 = 7 (3 bits)
# C = B + B + 1 = 7 + 7 + 1 = 15 (4 bits)
# etc.

s = SolverFor("QF_BV")  # Quantifier-Free Bit-Vectors
s.add(g)

satisfiable = s.check()
print("check:", satisfiable)
sys.stdout.flush()

if len(sys.argv) > 1 and sys.argv[1] == "check":
    sys.exit(0)

while satisfiable == sat:
    #satisfiable = s.check(plant_sum > previous)
    g.add(plant_sum > previous)
    satisfiable = s.check()
    print("check:", satisfiable)

    if satisfiable != sat:
        break

    model = s.model()
    #previous = model['Plants']  # doesn't work since this is *before* summing
    previous = sum(worth if model[var.name] else 0
                   for var, worth in plant_pairs)
    print("plants:", previous)

    water = sum(1 if var is not None and model[var.name] > 0 else 0 for row in levels for var in row)
    print("water: ", water)

    #print(model)
    #for k, v in model.items():
    #    if v != 0:
    #        print(k, v)

    lines = ""
    for rows in zip(levels, plants, blocks, shorts):
        all_row = [[get(el) for el in row] for row in rows]
        lines += "".join(render(level, plant, block)
                         for level, plant, block, _ in zip(*all_row))
        if no_preparation:  # shortest distance debugging
            lines += " "
            lines += "".join(render_short(short)
                             for _, _, _, short in zip(*all_row))
        lines += color_reset + "\n"
    print(lines)

    sys.stdout.flush()