# -*- coding: utf-8 -*-

"""Classes for working with Dicom studies
"""
import copy
import os
import pydicom as dicom
from pydicom.uid import generate_uid
from datetime import datetime
from tqdm import tqdm
import json
from multiprocessing import Pool
import numpy as np
import shutil
import matplotlib.pyplot as plt
# Local imports 
import spydcmtk.dcmTools as dcmTools
import spydcmtk.dcmVTKTK as dcmVTKTK
from spydcmtk.spydcm_config import SpydcmTK_config




## =====================================================================================================================
##        CLASSES
## =====================================================================================================================
class DicomSeries(list):
    """Extends a list of ds (pydicom dataset) objects.
    """
    def __init__(self, dsList=None, OVERVIEW=False, HIDE_PROGRESSBAR=False, FORCE_READ=False, SAFE_NAME_MODE=False):
        """Initialise DicomSeries with list of pydicom datasets.

        Args:
            dsList (list, optional): list of pydicom dataset. Defaults to None.
            OVERVIEW (bool, optional): Set True to not read pixel data, for efficiency in applications of meta 
                                        data only interest. 
                                        Pass to pydicom.filereader.dcmread parameter:stop_before_pixels
                                        Defaults to False.
            HIDE_PROGRESSBAR (bool, optional): Set True to hide tqdm progressbar. Defaults to False.
            FORCE_READ (bool, optional): Force read dicom files even if they do not conform to DICOM standard. 
                                            Pass to pydicom.filereader.dcmread parameter:force
                                            Defaults to False.
            SAFE_NAME_MODE (bool, optional): Set writing mode to use UIDs for file naming. Defaults to False.
        """
        if dsList is None:
            dsList = []
        self.OVERVIEW = OVERVIEW
        self.HIDE_PROGRESSBAR = HIDE_PROGRESSBAR
        self.FORCE_READ = FORCE_READ
        self.SAFE_NAME_MODE = SAFE_NAME_MODE
        self.NOT_FULLY_LOADED = False
        list.__init__(self, dsList)

    def __str__(self):
        """Return output of getSeriesOverview method
        """
        return ' '.join([str(i) for i in self.getSeriesOverview()[1]])


    def __getitem__(self, item):
        if isinstance(item, slice):
            result = super().__getitem__(item)
            return DicomSeries(result)
        else:
            return super().__getitem__(item)

    @classmethod
    def _setFromDictionary(cls, dicomDict, OVERVIEW=False, HIDE_PROGRESSBAR=False, FORCE_READ=False):
        dStudyList = ListOfDicomStudies.setFromDcmDict(dicomDict, OVERVIEW=OVERVIEW, HIDE_PROGRESSBAR=HIDE_PROGRESSBAR, FORCE_READ=FORCE_READ)
        if len(dStudyList) > 1:
            raise ValueError('More than one study found - use ListOfDicomStudies class')
        try:
            dStudy = dStudyList[0]
        except IndexError:
            raise ValueError('No DICOMS found - please check your inputs')
        if len(dStudy) > 1:
            raise ValueError('More than one series found - use DicomStudy class')
        return cls(dStudy[0], OVERVIEW=OVERVIEW, HIDE_PROGRESSBAR=HIDE_PROGRESSBAR, FORCE_READ=FORCE_READ)

    @classmethod
    def setFromDirectory(cls, dirName, OVERVIEW=False, HIDE_PROGRESSBAR=False, FORCE_READ=False, ONE_FILE_PER_DIR=False):
        """Initialise object from directory of dicom files.

        Args:
            dirName (str): Path to directory of dicom files
            OVERVIEW (bool, optional): Set attribute OVERVIEW. Defaults to False.
            HIDE_PROGRESSBAR (bool, optional): Set attribute HIDE_PROGRESSBAR. Defaults to False.
            FORCE_READ (bool, optional): Set attribute FORCE_READ. Defaults to False.
            ONE_FILE_PER_DIR (bool, optional): Read only one file per directory for fast applications
                                                with meta interest only. Defaults to False.

        Raises:
            ValueError: If dicom files belonging to more than one study are found.
            ValueError: If dicom files belonging to more than one series are found.

        Returns:
            DicomSeries: An instance of DicomSeries class.
        """
        dicomDict = dcmTools.organiseDicomHeirarchyByUIDs(dirName, HIDE_PROGRESSBAR=HIDE_PROGRESSBAR, FORCE_READ=FORCE_READ, ONE_FILE_PER_DIR=ONE_FILE_PER_DIR, OVERVIEW=OVERVIEW)
        obj = DicomSeries._setFromDictionary(dicomDict, OVERVIEW=OVERVIEW, HIDE_PROGRESSBAR=HIDE_PROGRESSBAR, FORCE_READ=FORCE_READ)
        obj.NOT_FULLY_LOADED = ONE_FILE_PER_DIR
        return obj

    @classmethod
    def setFromFileList(cls, fileList, OVERVIEW=False, HIDE_PROGRESSBAR=False, FORCE_READ=False):
        """Initialise object from list of files

        Args:
            fileList (list): List of full file names
            OVERVIEW (bool, optional): Set attribute OVERVIEW. Defaults to False.
            HIDE_PROGRESSBAR (bool, optional): Set attribute HIDE_PROGRESSBAR. Defaults to False.
            FORCE_READ (bool, optional): Set attribute FORCE_READ. Defaults to False.

        Raises:
            ValueError: If dicom files belonging to more than one study are found.
            ValueError: If dicom files belonging to more than one series are found.

        Returns:
            DicomSeries: An instance of DicomSeries class.
        """
        dicomDict = {}
        for iFile in fileList:
            dcmTools.readDicomFile_intoDict(iFile, dicomDict, FORCE_READ=FORCE_READ, OVERVIEW=OVERVIEW)
        return DicomSeries._setFromDictionary(dicomDict, OVERVIEW=OVERVIEW, HIDE_PROGRESSBAR=HIDE_PROGRESSBAR, FORCE_READ=FORCE_READ)

    def getRootDir(self):
        return os.path.split(self[0].filename)[0]

    def filterByTag(self, tagName, tagValue):
        return DicomSeries([i for i in self if i.getTag(tagName) == tagValue], 
                           OVERVIEW=self.OVERVIEW, 
                           HIDE_PROGRESSBAR=self.HIDE_PROGRESSBAR, 
                           FORCE_READ=self.FORCE_READ)

    def getDicomFullFileName(self, dsID=0):
        return self[dsID].filename

    def sortByInstanceNumber(self):
        self.sort(key=dcmTools.instanceNumberSortKey)

    def sortBySlice_InstanceNumber(self):
        self.sort(key=dcmTools.sliceLoc_InstanceNumberSortKey)

    def getTag(self, tag, dsID=0, ifNotFound='Unknown', convertToType=None):
        try:
            tt = self.getTagObj(tag, dsID)
            if convertToType is not None:
                return convertToType(tt.value)
            return tt.value
        except KeyError:
            return ifNotFound

    def setTags_all(self, tag, value):
        for ds in self:
            ds[tag].value = value

    def getTagObj(self, tag, dsID=0):
        try:
            if tag[:2] == '0x':
                tag = int(tag, 16)
        except TypeError:
            pass # Not string
        return self[dsID][tag]

    def getTagValuesList(self, tagList, RETURN_STR):
        """
        Build table with file name and value from each tag in list.
        Return list of lists
        """
        valueList = []
        for dsID, i in enumerate(self):
            subList = ['"%s"'%(str(i.filename))] + [str(self.getTag(t, dsID=dsID)) for t in tagList]
            valueList.append(subList)
        if RETURN_STR:
            return dcmTools._tagValuesListToString(tagList, valueList)
        return valueList

    def getTagListAndNames(self, tagList, dsID=0):
        names, vals = [], []
        for i in tagList:
            try:
                dataEle = self.getTagObj(i, dsID=dsID)
                iname = str(dataEle.keyword)
                if len(iname) == 0:
                    iname = str(dataEle.name)
                    iname = iname.replace('[','').replace(']','').replace(' ','')
                names.append(iname)
                vals.append(str(dataEle.value))
            except KeyError:
                names.append(str(i))
                vals.append('Unknown')
        return names, vals

    def tagsToJson(self, jsonFileOut):
        dOut = {}
        for ds in self:
            dd = ds.to_json_dict()
            dd.pop('7FE00010') # Remove actual PixelData
            dOut[ds.InstanceNumber] = dd
        dcmTools.writeDictionaryToJSON(jsonFileOut, dOut)
        return jsonFileOut


    def _loadToMemory(self):
        if self.NOT_FULLY_LOADED:
            rootDir = self.getRootDir()
            self.clear()
            self.extend([dicom.dcmread(os.path.join(rootDir, i), force=self.FORCE_READ) for i in tqdm(os.listdir(rootDir), disable=self.HIDE_PROGRESSBAR)])
            self.NOT_FULLY_LOADED = False


    def getSeriesOverview(self, tagList=SpydcmTK_config.SERIES_OVERVIEW_TAG_LIST):
        names, vals = self.getTagListAndNames(tagList)
        names.append('ImagesInSeries')
        if len(self) == 1: # ONLY READ ONE FILE PER DIR
            vals.append(dcmTools.countFilesInDir(self.getRootDir()))
        else:
            vals.append(len(self))
        return names, vals

    def getSeriesTimeAsDatetime(self):
        dos = self.getTag('SeriesDate', ifNotFound="19000101")
        tos = self.getTag('SeriesTime', ifNotFound="000000")
        try:
            return datetime.strptime(f"{dos} {tos}", "%Y%m%d %H%M%S.%f")
        except ValueError:
            return datetime.strptime(f"{dos} {tos}", "%Y%m%d %H%M%S")

    def getImagePositionPatient_np(self, dsID):
        ipp = self.getTag('ImagePositionPatient', dsID=dsID, ifNotFound=[0.0,0.0,0.0])
        return np.array(ipp)

    def getImageOrientationPatient_np(self, dsID=0):
        iop = self.getTag('ImageOrientationPatient', dsID=dsID)
        return np.array(iop)

    def yieldDataset(self):
        for ds in self:
            if self.OVERVIEW:
                yield  dicom.dcmread(ds.filename, stop_before_pixels=False, force=True)
            else:
                yield ds

    def isCompressed(self):
        return dcmTools._isCompressed(self[0])

    def getSeriesNumber(self):
        return int(self.getTag('SeriesNumber', ifNotFound=0))

    def getSeriesOutDirName(self, SE_RENAME={}):
        thisSeNum = self.getTag('SeriesNumber', ifNotFound='#')
        suffix = ''
        if (thisSeNum=="#") or self.SAFE_NAME_MODE:
            suffix += '_'+self.getTag("SeriesInstanceUID")
        for iTag in SpydcmTK_config.SERIES_NAMING_TAG_LIST:
            iVal = self.getTag(iTag, ifNotFound='')
            if len(iVal) > 0:
                suffix += '_'+iVal
        if thisSeNum in SE_RENAME.keys():
            return SE_RENAME[thisSeNum]
        return dcmTools.cleanString(f"SE{thisSeNum}{suffix}")

    # ----------------------------------------------------------------------------------------------------
    def removeTimes(self, timesIDs_ToRemove):
        K = int(self.getNumberOfSlicesPerVolume())
        self.sortBySlice_InstanceNumber()
        N = self.getNumberOfTimeSteps()
        c0 = 0
        dsToRm = []
        for k1 in range(K):
            for k2 in range(N):
                if k2 in timesIDs_ToRemove:
                    dsToRm.append(self[c0])
                c0 += 1
        self._rm_via_DS_List(dsToRm)


    def removeInstances(self, instanceIDs_to_remove):
        dsToRm = []
        for k1 in range(len(self)):
            if self.getTag("InstanceNumber", k1, convertToType=int) in instanceIDs_to_remove:
                dsToRm.append(self[k1])
        self._rm_via_DS_List(dsToRm)


    def _rm_via_DS_List(self, dsList):
        for iDS in dsList:
            self.remove(iDS)


    def deleteTag(self, tag, dsID):
        tagObj = self.getTagObj(tag, dsID)
        del tagObj

    # ----------------------------------------------------------------------------------------------------
    def resetUIDs(self, studyUID):
        """Reset SOPInstanceUID, SeriesInstanceUID and StudyInstaceUID (must pass last)

        Args:
            studyUID (str): UID - can use str(generate_uid())
        """
        seriesUID = str(generate_uid())
        for k1 in range(len(self)):
            self[k1].SOPInstanceUID = str(generate_uid())
            self[k1].SeriesInstanceUID = seriesUID
            self[k1].StudyInstanceUID = studyUID

    def anonymise(self, anonName, anonPatientID, anon_birthdate=True, remove_private_tags=True):
        """Anonymise series inplace

        Args:
            anonName (str): New anonymise name (can be empty str)
            anonPatientID (str): New anon patient ID, if None then no action taken.
            anon_birthdate (bool, optional): To remove birthdate from tags. Defaults to True.
            remove_private_tags (bool, optional): To remove private tags from dataset. Defaults to True. 
                In some cases this will need to be set False to keep tags necessary for post processing (e.g. some DTI tags). 
                It is up to the user to confirm complete anonymisation in these cases. 
        """
        def PN_callback(ds, data_element):
            """Called from the dataset "walk" recursive function for all data elements."""
            if data_element.VR == "PN":
                data_element.value = 'anonymous'
            if "Institution" in data_element.name:
                data_element.value = 'anonymous'
            if data_element.name == "Patient's Name":
                data_element.value = anonName
            if (data_element.name == "Patient ID") and (anonPatientID is not None):
                data_element.value = anonPatientID

        # Remove patient name and any other person names using a callback to walk over data elements for each ds
        for dataset in self:
            try:
                dataset.walk(PN_callback)
            except TypeError: 
                pass
            # Other Tags
            for name in ['OtherPatientIDs', 'OtherPatientIDsSequence', 'PatientAddress']:
                if name in dataset:
                    delattr(dataset, name)
            if anon_birthdate:
                for name in ['PatientBirthDate']:
                    if name in dataset:
                        dataset.data_element(name).value = ''
            if remove_private_tags:
                try:
                    dataset.remove_private_tags()
                except TypeError:
                    pass

    def writeToOrganisedFileStructure(self, studyOutputDir, SAFE_NAMING_CHECK=True, seriesOutDirName=None):
        """ Recurse down directory tree - grab dicoms and move to new
            hierarchical folder structure rooted at 'studyOutputDir'
        """
        self._loadToMemory()
        if SAFE_NAMING_CHECK:
            self.checkIfShouldUse_SAFE_NAMING()
        ADD_TRANSFERSYNTAX = False
        if seriesOutDirName is None:
            seriesOutDirName = self.getSeriesOutDirName()
        seriesOutputDir = os.path.join(studyOutputDir, seriesOutDirName)
        seriesOutputDirTemp = seriesOutputDir+".WORKING"
        if os.path.isdir(seriesOutputDir):
            os.rename(seriesOutputDir, seriesOutputDirTemp)
        if self.FORCE_READ:
            ADD_TRANSFERSYNTAX = True # need to add if required force read. 
        TO_DECOMPRESS = self.isCompressed()
        for ds in tqdm(self.yieldDataset(), total=len(self), disable=self.HIDE_PROGRESSBAR):
            if TO_DECOMPRESS:
                try:
                    ds.decompress()
                except AttributeError:
                    print(f'Error with file in {seriesOutputDirTemp}')
                    continue
            if ADD_TRANSFERSYNTAX:
                ds.file_meta.TransferSyntaxUID = '1.2.840.10008.1.2.1'
            destFile = dcmTools.writeOut_ds(ds, seriesOutputDirTemp, SAFE_NAMING=self.SAFE_NAME_MODE)
        # ON COMPLETION RENAME OUTPUTDIR
        os.rename(seriesOutputDirTemp, seriesOutputDir)
        return seriesOutputDir

    def _generateFileName(self, tagsToUse, extn):
        if type(tagsToUse) == str:
            fileName = tagsToUse
        else:
            fileName = '_'.join([str(self.getTag(i)) for i in tagsToUse])
        fileName = dcmTools.cleanString(fileName)
        if (len(extn) > 0) and (not extn.startswith('.')):
            extn = '.'+extn
        return fileName+extn

    def writeToNII(self, outputPath, outputNamingTags=('PatientName', 'SeriesNumber', 'SeriesDescription')):
        """Will write DicomSeries as nii.gz file (uses dcm2niix found on local system)

        Args:
            outputPath (str): Directory where to save output (if full pathname, with extn, is given then this is used)
            outputNamingTags (tuple, optional): Tags to use to name output. If outputPath contains extension then not used. Defaults to ('PatientName', 'SeriesNumber', 'SeriesDescription').

        Returns:
            str: Name of files saved
        """
        if outputPath.endswith('nii.gz'):
            outputPath, fileName = os.path.split(outputPath)
        else:
            fileName = self._generateFileName(outputNamingTags, '.nii.gz')
        return dcmTools.writeDirectoryToNII(self.getRootDir(), outputPath, fileName=fileName)

    def writeToVTI(self, outputPath, outputNamingTags=('PatientName', 'SeriesNumber', 'SeriesDescription'), TRUE_ORIENTATION=False):
        """Write DicomSeries as VTK ImageData (`*.vti`)

        Args:
            outputPath (str): Directory to save in or full filename to save
            outputNamingTags (tuple, optional): Tags to use for naming, only used if full outputpath not given. Defaults to ('PatientName', 'SeriesNumber', 'SeriesDescription').
            TRUE_ORIENTATION (bool, optional): To apply matrix to VTI data. Defaults to False.

        Returns:
            str: Name of file saved
        """
        if outputPath.endswith('vti') or outputPath.endswith('pvd'):
            outputPath, fileName = os.path.split(outputPath)
            fileName, _ = os.path.splitext(fileName)
        else:
            fileName = self._generateFileName(outputNamingTags, '')
        vtiDict = self.buildVTIDict(TRUE_ORIENTATION=TRUE_ORIENTATION)
        return dcmVTKTK.writeVTIDict(vtiDict, outputPath, fileName)

    def writeToVTS(self, outputPath, outputNamingTags=('PatientName', 'SeriesNumber', 'SeriesDescription')):
        """Write DicomSeries as VTK StructuredImageData (`*.vts`)

        Args:
            outputPath (str): Directory to save in or full filename to save.
            outputNamingTags (tuple, optional): Tags to use for naming, only used if full outputpath not given. Defaults to ('PatientName', 'SeriesNumber', 'SeriesDescription').

        Returns:
            str: Name of file saved
        """
        if outputPath.endswith('vts') or outputPath.endswith('pvd'):
            outputPath, fileName = os.path.split(outputPath)
            fileName, _ = os.path.splitext(fileName)
        else:
            fileName = self._generateFileName(outputNamingTags, '')
        vtsDict = self.buildVTSDict(outputPath)
        return dcmVTKTK.fIO.writeVTK_PVD_Dict(vtsDict, outputPath, filePrefix=fileName, fileExtn='vts', BUILD_SUBDIR=True)

    def buildVTSDict(self, outputPath=None):
        A, patientMeta = self.getPixelDataAsNumpy()
        return dcmVTKTK.arrToVTS(A, patientMeta, self[0], outputPath)

    def buildVTIDict(self, TRUE_ORIENTATION=False):
        A, patientMeta = self.getPixelDataAsNumpy()
        return dcmVTKTK.arrToVTI(A, patientMeta, self[0], TRUE_ORIENTATION=TRUE_ORIENTATION)

    @property
    def sliceLocations(self):
        return sorted([float(self.getTag('SliceLocation', i, ifNotFound=0.0)) for i in range(len(self))])

    def getNumberOfSlicesPerVolume(self):
        sliceLoc = self.sliceLocations
        return len(set(sliceLoc))

    def getNumberOfTimeSteps(self):
        sliceLoc = self.sliceLocations
        sliceLocS = set(sliceLoc)
        return sliceLoc.count(sliceLocS.pop())

    def getSliceNormalVector(self):
        self.sortBySlice_InstanceNumber()
        try:
            sliceVec = self.getImagePositionPatient_np(-1) - self.getImagePositionPatient_np(0)
        except:
            sliceVec = [0.0,0.0,0.0] # If no slice locations then assume no normal vector (one image)
        if np.linalg.norm(sliceVec) < 1e-9:
            iop = self.getTag('ImageOrientationPatient', ifNotFound=[1.0,0.0,0.0,0.0,1.0,0.0])
            sliceVec = np.cross(iop[:3], iop[3:6])
        else:
            sliceVec = sliceVec / np.linalg.norm(sliceVec)
        return sliceVec

    def doesSliceLocationNorm_Match_IOPNormal(self):
        iop = self.getTag('ImageOrientationPatient')
        iopN = np.cross(iop[:3], iop[3:6])
        sliceLocN = self.getSliceNormalVector()
        tf = [i*j>=0.0 for i,j in zip(iopN, sliceLocN)]
        return all(tf)

    def is3D(self):
        return self.getNumberOfSlicesPerVolume() > 1

    def is4D(self):
        return self.getNumberOfTimeSteps() > 1

    def getPixelDataAsNumpy(self):
        """Get pixel data as numpy array organised by slice and time(if present).
            Also return a PatientMatrix object containing meta
            NOTE: Some data is repeated to serve VTK and DICOM conventions

        Returns:
            tuple: 
                numpy array - shape [nRow, nCol, nSlice, nTime], 
                dictionary - keys: Spacing, Origin, ImageOrientationPatient, PatientPosition, Dimensions, Times
        """
        self._loadToMemory()
        I,J,K = int(self.getTag('Rows')), int(self.getTag('Columns')), int(self.getNumberOfSlicesPerVolume())
        self.sortBySlice_InstanceNumber() # This takes care of order - slices grouped, then time for each slice 
        N = self.getNumberOfTimeSteps()
        A = np.zeros((I, J, K, N))
        if (K*N) != len(self):
            print(f"DEBUG: Getting numpy array shape: [{I}, {J}, {K}, {N}] (K*N={K*N} == {len(self)})")
            raise FileNotFoundError(f"Missing some DICOM files.")
        c0 = 0
        for k1 in range(K):
            for k2 in range(N):
                iA = self[c0].pixel_array
                A[:, :, k1, k2] = iA
                c0 += 1
        patientMeta = dcmVTKTK.PatientMeta()
        patientMeta.initFromDicomSeries(self)
        return A, patientMeta

    def getScanDuration_secs(self):
        try:
            return self.getTag(0x0019105a, ifNotFound=0.0) / 1000000.0
        except AttributeError:
            return 0.0

    def _getPixelSpacing(self):
        return [float(i) for i in self.getTag('PixelSpacing', ifNotFound=[0.0,0.0])]

    def getDeltaRow(self):
        return self._getPixelSpacing()[0]

    def getDeltaCol(self):
        return self._getPixelSpacing()[1]

    def getDeltaSlice(self):
        self.sortByInstanceNumber()
        p0 = self.getImagePositionPatient_np(0)
        sliceLoc = [distBetweenTwoPts(p0, self.getImagePositionPatient_np(i)) for i in range(len(self))]
        sliceLocS = sorted(list(set(sliceLoc)))
        if len(sliceLocS) == 1: # May be CINE at same location
            dZ = self.getTag('SpacingBetweenSlices', ifNotFound=None)
            if dZ is None:
                dZ = self.getTag('SliceThickness')
            return float(dZ)
        return np.mean(np.diff(sliceLocS))

    def getTemporalResolution(self):
        try:
            return float(self.getTag('NominalInterval', ifNotFound=0.0)/self.getTag('CardiacNumberOfImages', ifNotFound=1))
        except ZeroDivisionError:
            return 0       

    def getTemporalResolution_TR_VPS(self):
        return float(self.getTag('RepetitionTime', ifNotFound=0.0)*self.getTag(0x00431007, ifNotFound=1.0))

    def getManufacturer(self):
        return self.getTag(0x00080070, ifNotFound='Unknown')

    def IS_GE(self):
        return self.getManufacturer().lower().startswith('ge')

    def IS_SIEMENS(self):
        return self.getManufacturer().lower().startswith('siemens')

    def IS_PHILIPS(self):
        return self.getManufacturer().lower().startswith('philips')

    def IS_BRUKER(self):
        return self.getManufacturer().lower().startswith('bruker')

    def getPulseSequenceName(self):
        if self.IS_GE():
            return self.getTag(0x0019109c)
        else:
            return self.getTag(0x00180024)

    def getVENC(self):
        """Get VENC value in mm/s
        """
        if self.IS_GE():
            return float(self.getTag(0x001910cc))
        elif self.IS_SIEMENS():
            vencStr = self.getTag(0x00511014)
            return float(vencStr.split("_")[0].replace("v", "")) * 10.0
        else: # TODO - add other vendors
            return None

    def getInternalPulseSequenceName(self):
        return self.getTag(0x0019109e)

    def getSeriesDescription(self):
        return self.getTag('SeriesDescription')

    def getSeriesInfoDict(self, EXTRA_TAGS=[]):
        # Default (standard tags):
        outDict = {'ScanDuration':self.getScanDuration_secs(),
            'nTime':self.getTag('CardiacNumberOfImages'),
            'nRow':self.getTag('Rows'),
            'nCols':self.getTag('Columns'),
            'dRow':self.getDeltaRow(),
            'dCol':self.getDeltaCol(),
            'dSlice':self.getTag('SliceThickness'),
            'dTime': self.getTemporalResolution(),
            'SpacingBetweenSlices':self.getTag('SpacingBetweenSlices'),
            'FlipAngle':self.getTag('FlipAngle'),
            'HeartRate':self.getTag('HeartRate'),
            'EchoTime':self.getTag('EchoTime'),
            'RepetitionTime':self.getTag('RepetitionTime'),
            'PulseSequenceName':self.getPulseSequenceName(),
            'MagneticFieldStrength': self.getTag('MagneticFieldStrength'),
            'InternalPulseSequenceName':self.getInternalPulseSequenceName(),
            'ReconstructionDiameter': self.getTag(0x00181100),
            'AcquisitionMatrix': str(self.getTag(0x00181310)),
            'Manufacturer': self.getTag("Manufacturer"),
            'ManufacturerModelName': self.getTag("ManufacturerModelName"),
            'SoftwareVersions': str(self.getTag(0x00181020)),}
        for extraTag in EXTRA_TAGS:
            outDict[extraTag] = self.getTag(extraTag)
        outDict['nSlice'] = len(self)
        try:
            outDict['AcquiredResolution'] = float(outDict['ReconstructionDiameter']) / float(max(self.getTag(0x00181310)))
        except ValueError:
            outDict['AcquiredResolution'] = f"{outDict['dRow']},{outDict['dCol']}"
        try:
            outDict['AcquiredTemporalResolution'] = self.getTemporalResolution_TR_VPS()
        except ValueError:
            outDict['AcquiredTemporalResolution'] = 0.0
        for i in outDict.keys():
            try:
                if ',' in outDict[i]:
                    outDict[i] = f'"{outDict[i]}"'
            except TypeError:
                pass
        return outDict

    def checkIfShouldUse_SAFE_NAMING(self, se_instance_set=None):
        """Checks is should name dicoms based on UIDs or if a SE-number, Instace-number naming convention is possible.
        Basically checks if there would be any overlaps. 

        Args:
            se_instance_set (set, optional): A set of instance file names (without extension). Defaults to None.
        """
        if self.SAFE_NAME_MODE:
            return
        if se_instance_set is None:
            se_instance_set = set()
        for k1 in range(len(self)):
            se = self.getTag('SeriesNumber', dsID=k1, ifNotFound='unknown')
            instance = self.getTag('InstanceNumber', dsID=k1, ifNotFound='unknown')
            se_instance_str = f"{se}_{instance}"
            if se_instance_str in se_instance_set:
                self.SAFE_NAME_MODE = True
                break
            se_instance_set.add(se_instance_str)

    def getStudyOutputDir(self, rootDir='', studyPrefix=''):
        """Generate a study output directory name based on config setting STUDY_NAMINF_TAG_LIST

        Args:
            rootDir (str, optional): The root directory to build study directory in. Defaults to ''.
            studyPrefix (str, optional): Optional prefix to add to naming. Defaults to ''.

        Returns:
            str: Path to directory at Study level where dicoms will be written
        """
        suffix = ''
        if self.SAFE_NAME_MODE:
            suffix += f"{self.getTag('StudyInstanceUID')}"
        for iTag in SpydcmTK_config.STUDY_NAMING_TAG_LIST:
            iVal = self.getTag(iTag, ifNotFound='', convertToType=str)
            if len(iVal) > 0:
                suffix += '_'+iVal
        return os.path.join(rootDir, dcmTools.cleanString(studyPrefix+suffix))


    def overviewImage(self, outputFileName=None, RETURN_FIG=False):
        """
        Build an overview image of the series.

        Args:
            outputFileName (str, optional): The name of the file to save the overview image to. Defaults to None.
            RETURN_FIG (bool, optional): Whether to return the figure. Defaults to False.

        Returns:
            if RETURN_FIG is True, returns the figure.
            if outputFileName is not None, saves the figure to the file and returns the file name.
            otherwise, displays the figure.
        """
        A = self.getPixelDataAsNumpy()[0]
        i,j,k,n = A.shape
        if k > 1:
            iSlice = k//2
            Image_list = [A[:,:,iSlice,0], np.rot90(A[i//2,:,:,0],3), np.rot90(A[:,j//2,:,0],3)]
        else:
            Image_list = [A[:,:,0,0]]
        fig, axs = plt.subplots(1, len(Image_list))
        if len(Image_list) == 1:
            axs = [axs]  # Convert single Axes to list for consistent indexing
        for i, iA in enumerate(Image_list):
            axs[i].imshow(iA, cmap='gray')
            axs[i].axis('off')
        fig.tight_layout()
        if RETURN_FIG:
            return fig
        if outputFileName is not None:
            plt.savefig(outputFileName)
            plt.close()
            return outputFileName
        else:
            plt.show()
        


# ================================================================================================================
# DicomStudy
# ================================================================================================================
class DicomStudy(list):
    """
    Extends list of DicomSeries objects (creating list of list of pydicom.dataset)
    """
    def __init__(self, dSeriesList, OVERVIEW=False, HIDE_PROGRESSBAR=False):
        """
        Set OVERVIEW = False to read pixel data as well (at a cost)
        """
        self.OVERVIEW = OVERVIEW
        self.HIDE_PROGRESSBAR = HIDE_PROGRESSBAR # FIXME - if set need ot pass down to series also - for the writing
        list.__init__(self, dSeriesList)

    @classmethod
    def setFromDictionary(cls, dicomDict, OVERVIEW=False, HIDE_PROGRESSBAR=False, FORCE_READ=False):
        dStudyList = ListOfDicomStudies.setFromDcmDict(dicomDict, OVERVIEW=OVERVIEW, HIDE_PROGRESSBAR=HIDE_PROGRESSBAR, FORCE_READ=FORCE_READ)
        if len(dStudyList) > 1:
            raise ValueError('More than one study found - use ListOfDicomStudies class')
        return cls(dStudyList[0], OVERVIEW=OVERVIEW, HIDE_PROGRESSBAR=HIDE_PROGRESSBAR)

    @classmethod
    def setFromDirectory(cls, dirName, OVERVIEW=False, HIDE_PROGRESSBAR=False, FORCE_READ=False, ONE_FILE_PER_DIR=False):
        dicomDict = dcmTools.organiseDicomHeirarchyByUIDs(dirName, HIDE_PROGRESSBAR=HIDE_PROGRESSBAR, FORCE_READ=FORCE_READ, ONE_FILE_PER_DIR=ONE_FILE_PER_DIR, OVERVIEW=OVERVIEW)
        obj = DicomStudy.setFromDictionary(dicomDict, OVERVIEW=OVERVIEW, HIDE_PROGRESSBAR=HIDE_PROGRESSBAR, FORCE_READ=FORCE_READ)
        for iSeries in obj:
            iSeries.NOT_FULLY_LOADED = ONE_FILE_PER_DIR
        return obj


    def __str__(self):
        return '%s: %s: %d series, %d dicoms'%(self.getTag('PatientName'),
                                                self.getTag('SeriesDescription'),
                                                len(self),
                                                self.getNumberOfDicoms())


    def setSafeNameMode(self):
        for iSeries in self:
            iSeries.SAFE_NAME_MODE = True

    def resetUIDs(self, studyUID=None):
        if studyUID is None:
            studyUID = str(generate_uid())
        for i in self:
            i.resetUIDs(studyUID)

    def filterByTag(self, tagName, tagValue):
        return DicomStudy([i for i in self if i.getTag(tagName) == tagValue], 
                           OVERVIEW=self.OVERVIEW, 
                           HIDE_PROGRESSBAR=self.HIDE_PROGRESSBAR)

    def anonymise(self, anonName, anonPatientID, anonBirthdate=True, removePrivateTags=True):
        for i in self:
            i.anonymise(anonName, anonPatientID, anonBirthdate, removePrivateTags)

    def getTag(self, tag, seriesID=0, instanceID=0, ifNotFound='Unknown', convertToType=None):
        return self[seriesID].getTag(tag, dsID=instanceID, ifNotFound=ifNotFound, convertToType=convertToType)

    def getTagListAndNames(self, tagList, seriesID=0, instanceID=0):
        return self[seriesID].getTagListAndNames(tagList, dsID=instanceID)

    def setTags_all(self, tag, value):
        for i in self:
            i.setTags_all(tag, value)

    def isCompressed(self):
        return self[0].isCompressed()

    def getNumberOfDicoms(self):
        return sum([len(i) for i in self])

    def getTopDir(self):
        return os.path.split(os.path.split(self[0][0].filename)[0])[0]

    def getSeriesMatchingDescription(self, matchingStrList, RETURN_SERIES_OBJ=False, REDUCE_MULTIPLE=False):
        matchStrList_lower = [i.lower() for i in matchingStrList]
        possibles = []
        for i in self:
            sDesc = i.getTag('SeriesDescription').lower()
            if any([j in sDesc  for j in matchStrList_lower ]):
                possibles.append(i)
        minID = None
        if len(possibles) == 0:
            return None
        elif (len(possibles) > 1) and REDUCE_MULTIPLE:
            seNums = [int(i.getTag('SeriesNumber')) for i in possibles]
            seNums = [i for i in seNums if i!=0]
            for k1 in range(1, len(seNums)):
                if seNums[k1] < seNums[minID]:
                    minID = k1
        if minID is not None:
            possibles = possibles[minID]
        if RETURN_SERIES_OBJ:
            return possibles
        return [f'SE{possibles[i].getTag("SeriesnNumber")}:{possibles[i].getTag("SeriesDescription")}' for i in range(len(possibles))]

    def getStudyOverview(self, tagList=SpydcmTK_config.STUDY_OVERVIEW_TAG_LIST):
        return self.getTagListAndNames(tagList)

    def getPatientOverview(self, tagList=SpydcmTK_config.SUBJECT_OVERVIEW_TAG_LIST):
        return self.getTagListAndNames(tagList)

    def getSeriesByID(self, ID):
        ID = int(ID)
        for iSeries in self:
            try:
                if int(iSeries.getTag('SeriesNumber')) == ID:
                    return iSeries
            except TypeError: # If SE# == None (report of Sec Capture etc)
                pass

    def getSeriesBySeriesNumber(self, seNum):
        return self.getSeriesByID(seNum)

    def getSeriesByUID(self, UID):
        for iSeries in self:
            if iSeries.getTag('SeriesInstanceUID') == UID:
                return iSeries


    def getSeriesByTag(self, tag, value, convertToType=None):
        for iSeries in self:
            if iSeries.getTag(tag, convertToType=convertToType) == value:
                return iSeries


    def mergeSeriesVolumesWithTime(self):
        nTimes = len(self)
        triggerTimes = sorted([self.getTag('TriggerTime', seriesID=i, instanceID=0, ifNotFound=0, convertToType=float) for i in range(nTimes)])
        nPerTime = len(self.getSeriesByTag('TriggerTime', triggerTimes[0], convertToType=float))
        sorted_ds_list = [None for _ in range(nTimes*nPerTime)]
        print(f"Merging series: n times: {nTimes}, N per time: {nPerTime}")
        for i in range(nTimes):
            iSeries = self.getSeriesByTag('TriggerTime', triggerTimes[i], convertToType=float)
            iSeries.sortByInstanceNumber()
            for k1, iDS in enumerate(iSeries):
                if nPerTime > 1:
                    # thisID = int(iDS['InstanceNumber'].value) + c1*(len(triggerTimes)-1) + c0 -1 # change to 0 index
                    thisID = i + (nTimes * k1)
                else:
                    thisID = i
                iDS.InstanceNumber = thisID + 1 # 1 index
                sorted_ds_list[thisID] = iDS
        if sorted_ds_list.count(None) != 0:
            raise ValueError('Missing some volumes')
        return DicomSeries(sorted_ds_list)


    def writeFDQ(self, seriesNumber_list, outputFileName, phase_factors=None, phase_offsets=None, working_dir=None, VERBOSE=True):
        """
        Writes a 4D-flow output as VTS format.

        Args:
            seriesNumber_list (list): A list of series numbers identifying mag and 3 phase series. Mag, X, Y, Z phases
            outputFileName (str): The name of the pvd file to write to.

        Returns:
            str: the pvd file written
        """
        # Check phase factors and offsets:
        if phase_factors is not None:
            if len(phase_factors) != 3:
                raise ValueError(f"phase_factors should be length 3")
        else:
            if self.getSeriesBySeriesNumber(seriesNumber_list[1]).IS_GE():
                phase_factors = [1.0, 1.0, -1.0]
                phase_offsets = [0.0, 0.0, 0.0]
            elif self.getSeriesBySeriesNumber(seriesNumber_list[1]).IS_SIEMENS():
                venc =  self.getSeriesBySeriesNumber(seriesNumber_list[1]).getVENC()
                c = -1.0 * venc
                m = (venc * 2.0) / 4096.0 # TODO: Check this
                phase_factors = [m, m, m]
                phase_offsets = [c, c, c]
            elif self.getSeriesBySeriesNumber(seriesNumber_list[1]).IS_PHILIPS(): # TODO: Check this
                phase_factors = [1.0, 1.0, 1.0]
                phase_offsets = [0.0, 0.0, 0.0]
            else:
                raise ValueError(f"Unknown scanner type")
        if phase_offsets is not None:
            if len(phase_offsets) != 3:
                raise ValueError(f"phase_offsets should be length 3")
        else:
            phase_offsets = [0.0, 0.0, 0.0]
        
        if working_dir is None:
            working_dir = os.path.dirname(outputFileName)
        intermediate_pvds = {}
        
        # Create process pool and process all directories in parallel
        origHIDE = self.HIDE_PROGRESSBAR
        self.HIDE_PROGRESSBAR = True
        with Pool() as pool:
            args = [(self.getSeriesBySeriesNumber(iSeriesNum), label, working_dir)
                    for iSeriesNum, label in zip(seriesNumber_list, ["MAG", "PX", "PY", "PZ"])]
            
            results = pool.map(_process_series_for_fdq, args)
        
        # Collect results into intermediate_pvds dictionary
        intermediate_pvds = dict(results)
        
        if VERBOSE: print(f"Combining to final 4D-flow PVD")
        fOut = dcmVTKTK.mergePhaseSeries4D(intermediate_pvds["MAG"], 
                                        [intermediate_pvds["PX"],
                                        intermediate_pvds["PY"],
                                        intermediate_pvds["PZ"]], 
                                        outputFileName, 
                                        phase_factors=phase_factors,
                                        phase_offsets=phase_offsets,
                                        scale_factor=0.001,
                                        DEL_ORIGINAL=True)
        self.HIDE_PROGRESSBAR = origHIDE
        if VERBOSE: print(f"Written {fOut}")
        return fOut


    def getStudySummaryDict(self, FORCE_STRING_KEYS=False):
        pt,pv = self.getPatientOverview()
        studySummaryDict = dict(zip(pt, pv))
        st,sv = self.getStudyOverview()
        stdyDict = dict(zip(st, sv))
        studySummaryDict.update(stdyDict)
        listSerDict = []
        for i in self:
            szt, szv = i.getSeriesOverview()
            listSerDict.append(dict(zip(szt, szv)))
        studySummaryDict['Series'] = listSerDict
        return studySummaryDict

    def getStudySummary(self, FULL=True):
        pt,pv = self.getPatientOverview()
        patStr = ','.join([' %s:%s'%(str(i), str(j)) for i,j in zip(pt,pv)])
        st,sv = self.getStudyOverview()
        studyStr = ','.join([' %s:%s'%(str(i), str(j)) for i,j in zip(st,sv)])
        if FULL:
            sst,_ = self[0].getSeriesOverview()
            seriesHeader = ','.join([str(i) for i in sst])
            seriesStr = [','.join([str(i2) for i2 in i.getSeriesOverview()[1]]) for i in self]
            try:
                seriesStr = sorted(seriesStr, key=lambda i:int(i.split(',')[0]))
            except ValueError:
                pass # Try to sort, but if fail (due to return not int somehow, then no sort)
            seriesStr_ = '\n\n'+seriesHeader+'\n    '+'\n    '.join(seriesStr)
            strOut = f"SUBJECT: {patStr}"
            strOut += f"\nSTUDY: {studyStr}"
            strOut += seriesStr_
            return strOut
        else:
            strOut = f"SUBJECT: {patStr}"
            strOut += f"\nSTUDY: {studyStr}"
            return strOut

    def getTagValuesList(self, tagList, RETURN_STR=False):
        output = []
        for i in self:
            output += i.getTagValuesList(tagList, False)
        if RETURN_STR:
            return dcmTools._tagValuesListToString(tagList, output)
        return output

    def writeToOrganisedFileStructure(self, patientOutputDir, studyPrefix=''):
        self.checkIfShouldUse_SAFE_NAMING()
        try:
            studyOutputDir = self[0].getStudyOutputDir(patientOutputDir, studyPrefix)
        except IndexError: # Study is empty
            return None
        studyOutputDirTemp = studyOutputDir+".WORKING"
        if os.path.isdir(studyOutputDir):
            os.rename(studyOutputDir, studyOutputDirTemp)
        for iSeries in self:
            iSeries.writeToOrganisedFileStructure(studyOutputDirTemp, 
                                                  SAFE_NAMING_CHECK=False)
        os.rename(studyOutputDirTemp, studyOutputDir)
        return studyOutputDir

    def writeToZipArchive(self, patientOutputDir, CLEAN_UP=True):
        studyOutputDir = self.writeToOrganisedFileStructure(patientOutputDir)
        if studyOutputDir is None:
            return
        r, f = os.path.split(studyOutputDir)
        shutil.make_archive(studyOutputDir, 'zip', os.path.join(r, f))
        fileOut = studyOutputDir+'.zip'
        if CLEAN_UP:
            shutil.rmtree(studyOutputDir)
        return fileOut

    def checkIfShouldUse_SAFE_NAMING(self):
        se_instance_set = set()
        for i in self:
            i.checkIfShouldUse_SAFE_NAMING(se_instance_set)


class ListOfDicomStudies(list):
    """
    Extends list of DicomStudies objects (creating list of list of list of pydicom.dataset)
    """
    def __init__(self, dStudiesList, OVERVIEW=False, HIDE_PROGRESSBAR=False):
        """
        Set OVERVIEW = False to read pixel data as well (at a cost)
        """
        self.OVERVIEW = OVERVIEW
        self.HIDE_PROGRESSBAR = HIDE_PROGRESSBAR
        list.__init__(self, dStudiesList)

    @classmethod
    def setFromInput(cls, input, OVERVIEW=False, HIDE_PROGRESSBAR=False, FORCE_READ=False, ONE_FILE_PER_DIR=False):
        if os.path.isdir(input):
            return ListOfDicomStudies.setFromDirectory(input, OVERVIEW=OVERVIEW, 
                                                                HIDE_PROGRESSBAR=HIDE_PROGRESSBAR, 
                                                                FORCE_READ=FORCE_READ, 
                                                                ONE_FILE_PER_DIR=ONE_FILE_PER_DIR)
        if os.path.isfile(input):
            if input.endswith('tar'):
                return ListOfDicomStudies.setFromTar(input, HIDE_PROGRESSBAR=HIDE_PROGRESSBAR, 
                                                            FORCE_READ=FORCE_READ)
            elif input.endswith('tar.gz'):
                return ListOfDicomStudies.setFromTar(input, HIDE_PROGRESSBAR=HIDE_PROGRESSBAR, 
                                                            FORCE_READ=FORCE_READ)
            elif input.endswith('zip'):
                return ListOfDicomStudies.setFromZip(input, HIDE_PROGRESSBAR=HIDE_PROGRESSBAR, 
                                                            FORCE_READ=FORCE_READ)
            else:
                dcmDict = {}
                try:
                    dcmTools.readDicomFile_intoDict(input, dcmDict, FORCE_READ=FORCE_READ, OVERVIEW=OVERVIEW)
                    return ListOfDicomStudies.setFromDcmDict(dcmDict, OVERVIEW, HIDE_PROGRESSBAR, FORCE_READ=FORCE_READ)
                except dicom.filereader.InvalidDicomError:
                    raise IOError("ERROR READING DICOMS: SPDCMTK capable to read dicom files from directory, zip, tar or tar.gz\n")

    @classmethod
    def setFromDirectory(cls, dirName, OVERVIEW=False, HIDE_PROGRESSBAR=False, FORCE_READ=False, ONE_FILE_PER_DIR=False, extn_filter=None):
        dicomDict = dcmTools.organiseDicomHeirarchyByUIDs(dirName, 
                                                         HIDE_PROGRESSBAR=HIDE_PROGRESSBAR, 
                                                         FORCE_READ=FORCE_READ, 
                                                         ONE_FILE_PER_DIR=ONE_FILE_PER_DIR, 
                                                         OVERVIEW=OVERVIEW,
                                                         extn_filter=extn_filter)
        obj = ListOfDicomStudies.setFromDcmDict(dicomDict, OVERVIEW, HIDE_PROGRESSBAR, FORCE_READ=FORCE_READ)
        for iStudy in obj:
            for iSeries in iStudy:
                iSeries.NOT_FULLY_LOADED = ONE_FILE_PER_DIR
        return obj

    @classmethod
    def setFromDcmDict(cls, dicomDict, OVERVIEW=False, HIDE_PROGRESSBAR=False, FORCE_READ=False):
        dStudiesList = []
        for studyUID in dicomDict.keys():
            dSeriesList = []
            for seriesUID in dicomDict[studyUID].keys():
                dSeriesList.append(DicomSeries(dicomDict[studyUID][seriesUID], OVERVIEW=OVERVIEW, HIDE_PROGRESSBAR=HIDE_PROGRESSBAR, FORCE_READ=FORCE_READ))
            dStudiesList.append(DicomStudy(dSeriesList, OVERVIEW, HIDE_PROGRESSBAR))
        return cls(dStudiesList, OVERVIEW, HIDE_PROGRESSBAR)

    @classmethod
    def setFromTar(cls, tarFileName, HIDE_PROGRESSBAR=False, FORCE_READ=False, FIRST_ONLY=False, matchTagPair=None):
        # Note need OVERVIEW = False as only get access to file (and pixels) on untaring (maybe only if tar.gz) 
        dicomDict = dcmTools.getDicomDictFromTar(tarFileName, FORCE_READ=FORCE_READ, FIRST_ONLY=FIRST_ONLY, OVERVIEW_ONLY=False,
                                    matchingTagValuePair=matchTagPair, QUIET=True)
        return ListOfDicomStudies.setFromDcmDict(dicomDict, HIDE_PROGRESSBAR=HIDE_PROGRESSBAR, FORCE_READ=FORCE_READ)

    @classmethod
    def setFromZip(cls, zipFileName, HIDE_PROGRESSBAR=False, FORCE_READ=False, FIRST_ONLY=False, matchTagPair=None):
        dicomDict = dcmTools.getDicomDictFromZip(zipFileName, FORCE_READ=FORCE_READ, FIRST_ONLY=FIRST_ONLY, OVERVIEW_ONLY=False,
                                    matchingTagValuePair=matchTagPair, QUIET=True)
        return ListOfDicomStudies.setFromDcmDict(dicomDict, HIDE_PROGRESSBAR=HIDE_PROGRESSBAR, FORCE_READ=FORCE_READ)


    def __str__(self):
        return '%d studies, %d dicoms'%(len(self), self.getNumberOfDicoms())

    def setSafeNameMode(self):
        for iStudy in self:
            for iSeries in iStudy:
                iSeries.SAFE_NAME_MODE = True

    def resetUIDs(self):
        for i in self:
            i.resetUIDs()

    def anonymise(self, anonName, anonPatientID, anonBirthdate=True, removePrivateTags=True):
        for i in self:
            i.anonymise(anonName, anonPatientID, anonBirthdate, removePrivateTags)

    def isCompressed(self):
        return self[0].isCompressed()

    def getNumberOfDicoms(self):
        return sum([i.getNumberOfDicoms() for i in self])

    def filterByTag(self, tagName, tagValue):
        return ListOfDicomStudies([i for i in self if i.getTag(tagName) == tagValue], 
                                  OVERVIEW=self.OVERVIEW, 
                                  HIDE_PROGRESSBAR=self.HIDE_PROGRESSBAR)

    def writeToOrganisedFileStructure(self, outputRootDir):
        outDirs = []
        for iStudy in self:
            suffix = ''
            for iTag in SpydcmTK_config.SUBJECT_NAMING_TAG_LIST:
                iVal = iStudy.getTag(iTag, ifNotFound='', convertToType=str)
                if len(iVal) > 0:
                    suffix += '_'+iVal
            if len(suffix) > 0:
                ioutputRootDir = os.path.join(outputRootDir, dcmTools.cleanString(suffix))
            else:
                ioutputRootDir = outputRootDir

            ooD = iStudy.writeToOrganisedFileStructure(ioutputRootDir)
            if ooD is not None:
                outDirs.append(ooD)
        return outDirs

    def writeToZipArchive(self, outputRootDir, CLEAN_UP=True):
        outDirs = []
        for iStudy in self:
            ooD = iStudy.writeToZipArchive(outputRootDir, CLEAN_UP=CLEAN_UP)
            outDirs.append(ooD)
        return outDirs

    def getSummaryString(self, FULL=True):
        output = [i.getStudySummary(FULL) for i in self]
        return '\n\n'.join(output)

    def getTableOfTagValues(self, DICOM_TAGS, RETURN_STR):
        output = []
        for i in self:
            output += i.getTagValuesList(DICOM_TAGS, False)
        if RETURN_STR:
            return dcmTools._tagValuesListToString(DICOM_TAGS, output)
        return output

    def printSummaryTable(self):
        overviewList = self.getTableOfTagValues(SpydcmTK_config.SUBJECT_OVERVIEW_TAG_LIST+SpydcmTK_config.STUDY_OVERVIEW_TAG_LIST, False)
        print(','.join(SpydcmTK_config.SUBJECT_OVERVIEW_TAG_LIST+SpydcmTK_config.STUDY_OVERVIEW_TAG_LIST))
        for i in overviewList:
            print(','.join(i))

    def getStudyByDate(self, date_str):
        return self.getStudyByTag('StudyDate', date_str)

    def getStudyByPID(self, PID):
        return self.getStudyByTag('PatientID', PID)

    def getStudyByTag(self, tagName, tagValue):
        for i in self:
            if i.getTag(tagName) == tagValue:
                return i

    def buildMSTable(self, DICOM_TAGS=SpydcmTK_config.SERIES_OVERVIEW_TAG_LIST):
        pass
        #TODO - need to pass series name or something to query and then calc mean / stdev etc

## =====================================================================================================================
##   BIDS
## =====================================================================================================================
class BIDSConst(object):
    anat = 'anat'
    func = 'func'
    dwi = 'dwi'
    fmap = 'fmap'
    code = 'code'
    configf = 'dcm2bids_config.json'


def getJsonDict(niiORjsonFile):
    if niiORjsonFile.endswith('.nii.gz'):
        niiORjsonFile = niiORjsonFile.replace('.nii.gz','.json')
    with open(niiORjsonFile, 'r') as f:
        jsonDict = json.load(f)
    return jsonDict

def getTagFromJson(niiORjsonFile, tag):
    jsonDict = getJsonDict(niiORjsonFile)
    return jsonDict[tag]


class BIDS_Subject(object):
    """
    Class to work with a single BIDS subject
    """
    def __init__(self, subjID, BIDSDirObj):
        if subjID.startswith('sub-'):
            self.subjID = subjID
        else:
            self.subjID = 'sub-%s'%(subjID)
        self.parent = BIDSDirObj

    def __gt__(self, other):
        return self.subjID > other.subjID

    def __eq__(self, other):
        return self.subjID == other.subjID

    def getSubjID(self):
        return self.subjID.replace('sub-', '')

    def getSessionDir(self, session):
        if not session.startswith('ses-'):
            session = 'ses-%s'%(session)
        return os.path.join(self.parent.rootDir, self.subjID, session)

    def getTypeDir(self, TYPE, session):
        return os.path.join(self.getSessionDir(session), TYPE)

    def _getAllTypeNII(self, TYPE, session):
        dd = self.getTypeDir(TYPE, session)
        return [os.path.join(dd, i) for i in os.listdir(dd) if i.endswith('nii.gz')]

    def getNIIMatchingModality(self, modality, TYPE, session):
        matchingF = [i for i in self._getAllTypeNII(TYPE, session) if i.endswith('_%s.nii.gz'%(modality))]
        return matchingF

    def getNIIMatchingModality_condition(self, modality, TYPE, session, tag, func):
        matchingNII = self.getNIIMatchingModality(modality, TYPE, session)
        if len(matchingNII) == 0:
            return None
        elif len(matchingNII) == 1:
            return matchingNII[0]
        else: # more than 1 run so need to make decision
            tagList = [getTagFromJson(i, tag) for i in matchingNII]
            ID = func(tagList)
            return matchingNII[ID]


class BIDS_Directory(object):
    """
    Class to work with BIDS directory
    """
    def __init__(self, rootDirectroy):
        self.rootDir = rootDirectroy

    def getConfigFile(self, configFName=BIDSConst.configf):
        return os.path.join(self.rootDir, BIDSConst.code, configFName)

    def getSubjList(self):
        subjIDs = [i for i in os.listdir(self.rootDir) if i.startswith('sub-')]
        return sorted([BIDS_Subject(i, self) for i in subjIDs])


def organiseDicoms(dcmDirectory, outputDirectory, anonName=None, anonID='', FORCE_READ=False, 
                   HIDE_PROGRESSBAR=False, REMOVE_PRIVATE_TAGS=False):
    """ Recurse down directory tree - grab dicoms and move to new
        hierarchical folder structure
    """
    if not HIDE_PROGRESSBAR:
        print('READING...')
    studies = ListOfDicomStudies.setFromInput(dcmDirectory, FORCE_READ=FORCE_READ, OVERVIEW=False, HIDE_PROGRESSBAR=HIDE_PROGRESSBAR)
    if not HIDE_PROGRESSBAR:
        print('WRITTING...')
    res = studies.writeToOrganisedFileStructure(outputDirectory, anonName=anonName, anonID=anonID, REMOVE_PRIVATE_TAGS=REMOVE_PRIVATE_TAGS)
    if res is None:
        print('    ## WARNING - No valid dicoms found')


def studySummary(pathToDicoms):
    """
    :param pathToDicoms: path to dicoms - anywhere in the tree above - will search down tree till find a dicom
    :return: a formatted string
    """
    dStudies = ListOfDicomStudies.setFromDirectory(pathToDicoms)
    return dStudies.getSummaryString()



## =====================================================================================================================
##   WRITE DICOMS
## =====================================================================================================================
def writeVTIToDicoms(vtiFile, dcmTemplateFile_or_ds, outputDir, arrayName=None, tagUpdateDict=None, patientMeta=None):
    if type(vtiFile) == str:
        vti = dcmVTKTK.fIO.readVTKFile(vtiFile)
        if vtiFile.endswith(".nii") or vtiFile.endswith(".nii.gz"):
            # TODO - need to also adjust patientMeta if this is nifti 
            vti.SetOrigin([i*0.001 for i in vti.GetOrigin()])
            vti.SetSpacing([i*0.001 for i in vti.GetSpacing()])
    else:
        vti = vtiFile
    if arrayName is None:
        A = dcmVTKTK.vtkfilters.getScalarsAsNumpy(vti)
    else:
        A = dcmVTKTK.vtkfilters.getArrayAsNumpy(vti, arrayName)
    A = np.reshape(A, vti.GetDimensions(), 'F')
    if patientMeta is None:
        patientMeta = dcmVTKTK.PatientMeta()
    patientMeta.initFromVTI(vti)
    return writeNumpyArrayToDicom(A, dcmTemplateFile_or_ds, patientMeta, outputDir, tagUpdateDict=tagUpdateDict)


def writeNumpyArrayToDicom(pixelArray, dcmTemplate_or_ds, patientMeta, outputDir, tagUpdateDict=None):
    if tagUpdateDict is None:
        tagUpdateDict = {}
    if dcmTemplate_or_ds is None:
        dsRAW = dcmTools.buildFakeDS()
    elif type(dcmTemplate_or_ds) == str:
        dsRAW = dicom.dcmread(dcmTemplate_or_ds)
    else:
        dsRAW = dcmTemplate_or_ds
    nRow, nCol, nSlice = pixelArray.shape
    if pixelArray.dtype != np.int16:
        pixelArray = pixelArray * (2**13 / np.max(pixelArray) )
        pixelArray = pixelArray.astype(np.int16)
    mx, mn = np.max(pixelArray), 0
    try:
        slice0 = tagUpdateDict.pop('SliceLocation0')
    except KeyError:
        slice0 = patientMeta.SliceLocation0*dcmVTKTK.m_to_mm
    
    sliceThick = patientMeta.SliceThickness*dcmVTKTK.m_to_mm
    ipp = [i*dcmVTKTK.m_to_mm for i in patientMeta.ImagePositionPatient]

    SeriesUID = dicom.uid.generate_uid()
    try:
        SeriesNumber = tagUpdateDict.pop('SeriesNumber')
    except KeyError:
        try:
            SeriesNumber = dsRAW.SeriesNumber * 100
        except AttributeError:
            SeriesNumber = 99
    dsList = []
    for k in range(nSlice):
        ds = copy.deepcopy(dsRAW)
        ds.SeriesInstanceUID = SeriesUID
        ds.SOPInstanceUID = dicom.uid.generate_uid()
        # ds.SpecificCharacterSet = 'ISO_IR 100'
        # ds.SOPClassUID = 'SecondaryCaptureImageStorage'
        ds.Rows = nRow
        ds.Columns = nCol
        ds.ImagesInAcquisition = nSlice
        ds.InStackPositionNumber = k+1
        # ds.RawDataRunNumber = k+1
        ds.SeriesNumber = SeriesNumber
        ds.InstanceNumber = k+1
        ds.SamplesPerPixel = 1
        ds.BitsAllocated = 16
        ds.BitsStored = 16
        ds.HighBit = 15
        ds.SliceThickness = sliceThick # Can no longer claim overlapping slices if have modified
        ds.SpacingBetweenSlices = sliceThick
        ds.SmallestImagePixelValue = int(mn)
        ds.LargestImagePixelValue = int(mx)
        ds.WindowCenter = int(mx / 2)
        ds.WindowWidth = int(mx / 2)
        ds.PixelSpacing = [i*dcmVTKTK.m_to_mm for i in list(patientMeta.PixelSpacing)]

        sliceVec = np.array(patientMeta.SliceVector)
        ImagePositionPatient = ipp + k*sliceVec*sliceThick
        ds.ImagePositionPatient = list(ImagePositionPatient)
        sliceLoc = slice0 + k*sliceThick
        ds.SliceLocation = sliceLoc
        ds.ImageOrientationPatient = list(patientMeta.ImageOrientationPatient)
        for iKey in tagUpdateDict.keys():
            if len(tagUpdateDict[iKey]) == 3:
                ds.add_new(tagUpdateDict[iKey][0], tagUpdateDict[iKey][1], tagUpdateDict[iKey][2])
            else:
                ds[iKey] = tagUpdateDict[iKey]
        ds.PixelData = pixelArray[:,:,k].tobytes()
        ds['PixelData'].VR = 'OW'
        dsList.append(ds)
    dcmSeries = DicomSeries(dsList, HIDE_PROGRESSBAR=True)
    return dcmSeries.writeToOrganisedFileStructure(outputDir)

def writeImageStackToDicom(images_sortedList, patientMeta, dcmTemplateFile_or_ds, 
                            outputDir, tagUpdateDict=None):

    combinedImage = dcmVTKTK.readImageStackToVTI(images_sortedList, patientMeta, CONVERT_TO_GREYSCALE=True)
    writeVTIToDicoms(combinedImage, 
                        dcmTemplateFile_or_ds=dcmTemplateFile_or_ds, 
                        outputDir=outputDir,
                        arrayName='PixelData',
                        tagUpdateDict=tagUpdateDict,
                        patientMeta=patientMeta)


def getResolution(dataVts):
    o,p1,p2,p3 = [0.0,0.0,0.0],[0.0,0.0,0.0],[0.0,0.0,0.0],[0.0,0.0,0.0]
    i0,i1,j0,j1,k0,k1 = dataVts.GetExtent()
    dataVts.GetPoint(i0,j0,k0, o)
    dataVts.GetPoint(i0+1,j0,k0, p1)
    dataVts.GetPoint(i0,j0+1,k0, p2)
    dataVts.GetPoint(i0,j0,k0+1, p3)
    di = abs(distBetweenTwoPts(p1, o))
    dj = abs(distBetweenTwoPts(p2, o))
    dk = abs(distBetweenTwoPts(p3, o))
    return [di, dj, dk]

def distBetweenTwoPts(a, b):
    return np.sqrt(np.sum((np.asarray(a) - np.asarray(b)) ** 2))

def _process_series_for_fdq(args):
    series, label, outDir = args
    series._loadToMemory()
    iOut = series.writeToVTS(os.path.join(outDir, f"{label}.pvd"))
    return (label, iOut)

