# -*- coding: utf-8 -*-
"""
Script to compute metrics from nifti files
"""
import matplotlib.pyplot as plt
import SimpleITK as sitk
import os

def calculate_metrics(nifti_folder, output_folder, metric_filename):
    '''
    Input:
        nifti_folder: Folder containing nifti files -patientpack_scani.nii for i=1,6;
        output_folder: Folder where you want the metric file saved to
        metric_filename: name of metric file e.g "abcde.txt"
    Function makes a metric file:
        patientpack,metric for scan3, m for scan 5
    '''
    metricfile=open(output_folder+metric_filename,"w+")

    registration_method = sitk.ImageRegistrationMethod()


    registration_method.SetMetricSamplingStrategy(registration_method.RANDOM)
    registration_method.SetMetricSamplingPercentage(0.01)
    registration_method.SetInterpolator(sitk.sitkLinear)

    #registration_method.SetMetricAsJointHistogramMutualInformation()
    registration_method.SetMetricAsMattesMutualInformation()



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

            scan1= nifti_folder + filename
            scan3= nifti_folder + patientpack + "_scan3.nii"
            scan5= nifti_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)

            metric_3= registration_method.MetricEvaluate(scan1_image, scan3_image)
            metric_5= registration_method.MetricEvaluate(scan1_image, scan5_image)

            metricfile.write(patientpack+","+str(metric_3)+","+str(metric_5)+"\n")



nifti_folder= "D:/data/MPhys2019/PETCT-HN/12-11/nonrigid_fullheads/nonrigid_nifti/"
output_folder= "D:/data/MPhys2019/PETCT-HN/12-11/nonrigid_fullheads/"
metric_filename= "Metric_sitk_MATTES.txt"

calculate_metrics(nifti_folder,output_folder,metric_filename)