from __future__ import annotations

import argparse
import time
from collections import deque

import numpy as np
import rootutils
from loguru import logger as log
from rich.logging import RichHandler
from stable_baselines3 import PPO
from stable_baselines3.common.callbacks import BaseCallback
from stable_baselines3.common.env_util import make_vec_env

import wandb
from wandb.integration.sb3 import WandbCallback

rootutils.setup_root(__file__, pythonpath=True)
log.configure(handlers=[{"sink": RichHandler(), "format": "{message}"}])

from metasim.cfg.scenario import ScenarioMetaCfg
from metasim.constants import SimType
from metasim.scripts.RL.wrapper_sb3 import MetaSimGymWrapper
from metasim.utils.setup_util import get_robot, get_task


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--num_envs", type=int, default=1)
    parser.add_argument("--sim", type=str, default="isaaclab", choices=["isaaclab", "isaacgym", "mujoco"])
    parser.add_argument(
        "--robot",
        type=str,
        default="h1",
        choices=[
            "franka",
            "ur5e_2f85",
            "sawyer",
            "franka_with_gripper_extension",
            "h1_2_without_hand",
            "h1",
            "h1_simple_hand",
        ],
    )
    parser.add_argument("--task", type=str, default="Stand")
    parser.add_argument("--add_table", action="store_true")
    parser.add_argument("--total_timesteps", type=int, default=20000000)
    parser.add_argument("--train_or_eval", type=str, default="train", choices=["train", "eval", "test"])
    parser.add_argument("--model_path", type=str, default=None)
    # PPO specific arguments
    parser.add_argument("--learning_rate", type=float, default=3e-5)
    parser.add_argument("--n_steps", type=int, default=2048)
    parser.add_argument("--batch_size", type=int, default=512)
    parser.add_argument("--n_epochs", type=int, default=10)
    # Wandb arguments
    parser.add_argument("--use_wandb", action="store_true", default=True)
    parser.add_argument("--wandb_project", type=str, default="metasim_rl_training")
    parser.add_argument("--wandb_entity", type=str, default="<your_wandb_entity>")
    args = parser.parse_args()
    return args


class EpisodeLogCallback(BaseCallback):
    """
    Custom callback for plotting additional values in tensorboard.
    """

    def __init__(self, verbose=0):
        super().__init__(verbose)
        self.returns_info = {
            "results/return": [],
            "results/episode_length": [],
            "results/success": [],
            "results/success_subtasks": [],
        }

    def _on_step(self) -> bool:
        infos = self.locals["infos"]
        for idx in range(len(infos)):
            curr_info = infos[idx]
            if "episode" in curr_info:
                self.returns_info["results/return"].append(curr_info["episode"]["r"])
                self.returns_info["results/episode_length"].append(curr_info["episode"]["l"])
                cur_info_success = curr_info.get("success", 0)
                self.returns_info["results/success"].append(cur_info_success)
                cur_info_success_subtasks = curr_info.get("success_subtasks", 0)
                self.returns_info["results/success_subtasks"].append(cur_info_success_subtasks)
        return True

    def _on_rollout_end(self) -> None:
        # Use the model's num_timesteps (cumulative across envs) as the global step.
        global_step = self.model.num_timesteps
        for key in self.returns_info.keys():
            if self.returns_info[key]:
                self.logger.record(key, np.mean(self.returns_info[key]), global_step)
                self.returns_info[key] = []


class EvalCallback(BaseCallback):
    def __init__(self, eval_every: int = 100000, scenario=None, sim_type=None, verbose: int = 0):
        super().__init__(verbose=verbose)
        self.eval_every = eval_every
        self.eval_env = None
        self.scenario = scenario
        self.sim_type = sim_type

    def _on_training_start(self) -> None:
        self.eval_env = MetaSimGymWrapper(scenario=self.scenario, sim_type=self.sim_type, num_envs=1)

    def _on_step(self) -> bool:
        if self.num_timesteps % self.eval_every == 0:
            self.record_video()
        return True

    def record_video(self) -> None:
        print("recording video")
        obs, info = self.eval_env.reset()
        video = [self.eval_env.render().transpose(2, 0, 1)]
        for _ in range(1000):
            action = self.model.predict(obs, deterministic=True)[0]
            obs, _, terminated, truncated, _ = self.eval_env.step(action)
            pixels = self.eval_env.render()
            if pixels is not None:
                video.append(pixels.transpose(2, 0, 1))
            if terminated or truncated:
                break

        if video:
            video = np.stack(video)
            wandb.log({"results/video": wandb.Video(video, fps=10, format="mp4")}, step=self.model.num_timesteps)
        else:
            log.warning("No video recorded")


