import Sofa
from Compliant import StructuralAPI, Tools
import SofaPython.Quaternion as quat
from SofaPython.Tools import listToStr as concat
from SofaPython.Tools import listListToStr as lconcat

USEPREFACTORIZATION = True
CONSTRAINALLPOINTS = False
USETANGENTPLANEPROJECTION = False
ROISize = 5000

def createSceneAndController(rootNode):
    
    rootNode.gravity=[0,0,0]
    rootNode.createObject('RequiredPlugin', pluginName="Compliant")
    rootNode.createObject('RequiredPlugin', pluginName="ContactMapping")
    
    rootNode.createObject('BackgroundSetting', color='1 1 1')
    rootNode.createObject('VisualStyle', displayFlags='showMechanicalMappings')

    rootNode.dt=1
    rootNode.gravity=[0,0,0]
    
    rootNode.createObject('CompliantAttachButtonSetting', compliance=1E-10)

    if USEPREFACTORIZATION:
        rootNode.createObject('ConstantCompliantPseudoStaticSolver', name="odesolver", iterations=1)
    else:
        rootNode.createObject('CompliantPseudoStaticSolver', name="odesolver", iterations=5)
#    rootNode.createObject('SequentialSolver', name="numsolver", iterations =10, precision=1e-14)
    rootNode.createObject('LDLTSolver', name="numsolver")
    # rootNode.createObject('MinresSolver', name="numsolver", iterations =10, precision=1e-14)

    t = StructuralAPI.RigidBody( rootNode, "target" )
    t.setManually(offset = [0,0,0,0,0,0,1] )
    # t.addCollisionMesh( 'mesh/raptor_35kp.obj', scale3d=[.2,.15,0.08] ) # stressing the registration a bit with different models
    t.addCollisionMesh( 'mesh/raptor_35kp.obj', scale3d=[.1,.1,.1] ) # similar models
    t.collision.addNormals()
    vt=t.collision.addVisualModel()
    vt.model.setColor(1,0,0,1)
    t.setFixed()

    s = StructuralAPI.RigidBody( rootNode, "source" )
    s.setManually(offset = [1,0,0]+quat.from_euler([1,1,1]).tolist() )
    s.addCollisionMesh( 'mesh/raptor_8kp.obj', scale3d=[.1,.1,.1] )
    s.collision.addNormals()
    vs=s.collision.addVisualModel()
    vs.model.setColor(0,1,0,1)

    # addConstraint(s,t,True)
    # addConstraint(s,t,False)
    addWeightedConstraint(s,t,0.5)

    return rootNode


