import sys
import os
import json
import pydicom
import importlib
import numpy
import Elastix 
import vtk
import DicomRtImportExportPlugin
import glob
import zipfile

#just load nix suite
sys.path.append(os.path.join(os.path.expanduser('~'),'software','src','nixSuite','wrapper'))
import nixWrapper

nixWrapper.loadLibrary('labkeyInterface')
import labkeyInterface

#DB access at ONKO-NIX
net=labkeyInterface.labkeyInterface()
fconfig=os.path.join(os.path.expanduser('~'),'.labkey','network.json')
net.init(fconfig)

#DMR are at merlin
netDMR=labkeyInterface.labkeyInterface()
fconfigDMR=os.path.join(os.path.expanduser('~'),'.labkey','merlin.json')
netDMR.init(fconfigDMR)

import labkeyDatabaseBrowser
db=labkeyDatabaseBrowser.labkeyDB(net)

import labkeyFileBrowser
fb=labkeyFileBrowser.labkeyFileBrowser(net)
fbDMR=labkeyFileBrowser.labkeyFileBrowser(netDMR)

#dicom tools
nixWrapper.loadLibrary('DICOMtools')
try:
    importlib.reload(importDicom)
except NameError:
    import importDicom
    
try:
    importlib.reload(loadDicom)
except NameError:
    import loadDicom

    
nixWrapper.loadLibrary('EMBRACE')
try:
    importlib.reload(loadPatient)
except NameError:
    import loadPatient

project='EMBRACE/studija'
iDicom=importDicom.importDicom()
lDicom=loadDicom.loadDicom()

try:
    del elastix
except NameError:
    pass

elastix = Elastix.ElastixLogic()

registrationIndex = 29 #targeted EMBRACE
registrationIndex = 30 #rigid EMBRACE
RegistrationPresets_ParameterFilenames = 5
parameterFilenames = elastix.getRegistrationPresets()[registrationIndex]\
    [RegistrationPresets_ParameterFilenames]

ids=range(1,87)
ids=range(6,7)

shn = slicer.vtkMRMLSubjectHierarchyNode.GetSubjectHierarchyNode(slicer.mrmlScene)
baseDir=os.path.join('z:\\','EMBRACE')
if not os.path.isdir(baseDir):
    os.makedirs(baseDir)

transformLogic = slicer.vtkSlicerTransformLogic()

exporter = DicomRtImportExportPlugin.DicomRtImportExportPluginClass()

for id in ids:

    label = 'LJU'+'{:03d}'.format(id)
    loadDicom.clearNodes()
    
    nodes=loadPatient.load(db,fb,fbDMR,iDicom,lDicom,project,label)
    print('Segmentations: {}'.format(len(slicer.util.getNodesByClass("vtkMRMLSegmentationNode"))))
    
    #debug 
    #seg=nodes['CTRS'].GetSegmentation()
    #n=seg.GetNumberOfSegments()
    #for i in numpy.arange(n):
    #    sId=seg.GetNthSegmentID(int(i))
    #    print(sId)
    #continue

    try:
        fixedVolumeNode = nodes['DMR']
        movingVolumeNode = nodes['CT']
    except KeyError:
        continue

    transformNodeName = label+'_T2_DF'
    #
    transformNode=slicer.vtkMRMLGridTransformNode()
    transformNode.SetName(transformNodeName)
    slicer.mrmlScene.AddNode(transformNode)
    
    elastix.registerVolumes(fixedVolumeNode, movingVolumeNode, \
                            parameterFilenames = parameterFilenames, outputVolumeNode = None, 
                            outputTransformNode = transformNode, \
                            fixedVolumeMaskNode = None, movingVolumeMaskNode = None, 
                            forceDisplacementFieldOutputTransform = True)

    
    #dataManager.export(network, 'EMBRACE/Studija', patientID)
    #create output dicom with MRI and migrated CTRS
    oDir=os.path.join(baseDir,label)
    if not os.path.isdir(oDir):
        os.makedirs(oDir)
    #transformNode=slicer.util.loadTransform(os.path.join(oDir,'Transform_R.h5'))
    
        
    #remove all prior dicom files
    dcmList=glob.glob(os.path.join(oDir,'*.dcm'))
    for f in dcmList:
        os.remove(f)
    
    for mode in ['CT','DMR']:
        slicer.util.saveNode(nodes[mode],os.path.join(oDir,label+'_'+mode+'.nrrd'))
    slicer.util.saveNode(transformNode,os.path.join(oDir,label+'_DF.nrrd'))
    seg=nodes['CTRS'].GetSegmentation()
    n=seg.GetNumberOfSegments()
    
    nodes['CTRS'].SetAndObserveTransformNodeID(transformNode.GetID())
    
    #transformLogic.hardenTransform(nodes['CTRS'])
    for i in numpy.arange(n):
        sId=seg.GetNthSegmentID(int(i))
        print(sId)
        seg.GetSegment(sId)
        fname=os.path.join(oDir,label+'_{}.nrrd'.format(sId))
        labelmapVolumeNode = slicer.mrmlScene.AddNewNodeByClass('vtkMRMLLabelMapVolumeNode')
        
        idArray = vtk.vtkStringArray()
        idArray.InsertNextValue(sId)
        slicer.vtkSlicerSegmentationsModuleLogic.ExportSegmentsToLabelmapNode(
                  nodes['CTRS'], idArray, labelmapVolumeNode, nodes['DMR'])
        slicer.util.saveNode(labelmapVolumeNode, fname)
    
    exportables=[]
    for mode in ['CTRS','DMR']:
        itemId=shn.GetItemByDataNode(nodes[mode])
        exportables.extend(exporter.examineForExport(itemId))
                                           
    for exp in exportables:
        exp.directory=oDir
        exp.setTag('PatientID',label)
        exp.setTag('PatientName',label)
        exp.setTag('StudyDescription','contoursAligned')
        exp.setTag('PatientID',label)
        if exp.tag('Modality')=='CT':
            exp.setTag('Modality','MR')
            exp.setTag('ScanningSequence',nodes['metadataDMR']['scanningSequence'])
            exp.setTag('SequenceVariant',nodes['metadataDMR']['sequenceVariant'])
            
    
    exporter.export(exportables)
   

ids=range(3,5)
for id in ids:
    label = 'LJU'+'{:03d}'.format(id)
    oDir=os.path.join(baseDir,label)
    dcmList=glob.glob(os.path.join(oDir,'*.dcm'))
    ozip=os.path.join(baseDir,label+'.zip')
    with zipfile.ZipFile(ozip, 'w') as zf:
        for f in dcmList:
            zf.write(f)

    rDir=fb.buildPathURL(project,[label])
    rZip='{}/{}.zip'.format(rDir,label)
    fb.writeFileToFile(ozip,rZip)