from tensorflow.keras.models import load_model
import os
import pickle
import urllib.request as http
from zipfile import ZipFile
from tensorflow.keras import utils

import tensorflow as tf
import numpy as np
from PIL import Image

from tensorflow.keras import layers as keras_layers
from tensorflow.keras import backend as K
from tensorflow.keras.datasets import cifar10
from tensorflow.keras.models import save_model, load_model


def load_cifar10(num_classes=3):
    """
    Downloads CIFAR-10 dataset, which already contains a training and test set,
    and return the first `num_classes` classes.
    Example of usage:

    >>> (x_train, y_train), (x_test, y_test) = load_cifar10()

    :param num_classes: int, default is 3 as required by the assignment.
    :return: the filtered data.
    """
    (x_train_all, y_train_all), (x_test_all, y_test_all) = cifar10.load_data()

    fil_train = tf.where(y_train_all[:, 0] < num_classes)[:, 0]
    fil_test = tf.where(y_test_all[:, 0] < num_classes)[:, 0]

    y_train = y_train_all[fil_train]
    y_test = y_test_all[fil_test]

    x_train = x_train_all[fil_train]
    x_test = x_test_all[fil_test]

    return (x_train, y_train), (x_test, y_test)

if __name__ == '__main__':
    
    _, (x_test, y_test) = load_cifar10()
    
    x_test_n = x_test / 255
    y_test_n = utils.to_categorical(y_test, 3)
    
    lrs = ["01", "0001"]
    nns = [16, 64]
    
    for lr in lrs:
        for nn in nns:
            # Load the trained models
            model_task1 = load_model('nn_task1_bonus/t1_bonus_%s_%d.h5' % (lr, nn))
            
            # Predict on the given samples
            y_pred_task1 = model_task1.predict(x_test_n)

            # Evaluate the missclassification error on the test set
            assert y_test_n.shape == y_pred_task1.shape
            test_loss, test_accuracy = model_task1.evaluate(x_test_n, y_test_n) # evaluate accuracy with proper function
            print("Accuracy model task 1 (%d neurons, 0.%s learning rate):" % (nn, lr), test_accuracy)