class LogCallback(BaseCallback):
    """
    Custom callback for plotting additional values in tensorboard.
    """

    def __init__(self, verbose=0, info_keywords=(), args=None):
        super().__init__(verbose)
        self.aux_rewards = {}
        self.aux_returns = {}
        for key in info_keywords:
            self.aux_rewards[key] = np.zeros(args.num_envs)
            self.aux_returns[key] = deque(maxlen=100)

    def _on_step(self) -> bool:
        infos = self.locals["infos"]
        for idx in range(len(infos)):
            for key in self.aux_rewards.keys():
                self.aux_rewards[key][idx] += infos[idx][key]

            if self.locals["dones"][idx]:
                for key in self.aux_rewards.keys():
                    self.aux_returns[key].append(self.aux_rewards[key][idx])
                    self.aux_rewards[key][idx] = 0
        return True

    def _on_rollout_end(self) -> None:
        global_step = self.model.num_timesteps
        for key in self.aux_returns.keys():
            self.logger.record(f"aux_returns_{key}/mean", np.mean(self.aux_returns[key]), global_step)


def main():
    args = parse_args()

    if args.use_wandb:
        run = wandb.init(
            project=args.wandb_project,
            entity=args.wandb_entity,
            config=vars(args),
            sync_tensorboard=True,
            monitor_gym=True,
            save_code=True,
            name=f"SB3-{args.sim}-{args.robot}-{args.task}-{time.strftime('%Y_%m_%d_%H_%M_%S')}",
        )
    else:
        from collections import namedtuple

        Run = namedtuple("Run", ["id"])
        run = Run(id=int(time.time()))

    # Create scenario config
    task = get_task(args.task)
    robot = get_robot(args.robot)
    scenario = ScenarioMetaCfg(task=task, robot=robot, try_add_table=args.add_table, sim=args.sim)

    env = make_vec_env(
        lambda: MetaSimGymWrapper(
            scenario=scenario,
            sim_type=SimType(args.sim),
            num_envs=1,  # Each instance is a single environment
            headless=args.train_or_eval == "train",
        ),
        n_envs=args.num_envs,
    )

    # Initialize PPO algorithm
    model = PPO(
        policy="MlpPolicy",
        env=env,
        learning_rate=args.learning_rate,
        n_steps=args.n_steps,
        batch_size=args.batch_size,
        n_epochs=args.n_epochs,
        verbose=1,
        tensorboard_log=f"./ppo_logs/{run.id}",
        device="cpu",
    )

    # Setup wandb callback with additional monitoring options
    wandb_callback = (
        WandbCallback(
            model_save_freq=100000,
            model_save_path=f"models/{run.id}",
            verbose=2,
            gradient_save_freq=100000,  # Added to track gradients
        )
        if args.use_wandb
        else None
    )

    if args.train_or_eval == "train":
        # Train the agent with additional callbacks
        log.info("Starting training...")
        model.learn(
            total_timesteps=args.total_timesteps,
            log_interval=1,
            callback=[
                EvalCallback(scenario=scenario, sim_type=SimType(args.sim)),
                EpisodeLogCallback(),
                LogCallback(args=args),
            ]
            + ([wandb_callback] if args.use_wandb else []),
            progress_bar=True,
        )

        # Save the trained model
        model_path = f"models/ppo_model_{args.robot}_{args.sim}_{run.id}"
        model.save(model_path)
        log.info(f"Model saved to {model_path}")
    elif args.train_or_eval == "eval":
        # Load the trained model
        model.load(args.model_path)

        # Evaluate the agent
        log.info("Starting evaluation...")
        obs, info = env.reset()
        rewards = []
        for _ in range(1000):
            action, _ = model.predict(obs, deterministic=True)
            obs, reward, terminated, truncated, info = env.step(action)
            rewards.append(reward)
            env.render()
            if terminated or truncated:
                obs, info = env.reset()
                import matplotlib.pyplot as plt

                plt.plot(rewards)
                plt.savefig("rewards.png")
                plt.clf()
                rewards = []

    elif args.train_or_eval == "test":
        # Test the joints
        log.info("Starting evaluation...")
        obs, info = env.reset()
        rewards = []

        step = 0
        target_joint = 0
        while True:
            step += 1
            action = np.array(
                [0] * 19,
                dtype=np.float32,
            )
            action[target_joint] = np.sin(step / 10)
            # action = env.normalize_action(np.array([0,0,-0.4,0.8,-0.4,0,0,-0.4,0.8,-0.4,0,0,0,0,0,0,0,0,0]))
            # action = env.normalize_action(obs[7:26])
            obs, reward, terminated, truncated, info = env.step(action)
            env.render()

    # Close environment and wandb
    env.close()
    wandb.finish()


if __name__ == "__main__":
    main()
