# -*- coding: utf-8 -*-
"""
Reads in a folder containing nifti scans and saves scan1,3,5 windowed into output folder
window means all pixels below a specified intensity are set to zero
"""


import SimpleITK as sitk
import numpy as np
import matplotlib.pyplot as plt
import os



folder= "D:/data/MPhys2019/PETCT-HN/26-11/NonRigid/BigDoseCropNonRigid/"
output_folder= "D:/data/MPhys2019/PETCT-HN/03-01/"



registration_method = sitk.ImageRegistrationMethod()
registration_method.SetMetricAsCorrelation()

for filename in os.listdir(folder):
    if "scan1" in filename:
        #print(filename.replace("_scan1.nii",""))
        patientpack= filename.replace("_scan1.nii","")
        print(patientpack)

        scan1= folder + filename
        scan3= folder + patientpack + "_scan3.nii"
        scan5= folder + patientpack + "_scan5.nii"

        scan1_image= sitk.ReadImage(scan1, sitk.sitkFloat32)
        scan3_image= sitk.ReadImage(scan3, sitk.sitkFloat32)
        scan5_image= sitk.ReadImage(scan5, sitk.sitkFloat32)

        window = sitk.IntensityWindowingImageFilter()
        #window.SetWindowMinimum(930)
        window.SetWindowMinimum(1224)
        window.SetWindowMaximum(5000)
        window.SetOutputMinimum(0)
        window.SetOutputMaximum(5000)

        scan1_image_window = window.Execute(scan1_image)
        scan3_image_window = window.Execute(scan3_image)
        scan5_image_window = window.Execute(scan5_image)

        writer = sitk.ImageFileWriter()
        writer.SetFileName(output_folder+filename)
        writer.Execute(scan1_image_window)

        writer.SetFileName(output_folder+patientpack+"_scan3.nii")
        writer.Execute(scan3_image_window)

        writer.SetFileName(output_folder+patientpack+"_scan5.nii")
        writer.Execute(scan5_image_window)