import sys
import numpy as np

import piper.anatomyDB as anatomyDB

import Sofa
import SofaPython.sml
import Anatomy.Tools
import SofaPython.units as units

printLog=True

def addAnatomicalLandmarkToFrameFactory(model, ignoreUnknownLandmark=False):
    meshCache = dict()
    pos=None
    ff = anatomyDB.FrameFactory.instance()
    for solid in model.solids.values():
        for mesh in solid.mesh:
            for id,group in mesh.group.iteritems():
                if not anatomyDB.exists(id):
                    continue
                if "landmark" in group.tags:
                    # TODO check mesh.format
                    m=None
                    if ignoreUnknownLandmark and not mesh.source in meshCache:
                        meshCache[mesh.source] = SofaPython.MeshLoader.loadOBJ(mesh.source)
                    m = meshCache[mesh.source]
                    if "sphere" in group.tags or "barycenter" in group.tags:
                        landmarkVertices=[]
                        for i in group.index:
                            landmarkVertices.append(m.vertices[i])
                        if "sphere" in group.tags:
                            pos = Anatomy.Tools.fitSphere(landmarkVertices,precision=units.length_from_SI(1E-5))[0]
                        elif "barycenter" in group.tags:
                            pos = Anatomy.Tools.barycenter(landmarkVertices)
                    elif len(group.index) == 1:
                        pos = m.vertices[group.index[0]]
                    else:
                        Sofa.msg_warning("Anatomy.sml", "invalid landmark id: {0}, size: {1}".format(id, len(group.index)))
                        continue
                    ff.landmark().add(id, np.array(pos))
                    if printLog:
                        Sofa.msg_info("Anatomy.sml","landmark {0} in solid {1} - position: {2}".format(id,solid.name,pos))

class ISBJointFactory:
    """ This class adds to a sml model JointGeneric corresponding to anatomical joints defined by ISB.
    The joints are created on the basis of anatomical frames computed from the landmarks available in the model.
    """

    def __init__(self):
        # to specify dof of the joints
        self.jointDof=dict()
        # default "ball and socket" joint dof
        self.jointDofDefault = [SofaPython.sml.Model.Dof(dof="rx"), SofaPython.sml.Model.Dof(dof="ry"), SofaPython.sml.Model.Dof(dof="rz")]
        # some basic defaults joint name must match anatomyDB name
        self.jointDof["Left_knee_joint"] = [SofaPython.sml.Model.Dof(dof="rz")]
        self.jointDof["Right_knee_joint"] = [SofaPython.sml.Model.Dof(dof="rz")]
        self.jointDof["Left_humeroulnar_joint"] = [SofaPython.sml.Model.Dof(dof="rz")]
        self.jointDof["Right_humeroulnar_joint"] = [SofaPython.sml.Model.Dof(dof="rz")]

    def process(self, model):
        """ Process the model to add JointGeneric to the sml model, as defined by ISB
        """
        # build a dictionnary to access solids by their reference name
        solidByRefName=dict()

        for solid in model.solids.values():
            refName = anatomyDB.getReferenceNameNoThrow(solid.id)
            if len(refName)>0:
                solidByRefName[refName]=solid
            else:
                solidByRefName[solid.id]=solid
                Sofa.msg_warning("Anatomy.sml.joint", "Unknown entity {0}, not using its reference name".format(solid.id))

        # let's gather all available known landmarks
        addAnatomicalLandmarkToFrameFactory(model, ignoreUnknownLandmark=True)
        # shortcut to the frame factory
        ff = anatomyDB.FrameFactory.instance()

        for jointName in anatomyDB.getSubClassOfFromBibliographyList("Joint", "ISB_JCS"):
            smlJoint = SofaPython.sml.Model.JointGeneric()
            smlJoint.id=jointName
            smlJoint.name=jointName
            frames = anatomyDB.getSubPartOfList(jointName,"Frame") # TODO check also for "ISB_JCS"
            if not len(frames) == 2:
                Sofa.msg_warning("Anatomy.sml.joint", "In joint {0} incorrect number of frames {1} - ignored".format(jointName, frames))
                continue
            smlJoint.solids=[] # initialized with [None,None]
            for frame in frames:
                # get bone to which the frame belongs, look for group of bones too
                bone = anatomyDB.getPartOfSubClassList(frame,"Bone")
                while len(bone) > 0 :
                    if bone[0] in solidByRefName:
                        # if found, append the solid to the smlJoint
                        smlJoint.solids.append(solidByRefName[bone[0]])
                        break
                    else:
                        bone = anatomyDB.getPartOfSubClassList(bone[0],"Bone")
            if not len(smlJoint.solids) == 2:
                Sofa.msg_warning("Anatomy.sml.joint", "In joint {0} articulated solids not found - ignored".format(jointName))
                continue
            # now compute offset
            smlJoint.offsets=[] # initialized with [None,None]
            for frame in frames:
                try:
                    vec = ff.computeFrame(frame)
                except:
                    Sofa.msg_warning("Anatomy.sml.joint", "In joint {0} frame {1} cannot be computed - error {2}".format(jointName,frame, sys.exc_value))
                else:
                    smlOffset = SofaPython.sml.Model.Offset()
                    smlOffset.name = frame
                    smlOffset.value = vec.toVec().flatten().tolist()
                    smlJoint.offsets.append(smlOffset)

            if not len(smlJoint.offsets) == 2:
                continue
            # and lastly the dof
            if jointName in self.jointDof:
                smlJoint.dofs = self.jointDof[jointName]
            else:
                smlJoint.dofs = self.jointDofDefault
            # we are done, let's add the joint to the model
            model.genericJoints[smlJoint.id]=smlJoint
            if printLog:
                Sofa.msg_info("Anatomy.sml.joint","JointGeneric {0}".format(smlJoint.id))
