import os
import time
from datetime import datetime
from os.path import join
from typing import Dict, Any, Tuple, List, Set

import gym
from gym import spaces

from isaacgym import gymtorch, gymapi
from isaacgymenvs.utils.torch_jit_utils import to_torch
from isaacgymenvs.utils.dr_utils import get_property_setter_map, get_property_getter_map, \
    get_default_setter_args, apply_random_samples, check_buckets, generate_random_samples

import torch
import numpy as np
import operator, random
from copy import deepcopy
from isaacgymenvs.utils.utils import nested_dict_get_attr, nested_dict_set_attr

from collections import deque

import sys

import abc
from abc import ABC

EXISTING_SIM = None
SCREEN_CAPTURE_RESOLUTION = (1027, 768)

def _create_sim_once(gym, *args, **kwargs):
    global EXISTING_SIM
    if EXISTING_SIM is not None:
        return EXISTING_SIM
    else:
        EXISTING_SIM = gym.create_sim(*args, **kwargs)
        return EXISTING_SIM

class Env(ABC):
    def __init__(self, config: Dict[str, Any], rl_device: str, sim_device: str, graphics_device_id: int, headless: bool):
        """Initialise the env.

        Args:
            config: the configuration dictionary.
            sim_device: the device to simulate physics on. eg. 'cuda:0' or 'cpu'
            graphics_device_id: the device ID to render with.
            headless: Set to False to disable viewer rendering.
        """

        split_device = sim_device.split(":")
        self.device_type = split_device[0]
        self.device_id = int(split_device[1]) if len(split_device) > 1 else 0

        self.device = "cpu"
        if config["sim"]["use_gpu_pipeline"]:
            if self.device_type.lower() == "cuda" or self.device_type.lower() == "gpu":
                self.device = "cuda" + ":" + str(self.device_id)
            else:
                print("GPU Pipeline can only be used with GPU simulation. Forcing CPU Pipeline.")
                config["sim"]["use_gpu_pipeline"] = False

        self.rl_device = rl_device

        # Rendering
        # if training in a headless mode
        self.headless = headless

        enable_camera_sensors = config["env"].get("enableCameraSensors", False)
        self.graphics_device_id = graphics_device_id
        if enable_camera_sensors == False and self.headless == True:
            self.graphics_device_id = -1

        self.num_environments = config["env"]["numEnvs"]
        self.num_agents = config["env"].get("numAgents", 1)  # used for multi-agent environments

        self.num_observations = config["env"].get("numObservations", 0)
        self.num_states = config["env"].get("numStates", 0)

        self.obs_space = spaces.Box(np.ones(self.num_obs) * -np.Inf, np.ones(self.num_obs) * np.Inf)
        self.state_space = spaces.Box(np.ones(self.num_states) * -np.Inf, np.ones(self.num_states) * np.Inf)

        self.num_actions = config["env"]["numActions"]
        self.control_freq_inv = config["env"].get("controlFrequencyInv", 1)

        self.act_space = spaces.Box(np.ones(self.num_actions) * -1., np.ones(self.num_actions) * 1.)

        self.clip_obs = config["env"].get("clipObservations", np.Inf)
        self.clip_actions = config["env"].get("clipActions", np.Inf)

        # Total number of training frames since the beginning of the experiment.
        # We get this information from the learning algorithm rather than tracking ourselves.
        # The learning algorithm tracks the total number of frames since the beginning of training and accounts for
        # experiments restart/resumes. This means this number can be > 0 right after initialization if we resume the
        # experiment.
        self.total_train_env_frames: int = 0

        # number of control steps
        self.control_steps: int = 0

        self.render_fps: int = config["env"].get("renderFPS", -1)
        self.last_frame_time: float = 0.0

        self.record_frames: bool = False
        self.record_frames_dir = join("recorded_frames", datetime.now().strftime("%Y-%m-%d_%H-%M-%S"))

    @abc.abstractmethod
    def allocate_buffers(self):
        """Create torch buffers for observations, rewards, actions dones and any additional data."""

    @abc.abstractmethod
    def step(self, actions: torch.Tensor) -> Tuple[Dict[str, torch.Tensor], torch.Tensor, torch.Tensor, Dict[str, Any]]:
        """Step the physics of the environment.
        Args:
            actions: actions to apply
        Returns:
            Observations, rewards, resets, info
            Observations are dict of observations (currently only one member called 'obs')
        """

    @abc.abstractmethod
    def reset(self)-> Dict[str, torch.Tensor]:
        """Reset the environment.
        Returns:
            Observation dictionary
        """

    @abc.abstractmethod
    def reset_idx(self, env_ids: torch.Tensor):
        """Reset environments having the provided indices.
        Args:
            env_ids: environments to reset
        """

    @property
    def observation_space(self) -> gym.Space:
        """Get the environment's observation space."""
        return self.obs_space

    @property
    def action_space(self) -> gym.Space:
        """Get the environment's action space."""
        return self.act_space

    @property
    def num_envs(self) -> int:
        """Get the number of environments."""
        return self.num_environments

    @property
    def num_acts(self) -> int:
        """Get the number of actions in the environment."""
        return self.num_actions

    @property
    def num_obs(self) -> int:
        """Get the number of observations in the environment."""
        return self.num_observations

    def set_train_info(self, env_frames, *args, **kwargs):
        """
        Send the information in the direction algo->environment.
        Most common use case: tell the environment how far along we are in the training process. This is useful
        for implementing curriculums and things such as that.
        """
        self.total_train_env_frames = env_frames
        # print(f'env_frames updated to {self.total_train_env_frames}')

    def get_env_state(self):
        """
        Return serializable environment state to be saved to checkpoint.
        Can be used for stateful training sessions, i.e. with adaptive curriculums.
        """
        return None

    def set_env_state(self, env_state):
        pass

