import random

import robots

from robots.algorithms import distance_relaxer
from robots.constants import City
from robots.utils import add_spawns

class CaptureSpawns(object):
    def __init__(self, variance=0.1):
        self.iterations = 10
        self.variance = variance

    def __call__(self, whoami, state):
        my_robots = state.robots_by_player[whoami]

        # Create a distance matrix.
        distances = []
        for y, row in enumerate(state.cities):
            output = []
            for x, city in enumerate(row):
                d = float('inf')
                if city == City.FACTORY:
                    if state.allegiances.get((x, y)) != whoami:
                        d = 0
                elif city == City.GHOST:
                    d = None
                output.append(d)
            distances.append(output)

        # Find the shortest path to a target from each cell.
        predecessors = distance_relaxer(distances)

        # Direct the robots to follow those paths.
        results = []
        for x, y, energy in my_robots:
            if random.random() < self.variance:
                result = random.choice('ULDR')
            elif predecessors[y][x]:
                result = random.choice(predecessors[y][x])
            elif state.allegiances.get((x, y)) != whoami:
                result = 'P'
            else:
                result = random.choice('UDLR')
            results.append(result)
        return results

if __name__ == '__main__':
#    random.seed(42)
    map_ = robots.border_map(40, 30, 0)

    add_spawns(map_, 300, 'X')
    add_spawns(map_, 20, '+')

    for y in range(20):
        map_[y][20] = 'X'

    add_spawns(map_, 6)

    game = robots.Game(map_)
    game.add_bot(CaptureSpawns(variance=0), 'Alice')
    game.add_bot(CaptureSpawns(variance=0.01), 'Bob')
    viewer = robots.CursesViewer(game)
    viewer.run()