#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from __future__ import (absolute_import, division, print_function,
                        unicode_literals)
import sys
import logging
import z3

log = logging.getLogger(__name__)


def guestpuzzle():
    """Main script.

    """
    # Define each room as a number.
    rooms = [11, 12, 13, 14, 15, 16, 17,
             21, 22, 23, 24, 25, 26, 27, 28]

    showers = [13, 17, 14, 15, 25, 28, 24, 23]
    baths = [11, 12, 16, 27, 26, 21, 22]

    # These rooms have a balcony (24 is broken)
    balcony = [27, 25, 23, 21]

    # Dir facing
    south_facing = [14, 12, 28, 26]
    east_facing = [16, 23, 22]

    # Twins available
    twins = [17, 12, 24, 21, 22]

    # Noisy rooms for disco
    disco = [11, 16, 24, 23]

    # View of pool. Not 1.4 or 2.8 due to tree
    pool_view = [26, 12, 21]

    # Primes
    primes = [11, 13, 17, 23]

    canhold2 = [17, 15, 13, 11, 16, 25, 23, 22, 26, 14, 27, 28, 12, 24, 21]
    canhold3 = [14, 27, 28, 12, 24, 21]
    canhold4 = [12, 24, 21]

    ground_floor = [11, 12, 13, 14, 15, 16, 17]

    # Now do the bookings!
    constraints = []
    bookings = []
    columns = []

    # Wilcotts want a balcony. There's 3 of them.
    booking1 = z3.Int("Wilcotts")
    constraints.append(isin(booking1, balcony))
    constraints.append(isin(booking1, canhold3))
    bookings.append(booking1)
    columns.append("Wilcotts")

    # Carnfore want a south facing room. There's 3 of them.
    booking2 = z3.Int("Carnfore")
    constraints.append(isin(booking2, south_facing))
    constraints.append(isin(booking2, canhold3))
    bookings.append(booking2)
    columns.append("Carnfore")

    # Osmonds want a bath. There's 4 of them.
    booking3 = z3.Int("Osmond")
    constraints.append(isin(booking3, baths))
    constraints.append(isin(booking3, canhold4))
    bookings.append(booking3)
    columns.append("Osmond")

    # Wallers have no requests. There's 2 of them.
    booking4 = z3.Int("Waller")
    constraints.append(isin(booking4, canhold2))
    bookings.append(booking4)
    columns.append("Waller")

    # PoitierVert want a twin.
    booking5 = z3.Int("PoitierVert")
    constraints.append(isin(booking5, twins))
    constraints.append(isin(booking5, canhold2))
    bookings.append(booking5)
    columns.append("PoitierVert")

    # Jennings want no balcony. x4. Also quiet.
    booking6 = z3.Int("Jennings")
    constraints.append(z3.Not(isin(booking6, balcony)))
    constraints.append(z3.Not(isin(booking6, disco)))
    constraints.append(isin(booking6, canhold4))
    bookings.append(booking6)
    columns.append("Jennings")

    # Tudor/Smith want no east. x2. And shower
    booking7 = z3.Int("TudorSmith")
    constraints.append(z3.Not(isin(booking7, east_facing)))
    constraints.append(isin(booking7, showers))
    constraints.append(isin(booking7, canhold2))
    bookings.append(booking7)
    columns.append("TudorSmith")

    # FoxMoran don't want primes x2.
    booking9 = z3.Int("FoxMoran")
    constraints.append(isin(booking9, canhold2))
    constraints.append(z3.Not(isin(booking9, primes)))
    bookings.append(booking9)
    columns.append("FoxMoran")

    # Proctors want a balcony. x2.
    booking10 = z3.Int("Proctor")
    constraints.append(isin(booking10, canhold2))
    constraints.append(isin(booking10, balcony))
    bookings.append(booking10)
    columns.append("Proctor")

    # Vieras want ground floor. x3
    booking11 = z3.Int("Vieras")
    constraints.append(isin(booking11, canhold3))
    constraints.append(isin(booking11, ground_floor))
    bookings.append(booking11)
    columns.append("Vieras")

    # WilliamsAdams want pool view. x2
    booking12 = z3.Int("WilliamsAdams")
    constraints.append(isin(booking12, canhold2))
    constraints.append(isin(booking12, pool_view))
    bookings.append(booking12)
    columns.append("WilliamsAdams")

    # InglebyCooper want twin. x2
    booking13 = z3.Int("InglebyCooper")
    constraints.append(isin(booking13, canhold2))
    constraints.append(isin(booking13, twins))
    bookings.append(booking13)
    columns.append("InglebyCooper")

    # Inglebys want to be adjoined to InglebyCooper. x2
    booking14 = z3.Int("Inglebys")
    constraints.append(isin(booking14, canhold2))
    constraints.append(adjacent(booking13, booking14))
    bookings.append(booking14)
    columns.append("Inglebys")

    # Cheshams don't care. x3
    booking15 = z3.Int("Cheshams")
    constraints.append(isin(booking15, canhold3))
    bookings.append(booking15)
    columns.append("Cheshams")

    # Grahams want no noise. x2.
    booking8 = z3.Int("Graham")
    constraints.append(isin(booking8, canhold2))
    constraints.append(z3.Not(isin(booking8, disco)))
    # No kids
    constraints.append(z3.Not(noiseadjacent(booking8, booking3)))
    constraints.append(z3.Not(noiseadjacent(booking8, booking4)))
    constraints.append(z3.Not(noiseadjacent(booking8, booking6)))
    constraints.append(z3.Not(noiseadjacent(booking8, booking11)))  # ???
    bookings.append(booking8)
    columns.append("Graham")

    # All bookings are distinct.
    constraints.append(z3.Distinct(*bookings))

    # Show our constraints
    log.debug("Constraints: %s", constraints)

    # Try and solve it!
    s = z3.Solver()
    s.add(constraints)

    sol_num = 0

    models = []

    while s.check() == z3.sat:
        sol_num += 1
        log.info("Solution %d found!", sol_num)

        model = s.model()
        log.info("%s", model)

        models.append(model)

        # Add a constraint such that we don't regenerate the same model.
        s.add(z3.Or(booking1 != model[booking1],
                    booking2 != model[booking2],
                    booking3 != model[booking3],
                    booking4 != model[booking4],
                    booking5 != model[booking5],
                    booking6 != model[booking6],
                    booking7 != model[booking7],
                    booking8 != model[booking8],
                    booking9 != model[booking9],
                    booking10 != model[booking10],
                    booking11 != model[booking11],
                    booking12 != model[booking12],
                    booking13 != model[booking13],
                    booking14 != model[booking14],
                    booking15 != model[booking15]))

    if sol_num == 0:
        log.error("Solution not found :(")
    else:
        with open("guestpuzzle3.csv", "wt") as f:
            f.write(",".join(columns))
            f.write("\n")

            for model in models:
                fields = [str(model[booking]) for booking in bookings]
                f.write(",".join(fields))
                f.write("\n")

    return ScriptRC.SUCCESS