class VecTask(Env):

    metadata = {"render.modes": ["human", "rgb_array"], "video.frames_per_second": 24}

    def __init__(self, config, rl_device, sim_device, graphics_device_id, headless, virtual_screen_capture: bool = False, force_render: bool = False):
        """Initialise the `VecTask`.

        Args:
            config: config dictionary for the environment.
            sim_device: the device to simulate physics on. eg. 'cuda:0' or 'cpu'
            graphics_device_id: the device ID to render with.
            headless: Set to False to disable viewer rendering.
            virtual_screen_capture: Set to True to allow the users get captured screen in RGB array via `env.render(mode='rgb_array')`.
            force_render: Set to True to always force rendering in the steps (if the `control_freq_inv` is greater than 1 we suggest stting this arg to True)
        """
        # super().__init__(config, rl_device, sim_device, graphics_device_id, headless, use_dict_obs)
        super().__init__(config, rl_device, sim_device, graphics_device_id, headless)
        self.virtual_screen_capture = virtual_screen_capture
        self.virtual_display = None
        if self.virtual_screen_capture:
            from pyvirtualdisplay.smartdisplay import SmartDisplay
            self.virtual_display = SmartDisplay(size=SCREEN_CAPTURE_RESOLUTION)
            self.virtual_display.start()
        self.force_render = force_render

        self.sim_params = self.__parse_sim_params(self.cfg["physics_engine"], self.cfg["sim"])
        if self.cfg["physics_engine"] == "physx":
            self.physics_engine = gymapi.SIM_PHYSX
        elif self.cfg["physics_engine"] == "flex":
            self.physics_engine = gymapi.SIM_FLEX
        else:
            msg = f"Invalid physics engine backend: {self.cfg['physics_engine']}"
            raise ValueError(msg)

        self.dt: float = self.sim_params.dt

        # optimization flags for pytorch JIT
        torch._C._jit_set_profiling_mode(False)
        torch._C._jit_set_profiling_executor(False)

        self.gym = gymapi.acquire_gym()

        self.first_randomization = True
        self.original_props = {}
        self.dr_randomizations = {}
        self.actor_params_generator = None
        self.extern_actor_params = {}
        self.last_step = -1
        self.last_rand_step = -1
        for env_id in range(self.num_envs):
            self.extern_actor_params[env_id] = None

        # create envs, sim and viewer
        self.sim_initialized = False
        self.create_sim()
        self.gym.prepare_sim(self.sim)
        self.sim_initialized = True

        self.set_viewer()
        self.allocate_buffers()

        self.obs_dict = {}

    def set_viewer(self):
        """Create the viewer."""

        # todo: read from config
        self.enable_viewer_sync = True
        self.viewer = None

        # if running with a viewer, set up keyboard shortcuts and camera
        if self.headless == False:
            # subscribe to keyboard shortcuts
            self.viewer = self.gym.create_viewer(
                self.sim, gymapi.CameraProperties())
            self.gym.subscribe_viewer_keyboard_event(
                self.viewer, gymapi.KEY_ESCAPE, "QUIT")
            self.gym.subscribe_viewer_keyboard_event(
                self.viewer, gymapi.KEY_V, "toggle_viewer_sync")
            self.gym.subscribe_viewer_keyboard_event(
                self.viewer, gymapi.KEY_R, "record_frames")

            # set the camera position based on up axis
            sim_params = self.gym.get_sim_params(self.sim)
            if sim_params.up_axis == gymapi.UP_AXIS_Z:
                cam_pos = gymapi.Vec3(20.0, 25.0, 3.0)
                cam_target = gymapi.Vec3(10.0, 15.0, 0.0)
            else:
                cam_pos = gymapi.Vec3(20.0, 3.0, 25.0)
                cam_target = gymapi.Vec3(10.0, 0.0, 15.0)

            self.gym.viewer_camera_look_at(
                self.viewer, None, cam_pos, cam_target)

    def allocate_buffers(self):
        """Allocate the observation, states, etc. buffers.

        These are what is used to set observations and states in the environment classes which
        inherit from this one, and are read in `step` and other related functions.

        """

        # allocate buffers
        self.obs_buf = torch.zeros(
            (self.num_envs, self.num_obs), device=self.device, dtype=torch.float)
        self.states_buf = torch.zeros(
            (self.num_envs, self.num_states), device=self.device, dtype=torch.float)
        self.rew_buf = torch.zeros(
            self.num_envs, device=self.device, dtype=torch.float)
        self.reset_buf = torch.ones(
            self.num_envs, device=self.device, dtype=torch.long)
        self.timeout_buf = torch.zeros(
             self.num_envs, device=self.device, dtype=torch.long)
        self.progress_buf = torch.zeros(
            self.num_envs, device=self.device, dtype=torch.long)
        self.randomize_buf = torch.zeros(
            self.num_envs, device=self.device, dtype=torch.long)
        self.extras = {}

    def create_sim(self, compute_device: int, graphics_device: int, physics_engine, sim_params: gymapi.SimParams):
        """Create an Isaac Gym sim object.

        Args:
            compute_device: ID of compute device to use.
            graphics_device: ID of graphics device to use.
            physics_engine: physics engine to use (`gymapi.SIM_PHYSX` or `gymapi.SIM_FLEX`)
            sim_params: sim params to use.
        Returns:
            the Isaac Gym sim object.
        """
        sim = _create_sim_once(self.gym, compute_device, graphics_device, physics_engine, sim_params)
        if sim is None:
            print("*** Failed to create sim")
            quit()

        return sim

    def get_state(self):
        """Returns the state buffer of the environment (the privileged observations for asymmetric training)."""
        return torch.clamp(self.states_buf, -self.clip_obs, self.clip_obs).to(self.rl_device)

    @abc.abstractmethod
    def pre_physics_step(self, actions: torch.Tensor):
        """Apply the actions to the environment (eg by setting torques, position targets).

        Args:
            actions: the actions to apply
        """

    @abc.abstractmethod
    def post_physics_step(self):
        """Compute reward and observations, reset any environments that require it."""

    def step(self, actions: torch.Tensor) -> Tuple[Dict[str, torch.Tensor], torch.Tensor, torch.Tensor, Dict[str, Any]]:
        """Step the physics of the environment.

        Args:
            actions: actions to apply
        Returns:
            Observations, rewards, resets, info
            Observations are dict of observations (currently only one member called 'obs')
        """

        # randomize actions
        if self.dr_randomizations.get('actions', None):
            actions = self.dr_randomizations['actions']['noise_lambda'](actions)

        action_tensor = torch.clamp(actions, -self.clip_actions, self.clip_actions)
        # apply actions
        self.pre_physics_step(action_tensor)

        # step physics and render each frame
        for i in range(self.control_freq_inv):
            if self.force_render:
                self.render()
            self.gym.simulate(self.sim)

        # to fix!
        if self.device == 'cpu':
            self.gym.fetch_results(self.sim, True)

        # compute observations, rewards, resets, ...
        self.post_physics_step()

        self.control_steps += 1

        # fill time out buffer: set to 1 if we reached the max episode length AND the reset buffer is 1. Timeout == 1 makes sense only if the reset buffer is 1.
        self.timeout_buf = (self.progress_buf >= self.max_episode_length - 1) & (self.reset_buf != 0)

        # randomize observations
        if self.dr_randomizations.get('observations', None):
            self.obs_buf = self.dr_randomizations['observations']['noise_lambda'](self.obs_buf)

        self.extras["time_outs"] = self.timeout_buf.to(self.rl_device)

        self.obs_dict["obs"] = torch.clamp(self.obs_buf, -self.clip_obs, self.clip_obs).to(self.rl_device)

        # asymmetric actor-critic
        if self.num_states > 0:
            self.obs_dict["states"] = self.get_state()

        return self.obs_dict, self.rew_buf.to(self.rl_device), self.reset_buf.to(self.rl_device), self.extras

    def zero_actions(self) -> torch.Tensor:
        """Returns a buffer with zero actions.

        Returns:
            A buffer of zero torch actions
        """
        actions = torch.zeros([self.num_envs, self.num_actions], dtype=torch.float32, device=self.rl_device)

        return actions

    def reset_idx(self, env_idx):
        """Reset environment with indces in env_idx.
        Should be implemented in an environment class inherited from VecTask.
        """
        pass

    def reset(self):
        """Is called only once when environment starts to provide the first observations.
        Doesn't calculate observations. Actual reset and observation calculation need to be implemented by user.
        Returns:
            Observation dictionary
        """
        self.obs_dict["obs"] = torch.clamp(self.obs_buf, -self.clip_obs, self.clip_obs).to(self.rl_device)

        # asymmetric actor-critic
        if self.num_states > 0:
            self.obs_dict["states"] = self.get_state()

        return self.obs_dict

    def reset_done(self):
        """Reset the environment.
        Returns:
            Observation dictionary, indices of environments being reset
        """
        done_env_ids = self.reset_buf.nonzero(as_tuple=False).flatten()
        if len(done_env_ids) > 0:
            self.reset_idx(done_env_ids)

        self.obs_dict["obs"] = torch.clamp(self.obs_buf, -self.clip_obs, self.clip_obs).to(self.rl_device)

        # asymmetric actor-critic
        if self.num_states > 0:
            self.obs_dict["states"] = self.get_state()

        return self.obs_dict, done_env_ids

    def render(self, mode="rgb_array"):
        """Draw the frame to the viewer, and check for keyboard events."""
        if self.viewer:
            # check for window closed
            if self.gym.query_viewer_has_closed(self.viewer):
                sys.exit()

            # check for keyboard events
            for evt in self.gym.query_viewer_action_events(self.viewer):
                if evt.action == "QUIT" and evt.value > 0:
                    sys.exit()
                elif evt.action == "toggle_viewer_sync" and evt.value > 0:
                    self.enable_viewer_sync = not self.enable_viewer_sync
                elif evt.action == "record_frames" and evt.value > 0:
                    self.record_frames = not self.record_frames

            # fetch results
            if self.device != 'cpu':
                self.gym.fetch_results(self.sim, True)

            # step graphics
            if self.enable_viewer_sync:
                self.gym.step_graphics(self.sim)
                self.gym.draw_viewer(self.viewer, self.sim, True)

                # Wait for dt to elapse in real time.
                # This synchronizes the physics simulation with the rendering rate.
                self.gym.sync_frame_time(self.sim)

                # it seems like in some cases sync_frame_time still results in higher-than-realtime framerate
                # this code will slow down the rendering to real time
                now = time.time()
                delta = now - self.last_frame_time
                if self.render_fps < 0:
                    # render at control frequency
                    render_dt = self.dt * self.control_freq_inv  # render every control step
                else:
                    render_dt = 1.0 / self.render_fps

                if delta < render_dt:
                    time.sleep(render_dt - delta)

                self.last_frame_time = time.time()

            else:
                self.gym.poll_viewer_events(self.viewer)

            if self.record_frames:
                if not os.path.isdir(self.record_frames_dir):
                    os.makedirs(self.record_frames_dir, exist_ok=True)

                self.gym.write_viewer_image_to_file(self.viewer, join(self.record_frames_dir, f"frame_{self.control_steps}.png"))

            if self.virtual_display and mode == "rgb_array":
                img = self.virtual_display.grab()
                return np.array(img)

    def __parse_sim_params(self, physics_engine: str, config_sim: Dict[str, Any]) -> gymapi.SimParams:
        """Parse the config dictionary for physics stepping settings.

        Args:
            physics_engine: which physics engine to use. "physx" or "flex"
            config_sim: dict of sim configuration parameters
        Returns
            IsaacGym SimParams object with updated settings.
        """
        sim_params = gymapi.SimParams()

        # check correct up-axis
        if config_sim["up_axis"] not in ["z", "y"]:
            msg = f"Invalid physics up-axis: {config_sim['up_axis']}"
            print(msg)
            raise ValueError(msg)

        # assign general sim parameters
        sim_params.dt = config_sim["dt"]
        sim_params.num_client_threads = config_sim.get("num_client_threads", 0)
        sim_params.use_gpu_pipeline = config_sim["use_gpu_pipeline"]
        sim_params.substeps = config_sim.get("substeps", 2)

        # assign up-axis
        if config_sim["up_axis"] == "z":
            sim_params.up_axis = gymapi.UP_AXIS_Z
        else:
            sim_params.up_axis = gymapi.UP_AXIS_Y

        # assign gravity
        sim_params.gravity = gymapi.Vec3(*config_sim["gravity"])

        # configure physics parameters
        if physics_engine == "physx":
            # set the parameters
            if "physx" in config_sim:
                for opt in config_sim["physx"].keys():
                    if opt == "contact_collection":
                        setattr(sim_params.physx, opt, gymapi.ContactCollection(config_sim["physx"][opt]))
                    else:
                        setattr(sim_params.physx, opt, config_sim["physx"][opt])
        else:
            # set the parameters
            if "flex" in config_sim:
                for opt in config_sim["flex"].keys():
                    setattr(sim_params.flex, opt, config_sim["flex"][opt])

        # return the configured params
        return sim_params

    """
    Domain Randomization methods
    """

    def get_actor_params_info(self, dr_params: Dict[str, Any], env):
        """Generate a flat array of actor params, their names and ranges.

        Returns:
            The array
        """

        if "actor_params" not in dr_params:
            return None
        params = []
        names = []
        lows = []
        highs = []
        param_getters_map = get_property_getter_map(self.gym)
        for actor, actor_properties in dr_params["actor_params"].items():
            handle = self.gym.find_actor_handle(env, actor)
            for prop_name, prop_attrs in actor_properties.items():
                if prop_name == 'color':
                    continue  # this is set randomly
                props = param_getters_map[prop_name](env, handle)
                if not isinstance(props, list):
                    props = [props]
                for prop_idx, prop in enumerate(props):
                    for attr, attr_randomization_params in prop_attrs.items():
                        name = prop_name+'_' + str(prop_idx) + '_'+attr
                        lo_hi = attr_randomization_params['range']
                        distr = attr_randomization_params['distribution']
                        if 'uniform' not in distr:
                            lo_hi = (-1.0*float('Inf'), float('Inf'))
                        if isinstance(prop, np.ndarray):
                            for attr_idx in range(prop[attr].shape[0]):
                                params.append(prop[attr][attr_idx])
                                names.append(name+'_'+str(attr_idx))
                                lows.append(lo_hi[0])
                                highs.append(lo_hi[1])
                        else:
                            params.append(getattr(prop, attr))
                            names.append(name)
                            lows.append(lo_hi[0])
                            highs.append(lo_hi[1])
        return params, names, lows, highs

    def apply_randomizations(self, dr_params):
        """Apply domain randomizations to the environment.

        Note that currently we can only apply randomizations only on resets, due to current PhysX limitations

        Args:
            dr_params: parameters for domain randomization to use.
        """

        # If we don't have a randomization frequency, randomize every step
        rand_freq = dr_params.get("frequency", 1)

        # First, determine what to randomize:
        #   - non-environment parameters when > frequency steps have passed since the last non-environment
        #   - physical environments in the reset buffer, which have exceeded the randomization frequency threshold
        #   - on the first call, randomize everything
        self.last_step = self.gym.get_frame_count(self.sim)
        if self.first_randomization:
            do_nonenv_randomize = True
            env_ids = list(range(self.num_envs))
        else:
            do_nonenv_randomize = (self.last_step - self.last_rand_step) >= rand_freq
            rand_envs = torch.where(self.randomize_buf >= rand_freq, torch.ones_like(self.randomize_buf), torch.zeros_like(self.randomize_buf))
            rand_envs = torch.logical_and(rand_envs, self.reset_buf)
            env_ids = torch.nonzero(rand_envs, as_tuple=False).squeeze(-1).tolist()
            self.randomize_buf[rand_envs] = 0

        if do_nonenv_randomize:
            self.last_rand_step = self.last_step

        param_setters_map = get_property_setter_map(self.gym)
        param_setter_defaults_map = get_default_setter_args(self.gym)
        param_getters_map = get_property_getter_map(self.gym)

        # On first iteration, check the number of buckets
        if self.first_randomization:
            check_buckets(self.gym, self.envs, dr_params)

        for nonphysical_param in ["observations", "actions"]:
            if nonphysical_param in dr_params and do_nonenv_randomize:
                dist = dr_params[nonphysical_param]["distribution"]
                op_type = dr_params[nonphysical_param]["operation"]
                sched_type = dr_params[nonphysical_param]["schedule"] if "schedule" in dr_params[nonphysical_param] else None
                sched_step = dr_params[nonphysical_param]["schedule_steps"] if "schedule" in dr_params[nonphysical_param] else None
                op = operator.add if op_type == 'additive' else operator.mul

                if sched_type == 'linear':
                    sched_scaling = 1.0 / sched_step * \
                        min(self.last_step, sched_step)
                elif sched_type == 'constant':
                    sched_scaling = 0 if self.last_step < sched_step else 1
                else:
                    sched_scaling = 1

                if dist == 'gaussian':
                    mu, var = dr_params[nonphysical_param]["range"]
                    mu_corr, var_corr = dr_params[nonphysical_param].get("range_correlated", [0., 0.])

                    if op_type == 'additive':
                        mu *= sched_scaling
                        var *= sched_scaling
                        mu_corr *= sched_scaling
                        var_corr *= sched_scaling
                    elif op_type == 'scaling':
                        var = var * sched_scaling  # scale up var over time
                        mu = mu * sched_scaling + 1.0 * \
                            (1.0 - sched_scaling)  # linearly interpolate

                        var_corr = var_corr * sched_scaling  # scale up var over time
                        mu_corr = mu_corr * sched_scaling + 1.0 * \
                            (1.0 - sched_scaling)  # linearly interpolate

                    def noise_lambda(tensor, param_name=nonphysical_param):
                        params = self.dr_randomizations[param_name]
                        corr = params.get('corr', None)
                        if corr is None:
                            corr = torch.randn_like(tensor)
                            params['corr'] = corr
                        corr = corr * params['var_corr'] + params['mu_corr']
                        return op(
                            tensor, corr + torch.randn_like(tensor) * params['var'] + params['mu'])

                    self.dr_randomizations[nonphysical_param] = {'mu': mu, 'var': var, 'mu_corr': mu_corr, 'var_corr': var_corr, 'noise_lambda': noise_lambda}

                elif dist == 'uniform':
                    lo, hi = dr_params[nonphysical_param]["range"]
                    lo_corr, hi_corr = dr_params[nonphysical_param].get("range_correlated", [0., 0.])

                    if op_type == 'additive':
                        lo *= sched_scaling
                        hi *= sched_scaling
                        lo_corr *= sched_scaling
                        hi_corr *= sched_scaling
                    elif op_type == 'scaling':
                        lo = lo * sched_scaling + 1.0 * (1.0 - sched_scaling)
                        hi = hi * sched_scaling + 1.0 * (1.0 - sched_scaling)
                        lo_corr = lo_corr * sched_scaling + 1.0 * (1.0 - sched_scaling)
                        hi_corr = hi_corr * sched_scaling + 1.0 * (1.0 - sched_scaling)

                    def noise_lambda(tensor, param_name=nonphysical_param):
                        params = self.dr_randomizations[param_name]
                        corr = params.get('corr', None)
                        if corr is None:
                            corr = torch.randn_like(tensor)
                            params['corr'] = corr
                        corr = corr * (params['hi_corr'] - params['lo_corr']) + params['lo_corr']
                        return op(tensor, corr + torch.rand_like(tensor) * (params['hi'] - params['lo']) + params['lo'])

                    self.dr_randomizations[nonphysical_param] = {'lo': lo, 'hi': hi, 'lo_corr': lo_corr, 'hi_corr': hi_corr, 'noise_lambda': noise_lambda}

        if "sim_params" in dr_params and do_nonenv_randomize:
            prop_attrs = dr_params["sim_params"]
            prop = self.gym.get_sim_params(self.sim)

            if self.first_randomization:
                self.original_props["sim_params"] = {
                    attr: getattr(prop, attr) for attr in dir(prop)}

            for attr, attr_randomization_params in prop_attrs.items():
                apply_random_samples(
                    prop, self.original_props["sim_params"], attr, attr_randomization_params, self.last_step)

            self.gym.set_sim_params(self.sim, prop)

        # If self.actor_params_generator is initialized: use it to
        # sample actor simulation params. This gives users the
        # freedom to generate samples from arbitrary distributions,
        # e.g. use full-covariance distributions instead of the DR's
        # default of treating each simulation parameter independently.
        extern_offsets = {}
        if self.actor_params_generator is not None:
            for env_id in env_ids:
                self.extern_actor_params[env_id] = \
                    self.actor_params_generator.sample()
                extern_offsets[env_id] = 0

        # randomise all attributes of each actor (hand, cube etc..)
        # actor_properties are (stiffness, damping etc..)

        # Loop over actors, then loop over envs, then loop over their props
        # and lastly loop over the ranges of the params

        for actor, actor_properties in dr_params["actor_params"].items():

            # Loop over all envs as this part is not tensorised yet
            for env_id in env_ids:
                env = self.envs[env_id]
                handle = self.gym.find_actor_handle(env, actor)
                extern_sample = self.extern_actor_params[env_id]

                # randomise dof_props, rigid_body, rigid_shape properties
                # all obtained from the YAML file
                # EXAMPLE: prop name: dof_properties, rigid_body_properties, rigid_shape properties
                #          prop_attrs:
                #               {'damping': {'range': [0.3, 3.0], 'operation': 'scaling', 'distribution': 'loguniform'}
                #               {'stiffness': {'range': [0.75, 1.5], 'operation': 'scaling', 'distribution': 'loguniform'}
                for prop_name, prop_attrs in actor_properties.items():
                    if prop_name == 'color':
                        num_bodies = self.gym.get_actor_rigid_body_count(
                            env, handle)
                        for n in range(num_bodies):
                            self.gym.set_rigid_body_color(env, handle, n, gymapi.MESH_VISUAL,
                                                          gymapi.Vec3(random.uniform(0, 1), random.uniform(0, 1), random.uniform(0, 1)))
                        continue

                    if prop_name == 'scale':
                        setup_only = prop_attrs.get('setup_only', False)
                        if (setup_only and not self.sim_initialized) or not setup_only:
                            attr_randomization_params = prop_attrs
                            sample = generate_random_samples(attr_randomization_params, 1,
                                                             self.last_step, None)
                            og_scale = 1
                            if attr_randomization_params['operation'] == 'scaling':
                                new_scale = og_scale * sample
                            elif attr_randomization_params['operation'] == 'additive':
                                new_scale = og_scale + sample
                            self.gym.set_actor_scale(env, handle, new_scale)
                        continue

                    prop = param_getters_map[prop_name](env, handle)
                    set_random_properties = True

                    if isinstance(prop, list):
                        if self.first_randomization:
                            self.original_props[prop_name] = [
                                {attr: getattr(p, attr) for attr in dir(p)} for p in prop]
                        for p, og_p in zip(prop, self.original_props[prop_name]):
                            for attr, attr_randomization_params in prop_attrs.items():
                                setup_only = attr_randomization_params.get('setup_only', False)
                                if (setup_only and not self.sim_initialized) or not setup_only:
                                    smpl = None
                                    if self.actor_params_generator is not None:
                                        smpl, extern_offsets[env_id] = get_attr_val_from_sample(
                                            extern_sample, extern_offsets[env_id], p, attr)
                                    apply_random_samples(
                                        p, og_p, attr, attr_randomization_params,
                                        self.last_step, smpl)
                                else:
                                    set_random_properties = False
                    else:
                        if self.first_randomization:
                            self.original_props[prop_name] = deepcopy(prop)
                        for attr, attr_randomization_params in prop_attrs.items():
                            setup_only = attr_randomization_params.get('setup_only', False)
                            if (setup_only and not self.sim_initialized) or not setup_only:
                                smpl = None
                                if self.actor_params_generator is not None:
                                    smpl, extern_offsets[env_id] = get_attr_val_from_sample(
                                        extern_sample, extern_offsets[env_id], prop, attr)
                                apply_random_samples(
                                    prop, self.original_props[prop_name], attr,
                                    attr_randomization_params, self.last_step, smpl)
                            else:
                                set_random_properties = False

                    if set_random_properties:
                        setter = param_setters_map[prop_name]
                        default_args = param_setter_defaults_map[prop_name]
                        setter(env, handle, prop, *default_args)

        if self.actor_params_generator is not None:
            for env_id in env_ids:  # check that we used all dims in sample
                if extern_offsets[env_id] > 0:
                    extern_sample = self.extern_actor_params[env_id]
                    if extern_offsets[env_id] != extern_sample.shape[0]:
                        print('env_id', env_id,
                              'extern_offset', extern_offsets[env_id],
                              'vs extern_sample.shape', extern_sample.shape)
                        raise Exception("Invalid extern_sample size")

        self.first_randomization = False


