# -*- coding: utf-8 -*-
'''
@File    : base_pipeline.py
@Time    : 2025/11/27 15:53:26
@Author  : wty-yy
@Version : 1.0
@Blog    : https://wty-yy.github.io/
@Desc    : Base Pipeline for Robogauge
'''
import yaml
import random
from copy import deepcopy
from pathlib import Path

from robogauge.utils.logger import logger
from robogauge.tasks.simulator import MujocoSimulator, MujocoConfig, SimData
from robogauge.tasks.robots import (
    BaseRobot, RobotConfig, Go2Config, Go2, Go2MoEConfig, Go2MoE
)
from robogauge.tasks.gauge import BaseGauge, BaseGaugeConfig
from robogauge.utils.helpers import class_to_dict

class BasePipeline:
    def __init__(self, 
        run_name: str,
        simulator_cfg: MujocoConfig,
        robot_cfg: RobotConfig,
        gauge_cfg: BaseGaugeConfig
    ):
        self.run_name = run_name
        self.sim_cfg = simulator_cfg
        self.robot_cfg = robot_cfg
        self.gauge_cfg = gauge_cfg

        self.sim: MujocoSimulator = eval(simulator_cfg.simulator_class)(simulator_cfg)
        self.robot: BaseRobot = eval(robot_cfg.robot_class)(robot_cfg)
        self.gauge: BaseGauge = eval(gauge_cfg.gauge_class)(gauge_cfg, robot_cfg)
    
        # save configs
        cfg = {}
        for name in ['sim_cfg', 'robot_cfg', 'gauge_cfg']:
            obj = getattr(self, name)
            obj_dict = class_to_dict(obj)
            cfg.update({name: obj_dict})
        with open(Path(logger.log_dir) / "configs.yaml", 'w') as file:
            yaml.dump(cfg, file)
    
    def load(self):
        self.sim.load(
            self.gauge_cfg.assets.terrain_xml,
            self.robot_cfg.assets.robot_xml,
            self.gauge_cfg.assets.terrain_spawn_pos,
            self.robot_cfg.control.default_dof_pos
        )
    
    def run(self):
        logger.info(f"🚀 Starting single run: {self.run_name}")
        try:
            self.load()
            sim_data = self.sim.step()
            frame_skip = int(self.robot_cfg.control.control_dt / self.sim_cfg.physics.simulation_dt)
            logger.info(f"Sim FPS: {1.0 / self.sim_cfg.physics.simulation_dt:.2f}, Control FPS: {1.0 / self.robot_cfg.control.control_dt:.2f}, Frame Skip: {frame_skip:d}")
            logger.info("Running pipeline...")
            while not self.gauge.is_done():
                goal_data = self.gauge.get_goal(sim_data)
                if goal_data is None:
                    continue
                obs = self.robot.build_observation(self.add_noise(sim_data), goal_data)
                action, p_gains, d_gains, control_type = self.robot.get_action(obs)

                if self.sim_cfg.domain_rand.action_delay:
                    actions_start_decimation = random.randint(0, frame_skip)
                else:
                    self.sim.setup_action(action, p_gains, d_gains, control_type)
                for i in range(frame_skip):
                    if self.sim_cfg.domain_rand.action_delay and i == actions_start_decimation:
                        self.sim.setup_action(action, p_gains, d_gains, control_type)
                    sim_data = self.sim.step()
                    self.gauge.update_metrics(sim_data, goal_data)
                if self.gauge.is_reset(sim_data):
                    self.sim.reset()
                    sim_data = self.sim.step()
        finally:
            self.sim.close_viewer()
            self.sim.close_video_writer()
            logger.info("✅ Pipeline execution finished.")
            logger.info(f"📁 Logging saved at: {logger.log_dir}")
        
        return logger.log_dir

    def add_noise(self, sim_data: SimData):
        sim_data = deepcopy(sim_data)
        noise_cfg = self.sim_cfg.noise
        if not noise_cfg.enabled:
            return sim_data
        proprio = sim_data.proprio
        def add_uniform_noise(data, noise_level):
            for i in range(len(data)):
                noise = random.uniform(-noise_level, noise_level)
                data[i] += noise
        add_uniform_noise(proprio.joint.pos, noise_cfg.joint_pos)
        add_uniform_noise(proprio.joint.vel, noise_cfg.joint_vel)
        add_uniform_noise(proprio.base.lin_vel, noise_cfg.lin_vel)
        add_uniform_noise(proprio.base.ang_vel, noise_cfg.ang_vel)
        add_uniform_noise(proprio.imu.lin_vel, noise_cfg.lin_vel)
        add_uniform_noise(proprio.imu.ang_vel, noise_cfg.ang_vel)
        return sim_data