#!/usr/bin/env python3
"""DreamWaQ training via the faithful rsl_rl-1.0.2 port (upstream-aligned).

Uses upstream's ActorCritic_DWAQ + PPO (joint VAE training) + OnPolicyRunner,
with MotrixLab's DreamWaQ env. Config matches upstream Go1RoughCfgPPO.

Usage:
    uv run scripts/train_dreamwaq_rsl.py --num-envs 2048 --iterations 3000
"""
import argparse
import os
import sys
import torch

import motrix_envs.locomotion.go1.dreamwaq  # noqa: F401  (registers env)
from motrix_envs import registry as env_registry
from motrix_rl.dwaq_rsl import OnPolicyRunner, DwaqVecEnv


# Upstream Go1RoughCfgPPO (legged_robot_config.py LeggedRobotCfgPPO + Go1 overrides)
TRAIN_CFG = {
    "runner": {
        "policy_class_name": "ActorCritic_DWAQ",
        "algorithm_class_name": "PPO",
        "num_steps_per_env": 24,
        "save_interval": 50,
    },
    "algorithm": {
        "value_loss_coef": 1.0,
        "use_clipped_value_loss": True,
        "clip_param": 0.2,
        "entropy_coef": 0.01,
        "num_learning_epochs": 5,
        "num_mini_batches": 4,
        "learning_rate": 1.0e-3,
        "schedule": "adaptive",
        "gamma": 0.99,
        "lam": 0.95,
        "desired_kl": 0.01,
        "max_grad_norm": 1.0,
    },
    "policy": {
        "init_noise_std": 1.0,
    },
}


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--num-envs", type=int, default=2048)
    p.add_argument("--iterations", type=int, default=3000)
    p.add_argument("--seed", type=int, default=1)
    p.add_argument("--init-noise-std", type=float, default=1.0)
    p.add_argument("--force-std", action="store_true",
                   help="force reset action std to --init-noise-std even when resuming")
    p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    p.add_argument("--resume", default=None,
                   help="checkpoint .pt to warm-start from (e.g. the flat policy before "
                        "the pyramid-terrain curriculum). Loads model + optimizer; the "
                        "action std comes from the checkpoint, not --init-noise-std.")
    p.add_argument("--level", type=int, default=None,
                   help="force ALL envs to this terrain level (0-9), skip curriculum")
    args = p.parse_args()

    torch.manual_seed(args.seed)
    import numpy as np
    np.random.seed(args.seed)

    TRAIN_CFG["policy"]["init_noise_std"] = args.init_noise_std

    raw_env = env_registry.make("go1-dreamwaq-walk", num_envs=args.num_envs)
    if args.level is not None:
        raw_env._force_level = args.level
        print(f"[DreamWaQ-rsl] forcing ALL envs at terrain level {args.level}")
    env = DwaqVecEnv(raw_env, device=args.device)
    print(f"[DreamWaQ-rsl] {args.num_envs} envs | obs={env.num_obs} priv={env.num_privileged_obs} "
          f"hist={env.num_obs_hist} act={env.num_actions} | device={args.device}")
    print(f"[DreamWaQ-rsl] actor_in={env.num_obs + 19} critic_in={env.num_privileged_obs} "
          f"cenet_in={env.num_obs_hist * env.num_obs}")

    log_dir = os.path.join("runs", "go1-dreamwaq-walk", "rsl_dwaq",
                           __import__("datetime").datetime.now().strftime("%m-%d_%H-%M-%S"))
    os.makedirs(log_dir, exist_ok=True)

    # Save config snapshot for later reference
    import json
    cfg_snapshot = {
        "command_line": sys.argv,
        "kp": raw_env.cfg.control_config.stiffness,
        "kd": raw_env.cfg.control_config.damping,
        "action_scale": raw_env.cfg.control_config.action_scale,
        "rewards": dict(raw_env.cfg.reward_config.scales),
        "sigma": raw_env.cfg.reward_config.tracking_sigma,
        "only_positive": raw_env.cfg.reward_config.only_positive_rewards,
        "force_level": args.level,
        "init_noise_std": TRAIN_CFG["policy"]["init_noise_std"],
        "entropy_coef": TRAIN_CFG["algorithm"]["entropy_coef"],
        "terrain_rows": raw_env.cfg.terrain_rows,
        "terrain_cols": raw_env.cfg.terrain_cols,
        "scene": raw_env.cfg.model_file,
    }
    with open(os.path.join(log_dir, "config.json"), "w") as f:
        json.dump(cfg_snapshot, f, indent=2, default=str)

    runner = OnPolicyRunner(env, TRAIN_CFG, log_dir=log_dir, device=args.device)
    if args.resume:
        runner.load(args.resume)
        runner.current_learning_iteration = 0
        if args.force_std:
            runner.alg.actor_critic.std.data.fill_(args.init_noise_std)
            print(f"[DreamWaQ-rsl] warm-start from {args.resume} "
                  f"(std FORCED to {args.init_noise_std})")
        else:
            print(f"[DreamWaQ-rsl] warm-start from {args.resume} "
                  f"(model+optimizer; std from checkpoint)")
    print(f"[DreamWaQ-rsl] log_dir={log_dir} | training {args.iterations} iterations...")
    runner.learn(args.iterations, init_at_random_ep_len=True)
    print("[DreamWaQ-rsl] done.")


if __name__ == "__main__":
    main()