# Copyright (c) 2020 All Rights Reserved
# Author: William H. Guss, Brandon Houghton

import collections
from jarvis.stark_tech.herobraine.hero.handler import Handler
import numpy as np
from functools import reduce

from typing import List, Tuple

from jarvis.stark_tech.herobraine.hero.spaces import Box, Dict, Enum, MineRLSpace


# TODO: Make a test.
# TODO: Refactor this. This iss unioning handlers, not sapces.
def union_spaces(hdls_1: List[Handler], hdls_2: List[Handler]) -> List[MineRLSpace]:
    # Merge action/observation spaces from two environments
    hdls = hdls_1 + hdls_2
    hdl_dict = collections.defaultdict(list)
    _ = [hdl_dict[hdl.to_string()].append(hdl) for hdl in hdls]  # Join matching handlers
    merged_hdls = [reduce(lambda a, b: a | b, matching) for matching in hdl_dict.values()]

    return merged_hdls


# TODO: make a test.
# TODO: Maybe this should be based on handlers as above.
# E.g. 1. Intersect the handlers
# E.g. 2. We then do best effort clipping for spaces space.intersect
# wg note: whenever you write isinstance for a bunch of classes 
# you are being a fucing idiot this is exactly why we have object 
# oriented programming.
def intersect_space(space, sample):
    if isinstance(space, Dict):
        new_sample = collections.OrderedDict()
        for key, value in space.spaces.items():
            new_sample[key] = intersect_space(value, sample[key])
        return new_sample
    elif isinstance(space, Enum):
        if sample not in space:
            return space.default
        else:
            return sample
    else:
        # TODO: SUPPORT SPACE INTERSECTION.
        return sample


# TODO: make a test
def flatten_spaces(hdls: List[Handler]) -> Tuple[list, List[Tuple[str, MineRLSpace]]]:
    return [hdl.space.flattened for hdl in hdls if hdl.space.is_flattenable()], \
           [(hdl.to_string(), hdl.space) for hdl in hdls if
            not hdl.space.is_flattenable()]