dx = [-1, 0, 1, 0]
dy = [0, -1, 0, 1]

def on_board(stu):
    temp = []
    for i in range(n):
        for j in range(n):
            cnt, blank = 0, 0
            if board[i][j] == 0:
                for k in range(4):
                    nx = i + dx[k]
                    ny = j + dy[k]
                    if nx < 0 or nx >= n or ny < 0 or ny >= n:
                        continue
                    if board[nx][ny] in stu:
                        cnt += 1
                    if board[nx][ny] == 0:
                        blank += 1
                    temp.append([cnt, blank, i, j])
    temp.sort(key=lambda x: (-x[0], -x[1], x[2], x[3]))
    board[temp[0][2]][temp[0][3]] = stu[0]

def sum_point():
    point = 0
    for i in range(n):
        for j in range(n):
            cnt = 0
            for k in range(4):
                nx = i + dx[k]
                ny = j + dy[k]
                if nx < 0 or nx >= n or ny < 0 or ny >= n:
                    continue
                if board[nx][ny] in dict_st[board[i][j]]:
                    cnt += 1
            if cnt > 0:
                point += 10**(cnt-1)
    return point

n = int(input())
board = [[0 for _ in range(n)] for _ in range(n)]
student = [list(map(int, input().split())) for _ in range(n**2)]
dict_st = [[] for _ in range(n**2+1)]
for st in student:
    on_board(st)
    # for i in range(n):
    #     print(board[i])
    # print()
    dict_st[st[0]] = [st[1], st[2], st[3], st[4]]
