# tensegrity.py : utility classes for helping define tensegrity structures in MuJoCo XML specs

import math
import xml.etree.ElementTree as ET
import numpy as np

#================================================================
# numerical utility functions

def vec_to_attr(v):
    return "%f %f %f" % (v[0], v[1], v[2])

def distance(pt1, pt2):
    diff = pt2 - pt1
    return math.sqrt(np.dot(diff, diff))

def deg_to_rad(angle):
    return math.pi * (angle/180.0)

#================================================================
class Strut:
    """Represent a tensegrity strut extending between two points represented in
    world coordinates as 3-element numpy vectors."""
    _next_id = 0

    def __init__(self, end0, end1):
        self.end0 = end0.copy()
        self.end1 = end1.copy()
        self.idnum = Strut._next_id
        Strut._next_id += 1
        self.length = distance(end0, end1)
        self.body_name = "body%03d" % self.idnum
        self.geom_name = "strut%03d" % self.idnum
        self.end0_name = "s_%03d_0" % (self.idnum)
        self.end1_name = "s_%03d_1" % (self.idnum)

    def translate(self, vec):
        """Move the end points by the given vector."""
        self.end0 += vec
        self.end1 += vec

    def as_element(self):
        """Create and return an Element object representing a tensegrity strut."""
        pt1    = vec_to_attr(self.end0)
        local1 = vec_to_attr(self.end1 - self.end0)

        body = ET.Element('body', attrib={'pos': pt1, 'name': self.body_name})
        body.append(ET.Element('joint', attrib={'type':'free'}))
        body.append(ET.Element('geom', attrib={'type': 'cylinder',
                                                   'name': self.geom_name,
                                                   'size': "0.005 1.0",
                                                   'fromto': "0.0 0.0 0.0 " + local1 }))
        body.append(ET.Element('site', attrib={'name' : self.end0_name,
                                               'pos'  : "0 0 0",
                                               'size' : "0.005" }))
        body.append(ET.Element('site', attrib={'name' : self.end1_name,
                                               'pos'  : local1,
                                               'size' : "0.005" }))
        return body

#================================================================
class Tendon:
    def __init__(self, strut1, end1_num, strut2, end2_num, ratio=0.95):
        """Represent a tendon connecting two given strut ends."""
        self.name = "t_%03d_%d_%03d_%d" % (strut1.idnum, end1_num, strut2.idnum, end2_num)

        # endpoint  in world coordinates
        self.end0 = strut1.end0.copy() if end1_num==0 else strut1.end1.copy()
        self.end1 = strut2.end0.copy() if end2_num==0 else strut2.end1.copy()
        self.distance = distance(self.end0, self.end1)
        self.length = ratio * self.distance

        # endpoint site names
        self.end0_name = strut1.end0_name if end1_num==0 else strut1.end1_name
        self.end1_name = strut2.end0_name if end2_num==0 else strut2.end1_name

    def translate(self, vec):
        """Move the end points by the given vector."""
        self.end0 += vec
        self.end1 += vec

    def as_element(self):
        """Create and return an Element object representing a tensegrity tendon."""
        tendon = ET.Element('tendon')
        spatial = ET.Element('spatial',
                             {'limited': 'true',
                              'range' : "0.0 %f" % self.length,
                              'width' : "0.001",
                              'name'  : self.name
                              })
        spatial.append(ET.Element('site', {'site': self.end0_name}))
        spatial.append(ET.Element('site', {'site': self.end1_name}))
        tendon.append(spatial)
        return tendon
#================================================================