import os
import numpy as np
import torch
import skvideo
import einops
import imageio
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
import gym
import mujoco_py as mjc
import warnings
import pdb

# -----------------------------------------------------------------------------#
# ------------------------------- helper structs ------------------------------#
# -----------------------------------------------------------------------------#

def to_np(x):
    if torch.is_tensor(x):
        x = x.detach().cpu().numpy()
    return x


def _make_dir(filename):
    folder = os.path.dirname(filename)
    if not os.path.exists(folder):
        os.makedirs(folder)


def save_video(filename, video_frames, fps=60, video_format='mp4'):
    assert fps == int(fps), fps
    _make_dir(filename)

    skvideo.io.vwrite(
        filename,
        video_frames,
        inputdict={
            '-r': str(int(fps)),
        },
        outputdict={
            '-f': video_format,
            '-pix_fmt': 'yuv420p',
            # '-pix_fmt=yuv420p' needed for osx https://github.com/scikit-video/scikit-video/issues/74
        }
    )


def save_videos(filename, *video_frames, axis=1, **kwargs):
    ## video_frame : [ N x H x W x C ]
    video_frames = np.concatenate(video_frames, axis=axis)
    save_video(filename, video_frames, **kwargs)


def env_map(env_name):
    '''
        map D4RL dataset names to custom fully-observed
        variants for rendering
    '''
    if 'halfcheetah' in env_name:
        return 'HalfCheetahFullObs-v2'
    elif 'hopper' in env_name:
        return 'HopperFullObs-v2'
    elif 'walker2d' in env_name:
        return 'Walker2dFullObs-v2'
    else:
        return env_name


# -----------------------------------------------------------------------------#
# ------------------------------ helper functions -----------------------------#
# -----------------------------------------------------------------------------#

def get_image_mask(img):
    background = (img == 255).all(axis=-1, keepdims=True)
    mask = ~background.repeat(3, axis=-1)
    return mask


def atmost_2d(x):
    while x.ndim > 2:
        x = x.squeeze(0)
    return x



# -----------------------------------------------------------------------------#
# ---------------------------------- renderers --------------------------------#
# -----------------------------------------------------------------------------#