def isin(z3int, intlist):
    constraints = [z3int == x for x in intlist]
    return z3.Or(*constraints)


def adjacent(first, second):
    constraints = [
        # First floor
        z3.And(first == 17, second == 15),
        z3.And(first == 15, second == 17),
        z3.And(first == 15, second == 13),
        z3.And(first == 13, second == 15),
        z3.And(first == 13, second == 11),
        z3.And(first == 11, second == 13),
        z3.And(first == 14, second == 12),
        z3.And(first == 12, second == 14),

        # Second floor
        z3.And(first == 27, second == 25),
        z3.And(first == 25, second == 27),
        z3.And(first == 25, second == 24),
        z3.And(first == 24, second == 25),
        z3.And(first == 24, second == 23),
        z3.And(first == 23, second == 24),
        z3.And(first == 23, second == 22),
        z3.And(first == 22, second == 23),
        z3.And(first == 28, second == 26),
        z3.And(first == 26, second == 28),
    ]
    return z3.Or(*constraints)


def noiseadjacent(first, second):
    constraints = [
        # First floor
        z3.And(first == 17, second == 15),
        z3.And(first == 17, second == 27),

        z3.And(first == 15, second == 17),
        z3.And(first == 15, second == 13),
        z3.And(first == 15, second == 27),
        z3.And(first == 15, second == 25),

        z3.And(first == 13, second == 15),
        z3.And(first == 13, second == 11),
        z3.And(first == 13, second == 25),

        z3.And(first == 11, second == 13),
        z3.And(first == 11, second == 25),
        z3.And(first == 11, second == 24),

        z3.And(first == 14, second == 12),
        z3.And(first == 14, second == 28),
        z3.And(first == 12, second == 14),
        z3.And(first == 12, second == 26),

        z3.And(first == 16, second == 22),

        # Second floor
        z3.And(first == 27, second == 25),
        z3.And(first == 27, second == 17),
        z3.And(first == 27, second == 15),

        z3.And(first == 25, second == 27),
        z3.And(first == 25, second == 24),
        z3.And(first == 25, second == 13),
        z3.And(first == 25, second == 15),
        z3.And(first == 25, second == 11),

        z3.And(first == 24, second == 25),
        z3.And(first == 24, second == 23),
        z3.And(first == 24, second == 11),

        z3.And(first == 23, second == 24),
        z3.And(first == 23, second == 22),

        z3.And(first == 22, second == 23),
        z3.And(first == 22, second == 16),

        z3.And(first == 28, second == 26),
        z3.And(first == 28, second == 14),
        z3.And(first == 26, second == 28),
        z3.And(first == 26, second == 12),
    ]
    return z3.Or(*constraints)


def main():
    """Main handling function. Wraps program."""

    # Set up basic logging
    logging.basicConfig(format="%(asctime)s %(levelname)-5.5s %(message)s",
                        stream=sys.stdout,
                        level=logging.DEBUG)

    # Run main script.
    try:
        rc = guestpuzzle()
    except Exception as e:
        log.exception(e)
        rc = ScriptRC.EXCEPTION

    log.info("Returning %d", rc)
    return rc


class ScriptRC(object):
    SUCCESS = 0
    FAILURE = 1
    EXCEPTION = 2


class ScriptException(Exception):
    pass


if __name__ == '__main__':
    sys.exit(main())