# -*- coding: utf-8 -*-
'''
@File    : helpers.py
@Time    : 2025/11/27 15:26:37
@Author  : wty-yy
@Version : 1.0
@Blog    : https://wty-yy.github.io/
@Desc    : Helpers for robogauge, include:
- Argument parsing
- Class to dict conversion
- Path parsing
'''
import yaml
from argparse import ArgumentParser
from pathlib import Path
from robogauge import ROBOGAUGE_ROOT_DIR

def parse_path(path):
    if "{ROBOGAUGE_ROOT_DIR}" in str(path):
        path = str(path).replace("{ROBOGAUGE_ROOT_DIR}", ROBOGAUGE_ROOT_DIR)
    return path

def class_to_dict(obj) -> dict:
    # From https://github.com/leggedrobotics/legged_gym/blob/master/legged_gym/utils/helpers.py
    if not hasattr(obj, "__dict__"):
        return obj
    result = {}
    for key in dir(obj):
        if key.startswith("_"):
            continue
        element = []
        val = getattr(obj, key)
        if isinstance(val, list):
            for item in val:
                element.append(class_to_dict(item))
        else:
            element = class_to_dict(val)
        result[key] = element
    return result

def set_seed(seed: int):
    import os
    import torch
    import random
    import numpy as np
    assert seed >= 0, "Seed must be non-negative."

    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)

def str2bool(v):
    if v.lower() in ('yes', 'true', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
    raise TypeError('Boolean value expected.')

def parse_args():
    parser = ArgumentParser()
    parameters = [
        # Single run parameters
        {"name": "--task-name", "type": str, "default": "base", "help": "Name of the task to run."},
        {"name": "--experiment-name", "type": str, "help": "Name of the experiment to run."},
        {"name": "--run-name", "type": str, "default": "run", "help": "Name of the run."},
        {"name": "--model-path", "type": str, "help": "Path to the model file."},
        {"name": "--headless", "action": "store_true", "default": False, "help": "Run in headless mode."},
        {"name": "--save-video", "action": "store_true", "default": False, "help": "Save video output."},
        {"name": "--seed", "type": int, "default": 42, "help": "Random seed."},
        {"name": "--write-tensorboard", "action": "store_true", "default": False, "help": "Write tensorboard logs."},
        {"name": "--plot-radar", "action": "store_true", "default": False, "help": "Plot radar charts for metrics."},
        {"name": "--base-mass", "type": float, "default": 0.0, "help": "Set the base mass of the robot."},
        {"name": "--friction", "type": float, "default": 1.0, "help": "Set the ground friction coefficient."},
        {"name": "--level", "type": int, "help": "Set the difficulty level of the environment, range 1-10 (flat is 0)."},

        # Multiprocessing parameters, with different seeds
        {"name": "--multi", "action": "store_true", "default": False, "help": "Enable multiprocessing."},
        {"name": "--num-processes", "type": int, "default": 2, "help": "Number of parallel processes."},
        {"name": "--seeds", "type": int, "nargs": "+", "default": [0], "help": "List of random seeds for multiple runs."},
        {"name": "--base-masses", "type": float, "nargs": "+", "default": [0], "help": "List of base masses for the model."},
        {"name": "--frictions", "type": float, "nargs": "+", "default": [0.4, 0.7, 1.0, 1.3, 1.6], "help": "List of friction coefficients for the model."},

        # Level pipeline parameters
        {"name": "--search-max-level", "action": "store_true", "default": False, "help": "Use level pipeline to search maximum level."},
    ]
    for param in parameters:
        parser.add_argument(param['name'], **{k: v for k, v in param.items() if k != 'name'})
    args = parser.parse_args()
    if args.experiment_name is not None:
        args.experiment_name = f"{args.task_name}_{args.experiment_name}"
    else:
        args.experiment_name = args.task_name
    return args

def snake_to_pascal(s: str) -> str:
    """ 'snake_case' to 'PascalCase' conversion """
    parts = [p for p in s.split('_') if p]
    return ''.join(p[0].upper() + p[1:] if p else '' for p in parts)