import os
import sys
import json
import pathlib
import slicer
import subprocess
import shutil

def main(configFile):
   print('Running with {}'.format(configFile))
   with open(configFile) as f:
      config=json.load(f)

   fsetup=os.path.join(os.path.expanduser('~'),'.labkey','setup.json')
   with open(fsetup) as f:
      setup=json.load(f)

   sys.path.append(setup['paths']['nixWrapper'])
   import nixWrapper

   nixWrapper.loadLibrary('labkeyInterface')
   import labkeyInterface
   import labkeyDatabaseBrowser
   import labkeyFileBrowser
   

   nixWrapper.loadLibrary('orthancInterface')
   import orthancInterface
   import orthancDatabaseBrowser
   import orthancFileBrowser


   net=labkeyInterface.labkeyInterface()
   fnet=os.path.join(os.path.expanduser('~'),'.labkey',config['network'])
   net.init(fnet)
   net.getCSRF()
   fb=labkeyFileBrowser.labkeyFileBrowser(net)
   db=labkeyDatabaseBrowser.labkeyDB(net)


   onet=orthancInterface.orthancInterface()
   onet.init(fnet)
   ofb=orthancFileBrowser.orthancFileBrowser(onet)
   odb=orthancDatabaseBrowser.orthancDB(onet)

   qFilter=[]  
   ds=db.selectRows(config['project'],config['schemaName'],config['queryName'],qFilter)


   try:
      rows=ds['rows']
   except KeyError:
      print('No rows returned')
      return
   for r in rows:
      print("Loading {}/{}".format(patientId(r,config),visitId(r,config)))
      nNames=nodeNames(r,config)
      rPath=fb.formatPathURL(config['project'],outputDir(r,config))
      rPath+='/'+nNames['CT']+'.nrrd'
      entryDone=fb.entryExists(rPath)
      if entryDone:
         try:
            if not config['recalculate']:
               print('Entry done')
               continue
         except KeyError:
            print('Entry done')
            continue

         print('Forced recalculation')

            
      print('{} available:{}'.format(nNames['CT'],fb.entryExists(rPath)))
      #loadPatient into slicer
      patient=loadPatient(ofb,r,config)
      #convert to nodes
      addCT(r,patient,config)
      addFrames(r,patient,config)
      addDummyInputFunction(r,patient,config)
      
      nodes=slicer.mrmlScene.GetNodesByClass('vtkMRMLScalarVolumeNode')
      print('Nodes')
      for node in nodes:
         print('\t{}'.format(node.GetName()))
         storeNode(fb,r,config,node)

      nodes=slicer.mrmlScene.GetNodesByClass('vtkMRMLDoubleArrayNode')
      print('Nodes (double array)')
      for node in nodes:
         print('\t{}'.format(node.GetName()))
         storeNode(fb,r,config,node)
      clearNodes(r,config)

      
def patientId(row,config):
   return row[config['ParticipantField']]

def visitId(row,config):
   return row['visitName']

def code(row,config):
   return '{}_{}'.format(patientId(row,config),visitId(row,config))

def pathList(row,config):
   return [config['baseDir'],patientId(row,config),visitId(row,config)]

def outputDir(row,config):
   return '/'.join(pathList(row,config))

def localDir(row,config):
   return os.path.join(config['tempDir'],code(row,config))

def nodeNames(row,config):
   names={}
   x=code(row,config)
   names['NM']=['{}_Volume{}'.format(x,i) for i in range(20)]
   names['CT']='{}_CT'.format(x)
   return names

def clearNodes(row,config):
   nodes=slicer.mrmlScene.GetNodesByClass('vtkMRMLScalarVolumeNode')
   nodes1=slicer.mrmlScene.GetNodesByClass('vtkMRMLDoubleArrayNode')
   for n in nodes1:
      nodes.AddItem(n)

   res=[slicer.mrmlScene.RemoveNode(f) for f in nodes]

def loadPatient(ofb,r,config):
   sys.path.append('../pythonScripts')
   import parseDicom
   import vtkInterface

   pd=parseDicom.parseDicom()

   masterPath=downloadAndUnzip(ofb,r,"nmMaster",config)
   pd.readMasterDirectory(masterPath,False)
   print('Time [{} .. {}]'.format(pd.frame_start,pd.frame_stop))
   nmPath=downloadAndUnzip(ofb,r,"nmCorrData",config)
   frame_data, frame_time, frame_duration, frame_origin, \
       frame_pixel_size, frame_orientation=\
       pd.readNMDirectory(nmPath,False)
   print('Frame time {}'.format(frame_time))
   ctPath=downloadAndUnzip(ofb,r,"ct",config)
   ct_data,ct_origin,ct_pixel_size, \
      ct_orientation=pd.readCTDirectory(ctPath,False)
   print('CT pixel {}'.format(ct_pixel_size))

   ct_orientation=vtkInterface.completeOrientation(ct_orientation)
   frame_orientation=vtkInterface.completeOrientation(frame_orientation)

   ct={'data':ct_data,'origin':ct_origin,'pixel_size':ct_pixel_size,
      'orientation':ct_orientation}
   nm={'data':frame_data,'time':frame_time,'duration':frame_duration,
      'origin':frame_origin,'pixel_size':frame_pixel_size,
      'orientation':frame_orientation}

   return {'CT':ct,'NM':nm}
   
   print('Done')

