import SimpleITK
import numpy
import os

def convertToITK(dt,label):
    #convert to SimpleITKimage to display it
    imgArray=dt[label]
    org=dt['{}origin'.format(label)]
    orient=dt['{}orientation'.format(label)]
    spacing=[numpy.linalg.norm(x) for x in orient]
    print(spacing)
    orient=[orient[i]/spacing[i] for i in range(3)]
    direction=numpy.reshape(orient,9)
    print(direction)
    img = SimpleITK.GetImageFromArray(imgArray)
    img.SetOrigin(org)
    img.SetDirection(direction)
    img.SetSpacing(spacing)
    return img

def generateNiftyFromNPZ(infile):
    #generate nifty files for PET and CT stored in NPZ
    outdir=os.path.dirname(infile)
    baseName=os.path.basename(infile).replace('.npz','')
    dt=numpy.load(infile)
    ct=convertToITK(dt,'CT')
    pet=convertToITK(dt,'PET')
    SimpleITK.WriteImage(ct,os.path.join(outdir,f'{baseName}_CT.nii.gz'))
    SimpleITK.WriteImage(pet,os.path.join(outdir,f'{baseName}_PET.nii.gz'))
    del pet
    del ct