import pydicom
import numpy
import sys
import os
import nibabel

# load the DICOM files
def load(dir):
    files = []
    print('Loading: {}'.format(dir))
    for fname in os.listdir(dir):
        print("loading: {}".format(fname))
        files.append(pydicom.dcmread(os.path.join(dir,fname)))

    print("file count: {}".format(len(files)))

    # skip files with no SliceLocation (eg scout views)
    slices = []
    skipcount = 0
    for f in files:
        if hasattr(f, 'SliceLocation'):
            slices.append(f)
        else:
            skipcount = skipcount + 1

    print("skipped, no SliceLocation: {}".format(skipcount))

    # ensure they are in the correct order
    slices = sorted(slices, key=lambda s: s.SliceLocation)
    return slices;

    # pixel aspects, assuming all slices are the same

    for s in slices:
        print("Pixel spacing: {}".format(s.PixelSpacing))

    return;

def convertToNIfTI(slices):

    # create 3D array
    img_shape = list(slices[0].pixel_array.shape)
    img_shape.append(len(slices))
    img3d = numpy.zeros(img_shape)

    # get position and orientation parameters

    x0=numpy.array([float(f) for f in slices[0].ImagePositionPatient])
    x1=numpy.array([float(f) for f in slices[-1].ImagePositionPatient])
    n=(x1-x0)/(len(slices)-1.);

    f=numpy.array([float(f) for f in slices[0].ImageOrientationPatient])

    s=numpy.array([float(f) for f in slices[0].PixelSpacing])
    
    #create affine 
    a=numpy.zeros((4,4))
    f1=numpy.zeros(4)
    f1[0:3]=f[0:3]
    #swap to account for row/column to (c,r) position mismatch
    a[:,1]=f1*s[1]
    f2=numpy.zeros(4)
    f2[0:3]=f[3:6]
    a[:,0]=f2*s[0]
    nn=numpy.zeros(4)
    nn[0:3]=n[0:3]
    a[:,2]=nn

    xn=numpy.ones(4)
    xn[0:3]=x0[0:3]
    a[:,3]=xn

     

    # fill 3D array with the images from the files
    for i, s in enumerate(slices):
        img2d = s.pixel_array
        img3d[:, :, i] = img2d
    #orientation and pixel spacing
    img = nibabel.Nifti1Image(img3d, a)

    return img