# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
"""DreamWaQ OnPolicyRunner — adapted from upstream rsl_rl-1.0.2.

Faithful port of Manaro-Alpha/DreamWaQ's on_policy_runner.py, with the
rsl_rl-package imports replaced by local (relative) imports and the isaacgym
noise removed. Drives the PPO + CENet (VAE) joint training loop.
"""
import time
import os
from collections import deque
import statistics

from torch.utils.tensorboard import SummaryWriter
import torch

from .ppo import PPO
from .actor_critic_dwaq import ActorCritic_DWAQ


class OnPolicyRunner:
    def __init__(self, env, train_cfg, log_dir=None, device="cpu"):
        self.cfg = train_cfg["runner"]
        self.alg_cfg = train_cfg["algorithm"]
        self.policy_cfg = train_cfg["policy"]
        self.device = device
        self.env = env

        num_critic_obs = env.num_privileged_obs if env.num_privileged_obs is not None else env.num_obs
        cenet_in_dim = env.num_obs_hist * env.num_obs   # 5*45 = 225
        cenet_out_dim = 19                               # vel(3) + latent(16)

        actor_critic = ActorCritic_DWAQ(
            env.num_obs + cenet_out_dim,   # actor input: obs(45) + code(19) = 64
            num_critic_obs,                # critic input: privileged_obs
            env.num_actions,               # 12
            cenet_in_dim,                  # 225
            cenet_out_dim,                 # 19
            **self.policy_cfg,
        ).to(self.device)

        self.alg = PPO(actor_critic, device=self.device, **self.alg_cfg)
        self.num_steps_per_env = self.cfg["num_steps_per_env"]
        self.save_interval = self.cfg["save_interval"]

        self.alg.init_storage(
            env.num_envs, self.num_steps_per_env,
            [env.num_obs], [env.num_privileged_obs],
            [env.num_obs_hist * env.num_obs], [env.num_actions],
        )

        self.log_dir = log_dir
        self.writer = None
        self.tot_timesteps = 0
        self.tot_time = 0
        self.current_learning_iteration = 0

        _, _, _, _ = self.env.reset()

    def learn(self, num_learning_iterations, init_at_random_ep_len=False):
        if self.log_dir is not None and self.writer is None:
            self.writer = SummaryWriter(log_dir=self.log_dir, flush_secs=10)
        if init_at_random_ep_len:
            self.env.episode_length_buf = torch.randint_like(
                self.env.episode_length_buf, high=int(self.env.max_episode_length))

        obs, obs_hist = self.env.get_observations()
        privileged_obs, prev_critic_obs = self.env.get_privileged_observations()
        critic_obs = privileged_obs if privileged_obs is not None else obs
        obs, critic_obs, prev_critic_obs, obs_hist = (
            obs.to(self.device), critic_obs.to(self.device),
            prev_critic_obs.to(self.device), obs_hist.to(self.device))
        self.alg.actor_critic.train()

        ep_infos = []
        rewbuffer = deque(maxlen=100)
        lenbuffer = deque(maxlen=100)
        cur_reward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device)
        cur_episode_length = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device)

        tot_iter = self.current_learning_iteration + num_learning_iterations
        for it in range(self.current_learning_iteration, tot_iter):
            start = time.time()
            with torch.inference_mode():
                for i in range(self.num_steps_per_env):
                    actions = self.alg.act(obs, critic_obs, prev_critic_obs, obs_hist)
                    obs, privileged_obs, prev_privileged_obs, obs_hist, rewards, dones, infos = self.env.step(actions)
                    critic_obs = privileged_obs if privileged_obs is not None else obs
                    prev_critic_obs = prev_privileged_obs
                    obs, critic_obs, prev_critic_obs, obs_hist, rewards, dones = (
                        obs.to(self.device), critic_obs.to(self.device),
                        prev_critic_obs.to(self.device), obs_hist.to(self.device),
                        rewards.to(self.device), dones.to(self.device))
                    self.alg.process_env_step(rewards, dones, infos)

                    if self.log_dir is not None:
                        if "episode" in infos:
                            ep_infos.append(infos["episode"])
                        cur_reward_sum += rewards
                        cur_episode_length += 1
                        new_ids = (dones > 0).nonzero(as_tuple=False)
                        rewbuffer.extend(cur_reward_sum[new_ids][:, 0].cpu().numpy().tolist())
                        lenbuffer.extend(cur_episode_length[new_ids][:, 0].cpu().numpy().tolist())
                        cur_reward_sum[new_ids] = 0
                        cur_episode_length[new_ids] = 0

                stop = time.time()
                collection_time = stop - start
                start = stop
                self.alg.compute_returns(critic_obs)

            mean_value_loss, mean_surrogate_loss, mean_autoenc_loss = self.alg.update()
            stop = time.time()
            learn_time = stop - start
            if self.log_dir is not None:
                self.log(locals())
            if it % self.save_interval == 0:
                self.save(os.path.join(self.log_dir, "model_{}.pt".format(it)))
            ep_infos.clear()

        self.current_learning_iteration += num_learning_iterations
        self.save(os.path.join(self.log_dir, "model_{}.pt".format(self.current_learning_iteration)))

    def log(self, locs, width=80, pad=35):
        self.tot_timesteps += self.num_steps_per_env * self.env.num_envs
        self.tot_time += locs["collection_time"] + locs["learn_time"]
        iteration_time = locs["collection_time"] + locs["learn_time"]

        if locs["ep_infos"]:
            for key in locs["ep_infos"][0]:
                infotensor = torch.tensor([], device=self.device)
                for ep_info in locs["ep_infos"]:
                    v = ep_info[key]
                    if not isinstance(v, torch.Tensor):
                        v = torch.Tensor([v])
                    if len(v.shape) == 0:
                        v = v.unsqueeze(0)
                    infotensor = torch.cat((infotensor, v.to(self.device)))
                self.writer.add_scalar("Episode/" + key, torch.mean(infotensor), locs["it"])
        mean_std = self.alg.actor_critic.std.mean()
        fps = int(self.num_steps_per_env * self.env.num_envs / (locs["collection_time"] + locs["learn_time"]))

        self.writer.add_scalar("Loss/value_function", locs["mean_value_loss"], locs["it"])
        self.writer.add_scalar("Loss/surrogate", locs["mean_surrogate_loss"], locs["it"])
        self.writer.add_scalar("Loss/autoenc_function", locs["mean_autoenc_loss"], locs["it"])
        self.writer.add_scalar("Loss/learning_rate", self.alg.learning_rate, locs["it"])
        self.writer.add_scalar("Policy/mean_noise_std", mean_std.item(), locs["it"])
        self.writer.add_scalar("Perf/total_fps", fps, locs["it"])
        if len(locs["rewbuffer"]) > 0:
            self.writer.add_scalar("Train/mean_reward", statistics.mean(locs["rewbuffer"]), locs["it"])
            self.writer.add_scalar("Train/mean_episode_length", statistics.mean(locs["lenbuffer"]), locs["it"])

        if it_is_log := (locs["it"] % 10 == 0):
            rew = statistics.mean(locs["rewbuffer"]) if len(locs["rewbuffer"]) > 0 else 0.0
            elen = statistics.mean(locs["lenbuffer"]) if len(locs["lenbuffer"]) > 0 else 0.0
            print(f"[it {locs['it']}/{self.current_learning_iteration + locs['num_learning_iterations']}] "
                  f"fps={fps} val_loss={locs['mean_value_loss']:.3f} surr={locs['mean_surrogate_loss']:.4f} "
                  f"vae={locs['mean_autoenc_loss']:.3f} std={mean_std.item():.2f} "
                  f"rew={rew:.2f} eplen={elen:.0f}")

    def save(self, path, infos=None):
        torch.save({
            "model_state_dict": self.alg.actor_critic.state_dict(),
            "optimizer_state_dict": self.alg.optimizer.state_dict(),
            "iter": self.current_learning_iteration,
            "infos": infos,
        }, path)

    def load(self, path, load_optimizer=True):
        loaded_dict = torch.load(path, map_location=self.device)
        self.alg.actor_critic.load_state_dict(loaded_dict["model_state_dict"])
        if load_optimizer:
            self.alg.optimizer.load_state_dict(loaded_dict["optimizer_state_dict"])
        self.current_learning_iteration = loaded_dict["iter"]
        return loaded_dict["infos"]

    def get_inference_policy(self, device=None):
        self.alg.actor_critic.eval()
        if device is not None:
            self.alg.actor_critic.to(device)
        return self.alg.actor_critic.act_inference
