import numpy as np
from numpy import array
from numpy.random import randint
import os
import sys
import PIL
from imageio import imread
from math import log10, sqrt
import matplotlib.pyplot as plt
plt.switch_backend('agg')

#to calculate PSNR and MSE
def PSNR_MSE(original, encoded):
    original = original.astype('int32')
    encoded = encoded.astype('int32')
    mse = np.mean((original - encoded) ** 2)
    if(mse == 0):  # MSE is zero means no noise is present in the signal .
                  # Therefore PSNR have no importance.
        return 100, 0.0
    max_pixel = 255.0
    psnr = 20 * log10(max_pixel / sqrt(mse))
    return psnr, mse


# Takes list of images and provide HR images in form of numpy array
def hr_images(images):
    images_hr = array(images)
    return images_hr

# Takes list of images and provide LR images in form of numpy array
def lr_images(images_real , downscale):
    
    images = []
    for img in  range(len(images_real)):
        images.append(np.array(PIL.Image.fromarray(images_real[img]).resize([images_real[img].shape[1]//downscale, images_real[img].shape[0]//downscale],resample=PIL.Image.BICUBIC)))
        images_lr = array(images)
        # print(images_lr.shape)
    return images_lr
    
def normalize(input_data):

    return (input_data.astype(np.float32) - 127.5)/127.5 
    
def denormalize(input_data):
    input_data = (input_data + 1) * 127.5
    return input_data.astype(np.uint8)
   
 
def load_path(path): # this module load all folders in the data folder
    directories = []
    if os.path.isdir(path):

        directories.append(path)
    for elem in os.listdir(path):
        if os.path.isdir(os.path.join(path,elem)):

            directories = directories + load_path(os.path.join(path,elem))
            directories.append(os.path.join(path,elem))

    return directories
    
def load_data_from_dirs(dirs, ext):
    files = []
    file_names = []
    count = 0
    for d in dirs:   
        for f in os.listdir(d):          
            if f.endswith(ext):
                image = imread(os.path.join(d,f))# this line change the color pallets of an image
                if len(image.shape) > 2:
                    files.append(image)
                    file_names.append(os.path.join(d,f))
                count = count + 1             
    return files     

def load_training_data(directory, ext, ds_factor, image_shape, number_of_images = 10, train_test_ratio = 0.8):
    print('loading images from %s'%directory)
    number_of_train_images = int(number_of_images * train_test_ratio)
    files = load_data_from_dirs(load_path(directory), ext)
    
    if len(files) < number_of_images:
        print(number_of_images)
        print("Number of image files are less then you specified")
        print("Please reduce number of images to %d" % len(files))
        sys.exit()
        
    test_array = array(files)
    # print(test_array.shape)
    if len(test_array.shape) < 3:
        print("Images are of not same shape")
        print("Please provide same shape images")
        sys.exit()
    print("Total number of images %d" % len(files))
    x_train = files[:number_of_train_images]
    x_test = files[number_of_train_images:number_of_images]
   
    x_train_hr = hr_images(x_train)
    x_train_hr = normalize(x_train_hr)
    
    x_train_lr = lr_images(x_train, ds_factor)
    x_train_lr = normalize(x_train_lr)
    
    x_test_hr = hr_images(x_test)
    x_test_hr = normalize(x_test_hr)
    
    x_test_lr = lr_images(x_test, ds_factor)
    x_test_lr = normalize(x_test_lr)
    
    x_train_lr = x_train_lr[:,:image_shape[0]//ds_factor, :image_shape[1]//ds_factor, :3]
    x_train_hr = x_train_hr[:,:image_shape[0], :image_shape[1], :3]
    x_test_lr = x_test_lr[:,:image_shape[0]//ds_factor, :image_shape[1]//ds_factor, :3]
    x_test_hr = x_test_hr[:,:image_shape[0], :image_shape[1], :3]
    return x_train_lr, x_train_hr, x_test_lr, x_test_hr

def load_validation_data(directory, ext, ds_factor, image_shape,  number_of_images): # by me
    
    files = load_data_from_dirs(load_path(directory), ext)
    x_val_hr = hr_images(files[:number_of_images])
    print(np.shape(x_val_hr))
    # x_test_hr = lr_images(files, 2)
    
    x_val_lr = lr_images(x_val_hr, ds_factor)
    x_val_hr = normalize(x_val_hr)
    
    x_val_lr = normalize(x_val_lr)
    x_val_lr = x_val_lr[:,:image_shape[0]//ds_factor, :image_shape[1]//ds_factor, :3]
    x_val_hr = x_val_hr[:,:image_shape[0], :image_shape[1], :3]
    return x_val_hr, x_val_lr
     

def val_patches_generator(x_val_hr, patch_sz):
    img_list = []
    x = x_val_hr[0].shape[0]//patch_sz # because the network was trained for only one upsampling, so we have to use the half of the image_shape
    y = x_val_hr[0].shape[1]//patch_sz
    for img in x_val_hr:
        for i in range(x):
            for j in range(y):
                img_list.append(img[i*patch_sz:i*patch_sz+patch_sz,j*patch_sz:j*patch_sz+patch_sz,:])
    
    return(np.array(img_list))
    
def val_mse(generator, val_patches_hr, val_patches_lr): # 
    image_hr = denormalize(val_patches_hr)
    image_sr = denormalize(generator.predict(val_patches_lr))
    mse = np.mean((image_hr - image_sr) ** 2)
    return mse


def plot_generated_images_psnr_olp3(output_dir, epoch, model, x_test_hr, x_test_lr, patch_hr, us, val = False, dim=(1, 3), figsize=(15, 5)): # this function is to avoide OOM error, one image is tested at a time
    
    examples = x_test_hr.shape[0] # number of test images
    print('total sample test images %d' %examples)
    idx = randint(0, examples)
    image_hr = denormalize(x_test_hr[idx])
    img_in = x_test_lr[idx]
   
    patch_lr = patch_hr[0]//us
    x = img_in.shape[0] 
    y = img_in.shape[1]    
    img_ouput = np.zeros((x*us, y*us, 3), 'float32') 

    olp = 16
    if patch_lr - olp <2:
        olp = patch_lr - 2
    
    pin = [] 
    ## first correct the borders and corners 
    pin.append(img_in[0:patch_lr, 0:patch_lr, :])
    pin.append(img_in[0:patch_lr, -patch_lr:, :])
    pin.append(img_in[-patch_lr:, 0:patch_lr, :])
    pin.append(img_in[-patch_lr:, -patch_lr:, :])
    #### correct top/bottom y rows
    for j in range(1,(y - patch_lr)//(patch_lr-olp)+1,1): 
        pin.append(img_in[0:patch_lr:, j*(patch_lr-olp):j*(patch_lr-olp)+patch_lr,:])
        pin.append(img_in[-patch_lr:, j*(patch_lr-olp):j*(patch_lr-olp)+patch_lr,:])
    for i in range(1,(x - patch_lr)//(patch_lr-olp)+1,1):
        ### correct first x colomns
        pin.append(img_in[i*(patch_lr-olp):i*(patch_lr-olp)+patch_lr, 0:patch_lr,:])
        ### correct last x colomns
        pin.append(img_in[i*(patch_lr-olp):i*(patch_lr-olp)+patch_lr, -patch_lr:,:])
        for j in range(1,(y - patch_lr)//(patch_lr-olp)+1,1):
            pin.append(img_in[i*(patch_lr-olp):i*(patch_lr-olp)+patch_lr, j*(patch_lr-olp):j*(patch_lr-olp)+patch_lr,:])
    ############ calculate the output  ################################
    pout = model.predict(np.array(pin)) 
    ##### Reconstructing the output image  ############################################
    p = 0                
    ## first correct the borders and corners 
    img_ouput[0:patch_hr[1], 0:patch_hr[1], :] = pout[p]
    p+=1
    img_ouput[0:patch_hr[1], -patch_hr[1]:, :] = pout[p]
    p+=1
    img_ouput[-patch_hr[1]:, 0:patch_hr[1], :] = pout[p]
    p+=1
    img_ouput[-patch_hr[1]:, -patch_hr[1]:, :] = pout[p]
    p+=1
    #### correct top/bottom y rows
    for j in range(1,(y - patch_lr)//(patch_lr-olp)+1,1): 
        img_ouput[0:patch_hr[1]-olp, j*(patch_hr[1]-olp*us) + olp : j*(patch_hr[1]-olp*us) -olp + patch_hr[1],:] = pout[p][0:-olp, olp:-olp,:]
        p+=1
        img_ouput[-patch_hr[1]+olp:, j*(patch_hr[1]-olp*us) + olp : j*(patch_hr[1]-olp*us) -olp + patch_hr[1],:] = pout[p][olp:, olp:-olp,:]
        p+=1
    for i in range(1,(x - patch_lr)//(patch_lr-olp)+1,1):
        ### correct first x colomns
        img_ouput[i*(patch_hr[0]-olp*us) + olp : i*(patch_hr[0]-olp*us) - olp + patch_hr[0], 0:patch_hr[0]-olp,:] = pout[p][olp:-olp, 0:-olp,:]
        p+=1
        ### correct last x colomns
        img_ouput[i*(patch_hr[0]-olp*us) + olp : i*(patch_hr[0]-olp*us) - olp + patch_hr[0], -patch_hr[0]+olp:,:] = pout[p][olp:-olp, olp:,:]
        p+=1
        for j in range(1,(y - patch_lr)//(patch_lr-olp)+1,1):
            img_ouput[i*(patch_hr[0]-olp*us) + olp : i*(patch_hr[0]-olp*us) - olp + patch_hr[0],   j*(patch_hr[1]-olp*us) + olp : j*(patch_hr[1]-olp*us) -olp + patch_hr[1],:] = pout[p][olp:-olp, olp:-olp,:]
            p+=1
                    
    img_ouput2 = denormalize(img_ouput)

    image_lr = denormalize(img_in)#
    img_bicubic = np.array(PIL.Image.fromarray(image_lr).resize([image_lr.shape[1]*us, image_lr.shape[0]*us],resample=PIL.Image.BICUBIC))#
    plt.figure(figsize=figsize)#
    psnrb, mseb = PSNR_MSE(image_hr[4:-4, 4:-4, :], img_bicubic[4:-4, 4:-4, :])#
    
    plt.subplot(dim[0], dim[1], 1, title = f"Bicubic, PSNR:{psnrb:0.2f}, MSE:{mseb:0.2f}")#
    plt.tight_layout()#
    plt.imshow(img_bicubic[4:-4, 4:-4, :], interpolation='nearest')#
    # plt.set_title('title', fontsize=16)
    plt.axis('off')#

    psnrg, mseg = PSNR_MSE(image_hr[4:-4, 4:-4, :], img_ouput2[4:-4, 4:-4, :])    
    plt.subplot(dim[0], dim[1], 2, title = f"SRResNetap, PSNR:{psnrg:0.2f}, MSE:{mseg:0.2f}")
    plt.tight_layout()
    plt.imshow(img_ouput2[4:-4, 4:-4, :], interpolation='nearest')
    plt.axis('off')
    
    plt.subplot(dim[0], dim[1], 3, title="Original")
    plt.imshow(image_hr[4:-4, 4:-4, :], interpolation='nearest')
    plt.axis('off')
    
    plt.tight_layout()
    if val:
        plt.savefig(output_dir + '\\val_%d_%0.2f_%d.png' %(idx, mseg, epoch))
        plt.close()
        return psnrg/psnrb, mseg/mseb, idx
    plt.savefig(output_dir + '\gen_%d_%0.2f_%d.png' %(idx, mseg, epoch))
    plt.close()
