from mlagents_envs.base_env import (
    ActionSpec,
    ObservationSpec,
    DimensionProperty,
    BehaviorSpec,
    DecisionSteps,
    TerminalSteps,
    ObservationType,
)
from mlagents_envs.exception import UnityObservationException
from mlagents_envs.timers import hierarchical_timer, timed
from mlagents_envs.communicator_objects.agent_info_pb2 import AgentInfoProto
from mlagents_envs.communicator_objects.observation_pb2 import (
    ObservationProto,
    NONE as COMPRESSION_TYPE_NONE,
)
from mlagents_envs.communicator_objects.brain_parameters_pb2 import BrainParametersProto
import numpy as np
import io
from typing import cast, List, Tuple, Collection, Optional, Iterable
from PIL import Image


PNG_HEADER = b"\x89PNG\r\n\x1a\n"


def behavior_spec_from_proto(
    brain_param_proto: BrainParametersProto, agent_info: AgentInfoProto
) -> BehaviorSpec:
    """
    Converts brain parameter and agent info proto to BehaviorSpec object.
    :param brain_param_proto: protobuf object.
    :param agent_info: protobuf object.
    :return: BehaviorSpec object.
    """
    observation_specs = []
    for obs in agent_info.observations:
        observation_specs.append(
            ObservationSpec(
                tuple(obs.shape),
                tuple(DimensionProperty(dim) for dim in obs.dimension_properties)
                if len(obs.dimension_properties) > 0
                else (DimensionProperty.UNSPECIFIED,) * len(obs.shape),
                ObservationType(obs.observation_type),
            )
        )

    # proto from communicator < v1.3 does not set action spec, use deprecated fields instead
    if (
        brain_param_proto.action_spec.num_continuous_actions == 0
        and brain_param_proto.action_spec.num_discrete_actions == 0
    ):
        if brain_param_proto.vector_action_space_type_deprecated == 1:
            action_spec = ActionSpec(
                brain_param_proto.vector_action_size_deprecated[0], ()
            )
        else:
            action_spec = ActionSpec(
                0, tuple(brain_param_proto.vector_action_size_deprecated)
            )
    else:
        action_spec_proto = brain_param_proto.action_spec
        action_spec = ActionSpec(
            action_spec_proto.num_continuous_actions,
            tuple(branch for branch in action_spec_proto.discrete_branch_sizes),
        )
    return BehaviorSpec(observation_specs, action_spec)


class OffsetBytesIO:
    """
    Simple file-like class that wraps a bytes, and allows moving its "start"
    position in the bytes. This is only used for reading concatenated PNGs,
    because Pillow always calls seek(0) at the start of reading.
    """

    __slots__ = ["fp", "offset"]

    def __init__(self, data: bytes):
        self.fp = io.BytesIO(data)
        self.offset = 0

    def seek(self, offset: int, whence: int = io.SEEK_SET) -> int:
        if whence == io.SEEK_SET:
            res = self.fp.seek(offset + self.offset)
            return res - self.offset
        raise NotImplementedError()

    def tell(self) -> int:
        return self.fp.tell() - self.offset

    def read(self, size: int = -1) -> bytes:
        return self.fp.read(size)

    def original_tell(self) -> int:
        """
        Returns the offset into the original byte array
        """
        return self.fp.tell()


@timed
def process_pixels(
    image_bytes: bytes, expected_channels: int, mappings: Optional[List[int]] = None
) -> np.ndarray:
    """
    Converts byte array observation image into numpy array, re-sizes it,
    and optionally converts it to grey scale
    :param image_bytes: input byte array corresponding to image
    :param expected_channels: Expected output channels
    :return: processed numpy array of observation from environment
    """
    image_fp = OffsetBytesIO(image_bytes)

    image_arrays = []
    # Read the images back from the bytes (without knowing the sizes).
    while True:
        with hierarchical_timer("image_decompress"):
            image = Image.open(image_fp)
            # Normally Image loads lazily, load() forces it to do loading in the timer scope.
            image.load()
        image_arrays.append(np.array(image, dtype=np.float32) / 255.0)

        # Look for the next header, starting from the current stream location
        try:
            new_offset = image_bytes.index(PNG_HEADER, image_fp.original_tell())
            image_fp.offset = new_offset
        except ValueError:
            # Didn't find the header, so must be at the end.
            break

    if mappings is not None and len(mappings) > 0:
        return _process_images_mapping(image_arrays, mappings)
    else:
        return _process_images_num_channels(image_arrays, expected_channels)


