import os
import pickle
import colorsys
from numpy import *

import xml.etree.ElementTree as etree

import Sofa
from SofaPython import Quaternion

class EntityProp:
    """ Store the properties for an entity
    """
    def __init__(self, density=None, youngModulus=None):
        self.roi=dict()
        self.density=density
        self.youngModulus=youngModulus

    def getLabel(self,roi):
        return self.roi[roi]

    def allLabels(self):
        return self.roi.values()
    
class Connection:
    def __init__(self):
        self.label1=None
        self.label2=None

class EntityProperties:
    """ This class handles anatomical entities properties, labels and connections, young module, density
    WARNING: a special "surface" roi is automatically added
    WARNING: a special "inside" roi is automatically added
    WARNING: a special "closing" roi is automatically added
    TODO: add save/load
    TODO: getDensityMap
    """

    def __init__(self):
        self._labelGenerator=0 # last used label, 0 will not be used
        self.entities=dict()
        self.connections=set()

    def _generateLabel(self):
        self._labelGenerator+=1
        if self._labelGenerator>255:
            Sofa.msg_error("Anatomy.Tools","generateLabel: more than 255 labels is not supported")
        return self._labelGenerator

    def addEntity(self, name, density=None, youngModulus=None):
        entity = EntityProp(density, youngModulus)
        entity.roi["surface"] = self._generateLabel()
        entity.roi["closing"] = self._generateLabel()
        entity.roi["inside"] = self._generateLabel()
        self.entities[name]=entity
        return entity

    def addROI(self, name, roi):
        if not roi in self.entities[name].roi:
            self.entities[name].roi[roi] = self._generateLabel()
        else:
            Sofa.msg_error("Anatomy.Tools","ROI {0} is already existing in solid {1}".format(roi,name))
        return self.entities[name].roi[roi]

    def connect(self, entity1, entity2, roi1=None, roi2=None):
        if roi1 is None and roi2 is None:
            self.connectAll(entity1, entity2)
        if not entity1 in self.entities:
            Sofa.msg_error("Anatomy.Tools","LabelFactory.connect: entity {0} does not exist".format(entity1))
            return
        if not roi1 is None and not roi1 in self.entities[entity1].roi:
            Sofa.msg_error("Anatomy.Tools","LabelFactory.connect: roi {0} in entity {1} does not exist".format(roi1,entity1))
            return
        if not entity2 in self.entities:
            Sofa.msg_error("Anatomy.Tools","LabelFactory.connect: entity {0} does not exist".format(entity2))
            return
        if not roi2 is None and not roi2 in self.entities[entity2].roi:
            Sofa.msg_error("Anatomy.Tools","LabelFactory.connect: roi {0} in entity {1} does not exist".format(roi2,entity2))
            return
        if roi1 is None:
            labels1=self.entities[entity1].allLabels()
        else:
            labels1=[self.entities[entity1].getLabel(roi1)]
        if roi2 is None:
            labels2=self.entities[entity2].allLabels()
        else:
            labels2=[self.entities[entity2].getLabel(roi2)]

        for l1 in labels1:
            for l2 in labels2:
                c=Connection()
                c.label1 = l1
                c.label2 = l2
                if c in self.connections:
                    Sofa.msg_warning("Anatomy.Tools","LabelFactory.connect: connection already exists: entity1:{0} roi1:{1} | entity2:{2} roi2:{3}".format(entity1, roi1, entity2, roi2))
                else:
                    self.connections.add(c)

    def connectAll(self, entity1, entity2):
        for roi1 in self.entities[entity1].roi.keys():
            for roi2 in self.entities[entity2].roi.keys():
                self.connect(entity1, entity2, roi1, roi2)

    def getConnectionString(self):
        """ Compute the connection string as expected by MergeBranchingImage
        """
        s = ""
        for c in self.connections:
            s+= str(c.label1) + " " + str(c.label2) + " "
        return s

    def getYoungMap(self):
        """ Compute label to young modulus map as expected by TransferFunction
        """
        m=""
        for entity in self.entities.values():
            for label in entity.allLabels() : # TODO: differentiate surface and inside young modulus
                m += str(label) + " " + str(entity.youngModulus) + " "
        return m

    def getDensityMap(self):
        """ Compute label to density map as expected by TransferFunction
        """
        m=""
        for entity in self.entities.values():
            for label in entity.allLabels() :
                m += str(label) + " " + str(entity.density) + " "
        return m

    def getFusionMap(self):
        """ Compute all labels to first label map as expected by TransferFunction
            This is used to fuse all labels of each entity
        """
        m=""
        for entity in self.entities.values():
            label0 = entity.allLabels()[0]
            for label in entity.allLabels() :
                m += str(label) + " " + str(label0) + " "
        return m

    def write(self, dataDir):
        with open(os.path.join(dataDir,"entitiesProp.pkl"), "w") as f:
            pickle.dump(self, f)

    def getSolidLabelsByTag(self, model, tag):
        """ return a set which contains the labels of all solids having the tag
        """
        labels = set()
        for id,entityProp in self.entities.iteritems():
            if tag in model.solids[id].tags:
                labels.update(entityProp.allLabels())
        return labels

    @staticmethod
    def read(dataDir):
        with open(os.path.join(dataDir,"entitiesProp.pkl"), "r") as f:
            return pickle.load(f)

