# Strategiser
# Goals

from collections import defaultdict
from functools import lru_cache
from random import choice

import robots

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

def get_close_spawns(whoami, state, threshold):
    distances = []
    for y, row in enumerate(state.cities):
        distances.append([])
        for x, cell in enumerate(row):
            if cell not in City.traversable:
                d = None
            elif (
                cell == City.FACTORY and
                state.allegiances.get((x, y)) != whoami
            ):
                d = 0
            else:
                d = float('inf')
            distances[-1].append(d)

    predecessors = distance_relaxer(distances, threshold)
    return predecessors, distances

@lru_cache(maxsize=16)
def closest_unpainted(whoami, state):
    distances = []
    for y, row in enumerate(state.cities):
        distances.append([])
        for x, cell in enumerate(row):
            if cell == City.GHOST:
                d = None
            elif state.allegiances.get((x, y)) != whoami:
                d = 0
            else:
                d = float('inf')
            distances[-1].append(d)

    return distance_relaxer(distances)

def good_moves(state, x, y, moves):
    return [
        move
        for move in moves
        if not state.robots[state.expected_position(x, y, move)]
    ]

class Napoleon:
    def __init__(self, capturers_frac=1, dampening=0.01):
        # (x, y) -> (goal, goal_state)
        self.bot_goals = {}
        self.capturers_frac = capturers_frac
        self.dampening = dampening

    def goal_capture(self, whoami, state, x, y, goal_state):
        moves = goal_state
        good = good_moves(state, x, y, moves)
        if good:
            return choice(good)
        return choice(moves or 'P')

    def goal_paint(self, whoami, state, x, y, goal_state):
        predecessors = closest_unpainted(whoami, state)
        moves = predecessors[y][x]
        good = good_moves(state, x, y, moves)
        if good:
            return choice(good)
        return choice(moves or 'P')

    def strategise(self, whoami, state):
        bot_goals = {}

        robots = state.robots_by_player[whoami]

        if len(state.factories_by_player[whoami]) <= 3:
            self.capturers_frac = 1.0
        if len(robots) <= 10:
            self.capturers_frac = 1.0

#        threshold = 10 if early_stages else None
        threshold = None
        spawns, spawns_dist = get_close_spawns(whoami, state, threshold)

        max_capturers = max(3, int(len(robots) * self.capturers_frac))
        self.capturers_frac = max(0.01, self.capturers_frac - self.dampening)
        capturers = {
            (x, y)
            for x, y, _ in sorted(
                robots,
                key=lambda r: spawns_dist[r[1]][r[0]],
            )[:max_capturers]
            if threshold is None or spawns_dist[y][x] <= threshold
        }
        assert len(capturers) <= max_capturers

        for x, y, energy in robots:
            if (x, y) in capturers:
                goal = 'capture'
                goal_state = spawns[y][x]
            else:
                goal = 'paint'
                goal_state = None

            bot_goals[x, y] = (goal, goal_state)

#        print(whoami, ''.join(v[0].upper() for v, _ in bot_goals.values()))
        return bot_goals

    def __call__(self, whoami, state):
        self.bot_goals = self.strategise(whoami, state)

        results = {}

        todo = self.bot_goals.keys()

        for _ in range(2):
            by_dest = defaultdict(list)

            for x, y in todo:
                goal, goal_state = self.bot_goals[x, y]
                goal_fn = getattr(self, 'goal_' + goal)
                action = goal_fn(whoami, state, x, y, goal_state)
                by_dest[state.expected_position(x, y, action)].append((x, y))
                results[x, y] = action

            todo = []
            for robots in by_dest.values():
                if len(robots) > 1:
                    todo.extend(robots)

        return ''.join(
            results[x, y]
            for x, y, _ in state.robots_by_player[whoami]
        )

if __name__ == '__main__':
#    import random

#    random.seed(42)
    map_ = robots.border_map(42, 22, 0)

    add_spawns(map_, 10, 'X')
    add_spawns(map_, 20, '+')
    add_spawns(map_, 6)

    game = robots.Game(map_)

    game.add_bot(Napoleon(), 'Napoleon')

    from capturer import CaptureSpawns
    game.add_bot(CaptureSpawns(), 'Alice')
    game.add_bot(CaptureSpawns(), 'Bob')
    game.add_bot(CaptureSpawns(), 'Charlie')
    game.add_bot(CaptureSpawns(), 'Doug')

    viewer = robots.CursesViewer(game)
    viewer.run()