# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

from typing import Optional, Tuple

import numpy as np
import torch
from scipy.spatial.transform import Rotation

from ardy.assets import skeleton_asset_path
from ardy.data_processing.bvh import Bvh

_DEFAULT_G1_XML = str(skeleton_asset_path("g1skel34", "xml", "g1.xml"))


class Bone:
    """Abstraction for a bone in character skeleton."""

    def __init__(self):
        # original bone info
        self.id = None
        self.name = None
        self.orient = np.identity(3)
        self.dof_index = []
        self.channels = []  # bvh only
        self.lb = []
        self.ub = []
        self.parent = None
        self.child = []

        # asf specific
        self.dir = np.zeros(3)
        self.len = 0
        # bvh specific
        self.offset = np.zeros(3)  # default offset for position
        self.offset_rot = None  # rotation for custom nv bvh

        # inferred info
        self.pos = np.zeros(3)
        self.end = np.zeros(3)

    def __repr__(self):
        return f"{self.name}"


class Skeleton:
    """Abstraction for character skeleton."""

    def __init__(self):
        self.bones = []
        self.name2bone = {}
        self.mass_scale = 1.0
        self.len_scale = 1.0
        self.dof_name = ["x", "y", "z"]
        self.root = None

    def get_bones_names(self):
        return [x.name for x in self.bones]

    def get_parent_indices(self):
        parent_indices = [-1] * len(self.bones)
        for bone in self.bones:
            if bone.parent:
                parent_indices[bone.id] = bone.parent.id
        return parent_indices

    def get_neutral_joints(self):
        joints = []
        for bone in self.bones:
            joints.append(bone.pos)
        joints = np.stack(joints, axis=0)
        return joints

    def load_from_bvh(self, fname, exclude_bones=None, spec_channels=None):
        if exclude_bones is None:
            exclude_bones = {}
        if spec_channels is None:
            spec_channels = dict()
        with open(fname) as f:
            mocap = Bvh(f.read())

        joint_names = list(
            filter(
                lambda x: all([t not in x for t in exclude_bones]),
                mocap.get_joints_names(),
            )
        )
        dof_ind = {"x": 0, "y": 1, "z": 2}
        self.len_scale = 1.0
        self.root = Bone()
        self.root.id = 0
        self.root.name = joint_names[0]
        self.root.channels = mocap.joint_channels(self.root.name)
        self.root.offset = np.array(mocap.joint_offset(self.root.name)) * self.len_scale
        self.root.offset_rot = mocap.joint_offset_rot(self.root.name)
        if self.root.offset_rot is not None:
            self.root.offset_rot = np.array(self.root.offset_rot)
        self.name2bone[self.root.name] = self.root
        self.bones.append(self.root)
        for i, joint in enumerate(joint_names[1:]):
            bone = Bone()
            bone.id = i + 1
            bone.name = joint
            bone.channels = spec_channels[joint] if joint in spec_channels.keys() else mocap.joint_channels(joint)
            bone.dof_index = [dof_ind[x[0].lower()] for x in bone.channels]
            bone.offset = np.array(mocap.joint_offset(joint)) * self.len_scale
            bone.offset_rot = mocap.joint_offset_rot(joint)
            if bone.offset_rot is not None:
                bone.offset_rot = np.array(bone.offset_rot)
            bone.lb = [-180.0] * 3
            bone.ub = [180.0] * 3
            self.bones.append(bone)
            self.name2bone[joint] = bone

        # for bone in self.bones:
        # print(bone.name, bone.channels, bone.offset)

        for bone in self.bones[1:]:
            parent_name = mocap.joint_parent(bone.name).name
            if parent_name in self.name2bone.keys():
                bone_p = self.name2bone[parent_name]
                bone_p.child.append(bone)
                bone.parent = bone_p

        self.forward_bvh(self.root)
        for bone in self.bones:
            if len(bone.child) == 0:
                child_vals = [str(node) for node in mocap.get_joint(bone.name).children]
                if "End Site" in child_vals:
                    end_site_idx = child_vals.index("End Site")
                    end_site_offset = mocap.get_joint(bone.name).children[end_site_idx]["OFFSET"]
                    bone.end = bone.pos + np.array([float(x) for x in end_site_offset]) * self.len_scale
                else:
                    pass
            else:
                bone.end = sum([bone_c.pos for bone_c in bone.child]) / len(bone.child)

    def load_from_xml(self, fname, exclude_bones=None, spec_channels=None):
        if exclude_bones is None:
            exclude_bones = {}
        if spec_channels is None:
            spec_channels = dict()

        import xml.etree.ElementTree as ET

        tree = ET.parse(fname)
        root = tree.getroot()

        def map_bone_name(body_name):
            if body_name == "torso_link":
                return "waist_pitch_skel"
            elif "_link" in body_name:
                return body_name.replace("_link", "_skel")
            else:
                return body_name + "_skel"

        bodies = root.find("worldbody").findall(".//body")
        name_to_body = {map_bone_name(body.get("name")): body for body in bodies}
        joint_names = [body.get("name") for body in bodies]

        R_zup_to_yup = Rotation.from_euler("x", -90, degrees=True)
        x_forward_to_y_forward = Rotation.from_euler("z", -90, degrees=True)
        mujoco_to_ardy = R_zup_to_yup * x_forward_to_y_forward

        def mujoco_joint_translation(body_name):
            body = name_to_body[body_name]
            if "pos" in body.attrib:
                result = np.array([100 * float(x) for x in body.get("pos").strip().split(" ")])
                return mujoco_to_ardy.apply(result)
            else:
                return np.zeros(3)

        parent_map = {
            map_bone_name(child.get("name")): map_bone_name(parent.get("name"))
            for parent in root.iter()
            if parent.get("name")
            for child in parent
            if child.get("name")
        }

        self.len_scale = 1.0
        self.root = Bone()
        self.root.id = 0
        self.root.name = map_bone_name(joint_names[0])
        self.root.offset = mujoco_joint_translation(self.root.name)
        self.root.offset_rot = np.eye(3)

        self.name2bone[self.root.name] = self.root
        self.bones.append(self.root)
        for i, joint in enumerate(joint_names[1:]):
            bone = Bone()
            bone.id = i + 1
            bone.name = map_bone_name(joint)
            bone.offset = mujoco_joint_translation(bone.name)
            bone.offset_rot = np.eye(3)
            self.bones.append(bone)
            self.name2bone[bone.name] = bone

        for bone in self.bones[1:]:
            parent_name = parent_map[bone.name]
            if parent_name in self.name2bone.keys():
                bone_p = self.name2bone[parent_name]
                bone_p.child.append(bone)
                bone.parent = bone_p

        self.forward_bvh(self.root)

    def forward_bvh(self, bone):
        if bone.parent:
            bone.pos = bone.parent.pos + bone.offset
        else:
            bone.pos = bone.offset
        for bone_c in bone.child:
            self.forward_bvh(bone_c)


