import shelve

class ImmutableDict:
    def __init__(self, d=None):
        self.d = {} if d is None else d
    def __getitem__(self, key):
        return self.d[key]
    def __hash__(self):
        return hash(tuple((k,v) for k,v in self.d.items()))
    def __eq__(self, other):
        return self.d == other.d
    def updated(self, key, value):
        d = self.d.copy()
        d[key] = value
        return ImmutableDict(d)

def other(s):
    if s == "X": return "O"
    if s == "O": return "X"
    raise Exception("Unknown player {}".format(s))

class State:
    def __init__(self, board, active_player, previous_moves=None):
        self.board = board
        self.active_player = active_player
        self.previous_moves = previous_moves or ImmutableDict({"X": None, "O": None})
    def iter_possible_moves(self):
        if self.winner() is not None:
            return
        right_column = [self.board[y][2] for y in range(3)]
        if any(cell == self.active_player for cell in right_column):
            can_slide = True
        elif all(cell is None for cell in right_column):
            can_slide = True
        else:
            can_slide = False
        for x in range(3):
            for y in range(3):
                if self.previous_moves[self.active_player] == (x,y):
                    continue
                if self.board[y][x] is None:
                    yield (False, x,y)
                if can_slide and (x == 0 or self.board[y][x-1] is None):
                    yield (True, x,y)
    def winner(self):
        raw_runs = "012, 345, 678, 036, 147, 258, 048, 246".split(", ")
        for raw_run in raw_runs:
            run = map(int, raw_run)
            cells = []
            for idx in run:
                x = idx % 3
                y = idx / 3
                cells.append(self.board[y][x])
            for c in "XO":
                if all(cell == c for cell in cells):
                    return c
        return None

    def apply(self, move):
        do_slide, x, y = move
        next_board = []
        for j in range(3):
            row = []
            for i in range(3):
                if x == i and y == j:
                    row.append(self.active_player)
                elif do_slide:
                    row.append(self.board[j][i-1] if i > 0 else None)
                else:
                    row.append(self.board[j][i])
            next_board.append(tuple(row))
        next_board = tuple(next_board)
        return State(
            next_board,
            other(self.active_player),
            self.previous_moves.updated(self.active_player, (x,y))
        )

    @staticmethod
    def from_idx(idx):
        x = idx % 3
        y = idx / 3
        return x,y
    @staticmethod
    def to_idx(x,y):
        return x + y*3

    def tuple(self):
        return (self.board, self.active_player, self.previous_moves)
    def __hash__(self):
        return hash(self.tuple())
    def __eq__(self, other):
        return self.tuple() == other.tuple()
    def __str__(self):
        ret = """
+---+---+---+
|   |   |   |
| Q | Q | Q |
|   |   |   |
+---+---+---+
|   |   |   |
| Q | Q | Q |
|   |   |   |
+---+---+---+
|   |   |   |
| Q | Q | Q |
|   |   |   |
+---+---+---+
+Q's turn   +
+-----------+
|last moves:|
|X: Q Q     |
|O: Q Q     |
+-----------+""".replace("Q", "{}")
        args = []
        for j in range(2,-1,-1):
            for i in range(3):
                cell = self.board[j][i]
                glyph = str(cell) if cell is not None else " "
                args.append(glyph)
        args.append(self.active_player)
        for p in "XO":
            coord = self.previous_moves[p]
            if coord is None:
                args.extend(("?", "?"))
            else:
                args.extend(map(str, coord))
        return ret.format(*args)

    def serialize(self):
        #pack the state into a string.
        ret= []
        #first nine characters: state of the board.
        for i in range(9):
            x,y = State.from_idx(i)
            c = self.board[y][x]
            if c is None: c = " "
            ret.append(c)
        #tenth character: active player.
        ret.append(self.active_player)
        #characters 11 through 15: last placed positions for players X and O.
        for player in "XO":
            if self.previous_moves[player] is None:
                ret.append(" ")
            else:
                ret.append(State.to_idx(*self.previous_moves[player]))
        return "".join(ret)

def is_round(x):
    if x < 10: return True
    if x % 10 != 0: return False
    return is_round(x/10)

def create_state_tree(start):
    d = {start: {"to": {}, "from": []}}
    to_visit = set([start])
    visited = set()

    milestones_seen = set()

    while to_visit:
        if is_round(len(visited)) and len(visited) not in milestones_seen:
            milestones_seen.add(len(visited))
            print len(visited)
            #if len(visited) == 100: return d

        cur = to_visit.pop()
        for move in cur.iter_possible_moves():
            neighbor = cur.apply(move)
            d[cur]["to"][move] = neighbor
            if neighbor not in d:
                d[neighbor] = {"to": {}, "from": []}
            d[neighbor]["from"].append(cur)
            if neighbor not in visited:
                to_visit.add(neighbor)
        visited.add(cur)
    return d

