#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
#  einstein_puzzle.py
#  
#  Copyright 2012 jackyspy <[email protected]>
#  
#  This program is free software; you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation; either version 2 of the License, or
#  (at your option) any later version.
#  
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#  
#  You should have received a copy of the GNU General Public License
#  along with this program; if not, write to the Free Software
#  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
#  MA 02110-1301, USA.
#  
#
 
from itertools import permutations
 
color    = ['Green', 'White', 'Yellow', 'Blue', 'Red']
national = ['English', 'Swedish', 'Norwegian', 'German', 'Dane']
drink    = ['Tea', 'Coffee', 'Milk', 'Beer', 'Water']
cigar    = ['PallMall', 'DunHill', 'Blends', 'BlueMaster', 'Prince']
pet      = ['Bird', 'Dog', 'Horse', 'Cat', 'Fish']
choices  = (color, national, drink, cigar, pet)
 
constraints=[
    ( 'Red',    'English', None,    None,        None   ), # The Englishman lives in the red house.
    ( None,     'Swedish', None,    None,        'Dog'  ), # The Swede keeps dogs.
    ( None,     'Dane',   'Tea',    None,        None   ), # The Dane drinks tea.
    ( 'Green',  None,     'Coffee', None,        None   ), # The owner of the green house drinks coffee.
    ( None,     None,     None,     'PallMall',  'Bird' ), # The Pall Mall smoker keeps birds.
    ( 'Yellow', None,     None,    'DunHill',    None   ), # The owner of the yellow house smokes Dunhills.
    ( None,     None,     'Beer',  'BlueMaster', None   ), # The man who smokes Blue Masters drinks bier.
    ( None,     'German', None,    'Prince',     None   ), # The German smokes Prince.
    ]
 
neighbors = [
    ( (3, 'Blends'),    (4, 'Cat')   ), # The Blend smoker has a neighbor who keeps cats.
    ( (3, 'DunHill'),   (4, 'Horse') ), # The man who keeps horses lives next to the Dunhill smoker.
    ( (1, 'Norwegian'), (0, 'Blue')  ), # The Norwegian lives next to the blue house.
    ( (3, 'Blends'),    (2, 'Water') ), # The Blend smoker has a neighbor who drinks water.
    ]
 
 
def find_iter(func, iterable):
    for i in iterable:
        if func(i):
            return i
 
 
def select_item(curr_person, idx, selected_persons):
    person_len = len(curr_person)
    assert idx < person_len
 
    # selected by the constraints already
    if curr_person[idx]:
        if idx == person_len - 1:
            yield tuple(curr_person)
        else:
            for i in select_item(curr_person, idx + 1, selected_persons):
                yield tuple(i)
        return
 
    for s in choices[idx]:
        if find_iter(lambda x: x[idx] == s, selected_persons):
            # chosen by others
            continue
 
        new_curr = curr_person[:]
        new_curr[idx] = s
 
        cons = find_iter(lambda x: x[idx] == s, constraints)
        if cons:
            is_conflict = False
 
            for i in range(person_len):
                if i < idx:
                    # check if conflict with prior
                    if cons[i]:
                        is_conflict = True
                        break
                elif i > idx:
                    # set the value indicated by constraints
                    if cons[i]:
                        new_curr[i] = cons[i]
 
            if is_conflict:
                continue
 
 
        if idx == person_len - 1:
            yield tuple(new_curr)
        else:
            for i in select_item(new_curr, idx + 1, selected_persons):
                yield tuple(i)
 
 
def select_all():
    for a1 in select_item(['Red', 'English', None, None, None], 1, []):
        for a2 in select_item(['Green', None, 'Coffee', None, None], 1, [a1]):
            for a3 in select_item(['Yellow', None, None, 'DunHill', None], 1, [a1, a2]):
                for a4 in select_item(['Blue', None, None, None, None], 1, [a1, a2, a3]):
                    for a5 in select_item(['White', None, None, None, None], 1, [a1, a2, a3, a4]):
                        yield [a1, a2, a3, a4, a5]
 
 
def check_result(res):
    def get_item(idx, val):
        return find_iter(lambda x: x[idx] == val, res)
 
    def get_house(idx, val, house):
        return house[get_item(idx,val)]
 
    def set_house(idx, val, house, house_id):
        item = get_item(idx, val)
 
        if house[item]:
            return house[item] == house_id # check if conflict with prior
 
        # if house_id used by other person, shout
        if house_id in house.values():
            return False
 
        house[item] = house_id
        return True
 
    def is_neighbor(cond1, cond2, house):
        house_id1 = get_house(cond1[0], cond1[1], house)
        house_id2 = get_house(cond2[0], cond2[1], house)
        return house_id1 - house_id2 in (1, -1)
 
    def check_neighbors(neighbors, house):
        for cond1, cond2 in neighbors:
            if not is_neighbor(cond1, cond2, house):
                return False
        return True
 
    house = {}
    for i in res:
        house[i] = 0 # init with house_id unset
 
    if not set_house(2, 'Milk', house, 3): # The man in the center house drinks milk.
        return False
 
    if not set_house(1, 'Norwegian', house, 1): # The Norwegian lives in the first house.
        return False
 
    for i in range(1,5): # since green cannot be the last one
        house2 = house.copy()
 
        if not set_house(0, 'Green', house2, i):
            continue
        if not set_house(0, 'White', house2, i + 1): # The green house is just to the left of the white one.
            continue
 
        for unset_ids in permutations(set(range(1,6)) - set((1, 3, i, i + 1))): # get full permutations of rest house ids
            house3 = house2.copy()
 
            # fill the houses without id one by one
            k = 0
            for j in res:
                if not house3[j]:                    
                    house3[j] = unset_ids[k]
 
                    k += 1
                    if k == len(unset_ids):
                        break
 
            if check_neighbors(neighbors, house3):
                return house3
 
    return False
 
 
 
def main():
    for i in select_all():
        res = check_result(i)
        if res:
            print '='*20 + ' solved ' + '='*20
            for _, person in sorted((v,k) for k,v in res.items()):
                print person
 
if __name__ == '__main__':
	main()