"""Generate a rope as an articulated chain, then convert it to USD.

Isaac Lab's rigid-body pipeline has no true deformable rope, and the particle/cloth path is a
different simulation stack from everything else in this suite. The standard substitute is a
CHAIN: many short rigid links joined end to end, each pair free to bend. With links this short
the chain drapes, coils and pulls taut like a rope, and -- importantly here -- it can only be
straightened by pulling both ends apart, which is the task.

Each joint is a UNIVERSAL joint, built as two perpendicular revolute joints separated by a
massless spacer link. A single revolute per segment would give a rope that bends in one plane
only, which reads as a hinge chain rather than a rope.

    python scripts/make_rope_urdf.py --links 14 --seg 0.022
"""
import argparse
import os

TEMPLATE_HEAD = """<?xml version="1.0"?>
<robot name="rope">
"""


def link(name, length, radius, mass, rgba):
    if mass <= 0.0:                       # massless spacer: geometry-free, just a frame
        return f"""  <link name="{name}"/>
"""
    ixx = mass*(3*radius**2+length**2)/12.0
    izz = mass*radius**2/2.0
    return f"""  <link name="{name}">
    <visual>
      <origin xyz="{length/2:.5f} 0 0" rpy="0 1.5708 0"/>
      <geometry><cylinder radius="{radius:.5f}" length="{length:.5f}"/></geometry>
      <material name="rope"><color rgba="{rgba}"/></material>
    </visual>
    <collision>
      <origin xyz="{length/2:.5f} 0 0" rpy="0 1.5708 0"/>
      <geometry><cylinder radius="{radius:.5f}" length="{length:.5f}"/></geometry>
    </collision>
    <inertial>
      <origin xyz="{length/2:.5f} 0 0"/>
      <mass value="{mass:.5f}"/>
      <inertia ixx="{ixx:.8f}" ixy="0" ixz="0" iyy="{ixx:.8f}" iyz="0" izz="{izz:.8f}"/>
    </inertial>
  </link>
"""


def joint(name, parent, child, axis, xyz, limit):
    return f"""  <joint name="{name}" type="revolute">
    <parent link="{parent}"/>
    <child link="{child}"/>
    <origin xyz="{xyz}" rpy="0 0 0"/>
    <axis xyz="{axis}"/>
    <limit lower="{-limit:.4f}" upper="{limit:.4f}" effort="4" velocity="12"/>
    <dynamics damping="0.02" friction="0.001"/>
  </joint>
"""


def build(n_links, seg, radius, mass_per_link, bend_limit, rgba):
    out = [TEMPLATE_HEAD]
    out.append(link("seg0", seg, radius, mass_per_link, rgba))
    for i in range(1, n_links):
        out.append(link(f"spacer{i}", 0, 0, 0.0, rgba))
        out.append(link(f"seg{i}", seg, radius, mass_per_link, rgba))
        # two perpendicular hinges at the same point == a universal joint
        out.append(joint(f"bend{i}_y", f"seg{i-1}", f"spacer{i}", "0 1 0",
                         f"{seg:.5f} 0 0", bend_limit))
        out.append(joint(f"bend{i}_z", f"spacer{i}", f"seg{i}", "0 0 1", "0 0 0", bend_limit))
    out.append("</robot>\n")
    return "".join(out)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--links", type=int, default=14)
    ap.add_argument("--seg", type=float, default=0.022, help="length of one link (m)")
    ap.add_argument("--radius", type=float, default=0.006)
    ap.add_argument("--mass", type=float, default=0.006, help="mass per link (kg)")
    ap.add_argument("--bend", type=float, default=1.05, help="bend limit per hinge (rad)")
    ap.add_argument("--rgba", default="0.85 0.72 0.45 1")
    ap.add_argument("--out", default=os.environ.get("ROBOTWIN_SRC",
                    "/home/yu/internship_yu/robotwin_object/objects")+"/custom_rope/rope")
    a = ap.parse_args()

    os.makedirs(a.out, exist_ok=True)
    path = os.path.join(a.out, "mobility.urdf")
    with open(path, "w") as fh:
        fh.write(build(a.links, a.seg, a.radius, a.mass, a.bend, a.rgba))
    with open(os.path.join(a.out, "semantics.txt"), "w") as fh:
        fh.write("\n".join(f"seg{i} chain rope_segment" for i in range(a.links))+"\n")
    total = a.links*a.seg
    print(f"[rope] {path}")
    print(f"[rope] {a.links} links x {a.seg*100:.1f} cm = {total*100:.1f} cm rope, "
          f"{a.links*a.mass*1000:.0f} g, {2*(a.links-1)} joints")


if __name__ == "__main__":
    main()