import os
from math import cos
from math import sin

import Sofa.Core
from splib.numerics import Quat, Vec3

from sofacontrol import measurement_models

path = os.path.dirname(os.path.abspath(__file__))


class TemplateEnvironment:
    def __init__(self, name='Template', rayleighMass=0.1, rayleighStiffness=0.1, dt=0.01):
        self.name = name
        self.robot = Sofa.Core.Node(name)
        # set-up solvers
        self.robot.addObject('EulerImplicitSolver', name='odesolver', firstOrder="0", rayleighMass=str(rayleighMass),
                             rayleighStiffness=str(rayleighStiffness))
        self.robot.addObject('SparseLDLSolver', name='preconditioner')
        self.robot.addObject('GenericConstraintCorrection', solverName="preconditioner")
        self.actuator_list = []
        self.nb_nodes = None
        self.gravity = [0., -9810., 0.]  # default
        self.dt = dt

    def get_measurement_model(self, nodes=None, pos=True, vel=True):
        if nodes is None:
            return measurement_models.linearModel(range(self.nb_nodes), self.nb_nodes, pos=pos, vel=vel)
        else:
            return measurement_models.linearModel(nodes, self.nb_nodes, pos=pos, vel=vel)


class Trunk(TemplateEnvironment):
    def __init__(self, name='Trunk', all_cables=True):
        super(Trunk, self).__init__(name=name)

        self.nb_nodes = 709

        self.gravity = [0., 0., 9810.]

        self.robot.min_force = [0.] * 8  # Without premultiplication with dt

        self.robot.addObject('MeshVTKLoader', name='loader', filename=path + '/mesh/trunk.vtk')
        self.robot.addObject('TetrahedronSetTopologyContainer', src='@loader', name='container')
        self.robot.addObject('TetrahedronSetTopologyModifier')
        self.robot.addObject('TetrahedronSetTopologyAlgorithms')
        self.robot.addObject('TetrahedronSetGeometryAlgorithms')
        # Option 1:
        self.robot.addObject('MechanicalObject', name='tetras', template='Vec3d', showIndices='false',
                             showIndicesScale='4e-5')
        # Option 2: Equivalent to option 1 (we believe)
        # self.robot.addObject('MechanicalObject', src='@loader')

        # Gives a mass to the model
        self.robot.addObject('UniformMass', totalMass=0.042)
        # Add a TetrahedronFEMForceField componant which implement an elastic material model solved using the Finite
        # Element Method on tetrahedrons.
        self.robot.addObject('TetrahedronFEMForceField', template='Vec3d', name='FEM', method='large',
                             poissonRatio=0.45,
                             youngModulus=450)
        # Fix the base of the trunk by adding constraints in a region of interest (ROI)
        self.robot.addObject('BoxROI', name='boxROI', box=[[-20, -20, 0], [20, 20, 20]], drawBoxes=False)
        self.robot.addObject('RestShapeSpringsForceField', points='@boxROI.indices', stiffness='1e12')

        ##########################################
        # Cable                                  #
        ##########################################
        actuator_names = ''

        length1 = 10.
        length2 = 2.
        lengthTrunk = 195.

        pullPoint = [[0., length1, 0.], [-length1, 0., 0.], [0., -length1, 0.], [length1, 0., 0.]]
        direction = Vec3(0., length2 - length1, lengthTrunk)
        direction.normalize()
        nbCables = 4
        actuators = self.robot.addChild('actuators')
        for i in range(0, nbCables):
            childname = 'cableL' + str(i)

            theta = 1.57 * i
            q = Quat(0., 0., sin(theta / 2.), cos(theta / 2.))
            position = [[0., 0., 0.]] * 20
            for k in range(0, 20, 2):
                v = Vec3(direction[0], direction[1] * 17.5 * (k / 2) + length1, direction[2] * 17.5 * (k / 2) + 21)
                position[k] = v.rotateFromQuat(q)
                v = Vec3(direction[0], direction[1] * 17.5 * (k / 2) + length1, direction[2] * 17.5 * (k / 2) + 27)
                position[k + 1] = v.rotateFromQuat(q)
            cableL = actuators.addChild(childname)
            cableL.addObject('MechanicalObject', name='meca',
                             position=pullPoint[i] + [pos.toList() for pos in position])
            cableL.addObject('CableConstraint', template='Vec3d', name="cable",
                             hasPullPoint="0",
                             indices=list(range(21)),
                             maxPositiveDisp='70',
                             maxDispVariation="1",
                             valueType='force',
                             minForce=self.robot.min_force[i] * self.robot.dt.value)

            cableL.addObject('BarycentricMapping', name='mapping', mapForces='false', mapMasses='false')
            actuator_names += childname + '/cable,'
            self.actuator_list.append(cableL.cable)

        if all_cables:
            for i in range(0, nbCables):
                childname = 'cableS' + str(i)
                theta = 1.57 * i
                q = Quat(0., 0., sin(theta / 2.), cos(theta / 2.))

                position = [[0., 0., 0.]] * 10
                for k in range(0, 9, 2):
                    v = Vec3(direction[0], direction[1] * 17.5 * (k / 2) + length1, direction[2] * 17.5 * (k / 2) + 21)
                    position[k] = v.rotateFromQuat(q)
                    v = Vec3(direction[0], direction[1] * 17.5 * (k / 2) + length1, direction[2] * 17.5 * (k / 2) + 27)
                    position[k + 1] = v.rotateFromQuat(q)

                cableS = actuators.addChild(childname)
                cableS.addObject('MechanicalObject', name='meca',
                                 position=pullPoint[i] + [pos.toList() for pos in position])
                cableS.addObject('CableConstraint', template='Vec3d', name="cable",
                                 hasPullPoint="0",
                                 indices=list(range(10)),
                                 maxPositiveDisp='40',
                                 maxDispVariation="1",
                                 valueType='force',
                                 minForce=self.robot.min_force[i + 4] * self.robot.dt.value)
                cableS.addObject('BarycentricMapping', name='mapping', mapForces='false', mapMasses='false')
                actuator_names += childname + '/cable,'
                self.actuator_list.append(cableS.cable)

        self.robot.actuator_list = self.actuator_list
        ##########################################
        # Visualization                          #
        ##########################################
        trunkVisu = self.robot.addChild('VisualModel')
        trunkVisu.addObject('MeshSTLLoader', filename=path + "/mesh/trunk.stl")
        trunkVisu.addObject('OglModel', template='Vec3d', color=[1., 1., 1., 0.8])
        trunkVisu.addObject('BarycentricMapping')


