import pydicom
import numpy as np
import scipy

class CTimage:

  def __init__(self):
    self.SeriesInstanceUID = ""
    self.PatientInfo = {}
    self.StudyInfo = {}
    self.FrameOfReferenceUID = ""
    self.ImgName = ""
    self.SOPClassUID = ""
    self.DcmFiles = []
    self.isLoaded = 0
    
    
    
  def print_CT_info(self, prefix=""):
    print(prefix + "CT series: " + self.SeriesInstanceUID)
    for ct_slice in self.DcmFiles:
      print(prefix + "   " + ct_slice)
      
      
  def resample_CT(self, newvoxelsize):
    ct = self.Image
    # Rescaling to the newvoxelsize if given in parameter

    source_shape = self.GridSize
    voxelsize = self.PixelSpacing
    #print("self.ImagePositionPatient",self.ImagePositionPatient, "source_shape",source_shape,"voxelsize",voxelsize)
    VoxelX_source = self.ImagePositionPatient[0] + np.arange(source_shape[0])*voxelsize[0]
    VoxelY_source = self.ImagePositionPatient[1] + np.arange(source_shape[1])*voxelsize[1]
    VoxelZ_source = self.ImagePositionPatient[2] + np.arange(source_shape[2])*voxelsize[2]
  
    target_shape = np.ceil(np.array(source_shape).astype(float)*np.array(voxelsize).astype(float)/newvoxelsize).astype(int)
    VoxelX_target = self.ImagePositionPatient[0] + np.arange(target_shape[0])*newvoxelsize[0]
    VoxelY_target = self.ImagePositionPatient[1] + np.arange(target_shape[1])*newvoxelsize[1]
    VoxelZ_target = self.ImagePositionPatient[2] + np.arange(target_shape[2])*newvoxelsize[2]
    #print("source_shape",source_shape,"target_shape",target_shape)
    if(all(source_shape == target_shape) and np.linalg.norm(np.subtract(voxelsize, newvoxelsize) < 0.001)):
      print("Image does not need filtering")
    else:
      # anti-aliasing filter
      sigma = [0, 0, 0]
      if(newvoxelsize[0] > voxelsize[0]): sigma[0] = 0.4 * (newvoxelsize[0]/voxelsize[0])
      if(newvoxelsize[1] > voxelsize[1]): sigma[1] = 0.4 * (newvoxelsize[1]/voxelsize[1])
      if(newvoxelsize[2] > voxelsize[2]): sigma[2] = 0.4 * (newvoxelsize[2]/voxelsize[2])
      
      if(sigma != [0, 0, 0]):
          print("Image is filtered before downsampling")
          ct = scipy.ndimage.gaussian_filter(ct, sigma)
      
              
          
      xi = np.array(np.meshgrid(VoxelX_target, VoxelY_target, VoxelZ_target))
      xi = np.rollaxis(xi, 0, 4)
      xi = xi.reshape((xi.size // 3, 3))
      
      # get resized ct
      ct = scipy.interpolate.interpn((VoxelX_source,VoxelY_source,VoxelZ_source), ct, xi, method='linear', fill_value=-1000, bounds_error=False).reshape(target_shape).transpose(1,0,2)
  
    self.PixelSpacing = newvoxelsize      
    self.GridSize = list(ct.shape) 
    self.NumVoxels = self.GridSize[0] * self.GridSize[1] * self.GridSize[2]
    self.Image = ct
    #print("self.ImagePositionPatient",self.ImagePositionPatient, "self.GridSize[0]",self.GridSize[0],"self.PixelSpacing",self.PixelSpacing)
      
    self.VoxelX = self.ImagePositionPatient[0] + np.arange(self.GridSize[0])*self.PixelSpacing[0]
    self.VoxelY = self.ImagePositionPatient[1] + np.arange(self.GridSize[1])*self.PixelSpacing[1]
    self.VoxelZ = self.ImagePositionPatient[2] + np.arange(self.GridSize[2])*self.PixelSpacing[2]
    self.isLoaded = 1

  def import_Dicom_CT(self):
    
    if(self.isLoaded == 1):
      print("Warning: CT serries " + self.SeriesInstanceUID + " is already loaded")
      return
  
    images = []
    SOPInstanceUIDs = []
    SliceLocation = np.zeros(len(self.DcmFiles), dtype='float')

    for i in range(len(self.DcmFiles)):
      file_path = self.DcmFiles[i]
      dcm = pydicom.dcmread(file_path)

      if(hasattr(dcm, 'SliceLocation') and abs(dcm.SliceLocation - dcm.ImagePositionPatient[2]) > 0.001):
        print("WARNING: SliceLocation (" + str(dcm.SliceLocation) + ") is different than ImagePositionPatient[2] (" + str(dcm.ImagePositionPatient[2]) + ") for " + file_path)

      SliceLocation[i] = float(dcm.ImagePositionPatient[2])
      images.append(dcm.pixel_array * dcm.RescaleSlope + dcm.RescaleIntercept)
      SOPInstanceUIDs.append(dcm.SOPInstanceUID)

    # sort slices according to their location in order to reconstruct the 3d image
    sort_index = np.argsort(SliceLocation)
    SliceLocation = SliceLocation[sort_index]
    SOPInstanceUIDs = [SOPInstanceUIDs[n] for n in sort_index]
    images = [images[n] for n in sort_index]
    ct = np.dstack(images).astype("float32")

    if ct.shape[0:2] != (dcm.Rows, dcm.Columns):
      print("WARNING: GridSize " + str(ct.shape[0:2]) + " different from Dicom Rows (" + str(dcm.Rows) + ") and Columns (" + str(dcm.Columns) + ")")

    MeanSliceDistance = (SliceLocation[-1] - SliceLocation[0]) / (len(images)-1)
    if(abs(MeanSliceDistance - dcm.SliceThickness) > 0.001):
      print("WARNING: MeanSliceDistance (" + str(MeanSliceDistance) + ") is different from SliceThickness (" + str(dcm.SliceThickness) + ")")

    self.SOPClassUID = dcm.SOPClassUID
    self.FrameOfReferenceUID = dcm.FrameOfReferenceUID
    self.ImagePositionPatient = [float(dcm.ImagePositionPatient[0]), float(dcm.ImagePositionPatient[1]), SliceLocation[0]]
    self.PixelSpacing = [float(dcm.PixelSpacing[0]), float(dcm.PixelSpacing[1]), MeanSliceDistance]
    self.GridSize = list(ct.shape) 
    self.NumVoxels = self.GridSize[0] * self.GridSize[1] * self.GridSize[2]
    self.Image = ct
    self.SOPInstanceUIDs = SOPInstanceUIDs
    self.VoxelX = self.ImagePositionPatient[0] + np.arange(self.GridSize[0])*self.PixelSpacing[0]
    self.VoxelY = self.ImagePositionPatient[1] + np.arange(self.GridSize[1])*self.PixelSpacing[1]
    self.VoxelZ = self.ImagePositionPatient[2] + np.arange(self.GridSize[2])*self.PixelSpacing[2]
    self.isLoaded = 1