# print(dict_st)
print(sum_point())

# 문제 접근 방법
# # cnt와 blank를 이용하여 주위의 값을 파악하고
# # 이를 i, j와 함께 배열에 담아 마지막에 sort하여 0번째 idx 값을 추출


r"""
STB dataset
A Hand joint Tracking Benchmark from Stereo Matching, ICIP 2017
"""

import math
import os
import pickle

import PIL
import cv2
import matplotlib.pyplot as plt
import numpy as np
import scipy.io as sio
import torch.utils.data
from PIL import Image
from termcolor import colored
from tqdm import tqdm

import config as cfg
import utils.handutils as handutils

CACHE_HOME = os.path.expanduser(cfg.DEFAULT_CACHE_DIR)

# some globals, ugly but work
sk_fx_color = 607.92271
sk_fy_color = 607.88192
sk_tx_color = 314.78337
sk_ty_color = 236.42484

bb_fx = 822.79041
bb_fy = 822.79041
bb_tx = 318.47345
bb_ty = 250.31296

sk_rot_vec = [0.00531, -0.01196, 0.00301]
sk_trans_vec = [-24.0381, -0.4563, -1.2326]  # mm

snap_joint_name2id = {w: i for i, w in enumerate(cfg.snap_joint_names)}
stb_joint_name2id = {w: i for i, w in enumerate(cfg.stb_joints)}