class Trunk4Cables(Trunk):
    def __init__(self, name='Trunk4Cables'):
        super(Trunk4Cables, self).__init__(name=name, all_cables=False)
        self.robot.min_force = [0, 0, 0, 0]  # Without premultiplication with dt


class Finger(TemplateEnvironment):
    def __init__(self, name='Finger'):
        super(Finger, self).__init__(name=name)

        self.nb_nodes = 158

        self.robot.min_force = [0.]  # Without premultiplication with dt

        self.robot.addObject('MeshVTKLoader', name='loader', filename=path + '/mesh/finger.vtk')
        self.robot.addObject('TetrahedronSetTopologyContainer', src='@loader', name='container')
        self.robot.addObject('TetrahedronSetTopologyModifier')
        self.robot.addObject('TetrahedronSetTopologyAlgorithms')
        self.robot.addObject('TetrahedronSetGeometryAlgorithms')
        self.robot.addObject('MechanicalObject', name='tetras', template='Vec3d', showIndices='false',
                             showIndicesScale='4e-5')

        self.robot.addObject('UniformMass', totalMass=0.075)

        # Add a TetrahedronFEMForceField componant which implement an elastic material model solved using the Finite Element Method on tetrahedrons.
        self.robot.addObject('TetrahedronFEMForceField', template='Vec3d', name='FEM', method='large',
                             poissonRatio=0.45,
                             youngModulus=600)

        # Fix the base of the trunk by adding constraints in a region of interest (ROI)
        self.robot.addObject('BoxROI', name='boxROI', box=[[-15, 0, 0], [5, 10, 15]], drawBoxes=False)
        self.robot.addObject('RestShapeSpringsForceField', points='@boxROI.indices', stiffness='1e12')

        ##########################################
        # Cable                                 #
        ##########################################

        #  This creates a new node in the scene. This node is appended to the finger's node.
        actuators = self.robot.addChild('actuators')

        cable = actuators.addChild('cable')

        #  This create a MechanicalObject, a componant holding the degree of freedom of our
        # mechanical modelling. In the case of a cable it is a set of positions specifying
        #  the points where the cable is passing by.
        cable.addObject('MechanicalObject', name='meca',
                        position=(
                                "-17.5 12.5 2.5 " +
                                "-32.5 12.5 2.5 " +
                                "-47.5 12.5 2.5 " +
                                "-62.5 12.5 2.5 " +
                                "-77.5 12.5 2.5 " +

                                "-83.5 12.5 4.5 " +
                                "-85.5 12.5 6.5 " +
                                "-85.5 12.5 8.5 " +
                                "-83.5 12.5 10.5 " +

                                "-77.5 12.5 12.5 " +
                                "-62.5 12.5 12.5 " +
                                "-47.5 12.5 12.5 " +
                                "-32.5 12.5 12.5 " +
                                "-17.5 12.5 12.5 "))

        # Create a CableConstraint object with a name.
        # the indices are referring to the MechanicalObject's positions.
        # The last indice is where the pullPoint is connected.
        cable.addObject('CableConstraint', name="cable",
                        indices=list(range(14)),
                        pullPoint="0.0 12.5 2.5", valueType='force',
                        minForce=self.robot.min_force[0] * self.robot.dt.value)
        # This create a BarycentricMapping. A BarycentricMapping is a key element as it will create a bi-directional link
        #  between the cable's DoFs and the finger's ones so that movements of the cable's DoFs will be mapped
        #  to the finger and vice-versa;
        cable.addObject('BarycentricMapping', name='mapping', mapForces='false', mapMasses='false')

        self.actuator_list.append(cable.cable)

        self.robot.actuator_list = self.actuator_list
        ##########################################
        # Visualization                          #
        ##########################################
        # In Sofa, visualization is handled by adding a rendering model.
        #  Create an empty child node to store this rendering model.
        fingerVisu = self.robot.addChild('VisualModel')

        # Add to this empty node a rendering model made of triangles and loaded from an stl file.
        fingerVisu.addObject('MeshSTLLoader', filename=path + "/mesh/finger.stl")
        fingerVisu.addObject('OglModel', template='Vec3d', color=[1., 1., 1., 0.8])

        # Add a BarycentricMapping to deform rendering model in way that follow the ones of the parent mechanical model.
        fingerVisu.addObject('BarycentricMapping')


