import numpy as np
import csv
import os
import cv2

class EpisodeMonitor:
    
    def __init__(self, env, logs_dir=None, columns=None, record=False, render_freq=0, num_envs=1):
        self.env = env
        self.logs_dir = logs_dir
        self.record = record
        self.mujoco_step_counter = 0
        self.columns = columns or [
            "episode", "max_reward", "object_name", "texture_name", 
            "grasped_bodies", "steps", "initial_robot_pose", 
            "initial_object_pose", "is_grasping",
            "gripper_current_position", "grasping_object"
        ]
        
        self.is_vector = num_envs > 1
        self.render_freq = render_freq
        
        if self.is_vector:
            self.num_envs = num_envs
            self.global_episode_counter = 0
            self.episode_idx = np.zeros(self.num_envs, dtype=int)
            self.max_reward = np.zeros(self.num_envs)
            self.steps = np.zeros(self.num_envs, dtype=int)
            self.episode_info = [{} for _ in range(self.num_envs)]
            if self.record:
                self.video_writers = [None] * self.num_envs
                self.video_frames = [[] for _ in range(self.num_envs)]
        else:
            self.episode_idx = 0
            self.max_reward = 0
            self.steps = 0
            self.episode_info = {}
            if self.record:
                self.video_writer = None
                self.video_frames = []
        
        self.csv_writer = None
        self.csv_file = None

        os.makedirs(self.logs_dir, exist_ok=True)
        csv_path = os.path.join(self.logs_dir, "log.csv")
        self.csv_file = open(csv_path, 'w', newline='')
        self.csv_writer = csv.DictWriter(self.csv_file, fieldnames=self.columns, delimiter='\t')
        self.csv_writer.writeheader()
        
        if self.record:
            self.video_dir = os.path.join(self.logs_dir, "videos")
            os.makedirs(self.video_dir, exist_ok=True)
        
        self._callbacks_initialized = False
        self._callbacks_paused = False
        self._callbacks_paused = False
    
    def _setup_callbacks(self):
        if self._callbacks_initialized or self.render_freq <= 0:
            return
        
        if self.is_vector:
            # AsyncVectorEnv doesn't provide direct access to individual envs
            # Callbacks cannot be set up for vectorized environments
            pass
        else:
            def callback():
                self._capture_frame_single()
            actual_env = self.env.unwrapped
            actual_env._render_callback = callback
            actual_env._render_freq = self.render_freq
            actual_env._should_render = actual_env.render_mode == 'human'
        
        self._callbacks_initialized = True
    
    def _capture_frame_single(self):
        if self.render_freq <= 0 or self._callbacks_paused:
            return
        
        try:
            actual_env = self.env.unwrapped
            
            if actual_env._should_render:
                actual_env.render()
            
            if self.record:
                frame_ego = actual_env.robot.get_camera_view(actual_env.robot.camera_names[0])
                frame_exo = actual_env.robot.get_camera_view(actual_env.robot.camera_names[1])
                frame = np.concatenate((frame_ego, frame_exo), axis=1)
                frame = self._process_frame(frame)
                self.video_frames.append(frame)
        except Exception as e:
            pass
    
    def _capture_frame_vectorized(self, env_idx):
        if self.render_freq <= 0 or self._callbacks_paused:
            return
        
        try:
            env = self.env.envs[env_idx]
            actual_env = env.unwrapped
            
            if actual_env._should_render:
                actual_env.render()
            
            if self.record:
                frame_ego = actual_env.robot.get_camera_view(actual_env.robot.camera_names[0])
                frame_exo = actual_env.robot.get_camera_view(actual_env.robot.camera_names[1])
                frame = np.concatenate((frame_ego, frame_exo), axis=1)
                frame = self._process_frame(frame)
                self.video_frames[env_idx].append(frame)
        except Exception as e:
            pass
    
    def _process_frame(self, frame):
        if isinstance(frame, np.ndarray):
            if frame.dtype != np.uint8:
                frame = (frame * 255).astype(np.uint8) if frame.max() <= 1.0 else frame.astype(np.uint8)
            if frame.shape[-1] == 3:
                frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
        return frame
    
    def reset(self, **kwargs):
        self._callbacks_paused = True
        
        if self.is_vector:
            options = kwargs.get('options', {})
            reset_mask = options.get('reset_mask', np.ones(self.num_envs, dtype=bool))
            
            obs, info = self.env.reset(**kwargs)
            self._setup_callbacks()
            
            for i in np.where(reset_mask)[0]:
                if self.steps[i] == 0 and self.record:
                    self.video_frames[i] = []
                self.max_reward[i] = 0
                self.steps[i] = 0
                self.episode_info[i] = {}
            
            self._callbacks_paused = False
        else:
            if self.steps > 0:
                if self.csv_writer:
                    self._log_episode()
                if self.record:
                    self._save_video()
                self.episode_idx += 1
            obs, info = self.env.reset(**kwargs)
            self._setup_callbacks()
            
            self.max_reward = 0
            self.steps = 0
            self.episode_info = {}
            if self.record:
                self.video_frames = []
            self._callbacks_paused = False
        return obs, info
    
    def step(self, action):
        obs, reward, terminated, truncated, info = self.env.step(action)
        
        if self.render_freq == 0:
            if self.is_vector:
                # AsyncVectorEnv doesn't provide direct access to individual envs
                # Recording from observations if available
                if self.record and 'rgb_ego' in obs and 'rgb_exo' in obs:
                    for i in range(self.num_envs):
                        frame_ego = obs['rgb_ego'][i]
                        frame_exo = obs['rgb_exo'][i]
                        frame = np.concatenate((frame_ego, frame_exo), axis=1)
                        frame = self._process_frame(frame)
                        self.video_frames[i].append(frame)
            else:
                actual_env = self.env.unwrapped
                if actual_env._should_render:
                    actual_env.render()
                if self.record and 'rgb_ego' in obs and 'rgb_exo' in obs:
                    frame_ego = obs['rgb_ego']
                    frame_exo = obs['rgb_exo']
                    frame = np.concatenate((frame_ego, frame_exo), axis=1)
                    frame = self._process_frame(frame)
                    self.video_frames.append(frame)
        
        if self.is_vector:
            self.max_reward = np.maximum(self.max_reward, reward)
            self.steps += 1
            for i in range(self.num_envs):
                if isinstance(info, dict):
                    env_info = {}
                    for k, v in info.items():
                        if k.startswith('_'):
                            continue
                        if isinstance(v, (list, np.ndarray)) and len(v) > i:
                            env_info[k] = v[i]
                        else:
                            env_info[k] = v
                else:
                    env_info = info[i] if isinstance(info, (list, tuple)) else {}
                self.episode_info[i].update(env_info)
            
            dones = np.logical_or(terminated, truncated)
            if isinstance(dones, bool):
                dones = np.array([dones] * self.num_envs)
        else:
            self.max_reward = max(self.max_reward, reward)
            self.steps += 1
            self.episode_info.update(info)
            info['episode_stats'] = {'episode': self.episode_idx, 'max_reward': self.max_reward, 'steps': self.steps}
        
        return obs, reward, terminated, truncated, info
    
    def _log_episode(self, env_idx=None):
        if env_idx is not None:
            log_data = {"episode": self.global_episode_counter, "max_reward": float(self.max_reward[env_idx]), "steps": int(self.steps[env_idx]), **self.episode_info[env_idx]}
            self.global_episode_counter += 1
        else:
            log_data = {"episode": self.episode_idx, "max_reward": self.max_reward, "steps": self.steps, **self.episode_info}
        
        filtered_data = {}
        for col in self.columns:
            value = log_data.get(col, "")
            if isinstance(value, np.ndarray):
                value = value.flatten().tolist()
            filtered_data[col] = value
        
        self.csv_writer.writerow(filtered_data)
        self.csv_file.flush()
    
    def log_episodes(self, env_indices):
        if self.is_vector:
            for i in env_indices:
                if self.csv_writer:
                    self._log_episode(i)
                if self.record:
                    self._save_video(i)
                    self.video_frames[i] = []
                self.episode_idx[i] += 1
    
    def _save_video(self, env_idx=None):
        if env_idx is not None:
            frames = self.video_frames[env_idx]
            episode_num = int(self.episode_idx[env_idx])
            video_filename = f"env_{env_idx}_episode_{episode_num}.mp4"
        else:
            frames = self.video_frames
            episode_num = self.episode_idx
            video_filename = f"episode_{episode_num}.mp4"
        
        if not frames:
            return
        
        video_path = os.path.join(self.video_dir, video_filename)
        height, width = frames[0].shape[:2]
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(video_path, fourcc, 30, (width, height))
        
        for frame in frames:
            out.write(frame)
        
        out.release()
    
    def close(self):
        if self.is_vector:
            for i in range(self.num_envs):
                if self.steps[i] > 0:
                    if self.csv_writer:
                        self._log_episode(i)
                    if self.record:
                        self._save_video(i)
        else:
            if self.steps > 0:
                if self.csv_writer:
                    self._log_episode()
                if self.record:
                    self._save_video()
        if self.csv_file:
            self.csv_file.close()
        self.env.close()
    
    def __getattr__(self, name):
        return getattr(self.env, name)
