# imports
import h5py # type: ignore
import json
import numpy as np
import pandas as pd
from pathlib import Path
import mujoco # type: ignore
import argparse
from scipy.io import loadmat # type: ignore

def recursive_load_from_h5py(file, group_path="/"):
    group = file[group_path]
    if isinstance(group, h5py.Dataset):
        return group[()]  # Read dataset value
    elif isinstance(group, h5py.Group):
        if all(k.isdigit() for k in group.keys()):  # Likely a list
            return [
                recursive_load_from_h5py(file, f"{group_path}/{k}")
                for k in sorted(group.keys(), key=int)
            ]
        else:  # Dictionary-like group
            return {
                k: recursive_load_from_h5py(file, f"{group_path}/{k}")
                for k in group.keys()
            }
    else:
        raise TypeError(f"Unsupported group type: {type(group)}")

def load_from_h5py(file_path: str | Path, group_path="/"):
    """
    Load a pytree structure from an HDF5 file.

    Args:
        file (h5py.File): An open HDF5 file object.
        group_path (str): The HDF5 group path to read data from.

    Returns:
        The reconstructed data structure.
    """
    with h5py.File(file_path, "r") as file:
        return recursive_load_from_h5py(file, group_path)

def get_h5_data(file):
    data = load_from_h5py(file)
    kp_data = np.array(data['kp_data'][0], dtype=np.float32)
 
    return kp_data, np.array(data['kp_names'], dtype=np.dtypes.StringDType())

def get_json_data(file):
    data = None
    with open (file, 'r') as f:
        data = json.load(file)
    
    return np.array(data['kp_data'][0], dtype=np.float32), np.array(data['kp_names'], dtype=np.dtypes.StringDType())

def get_csv_data(file):
    kp_data = []
    kp_names = []
    xyz_coords = []

    df = pd.read_csv(file)
    kp_df = df.filter(regex='_x|_y|_z', axis=1)
    kp_df_data = kp_df.iloc[0].to_dict()

    for each in kp_df_data:
        xyz_coords.append(kp_df_data[each])
        name = each[:-2]
        if not (name in kp_names):
            kp_names.append(name)
    for i in range(2, len(xyz_coords), 3):
        coords_compressed = [xyz_coords[i-2], xyz_coords[i-1], xyz_coords[i]]
        kp_data.append(coords_compressed)
    
    return np.array(kp_data, dtype=np.float32), np.array(kp_names, dtype=np.dtypes.StringDType())

def get_xml(file):
    xml_path = file
    model = mujoco.MjModel.from_xml_path(xml_path)
    model_data = mujoco.MjData(model)
    mujoco.mj_kinematics(model, model_data)

    geom_names = []
    body_names = []
    geoms_to_body = {}
    body_pos_list = []

    for i in range(model.ngeom):
        geom_name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_GEOM, i)

        geom_names.append(geom_name)

        body_id = model.geom_bodyid[i]
        body_pos = model_data.xpos[body_id].copy()
        body_name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_BODY, body_id)

        geoms_to_body[geom_name] = body_name
        body_pos_list.append(body_pos)
        body_names.append(body_name)

    absolute_coordinates_geoms = np.array(list(model_data.geom_xpos), dtype=np.float32)

    return absolute_coordinates_geoms, geom_names, geoms_to_body
  
def get_both_coordinates(file1: str, file2: str):
    kp_data = None
    kp_names = None

    if file1[-3:] == 'csv':
        kp_data, kp_names = get_csv_data(file1)
    elif file1[-3:] == '.h5':
        kp_data, kp_names = get_h5_data(file1)
    elif file1[-4:] == 'json':
        kp_data, kp_names = get_json_data(file1)

    model_abs_coords, geom_names, geoms_to_body = get_xml(file2)

    return kp_data, model_abs_coords, kp_names, geom_names, geoms_to_body