import os
import shutil
import sys
from unittest import result

import SofaPython.Tools
from SofaPython.Tools import listToStr as concat
import SofaPython.SofaNumpy
import numpy
import SofaPython.units as units
import Compliant.sml
import SofaImage.API
import Flexible.API

import Anatomy.Tools
import tools
import check

import Sofa
printLog = True

class Scene(Compliant.sml.SceneArticulatedRigid):
    """ Builds an Hexa FEM simulation from anatomy model
    rigid [bone] are simulated as rigids
    Compliant joints are setup between the rigids
    SurfaceLink [sliding] are simulated as bilateral contact
    SurfaceLink [collision] are simulated as unilateral contact
    """

    def addMeshExporters(self, dir, ExportAtEnd=False):
        Compliant.sml.SceneArticulatedRigid.addMeshExporters(self,dir,ExportAtEnd=ExportAtEnd) # export rigids
        tools.addMeshExporters(self,dir,ExportAtEnd=ExportAtEnd) # export deformables


    def __init__(self, parentNode, model):

        Compliant.sml.SceneArticulatedRigid.__init__(self, parentNode, model)

        check.solidAnatomicalTags(model)

        # load default anatomy material
        self.material.load(os.path.join(os.path.dirname(__file__), "..", "material.json"))
        self.setMaterialByTag("ligament", "ligament")
        self.setMaterialByTag("muscle", "muscle")
        self.setMaterialByTag("bone", "bone")
        self.setMaterialByTag("skin", "fat")
        self.setMaterialByTag("flesh", "flesh")
        self.setMaterialByTag("ligament_1d", "ligament_1d")
        self.setMaterialByTag("muscle_1d", "muscle_1d")

        # variables
        self.entitiesProp = Anatomy.Tools.EntityProperties()

        self.Images=dict()
        self.behaviors=dict()
        self.behaviors_1d=dict()
        self.collisions=dict()
        self.visuals=dict()

        self.slidingContacts = dict()
        self.collisionContacts = dict() # when symmetricContact=True, for each contact contains two StrucTuralAPI.Contact index by id and id_inverted
        check.checkContactsModel(self.model)

        # FEM
        self.labelImage=None
        self.sampler=None
        self.mappedFEMDof=dict() # one for each rigid
        self.freeFEMDof=None
        self.FEMDof=None
        self.cells=dict()

        # parameters
        self.param.color = {"default":[.8,.8,.8,1], "bone":[.8,.8,.8,1],  "flesh":[1.,0.5,0.5,0.5], "skin":[1.,0.75,0.75,0.5], "fat":[1.,0.75,0.75,0.5], "ligament":[.2,.2,1,1], "muscle":[1,.2,.2,1] }
        self.param.parallel = False # activate parallel computations for component supporting it

        # sliding contacts
        self.param.contactCompliance = 1e-2 # s^2/kg
        self.param.slidingContactRejectBoundary = True
        # collision contacts
        self.param.collisionCompliance = 1e-14 # compliance value for bone collision constraint
        self.param.collisionOffset = 0 # m collisions are detected when the distance is as small as offset
        self.param.collisionNbClosestVertex = None # number of vertices used to detect and react to a collision, None for all vertices

        # FEM
        self.param.autoConnectSkin = True # attach any solid to the flesh
        self.param.voxelSize = 5E-3 # m
        self.param.coarseningFactor = 1
        self.param.gridType = 'branching' # 'regular'

        self.param.softTissueTags={"ligament","muscle","skin","flesh"}
        self.param.boneTags={"bone"}

        self.param.showHexa = False
        self.param.showGPScale = 0 # m

        # springs
        self.param.springTags={"ligament_1d","muscle_1d"}
        self.param.color["ligament_1d"]=self.param.color["ligament"]
        self.param.color["muscle_1d"]=self.param.color["muscle"]
        self.param.showSpringScale= 1E-3 # m
        self.param.enableSpringLimits = False

        # material
        self.param.behaviorLaw = 'hooke'  # 'projective'
        self.param.strainMeasure="Corotational" # 'Green'
        self.param.useStrainOffset = False

        self.param.exportDir=''

    def createScene(self):

        self.param.rigidTags.update(self.param.boneTags) # add bones to compliant tags
        for t in self.param.boneTags:
            self.setMaterialByTag(t, "bone")
        Compliant.sml.SceneArticulatedRigid.createScene(self)

        self.node.createObject("RequiredPlugin", name="Image")
        self.node.createObject("RequiredPlugin", name="BranchingImage")
        self.node.createObject("RequiredPlugin", name="Flexible")
        if len(self.model.getSurfaceLinksByTags({"sliding", "collision"})) > 0:
            self.node.createObject("RequiredPlugin", name="ContactMapping")

        # self.node.createObject('VisualStyle', displayFlags='showBehaviorModels showVisual')
        # self.node.createObject('BackgroundSetting',color='1 1 1')

        tools.insertImagesFromSolids(self,self.param.boneTags | self.param.softTissueTags, self.param.voxelSize)
        tools.connectImages(self)
        self.labelImage = tools.mergeImages(self, self.param.boneTags | self.param.softTissueTags, 'label', self.param.coarseningFactor, True)

        # build grid mesh
        if self.param.gridType == 'branching':
            self.sampler = SofaImage.API.Sampler(self.labelImage.node)
            self.sampler.addImageRegularSampler(self.labelImage,showEdges="0", printLog=printLog)
        else : #  @todo: generate directly a regular image, and handle rigid cells
            regularIm = self.labelImage.toImage(self.labelImage.node,'',conversionType=0)
            self.sampler = SofaImage.API.Sampler(regularIm.node)
            self.sampler.addImageRegularSampler(regularIm,showEdges="0", printLog=printLog)