class Diamond(TemplateEnvironment):
    def __init__(self, name='Diamond', totalMass=0.5, poissonRatio=0.45, youngModulus=450, rayleighMass=0.1, rayleighStiffness=0.1, dt=0.01):
        super(Diamond, self).__init__(name=name, rayleighMass=rayleighMass, rayleighStiffness=rayleighStiffness, dt=dt)

        self.nb_nodes = 1628
        self.gravity = [0., 0., -9810.]

        rotation = [90, 0.0, 0.0]
        translation = [0.0, 0.0, 35]

        self.robot.min_force = [0, 0, 0, 0]  # Without premultiplication with dt

        self.robot.addObject('MeshVTKLoader', name='loader', filename=path + "/mesh/diamond.vtu", rotation=rotation,
                             translation=translation)
        self.robot.addObject('TetrahedronSetTopologyContainer', src='@loader', name='container')
        self.robot.addObject('TetrahedronSetTopologyModifier')
        self.robot.addObject('TetrahedronSetTopologyAlgorithms')
        self.robot.addObject('TetrahedronSetGeometryAlgorithms')
        self.robot.addObject('MechanicalObject', template='Vec3d', name='tetras', showIndices='false',
                             showIndicesScale='4e-5')
        self.robot.addObject('UniformMass', totalMass=totalMass, name='mass')
        self.robot.addObject('TetrahedronFEMForceField', template='Vec3d',
                             method='large', name='forcefield',
                             poissonRatio=poissonRatio, youngModulus=youngModulus)

        # Fix the base of the trunk by adding constraints in a region of interest (ROI)
        self.robot.addObject('BoxROI', name='boxROI', box=[-15, -15, -40, 15, 15, 10], drawBoxes=True)
        self.robot.addObject('RestShapeSpringsForceField', points='@boxROI.indices', stiffness='1e12')

        ##########################################
        # Cable                                 #
        ##########################################
        self.actuatorsParam = [
            {'withName': 'A',
             'withCableGeometry': [[0, 97, 45]],
             'withAPullPointLocation': [0, 10, 30]
             },
            {'withName': 'B',
             'withCableGeometry': [[-97, 0, 45]],
             'withAPullPointLocation': [-10, 0, 30]
             },
            {'withName': 'C',
             'withCableGeometry': [[0, -97, 45]],
             'withAPullPointLocation': [0, -10, 30]
             },
            {'withName': 'D',
             'withCableGeometry': [[97, 0, 45]],
             'withAPullPointLocation': [10, 0, 30]
             }
        ]

        actuators = self.robot.addChild('actuators')

        for i in range(len(self.actuatorsParam)):
            cable = actuators.addChild(self.actuatorsParam[i]['withName'])
            cable.addObject('MechanicalObject', position=self.actuatorsParam[i]['withCableGeometry'])
            cable.addObject('CableConstraint',
                            name='cable',
                            indices=list(range(len(self.actuatorsParam[i]['withCableGeometry']))),
                            pullPoint=self.actuatorsParam[i]['withAPullPointLocation'],
                            valueType='force',
                            hasPullPoint=True,
                            minForce=self.robot.min_force[i] * self.robot.dt.value
                            )

            cable.addObject('BarycentricMapping', name="Mapping", mapForces=False, mapMasses=False)
            self.actuator_list.append(cable.cable)

        self.robot.actuator_list = self.actuator_list

        ##########################################
        # Visualization                          #
        ##########################################
        diamondVisu = self.robot.addChild('VisualModel')
        diamondVisu.addObject('MeshSTLLoader', filename=path + "/mesh/diamond.stl")
        diamondVisu.addObject('OglModel', template='Vec3d', color=[0.7, 0.7, 0.7, 0.7], updateNormals=False)
        diamondVisu.addObject('BarycentricMapping')

