import numpy as np
from attr import dataclass

from environments.factory.base_factory import BaseFactory
from collections import namedtuple
from typing import Iterable
from environments import helpers as h

DIRT_INDEX = -1
@dataclass
class DirtProperties:
    clean_amount = 0.25
    max_spawn_ratio = 0.1
    gain_amount = 0.1


class GettingDirty(BaseFactory):

    @property
    def _clean_up_action(self):
        return self.movement_actions + 1 - 1

    def __init__(self, *args, dirt_properties: DirtProperties, **kwargs):
        self._dirt_properties = dirt_properties
        super(GettingDirty, self).__init__(*args, **kwargs)
        self.slice_strings.update({self.state.shape[0]-1: 'dirt'})

    def spawn_dirt(self) -> None:
        free_for_dirt = self.free_cells
        # randomly distribute dirt across the grid
        n_dirt_tiles = int(random.uniform(0, self._dirt_properties.max_spawn_ratio) * len(free_for_dirt))
        for x, y in free_for_dirt[:n_dirt_tiles]:
            self.state[DIRT_INDEX, x, y] += self._dirt_properties.gain_amount

    def clean_up(self, pos: (int, int)) -> ((int, int), bool):
        new_dirt_amount = self.state[DIRT_INDEX][pos] - self._dirt_properties.clean_amount
        cleanup_was_sucessfull: bool
        if self.state[DIRT_INDEX][pos] == h.IS_FREE_CELL:
            cleanup_was_sucessfull = False
            return pos, cleanup_was_sucessfull
        else:
            cleanup_was_sucessfull = True
            self.state[DIRT_INDEX][pos] = max(new_dirt_amount, h.IS_FREE_CELL)
            return pos, cleanup_was_sucessfull

    def additional_actions(self, agent_i, action) -> ((int, int), bool):
        if action != self._is_moving_action(action):
            if action == self._clean_up_action:
                agent_i_pos = self.agent_i_position(agent_i)
                _, valid = self.clean_up(agent_i_pos)
                if valid:
                    print(f'Agent {agent_i} did just clean up some dirt at {agent_i_pos}.')
                else:
                    print(f'Agent {agent_i} just tried to clean up some dirt at {agent_i_pos}, but was unsucsessfull.')
                return agent_i_pos, valid
        else:
            raise RuntimeError('This should not happen!!!')

    def reset(self) -> None:
        # ToDo: When self.reset returns the new states and stuff, use it here!
        super().reset()  # state, agents, ... =
        dirt_slice = np.zeros((1, *self.state.shape[1:]))
        self.state = np.concatenate((self.state, dirt_slice))  # dirt is now the last slice
        self.spawn_dirt()

    def calculate_reward(self, collisions_vecs: np.ndarray, actions: Iterable[int]) -> (int, dict):
        for agent_i, cols in enumerate(collisions_vecs):
            cols = np.argwhere(cols != 0).flatten()
            print(f't = {self.steps}\tAgent {agent_i} has collisions with '
                  f'{[self.slice_strings[entity] for entity in cols if entity != self.state.shape[0]]}')
        return 0, {}


if __name__ == '__main__':
    import random
    dirt_props = DirtProperties()
    factory = GettingDirty(n_agents=1, dirt_properties=dirt_props)
    random_actions = [random.randint(0, 8) for _ in range(200)]
    for action in random_actions:
        state, r, done, _ = factory.step(action)