#        self.labelImage.addViewer()

        self.sampler.addMesh()
        self.createChild(self.sampler.node, "dof")

        # compute the cell indices for all solids
        for solid in self.model.getSolidsByTags(self.param.boneTags | self.param.softTissueTags):
            self.cells[solid.id] = self.labelImage.node.createObject('BranchingCellIndicesFromLabels', name="hexaIndices_"+solid.name, template=self.labelImage.template(), labels = concat(self.entitiesProp.entities[solid.id].allLabels()), image=self.labelImage.getImagePath()+".image", transform=self.labelImage.getImagePath()+".transform")

        if len(self.rigids):  # partially rigid map FEM nodes
            # get cells that should be rigid mapped
            splitArgs = dict()
            for solid in self.model.getSolidsByTags(self.param.boneTags):
                splitArgs['hexahedronIndices'+str(len(splitArgs)+1)]=self.cells[solid.id].getLinkPath()+".cell"

            # segment complete mesh into multiple parts
            split = self.sampler.node.createObject('MeshSplittingEngine', name="split", src=self.sampler.mesh.getLinkPath(), nbInputs=len(splitArgs), **splitArgs)

            self.createChild(self.nodes["dof"], "FEM")

            # map rigid parts
            count=0
            mapinput=' '
            for solid in self.model.getSolidsByTags(self.param.boneTags):
                node = self.rigids[solid.id].node.createChild( "hexa" )
                self.mappedFEMDof[solid.id] = Flexible.API.FEMDof(node)
                self.mappedFEMDof[solid.id].addMechanicalObject(position=split.getLinkPath()+".position"+str(count+1))
                self.mappedFEMDof[solid.id].mapping = self.mappedFEMDof[solid.id].node.createObject('RigidMapping', name="mapping", globalToLocalCoords="true")
                self.sampler.node.addChild(node) # @todo: this generates a DAG validation error, but is needed because it depends on split
                node.addChild(self.nodes["FEM"])
                mapinput += '@'+node.getPathName()+' '
                count+=1

            # remaining free nodes
            self.freeFEMDof = Flexible.API.FEMDof(self.nodes["dof"])
            self.freeFEMDof.addMechanicalObject(position=split.getLinkPath()+".position"+str(count+1))
            self.sampler.node.addChild(self.nodes["dof"])
            mapinput += '@'+self.nodes["dof"].getPathName()+' '

            # full nodes
            self.FEMDof = Flexible.API.FEMDof(self.nodes["FEM"])
            self.FEMDof.addMesh(src=self.sampler.sampler.getLinkPath())
            self.FEMDof.addMechanicalObject()
            self.FEMDof.mapping = self.FEMDof.node.createObject('SubsetMultiMapping', name='mapping', input=mapinput, output = '@./', indexPairs=split.getLinkPath()+".indexPairs")


        else :  # the FEM grid has only independent nodes

            self.nodes["FEM"] = self.nodes["dof"]
            self.FEMDof = Flexible.API.FEMDof( self.nodes["dof"] )
            self.FEMDof.addMesh(src=self.sampler.sampler.getLinkPath())
            self.FEMDof.addMechanicalObject()


        # mass
        # @todo: improve this using non uniform mass
        tools.computeVolumes(self,self.param.softTissueTags)
        totalMass=0
        for solid in self.model.getSolidsByTags(self.param.softTissueTags):
            totalMass += solid.volume * self.material.density(self.getMaterial(solid.id))
        self.FEMDof.addUniformMass(totalMass)

        # display FEM
        if self.param.showHexa:
            self.createChild( self.nodes["FEM"],'VisuHexa' )
            self.nodes["VisuHexa"].createObject('VisualModel', color="0.8 0.8 1 1",edges="@../mesh.edges",position="@../mesh.position",useVBO="0")
            self.nodes["VisuHexa"].createObject('IdentityMapping')

        # soft tissues
        self.FEMDof.addShapeFunction()

        #   possibly improves accuracy of integration volume using finer image
        fineVolumes = self.labelImage.getImagePath()+".volumes" if self.param.gridType == 'branching' and self.param.coarseningFactor>1 else ''


        # @todo: add 'skin' layer
        for solid in self.model.getSolidsByTags(self.param.softTissueTags):
            if printLog:
                Sofa.msg_info('Anatomy.sml',"Anatomy.sml creating solid "+solid.name)

            mat = self.getMaterial(solid.id)
            self.behaviors[solid.id] = tools.insertFEMBehavior(solid.name, self.sampler.mesh, self.nodes["FEM"], self.nodes["FEM"],
                                                                  fineVolumes=fineVolumes , indices= self.cells[solid.id].getLinkPath()+".cell",
                                                                  showSamplesScale=units.length_from_SI(self.param.showGPScale), drawMode=1)

            if self.param.behaviorLaw == 'projective':
                self.behaviors[solid.id].addProjective( youngModulus=self.material.youngModulus(mat) )
            else:
                self.behaviors[solid.id].addHooke( strainMeasure=self.param.strainMeasure, youngModulus=self.material.youngModulus(mat), poissonRatio= self.material.poissonRatio(mat),useOffset=self.param.useStrainOffset )

            if self.param.gridType == 'branching':
                (self.collisions[solid.id],self.visuals[solid.id]) = tools.insertDeformable(solid, self.nodes["FEM"], dofAffineNode=self.nodes["FEM"], labelImage=self.labelImage, labels=self.entitiesProp.entities[solid.id].allLabels(), useGlobalIndices=True, color=solid.getValueByTag(self.param.color))
            else:
                (self.collisions[solid.id],self.visuals[solid.id]) = tools.insertDeformable(solid, self.nodes["FEM"], dofAffineNode=self.nodes["FEM"], color=solid.getValueByTag(self.param.color))

        # add springs
        for solid in self.model.getSolidsByTags(self.param.springTags):
            if printLog:
                Sofa.msg_info('Anatomy.sml',"creating solid "+solid.name)

            if self.param.gridType == 'branching':
                indexLabelPairs=tools.getSpringAttachments(solid, self)
                (self.collisions[solid.id],self.visuals[solid.id]) = tools.insertSpring(solid, self.nodes["FEM"], dofAffineNode=self.nodes["FEM"],
                                                                                       labelImage=self.labelImage,indexLabelPairs=indexLabelPairs,
                                                                                       radius=units.length_from_SI(self.param.showSpringScale), useGlobalIndices=True, color=solid.getValueByTag(self.param.color))
            else:
                (self.collisions[solid.id],self.visuals[solid.id]) = tools.insertSpring(solid, self.nodes["FEM"], dofAffineNode=self.nodes["FEM"],
                                                                                       radius=units.length_from_SI(self.param.showSpringScale),  color=solid.getValueByTag(self.param.color))

            # add behavior
            mat = self.material.data[self.getMaterial(solid.id)]
            self.behaviors_1d[solid.id]=dict()
            for meshid, c in self.collisions[solid.id].iteritems():
                self.behaviors_1d[solid.id][meshid] = tools.insertSpringBehavior(c, mat, self.param.enableSpringLimits)

        # contacts
        tools.insertSlidingContactsFromModel(self)
        tools.insertCollisionContactsFromModel(self)

        # obj exporters
        if len(self.param.exportDir):
            self.addMeshExporters(self.param.exportDir,ExportAtEnd=False)


    def export(self, dataDir):
        if os.path.exists(dataDir):
            shutil.rmtree(dataDir)
        os.makedirs(dataDir)

        # export dofs
        if not self.freeFEMDof is None:
            self.freeFEMDof.write('freeFEMDof',directory=dataDir)
        if not self.FEMDof is None:
            self.FEMDof.write('FEMDof',directory=dataDir)
        for s,fem in self.mappedFEMDof.iteritems():
            fem.write('mappedFEMDof_'+s,directory=dataDir)

        # export behavior (shape function values + gradients + Hessian + volume)
        for b in self.behaviors.values():
            b.write(directory=dataDir)

        # export deformables (shape function values)
        for solid in self.model.getSolidsByTags(self.param.softTissueTags | self.param.springTags):
            for d in self.collisions[solid.id].values() + self.visuals[solid.id].values():
                d.write(directory=dataDir)



    def loadScene(self, dataDir):

        self.param.rigidTags.update(self.param.boneTags) # add bones to compliant tags
        for t in self.param.boneTags:
            self.setMaterialByTag(t, "bone")
        Compliant.sml.SceneArticulatedRigid.createScene(self)

        self.node.createObject("RequiredPlugin", name="Flexible")
        if len(self.model.getSurfaceLinksByTags({"sliding", "collision"})) > 0:
            self.node.createObject("RequiredPlugin", name="ContactMapping")

        self.createChild(self.node, "dof")

        if len(self.rigids):  # partially rigid map FEM nodes

            self.createChild(self.nodes["dof"], "FEM")

            # map rigid parts
            count=0
            mapinput=' '
            for solid in self.model.getSolidsByTags(self.param.boneTags):
                node = self.rigids[solid.id].node.createChild( "hexa" )
                self.mappedFEMDof[solid.id] = Flexible.API.FEMDof(node)
                self.mappedFEMDof[solid.id].read("mappedFEMDof_"+solid.id,directory=dataDir)
                self.mappedFEMDof[solid.id].mapping = self.mappedFEMDof[solid.id].node.createObject('RigidMapping', name="mapping", globalToLocalCoords="true")
                node.addChild(self.nodes["FEM"])
                mapinput += '@'+node.getPathName()+' '
                count+=1

            # remaining free nodes
            self.freeFEMDof = Flexible.API.FEMDof(self.nodes["dof"])
            self.freeFEMDof.read("freeFEMDof",directory=dataDir)
            mapinput += '@'+self.nodes["dof"].getPathName()+' '

            # full nodes
            self.FEMDof = Flexible.API.FEMDof(self.nodes["FEM"])
            self.FEMDof.read("FEMDof",directory=dataDir)
            self.FEMDof.mapping = self.FEMDof.node.createObject('SubsetMultiMapping', name='mapping', input=mapinput, output = '@./')
            self.FEMDof.readMapping("FEMDof",directory=dataDir)

        else :  # the FEM grid has only independent nodes

            self.nodes["FEM"] = self.nodes["dof"]
            self.FEMDof = Flexible.API.FEMDof( self.nodes["dof"] )
            self.FEMDof.read("FEMDof",directory=dataDir)

        # display FEM
        if self.param.showHexa:
            self.createChild( self.nodes["FEM"],'VisuHexa' )
            self.nodes["VisuHexa"].createObject('VisualModel', color="0.8 0.8 1 1",edges="@../mesh.edges",position="@../mesh.position",useVBO="0")
            self.nodes["VisuHexa"].createObject('IdentityMapping')

        # soft tissues
        # @todo: add 'skin' layer
        for solid in self.model.getSolidsByTags(self.param.softTissueTags):
            if printLog:
                Sofa.msg_info('Anatomy.sml',"Anatomy.sml creating solid "+solid.name)

            # add collisions/visuals
            (self.collisions[solid.id],self.visuals[solid.id]) = tools.insertDeformable(solid, self.nodes["FEM"], dofAffineNode=self.nodes["FEM"], color=solid.getValueByTag(self.param.color))
            for d in self.collisions[solid.id].values() + self.visuals[solid.id].values():
                d.read(directory=dataDir)

            # add behavior
            mat = self.getMaterial(solid.id)
            self.behaviors[solid.id] = tools.loadBehavior(solid, self.nodes["FEM"], dataDir, dofAffineNode = self.nodes["FEM"])
            self.behaviors[solid.id].readWeights(directory=dataDir)
            if self.param.behaviorLaw == 'projective':
                self.behaviors[solid.id].addProjective( youngModulus=self.material.youngModulus(mat) )
            else:
                self.behaviors[solid.id].addHooke( strainMeasure=self.param.strainMeasure, youngModulus=self.material.youngModulus(mat), poissonRatio= self.material.poissonRatio(mat),useOffset=self.param.useStrainOffset )

        # add springs
        for solid in self.model.getSolidsByTags(self.param.springTags):
            if printLog:
                Sofa.msg_info('Anatomy.sml',"creating solid "+solid.name)

            # add collisions/visuals
            (self.collisions[solid.id],self.visuals[solid.id]) = tools.insertSpring(solid, self.nodes["FEM"], dofAffineNode = self.nodes["FEM"], radius=SofaPython.units.length_from_SI(self.param.showSpringScale),  useGlobalIndices=True, color=solid.getValueByTag(self.param.color))
            for d in self.collisions[solid.id].values() + self.visuals[solid.id].values():
                d.read(directory=dataDir)

            # add behavior
            mat = self.material.data[self.getMaterial(solid.id)]
            self.behaviors_1d[solid.id]=dict()
            for meshid, c in self.collisions[solid.id].iteritems():
                self.behaviors_1d[solid.id][meshid] = tools.insertSpringBehavior(c, mat, self.param.enableSpringLimits)

        # contacts
        tools.insertSlidingContactsFromModel(self)
        tools.insertCollisionContactsFromModel(self)

        # obj exporters
        if len(self.param.exportDir):
            self.addMeshExporters(self.param.exportDir,ExportAtEnd=False)



    def getCellNodeIndices(self, solidId):
        ## returns a set of the indices of the nodes of the cells/hexa containing the given solid
        ## Of course, it must be called once everything is initialized, and Data are filled-up

        # Note that if the information would be store internally in python, it would not be necessary to use SofaPython
        # but since everything is stored directly in component's Data it is mode efficient not to copy these data in python

        # get the indices of the cells containing the given solid
        cells = SofaPython.SofaNumpy.numpy_data( self.cells[solidId], "cell" )
        # get the indices of the points per hexa
        hexa = SofaPython.SofaNumpy.numpy_data( self.FEMDof.mesh, "hexahedra" )

        result = set() # unique
        for c in cells:
            h = hexa[c[0]]
            for i in xrange(8):
                result.add(h[i])

        return result