class SkeletonG1(Skeleton):
    """Add functionality specific to the G1 including loading from XML specific to G1 and adding
    dummy joints."""

    def load_from_xml(self, fname, exclude_bones=None, spec_channels=None):
        if exclude_bones is None:
            exclude_bones = {}
        if spec_channels is None:
            spec_channels = dict()

        import xml.etree.ElementTree as ET

        tree = ET.parse(fname)
        root = tree.getroot()

        def map_bone_name(body_name):
            if body_name == "torso_link":
                return "waist_pitch_skel"
            elif "_link" in body_name:
                return body_name.replace("_link", "_skel")
            else:
                return body_name + "_skel"

        bodies = root.find("worldbody").findall(".//body")
        name_to_body = {map_bone_name(body.get("name")): body for body in bodies}
        joint_names = [body.get("name") for body in bodies]

        R_zup_to_yup = Rotation.from_euler("x", -90, degrees=True)
        x_forward_to_y_forward = Rotation.from_euler("z", -90, degrees=True)
        mujoco_to_ardy = R_zup_to_yup * x_forward_to_y_forward

        def mujoco_joint_translation(body_name):
            body = name_to_body[body_name]
            if "pos" in body.attrib:
                result = np.array([100 * float(x) for x in body.get("pos").strip().split(" ")])
                return mujoco_to_ardy.apply(result)
            else:
                return np.zeros(3)

        parent_map = {
            map_bone_name(child.get("name")): map_bone_name(parent.get("name"))
            for parent in root.iter()
            if parent.get("name")
            for child in parent
            if child.get("name")
        }

        self.len_scale = 1.0
        self.root = Bone()
        self.root.id = 0
        self.root.name = map_bone_name(joint_names[0])
        self.root.offset = mujoco_joint_translation(self.root.name)
        self.root.offset_rot = np.eye(3)

        self.name2bone[self.root.name] = self.root
        self.bones.append(self.root)
        for i, joint in enumerate(joint_names[1:]):
            bone = Bone()
            bone.id = i + 1
            bone.name = map_bone_name(joint)
            bone.offset = mujoco_joint_translation(bone.name)
            bone.offset_rot = np.eye(3)
            self.bones.append(bone)
            self.name2bone[bone.name] = bone

        for bone in self.bones[1:]:
            parent_name = parent_map[bone.name]
            if parent_name in self.name2bone.keys():
                bone_p = self.name2bone[parent_name]
                bone_p.child.append(bone)
                bone.parent = bone_p

        self.forward_bvh(self.root)

    def add_joint(self, parent_name, bone_name, offset):
        new_bone = Bone()
        new_bone.name = bone_name
        new_bone.parent = next((x for x in self.bones if x.name == parent_name))
        new_bone.parent.child.append(new_bone)
        new_bone.offset = offset
        new_bone.offset_rot = np.eye(3)

        parent_index = next((i for i, x in enumerate(self.bones) if x.name == parent_name))
        self.bones.insert(parent_index + 1, new_bone)
        self.name2bone[bone_name] = new_bone