class ColorFactory:
    """Simple palette generator using the HSV color space"""

    def __init__(self, nbColors=10, s=1., v=1.):
        self.palette=list()
        for i in range(0, nbColors):
            self.palette.append(colorsys.hsv_to_rgb(float(i)/nbColors, s, v))

    def getColor(self, i):
        return self.palette[i%len(self.palette)]

def transform(p,offset):
    return array(offset[:3]) + Quaternion.rotate(offset[3:],p)


def getFirstVertex(file):
    ext = os.path.splitext(file)[1][1:]
    if ext=="obj":
        with open(file,'r') as f:
            for line in f:
                if "v " in line:
                    return map(float,line[1:-1].strip().split())[0:3]
    elif ext=="vtu":
        vtu = readVtuIdVertices(file)
        return vtu.values()[0]

def readObjVertices(objFile):
    """ read vertices from an obj file
        @deprecated have a look to SofaPython.MeshLoader
        """
    Sofa.msg_deprecated("Anatomy.Tools","readObjVertices is deprecated, have a look to SofaPython.MeshLoader")
    assert os.path.exists(objFile)
    pos=[]
    with open(objFile,'r') as f:
        for line in f:
            if "v " in line:
                pos.append(map(float,line[1:-1].strip().split())[0:3])
    return pos

def readVtuIdVertices(vtuFile):
    """ read a vtu file and return a {id:[vertex]} dictionary of the vertices of <Points>
    """
    vertices={}
    with open(vtuFile,'r') as f:
        vtuTree = etree.parse(f)
        nbPoints = int(vtuTree.find("./UnstructuredGrid/Piece").attrib["NumberOfPoints"])
        points = map(float,vtuTree.find("./UnstructuredGrid/Piece/Points/DataArray[@Name='Points']").text.split())
        ids = map(int,vtuTree.find("./UnstructuredGrid/Piece/PointData/DataArray[@Name='nid']").text.split())
        for i in range(nbPoints):
            vertices[ids[i]]=points[3*i:3*i+3]
    return vertices

def readVtmIdVertices(vtmFile):
    """ read a vtm file and return a {id:[vertex]} dictionary of the vertices of all <DataSet>
    """
    vertices={}
    basePath = os.path.dirname(vtmFile)
    with open(vtmFile,'r') as f:
        vtmTree = etree.parse(f)
        for dataSet in vtmTree.findall("./vtkMultiBlockDataSet/DataSet"):
            vertices.update(readVtuIdVertices(os.path.join(basePath,dataSet.attrib["file"])))
    return vertices

def barycenter(pos):
    """ compute the barycenter of the set of points in \a pos
    """
    M=array([0,0,0])
    for p in pos:
        M=M+array(p)
    M=M*(1./len(pos))
    return M.tolist()

def fitSphere(pos, precision = 1E-3, maxIter=100):
    """ approximate a set of points by a sphere
        output: sphere center and radius
    """
    precision2=precision*precision
    nbinv=1./len(pos)
    #average
    M=array([0,0,0])
    for p in pos:
        M=M+array(p)
    M=M*nbinv
    #newton
    C=M
    for it in xrange(maxIter):
        L=array([0,0,0])
        Lm=0
        for p in pos:
            li=linalg.norm(C-array(p))
            if li !=0 :
                L=L+(C-array(p))/li
                Lm=Lm+li
        R=Lm*nbinv
        dC=M+L*R*nbinv-C
        C=C+dC
        if dot(dC,dC)<precision2:
            return [C.tolist(),R]
    return [C.tolist(),R]

def getMaxCartesianVelocity2(rigidDof=None, affineDof=None):
    """ @return the maximal cartesian velocity square norl from rigidDof and affineDof
    rigidDof, affineDof are MechanicalObject
    """
    vMax = 0.0
    if not rigidDof is None:
        for v in rigidDof.velocity:
            vNorm2 = v[0]**2 + v[1]**2  + v[2]**2
            if vNorm2>vMax: vMax = vNorm2
    if not affineDof is None:
        for v in affineDof.velocity:
            vNorm2 = v[0]**2 + v[1]**2  + v[2]**2
            if vNorm2>vMax: vMax = vNorm2
    return vMax
