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.double_pos_ctrl_env import make_sim_env
# from visualize_episodes import save_videos
from roboimi.utils.act_ex_utils import sample_transfer_pose


class ActionSmoother:
    """
    动作平滑器，支持多种平滑策略
    """
    def __init__(self, action_dim, method='ema', alpha=0.3, window_size=5):
        """
        Args:
            action_dim: 动作维度
            method: 平滑方法 ('ema', 'moving_avg', 'lowpass', 'none')
            alpha: EMA 平滑系数 (0-1)，越小越平滑
            window_size: 滑动窗口大小
        """
        self.action_dim = action_dim
        self.method = method
        self.alpha = alpha
        self.window_size = window_size
        self.history = []
        self.prev_action = None

    def smooth(self, action):
        """
        对动作进行平滑处理

        Args:
            action: 当前动作 [action_dim]

        Returns:
            smoothed_action: 平滑后的动作
        """
        if self.method == 'none':
            return action

        if self.method == 'ema':
            # 指数移动平均
            if self.prev_action is None:
                smoothed = action
            else:
                smoothed = self.alpha * action + (1 - self.alpha) * self.prev_action
            self.prev_action = smoothed
            return smoothed

        elif self.method == 'moving_avg':
            # 滑动平均
            self.history.append(action.copy())
            if len(self.history) > self.window_size:
                self.history.pop(0)
            return np.mean(self.history, axis=0)

        elif self.method == 'lowpass':
            # 一阶低通滤波器
            if self.prev_action is None:
                smoothed = action
            else:
                smoothed = self.prev_action + self.alpha * (action - self.prev_action)
            self.prev_action = smoothed
            return smoothed

        else:
            raise ValueError(f"Unknown smoothing method: {self.method}")

    def reset(self):
        """重置平滑器状态"""
        self.history = []
        self.prev_action = None


#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)
    model_interface.setup()
    policy = IOUtils.load_policy(config, ckpt_name)
    stats = IOUtils.load_stats(config['ckpt_dir'])
    num_rollouts = 3
    episode_returns = []
    highest_rewards = []
    
    
    
   
    
    run_episode(config, policy, stats,
                save_episode,num_rollouts)
    # episode_return, episode_highest_reward = run_episode(config, policy, stats,
    #                                                           save_episode,num_rollouts)
    

    

def run_episode(config, policy, stats, save_episode,num_rollouts):

    if 'sim_transfer' in config['task_name']:
        task_name =  'sim_transfer'  #config['task_name']
        env = make_sim_env(task_name)
    
    max_timesteps = config['episode_len']
    max_timesteps = int(max_timesteps * 1)
    pre_process = lambda s_qpos: (s_qpos - stats['qpos_mean']) / stats['qpos_std']
    post_process = lambda a: a * stats['action_std'] + stats['action_mean']
    box_pos = sample_transfer_pose()

    # 初始化动作平滑器
    action_dim = config['action_dim']
    use_smoothing = config.get('use_action_smoothing', False)
    smooth_method = config.get('smooth_method', 'ema')
    smooth_alpha = config.get('smooth_alpha', 0.3)

    if use_smoothing and config['policy_class'] == "GR00T":
        smoother = ActionSmoother(action_dim, method=smooth_method, alpha=smooth_alpha)
        print(f"Action smoothing enabled: method={smooth_method}, alpha={smooth_alpha}")
    else:
        smoother = None

    for rollout_id in range(num_rollouts):
        print("\nrollout_id===",rollout_id,"\n")
        image_list = []
        rewards = []
        query_frequency = config['policy_config'].get('num_queries', 1)
        print("query_freq =====",query_frequency)
        env.reset(box_pos)

        # 重置平滑器
        if smoother is not None:
            smoother.reset()

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


                action = post_process(raw_action)

                # 应用动作平滑（仅对 GR00T）
                if smoother is not None:
                    action = smoother.smooth(action)

                print("action == ",action)
                env.step_jnt(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_transfer')
        env.reset()
        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)


