import os
import time
import cv2
import gymnasium as gym
import numpy as np
from tqdm import tqdm
from itertools import cycle, islice
import multiprocessing
import json

import egogym
from egogym.wrappers import EpisodeMonitor, get_vlm_wrapper
import egogym.assets.constants as constants

def evaluate_policy(
        task_name="pick",
        policy=None,
        robot="rum",
        action_space="delta",
        num_objs=1,
        num_episodes=100,
        max_steps=50,
        render=False,
        record=False,
        render_freq=0,
        render_size=(299, 224),
        num_envs=1,
        seed=42,
        log_file=None,
        objects_set=None,
        reward_threshold=None,
        use_unprivileged_vlm=None,
        logs_dir="logs"
    ):

    num_envs = max(1, min(num_envs, num_episodes))
    if num_envs > 1 and render_freq > 0:
        print(f"Warning: render_freq={render_freq} is not supported with multiple environments. Setting render_freq=0.")
        render_freq = 0
    task_name = task_name.lower()
    if objects_set is None:
        if task_name == "pick":
            objects_set = constants.lite_pick_objects_set
        elif task_name == "open":
            objects_set = constants.all_open_objects_set
    else:
        objects_set_options = {
            'all_pick': constants.all_pick_objects_set,
            'lite_pick': constants.lite_pick_objects_set,
            'diverse_pick': constants.diverse_pick_objects_set,
            'full_eval': constants.full_eval_objects_set,
            'all_open': constants.all_open_objects_set,
            "all_close": constants.all_close_objects_set,
            'cabinet': constants.cabinet_objects_set,
            'drawer': constants.drawer_objects_set,
        }
        if objects_set in objects_set_options:
            objects_set = objects_set_options[objects_set]
        else:
            raise ValueError(f"Unknown objects_set: {objects_set}")
        
    if reward_threshold is None:
        if task_name == "pick":
            reward_threshold = 0.05
        elif task_name == "open":
            reward_threshold = 0.5

    if not log_file:
        log_file = f"evaluation_{int(time.time())}"

    logs_dir = f"{logs_dir}/{log_file}"
    os.makedirs(logs_dir, exist_ok=True)
    
    eval_args = {
        "task_name": task_name,
        "robot": robot,
        "action_space": action_space,
        "num_objs": num_objs,
        "num_episodes": num_episodes,
        "max_steps": max_steps,
        "render": render,
        "record": record,
        "render_freq": render_freq,
        "render_size": render_size,
        "num_envs": num_envs,
        "seed": seed,
        "log_file": log_file,
        "objects_set": objects_set,
        "reward_threshold": reward_threshold,
    }
    
    with open(os.path.join(logs_dir, "eval_args.json"), "w") as f:
        json.dump(eval_args, f, indent=4)

    if num_envs > 1:
        multiprocessing.set_start_method("spawn", force=True)
        
        expanded_objects = list(islice(cycle(objects_set), num_episodes))
        env_objects = [expanded_objects[i::num_envs] for i in range(num_envs)]
        
        def make_env(env_id):
            def _init():
                env = gym.make(
                    f"Egogym-{task_name.capitalize()}-v0",
                    robot=robot,
                    action_space=action_space,
                    render_mode='human' if render else None,
                    render_size=render_size,
                    num_objs=num_objs,
                    seed=seed + env_id,
                    objects_set=env_objects[env_id]
                )
                if use_unprivileged_vlm is not None:
                    wrapper_cls = get_vlm_wrapper(use_unprivileged_vlm)
                    env = wrapper_cls(env)
                return env
            return _init
        
        env = gym.vector.AsyncVectorEnv(
            [make_env(i) for i in range(num_envs)],
            autoreset_mode=gym.vector.AutoresetMode.DISABLED
        )
        env = EpisodeMonitor(env, logs_dir=logs_dir, record=record, render_freq=render_freq, num_envs=num_envs)
        obs, _ = env.reset()
        env_episode_count = np.zeros(num_envs, dtype=int)
        env_step_count = np.zeros(num_envs, dtype=int)
        env_success_count = 0
        episodes_per_env = np.array([len(env_objects[i]) for i in range(num_envs)])
        
        pbar = tqdm(total=num_episodes, desc="Evaluating | Success Rate: 0.0%")
        
        while np.any(env_episode_count < episodes_per_env):
            actions = policy.get_action(obs)
            inactive = env_episode_count >= episodes_per_env
            actions[inactive] = None

            obs, rewards, terminated, truncated, info = env.step(actions)
            env_step_count += 1
            
            if render:
                ego_exo = np.concatenate([obs["rgb_ego"], obs["rgb_exo"]], axis=2)
                combined_view = np.concatenate(ego_exo, axis=0)
                cv2.imshow("Ego View", cv2.cvtColor(combined_view, cv2.COLOR_RGB2BGR))
                cv2.waitKey(1)
            
            dones = (rewards > reward_threshold) | terminated | (env_step_count >= max_steps)
            done_indices = np.where(dones)[0]

            active_done = done_indices[env_episode_count[done_indices] < episodes_per_env[done_indices]]
            
            if len(active_done) > 0:
                env.log_episodes(active_done)
                
                successes = rewards[active_done] > reward_threshold
                env_success_count += np.sum(successes)
                
                env_episode_count[active_done] += 1
                pbar.update(len(active_done))
                
                total_completed = np.sum(env_episode_count)
                success_rate = (env_success_count / total_completed) * 100 if total_completed > 0 else 0
                pbar.set_description(f"Evaluating | Success Rate: {success_rate:.1f}%")
                
                policy.reset(active_done.tolist())
                
                reset_mask = np.zeros(num_envs, dtype=bool)
                reset_mask[active_done] = True
                obs, _ = env.reset(options={"reset_mask": reset_mask})
                env_step_count[active_done] = 0
            
        pbar.close()
        success_rate = (env_success_count / num_episodes) * 100
        return success_rate
        
    else:

        env = gym.make(
            f"Egogym-{task_name.capitalize()}-v0",
            robot=robot,
            action_space=action_space,
            render_mode='human' if render else None,
            render_size=render_size,
            num_objs=num_objs,
            seed=seed,
            objects_set=objects_set
        )
        if use_unprivileged_vlm is not None:
            wrapper_cls = get_vlm_wrapper(use_unprivileged_vlm)
            env = wrapper_cls(env)
        env = EpisodeMonitor(env, logs_dir=logs_dir, record=record, render_freq=render_freq, num_envs=num_envs)
        
        success_count = 0
        pbar = tqdm(range(num_episodes), desc="Evaluating | Success Rate: 0.0%")
        for episode_idx in pbar:
            obs, _ = env.reset()
            policy.reset()
            
            success = False
            for _ in range(max_steps):
                action = policy.get_action(obs)
                obs, reward, terminated, truncated, _ = env.step(action[0])
                    
                if reward > reward_threshold:
                    success = True
                    break
                if terminated or truncated:
                    break
            
            if success:
                success_count += 1
            
            success_rate = (success_count / (episode_idx + 1)) * 100
            pbar.set_description(f"Evaluating | Success Rate: {success_rate:.1f}%")

    env.close()
    return (success_count / num_episodes) * 100 if num_envs == 1 else success_rate