import pydicom
import datetime
import numpy as np
import scipy
import nibabel as nib

class PatientInfo:

  def __init__(self):
    self.PatientID = ''
    self.PatientName = ''
    self.PatientBirthDate = ''
    self.PatientSex = ''

class StudyInfo:

  def __init__(self):
    self.StudyInstanceUID = ''
    self.StudyID = ''
    self.StudyDate = ''
    self.StudyTime = ''

class RTdose:

  def __init__(self):
    self.SeriesInstanceUID = ""
    self.SOPInstanceUID = ""
    self.PatientInfo = PatientInfo()
    self.StudyInfo = StudyInfo()
    self.CT_SeriesInstanceUID = ""
    self.Plan_SOPInstanceUID = ""
    self.FrameOfReferenceUID = ""
    self.ImgName = ""
    self.beam_number = "" # Beam number (str) or PLAN if sum of all
    self.DcmFile = ""
    self.isLoaded = 0
    
  def print_dose_info(self, prefix=""):
    print(prefix + "Dose: " + self.SOPInstanceUID)
    print(prefix + "   " + self.DcmFile)
    
    
    
  def import_Dicom_dose(self, CT):
    if(self.isLoaded == 1):
      print("Warning: Dose image " + self.SOPInstanceUID + " is already loaded")
      return
      
    dcm = pydicom.dcmread(self.DcmFile)
    
    self.CT_SeriesInstanceUID = CT.SeriesInstanceUID
    # self.Plan_SOPInstanceUID = dcm.ReferencedRTPlanSequence[0].ReferencedSOPInstanceUID
    
    
    if(dcm.BitsStored == 16 and dcm.PixelRepresentation == 0):
      dt = np.dtype('uint16')
    elif(dcm.BitsStored == 16 and dcm.PixelRepresentation == 1):
      dt = np.dtype('int16')
    elif(dcm.BitsStored == 32 and dcm.PixelRepresentation == 0):
      dt = np.dtype('uint32')
    elif(dcm.BitsStored == 32 and dcm.PixelRepresentation == 1):
      dt = np.dtype('int32')
    else:
      print("Error: Unknown data type for " + self.DcmFile)
      return
      
    if(dcm.HighBit == dcm.BitsStored-1):
      dt = dt.newbyteorder('L')
    else:
      dt = dt.newbyteorder('B')
      
    dose_image = np.frombuffer(dcm.PixelData, dtype=dt) 
    dose_image = dose_image.reshape((dcm.Columns, dcm.Rows, dcm.NumberOfFrames), order='F').transpose(1,0,2)
    dose_image = dose_image * dcm.DoseGridScaling
    
    self.Image = dose_image
    self.FrameOfReferenceUID = dcm.FrameOfReferenceUID
    self.ImagePositionPatient = dcm.ImagePositionPatient
    if dcm.SliceThickness is not None:
        self.PixelSpacing = [dcm.PixelSpacing[0], dcm.PixelSpacing[1], dcm.SliceThickness]
    else:
        self.PixelSpacing = [dcm.PixelSpacing[0], dcm.PixelSpacing[1], dcm.GridFrameOffsetVector[1]-dcm.GridFrameOffsetVector[0]]
    self.GridSize = [dcm.Columns, dcm.Rows, dcm.NumberOfFrames]
    self.NumVoxels = self.GridSize[0] * self.GridSize[1] * self.GridSize[2]
    
    if hasattr(dcm, 'GridFrameOffsetVector'):
      if(dcm.GridFrameOffsetVector[1] - dcm.GridFrameOffsetVector[0] < 0):
        self.Image = np.flip(self.Image, 2)
        self.ImagePositionPatient[2] = self.ImagePositionPatient[2] - self.GridSize[2]*self.PixelSpacing[2]
    
    self.resample_to_CT_grid(CT)
    self.isLoaded = 1
        
    
    
  def euclidean_dist(self, v1, v2):
    return sum((p-q)**2 for p, q in zip(v1, v2)) ** .5
    
    
      
  def resample_to_CT_grid(self, CT):
    if(self.GridSize == CT.GridSize and self.euclidean_dist(self.ImagePositionPatient, CT.ImagePositionPatient) < 0.001 and self.euclidean_dist(self.PixelSpacing, CT.PixelSpacing) < 0.001):
      return
    else:
      # anti-aliasing filter
      sigma = [0, 0, 0]
      if(CT.PixelSpacing[0] > self.PixelSpacing[0]): sigma[0] = 0.4 * (CT.PixelSpacing[0]/self.PixelSpacing[0])
      if(CT.PixelSpacing[1] > self.PixelSpacing[1]): sigma[1] = 0.4 * (CT.PixelSpacing[1]/self.PixelSpacing[1])
      if(CT.PixelSpacing[2] > self.PixelSpacing[2]): sigma[2] = 0.4 * (CT.PixelSpacing[2]/self.PixelSpacing[2])
      if(sigma != [0, 0, 0]):
        print("Image is filtered before downsampling")
        self.Image = scipy.ndimage.gaussian_filter(self.Image, sigma)

      
      print('Resample dose image to CT grid.')
    
      x = self.ImagePositionPatient[1] + np.arange(self.GridSize[1]) * self.PixelSpacing[1]
      y = self.ImagePositionPatient[0] + np.arange(self.GridSize[0]) * self.PixelSpacing[0]
      z = self.ImagePositionPatient[2] + np.arange(self.GridSize[2]) * self.PixelSpacing[2]
  
      xi = np.array(np.meshgrid(CT.VoxelY, CT.VoxelX, CT.VoxelZ))
      xi = np.rollaxis(xi, 0, 4)
      xi = xi.reshape((xi.size // 3, 3))
  
      self.Image = scipy.interpolate.interpn((x,y,z), self.Image, xi, method='linear', fill_value=0, bounds_error=False)
      self.Image = self.Image.reshape((CT.GridSize[0], CT.GridSize[1], CT.GridSize[2])).transpose(1,0,2)
  
      self.ImagePositionPatient = CT.ImagePositionPatient
      self.PixelSpacing = CT.PixelSpacing
      self.GridSize = CT.GridSize
      self.NumVoxels = CT.NumVoxels

 
  def load_from_nii(self, dose_nii):   
      
    # load the nii image
    img = nib.load(dose_nii)
    
    self.Image = img.get_fdata() ### SHOULD I TRANSPOSE?
    self.GridSize = self.Image.shape
    self.PixelSpacing = [img.header['pixdim'][1], img.header['pixdim'][2], img.header['pixdim'][3]]
    self.ImagePositionPatient = [ img.affine[0][3], img.affine[1][3], img.affine[2][3]]
    
  def export_Dicom(self, refCT, OutputFile):

    # meta data
    SOPInstanceUID = pydicom.uid.generate_uid() 
    meta = pydicom.dataset.FileMetaDataset()
    meta.MediaStorageSOPClassUID = '1.2.840.10008.5.1.4.1.1.481.2' # UID class for RTDOSE
    meta.MediaStorageSOPInstanceUID = SOPInstanceUID
    #meta.ImplementationClassUID = '1.2.826.0.1.3680043.1.2.100.5.7.0.47' # from RayStation
    #meta.ImplementationClassUID = '1.2.826.0.1.3680043.5.5.100.5.7.0.03' # modified OpenTPS
    meta.ImplementationClassUID = '1.2.826.0.1.3680043.1.2.100.6.40.0.76'# from Halcyon? st. luc breast patients
    #meta.TransferSyntaxUID = '1.2.840.10008.1.2'

    meta.FileMetaInformationGroupLength = 200
    #meta.FileMetaInformationVersion = 
    #meta.ImplementationVersionName = 'DicomObjects.NET'
    # dcm_file.ImplementationVersionName =
    # dcm_file.SoftwareVersion = 

    # dicom dataset
    dcm_file = pydicom.dataset.FileDataset(OutputFile, {}, file_meta=meta, preamble=b"\0" * 128) # CONFIRM WHAT IS OUTPUTFILE AND WHAT THIS LINE IS DOING

    # transfer syntax  
    dcm_file.file_meta.TransferSyntaxUID = pydicom.uid.ExplicitVRLittleEndian
    print(dcm_file.file_meta.TransferSyntaxUID)
    dcm_file.is_little_endian = True
    dcm_file.is_implicit_VR = False

    # Patient
    dcm_file.PatientName = refCT.PatientInfo.PatientName #self.PatientInfo.PatientName
    dcm_file.PatientID = refCT.PatientInfo.PatientID #self.PatientInfo.PatientID
    dcm_file.PatientBirthDate = refCT.PatientInfo.PatientBirthDate #self.PatientInfo.PatientBirthDate
    dcm_file.PatientSex = refCT.PatientInfo.PatientSex #self.PatientInfo.PatientSex
    
    # General Study
    dt = datetime.datetime.now()
    dcm_file.StudyDate = dt.strftime('%Y%m%d')
    dcm_file.StudyTime = dt.strftime('%H%M%S.%f')
    dcm_file.AccessionNumber = '1' # A RIS/PACS (Radiology Information System/picture archiving and communication system) generated number that identifies the order for the Study.
    dcm_file.ReferringPhysicianName = 'NA'
    dcm_file.StudyInstanceUID = refCT.StudyInfo.StudyInstanceUID # get from reference CT to indicate that they belong to the same study#self.StudyInfo.StudyInstanceUID
    dcm_file.StudyID = refCT.StudyInfo.StudyID # get from reference CT to indicate that they belong to the same study

    # RT Series
    dcm_file.Modality = 'RTDOSE'
    dcm_file.SeriesDescription = 'AI-predicted' + dt.strftime('%Y%m%d') + dt.strftime('%H%M%S.%f')#self.ImgName
    dcm_file.OperatorsName = 'MIRO'
    dcm_file.SeriesInstanceUID = pydicom.uid.generate_uid() # if we have a uid_base --> pydicom.uid.generate_uid(uid_base)
    dcm_file.SeriesNumber = 1

    # Frame of Reference
    dcm_file.FrameOfReferenceUID = refCT.FrameOfReferenceUID # pydicom.uid.generate_uid()
    dcm_file.PositionReferenceIndicator = '' #empty if unknown https://dicom.innolitics.com/ciods/rt-dose/frame-of-reference/00201040

    # General Equipment
    dcm_file.Manufacturer = 'echarp'
    #dcm_file.ManufacturerModelName = 'echarp'
    #dcm_file.PixelPaddingValue = # conditionally required! https://dicom.innolitics.com/ciods/rt-dose/general-equipment/00280120
    
    # General Image
    dcm_file.ContentDate = dt.strftime('%Y%m%d')
    dcm_file.ContentTime = dt.strftime('%H%M%S.%f')
    dcm_file.InstanceNumber = 1
    dcm_file.PatientOrientation = ''
    
    # Image Plane
    dcm_file.SliceThickness = self.PixelSpacing[2]
    dcm_file.ImagePositionPatient = self.ImagePositionPatient
    dcm_file.ImageOrientationPatient = [1, 0, 0, 0, 1, 0] # HeadFirstSupine=1,0,0,0,1,0  FeetFirstSupine=-1,0,0,0,1,0  HeadFirstProne=-1,0,0,0,-1,0  FeetFirstProne=1,0,0,0,-1,0
    dcm_file.PixelSpacing = self.PixelSpacing[0:2]
    
    # Image pixel
    dcm_file.SamplesPerPixel = 1
    dcm_file.PhotometricInterpretation = 'MONOCHROME2'
    dcm_file.Rows = self.GridSize[1]
    dcm_file.Columns = self.GridSize[0]
    dcm_file.BitsAllocated = 16
    dcm_file.BitsStored = 16
    dcm_file.HighBit = 15
    dcm_file.BitDepth = 16
    dcm_file.PixelRepresentation = 0 # 0=unsigned, 1=signed
    #dcm_file.ColorType = 'grayscale'
    
    # multi-frame
    dcm_file.NumberOfFrames = self.GridSize[2]
    dcm_file.FrameIncrementPointer = pydicom.tag.Tag((0x3004, 0x000c))
    
    # RT Dose
    dcm_file.DoseUnits = 'GY'
    dcm_file.DoseType = 'PHYSICAL' # or 'EFFECTIVE' for RBE dose (but RayStation exports physical dose even if 1.1 factor is already taken into account)
    dcm_file.DoseSummationType = 'PLAN'
    dcm_file.GridFrameOffsetVector = list(np.arange(0, self.GridSize[2]*self.PixelSpacing[2], self.PixelSpacing[2]))
    dcm_file.DoseGridScaling = self.Image.max()/(2**dcm_file.BitDepth - 1)
    # pixel data
    dcm_file.PixelData = (self.Image/dcm_file.DoseGridScaling).astype(np.uint16).transpose(2,0,1).tostring() # ALTERNATIVE: self.Image.tobytes()

    #dcm_file.TissueHeterogeneityCorrection = 'IMAGE,ROI_OVERRIDE'
    # ReferencedPlan = pydicom.dataset.Dataset()
    # ReferencedPlan.ReferencedSOPClassUID = "1.2.840.10008.5.1.4.1.1.481.8" # ion plan
    # if(plan_uid == []): ReferencedPlan.ReferencedSOPInstanceUID = self.Plan_SOPInstanceUID
    # else: ReferencedPlan.ReferencedSOPInstanceUID = plan_uid
    # dcm_file.ReferencedRTPlanSequence = pydicom.sequence.Sequence([ReferencedPlan])
    
    # SOP common
    dcm_file.SpecificCharacterSet = 'ISO_IR 100'
    dcm_file.InstanceCreationDate = dt.strftime('%Y%m%d')
    dcm_file.InstanceCreationTime = dt.strftime('%H%M%S.%f')
    dcm_file.SOPClassUID = meta.MediaStorageSOPClassUID
    dcm_file.SOPInstanceUID = SOPInstanceUID

    # save dicom file
    print("Export dicom RTDOSE: " + OutputFile)
    dcm_file.save_as(OutputFile)