import h5py
import numpy as np
import yaml
import os

class MyDumper(yaml.Dumper):
    def represent_float(self, data):
        return self.represent_scalar('tag:yaml.org,2002:float', f"{data:.6f}")

MyDumper.add_representer(float, MyDumper.represent_float)


def collect_qpos_from_tasks(root_dir, task_list, task_config, output_yaml_path):

    assert task_config in ['demo_clean', 'demo_randomized'], "task_config must be 'demo_clean' or 'demo_randomized'"

    all_qpos_data = []

    for task_name in task_list:
        data_folder = os.path.join(root_dir, task_name, task_config, 'data')
        if not os.path.isdir(data_folder):
            print(f"??  Skipping missing folder: {data_folder}")
            continue

        for file_name in os.listdir(data_folder):
            if not file_name.endswith('.hdf5'):
                continue
            file_path = os.path.join(data_folder, file_name)
            try:
                with h5py.File(file_path, 'r') as file:
                    qpos_data = file['joint_action/vector'][:]
                    all_qpos_data.append(qpos_data)
            except KeyError as e:
                print(f"??  Dataset {e} not found in {file_name}, skipping")
            except Exception as e:
                print(f"? Error processing {file_name}: {e}")

    if not all_qpos_data:
        raise ValueError(f"No valid 'joint_action/vector' data found for tasks: {task_list}")

    all_qpos_data = np.concatenate(all_qpos_data, axis=0)
    mean_values = np.mean(all_qpos_data, axis=0).tolist()
    std_values = np.std(all_qpos_data, axis=0).tolist()
    min_bounds = np.min(all_qpos_data, axis=0).tolist()
    max_bounds = np.max(all_qpos_data, axis=0).tolist()

    yaml_data = {
        'robot_obs': {
            'mean': mean_values,
            'std': std_values
        },
        'act_min_bound': min_bounds,
        'act_max_bound': max_bounds
    }

    with open(output_yaml_path, 'w') as f:
        yaml.dump(yaml_data, f, default_flow_style=False, Dumper=MyDumper)

    print(f"? YAML generated for {task_list[0]} at {output_yaml_path}")


def process_all_tasks(root_data_dir, output_dir, task_config):

    os.makedirs(output_dir, exist_ok=True)
    print(f"Output directory set to: {output_dir}")

    for task_name in os.listdir(root_data_dir):
        task_path = os.path.join(root_data_dir, task_name)
        if os.path.isdir(task_path):
            print(f"\n--- Processing task: {task_name} ---")

            output_filename = f"{task_name}_obs_statistics.yaml"
            output_yaml_path = os.path.join(output_dir, output_filename)

            try:
                collect_qpos_from_tasks(
                    root_dir=root_data_dir,
                    task_list=[task_name],
                    task_config=task_config,
                    output_yaml_path=output_yaml_path
                )
            except ValueError as e:
                print(f"??  Skipping task '{task_name}' due to an error: {e}")
            except Exception as e:
                print(f"? An unexpected error occurred while processing task '{task_name}': {e}")
    
    print("\n All tasks processed.")

if __name__ == "__main__":
    root_data_dir = "/yourpath/RoboTwin/data/"
    output_directory = "./yaml_statistics"
    task_config = "demo_clean" #or demo_randomized
    process_all_tasks(root_data_dir, output_directory, task_config)