import Sofa

import SofaPython.Tools



class ShapeFunctionMesh:
    """ To create a Mesh suited for creating a ShapeFunction
    """

    def __init__(self, node):
        self.node = node
        self.envelopMeshes = dict()
        self.innerPoints = dict()
        self.closing = None
        self.merge = None
        self.topology = None
        self.pointHandle = dict()
        self.affineHandle = dict()

    def addEnvelop(self, name, mesh):
        self.pointHandle[name] = list()
        self.envelopMeshes[name] = mesh

    def addInnerPoints(self, name, mesh):
        self.pointHandle[name] = list()
        self.innerPoints[name] = mesh

    def addPointHandle(self, name, index):
        # TODO check name
        self.pointHandle[name].append(index)

#    def addMeshClosing(self):
#        self.node.createObject("MeshClosingEngine", name="closing")

    def _addMergeMeshes(self):
        pass

    def getEnvelopPath(self):
        envelopPath = "@"
        if self.closing is None:
            if len(self.envelopMeshes) == 1:
                envelopPath += self.envelopMeshes.values()[0].getPathName()
        return envelopPath

    def addMeshTetraStuffing(self):
        """ Wrapper for MeshTetraStuffing.
        Deprecated : use SofaTetGen plugin
        """

        envelopPath = self.getEnvelopPath()
        self.node.createObject("MeshTetraStuffing", name="stuffing", inputPoints=envelopPath+".position", inputTriangles=envelopPath+".triangles", inputQuads=envelopPath+".quads", splitTetrahedra=True, snapPoints=True)
        self.topology = self.node.createObject("MeshTopology", name="topology", position="@stuffing.outputPoints", tetrahedra="@stuffing.outputTetrahedra")

    def addTetGenTetrahedrization(self):
        envelopPath = self.getEnvelopPath()
        self.node.createObject("MeshTetrahedrizationEngine", name="tetrahedrizer", in_position=envelopPath+".position", in_triangle=envelopPath+".triangles", plc=True, quality=True, minratio=1.414)
        self.topology = self.node.createObject("MeshTopology", name="topology", position="@tetrahedrizer.out_position", tetrahedra="@tetrahedrizer.out_tetra")

    def showTetra(self):
        self.topology.drawTetrahedra=True



class LSDof:
    def __init__(self, sfMesh):
        self.node = sfMesh.node
        self.mesh = sfMesh
        self.affineNode = None
        self.pointNode = None
        self.pointHandles = [] # list of indices in tetra topology of point dof, or path to data
        self.pointHandlesSize = 0
        self.affineHandles = [] # list of list of indices in tetra topology of affine dof, or path to data

    def addAffineFromBoxRoi(self, box, showObject=False):
        """ Wrapper for BoxROI.
        """

        name="affine_"+str(len(self.affineHandles))
        roi = self.node.createObject("BoxROI", name = "roi_"+name, position="@"+self.mesh.topology.getPathName()+".position", triangles="@"+self.mesh.topology.getPathName()+".triangles", tetrahedra="@"+self.mesh.topology.getPathName()+".tetrahedra",  box=SofaPython.Tools.listToStr(box), drawBoxes=showObject, drawPoints=showObject)
#        avg = self.node.createObject("AverageCoord", template="Vec3", name="position_"+name, position="@"+roi.getPathName()+".position", indices="@"+roi.getPathName()+".indices")
        self.affineHandles.append(roi.getPathName()+".indices")
        return len(self.affineHandles)-1


    def addAffineFromMeshRoi(self, mesh, showObject=False):
        """ Wrapper for MeshROI.
        """

        name="affine_"+str(len(self.affineHandles))
        roi = self.node.createObject("MeshROI", name = "roi_"+name, position="@"+self.mesh.topology.getPathName()+".position", triangles="@"+self.mesh.topology.getPathName()+".triangles", tetrahedra="@"+self.mesh.topology.getPathName()+".tetrahedra", ROIposition="@"+mesh+".position", ROItriangles="@"+mesh+".triangles")