def determine_ultimate_winners(d):
    ultimate_winner = {state: None for state in d.keys()}
    time_to_victory = {state: None for state in d.keys()}
    while True:
        updates_made_this_round = 0
        for state in ultimate_winner:
            if ultimate_winner[state] is not None: #already found the result for this one
                continue
            if state.winner() is not None:
                ultimate_winner[state] = state.winner()
                time_to_victory[state] = 0
                updates_made_this_round += 1
            neighbors = list(d[state]["to"].values())
            if len(neighbors) == 0: #shouldn't happen when exploring full tree
                continue
            winning_neighbors = [neighbor for neighbor in neighbors if ultimate_winner[neighbor] == state.active_player]
            losing_neighbors = [neighbor for neighbor in neighbors if ultimate_winner[neighbor] == other(state.active_player)]
            if winning_neighbors:
                ultimate_winner[state] = state.active_player
                time_to_victory[state] = 1 + min(time_to_victory[neighbor] for neighbor in winning_neighbors)
                updates_made_this_round += 1
            elif len(losing_neighbors) == len(neighbors):
                ultimate_winner[state] = other(state.active_player)
                time_to_victory[state] = 1 + max(time_to_victory[neighbor] for neighbor in losing_neighbors)
                updates_made_this_round += 1
        if updates_made_this_round == 0:
            break
    return ultimate_winner, time_to_victory

def prompt_choice(items):
    for i, item in enumerate(items):
        print "{:3}: {}".format(i, item)
    while True:
        response = raw_input("Choose one: ")
        try:
            response = int(response)
        except:
            print "Sorry, did not understand that response."
            continue
        if response < 0 or response >= len(items):
            print "Sorry, that is not a valid choice."
            continue
        return response

starting_state = State(
    ((None, None, None),(None, None, None),(None, None, None)),
    "X")

shelf = shelve.open("results.dat")

force_generate = True
print "Attempting to load data from file..."
if force_generate:
    print "Made to fail by configuration setting."
    generate = True
elif "value" not in shelf:
    generate = True
else:
    generate = False

if generate:
    print "Failed. generating data from scratch..."
    d = create_state_tree(starting_state)
    print "Generated."
    #shelf["value"] = d
else:
    print "Success."
    d = shelf["value"]

shelf.close()

print "Determining winners from all states..."
winners, time_to_victory = determine_ultimate_winners(d)
print "Determined."
print "\a"
print "Winner of starting state:", winners[starting_state]
print "Moves to end of game, assuming optimal play:", time_to_victory[starting_state]

print "Proportion of state space with known winners: {} out of {}".format(sum(1 for state in winners if winners[state] is not None), len(winners))

print "Moves makable by active player:"
for move in starting_state.iter_possible_moves():
    neighbor = starting_state.apply(move)
    state = {"X": "Winner", "O": "Loser", None: "Tie?"}[winners[neighbor]]
    print move, state

# import pdb; pdb.set_trace()

#returns a (state, commentary) tuple
def best_choice(state):
    if state not in winners: #shouldn't happen
        return next(state.iter_possible_moves()), "I never explored this possibility."
        #return random.choice(list(state.iter_possible_moves()))
    if winners[state] is None: #stalemate or unexplored state
        return next(state.iter_possible_moves()), "Not sure who will win this."
    if winners[state] != state.active_player: #hopeless
        choice = max(state.iter_possible_moves(), key=lambda move: time_to_victory[state.apply(move)])
        return choice, "I can be defeated in {} moves.".format(time_to_victory[state.apply(choice)])
    candidates = []
    for move in state.iter_possible_moves():
        neighbor = state.apply(move)
        if winners[neighbor] == state.active_player:
            candidates.append(move)
    choice = min(candidates, key=lambda move: time_to_victory[state.apply(move)])
    return choice, "I win in no more than {} moves.".format(time_to_victory[state.apply(choice)])

def construct_minimal_tree(state):
    move, _ = best_choice(state)
    neighbor = state.apply(move)
    if neighbor.winner() is not None:
        return (move, None)
    responses = {}
    for next_move in neighbor.iter_possible_moves():
        responses[next_move] = construct_minimal_tree(neighbor.apply(next_move))
    return (move, responses)

def tree_size(t):
    if t[1] is None: return 1
    return 1 + sum(tree_size(branch) for branch in t[1].values())

tree = construct_minimal_tree(starting_state)
print "Size of optimal move tree:", tree_size(tree)



import random

try:
    while True:
        print "Let's play a game. you can be O."
        move_repr = lambda move: ("Move board right" if move[0] else "Leave board as-is") + " and take spot {} {}".format(move[1], move[2])
        state = starting_state
        while True:
            print state
            possible_moves = list(sorted(state.iter_possible_moves()))
            if len(possible_moves) == 0:
                print "Looks like a tie."
                break
            if state.active_player == "O":
                reply = prompt_choice(map(move_repr, possible_moves))
                state = state.apply(possible_moves[reply])
            else:
                move, comment = best_choice(state)
                print "I'll {}. {}".format(move_repr(move).lower(), comment)
                state = state.apply(move)
            if state.winner() is not None:
                print state
                print "{} is the winner!".format(state.winner())
                break
        reply = raw_input("Play again? ")
        if reply.lower() == "n": break
except KeyboardInterrupt:
    import pdb; pdb.set_trace()