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.distances = None
        self.iterations = 10
        self.variance = variance

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

        self.distances = [
            [
                {
                    City.NORMAL: (
                        float('inf')
#                        if self.distances is None or self.distances[y][x] is None
#                        else self.distances[y][x] + self.iterations // 2
                    ),
                    City.FACTORY: (
                        0
                        if state.allegiances.get((x, y)) != whoami else
                        float('inf')
                    ),
                    City.GHOST: None,
                }[city]
                for x, city in enumerate(row)
            ]
            for y, row in enumerate(state.cities)
        ]

        predecessors = distance_relaxer(self.distances, None)

        results = []
        for x, y, energy in my_robots:
            if random.random() < self.variance:
                result = random.choice('ULDR')
            elif predecessors[y][x]:
                result = 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()