def _process_images_mapping(image_arrays, mappings):
    """
    Helper function for processing decompressed images with compressed channel mappings.
    """
    image_arrays = np.concatenate(image_arrays, axis=2).transpose((2, 0, 1))

    if len(mappings) != len(image_arrays):
        raise UnityObservationException(
            f"Compressed observation and its mapping had different number of channels - "
            f"observation had {len(image_arrays)} channels but its mapping had {len(mappings)} channels"
        )
    if len({m for m in mappings if m > -1}) != max(mappings) + 1:
        raise UnityObservationException(
            f"Invalid Compressed Channel Mapping: the mapping {mappings} does not have the correct format."
        )
    if max(mappings) >= len(image_arrays):
        raise UnityObservationException(
            f"Invalid Compressed Channel Mapping: the mapping has index larger than the total "
            f"number of channels in observation - mapping index {max(mappings)} is"
            f"invalid for input observation with {len(image_arrays)} channels."
        )

    processed_image_arrays: List[np.array] = [[] for _ in range(max(mappings) + 1)]
    for mapping_idx, img in zip(mappings, image_arrays):
        if mapping_idx > -1:
            processed_image_arrays[mapping_idx].append(img)

    for i, img_array in enumerate(processed_image_arrays):
        processed_image_arrays[i] = np.mean(img_array, axis=0)
    img = np.stack(processed_image_arrays, axis=2)
    return img


def _process_images_num_channels(image_arrays, expected_channels):
    """
    Helper function for processing decompressed images with number of expected channels.
    This is for old API without mapping provided. Use the first n channel, n=expected_channels.
    """
    if expected_channels == 1:
        # Convert to grayscale
        img = np.mean(image_arrays[0], axis=2)
        img = np.reshape(img, [img.shape[0], img.shape[1], 1])
    else:
        img = np.concatenate(image_arrays, axis=2)
        # We can drop additional channels since they may need to be added to include
        # numbers of observation channels not divisible by 3.
        actual_channels = list(img.shape)[2]
        if actual_channels > expected_channels:
            img = img[..., 0:expected_channels]
    return img


@timed
def _observation_to_np_array(
    obs: ObservationProto, expected_shape: Optional[Iterable[int]] = None
) -> np.ndarray:
    """
    Converts observation proto into numpy array of the appropriate size.
    :param obs: observation proto to be converted
    :param expected_shape: optional shape information, used for sanity checks.
    :return: processed numpy array of observation from environment
    """
    if expected_shape is not None:
        if list(obs.shape) != list(expected_shape):
            raise UnityObservationException(
                f"Observation did not have the expected shape - got {obs.shape} but expected {expected_shape}"
            )
    expected_channels = obs.shape[2]
    if obs.compression_type == COMPRESSION_TYPE_NONE:
        img = np.array(obs.float_data.data, dtype=np.float32)
        img = np.reshape(img, obs.shape)
        return img
    else:
        img = process_pixels(
            obs.compressed_data, expected_channels, list(obs.compressed_channel_mapping)
        )
        # Compare decompressed image size to observation shape and make sure they match
        if list(obs.shape) != list(img.shape):
            raise UnityObservationException(
                f"Decompressed observation did not have the expected shape - "
                f"decompressed had {img.shape} but expected {obs.shape}"
            )
        return img


@timed
def _process_maybe_compressed_observation(
    obs_index: int,
    shape: Tuple[int, int, int],
    agent_info_list: Collection[AgentInfoProto],
) -> np.ndarray:
    if len(agent_info_list) == 0:
        return np.zeros((0, shape[0], shape[1], shape[2]), dtype=np.float32)

    batched_visual = [
        _observation_to_np_array(agent_obs.observations[obs_index], shape)
        for agent_obs in agent_info_list
    ]
    return np.array(batched_visual, dtype=np.float32)


def _raise_on_nan_and_inf(data: np.array, source: str) -> np.array:
    # Check for NaNs or Infinite values in the observation or reward data.
    # If there's a NaN in the observations, the np.mean() result will be NaN
    # If there's an Infinite value (either sign) then the result will be Inf
    # See https://stackoverflow.com/questions/6736590/fast-check-for-nan-in-numpy for background
    # Note that a very large values (larger than sqrt(float_max)) will result in an Inf value here
    # Raise a Runtime error in the case that NaNs or Infinite values make it into the data.
    if data.size == 0:
        return data

    d = np.mean(data)
    has_nan = np.isnan(d)
    has_inf = not np.isfinite(d)

    if has_nan:
        raise RuntimeError(f"The {source} provided had NaN values.")
    if has_inf:
        raise RuntimeError(f"The {source} provided had Infinite values.")