def mapFineScene(rootNode, sceneCoarse, sceneFine):
    """  map newly created fine dofs into a existing coarse scene and suppress it
    """
    dt = rootNode.dt

    # sceneFine.node.init()
    rootNode.simulationStep(1E-100)  # need one time step to create the fine scene. init() does not seem to do the job (some data remain dirty..)

    # transfer rigid position
    for rname,r in sceneCoarse.rigids.iteritems() :
        sceneFine.rigids[rname].dofs.position = r.dofs.position

    # map fine node positions in coarse simu
    node = sceneCoarse.nodes["FEM"].createChild("fineHexa")
    dofs = node.createObject('MechanicalObject', name='dofs', template="Vec3d", position=SofaPython.Tools.listListToStr(sceneFine.nodes["dof"].getObject('dofs').position),showObject=True)

    # retrieve coarse cell index for each fine point
    cells=[]
    if sceneFine.param.gridType == 'branching' :
        if sceneFine.param.voxelSize==sceneCoarse.param.voxelSize:
            voxToCoarseInd = sceneCoarse.labelImage.image.voxelIndices
            hexToCoarseInd = [0] * len(sceneFine.sampler.mesh.hexahedra)
            if sceneFine.param.coarseningFactor ==1: # voxel index = fine resolution index
                for nh,h in enumerate(voxToCoarseInd):
                    hexToCoarseInd[nh]=voxToCoarseInd[nh][0]
            else:  # apply additional indirection to convert voxel index to fine resolution index
                voxToFineInd = sceneFine.labelImage.image.voxelIndices
                for nh,h in enumerate(voxToFineInd):
                    hexToCoarseInd[voxToFineInd[nh][0]] = voxToCoarseInd[nh][0]
            cells = [-1]*len(sceneFine.sampler.mesh.position)
            for nh,h in enumerate(sceneFine.sampler.mesh.hexahedra):
                for i in h:
                    cells[i]=hexToCoarseInd[nh]

            if len(sceneFine.rigids): # apply additional indirection to convert FEM node indices to free node indices
                nrigids = len(sceneFine.rigids)
                indexPairs = sceneFine.sampler.node.getObject('split').indexPairs
                cells2 = [-1]*len(sceneFine.nodes["dof"].getObject('dofs').position)
                for i in xrange(len(indexPairs)/2):
                    if indexPairs[2*i][0]==nrigids: # free nodes
                        cells2[indexPairs[2*i+1][0]] = cells[i]
                cells = cells2
        else:
            Sofa.msg_warning('Anatomy.sml',"fine and coarse simulations have different voxel sizes, superimposed hexa might not be well mapped")

    node.createObject('LinearMapping', name='mapping',cell=SofaPython.Tools.listToStr(cells) )
    node.init() # apply to compute mapped fine dofs

    # transfer point position
    sceneFine.nodes["dof"].getObject('dofs').position = SofaPython.Tools.listListToStr(dofs.position)

    # sceneCoarse.node.getParents()[0].detachFromGraph()

    rootNode.propagatePositionAndVelocity()

    rootNode.findData('dt').value=dt

    return