stb_to_snap_id = [snap_joint_name2id[joint_name] for joint_name in cfg.stb_joints]

def sk_rot_mx(rot_vec):
    """
    use Rodrigues' rotation formula to transform the rotation vector into rotation matrix
    :param rot_vec:
    :return:
    """
    theta = np.linalg.norm(rot_vec)
    vector = np.array(rot_vec) * math.sin(theta / 2.0) / theta
    a = math.cos(theta / 2.0)
    b = -vector[0]
    c = -vector[1]
    d = -vector[2]
    return np.array(
        [
            [
                a * a + b * b - c * c - d * d,
                2 * (b * c + a * d),
                2 * (b * d - a * c)
            ],
            [
                2 * (b * c - a * d),
                a * a + c * c - b * b - d * d,
                2 * (c * d + a * b)
            ],
            [
                2 * (b * d + a * c),
                2 * (c * d - a * b),
                a * a + d * d - b * b - c * c
            ]
        ]
    )

def sk_xyz_depth2color(depth_xyz, trans_vec, rot_mx):
    """
    in the STB dataset: 'rotation and translation vector can transform the coordinates
                         relative to color camera to those relative to depth camera'.
    however here we want depth_xyz -> color_xyz
    a inverse transformation happen:
    T = [rot_mx | trans_vec | 0  1], Tinv = T.inv, then output Tinv * depth_xyz

    :param depth_xyz: N x 21 x 3, trans_vec: 3, rot_mx: 3 x 3
    :return: color_xyz: N x 21 x 3
    """
    color_xyz = depth_xyz - np.tile(trans_vec, [depth_xyz.shape[0], depth_xyz.shape[1], 1])
    return color_xyz.dot(rot_mx)

