
import json
import numpy as np
import pickle
import random
import tensorflow as tf
import tensorflow_datasets as tfds

from tensorflow import keras
from tensorflow.keras.models import Sequential, Model
from tqdm import tqdm

from PIL import Image


# model_path = "/home/ec2-user/gan_submission_1/mnist/mnist_v2/model_keras.h5"
# file_path = '/home/ec2-user/gan_submission_1/mnist/mnist_v2/'
# tfds_name = 'mnist'
# im_size = 28
# n_channels = 1
# n_samples = 10000  # 10000
# n_classes = 10

# model_path = "/home/ec2-user/gan_submission_1/svhn/svhn_v2/model_keras.h5"
# file_path = '/home/ec2-user/gan_submission_1/svhn/svhn_v2/'
# tfds_name = 'svhn_cropped'
# im_size = 32
# n_channels = 3
# n_samples = 26032
# n_classes = 10

model_path = "/home/ec2-user/gan_submission_1/cifar10/cifar10_v2/model_keras.h5"
file_path = '/home/ec2-user/gan_submission_1/cifar10/cifar10_v2/'
tfds_name = 'cifar10'
im_size = 32
n_channels = 3
n_samples = 10000  # 10000
n_classes = 10


baseline = 'closest'  # 'gan', 'closest', 'ranch'


def get_evaluation_curve(model, xs, ys_label, ys_target, saliency):
    # data_struct = {key: np.copy(xs) for key in ['zero', 'max', 'mid', 'flip', 'rand', 'rand33', 'rand55']}
    data_struct = {key: np.copy(xs) for key in ['flip']}


    saliency = np.mean(saliency, axis=-1)
    instructions = [
        sorted([(saliency[i, x, y], x, y) for x in range(im_size) for y in range(im_size)], reverse=True)
        for i in range(n_samples)]

    curve_prob_map = {k: [] for k in data_struct.keys()}
    mask33 = [[[True] * im_size for _ in range(im_size)] for _ in range(n_samples)]
    mask55 = [[[True] * im_size for _ in range(im_size)] for _ in range(n_samples)]

    for idx in tqdm(range(im_size * im_size)):
        for i in range(n_samples):
            _, x, y = instructions[i][idx]
            # data_struct['zero'][i, x, y, :] = 0.0
            # data_struct['max'][i, x, y, :] = 1.0
            # data_struct['mid'][i, x, y, :] = 0.5
            data_struct['flip'][i, x, y, :] = 1.0 - data_struct['flip'][i, x, y, :]
            # data_struct['rand'][i, x, y, :] = np.random.rand(n_channels)
            #
            # for dx in range(-1, 2):
            #     for dy in range(-1, 2):
            #         if 0 <= x + dx < 28 and 0 <= y + dy < 28 and mask33[i][x + dx][y + dy]:
            #             mask33[i][x + dx][y + dy] = False
            #             data_struct['rand33'][i, x + dx, y + dy, :] = np.random.rand(n_channels)
            #
            # for dx in range(-2, 3):
            #     for dy in range(-2, 3):
            #
            #         if 0 <= x + dx < 28 and 0 <= y + dy < 28 and mask55[i][x + dx][y + dy]:
            #             mask55[i][x + dx][y + dy] = False
            #             data_struct['rand55'][i, x + dx, y + dy, :] = np.random.rand(n_channels)

        for eval_type, xs_this in data_struct.items():
            ys_pred = model.predict(xs_this)

            orig_score = np.mean([ys_pred[i, l] for i, l in enumerate(ys_label)])
            target_score = np.mean([ys_pred[i, l] for i, l in enumerate(ys_target)])

            curve_prob_map[eval_type].append(orig_score - target_score)

    output = {}
    for eval_type, curve in curve_prob_map.items():
        output[eval_type] = curve
    return output


if __name__ == '__main__':
    # Load model
    model = keras.models.load_model(model_path)

    # Load data
    ds = tfds.load(tfds_name, split='test')
    dslist = list(tfds.as_numpy(ds.take(n_samples)))

    ys_target = np.random.RandomState(seed=222).randint(n_classes - 1, size=n_samples)
    xs, ys_label = [], []
    for ind, sample in enumerate(dslist):
        xs.append(sample['image'])
        ys_label.append(sample['label'])
        if ys_target[ind] >= sample['label']:
            ys_target[ind] += 1

    xs = np.array(xs)
    xs = xs / 255.5
    ys_label = np.array(ys_label)

    assert xs.shape == (n_samples, im_size, im_size, n_channels)
    assert ys_label.shape == (n_samples,)
    assert ys_label.shape == (n_samples,)

    # Load explanations
    if baseline == 'gan':
        explain_types = ["expected_intgrad", "intgrad", "intgrad_base", "deeplift", "deeplift_base", "occlusion", "occlusion_base"]
    else:
        explain_types = ["intgrad_base", "deeplift_base", "occlusion_base"]

    explain_map = {et: np.load('%s%s_%s_%s.npy' % (file_path, et, baseline, n_samples)) for et in explain_types}

    # Evaluation
    all_curves = {}
    for et, saliency in explain_map.items():
        all_curves[et] = get_evaluation_curve(model, xs, ys_label, ys_target, saliency)

        with open('%sevaluation_%s_%d_flip' % (file_path, baseline, n_samples), "wb") as fp:
            pickle.dump(all_curves, fp)


# with open('%sevaluation_%d.json' % (file_path, n_samples), 'r') as fp:
#     data = json.load(fp)
