import dicom
import vtkInterface as vi
import os
import vtk, qt, ctk, slicer
from slicer.ScriptedLoadableModule import *
import slicerNetwork
import json
import loadDicom
import DICOMLib

class importDicom(slicer.ScriptedLoadableModule.ScriptedLoadableModule):
    def __init__(self,parent):
        slicer.ScriptedLoadableModule.ScriptedLoadableModule.__init__(self, parent)
        self.className="loadPatient"
        self.parent.title="loadPatient"

#equivalent of loadable + labkey interface
class dicomSeries():
    def __init__(self):
        self.data = []
        self.idx = []
        self.z = []
        self.pixel_size = [0,0,0]
        self.lpsOrigin = [0,0,0]
        self.lpsOrigin[2]=1e30
        self.lpsOrientation=[0,0,0,0,0,0]


    def getfile(self,net,relativePath):
        return net.readFileToBuffer(relativePath)

    def addFile(self,f):
        try:
            self.files.append(f)
        except:
            self.files=[f]

    def setLabel(self,label):
        self.label=label

    def getLabel(self):
        try:
            return self.label
        except:
            return None

    def setMetadata(self,key,value):
        try:
            self.metadata[key]=value
        except:
            self.metadata={key:value}

    def getMetadata(self):
        try:
            return self.metadata
        except:
            return {}

    def load(self,net):

        for f in self.files:
            print '{}:'.format(f)
            fileBuffer=self.getfile(net,f)
            self.loadFile(fileBuffer)

        nz=len(self.idx)
        sh=self.data[-1].shape
        sh_list=list(sh)
        sh_list.append(nz)
        data_array=np.zeros(sh_list)

        for k in range(0,nz):
            kp=int(np.round((self.z[k]-self.center[2])/self.pixel_size[2]))
            data_array[:,:,kp]=np.transpose(self.data[k])

        try:
            nodeName='Series'+self.label
        except:
            print('Could not set series label')
            nodeName='UnknownSeries'

        newNode=slicer.vtkMRMLScalarVolumeNode()
        newNode.SetName(nodeName)

        ijkToRAS = vtk.vtkMatrix4x4()
        #think how to do this with image orientation

        rasOrientation=[-self.lpsOrientation[i] if (i%3 < 2) else self.lpsOrientation[i]
          for i in range(0,len(self.lpsOrientation))]
        rasOrigin=[-self.lpsOrigin[i] if (i%3<2) else self.lpsOrigin[i] for i in range(0,len(self.lpsOrigin))]

        for i in range(0,3):
             for j in range(0,3):
                 ijkToRAS.SetElement(i,j,self.pixel_size[i]*rasOrientation[3*j+i])

             ijkToRAS.SetElement(i,3,rasOrigin[i])

        newNode.SetIJKToRASMatrix(ijkToRAS)

        v=vtk.vtkImageData()
        v.GetPointData().SetScalars(
            vtk.util.numpy_support.numpy_to_vtk(
                np.ravel(self.data,order='F'),deep=True, array_type=vtk.VTK_FLOAT))
        v.SetOrigin(0,0,0)
        v.SetSpacing(1,1,1)
        v.SetDimensions(self.data.shape)

        newNode.SetAndObserveImageData(v)
        slicer.mrmlScene.AddNode(newNode)
        volume={'node':newNode,'metadata':self.metadata}
        return volume

    def loadFile(self,fileBuffer):

        plan=dicom.read_file(fileBuffer)

        self.data.append(plan.pixel_array)
        self.idx.append(plan.InstanceNumber)
        self.z.append(plan.ImagePositionPatient[2])


        #pixelSize
        pixel_size=[plan.PixelSpacing[0],plan.PixelSpacing[1],
                    plan.SliceThickness]

        for i in range(0,3):
            if self.pixel_size[i] == 0:
                self.pixel_size[i] = float(pixel_size[i])
            if abs(self.pixel_size[i]-pixel_size[i]) > 1e-3:
                print 'Pixel size mismatch {.2f}/{.2f}'.format(self.pixel_size[i],
                        pixel_size[i])

        #origin
        for i in range(0,2):
            if self.lpsOrigin[i] == 0:
                self.lpsOrigin[i] = float(plan.ImagePositionPatient[i])
            if abs(self.lpsOrigin[i]-plan.ImagePositionPatient[i]) > 1e-3:
                print 'Image center mismatch {.2f}/{.2f}'.format(self.lpsOrigin[i],
                    plan.ImagePositionPatient[i])
        #not average, but minimum (!) why??

        if plan.ImagePositionPatient[2]<self.lpsOrigin[2]:
            self.lpsOrigin[2]=plan.ImagePositionPatient[2]


        #orientation
        for i in range(0,6):
            if self.lpsOrientation[i] == 0:
                self.lpsOrientation[i] = float(plan.ImageOrientationPatient[i])
            if abs(self.lpsOrientation[i]-plan.ImageOrientationPatient[i]) > 1e-3:
                print 'Image orientation mismatch {0:.2f}/{1:.2f}'.format(self.lpsOrientation[i],
                        plan.ImageOrientationPatient[i])

        return True

