import math
import sys

import Sofa

from SofaPython import Quaternion
import Compliant.Frame
import Compliant.StructuralAPI
import Anatomy.Constraint

def setShowDofs(dofs, scale=0.5, mode=0):
    dofs.showObject = True
    dofs.showObjectScale = scale
    dofs.drawMode = mode

class Point:
    pass

def createScene(rootNode):
    rootNode.createObject('RequiredPlugin', pluginName="Compliant")
    rootNode.createObject('RequiredPlugin', pluginName="Anatomy")
    
    rootNode.createObject('CompliantAttachButtonSetting' )
    rootNode.createObject('CompliantImplicitSolver', name='odesolver', stabilization="true")
    
    rootNode.createObject('SequentialSolver', name='numsolver', iterations=250, precision=1e-14, iterateOnBilaterals=True)
    #rootNode.createObject('LDLTSolver')
    rootNode.createObject('LDLTResponse', name='response')

    rootNode.gravity=[0,0,0]

    #
    # absolute translation constraint for a point
    #
    pointTranslationNode = rootNode.createChild("PointTranslation")
    pointI = Point()
    pointI.node = pointTranslationNode
    pointI.dofs = pointTranslationNode.createObject("MechanicalObject", template="Vec3", position="0 7 0")
    setShowDofs(pointI.dofs, scale=0.1, mode=1)
    constraintI = Anatomy.Constraint.PointTranslation("point_I", pointI, [1, 1, 0], progressive=True)
    constraintI.progressiveTarget.step = 0.005
    constraintI.setTarget([1, 7])

    #
    # rigid constraint between two rigids
    #
    
    rigidRigidTransformationNode = rootNode.createChild("RigidRigidTransformation")
    
    nodeG = rigidRigidTransformationNode.createChild("G")
    G = Compliant.Frame.Frame([0,6,0,0,0,0,1])
    G.dofs = G.insert(nodeG, name="dofs_green")
    nodeG.createObject('RigidMass', name = 'mass', mass = 1, inertia="1 1 1")
    setShowDofs(G.dofs,.5,1)
    nodeG.createObject('FixedConstraint')
    
    nodeH = rigidRigidTransformationNode.createChild("H")
    H = Compliant.Frame.Frame([1,6,0,0,0,0,1])
    H.dofs = H.insert(nodeH, name="dofs_red")
    nodeH.createObject('RigidMass', name = 'mass', mass = 1, inertia="1 1 1")
    setShowDofs(H.dofs,.5,2)
    
    constraintGToH = Anatomy.Constraint.RigidRigidTransformation("G_to_H", G, H, [0,0,0,1,1,1])
    constraintGToH.setTarget([0, math.radians(45), math.radians(15)])
    

    #
    # constraint between two rigids
    #

    rigidTransformationNode = rootNode.createChild("RigidTransformation")

    nodeA = rigidTransformationNode.createChild("A")
    A = Compliant.Frame.Frame([0,4,0,0,0,0,1])
    dofsA = A.insert(nodeA, name="dofs_green")
    nodeA.createObject('RigidMass', name = 'mass', mass = 1, inertia="1 1 1")
    setShowDofs(dofsA,.5,1)
    nodeA.createObject('FixedConstraint')
    
    nodeB = rigidTransformationNode.createChild("B")
    B = Compliant.Frame.Frame([1,4,0,0,0,0,1])
    dofsB = B.insert(nodeB, name="dofs_red")
    nodeB.createObject('RigidMass', name = 'mass', mass = 1, inertia="1 1 1")
    setShowDofs(dofsB,.5,2)
    
    constraintAToB = Anatomy.Constraint.RigidTransformation("A_to_B", A, B, [0,0,0,0,1,0])
    constraintAToB.setTarget([math.radians(45)])
    
    #
    # constraint between two rigids linked by a cinematic joint
    #
    rigidRigidRelativeRotationNode = rootNode.createChild("RigidRigidRelativeRotation")

    nodeC = rigidRigidRelativeRotationNode.createChild("C")
    C = Compliant.StructuralAPI.RigidBody(nodeC, "rigidC_green")
    C.setManually(offset=[0,2,0,0,0,0,1])
    setShowDofs(C.dofs,.5,1)
    rigidCJointCenter = C.addAbsoluteOffset("jointCenter_blue", [1,2,0,0,0,0,1])
    setShowDofs(rigidCJointCenter.dofs, 0.2, 3)
    C.node.createObject('FixedConstraint')
    
    nodeD = rigidRigidRelativeRotationNode.createChild("D")
    D = Compliant.StructuralAPI.RigidBody(nodeD, "rigidD_red")
    D.setManually(offset=[2,2,0]+Quaternion.axisToQuat([1.,1.5,1.2], math.radians(60)))
    setShowDofs(D.dofs,.5,2)
    rigidDJointCenter = D.addAbsoluteOffset("jointCenter_blue", [1,2,0,0,0,0,1])
    setShowDofs(rigidDJointCenter.dofs, 0.2, 3)
    
    sphericalJoint = Compliant.StructuralAPI.BallAndSocketRigidJoint("sphericalJoint", rigidCJointCenter.node, rigidDJointCenter.node)
    
    # add physics
    sphericalJoint.addSpring(100)
    sphericalJoint.addDamper(10)
    
    targetAngleDeg = [60, None, 10]
    mask=list()
    for t in targetAngleDeg:
        if not t is None:
            mask+=[1]
        else:
            mask+=[0]
    constraintCToD = Anatomy.Constraint.RigidRigidRelativeRotation("C_to_D", C, D, mask, compliance=1e-6)
    setShowDofs(constraintCToD.frame1AttachedToFrame2.dofs, 0.3)
    
    target = list()
    for t in targetAngleDeg:
        if not t is None:
            target += [math.radians(t)]
    # print "target:", target
    constraintCToD.setTarget(target)

    #
    # absolute rigid constraint between a Rigid and the world frame
    #
    rigidWorldTransformationNode = rootNode.createChild("RigidWorldTransformation")

    nodeE = rigidWorldTransformationNode.createChild("E")
    E = Compliant.StructuralAPI.RigidBody(nodeE, "rigidE_green")
    E.setManually(offset=[-.5,0.1,.8,0.4619397662556433, -0.19134171618254492, 0.4619397662556433, 0.7325378163287419])
    setShowDofs(E.dofs,.5,1)

    Anatomy.Constraint.RigidWorldTransformation("E_to_world",E,[0,0,0,1,0,1],compliance=1e-6)


    #
    # relative rotation constraint between a Rigid and the world frame
    #
    rigidWorldTransformationNode = rootNode.createChild("RelativeRigidWorldRotation")

    nodeF = rigidWorldTransformationNode.createChild("F")
    F = Compliant.StructuralAPI.RigidBody(nodeF, "rigidF_green")
    F.setManually(offset=[-.5,-2,.8,0.4619397662556433, -0.19134171618254492, 0.4619397662556433, 0.7325378163287419])
    setShowDofs(F.dofs,.5,1)

    constraintF = Anatomy.Constraint.RigidWorldRelativeRotation("E_to_world",F,[0,0,0,1,0,1],compliance=1e-6)
    setShowDofs(constraintF.frame1AttachedToFrame2.dofs, 0.3)
    constraintF.setTarget(target)


    
    return rootNode
    