@timed
def _process_rank_one_or_two_observation(
    obs_index: int, shape: Tuple[int, ...], agent_info_list: Collection[AgentInfoProto]
) -> np.ndarray:
    if len(agent_info_list) == 0:
        return np.zeros((0,) + shape, dtype=np.float32)
    np_obs = np.array(
        [
            agent_obs.observations[obs_index].float_data.data
            for agent_obs in agent_info_list
        ],
        dtype=np.float32,
    ).reshape((len(agent_info_list),) + shape)
    _raise_on_nan_and_inf(np_obs, "observations")
    return np_obs


@timed
def steps_from_proto(
    agent_info_list: Collection[AgentInfoProto], behavior_spec: BehaviorSpec
) -> Tuple[DecisionSteps, TerminalSteps]:
    decision_agent_info_list = [
        agent_info for agent_info in agent_info_list if not agent_info.done
    ]
    terminal_agent_info_list = [
        agent_info for agent_info in agent_info_list if agent_info.done
    ]
    decision_obs_list: List[np.ndarray] = []
    terminal_obs_list: List[np.ndarray] = []
    for obs_index, observation_specs in enumerate(behavior_spec.observation_specs):
        is_visual = len(observation_specs.shape) == 3
        if is_visual:
            obs_shape = cast(Tuple[int, int, int], observation_specs.shape)
            decision_obs_list.append(
                _process_maybe_compressed_observation(
                    obs_index, obs_shape, decision_agent_info_list
                )
            )
            terminal_obs_list.append(
                _process_maybe_compressed_observation(
                    obs_index, obs_shape, terminal_agent_info_list
                )
            )
        else:
            decision_obs_list.append(
                _process_rank_one_or_two_observation(
                    obs_index, observation_specs.shape, decision_agent_info_list
                )
            )
            terminal_obs_list.append(
                _process_rank_one_or_two_observation(
                    obs_index, observation_specs.shape, terminal_agent_info_list
                )
            )
    decision_rewards = np.array(
        [agent_info.reward for agent_info in decision_agent_info_list], dtype=np.float32
    )
    terminal_rewards = np.array(
        [agent_info.reward for agent_info in terminal_agent_info_list], dtype=np.float32
    )

    decision_group_rewards = np.array(
        [agent_info.group_reward for agent_info in decision_agent_info_list],
        dtype=np.float32,
    )
    terminal_group_rewards = np.array(
        [agent_info.group_reward for agent_info in terminal_agent_info_list],
        dtype=np.float32,
    )

    _raise_on_nan_and_inf(decision_rewards, "rewards")
    _raise_on_nan_and_inf(terminal_rewards, "rewards")
    _raise_on_nan_and_inf(decision_group_rewards, "group_rewards")
    _raise_on_nan_and_inf(terminal_group_rewards, "group_rewards")

    decision_group_id = [agent_info.group_id for agent_info in decision_agent_info_list]
    terminal_group_id = [agent_info.group_id for agent_info in terminal_agent_info_list]

    max_step = np.array(
        [agent_info.max_step_reached for agent_info in terminal_agent_info_list],
        dtype=np.bool,
    )
    decision_agent_id = np.array(
        [agent_info.id for agent_info in decision_agent_info_list], dtype=np.int32
    )
    terminal_agent_id = np.array(
        [agent_info.id for agent_info in terminal_agent_info_list], dtype=np.int32
    )
    action_mask = None
    if behavior_spec.action_spec.discrete_size > 0:
        if any(
            [agent_info.action_mask is not None]
            for agent_info in decision_agent_info_list
        ):
            n_agents = len(decision_agent_info_list)
            a_size = np.sum(behavior_spec.action_spec.discrete_branches)
            mask_matrix = np.ones((n_agents, a_size), dtype=np.bool)
            for agent_index, agent_info in enumerate(decision_agent_info_list):
                if agent_info.action_mask is not None:
                    if len(agent_info.action_mask) == a_size:
                        mask_matrix[agent_index, :] = [
                            False if agent_info.action_mask[k] else True
                            for k in range(a_size)
                        ]
            action_mask = (1 - mask_matrix).astype(np.bool)
            indices = _generate_split_indices(
                behavior_spec.action_spec.discrete_branches
            )
            action_mask = np.split(action_mask, indices, axis=1)
    return (
        DecisionSteps(
            decision_obs_list,
            decision_rewards,
            decision_agent_id,
            action_mask,
            decision_group_id,
            decision_group_rewards,
        ),
        TerminalSteps(
            terminal_obs_list,
            terminal_rewards,
            max_step,
            terminal_agent_id,
            terminal_group_id,
            terminal_group_rewards,
        ),
    )


def _generate_split_indices(dims):
    if len(dims) <= 1:
        return ()
    result = (dims[0],)
    for i in range(len(dims) - 2):
        result += (dims[i + 1] + result[i],)
    return result