import torch
import os
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm
from einops import rearrange
from roboimi.utils.utils import set_seed
from roboimi.utils.io_utils import IOUtils
from roboimi.utils.model_interface import ModelInterface
from roboimi.envs.vx300s_jnt import make_sim_env
import time

# from visualize_episodes import save_videos
from roboimi.utils.utils import sample_box_pose, sample_insertion_pose



#should be added into IOUtils
def get_image(obs,camera_names):
    curr_images = []
    for cam_name in camera_names:
        curr_image = rearrange(obs['images'][cam_name], 'h w c -> c h w')
        curr_images.append(curr_image)
    curr_image = np.stack(curr_images, axis=0)
    curr_image = torch.from_numpy(curr_image / 255.0).float().cuda().unsqueeze(0)
    return curr_image


def eval_bc(config, ckpt_name='policy_best.ckpt', save_episode=True):
    set_seed(1)
    model_interface = ModelInterface(config)
    task_name = 'sim_insertion' #config['task_name']
    model_interface.setup()
    policy = IOUtils.load_policy(config, ckpt_name)
    stats = IOUtils.load_stats(config['ckpt_dir'])
    num_rollouts = 3
    episode_returns = []
    highest_rewards = []
    for rollout_id in range(num_rollouts):
        episode_return, episode_highest_reward = run_episode(config, policy, stats,
                                                              save_episode,rollout_id)
    

    

def run_episode(config, policy, stats, save_episode,rollout_id):
    print("\nrollout_id===",rollout_id,"\n")
    pre_process = lambda s_qpos: (s_qpos - stats['qpos_mean']) / stats['qpos_std']
    post_process = lambda a: a * stats['action_std'] + stats['action_mean']
    if 'sim_insertion' in config['task_name']:
        peg_pose, socket_pose = sample_insertion_pose()
        box_pose = np.hstack((peg_pose[:3],socket_pose[:3])) # used in sim reset
    task_name =  'sim_insertion'  #config['task_name']
    env = make_sim_env(task_name)
    env.reset(box_pose)
    max_timesteps = config['episode_len']
    max_timesteps = int(max_timesteps * 1)
 
    image_list = []
    rewards = []
    query_frequency = config['policy_config'].get('num_queries', 1)

    with torch.inference_mode():
        for t in range(700):
            # print("obs_img",env.obs['images'])
            image_list.append(env.obs['images'] if 'images' in env.obs else {print("img error")})
            qpos_numpy = np.array(env.obs['qpos'])
            qpos = pre_process(qpos_numpy)
            qpos = torch.from_numpy(qpos).float().cuda().unsqueeze(0)
            curr_image = get_image(env.obs, config['camera_names'])
            if config['policy_class'] == "ACT" or "ACTTV":
                if t % query_frequency == 0:
                    all_actions = policy(qpos, curr_image)
            elif config['policy_class'] == "CNNMLP":
                raw_action = policy(qpos, curr_image)
            else:
                raise NotImplementedError
            raw_action = all_actions[:, t % query_frequency]
            raw_action = raw_action.squeeze(0).cpu().numpy()
            action = post_process(raw_action)
            
            env.step(action)
            rewards.append(env.rew)
            env.render()


    rewards = np.array(rewards)
    episode_return = np.sum(rewards[rewards != None])
    episode_highest_reward = np.max(rewards)
    env.viewer.close()

    del env
    return episode_return, episode_highest_reward


def test_env():
    try:
        env = make_sim_env('sim_insertion')
        box_pos = np.concatenate(sample_insertion_pose())
        env.reset(box_pos)
        while True: pass
    except KeyboardInterrupt:
        del env
        print("stop")
    

if __name__ == '__main__':
    test_env()
    # io_utils = IOUtils()
    # config = io_utils.load_config()
    # eval_bc(config)




# config===== {'onscreen_render': False,
#  'eval': 1, 
# 'ckpt_dir': 'ckpt_models', 
# 'num_epochs': 3000, 
# 'temporal_agg': False, 
# 'policy_class': 'ACT', 
# 'backbone': 'resnet18', 
# 'seed': 0, 'real_robot': 0,
#  'task_name': 'sim_insertion', 
# 'images_render_height': 480, 
# 'images_render_width': 640, 
# 'left_arm_DOF_number': 6, 
# 'right_arm_DOF_number': 6, 
# 'left_qpos_raw': 8, 
# 'right_qpos_raw': 8, 
# 'left_qvel_raw': 8, 
# 'right_qvel_raw': 8, 
# 'dataset_dir': '/home/arm/lzd/act_env/dataset/sim_insertion', 
# 'num_episodes': 7, 
# 'episode_len': 400, 
# 'camera_names': ['top'], 
# 'xml_dir': None, 
# 'batch_size': 8, 
# 'state_dim': 14, 
# 'action_dim': 14, 
# 'lr_backbone': 1e-05, 
# 'enc_layers': 4, 
# 'dec_layers': 7, 
# 'nheads': 8, 
# 'qpos_noise_std': 0, 
# 'DT': 0.02, 
# 'lr': 1e-05, 
# 'kl_weight': 10, 
# 'chunk_size': 100, 
# 'hidden_dim': 512, 
# 'dim_feedforward': 3200, 
# 'policy_config': {'lr': 1e-05, 'num_queries': 100, 'kl_weight': 10, 'hidden_dim': 512, 'dim_feedforward': 3200, 'lr_backbone': 1e-05, 'backbone': 'resnet18', 'enc_layers': 4, 'dec_layers': 7, 'nheads': 8, 'camera_names': ['top']}} 