from typing import Dict, Union

import gym
import numpy as np

from stable_baselines3.common.type_aliases import GymStepReturn


class SimpleMultiObsEnv(gym.Env):
    def __init__(
        self,
        num_col: int = 4,
        num_row: int = 4,
        random_start: bool = True,
        discrete_actions: bool = True,
        channel_last: bool = True,
    ):
        super(SimpleMultiObsEnv, self).__init__()
        self.vector_size = 5
        if channel_last:
            self.img_size = [64, 64, 1]
        else:
            self.img_size = [1, 64, 64]
        self.random_start = random_start
        self.discrete_actions = discrete_actions
        if discrete_actions:
            self.action_space = gym.spaces.Discrete(4)
        else:
            self.action_space = gym.spaces.Box(0, 1, (4,))
        self.observation_space = gym.spaces.Dict(
            spaces={
                "vec": gym.spaces.Box(0, 1, (self.vector_size,)),
                "img": gym.spaces.Box(0, 255, self.img_size, dtype=np.uint8),
            }
        )
        self.count = 0
        self.max_count = 100
        self.log = ""
        self.state = 0
        self.action2str = ["left", "down", "right", "up"]
        self.init_possible_transitions()
        self.num_col = num_col
        self.state_mapping = []
        self.init_state_mapping(num_col, num_row)
        self.max_state = len(self.state_mapping) - 1

    def init_state_mapping(self, num_col: int, num_row: int) -> None:
        col_vecs = np.random.random((num_col, self.vector_size))
        row_imgs = np.random.randint(0, 255, (num_row, 64, 64), dtype=np.int32)
        for i in range(num_col):
            for j in range(num_row):
                self.state_mapping.append(
                    {"vec": col_vecs[i], "img": row_imgs[j].reshape(self.img_size)}
                )

    def get_state_mapping(self) -> Dict[str, np.ndarray]:
        return self.state_mapping[self.state]

    def init_possible_transitions(self) -> None:
        self.left_possible = [1, 2, 3, 13, 14, 15]
        self.down_possible = [0, 4, 8, 3, 7, 11]
        self.right_possible = [0, 1, 2, 12, 13, 14]
        self.up_possible = [4, 8, 12, 7, 11, 15]

    def step(self, action: Union[int, float, np.ndarray]) -> GymStepReturn:
        if not self.discrete_actions:
            action = np.argmax(action)
        else:
            action = int(action)
        self.count += 1
        prev_state = self.state
        reward = -0.1
        if self.state in self.left_possible and action == 0:
            self.state -= 1
        elif self.state in self.down_possible and action == 1:
            self.state += self.num_col
        elif self.state in self.right_possible and action == 2:
            self.state += 1
        elif self.state in self.up_possible and action == 3:
            self.state -= self.num_col
        got_to_end = self.state == self.max_state
        reward = 1 if got_to_end else reward
        done = self.count > self.max_count or got_to_end
        self.log = f"Went {self.action2str[action]} in state {prev_state}, got to state {self.state}"
        return self.get_state_mapping(), reward, done, {"got_to_end": got_to_end}

    def render(self, mode: str = "human") -> None:
        print(self.log)

    def reset(self) -> Dict[str, np.ndarray]:
        self.count = 0
        if not self.random_start:
            self.state = 0
        else:
            self.state = np.random.randint(0, self.max_state)
        return self.state_mapping[self.state]