def stb_palm2wrist(joint_xyz):
    root = snap_joint_name2id['loc_bn_palm_L']  # 0

    index = snap_joint_name2id['loc_bn_index_L_01']  # 5
    mid = snap_joint_name2id['loc_bn_mid_L_01']  # 9
    ring = snap_joint_name2id['loc_bn_ring_L_01']  # 13
    pinky = snap_joint_name2id['loc_bn_pinky_L_01']  # 17

    def _new_root(joint_xyz, id, root_id):
        return joint_xyz[:, id, :] + \
               2.25 * (joint_xyz[:, root_id, :] - joint_xyz[:, id, :])  # N x K x 3

    joint_xyz[:, root, :] = \
        _new_root(joint_xyz, index, root) + \
        _new_root(joint_xyz, mid, root) + \
        _new_root(joint_xyz, ring, root) + \
        _new_root(joint_xyz, pinky, root)
    joint_xyz[:, root, :] = joint_xyz[:, root, :] / 4.0

    return joint_xyz

def _stb_palm2wrist(joint_xyz):
    root_id = snap_joint_name2id['loc_bn_palm_L']  # 0
    mid_root_id = snap_joint_name2id['loc_bn_mid_L_01']  # 9
    joint_xyz[:, root_id, :] = \
        joint_xyz[:, mid_root_id, :] + \
        2.2 * (joint_xyz[:, root_id, :] - joint_xyz[:, mid_root_id, :])  # N x K x 3
    return joint_xyz

def ge_palm2wrist(pose_xyz):
    root_id = snap_joint_name2id['loc_bn_palm_L']
    ring_root_id = snap_joint_name2id['loc_bn_ring_L_01']
    pose_xyz[:, root_id, :] = pose_xyz[:, ring_root_id, :] + \
                              2.0 * (pose_xyz[:, root_id, :] - pose_xyz[:, ring_root_id, :])  # N x K x 3
    return pose_xyz

