from __future__ import annotations
from collections import deque
from itertools import product
import time
from typing import List
import numpy as np
from luxai2022.config import EnvConfig
from luxai2022.map.board import Board
from luxai2022.map.position import Position
from luxai2022.team import Team
from luxai2022.unit import UnitCargo
from luxai2022.actions import move_deltas
from luxai2022.globals import TERM_COLORS
try:
    from termcolor import colored
except:
    pass


def compute_water_info(init: np.ndarray, MIN_LICHEN_TO_SPREAD: int, lichen: np.ndarray, lichen_strains: np.ndarray, factory_occupancy_map: np.ndarray, strain_id: int, forbidden: np.ndarray):
    # TODO - improve the performance here with cached solution
    frontier = deque(init)
    seen = set(map(tuple, init))
    grow_lichen_positions = set()
    H, W = lichen.shape
    ct = 0
    while len(frontier) > 0:
        ct += 1
        if ct > 1_000_000:
            print("Error! Lichen Growth calculation took too long")
            break
        pos = frontier.popleft()
        if pos[0] < 0 or pos[1] < 0 or pos[0] >= forbidden.shape[0] or pos[1] >= forbidden.shape[1]:
            continue

        if forbidden[pos[0], pos[1]]:
            continue
        pos_lichen = lichen[pos[0], pos[1]]
        pos_strain = lichen_strains[pos[0], pos[1]]
        # check for surrounding tiles with lichen and no incompatible lichen strains, grow on those
        can_grow = True
        for move_delta in move_deltas[1:]:
            check_pos = pos + move_delta
            # check surrounding tiles on the map
            if check_pos[0] < 0 or check_pos[1] < 0 or check_pos[0] >= H or check_pos[1] >= W: continue

            # If any neighbor 1. has a different strain, or 2. is a different factory, 
            # then the current pos cannot grow
            adj_strain = lichen_strains[check_pos[0], check_pos[1]]
            adj_factory = factory_occupancy_map[check_pos[0], check_pos[1]]
            if (adj_strain != -1 and adj_strain != strain_id) \
                or (adj_factory != -1 and adj_factory != strain_id):
                can_grow = False

            # if seen, skip
            if (check_pos[0], check_pos[1]) in seen:
                continue

            # we add it to the frontier only in two cases:
            #  1. it is an empty tile, and current pos has enough lichen to expand.
            #  2. both current tile and check_pos are of our strain.
            if (adj_strain == -1 and pos_lichen >= MIN_LICHEN_TO_SPREAD) \
                or (adj_strain == strain_id and pos_strain == strain_id):
                seen.add(tuple(check_pos))
                frontier.append(check_pos)

        if can_grow or (lichen_strains[pos[0], pos[1]] == strain_id):
            grow_lichen_positions.add((pos[0], pos[1]))
    return grow_lichen_positions

