# Manuel A. Morales (moralesq@mit.edu)

# Harvard-MIT Department of Health Sciences & Technology  

# Athinoula A. Martinos Center for Biomedical Imaging

import os

import h5py

import glob

import pydicom

import warnings

import numpy as np

import pandas as pd

import nibabel as nib

from data.base_dataset import BaseDataset, Transforms

from data.image_folder import make_dataset

class DICOMDataset(BaseDataset):

    def __init__(self, opt):

        BaseDataset.__init__(self, opt)

        self.filenames = sorted(make_dataset(opt.dataroot, opt.max_dataset_size, 'DICOM'))

        self.metadata, self.acquisitions  = self.read_metadata()

    def __len__(self):

        return len(self.acquisitions)

    def __getitem__(self, idx): 

        pid = '_'.join(self.acquisitions[idx].split('_')[:-1]) #PatientName

        uid = self.acquisitions[idx].split('_')[-1] #AcquisitionInstanceUID

        df  = self.metadata[(self.metadata.PatientName==pid)&(self.metadata.AcquisitionInstanceUID==uid)]

        return self.load_acquisition(df)

    def read_metadata(self):       

        metadata = {'FileName':[],

                    'PatientName':[], 

                    'SeriesInstanceUID':[], 

                    'StudyInstanceUID':[], 

                    'ProtocolName':[], 

                    'SeriesTime':[], 

                    'TriggerTime':[], 

                    'InstanceNumber':[], 

                    'ImageOrientationPatient':[],  

                    'ImagePositionPatient':[],

                    'SliceLocation':[], 

                    'PixelSpacing':[], 

                    'SliceThickness':[], 

                    'AcquisitionInstanceUID':[], 

                    'SliceInstanceUID':[]}

        for filename in self.filenames:

            dicom = pydicom.read_file(filename)

            metadata['FileName']                += [filename]

            metadata['PatientName']             += [str(dicom[0x0010, 0x0010].value)]

            metadata['SeriesInstanceUID']       += [dicom.SeriesInstanceUID]

            metadata['StudyInstanceUID']        += [dicom.StudyInstanceUID]

            metadata['ProtocolName']            += [dicom.ProtocolName]

            metadata['SeriesTime']              += [dicom.SeriesTime]

            metadata['TriggerTime']             += [dicom.TriggerTime]

            metadata['InstanceNumber']          += [dicom.InstanceNumber]

            metadata['ImageOrientationPatient'] += [dicom.ImageOrientationPatient]

            metadata['ImagePositionPatient']    += [dicom.ImagePositionPatient]

            metadata['SliceLocation']           += [dicom.SliceLocation]

            metadata['PixelSpacing']            += [dicom.PixelSpacing]

            metadata['SliceThickness']          += [dicom.SliceThickness]

            metadata['AcquisitionInstanceUID']  += [dicom.SeriesInstanceUID.split('.')[9]]

            metadata['SliceInstanceUID']        += [dicom.SeriesInstanceUID.split('.')[10]]

        metadata = pd.DataFrame(metadata)   

        acquisitions = []

        print('Found %d patient(s):'%(len(metadata.PatientName.unique())))

        for patient in metadata.PatientName.unique():

            acqs = metadata[metadata.PatientName==patient].AcquisitionInstanceUID.unique().tolist()

            print(patient, ': with', len(acqs), 'acquisitions:')

            for acq in sorted(acqs):

                acquisitions += [patient+'_'+acq]

                print('  ', acquisitions[-1])

        return metadata, acquisitions

    def load_acquisition(self, df):

        slices = df.SeriesInstanceUID.unique().tolist()

        phases = df.SeriesInstanceUID.value_counts().unique()

        assert len(phases)==1, 'Number of phases does not match!'

        number_of_slices = len(slices) 

        number_of_phases = int(phases)

        pixel_array = pydicom.read_file(df.iloc[0].FileName).pixel_array

        sax_4D = np.zeros((pixel_array.shape +(number_of_slices, number_of_phases)), dtype=pixel_array.dtype)

        for z_slice, series in enumerate(slices):

            for phase in range(number_of_phases):

                filename = df[df.SeriesInstanceUID==series].sort_values('InstanceNumber').iloc[phase].FileName

                dicom = pydicom.read_file(filename)

                sax_4D[:,:,z_slice,phase] += dicom.pixel_array

        affine = read_affine(df.iloc[df.SliceLocation.argmin()])

        return nib.Nifti1Image(sax_4D, affine)

def extract_cosines(ImageOrientationPatient):

    row_cosine    = np.array(ImageOrientationPatient[:3])

    column_cosine = np.array(ImageOrientationPatient[3:])

    slice_cosine  = np.cross(row_cosine, column_cosine)

    return np.stack((row_cosine, column_cosine, slice_cosine))

def read_affine(df, viewer='slicer'):

    Zooms = np.array(list(df.PixelSpacing)+[df.SliceThickness], dtype=float)

    ImageOrientationPatient = np.array(df.ImageOrientationPatient, dtype=float)

    ImagePositionPatient    = np.array(df.ImagePositionPatient, dtype=float)

    ijk2ras = extract_cosines(ImageOrientationPatient)

    if viewer == "slicer":

        ijk2ras = (ijk2ras*np.array([-1,-1,1])).T

        ImagePositionPatient = ImagePositionPatient*np.array([-1,-1,1])

    affine  = np.stack((ijk2ras[:,0]*Zooms[0],

                        ijk2ras[:,1]*Zooms[1],

                        ijk2ras[:,2]*Zooms[2],

                        ImagePositionPatient), axis=1)

    return np.vstack((affine,[[0,0,0,1]]))