class STBDataset(torch.utils.data.Dataset):
    def __init__(
            self,
            data_root,
            data_split='train',
            hand_side='right',
            njoints=21,
            use_cache=True,
            visual=False
    ):
        if not os.path.exists(data_root):
            raise ValueError("data_root: %s not exist" % data_root)
        self.name = 'stb'
        self.data_split = data_split
        self.hand_side = hand_side
        self.img_paths = []
        self.dep_paths = []
        self.joints = []
        self.kp2ds = []
        self.centers = []
        self.my_scales = []
        self.njoints = njoints  # total 21 hand parts
        self.visual = visual

        self.root_id = snap_joint_name2id['loc_bn_palm_L']
        self.mid_mcp_id = snap_joint_name2id['loc_bn_mid_L_01']
        ann_base = os.path.join(data_root, "labels")
        img_base = os.path.join(data_root, "images")
        sk_rot = sk_rot_mx(sk_rot_vec)

        self.sk_intr = np.array([
            [sk_fx_color, 0.0, sk_tx_color],
            [0.0, sk_fy_color, sk_ty_color],
            [0.0, 0.0, 1.0],
        ], dtype=np.float32)  # (3,3)

        self.sequence = []
        if data_split == 'train':
            self.sequence = [
                "B2Counting",
                "B2Random",
                "B3Counting",
                "B3Random",
                "B4Counting",
                "B4Random",
                "B5Counting",
                "B5Random",
                "B6Counting",
                "B6Random"
            ]
        elif data_split == 'test':
            self.sequence = [
                "B1Counting",
                "B1Random"
            ]
        elif data_split == 'val':
            self.sequence = [
                "B2Counting",
                "B2Random"
            ]
        elif data_split == "train_val":
            self.sequence = [
                "B3Counting",
                "B3Random",
                "B4Counting",
                "B4Random",
                "B5Counting",
                "B5Random",
                "B6Counting",
                "B6Random"
            ]
        elif data_split == "all":
            self.sequence = [
                "B1Counting",
                "B1Random",
                "B2Counting",
                "B2Random",
                "B3Counting",
                "B3Random",
                "B4Counting",
                "B4Random",
                "B5Counting",
                "B5Random",
                "B6Counting",
                "B6Random"
            ]
        else:
            raise ValueError("split {} not in [train|test|val|train_val|all]")

        self.cache_folder = os.path.join(CACHE_HOME, "my-{}".format(data_split), "stb")
        os.makedirs(self.cache_folder, exist_ok=True)
        cache_path = os.path.join(
            self.cache_folder, "{}.pkl".format(self.data_split)
        )
        if os.path.exists(cache_path) and use_cache:
            with open(cache_path, "rb") as fid:
                annotations = pickle.load(fid)
                self.img_paths = annotations["img_paths"]
                self.dep_paths = annotations["dep_paths"]
                self.joints = annotations["joints"]
                self.kp2ds = annotations["kp2ds"]
                self.centers = annotations["centers"]
                self.my_scales = annotations["my_scales"]
            print("stb {} gt loaded from {}".format(self.data_split, cache_path))
            return

        self.imgpath_list = [
            os.path.join(img_base, seq) for seq in self.sequence
        ]

        imgsk_prefix = "SK_color"
        depsk_prefix = "SK_depth_seg"

        annsk_list = [
            os.path.join(
                ann_base,
                "{}_{}.mat".format(seq, imgsk_prefix[:2])
            ) for seq in self.sequence
        ]

        self.ann_list = annsk_list

        for imgpath, ann in zip(self.imgpath_list, self.ann_list):
            ''' we only use SK image '''
            assert "SK" in ann
            ''' 1. load joint '''
            rawmat = sio.loadmat(ann)
            rawjoint = rawmat["handPara"].transpose((2, 1, 0))  # N x K x 3
            num = rawjoint.shape[0]  # N

            rawjoint = sk_xyz_depth2color(rawjoint, sk_trans_vec, sk_rot)
            # reorder idx
            joint = rawjoint[:, stb_to_snap_id, :]
            # scale to meter
            joint = joint / 1000.0
            # root from palm to wrist
            # joint = _stb_palm2wrist(joint)  # N x K x 3 # yang lixin
            joint = ge_palm2wrist(joint)  # N x K x 3  #liu hao ge // vae//
            self.joints.append(joint)

            ''' 4. load images pth '''
            for idx in range(joint.shape[0]):
                self.img_paths.append(os.path.join(
                    imgpath, "{}_{}.png".format(imgsk_prefix, idx)
                ))
                self.dep_paths.append(os.path.join(
                    imgpath, "{}_{}.png".format(depsk_prefix, idx)
                ))

        self.joints = np.concatenate(self.joints, axis=0).astype(np.float32)  ##(30000, 21, 3)

        for i in tqdm(range(len(self.img_paths))):
            joint = self.joints[i]
            kp2d_homo = self.sk_intr.dot(joint.T).T
            kp2d = kp2d_homo / kp2d_homo[:, 2:3]
            kp2d = kp2d[:, :2]
            center = handutils.get_annot_center(kp2d)

            # caculate my_scale
            dep = Image.open(self.dep_paths[i]).convert("RGB")
            rel_dep = self.real_dep_img(dep)
            mask_rel_dep = np.argwhere(rel_dep > 1e-6)
            # my_scale = handutils.get_ori_crop_scale(mask_rel_dep, side=0, kp2d=kp2d) # ori
            my_scale = handutils.get_ori_crop_scale(mask_rel_dep, side=0, kp2d=kp2d,
                                                    mask_flag=False)  # get bbx only from kp2d ,比起ori差距不大,略好一点点一点点
            my_scale = (np.atleast_1d(my_scale))[np.newaxis, :]
            self.my_scales.append(my_scale)

            self.kp2ds.append(kp2d[np.newaxis, :, :])
            self.centers.append(center[np.newaxis, :])
            # self.scales.append((np.atleast_1d(scale))[np.newaxis, :])

        self.kp2ds = np.concatenate(self.kp2ds, axis=0).astype(np.float32)  # (N, 21, 2)
        self.centers = np.concatenate(self.centers, axis=0).astype(np.float32)  # (N, 2)
        # self.scales = np.concatenate(self.scales, axis=0).astype(np.float32)  # (N, 1)
        self.my_scales = np.concatenate(self.my_scales, axis=0).astype(np.float32)  # (N, 1)
        if use_cache:
            full_info = {
                "img_paths": self.img_paths,
                "dep_paths": self.dep_paths,
                "joints": self.joints,
                "kp2ds": self.kp2ds,
                "centers": self.centers,
                # "scales": self.scales,
                "my_scales": self.my_scales,
            }
            with open(cache_path, "wb") as fid:
                pickle.dump(full_info, fid)
                print("Wrote cache for dataset stb {} to {}".format(
                    self.data_split, cache_path
                ))
        return

    def __len__(self):
        """for STB dataset total (1,500 * 2) * 2 * 6 = 36,000 samples

        :return - if is train: 30,000 samples
        :return - if is eval:   6,000 samples
        """
        return len(self.img_paths)

    def __str__(self):
        info = "STB {} set. lenth {}".format(
            self.data_split, len(self.img_paths)
        )
        return colored(info, 'blue', attrs=['bold'])

    def _is_valid(self, clr, index):
        valid_data = isinstance(clr, (np.ndarray, PIL.Image.Image))

        if not valid_data:
            raise Exception("Encountered error processing stb[{}]".format(index))
        return valid_data

    def get_sample(self, index):  # replace __getitem__
        flip = True if self.hand_side != 'left' else False

        intr = self.sk_intr

        # prepare color image
        clr = Image.open(self.img_paths[index]).convert("RGB")
        self._is_valid(clr, index)

        # prepare joint
        joint = self.joints[index].copy()  # (21, 3)

        # prepare kp2d
        kp2d = self.kp2ds[index].copy()

        center = self.centers[index].copy()
        # scale = self.scales[index].copy()

        my_scale = self.my_scales[index].copy()

        if self.dep_paths[index]:
            dep = Image.open(self.dep_paths[index]).convert("RGB")
            ### dep values now are stored as |mod|div|0| (RGB)
            self._is_valid(dep, index)
            valid_dep = True
        else:
            dep = None
            valid_dep = False

        if flip:
            clr = clr.transpose(Image.FLIP_LEFT_RIGHT)
            center[0] = clr.size[0] - center[0]
            kp2d[:, 0] = clr.size[0] - kp2d[:, 0]
            joint[:, 0] = -joint[:, 0]
            if valid_dep:
                dep = dep.transpose(Image.FLIP_LEFT_RIGHT)

            # visualization
        if self.visual:
            clr_ = np.array(clr)
            dep_ = np.array(dep)

            fig = plt.figure(figsize=(20, 20))

            plt.subplot(2, 4, 1)
            clr1 = clr_.copy()
            rel_dep = self.real_dep_img(dep_)
            mask_rel_dep = np.argwhere(rel_dep > 1e-6)
            rmin, cmin = mask_rel_dep.min(0)
            rmax, cmax = mask_rel_dep.max(0)
            cv2.rectangle(clr1, (cmin, rmin), (cmax, rmax), (255), thickness=3)
            plt.imshow(clr1)
            plt.title('Color+BounduingBox')

            plt.subplot(2, 4, 2)
            dep1 = dep_.copy()
            plt.imshow(dep1)
            plt.title('Depth')

            clr_dep = clr_ + dep_
            plt.subplot(2, 4, 3)
            plt.imshow(clr_dep)
            plt.title('Color+Depth')

            rel_dep_ = rel_dep.copy()
            plt.subplot(2, 4, 4)
            plt.imshow(rel_dep_)
            plt.title('real_Depth')

            plt.subplot(2, 4, 5)
            plt.imshow(clr.copy())
            plt.title('Color')

            plt.subplot(2, 4, 6)
            plt.imshow(clr.copy())
            plt.plot(kp2d[:, :1], kp2d[:, 1:], 'ro')
            plt.title('Color+2D annotations')

            ax = fig.add_subplot(247, projection='3d')
            plt.plot(joint[:, 0], joint[:, 1], joint[:, 2], 'yo', label='keypoint')
            plt.plot(joint[:5, 0], joint[:5, 1],
                     joint[:5, 2],
                     'r',
                     label='thumb')
            plt.plot(joint[[0, 5, 6, 7, 8, ], 0], joint[[0, 5, 6, 7, 8, ], 1],
                     joint[[0, 5, 6, 7, 8, ], 2],
                     'b',
                     label='index')
            plt.plot(joint[[0, 9, 10, 11, 12, ], 0], joint[[0, 9, 10, 11, 12], 1],
                     joint[[0, 9, 10, 11, 12], 2],
                     'b',
                     label='middle')
            plt.plot(joint[[0, 13, 14, 15, 16], 0], joint[[0, 13, 14, 15, 16], 1],
                     joint[[0, 13, 14, 15, 16], 2],
                     'b',
                     label='ring')
            plt.plot(joint[[0, 17, 18, 19, 20], 0], joint[[0, 17, 18, 19, 20], 1],
                     joint[[0, 17, 18, 19, 20], 2],
                     'b',
                     label='pinky')
            # snap convention
            plt.plot(joint[4][0], joint[4][1], joint[4][2], 'rD', label='thumb')
            plt.plot(joint[8][0], joint[8][1], joint[8][2], 'ro', label='index')
            plt.plot(joint[12][0], joint[12][1], joint[12][2], 'ro', label='middle')
            plt.plot(joint[16][0], joint[16][1], joint[16][2], 'ro', label='ring')
            plt.plot(joint[20][0], joint[20][1], joint[20][2], 'ro', label='pinky')
            # plt.plot(joint [1:, 0], joint [1:, 1], joint [1:, 2], 'o')

            plt.title('3D annotations')
            ax.set_xlabel('x')
            ax.set_ylabel('y')
            ax.set_zlabel('z')
            plt.legend()
            ax.view_init(-90, -90)
            plt.show()

        sample = {
            'index': index,
            'clr': clr,
            # 'dep': dep,  # if has renturn PIL image
            'kp2d': kp2d,
            'center': center,
            # 'scale': scale,
            'my_scale': my_scale,
            'joint': joint,
            'intr': intr,
            # 'valid_dep': valid_dep,
        }

        return sample

    def norm_dep_img(self, dep_, joint_z):
        if isinstance(dep_, PIL.Image.Image):
            dep_ = np.array(dep_)
            assert (dep_.shape[-1] == 3)  # used as "RGB"

        ''' Converts a RGB-coded depth into float valued depth. '''
        ''' dep values now are stored as |mod|div|0| (RGB) '''
        dep = (dep_[:, :, 1] * 2 ** 8 + dep_[:, :, 0]).astype('float32')
        dep /= 1000.0  # depth now in meter

        lower_bound = joint_z.min() - 0.05  # meter
        upper_bound = joint_z.max() + 0.05

        np.putmask(dep, dep <= lower_bound, upper_bound)
        min_dep = dep.min() - 1e-3  # slightly compensate
        np.putmask(dep, dep >= upper_bound, 0.0)
        max_dep = dep.max() + 1e-3
        np.putmask(dep, dep <= min_dep, max_dep)
        range_dep = max_dep - min_dep
        dep = (-1 * dep + max_dep) / range_dep
        return dep

    def real_dep_img(self, dep_):
        if isinstance(dep_, PIL.Image.Image):
            dep_ = np.array(dep_)
            assert (dep_.shape[-1] == 3)  # used as "RGB"

        ''' Converts a RGB-coded depth into float valued depth. '''
        ''' dep values now are stored as |mod|div|0| (RGB) '''
        dep = (dep_[:, :, 1] * 2 ** 8 + dep_[:, :, 0]).astype('float32')
        dep /= 1000.0  # depth now in meter

        return dep

def main():
    # go through the whole dataset
    data_split = "test"
    stb = STBDataset(
        data_root="/home/chen/datasets/STB",
        data_split=data_split,
        hand_side="right",
        visual=True
    )
    print("len(stb)=", len(stb))

    for i in tqdm(range(len(stb))):
        print("i=", i)
        data = stb.get_sample(i)

if __name__ == "__main__":
    main()