# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

"""Retargeting for humand-object interaction with HDMI simulator.

Author: Chaoyi Pan
Date: 2025-10-18
"""

from __future__ import annotations

import time
from pathlib import Path

import hydra
import imageio
import loguru
import mujoco
import numpy as np
import torch
from omegaconf import DictConfig

from spider.config import Config, process_config
from spider.interp import get_slice
from spider.optimizers.sampling import (
    make_optimize_fn,
    make_optimize_once_fn,
    make_rollout_fn,
)
from spider.simulators.hdmi import (
    copy_sample_state,
    get_reference,
    get_reward,
    get_terminal_reward,
    get_trace,
    load_env_params,
    load_state,
    save_env_params,
    save_state,
    setup_env,
    step_env,
    sync_env,
)
from spider.viewers import (
    log_frame,
    render_image,
    setup_renderer,
    setup_viewer,
    update_viewer,
)


def main(config: Config):
    """Run the SPIDER using HDMI backend."""
    # Setup env (ref_data set to None since environment has built-in reference)
    env = setup_env(config, None)
    if config.max_sim_steps == -1:
        config.max_sim_steps = env.max_episode_length
        loguru.logger.info(f"Max simulation steps set to {config.max_sim_steps}")

    config.nu = env.action_spec.shape[-1]

    # Process config, set defaults and derived fields
    config = process_config(config)

    # Create placeholder reference data for compatibility
    ref_data = (
        torch.zeros(
            config.max_sim_steps + config.horizon_steps + config.ctrl_steps,
            config.nu,
            device=config.device,
        ),
    )

    # Setup env params (empty for HDMI, no domain randomization)
    env_params_list = []
    for _ in range(config.max_num_iterations):
        env_params = [{}] * config.num_dr
        env_params_list.append(env_params)
    config.env_params_list = env_params_list

    # Get reference data (states and controls)
    qpos_ref, qvel_ref, ctrl_ref = get_reference(config, env)
    np.savez(
        f"{config.output_dir}/trajectory_kinematic.npz",
        qpos=qpos_ref.detach().cpu().numpy(),
        qvel=qvel_ref.detach().cpu().numpy(),
        ctrl=ctrl_ref.detach().cpu().numpy(),
    )
    # optional: also save env xml
    # env.scene.to_zip(Path(config.output_dir) / "../scene.zip")

    # Setup mujoco model and data from HDMI env (for rendering)
    default_xml_path = Path(
        "example_datasets/processed/hdmi/unitree_g1/humanoid_object/move_suitcase/scene/mjlab scene.xml"
    )
    if default_xml_path.exists():
        mj_model = mujoco.MjModel.from_xml_path(str(default_xml_path))
        config.model_path = default_xml_path
    else:
        loguru.logger.warning(
            f"Default XML path {default_xml_path} does not exist, using env model"
        )
        mj_model = env.sim.mj_model
        config.model_path = "hdmi_scene_from_env"
    mj_data = mujoco.MjData(mj_model)
    mj_data_ref = mujoco.MjData(mj_model)

    # Initialize mj_data with current env state
    sim_data = env.sim.data
    mj_data.qpos[:] = sim_data.qpos[0].detach().cpu().numpy()
    mj_data.qvel[:] = sim_data.qvel[0].detach().cpu().numpy()
    mujoco.mj_step(mj_model, mj_data)
    mj_data.time = 0.0

    # Initialize reference mj_data
    mj_data_ref.qpos[:] = qpos_ref[0].detach().cpu().numpy()
    mj_data_ref.qvel[:] = qvel_ref[0].detach().cpu().numpy()
    mujoco.mj_step(mj_model, mj_data_ref)

    # Setup for video rendering
    images = []

    # Setup viewer and renderer
    # Note: mjlab viewer is disabled via cfg.viewer.headless in setup_env
    # For HDMI with rerun, build scene directly from spec and model
    # if "rerun" in config.viewer:
    #     # Build and log 3D scene from HDMI's spec and model
    #     if default_xml_path.exists():
    #         xml_path = default_xml_path
    #     else:
    #         loguru.logger.info(
    #             f"Default scene path {default_xml_path} does not exist, using None"
    #         )
    #         xml_path = None
    #     loguru.logger.info("Building Rerun scene from HDMI spec and model...")
    #     config.viewer_body_entity_and_ids = build_and_log_scene_from_spec(
    #         spec=mj_spec,
    #         model=mj_model,
    #         xml_path=xml_path,
    #         entity_root="mujoco",
    #     )
    #     loguru.logger.info(
    #         f"Rerun scene built with {len(config.viewer_body_entity_and_ids)} body entities"
    #     )
    #     # Set model_path to dummy value to indicate scene is already built
    #     config.model_path = "hdmi_scene_from_spec"

    run_viewer = setup_viewer(config, mj_model, mj_data)
    renderer = setup_renderer(config, mj_model)

    # Setup optimizer
    rollout = make_rollout_fn(
        step_env,
        save_state,
        load_state,
        get_reward,
        get_terminal_reward,
        get_trace,
        save_env_params,
        load_env_params,
        copy_sample_state,
    )
    optimize_once = make_optimize_once_fn(rollout)
    optimize = make_optimize_fn(optimize_once)

    # Initial controls - first horizon_steps from reference
    ctrls = ctrl_ref[: config.horizon_steps]

    # Buffers for saving info and trajectory
    info_list = []

    # Run viewer + control loop
    t_start = time.perf_counter()
    sim_step = 0
    with run_viewer() as viewer:
        while viewer.is_running():
            t0 = time.perf_counter()

            # Optimize using future reference window at control-rate (+1 lookahead)
            ref_slice = get_slice(
                ref_data, sim_step + 1, sim_step + config.horizon_steps + 1
            )
            if config.max_num_iterations > 0:
                ctrls, infos = optimize(config, env, ctrls, ref_slice)
            else:
                infos = {"opt_steps": [0], "improvement": 0.0}
            infos["sim_step"] = sim_step

            # Step environment for ctrl_steps
            step_info = {"qpos": [], "qvel": [], "time": [], "ctrl": [], "ctrl_ref": []}
            for i in range(config.ctrl_steps):
                ctrl = ctrls[i]
                ctrl_repeat = ctrl.unsqueeze(0).repeat(
                    int(config.num_samples), 1
                )  # (batch_size, num_actions)
                step_env(config, env, ctrl_repeat)

                # Update mj_data with current state
                sim_data = env.sim.data
                mj_data.qpos[:] = sim_data.qpos[0].detach().cpu().numpy()
                mj_data.qvel[:] = sim_data.qvel[0].detach().cpu().numpy()
                mj_data.time = (sim_step + 1) * config.sim_dt

                # Render video if enabled
                should_render = (
                    config.save_video
                    and renderer is not None
                    and i % int(np.round(config.render_dt / config.sim_dt)) == 0
                )
                if should_render:
                    # Get reference state from reference data
                    ref_idx = min(sim_step, qpos_ref.shape[0] - 1)
                    mj_data_ref.qpos[:] = qpos_ref[ref_idx].detach().cpu().numpy()
                    mj_data_ref.qvel[:] = qvel_ref[ref_idx].detach().cpu().numpy()
                    mujoco.mj_step(mj_model, mj_data_ref)
                    image = render_image(
                        config, renderer, mj_model, mj_data, mj_data_ref
                    )
                    images.append(image)
                if "rerun" in config.viewer or "viser" in config.viewer:
                    # manually log the state
                    log_frame(
                        mj_data,
                        sim_time=mj_data.time,
                        viewer_body_entity_and_ids=config.viewer_body_entity_and_ids,
                    )

                # Record state info
                step_info["qpos"].append(mj_data.qpos.copy())
                step_info["qvel"].append(mj_data.qvel.copy())
                step_info["time"].append(mj_data.time)
                step_info["ctrl"].append(ctrl.detach().cpu().numpy())
                step_info["ctrl_ref"].append(ctrl_ref[sim_step].detach().cpu().numpy())

                sim_step += 1

            # Stack step info
            for k in step_info:
                step_info[k] = np.stack(step_info[k], axis=0)
            infos.update(step_info)

            # Sync env state (broadcast from first env to all)
            env = sync_env(config, env)

            # Receding horizon update
            prev_ctrl = ctrls[config.ctrl_steps :]
            new_ctrl = ctrl_ref[
                sim_step + prev_ctrl.shape[0] : sim_step
                + prev_ctrl.shape[0]
                + config.ctrl_steps
            ]
            ctrls = torch.cat([prev_ctrl, new_ctrl], dim=0)

            # Sync viewer state and render
            sim_data = env.sim.data
            mj_data.qpos[:] = sim_data.qpos[0].detach().cpu().numpy()
            mj_data.qvel[:] = sim_data.qvel[0].detach().cpu().numpy()
            # Update reference state
            ref_idx = min(sim_step, qpos_ref.shape[0] - 1)
            mj_data_ref.qpos[:] = qpos_ref[ref_idx].detach().cpu().numpy()
            mj_data_ref.qvel[:] = qvel_ref[ref_idx].detach().cpu().numpy()
            mujoco.mj_step(mj_model, mj_data_ref)
            update_viewer(config, viewer, mj_model, mj_data, mj_data_ref, infos)

            # Progress
            t1 = time.perf_counter()
            rtr = config.ctrl_dt / (t1 - t0)
            print(
                f"Realtime rate: {rtr:.2f}, plan time: {t1 - t0:.4f}s, sim_steps: {sim_step}/{config.max_sim_steps}, opt_steps: {infos['opt_steps'][0]}",
                end="\r",
            )

            # Record info/trajectory at control tick
            info_list.append({k: v for k, v in infos.items() if k != "trace_sample"})

            if sim_step >= config.max_sim_steps:
                break

        t_end = time.perf_counter()
        print(f"\nTotal time: {t_end - t_start:.4f}s")

    # Save retargeted trajectory
    if config.save_info and len(info_list) > 0:
        info_aggregated = {}
        for k in info_list[0]:
            info_aggregated[k] = np.stack([info[k] for info in info_list], axis=0)
        np.savez(f"{config.output_dir}/trajectory_hdmi.npz", **info_aggregated)
        loguru.logger.info(f"Saved info to {config.output_dir}/trajectory_hdmi.npz")

    # Save video
    if config.save_video and len(images) > 0:
        video_path = f"{config.output_dir}/visualization_hdmi.mp4"
        imageio.mimsave(
            video_path,
            images,
            fps=int(1 / config.render_dt),
        )
        loguru.logger.info(f"Saved video to {video_path}")

    return


@hydra.main(version_base=None, config_path="config", config_name="hdmi")
def run_main(cfg: DictConfig) -> None:
    """Main entry point for HDMI retargeting."""
    # Convert DictConfig to Config dataclass, handling special fields
    config_dict = dict(cfg)

    # Handle special conversions
    if "noise_scale" in config_dict and config_dict["noise_scale"] is None:
        config_dict.pop("noise_scale")  # Let the default factory handle it

    # Convert lists to tuples where needed
    if "pair_margin_range" in config_dict:
        config_dict["pair_margin_range"] = tuple(config_dict["pair_margin_range"])
    if "xy_offset_range" in config_dict:
        config_dict["xy_offset_range"] = tuple(config_dict["xy_offset_range"])

    config = Config(**config_dict)
    main(config)


if __name__ == "__main__":
    run_main()