def addConstraint(r1,r2,flip):
    """ attraction of projection constraints
    """

    name = 'attraction' if flip else 'projection'
    node = r1.collision.node.createChild(name)
    node.createObject('MechanicalObject', template="Vec3d", name="dof")

    global cp
    if flip:
        cp = node.createObject('ClosestPointEngine', name="cp", parallel=True, sourcePosition="@"+r2.collision.node.getPathName()+"/topology.position", sourceRigidPosition="@"+r2.node.getPathName()+"/dofs.position", targetPosition="@"+r1.collision.node.getPathName()+"/topology.position", targetRigidPosition="@"+r1.node.getPathName()+"/dofs.position", printLog=True,drawMode="0",method="1",useBV=False, cacheSize=5, targetROISize=ROISize, sourceROISize=ROISize, revertPairs=True)
    else:
        cp = node.createObject('ClosestPointEngine', name="cp", parallel=True, sourcePosition="@"+r1.collision.node.getPathName()+"/topology.position", sourceRigidPosition="@"+r1.node.getPathName()+"/dofs.position", targetPosition="@"+r2.collision.node.getPathName()+"/topology.position", targetRigidPosition="@"+r2.node.getPathName()+"/dofs.position", printLog=True,drawMode="0",method="1",useBV=False, cacheSize=5, targetROISize=ROISize, sourceROISize=ROISize)
    # node.createObject('ClosestPointEngine', name="cp", parallel=True, sourcePosition="@"+r1.collision.node.getPathName()+"/dofs.position", targetPosition="@"+r2.collision.node.getPathName()+"/dofs.position",  printLog=True,drawMode="0",method="1",useBV=False, cacheSize=100, targetROISize=0, sourceROISize=0)

    if not USETANGENTPLANEPROJECTION:
        # projection to mesh -> preserves momentum
        node.createObject('ProjectionToMeshEngine', name="pm", printLog=True, template="Vec3d",
                        sourcePosition="@"+r1.collision.dofs.getPathName()+".position",
                        targetPosition="@"+r2.collision.dofs.getPathName()+".position",
                        indexPairs="@cp.indexPairs" ,
                        triangles2="@"+r2.collision.node.getPathName()+"/topology.triangles", quads2="@"+r2.collision.node.getPathName()+"/topology.quads",
                        useRestLength=True, restLength='0' )
    else:
        # projection to tangent planes -> ghost momentum
        node.createObject('ProjectionToVertexPlaneEngine', name="pm",
                        sourcePosition="@"+r1.collision.dofs.getPathName()+".position",
                        targetPosition="@"+r2.collision.dofs.getPathName()+".position",
                        indexPairs="@cp.indexPairs" ,
                        normals2="@"+r2.collision.node.getPathName()+"/normalsFromPoints.normals",
                        useRestLength=True, restLength='0' )


    if USEPREFACTORIZATION:
        node.createObject('TargetFromProjectionEngine', template="Vec3d", name="targetEngine", position="@"+r1.collision.dofs.getPathName()+".position", indexPairs="@cp.indexPairs", closests="@pm.output")
        node.createObject('DifferenceFromTargetMapping', indices="@targetEngine.indices", targetIndices="@targetEngine.indices", template="Vec3d,Vec3d", targets="@targetEngine.targets",showObjectScale="0")
        # node.createObject('DifferenceFromTargetMapping', template="Vec3d,Vec3d", targets="@targetEngine.targets",showObjectScale="0")
        node.createObject('UniformCompliance', compliance=1E-2,rayleighStiffness=0,isCompliance='0',resizable=True)

    else:
        node.createObject('ProjectionToMeshMapping', template="Vec3d,Vec3d", name="mapping",  input="@"+r1.collision.node.getPathName(), output="@.", indexPairs="@cp.indexPairs", position="@"+r2.collision.node.getPathName()+"/topology.position", closests="@pm.output", printLog=True)
        diffnode = node.createChild("diff")
        diffnode.createObject('MechanicalObject', template="Vec3d", name="dof")
        global diffmapping
        diffmapping = diffnode.createObject('DifferenceMultiMapping', template="Vec3d,Vec3d",  input="@"+r1.collision.node.getPathName()+" @../", output="@.")
        diffnode.createObject('UniformCompliance', compliance=1E-2,rayleighStiffness=0,isCompliance='0',resizable=True)

    # r2.collision.node.addChild(node)

def bwdInitGraph(node):
    try:
        diffmapping
    except:
        diffmapping=None
    if not diffmapping is None:
        print "set pairs"
        roi = cp.sourceROI
        pairs=list()
        for i,r in enumerate(roi):
            pairs.append(r[0])
            pairs.append(i)
        diffmapping.pairs= concat(pairs)
        diffmapping.reinit()

