import random
from copy import deepcopy

PT = 0.8
PO = 0.8
C = 4

RED = 1
GREEN = 2
BLUE = 3
BLACK = 0

COLORS = [BLACK, RED, GREEN, BLUE]

def weighted_choice(choices):
   total = sum(w for c, w in choices)
   r = random.uniform(0, total)
   upto = 0
   for c, w in choices:
      if upto + w >= r:
         return c
      upto += w

class Grid(object):
    def __init__(self, name, heights, colors):
        self.name = name
        self.heights = heights
        self.colors = colors
        self.size = len(heights)

    def height(self, state):
        return self.heights[state[0]][state[1]]

    def color(self, state):
        return self.colors[state[0]][state[1]]

    def get_neighbors(self, state):
        neighbors = [self.left(state), self.right(state), self.up(state), self.down(state)]
        neighbors = [(n, self.height(n)) for n in neighbors if n is not None]
        max_height = max([n[1] for n in neighbors])
        n_max = [n[1] for n in neighbors].count(max_height)

        prob = (1 - PT) / len(neighbors)
        return {n[0]:(prob + PT / n_max if (n[1] == max_height) else prob) for n in neighbors}

    def get_colors(self, state):
        color = self.color(state)

        prob = (1 - PO) / C
        return {c:(PO + prob if c == color else prob) for c in COLORS}

    def transition_prob(self, src, dst):
        return self.get_neighbors(src).get(dst, 0)

    def emission_prob(self, state, color):
        return self.get_colors(state).get(color, 0)

    def get_sequence(self, length):
        states = []
        observations = []
        state = None
        for t in range(length):
            if t == 0:
                state = (random.randint(0, self.size - 1), random.randint(0, self.size - 1))
            else:
                state = weighted_choice([(k,v) for k,v in self.get_neighbors(state).iteritems()])
            observation = weighted_choice([(k,v) for k,v in self.get_colors(state).iteritems()])

            states.append(state)
            observations.append(observation)

        return states, observations

    def forward(self, observations):
        alpha = []
        for _ in range(self.size):
            alpha.append([1 / (1.0 * self.size * self.size)] * self.size)
        for row in range(self.size):
            for col in range(self.size):
                alpha[row][col] *= self.emission_prob((row, col), observations[0])

        for idx in range(len(observations)):
            prev_alpha = deepcopy(alpha)
            alpha = []
            for _ in range(self.size):
                alpha.append([0] * self.size)

            for row in range(self.size):
                for col in range(self.size):
                    for state, prob in self.get_neighbors((row, col)).iteritems():
                        alpha[row][col] += (prev_alpha[state[0]][state[1]] * prob * self.emission_prob((row, col), observations[idx]))

        return sum([sum(s) for s in prev_alpha]), prev_alpha


    def left(self, state):
        if state[1] == 0:
            return None
        return state[0], state[1] - 1

    def right(self, state):
        if state[1] == self.size - 1:
            return None
        return state[0], state[1] + 1

    def up(self, state):
        if state[0] == 0:
            return None
        return state[0] - 1, state[1]

    def down(self, state):
        if state[0] == self.size - 1:
            return None
        return state[0] + 1, state[1]


def main():
    grid1 = Grid("Grid 1",
                 [[1, 2, 3, 5], [2, 2, 1, 2], [3, 2, 1, 1], [0, 0, 0, 0]],  # elevation
                 [[0, 3, 1, 2], [3, 1, 2, 0], [2, 2, 0, 0], [3, 0, 3, 1]])  # color
    grid2 = Grid("Grid 2",
                 [[0, 0, 1, 1], [2, 1, 0, 2], [1, 0, 0, 2], [4, 4, 3, 3]],  # elevation
                 [[0, 3, 1, 2], [3, 1, 2, 0], [2, 2, 0, 0], [3, 0, 3, 1]])  # color
    grid3 = Grid("Grid 3",
                 [[2, 1, 2, 3], [1, 1, 2, 2], [1, 0, 1, 1], [2, 1, 1, 2]],  # elevation
                 [[2, 3, 1, 0], [1, 3, 3, 1], [0, 2, 0, 2], [2, 1, 1, 2]])  # color
    GRIDS = [grid1, grid2, grid3]

    # print grid.get_neighbors((3,1))
    # print grid.get_colors((0, 0))
    # print grid.transition_prob((3, 3), (3, 2))
    # print grid.emission_prob((0, 0), GREEN)
    # print grid.get_sequence(2)
    for grid in GRIDS:
        print grid.forward([0, 0])
        print ''

if __name__ == '__main__':
    main()