import xml.etree.ElementTree as ET
import math

def round_axis(x, y, z):
    # Normalize vector
    mag = math.sqrt(x**2 + y**2 + z**2)
    if mag == 0:
        return (0, 0, 0)
    x, y, z = x/mag, y/mag, z/mag

    # Compare to cardinal directions
    directions = {
        (1, 0, 0): 'X+',
        (-1, 0, 0): 'X-',
        (0, 1, 0): 'Y+',
        (0, -1, 0): 'Y-',
        (0, 0, 1): 'Z+',
        (0, 0, -1): 'Z-'
    }

    best = max(directions.keys(), key=lambda d: x*d[0] + y*d[1] + z*d[2])
    return best

def process_urdf(file_path, output_path):
    tree = ET.parse(file_path)
    root = tree.getroot()

    for joint in root.findall('joint'):
        axis_elem = joint.find('axis')
        if axis_elem is not None:
            axis_str = axis_elem.attrib.get('xyz', '')
            try:
                x, y, z = map(float, axis_str.strip().split())
                rx, ry, rz = round_axis(x, y, z)
                axis_elem.set('xyz', f"{rx} {ry} {rz}")
            except:
                continue

    tree.write(output_path)
    print(f"Rounded axes written to {output_path}")

# Example usage
process_urdf("./urdf/sample.urdf", "./urdf/output.urdf")