class importDicomLogic(slicer.ScriptedLoadableModule.ScriptedLoadableModuleLogic):
    def __init__(self,parent):
       slicer.ScriptedLoadableModule.ScriptedLoadableModuleLogic.__init__(self, parent)
       self.tag={
        'studyInstanceUid':0x0020000d,
        'seriesInstanceUid':0x0020000e,
        'patientId':0x00100020,
        'patientName':0x00100010,
        'sequenceName':0x00180024,
        'seriesNumber':0x00200011,
        'percentPhaseFieldOfView':0x00180094,
        'modality': 0x00080060,
        'patientSex': 0x00100040,
        'patientBirthDate': 0x00100030,
        'patientComments': 0x00104000,
        'studyDescription': 0x00081030,
        'studyDate': 0x00080020,
        'studyId': 0x00200010,
        'studyTime': 0x00080030,
        'frameOfReferenceInstanceUid':0x00200052}

    def loadVolumes(self,net,directory,filter):
        #mimic examineForImport
        seriesList=self.examineForImport(net,directory,filter)
        print("Got {} series").format(len(seriesList))
        volumes=[]

        for s in seriesList:
            try:
                volumes.append(s.load(net))
                #often fails, e.g. JPEGLossles
            except:
                loadable=DICOMLib.DICOMLoadable()
                loadable.name='Series'+str(s.getLabel())
                print("Loading for {} number of files (pre-load) {}").format(loadable.name,len(s.files))
                loadable.files=[net.DownloadFileToCache(f) for f in s.files]
                print("Loading for {} number of files (pre-sort) {}").format(loadable.name,len(loadable.files))

                loadable.files,distances,loadable.warning=DICOMLib.DICOMUtils.getSortedImageFiles(loadable.files,1e-3)

                print("Loading for {} number of files {}").format(loadable.name,len(loadable.files))
                try:
                    volumeNode=self.volumePlugin.load(loadable)
                except:
                    self.volumePlugin=slicer.modules.dicomPlugins['DICOMScalarVolumePlugin']()
                    volumeNode=self.volumePlugin.load(loadable)
                volume={'node':volumeNode,'metadata':s.getMetadata()}
                volumes.append(volume)


        return volumes


    def listdir(self,net,relativeDirectory):
        return net.listRelativeDir(relativeDirectory)

    def getfile(self,net,relativePath):
        return net.readFileToBuffer(relativePath)

    def examineForImport(self,net,directory,filter):
        #split by series
        files=self.listdir(net,directory)
        seriesList=[]

        seriesList.append(dicomSeries())
        series=seriesList[0]

        for f in files:
            fileBuffer=self.getfile(net,f)
            #validate
            try:
                plan = dicom.read_file(fileBuffer)
            except:
                print ("{}: Not a dicom file")
                continue

            #determine validity first
            fileValid=True
            seriesKey=None
            for key in filter:
                if filter[key]==None:
                    continue
                if filter[key]=='SeriesLabel':
                    seriesTag=self.tag[key]
                    continue

                v=plan[self.tag[key]].value
                if not v==filter[key]:
                    print('Filter mismatch {}{:x}: {}/{}').format(key,self.tag[key],v,filter[key])
                    fileValid=False

            if not fileValid:
                continue

            #determine serieslabel second
            seriesLabel=plan[seriesTag].value
            if series.getLabel()==seriesLabel:
                series.addFile(f)
                continue

            #add new series
            if not series.getLabel()==None:
                seriesList.append(dicomSeries())
                series=seriesList[-1]

            #set series parameters
            series.addFile(f)
            series.setLabel(seriesLabel)
            for key in filter:
                if not filter[key]==None:
                    continue
                v=plan[self.tag[key]].value
                series.setMetadata(key,v)

        return seriesList