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()
    
    # Load the trained models
    model_task1 = load_model('nn_task1/nn_task1.h5')
    x_test_n = x_test / 255
    y_test_n = utils.to_categorical(y_test, 3)

    # Predict on the given samples
    #for example
    y_pred_task1 = model_task1.predict(x_test_n)

    # Evaluate the missclassification error on the test set
    # for example
    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:", test_accuracy)