def load_bvh_animation(
    fname: str,
    skeleton: Skeleton,
    rot_order: Optional[str] = "native",
    backend: Optional[str] = "np",
    return_quat: Optional[bool] = False,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Given a bvh file for one motion sequence, and Skeleton object containing the neutral
    skeleton, return root translation and joint rotation matrics for the motion sequence.

    Args:
        fname (str): full path of the bvh file
        skeleton (Skeleton): Skeleton object initialized by the bvh file, containing the neutral skeleton
        rot_order (str): one of "native", "ZXY", "YXZ", ....
        return_quat (bool): if True, return joint rotation quaternions instead of rotation matrices
    Returns:
        torch.Tensor: [T, 3] root translations
        torch.Tensor: [T, J, 3, 3] joint rotation matrix
    """
    with open(fname) as f:
        mocap = Bvh(f.read(), backend=backend)

    # assume all joints are same ordering, load in with native ordering
    root_channels = mocap.joint_channels(skeleton.root.name)
    pos_channels = [channel for channel in root_channels if channel.endswith("position")]
    rot_channels = [channel for channel in root_channels if channel.endswith("rotation")]

    root_trans = np.array(mocap.frames_joint_channels(skeleton.root.name, pos_channels))

    if backend == "np":
        # NOTE: assumes rot channel ordering is the same for all joints
        joint_eulers = mocap.frames_joints_channels(skeleton.get_bones_names(), rot_channels)
        joint_eulers = np.deg2rad(joint_eulers)
    elif backend == "graph":
        joint_eulers = []
        for bone in skeleton.bones:
            bone_channels = mocap.joint_channels(bone.name)
            bone_rot_channels = [channel for channel in bone_channels if channel.endswith("rotation")]
            assert bone_rot_channels == rot_channels, "Rotation channel ordering is not consistent across joints!"
            # use native rotation order
            euler = np.deg2rad(np.array(mocap.frames_joint_channels(bone.name, rot_channels)))
            joint_eulers.append(euler)
        joint_eulers = np.stack(joint_eulers, axis=1)
    else:
        raise ValueError(f"Unknown backend for BVH loading: {backend}")

    if rot_order == "native":
        rot_order = ""
        for axis in rot_channels:
            rot_order += axis[0]
    else:
        # need to reorder dims
        ordered_joint_eulers = []
        for axis in rot_order:
            i = rot_channels.index(axis + "rotation")
            ordered_joint_eulers.append(joint_eulers[..., i])
        joint_eulers = np.stack(ordered_joint_eulers, axis=-1)

    rotations = Rotation.from_euler(rot_order, joint_eulers.reshape(-1, 3))
    if return_quat:
        joint_rots = rotations.as_quat(scalar_first=True).reshape(joint_eulers.shape[:-1] + (4,))
    else:
        joint_rots = rotations.as_matrix().reshape(joint_eulers.shape[:-1] + (3, 3))

    return root_trans, joint_rots


def load_csv_g1_animation(
    fname: str,
    skeleton: Skeleton,
    csv_format: str = "retargeted_from_gmr",
    xml_path: str = _DEFAULT_G1_XML,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Load a csv file that's in the format of mujoco qpos (32 joints, dof 29) return root
    translation and joint rotation matrics for the motion sequence.

    Args:
        fname (str): full path of the csv file
        skeleton (Skeleton): Skeleton object initialized by the bvh file, containing the neutral skeleton
        csv_format (str): format of the csv file ["retargeted_from_gmr", "retargeted_from_unitree"]
        xml_path (str): full path of the xml file for the skeleton
    Returns:
        torch.Tensor: [T, 3] root translations
        torch.Tensor: [T, J, 3, 3] joint rotation matrix
    """
    assert csv_format in ["retargeted_from_gmr", "retargeted_from_unitree"]
    import xml.etree.ElementTree as ET

    with open(fname) as f:
        csv_raw_data = f.readlines()
    if csv_format == "retargeted_from_gmr":
        csv_raw_data = np.array(
            [np.array([float(i) for i in row.split(",")[1:]]) for row in csv_raw_data[1:]]
        )  # drop the title row, and the frame column (first column)
        root_trans = csv_raw_data[:, :3]  # cm since it will be converted back to m outside this func
        xyz = csv_raw_data[:, 3:6]
        root_rot = Rotation.from_euler("xyz", xyz, degrees=True)

    elif csv_format == "retargeted_from_unitree":
        csv_raw_data = np.array([np.array([float(i) for i in row.split(",")]) for row in csv_raw_data])
        root_trans = (
            np.array(csv_raw_data[:, :3]) * 100
        )  # m -> cm since it will be converted back to m outside this func
        root_rot = Rotation.from_quat(csv_raw_data[:, 3:7], scalar_first=False)  # [x, y, z, w] is the order in csv file

    # there are two transformations from the mujoco coordinate to the ardy coordinate (both are right hand system):
    # 1) the mujoco is x facing, and ardy is y facing in the mujoco coordinate (or z facing in the ardy system).
    # 2) the source csv uses z up, and the ardy uses y up;
    R_zup_to_yup = Rotation.from_euler("x", -90, degrees=True)
    x_forward_to_y_forward = Rotation.from_euler("z", -90, degrees=True)
    mujoco_to_ardy = R_zup_to_yup * x_forward_to_y_forward

    # for root transform, we do the similar transformation for the root rotations.
    root_trans = mujoco_to_ardy.apply(root_trans)
    root_rot = mujoco_to_ardy * root_rot * mujoco_to_ardy.inv()

    # get the joint rotations from the csv file
    rotations = Rotation.identity(csv_raw_data.shape[0] * len(skeleton.bones))
    joint_rots = rotations.as_matrix().reshape(csv_raw_data.shape[0], len(skeleton.bones), 3, 3)
    joint_rots[:, 0] = root_rot.as_matrix()  # the rotation of pelvis

    # load the mujoco xml file for joint axis; find all joint, get name and axis
    tree = ET.parse(xml_path)
    root = tree.getroot()

    xml_classes = [x for x in tree.findall(".//default") if "class" in x.attrib]
    joint_axes = dict()
    for xml_class in xml_classes:
        j = xml_class.findall("joint")
        if j:
            joint_axes[xml_class.get("class")] = j[0].get("axis")

    parent_map = {child: parent for parent in root.iter() for child in parent}

    for joint_id_in_csv, joint in enumerate(root.find("worldbody").findall(".//joint")):  # skip the base joint
        # the order in the xml file is the order of the data shows up; note that in the xml file, the joint ends with
        # "_joint", but the skeleton.get_bones_names() joints end with "_skel"
        assert joint.get("name").endswith("_joint")
        joint_name_in_skeleton = joint.get("name").replace("_joint", "_skel")
        assert joint_name_in_skeleton in skeleton.get_bones_names()
        axis_values = [float(x) for x in (joint.get("axis") or joint_axes[joint.get("class")]).split(" ")]
        assert sum(axis_values) == 1 and all(x in [0, 1] for x in axis_values), (
            f"Invalid axis: {axis_values}: one and only one of the axis values should be 1; others should be 0s."
        )

        # the mapped axis in the ardy's g1 skeleton space is calculated as bones_axis = mujoco_to_ardy.apply(axis_values)
        # [1, 0, 0] -> [0, 0, 1]; [0, 1, 0] -> [1, 0, 0]; [0, 0, 1] -> [0, 1, 0]
        axis_in_ardy = ["x", "y", "z"][np.argmax(axis_values)]
        if csv_format == "retargeted_from_gmr":
            joint_dof = csv_raw_data[:, joint_id_in_csv + 6] * np.pi / 180
        elif csv_format == "retargeted_from_unitree":
            joint_dof = csv_raw_data[:, joint_id_in_csv + 7]
        else:
            raise ValueError(f"Unknown csv format: {csv_format}")

        # Single-axis sequence needs angles of shape (N, 1) under SciPy's array-API
        # from_euler, which reads the last axis as the axis count (not a bare (N,)).
        joint_rot = Rotation.from_euler(axis_in_ardy, joint_dof[:, None], degrees=False)

        body = parent_map[joint]
        if "quat" in body.attrib:
            joint_extra_rot = Rotation.from_quat(
                [float(x) for x in body.get("quat").strip().split(" ")],
                scalar_first=True,
            )
            joint_rot = joint_extra_rot * joint_rot

        joint_rots[:, skeleton.get_bones_names().index(joint_name_in_skeleton)] = (
            mujoco_to_ardy * joint_rot * mujoco_to_ardy.inv()
        ).as_matrix()

    return root_trans, joint_rots