[((10, 23, 10, 48), 'os.path.abspath', 'os.path.abspath', ({(10, 39, 10, 47): '__file__'}, {}), '(__file__)', False, 'import os\n'), ((76, 20, 76, 60), 'splib.numerics.Vec3', 'Vec3', ({(76, 25, 76, 27): '0.0', (76, 29, 76, 46): 'length2 - length1', (76, 48, 76, 59): 'lengthTrunk'}, {}), '(0.0, length2 - length1, lengthTrunk)', False, 'from splib.numerics import Quat, Vec3\n'), ((31, 19, 31, 89), 'sofacontrol.measurement_models.linearModel', 'measurement_models.linearModel', (), '', False, 'from sofacontrol import measurement_models\n'), ((84, 29, 84, 44), 'math.sin', 'sin', ({(84, 33, 84, 43): 'theta / 2.0'}, {}), '(theta / 2.0)', False, 'from math import sin\n'), ((84, 46, 84, 61), 'math.cos', 'cos', ({(84, 50, 84, 60): 'theta / 2.0'}, {}), '(theta / 2.0)', False, 'from math import cos\n'), ((87, 20, 87, 115), 'splib.numerics.Vec3', 'Vec3', ({(87, 25, 87, 37): 'direction[0]', (87, 39, 87, 78): 'direction[1] * 17.5 * (k / 2) + length1', (87, 80, 87, 114): 'direction[2] * 17.5 * (k / 2) + 21'}, {}), '(direction[0], direction[1] * 17.5 * (k / 2) + length1, direction[2] * \n    17.5 * (k / 2) + 21)', False, 'from splib.numerics import Quat, Vec3\n'), ((89, 20, 89, 115), 'splib.numerics.Vec3', 'Vec3', ({(89, 25, 89, 37): 'direction[0]', (89, 39, 89, 78): 'direction[1] * 17.5 * (k / 2) + length1', (89, 80, 89, 114): 'direction[2] * 17.5 * (k / 2) + 27'}, {}), '(direction[0], direction[1] * 17.5 * (k / 2) + length1, direction[2] * \n    17.5 * (k / 2) + 27)', False, 'from splib.numerics import Quat, Vec3\n'), ((110, 33, 110, 48), 'math.sin', 'sin', ({(110, 37, 110, 47): 'theta / 2.0'}, {}), '(theta / 2.0)', False, 'from math import sin\n'), ((110, 50, 110, 65), 'math.cos', 'cos', ({(110, 54, 110, 64): 'theta / 2.0'}, {}), '(theta / 2.0)', False, 'from math import cos\n'), ((114, 24, 114, 119), 'splib.numerics.Vec3', 'Vec3', ({(114, 29, 114, 41): 'direction[0]', (114, 43, 114, 82): 'direction[1] * 17.5 * (k / 2) + length1', (114, 84, 114, 118): 'direction[2] * 17.5 * (k / 2) + 21'}, {}), '(direction[0], direction[1] * 17.5 * (k / 2) + length1, direction[2] * \n    17.5 * (k / 2) + 21)', False, 'from splib.numerics import Quat, Vec3\n'), ((116, 24, 116, 119), 'splib.numerics.Vec3', 'Vec3', ({(116, 29, 116, 41): 'direction[0]', (116, 43, 116, 82): 'direction[1] * 17.5 * (k / 2) + length1', (116, 84, 116, 118): 'direction[2] * 17.5 * (k / 2) + 27'}, {}), '(direction[0], direction[1] * 17.5 * (k / 2) + length1, direction[2] * \n    17.5 * (k / 2) + 27)', False, 'from splib.numerics import Quat, Vec3\n')]