from typing import List, Optional

from cs11puzzles.grid import ORTHO_DIRS, VEC_TO_DIR_NAME, Grid, Point
from cs11puzzles.str_helpers import ortho_dirs_sym, str_matrix_to_text
from genres.moonsun.moonsun_util import MOON, SUN, MoonSunInstance, load_puzzle
from z3.z3 import Bool, Int, ModelRef, sat


class MoonSunCell:
    def __init__(self, p: Point, u: int):
        self.point = p
        self.uniquifier = u

        # Draw connections between cells
        self.dirs = {
            d: Bool(f"moonsun-dir-u{u}-d{VEC_TO_DIR_NAME[d]}-r{p.y}-c{p.x}")
            for d in ORTHO_DIRS
        }

        self.order = Int(f"moonsun-order-u{u}-r{p.y}-c{p.x}")
        self.moon_region = Bool(f"moonsun-moon_reg-u{u}-r{p.y}-c{p.x}")

    def cell_to_str(self, m: ModelRef):
        return ortho_dirs_sym(
            **{VEC_TO_DIR_NAME[vec]: bool(m[val]) for vec, val in self.dirs.items()}
        )


def solve_moonsun(puzzle: MoonSunInstance) -> Optional[List[List[str]]]:
    areas = puzzle.areas
    moonsun = puzzle.moonsun

    height = len(areas)
    width = len(areas[0])
    num_cells = height * width

    g = Grid[MoonSunCell](height, width, MoonSunCell)

    # TODO: constraints

    if g.solver.check() == sat:

        def label_fn(cell: MoonSunCell, m: ModelRef):
            return cell.cell_to_str(m)

        return g.to_str_matrix(label_fn)

    else:
        print(g.solver.unsat_core())

    return None


if __name__ == "__main__":
    filename = "tests/moonsun/data/moonsun_1.txt"
    puzzle = load_puzzle(filename)
    sol = solve_moonsun(puzzle)

    if sol:
        print(str_matrix_to_text(sol))
    else:
        print("No solution.")