#        avg = self.node.createObject("AverageCoord", template="Vec3", name="position_"+name, position="@"+roi.getPathName()+".position", indices="@"+roi.getPathName()+".indices")
        self.affineHandles.append(roi.getPathName()+".indices")

        visuNode = self.node.createChild(name+"_node")
        visuNode.createObject("VisualModel", name=name+"_visual", src="@../"+mesh, color="#aa000088")
        if not showObject:
            visuNode.activated = False

        return len(self.affineHandles)-1


    def addAffineFromIndices(self, indices, showObject=False):
        """ Directly adds a region handle.
        """

        self.affineHandles.append(indices)


    def addPointFromSampler(self, number):
        """ Wrapper for MeshSampler.
        """
        name="sampler_"+str(len(self.pointHandles))
        # TODO implement fixed position in MeshSampler, sampler in a boxROI, sampler in a MeshROI
        sampler = self.node.createObject("MeshSampler", name=name, position="@"+self.mesh.topology.getPathName()+".position", edges="@"+self.mesh.topology.getPathName()+".edges", number=number)
        self.pointHandles = sampler.getPathName()+".outputIndices"
        self.pointHandlesSize += number


    def addPointIndex(self, index):
        """ Directly adds a point handle.
        """

        self.pointHandles.append(index)
        self.pointHandlesSize += 1
        return len(self.pointHandles)-1


    def addPointsNear(self, targets):
        """ Wrapper for ProximityROI.
        """

        name="proximity_"+str(len(self.pointHandles))
        N=len(targets)
        roi = self.node.createObject("ProximityROI", name = name, position="@"+self.mesh.topology.getPathName()+".position", N=N, centers=SofaPython.Tools.listListToStr(targets), radii=SofaPython.Tools.listToStr([1.0 for i in range(N)]))
        self.pointHandles.append(roi.getPathName()+".indices")
        self.pointHandlesSize += N

    def havePointHandles(self):
        return len(self.pointHandles)>0

    def haveAffineHandles(self):
        return len(self.affineHandles)>0

    def addAffineIndices(self):
        """ return path to affine indices data
        """
        args=dict()
        for i,affine in enumerate(self.affineHandles):
            if type(affine) == list :
                args["indices"+str(i+1)] = SofaPython.Tools.listToStr(affine)
            else:
                args["indices"+str(i+1)] = "@"+affine
        mergeRoi = self.node.createObject("MergeROIs", name="affineIndices", nbROIs=len(self.affineHandles), **args)
        return mergeRoi.getPathName()+".roiIndices"

    def addPointIndices(self):
        """ return path to point indices data
        """
        args=dict()
        for i,point in enumerate(self.pointHandles):
            if type(point) == int :
                args["input"+str(i+1)] = SofaPython.Tools.listToStr([point])
            else:
                args["input"+str(i+1)] = "@"+point
        mergeVectors = self.node.createObject("MergeVectors", name="pointIndices", nbInputs=len(self.pointHandles), **args)
        return mergeVectors.getPathName()+".output"

    def getPointIndicesData(self):
        return self.pointNode.getObject("initPosition").getPathName()+".indices"

    def addMechanicalObjects(self):
        """ Creates both pointNode and affineNode nodes, and adds and initializes corresponding MechanicalObjects.
        """
        self.pointNode = self.node.createChild("dofPoint")
        self.pointNode.createObject("PointsFromIndices", name="initPosition", position="@"+self.mesh.topology.getPathName()+".position", indices="@"+self.addPointIndices())
        self.pointNode.createObject("MechanicalObject", template="Vec3", name="dof", position="@initPosition.indices_position")

        self.affineNode = self.node.createChild("dofAffine")
        # initial affine value is identity
        self.affineNode.createObject("MechanicalObject", template="Affine", name="dof", size=len(self.affineHandles))

    def createChild(self, name):
        """ Creates a node having both pointNode and affineNode as parents.
        """
        node = self.pointNode.createChild(name)
        self.affineNode.addChild(node)
        return node

    def insertLSMapping(self, node, isMechanical=True):
        """ Creates an LSMultiMapping node which parent is passed in parameter, using adequate input (pointNode and affineNode).
        """
        return node.createObject("LSMultiMapping", name="mapping", input1="@"+self.pointNode.getPathName(), input2="@"+self.affineNode.getPathName(), output="@.", mapForces=isMechanical, mapConstraints=isMechanical, mapMasses=isMechanical)
#        return node.createObject("LinearLSMultiMapping", name="mapping", input1="@"+self.pointNode.getPathName(), input2="@"+self.affineNode.getPathName(), output="@.", mapForces=isMechanical, mapConstraints=isMechanical, mapMasses=isMechanical)




class ShapeFunction:
    """ High-level API to manipulate an LS ShapeFunction
    """

    def __init__(self, sfMesh, lsDof):
        self.node = sfMesh.node
        self.mesh = sfMesh
        self.lsDof = lsDof
        self.shapeFunction = None


    def addJaconsonShapeFunction(self, stiffness):
        """ Adds energy and weight computation engines as well as the new shape functions, and connects them accordingly.
        """

        topologyPath = "@"+self.mesh.topology.getPathName()
        if len(stiffness) == 0:
            self.node.createObject("JacobsonEnergyComputationEngine", template="Vec3", name="energy", position=topologyPath+".position", laplacianScheme="3", massScheme="0")
        else:
            self.node.createObject("JacobsonEnergyComputationEngine", template="Vec3", name="energy", position=topologyPath+".position", laplacianScheme="3", massScheme="0", stiffness = stiffness)


        argsWeight = dict()
        if self.lsDof.havePointHandles():
            argsWeight["posConstraints"] = "@"+self.node.getObject("pointIndices").getPathName()+".output"
        if self.lsDof.haveAffineHandles():
            argsWeight["regionConstraints"] = "@"+self.lsDof.addAffineIndices()

        self.node.createObject("JacobsonWeightComputationEngine", template="Vec3", name="weights", position="@energy.position", Q="@energy.Q", **argsWeight)
        self.node.createObject("LSJacobsonShapeFunction", name="shapeFunction", W="@weights.W", position="@energy.position", nbPointParent=self.lsDof.pointHandlesSize)



