import numpy as np
import glob
import os
from scipy import misc
from skimage.measure import compare_ssim
from skimage.color import rgb2ycbcr, rgb2yuv

from skimage.measure import compare_psnr
from utilities import preprocess, downsample


class Benchmark:
    """A collection of images to test a model on."""

    def __init__(self, path, name):
        self.path = path
        self.name = name
        # self.images_lr, self.names = self.load_images_by_model(model='LR')
        self.images_hr, self.names = self.load_images_by_model(model='HR')
        self.images_lr = []
        for img in self.images_hr:
            self.images_lr.append(downsample(img, 4))

    def load_images_by_model(self, model, file_format='png'):
        """Loads all images that match '*_{model}.{file_format}' and returns sorted list of filenames and names"""
        # Get files that match the pattern
        filenames = sorted(glob.glob(os.path.join(self.path, '*_' + model + '.' + file_format)))
        # Extract name/prefix eg: '/.../baby_LR.png' -> 'baby'
        names = [os.path.basename(x).split('_')[0] for x in filenames]
        return self.load_images(filenames), names

    def load_images(self, images):
        """Given a list of file names, return a list of images"""
        out = []
        for image in images:
            out.append(misc.imread(image, mode='RGB').astype(np.uint8))
        return out

    def deprocess(self, image):
        """Deprocess image output by model (from -1 to 1 float to 0 to 255 uint8)"""
        image = np.clip(255 * 0.5 * (image + 1.0), 0.0, 255.0).astype(np.uint8)
        return image

    def luminance(self, image):
        # Get luminance
        lum = rgb2ycbcr(image)[:, :, 0]
        # Crop off 4 border pixels
        lum = lum[4:lum.shape[0] - 4, 4:lum.shape[1] - 4]
        # lum = lum.astype(np.float64)
        return lum

    def PSNR(self, gt, pred):
        # gt = gt.astype(np.float64)
        # pred = pred.astype(np.float64)
        # mse = np.mean((pred - gt)**2)
        # psnr = 10*np.log10(255*255/mse)
        # return psnr
        return compare_psnr(gt, pred, data_range=255)

    def SSIM(self, gt, pred):
        ssim = compare_ssim(gt, pred, data_range=255, gaussian_weights=True)
        return ssim

    def test_images(self, gt, pred):
        """Applies metrics to compare image lists pred vs gt"""
        avg_psnr = 0
        avg_ssim = 0
        individual_psnr = []
        individual_ssim = []

        for i in range(len(pred)):
            # compare to gt
            psnr = self.PSNR(self.luminance(gt[i]), self.luminance(pred[i]))
            ssim = self.SSIM(self.luminance(gt[i]), self.luminance(pred[i]))
            # save results to log_path ex: 'results/experiment1/Set5/baby/1000.png'
            # if save_images:
            #  path = os.path.join(log_path, self.name, self.names[i])
            # gather results
            individual_psnr.append(psnr)
            individual_ssim.append(ssim)
            avg_psnr += psnr
            avg_ssim += ssim

        avg_psnr /= len(pred)
        avg_ssim /= len(pred)
        return avg_psnr, avg_ssim, individual_psnr, individual_ssim

    def validate(self):
        """Tests metrics by using images output by other models"""
        for model in ['bicubic', 'SRGAN-MSE', 'SRGAN-VGG22', 'SRGAN-VGG54', 'SRResNet-MSE', 'SRResNet-VGG22']:
            model_output, _ = self.load_images_by_model(model)
            psnr, ssim, _, _ = self.test_images(self.images_hr, model_output)
            print('Validate %-6s for %-14s: PSNR: %.2f, SSIM: %.4f' % (self.name, model, psnr, ssim))

    def save_image(self, image, path):
        if not os.path.exists(os.path.split(path)[0]):
            os.makedirs(os.path.split(path)[0])
        misc.toimage(image, cmin=0, cmax=255).save(path)

    def save_images(self, images, log_path, iteration):
        count = 0
        for output, lr, hr, name in zip(images, self.images_lr, self.images_hr, self.names):
            # Save output
            path = os.path.join(log_path, self.name, name, '%d_out.png' % iteration)
            self.save_image(output, path)
            # Save ground truth
            path = os.path.join(log_path, self.name, name, '%d_hr.png' % iteration)
            self.save_image(hr, path)
            # Save low res
            path = os.path.join(log_path, self.name, name, '%d_lr.png' % iteration)
            self.save_image(lr, path)

            # Hack so that we only do first 14 images in BSD100 instead of the whole thing
            count += 1
            if count >= 14:
                break

    def evaluate(self, sess, g_y_pred, log_path=None, iteration=0):
        """Evaluate benchmark, returning the score and saving images."""
        pred = []
        for i, lr in enumerate(self.images_lr):
            # feed images 1 by 1 because they have different sizes
            lr = lr / 255.0
            # print(lr.shape)
            # print(lr[np.newaxis].shape)
            output = sess.run(g_y_pred, feed_dict={'d_training:0': False, 'g_training:0': False,
                                                   'input_lowres:0': lr[np.newaxis]})
            '''
            e.g. lr.shape=(128,128,3)
            lr[np.newaxis].shape=(1,128,128,3)
            '''
            # deprocess output
            pred.append(self.deprocess(np.squeeze(output, axis=0)))
        # save images
        if log_path:
            self.save_images(pred, log_path, iteration)
        return self.test_images(self.images_hr, pred)
    
    
    def eval(self, sess, g_y_pred, log_path=None, iteration=0):
        pred = []
        
        for i, input in enumerate(self.images_hr):
            input = input/255.0
            output = sess.run(g_y_pred, feed_dict={'d_training:0': False, 'g_training:0': False,
                                                   'input_lowres:0': input[np.newaxis]})
            pred.append(self.deprocess(np.squeeze(output, axis=0)))
                # save images
        if log_path:
            self.save_images(pred, log_path, iteration)
        return self.test_images(self.images_hr, pred)
                                                   