import os
import statistics

import imageio

from PIL import Image, ImageFilter, ImageMath

import numpy as np
import SimpleITK as sitk

STUDY_PATH = "/media/nfs/SRS/storage/0/CT Without Contrast-Brain_55121720"
STUDY_PATH = '/media/nfs/SRS/storage/0/MRI With_Without Contrast--Brain_54141890'

MODEL_PATH = '/home/xfr/nni/model-5-64/TwNuKtj7/best_zdoyO.pth'


# Write image series using SimpleITK
def flush_file(shape, fileNames):
    if len(fileNames) > 1:

        xy = min(shape)

        outfile = '%s.nii.gz' % os.path.basename(fileNames[0]).split('.')[0]
        img = sitk.ReadImage(fileNames)
        img.SetSpacing([1.0,1.0, 1.0*xy/len(fileNames)])
        sitk.WriteImage(img, outfile)


COR_ABS_THRESHOLD = 0.5
COR_REL_THRESHOLD = 0.8

def lower_bound(cors):
    THRESHOLD = 3
    if len(cors) < 2:
        return 0
    
    return min(statistics.mean(cors) - statistics.stdev(cors) * THRESHOLD, min(cors[:-1]))

def lower_bound2(cors):
    THRESHOLD = 1.5
    if len(cors) < 1:
        return 0
    Q1 = np.percentile(cors, 25, interpolation = 'lower')
    Q3 = np.percentile(cors, 25, interpolation = 'higher')
    IQR = Q3 - Q1
    return min(Q1 - THRESHOLD * IQR, min(cors[:-1]))

 
def main():
    old_shape = None
    old_array = None
    old_cor = COR_ABS_THRESHOLD
    fileNames = []
    cors = []

    for jpg_file in sorted(os.listdir(STUDY_PATH)):
        jpg_path = os.path.join(STUDY_PATH, jpg_file)

        array = np.asarray(Image.open(jpg_path).convert('L'))
        shape = array.shape

#        LB = lower_bound(cors)

        if not fileNames:
            cor = COR_ABS_THRESHOLD
        else:
            if old_shape != shape:
                cor =0
            else:
                # cor = correlate (old_array, array)
                cor = np.corrcoef(old_array.flat, array.flat)[0,1]
                cors.append(cor)

        # if cor < COR_ABS_THRESHOLD or cor < old_cor * COR_REL_THRESHOLD:
        LB = lower_bound(cors)
        if cor < COR_ABS_THRESHOLD or cor < LB:
            flush_file(old_shape, fileNames)
            fileNames = [jpg_path]
            cors = []
            mark = '**********'
        else:
            fileNames.append(jpg_path)
            mark = len(fileNames)
        print('%s %.4f %.4f %s' %(jpg_file,cor, LB, mark))

        old_array = array
        old_shape = shape
        old_cor = cor

    flush_file(old_shape, fileNames)




if __name__ == '__main__':
    main()