import math

import Sofa

from SofaPython import Quaternion
from SofaPython.Tools import listToStr as concat
import Compliant.Frame
import Compliant.StructuralAPI
import Anatomy.Constraint

Compliant.StructuralAPI.geometric_stiffness = 2

progressive=True

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

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('CompliantPseudoStaticSolver', name='odesolver', iterations=1, stabilization="pre-stabilization", neglecting_compliance_forces_in_geometric_stiffness=False)
    
    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]
    
    #
    # full rigid constraint between two rigids A-B
    #
    
    rigidRigidTransformationNode = rootNode.createChild("RigidRigid_full")
    
    nodeA = rigidRigidTransformationNode.createChild("A")
    A = Compliant.Frame.Frame([0,0,0,0,0,0,1])
    A.dofs = A.insert(nodeA, name="dofs_green")
    nodeA.createObject('RigidMass', name = 'mass', mass = 1, inertia="1 1 1")
    setShowDofs(A.dofs,.5,1)
    nodeA.createObject('FixedConstraint')
    
    nodeB = rigidRigidTransformationNode.createChild("B")
    B = Compliant.Frame.Frame([0,0,0,0,0,0,1])
    B.dofs = B.insert(nodeB, name="dofs_red")
    nodeB.createObject('RigidMass', name = 'mass', mass = 1, inertia="1 1 1")
    setShowDofs(B.dofs,.5,2)
    
    constraintAToB = Anatomy.Constraint.RigidRigidTransformation("A_to_B", A, B, [1,1,1,1,1,1], progressive=progressive)
    if progressive:
        constraintAToB.progressiveTarget.step = concat([0.005, math.radians(5)])
    constraintAToB.setTarget([1,0,0,math.radians(0), math.radians(220), math.radians(15)])
    
    #
    # rotation only rigid constraint between two rigids C-D
    #
    
    rigidRigidTransformationNode = rootNode.createChild("RigidRigid_rotOnly")
    
    nodeC = rigidRigidTransformationNode.createChild("C")
    C = Compliant.Frame.Frame([0,1,0,0,0,0,1])
    C.dofs = C.insert(nodeC, name="dofs_green")
    nodeC.createObject('RigidMass', name = 'mass', mass = 1, inertia="1 1 1")
    setShowDofs(C.dofs,.5,1)
    nodeC.createObject('FixedConstraint')
    
    nodeD = rigidRigidTransformationNode.createChild("D")
    D = Compliant.Frame.Frame([0,1,0,0,0,0,1])
    D.dofs = D.insert(nodeD, name="dofs_red")
    nodeD.createObject('RigidMass', name = 'mass', mass = 1, inertia="1 1 1")
    setShowDofs(D.dofs,.5,2)
    
    constraintCToD = Anatomy.Constraint.RigidRigidTransformation("C_to_D", C, D, [0,0,0,1,1,1], progressive=progressive)
    if progressive:
        constraintCToD.progressiveTarget.step = concat([0.005, math.radians(5)])
    constraintCToD.setTarget([math.radians(0), math.radians(220), math.radians(15)])
    
    #
    # translation only rigid constraint between two rigids E-F
    #
    
    rigidRigidTransformationNode = rootNode.createChild("RigidRigid_transOnly")
    
    nodeE = rigidRigidTransformationNode.createChild("E")
    E = Compliant.Frame.Frame([0,2,0] + Quaternion.from_euler([math.radians(45),0,0]).tolist())
    E.dofs = E.insert(nodeE, name="dofs_green")
    nodeE.createObject('RigidMass', name = 'mass', mass = 1, inertia="1 1 1")
    setShowDofs(E.dofs,.5,1)
    nodeE.createObject('FixedConstraint')
    
    nodeF = rigidRigidTransformationNode.createChild("F")
    F = Compliant.Frame.Frame([0,2,0,0,0,0,1])
    F.dofs = F.insert(nodeF, name="dofs_red")
    nodeF.createObject('RigidMass', name = 'mass', mass = 1, inertia="1 1 1")
    setShowDofs(F.dofs,.5,2)
    
    constraintEToF = Anatomy.Constraint.RigidRigidTransformation("E_to_F", E, F, [1,1,1,0,0,0], progressive=progressive)
    if progressive:
        constraintEToF.progressiveTarget.step = concat([0.005, math.radians(5)])
    constraintEToF.setTarget([1,0.5,0])
    
    #
    # partial translation only rigid constraint between two rigids G-H
    #
    
    rigidRigidTransformationNode = rootNode.createChild("RigidRigid_partialTransOnly")
    
    nodeG = rigidRigidTransformationNode.createChild("G")
    
    G = Compliant.Frame.Frame([0,3,0] + Quaternion.from_euler([math.radians(45),0,0]).tolist())
    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([0,3,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, [1,1,0,0,0,0], progressive=progressive)
    if progressive:
        constraintGToH.progressiveTarget.step = concat([0.005, math.radians(5)])
    constraintGToH.setTarget([1,0.5])

    #
    # rotation only rigid constraint between two rigids I-J
    #
    
    rigidRigidTransformationNode = rootNode.createChild("RigidRigid_rotOnly_relative")
    
    nodeI = rigidRigidTransformationNode.createChild("I")
    I = Compliant.Frame.Frame([0,4,0,0,0,0,1])
    I.dofs = I.insert(nodeI, name="dofs_green")
    nodeI.createObject('RigidMass', name = 'mass', mass = 1, inertia="1 1 1")
    setShowDofs(I.dofs,.5,1)
    nodeI.createObject('FixedConstraint')
    
    nodeJ = rigidRigidTransformationNode.createChild("J")
    J = Compliant.Frame.Frame([0,4,0] + Quaternion.from_euler([math.radians(30),0,0]).tolist())
    J.dofs = J.insert(nodeJ, name="dofs_red")
    nodeJ.createObject('RigidMass', name = 'mass', mass = 1, inertia="1 1 1")
    setShowDofs(J.dofs,.5,2)
    
    constraintIToJ = Anatomy.Constraint.RigidRigidTransformation("I_to_J", I, J, [0,0,0,1,1,1], progressive=progressive, isRelative=True)
    if progressive:
        constraintIToJ.progressiveTarget.step = concat([0.005, math.radians(5)])
    constraintIToJ.setTarget([math.radians(0), math.radians(220), math.radians(15)])
    
    return rootNode