class MuJoCoRenderer:
    '''
        default mujoco renderer
    '''

    def __init__(self, env):
        if type(env) is str:
            self.env_name = env
            env = env_map(env)
            self.env = gym.make(env)
        else:
            self.env = env
        ## Core "- 1" because the envs in renderer are fully-observed
        self.observation_dim = np.prod(self.env.observation_space.shape) - 1
        self.action_dim = np.prod(self.env.action_space.shape)
        try:
            self.viewer = mjc.MjRenderContextOffscreen(self.env.sim)
        except:
            print('[ utils/rendering ] Warning: could not initialize offscreen renderer')
            self.viewer = None

    def get_mass_center(self, image):
        if "hopper" in self.env_name:
            seg_total = 4
        else:
            seg_total = 3

        mass_center = []
        for i in range(1, seg_total+1):
            idx = np.where(image == i)
            y_center = np.sum(idx[0])/idx[0].shape[0]
            x_center = np.sum(idx[1])/idx[1].shape[0]
            mass_center.append(np.array([y_center, x_center]))
        mass_center = np.stack(mass_center, axis=0)
        return mass_center


    def pad_observation(self, observation):
        state = np.concatenate([
            np.zeros(1),
            observation,
        ])
        return state

    def pad_observations(self, observations):
        qpos_dim = self.env.sim.data.qpos.size
        ## xpos is hidden
        xvel_dim = qpos_dim - 1
        xvel = observations[:, xvel_dim]
        xpos = np.cumsum(xvel) * self.env.dt
        states = np.concatenate([
            xpos[:, None],
            observations,
        ], axis=-1)
        return states

    def render(self, observation, dim=256, partial=False, qvel=True, render_kwargs=None, conditions=None):

        if type(dim) == int:
            dim = (dim, dim)

        if self.viewer is None:
            return np.zeros((*dim, 3), np.uint8)

        if render_kwargs is None:
            xpos = observation[0] if not partial else 0
            render_kwargs = {
                'trackbodyid': 2,
                'distance': 3,
                'lookat': [xpos, -0.5, 1],
                'elevation': -20
            }

        for key, val in render_kwargs.items():
            if key == 'lookat':
                self.viewer.cam.lookat[:] = val[:]
            else:
                setattr(self.viewer.cam, key, val)

        if partial:
            state = self.pad_observation(observation)
        else:
            state = observation

        qpos_dim = self.env.sim.data.qpos.size
        if not qvel or state.shape[-1] == qpos_dim:
            qvel_dim = self.env.sim.data.qvel.size
            state = np.concatenate([state, np.zeros(qvel_dim)])

        set_state(self.env, state)
        saving_mass_center = True
        if saving_mass_center:
            self.viewer.render(*dim, camera_id=-1, segmentation=True)
            seg_image = self.viewer.read_pixels(*dim, depth=False)
            mass_center = self.get_mass_center(seg_image[::-1, :, 0])

        self.viewer.render(*dim, camera_id=-1, segmentation=False)
        data = self.viewer.read_pixels(*dim, depth=False)
        data = data[::-1, :, :]
        # data[round(mass_center[0, 0]), round(mass_center[0, 1]), 0] = 255
        # data[round(mass_center[0, 0]), round(mass_center[0, 1]), 1] = 0
        # data[round(mass_center[0, 0]), round(mass_center[0, 1]), 2] = 0
        # data[round(mass_center[1, 0]), round(mass_center[1, 1]), 0] = 0
        # data[round(mass_center[1, 0]), round(mass_center[1, 1]), 1] = 255
        # data[round(mass_center[1, 0]), round(mass_center[1, 1]), 2] = 0
        # data[round(mass_center[2, 0]), round(mass_center[2, 1]), 0] = 0
        # data[round(mass_center[2, 0]), round(mass_center[2, 1]), 1] = 0
        # data[round(mass_center[2, 0]), round(mass_center[2, 1]), 2] = 255
        # data[round(mass_center[3, 0]), round(mass_center[3, 1]), 0] = 0
        # data[round(mass_center[3, 0]), round(mass_center[3, 1]), 1] = 255
        # data[round(mass_center[3, 0]), round(mass_center[3, 1]), 2] = 255
        if saving_mass_center:
            return data, mass_center
        else:
            return data

    def _renders(self, observations, **kwargs):
        images = []
        centers = []
        for observation in observations:
            img, mass_center = self.render(observation, **kwargs)
            images.append(img)
            centers.append(mass_center)
        return np.stack(images, axis=0)

    def renders(self, samples, partial=False, **kwargs):
        if partial:
            samples = self.pad_observations(samples)
            partial = False

        sample_images = self._renders(samples, partial=partial, **kwargs)

        composite = np.ones_like(sample_images[0]) * 255

        step = 0
        if "hopper" in self.env_name or "walker2d" in self.env_name:
            interval = 4
        else:
            interval = 1
        for img in sample_images:
            if step % interval == 0:
                mask = get_image_mask(img)
                composite[mask] = img[mask]
            else:
                pass
            step += 1

        return composite

    def composite(self, savepath, paths, dim=(1024, 256), **kwargs):

        if "hopper" in self.env_name:
            dim=(2048, 256)
            render_kwargs = {
                'trackbodyid': 0,
                'distance': 3,
                'lookat': [10, -0.5, 1],
                'elevation': -10
            }
        elif "walker2d" in self.env_name:
            dim = (2200, 200)
            render_kwargs = {
                'trackbodyid': 2,
                'distance': 3.5,
                'lookat': [17, -0.5, 1],
                'elevation': -10
            }
        elif "halfcheetah" in self.env_name:
            dim = (2500, 150)
            render_kwargs = {
                'trackbodyid': 2,
                'distance': 2,
                'lookat': [17, -0.5, 1],
                'elevation': -10
            }
        else:
            render_kwargs = {
                'trackbodyid': 2,
                'distance': 10,
                'lookat': [5, 2, 0.5],
                'elevation': 0
            }

        images = []
        for path in paths:
            ## [ H x obs_dim ]
            path = atmost_2d(path)
            img = self.renders(to_np(path), dim=dim, partial=True, qvel=True, render_kwargs=render_kwargs, **kwargs)
            images.append(img)
        images = np.concatenate(images, axis=0)

        if savepath is not None:
            imageio.imsave(savepath, images)
            print(f'Saved {len(paths)} samples to: {savepath}')

        return images

    def render_rollout(self, savepath, states, **video_kwargs):
        if type(states) is list: states = np.array(states)
        images = self._renders(states, partial=True)
        save_video(savepath, images, **video_kwargs)

    def render_plan(self, savepath, actions, observations_pred, state, fps=30):
        ## [ batch_size x horizon x observation_dim ]
        observations_real = rollouts_from_state(self.env, state, actions)

        ## there will be one more state in `observations_real`
        ## than in `observations_pred` because the last action
        ## does not have an associated next_state in the sampled trajectory
        observations_real = observations_real[:, :-1]

        images_pred = np.stack([
            self._renders(obs_pred, partial=True)
            for obs_pred in observations_pred
        ])

        images_real = np.stack([
            self._renders(obs_real, partial=False)
            for obs_real in observations_real
        ])

        ## [ batch_size x horizon x H x W x C ]
        images = np.concatenate([images_pred, images_real], axis=-2)
        save_videos(savepath, *images)

    def render_diffusion(self, savepath, diffusion_path, **video_kwargs):
        '''
            diffusion_path : [ n_diffusion_steps x batch_size x 1 x horizon x joined_dim ]
        '''
        render_kwargs = {
            'trackbodyid': 2,
            'distance': 10,
            'lookat': [10, 2, 0.5],
            'elevation': 0,
        }

        diffusion_path = to_np(diffusion_path)

        n_diffusion_steps, batch_size, _, horizon, joined_dim = diffusion_path.shape

        frames = []
        for t in reversed(range(n_diffusion_steps)):
            print(f'[ utils/renderer ] Diffusion: {t} / {n_diffusion_steps}')

            ## [ batch_size x horizon x observation_dim ]
            states_l = diffusion_path[t].reshape(batch_size, horizon, joined_dim)[:, :, :self.observation_dim]

            frame = []
            for states in states_l:
                img = self.composite(None, states, dim=(1024, 256), partial=True, qvel=True,
                                     render_kwargs=render_kwargs)
                frame.append(img)
            frame = np.concatenate(frame, axis=0)

            frames.append(frame)

        save_video(savepath, frames, **video_kwargs)

    def __call__(self, *args, **kwargs):
        return self.renders(*args, **kwargs)


# -----------------------------------------------------------------------------#
# ---------------------------------- rollouts ---------------------------------#
# -----------------------------------------------------------------------------#

def set_state(env, state):
    qpos_dim = env.sim.data.qpos.size
    qvel_dim = env.sim.data.qvel.size
    if not state.size == qpos_dim + qvel_dim:
        warnings.warn(
            f'[ utils/rendering ] Expected state of size {qpos_dim + qvel_dim}, '
            f'but got state of size {state.size}')
        state = state[:qpos_dim + qvel_dim]

    env.set_state(state[:qpos_dim], state[qpos_dim:])


def rollouts_from_state(env, state, actions_l):
    rollouts = np.stack([
        rollout_from_state(env, state, actions)
        for actions in actions_l
    ])
    return rollouts


def rollout_from_state(env, state, actions):
    qpos_dim = env.sim.data.qpos.size
    env.set_state(state[:qpos_dim], state[qpos_dim:])
    observations = [env._get_obs()]
    for act in actions:
        obs, rew, term, _ = env.step(act)
        observations.append(obs)
        if term:
            break
    for i in range(len(observations), len(actions) + 1):
        ## if terminated early, pad with zeros
        observations.append(np.zeros(obs.size))
    return np.stack(observations)
