import Sofa

mapBoundary = True

def createScene(rootNode):
    
    rootNode.createObject('RequiredPlugin',pluginName="Flexible")
    rootNode.createObject('RequiredPlugin',pluginName="image")
    rootNode.createObject('RequiredPlugin',pluginName="BranchingImage")

    rootNode.gravity="0 0 -10"
    rootNode.dt="0.1"
    rootNode.createObject('BackgroundSetting',color='1 1 1')

    rootNode.createObject('VisualStyle',displayFlags="showVisual showBehaviorModels showForceFields hideWireframe")

    # images
    rootNode.createObject('MeshObjLoader',filename="mesh/cubeQuad.obj",translation="0 -1 0",triangulate="1")
    rootNode.createObject('MeshObjLoader',filename="mesh/cubeQuad.obj",translation="0 11 0",triangulate="1")
    rootNode.createObject('MeshObjLoader',filename="mesh/cubeQuad.obj",translation="2 5 0",scale3d="1 7 1",triangulate="1")
    rootNode.createObject('MergeMeshes',name="fr",nbMeshes="3",position1="@[-1].position",triangles1="@[-1].triangles",position2="@[-2].position",triangles2="@[-2].triangles",position3="@[-3].position",triangles3="@[-3].triangles")
    rootNode.createObject('MeshToImageEngine',name="fr_im",template="ImageUC",position="@fr.position",triangles="@fr.triangles",voxelSize=".18",padSize="1",rotateImage="false",value="3",insideValue="3")
    rootNode.createObject('ImageToBranchingImageConverter',name="fr_biconv",inputImage="@fr_im.image",inputTransform="@fr_im.transform",coarseningLevels="0",superimpositionType="0")
    label = rootNode.createObject('ImageContainer',template="BranchingImageUC",name="fr_bi",branchingImage="@fr_biconv.branchingImage",transform="@fr_biconv.transform",drawBB="false")

    rootNode.createObject('MeshObjLoader',name="cyl",filename="mesh/cylinder_open.obj",triangulate="1")
    rootNode.createObject('MeshClosingEngine',name="closedcyl",inputPosition="@cyl.position",inputTriangles="@cyl.triangles")
    rootNode.createObject('MeshToImageEngine',name="cyl_im",template="ImageUC",position="@closedcyl.position",triangles="@closedcyl.triangles",voxelSize=".18",padSize="1",rotateImage="false",value="1",insideValue="1",roiValue="2",roiIndices="@closedcyl.indices")
    rootNode.createObject('ImageToBranchingImageConverter',name="cyl_biconv",inputImage="@cyl_im.image",inputTransform="@cyl_im.transform",coarseningLevels="0",superimpositionType="0")
    rootNode.createObject('ImageContainer',template="BranchingImageUC",name="cyl_bi",branchingImage="@cyl_biconv.branchingImage",transform="@cyl_biconv.transform",drawBB="false")

    # fem mesh generation and spliting
    rootNode.createObject('BranchingImageSubsampler',template="BranchingImageUC",name="subsampled",inputImage="@cyl_bi.branchingImage",inputTransform="@cyl_bi.transform",coarseningFactor="4",splitLabels="false")
    sampler = rootNode.createObject('ImageSampler',template="BranchingImageUC",name="sampler",src="@subsampled",param="1",showEdges="False",printLog="true")
    rootNode.createObject('MeshLabeler',template="BranchingImageUC",name="labeler",image="@cyl_bi.image",hexahedra="@sampler.hexahedra",voxelToHexahedra="@subsampled.voxelIndices",printLog="true")

    if mapBoundary :
        # use merged image for shapefunction computation
        label = rootNode.createObject('MergeImagesIntoBranching',name="merged_bi",nbImages="2",image1="@cyl_im.image",transform1="@cyl_im.transform",image2="@fr_im.image",transform2="@fr_im.transform",connectLabels="2 3")

        rootNode.createObject('ImageMeshIntersection',template="BranchingImageUC",name="imageMeshIntersection",computeImage=False,image="@fr_bi.image",transform="@fr_bi.transform",hexahedra="@sampler.hexahedra",position="@sampler.position",printLog="true")
        rootNode.createObject('MeshLabeler',template="BranchingImageUC",name="labeler2",image="@fr_bi.image",hexahedra="@sampler.hexahedra",voxelToHexahedra="@imageMeshIntersection.voxelIndices",printLog="true")
        # visu output (set computeImage=true)
        # rootNode.createObject('BranchingImageToImageConverter',template="BranchingImageUI,ImageUI",name="conv",inputImage="@imageMeshIntersection.outputImage",conversionType=0)
        # rootNode.createObject('ImageContainer',template="ImageUI",name="im",image="@conv.image",transform="@cyl_bi.transform")
        # rootNode.createObject('ImageViewer',template="ImageUI",name="v",image="@im.image",transform="@im.transform")
        # test ImageMeshIntersection using regular images
        # rootNode.createObject('ImageMeshIntersection',template="ImageUC",name="imageMeshIntersection_i",computeImage=True,image="@cyl_im.image",transform="@cyl_im.transform",hexahedra="@sampler.hexahedra",position="@sampler.position",printLog="true")
        # rootNode.createObject('ImageViewer',template="ImageUI",name="v",image="@imageMeshIntersection_i.outputImage",transform="@cyl_im.transform")

        rootNode.createObject('SelectConnectedLabelsROI',name="connectROI",template="unsigned char",nbLabels=2,labels1="@labeler.pointLabels",labels2="@labeler2.pointLabels",connectLabels="2 3")
        split = rootNode.createObject('MeshSplittingEngine',name="split",position="@sampler.position",hexahedra="@sampler.hexahedra",nbInputs="1",indices1="@connectROI.indices")
        offset = rootNode.createObject('BranchingCellOffsetsFromPositions',template="BranchingImageUC",name="connect",image="@merged_bi.image",transform="@merged_bi.transform",position="@split.position1",labels="1 2 3")
        # visu mapped nodes
        # rootNode.createObject('MechanicalObject',template="Vec3",name="d",showObject="true",showObjectScale="0.2",position="@split.position1", drawMode=1,showColor="0 1 0 1")

    else:
        offset = rootNode.createObject('BranchingCellOffsetsFromConnectLabels',template="BranchingImageUC",name="offset",image="@fr_bi.image",transform="@fr_bi.transform",pointLabels="@labeler.pointLabels",position="@sampler.position",connectLabels="2 3")
        split = rootNode.createObject('MeshSplittingEngine',name="split",nbInputs="1",position="@sampler.position",indices1="@offset.indices")

    # global solver
    rootNode.createObject('EulerImplicit')
    rootNode.createObject('CGLinearSolver',iterations="25",tolerance="1E-5",threshold="1E-5")

    # Frame simu
    frameNode = rootNode.createChild("Frames")

    frameNode.createObject('ImageSampler',template="BranchingImageUC",name="sampler",src="@../../fr_bi",method="1",param="5 1")
    frameNode.createObject('Mesh',name="mesh",position="@sampler.position")
    frameNode.createObject('MechanicalObject',template="Affine",name="dof",showObject="true",showObjectScale="1.2",src="@mesh")
    frameNode.createObject('FixedConstraint',indices="0")

    frameNode.createObject('VoronoiShapeFunction',name="SF",template="ShapeFunctiond,BranchingImageUC",position="@dof.rest_position",src=label.getLinkPath(),method="0",nbRef="4",bias="false",clearImages=True)
    # visu
    # frameNode.createObject('BranchingImageToImageConverter',template="BranchingImageUI,ImageUI",inputBranchingImage="@SF.indices",conversionType="0")
    # frameNode.createObject('ImageViewer',template="ImageUI",image="@[-1].image",transform=label.getLinkPath()+".transform")

    bnode = frameNode.createChild("behavior")
    bnode.createObject('ImageGaussPointSampler',name="sampler",template="BranchingImageD,BranchingImageUC",mask=label.getLinkPath()+".image",maskLabels="3",indices="@../SF.indices",weights="@../SF.weights",transform="@../SF.transform",method="2",order="4",targetNumber="1")
    bnode.createObject('MechanicalObject',template="F332")
    bnode.createObject('LinearMapping',template="Affine,F332")
    enode = bnode.createChild("E")
    enode.createObject('MechanicalObject',template="E332",name="E")
    enode.createObject('CorotationalStrainMapping',template="F332,E332")
    enode.createObject('HookeForceField',template="E332",name="ff",youngModulus="10000",poissonRatio="0",viscosity="0")

    cnode = frameNode.createChild("collision2")
    cnode.createObject('VisualStyle',displayFlags="hideBehaviorModels")
    cnode.createObject('Mesh',name="mesh",src="@../../fr")
    cnode.createObject('MechanicalObject',template="Vec3d",name="pts")
    cnode.createObject('UniformMass',totalMass="20")
    cnode.createObject('BranchingCellOffsetsFromPositions',name="cell",template="BranchingImageUC",position="@mesh.position",src=label.getLinkPath(),labels="3")
    cnode.createObject('LinearMapping',template="Affine,Vec3d",cell="@cell.cell")
    vnode = cnode.createChild("visual")
    vnode.createObject('VisualModel',color="0.50.511")
    vnode.createObject('IdentityMapping')

    mappedNode = frameNode.createChild("mapped_FEMnodes")
    mappedDofs = mappedNode.createObject('MechanicalObject',template="Vec3d",position=split.getLinkPath()+".position1")
    mappedNode.createObject('LinearMapping',template="Affine,Vec3d",cell=offset.getLinkPath()+".cell")

    # FEM simu
    femNode = rootNode.createChild("FEM")

    freeNode = femNode.createChild("free_FEMnodes")
    freeNode.createObject('MechanicalObject',template="Vec3d",position=split.getLinkPath()+".position2")

    allNode = freeNode.createChild("all_FEMnodes")
    mappedNode.addChild(allNode)
    allNode.createObject('Mesh',name="mesh",position=sampler.getLinkPath()+".position",hexahedra=sampler.getLinkPath()+".hexahedra")
    allNode.createObject('MechanicalObject',template="Vec3d")
    allNode.createObject('SubsetMultiMapping', name='mapping', input='@'+mappedNode.getPathName()+' @'+freeNode.getPathName(), output = '@./', indexPairs=split.getLinkPath()+".indexPairs")
    allNode.createObject('UniformMass',totalMass="200")
    # allNode.createObject('HexahedronFEMForceField',younModulus="1000")
    allNode.createObject('FlexibleCorotationalFEMForceField', youngModulus="1000.0",poissonRatio="0.3",method="polar",order="2")

    VisuHexa = allNode.createChild('VisuHexa')
    VisuHexa.createObject('VisualModel', color="0.8 0.8 1 1",edges="@../mesh.edges",position="@../mesh.position")
    VisuHexa.createObject('IdentityMapping')

    cnode = allNode.createChild("collision1")
    cnode.createObject('VisualStyle',displayFlags="hideBehaviorModels")
    cnode.createObject('Mesh',name="mesh",src="@/cyl")
    cnode.createObject('MechanicalObject',template="Vec3d",name="pts")
    # cnode.createObject('BranchingCellOffsetsFromPositions',name="cell",template="BranchingImageUC",position="@mesh.position",src="@/cyl_bi",labels="1 2")
    # cnode.createObject('LinearMapping',template="Vec3d,Vec3d",cell="@cell.cell")
    cnode.createObject('BarycentricMapping',input="@..",output="@.")
    vnode = cnode.createChild("visual")
    vnode.createObject('VisualModel',color="1 0.5 0.5")
    vnode.createObject('IdentityMapping')

    return rootNode
    
    
