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, begin, end):
    # THRESHOLD = 3
    THRESHOLD = 5
    return np.mean(cors[begin+1:end]) - np.std(cors[begin+1:end]) * THRESHOLD

def lower_bound(cors, begin, end):
    # Not so good
    # THRESHOLD = 1.5
    THRESHOLD = 2
    Q1 = np.percentile(cors[begin+1:end], 25, interpolation = 'lower')
    Q3 = np.percentile(cors[begin+1:end], 75, interpolation = 'higher')
    IQR = Q3 - Q1
    return Q1 - THRESHOLD * IQR

NewSer = None

def check_low(cors, begin, end):
    if  end - begin < 2:
        return
    mini = np.min(cors[begin+1:end])
    if mini > COR_ABS_THRESHOLD and mini > lower_bound(cors, begin, end):
        #  exit()
        return
    argmin =  np.argmin(cors[begin+1:end]) + begin+1
    print(begin, end, lower_bound(cors, begin, end), mini, argmin)
    NewSer[argmin] = 1
    check_low(cors, begin, argmin)
    check_low(cors, argmin, end)
    #  exit()


def main():
    global NewSer
    old_shape = None
    old_array = None
    old_cor = COR_ABS_THRESHOLD
    fileNames = []
    shapes = []
    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 = 0
        else:
            if old_shape != shape:
                cor =0
            else:
                cor = np.corrcoef(old_array.flat, array.flat)[0,1]

        fileNames.append(jpg_path)
        shapes.append(shape)
        cors.append(cor)

        old_array = array
        old_shape = shape
        old_cor = cor

    length = len(fileNames)
    # print(length)
    # exit()

    for i in range(length):
        # print(i)
        print(fileNames[i], i, shapes[i], cors[i])

    NewSer =  np.zeros(length)
    check_low(cors, 0, length)

    for i in range(length):
        # print(i)
        print(fileNames[i], i, shapes[i], cors[i], '***' if NewSer[i] else '')



if __name__ == '__main__':
    main()