class Factory:
    def __init__(self, team: Team, unit_id: str, num_id: int) -> None:
        self.team_id = team.team_id
        self.team = team
        self.unit_id = unit_id
        self.pos = Position(np.zeros(2, dtype=int))
        self.power = 0
        self.cargo = UnitCargo()
        self.num_id = num_id
        self.action_queue = []
        self.grow_lichen_positions = None

    @property
    def pos_slice(self):
        return slice(self.pos.x - 1, self.pos.x + 2), slice(self.pos.y - 1, self.pos.y + 2)

    @property
    def min_dist_slice(self):
        deltas = np.array([(4, 0), (5, 1), (0, -4), (0, 5), (2, 2), (-4, 1), (-1, -1), (-1, -2), (-2, -1), (-2, -2), (4, 2), (0, -2), (0, -1), (2, 4), (-1, 0), (-2, 0), (-3, 1), (0, 0), (2, -3), (-2, 2), (-1, 2), (3, 1), (5, -1), (-3, 3), (0, 2), (1, 3), (-4, -1), (-1, -5), (-4, -2), (-1, 4), (-2, 4), (3, 3), (5, 0), (-6, 0), (0, -5), (1, -4), (1, 5), (-4, 0), (-1, -3), (-2, -3), (-5, 1), (-3, -1), (0, -3), (-3, -2), (1, -1), (1, -2), (-4, 2), (3, -1), (3, -2), (-3, 0), (1, 0), (3, 0), (-3, 2), (1, 2), (0, 4), (2, 1), (3, 2), (4, 1), (-5, -1), (1, -5), (1, 4), (0, 6), (2, 3), (-1, -4), (-5, 0), (-3, -3), (1, -3), (2, -4), (-1, 1), (-2, 1), (3, -3), (0, 1), (2, -1), (2, -2), (-1, 3), (-2, 3), (4, -1), (4, -2), (0, -6), (1, 1), (0, 3), (2, 0), (6, 0), (-2, -4), (-1, 5)])
        return self.pos.pos + deltas


    def refine_step(self, config: EnvConfig):
        max_consumed_ice = min(self.cargo.ice, config.FACTORY_PROCESSING_RATE_WATER)
        max_consumed_ore = min(self.cargo.ore, config.FACTORY_PROCESSING_RATE_METAL)
        # permit refinement of blocks of resources, no floats.
        produced_water = max_consumed_ice // config.ICE_WATER_RATIO
        produced_metal = max_consumed_ore // config.ORE_METAL_RATIO
        self.cargo.ice -= produced_water * config.ICE_WATER_RATIO
        self.cargo.ore -= produced_metal * config.ORE_METAL_RATIO

        self.cargo.water += produced_water
        self.cargo.metal += produced_metal

    def cache_water_info(self, board: Board, env_cfg: EnvConfig):
        # TODO this can easily be a fairly slow function, can we make it much faster?
        # Caches information about which tiles lichen can grow on for this factory

        # perform a BFS from the factory position and look for non rubble, non factory tiles.
        # find the current frontier from 12 starting positions x marked below
        """
           x x x
           _ _ _
        x |     | x
        x |     | x
        x |_ _ _| x
           x x x

        """
        forbidden = (board.rubble > 0) | (board.factory_occupancy_map != -1) | (board.ice > 0) | (board.ore > 0)
        deltas = [np.array([0, -2]),  np.array([-1, -2]),  np.array([1, -2]),  np.array([0, 2]),  np.array([-1, 2]),  np.array([1, 2]),
                    np.array([2, 0]),  np.array([2, -1]),  np.array([2, 1]),  np.array([-2, 0]),  np.array([-2, -1]),  np.array([-2, 1])]
        init_arr = np.stack(deltas) + self.pos.pos
        self.grow_lichen_positions = compute_water_info(init_arr, env_cfg.MIN_LICHEN_TO_SPREAD, board.lichen, board.lichen_strains, board.factory_occupancy_map, self.num_id, forbidden)
    def water_cost(self, config: EnvConfig):
        return int(np.ceil(len(self.grow_lichen_positions) / config.LICHEN_WATERING_COST_FACTOR))

    ### Add and sub resource functions copied over from unit.py code, can we consolidate them somewhere?
    def add_resource(self, resource_id, transfer_amount):
        if transfer_amount < 0: transfer_amount = 0
        if resource_id == 0:
            self.cargo.ice += transfer_amount
        elif resource_id == 1:
            self.cargo.ore += transfer_amount
        elif resource_id == 2:
            self.cargo.water += transfer_amount
        elif resource_id == 3:
            self.cargo.metal += transfer_amount
        elif resource_id == 4:
            self.power += transfer_amount
        return int(transfer_amount)
    def sub_resource(self, resource_id, amount):
        # subtract/transfer out as much as you min(have, request)
        if amount < 0: amount = 0
        if resource_id == 0:
            transfer_amount = min(self.cargo.ice, amount)
            self.cargo.ice -= transfer_amount
        elif resource_id == 1:
            transfer_amount = min(self.cargo.ore, amount)
            self.cargo.ore -= transfer_amount
        elif resource_id == 2:
            transfer_amount = min(self.cargo.water, amount)
            self.cargo.water -= transfer_amount
        elif resource_id == 3:
            transfer_amount = min(self.cargo.metal, amount)
            self.cargo.metal -= transfer_amount
        elif resource_id == 4:
            transfer_amount = min(self.power, amount)
            self.power -= transfer_amount
        return int(transfer_amount)




    def state_dict(self):
        return dict(
            pos=self.pos.pos,
            power=self.power,
            cargo=self.cargo.state_dict(),
            unit_id=self.unit_id,
            strain_id=self.num_id, # number version of unit_id
            team_id=self.team_id

        )

    def __str__(self) -> str:
        out = f"[{self.team_id}] {self.unit_id} Factory at {self.pos}"
        if TERM_COLORS:
            return colored(out, self.team.faction.value.color)
        return out