def addWeightedConstraint(r1,r2,weight):
    """ weighted combination of attraction and projection constraints
    """

    name = 'weightedAttractionProjection'
    node = r1.collision.node.createChild(name)
    node.createObject('MechanicalObject', template="Vec3d", name="dof")

    node.createObject('ClosestPointEngine', name="cp_attraction", parallel=True, sourcePosition="@"+r2.collision.node.getPathName()+"/topology.position", sourceRigidPosition="@"+r2.node.getPathName()+"/dofs.position", targetPosition="@"+r1.collision.node.getPathName()+"/topology.position", targetRigidPosition="@"+r1.node.getPathName()+"/dofs.position", printLog=True,drawMode="0",method="1",useBV=False, cacheSize=5, targetROISize=ROISize, sourceROISize=ROISize, revertPairs=True)
    node.createObject('ClosestPointEngine', name="cp_projection", parallel=True, sourcePosition="@"+r1.collision.node.getPathName()+"/topology.position", sourceRigidPosition="@"+r1.node.getPathName()+"/dofs.position", targetPosition="@"+r2.collision.node.getPathName()+"/topology.position", targetRigidPosition="@"+r2.node.getPathName()+"/dofs.position", printLog=True,drawMode="0",method="1",useBV=False, cacheSize=5, targetROISize=ROISize, sourceROISize=ROISize)

    if not USETANGENTPLANEPROJECTION:
        # projection to mesh -> preserves momentum
        node.createObject('ProjectionToMeshEngine', name="pm_attraction", printLog=True, template="Vec3d",
                        sourcePosition="@"+r1.collision.dofs.getPathName()+".position",
                        targetPosition="@"+r2.collision.dofs.getPathName()+".position",
                        indexPairs="@cp_attraction.indexPairs" ,
                        triangles2="@"+r2.collision.node.getPathName()+"/topology.triangles", quads2="@"+r2.collision.node.getPathName()+"/topology.quads",
                        useRestLength=True, restLength='0' )
        node.createObject('ProjectionToMeshEngine', name="pm_projection", printLog=True, template="Vec3d",
                        sourcePosition="@"+r1.collision.dofs.getPathName()+".position",
                        targetPosition="@"+r2.collision.dofs.getPathName()+".position",
                        indexPairs="@cp_projection.indexPairs" ,
                        triangles2="@"+r2.collision.node.getPathName()+"/topology.triangles", quads2="@"+r2.collision.node.getPathName()+"/topology.quads",
                        useRestLength=True, restLength='0' )
    else:
        # projection to tangent planes -> ghost momentum
        node.createObject('ProjectionToVertexPlaneEngine', name="pm_attraction",
                        sourcePosition="@"+r1.collision.dofs.getPathName()+".position",
                        targetPosition="@"+r2.collision.dofs.getPathName()+".position",
                        indexPairs="@cp_attraction.indexPairs" ,
                        normals2="@"+r2.collision.node.getPathName()+"/normalsFromPoints.normals",
                        useRestLength=True, restLength='0' )
        node.createObject('ProjectionToVertexPlaneEngine', name="pm_projection",
                        sourcePosition="@"+r1.collision.dofs.getPathName()+".position",
                        targetPosition="@"+r2.collision.dofs.getPathName()+".position",
                        indexPairs="@cp_projection.indexPairs" ,
                        normals2="@"+r2.collision.node.getPathName()+"/normalsFromPoints.normals",
                        useRestLength=True, restLength='0' )

    node.createObject('TargetFromProjectionEngine', template="Vec3d", name="targetEngine_attraction", position="@"+r1.collision.dofs.getPathName()+".position", indexPairs="@cp_attraction.indexPairs", closests="@pm_attraction.output")
    node.createObject('TargetFromProjectionEngine', template="Vec3d", name="targetEngine_projection", position="@"+r1.collision.dofs.getPathName()+".position", indexPairs="@cp_projection.indexPairs", closests="@pm_projection.output")

    node.createObject('TargetInterpolationEngine', template="Vec3d", name="targetEngine", weight=weight,
                      indices1="@targetEngine_projection.indices",indices2="@targetEngine_attraction.indices",
                      targets1="@targetEngine_projection.targets",targets2="@targetEngine_attraction.targets")

    if CONSTRAINALLPOINTS:
        node.createObject('DifferenceFromTargetMapping', template="Vec3d,Vec3d", targets="@targetEngine.targets",showObjectScale="0")
    else:
        if USEPREFACTORIZATION :
            Sofa.msg_warning("the set of constrained points should be constant for prefactorization -> use sourceROI")
            node.createObject('DifferenceFromTargetMapping', indices="@cp_projection.sourceROI", targetIndices="@cp_projection.sourceROI", template="Vec3d,Vec3d", targets="@targetEngine.targets",showObjectScale="0")
        else:
            node.createObject('DifferenceFromTargetMapping', indices="@targetEngine.indices", targetIndices="@targetEngine.indices", template="Vec3d,Vec3d", targets="@targetEngine.targets",showObjectScale="0")
    node.createObject('UniformCompliance', compliance=1E-2,rayleighStiffness=0,isCompliance='0',resizable=True)

    # r2.collision.node.addChild(node)