def downloadAndUnzip(ofb,r,xcode,config):
   orthancId=r[xcode+'OrthancId']
   fileCode='{}_{}'.format(code(r,config),xcode)
   zipFile=os.path.join(config['tempDir'],fileCode+'.zip')
   ofb.getZip('series',orthancId,zipFile)
   zipDir=localDir(r,config)
   try:
      os.mkdir(zipDir)
   except FileExistsError:
      shutil.rmtree(zipDir)
   try:
      outTxt=subprocess.check_output(["unzip","-d",zipDir,"-xj",zipFile])
   except subprocess.CalledProcessError:
      print("unzip failed for {}".format(zipFile))
      return ""

   return zipDir

def addCT(r,patient,config):
   sys.path.append('../pythonScripts')
   import vtkInterface
   ct=patient['CT']
   vtkData=vtkInterface.numpyToVTK(ct['data'],ct['data'].shape)
   nodeName=code(r,config)+'_CT'
   addNode(nodeName,vtkData,ct['origin'],ct['pixel_size'],ct['orientation'],0)

def addFrames(r,patient,config):
   sys.path.append('../pythonScripts')
   import vtkInterface
   #convert data from numpy.array to vtkImageData
   #use time point it
   x=code(r,config)
   nm=patient['NM']
   print("NFrames: {}".format(nm['data'].shape[3]))
   for it in range(0,nm['data'].shape[3]):
      frame_data=nm['data'][:,:,:,it];
      nodeName=x+'_Volume'+str(it)
      vtkData=vtkInterface.numpyToVTK(frame_data,frame_data.shape)
      addNode(nodeName,vtkData,nm['origin'],nm['pixel_size'],nm['orientation'],1)



def addNode(nodeName,v,origin,pixel_size,orientation,dataType):
   #origin,orientation in lps
   #dataType=0 is CT (to background)
   #dataType=1 is SPECT, view not adjusted, foreground, 
   newNode=slicer.vtkMRMLScalarVolumeNode()
   newNode.SetName(nodeName)
   v.SetOrigin([0,0,0])
   v.SetSpacing([1,1,1])
   ijkToRAS = vtk.vtkMatrix4x4()
   #think how to do this with image orientation
   #orientation from lps to ras
   rasOrientation=[-orientation[i] if (i%3 < 2) else orientation[i]
        for i in range(0,len(orientation))]
   #origin from lps to ras
   rasOrigin=[-origin[i] if (i%3<2) else origin[i] for i in range(0,len(origin))]

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

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

   newNode.SetIJKToRASMatrix(ijkToRAS)
   newNode.SetAndObserveImageData(v)
   slicer.mrmlScene.AddNode(newNode)

def addDummyInputFunction(r,patient,config):

   nm=patient['NM']
   n=nm['data'].shape[3]

   dnsNodeName=code(r,config)+'_Dummy'
   dns = slicer.mrmlScene.GetNodesByClassByName('vtkMRMLDoubleArrayNode',dnsNodeName)
   if dns.GetNumberOfItems() == 0:
      dn = slicer.mrmlScene.AddNode(slicer.vtkMRMLDoubleArrayNode())
      dn.SetName(dnsNodeName)
   else:
      dn = dns.GetItemAsObject(0)

   dn.SetSize(n)

   ft=nm['time']
   dt=nm['duration']
   for i in range(0,n):
      fx=ft[i]
      fy=dt[i]
      dn.SetValue(i, 0, fx)
      dn.SetValue(i, 1, fy)
      dn.SetValue(i, 2, 0)
      print('{} ({},{})'.format(i,fx,fy))

def storeNode(fb,row,config,node):

   suffix=".nrrd"
   if node.__class__.__name__=="vtkMRMLDoubleArrayNode":
      suffix=".mcsv"
   if (node.__class__.__name__=="vtkMRMLTransformNode" or \
         node.__class__.__name__=="vtkMRMLGridTransformNode"):
      suffix=".h5"

   fileName=node.GetName()+suffix

   localPath=os.path.join(localDir(row,config),fileName)

   slicer.util.saveNode(node,localPath)
   print("Stored to: {}".format(localPath))
   labkeyPath=fb.buildPathURL(config['project'],pathList(row,config))
   print ("Remote: {}".format(labkeyPath))
   #checks if exists
   remoteFile=labkeyPath+'/'+fileName
   fb.writeFileToFile(localPath,remoteFile)


   






def readPatient(fb,localDir,project,patientId):

    rDir=fb.formatPathURL(project,'/'.join([patientId]))
    lDir=os.path.join(localDir,patientId)
    if not os.path.isdir(lDir):
        os.makedirs(lDir)

    ok,files=fb.listRemoteDir(rDir)
    locFiles=[]
    for f in files:
        print(f)
        p=pathlib.Path(f)
        localFile=os.path.join(lDir,p.name)
        fb.readFileToFile(f,localFile)
        locFiles.append(localFile)
    print('Done')
    return locFiles

if __name__=='__main__':
    main(sys.argv[1])
    quit()

