import numpy as np
import torch
import torch.nn as nn
import torchvision.datasets as dsets
import torchvision.transforms as transforms
from torch.autograd import Variable
from tqdm import tqdm
from typing import Optional, Union, Tuple, List, Sequence, Iterable
import math
from scipy.spatial.distance import euclidean
from torch.nn.modules.utils import _pair
from torchvision import models
from sklearn.metrics import jaccard_score
import matplotlib.pyplot as plt
import os
from models.Early_exit_models import Early_exit_conv_layer, Early_exit_classifier, Early_exit_confidence_layer


def conv_run_test(conv_layers, classifiers, confidences, threshold=None):

    batch_size = 1
    n_layers = len(conv_layers)

    transform = transforms.Compose(
        [transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

    test_data = dsets.CIFAR10(root = './data', train = False,
                    transform = transform)
    
    test_gen = torch.utils.data.DataLoader(dataset = test_data,
                                    batch_size = batch_size,
                                    shuffle = False)

# Run test over all the test data #

    exit_counts = np.zeros(n_layers)

    correct = np.zeros(n_layers)

    sum_confidences = np.zeros(n_layers)


    for i, (image, label) in tqdm(enumerate(test_gen)):

        # Called main_outs because this will be the inputs/outputs of the main body (convear layers)
        main_outs = image

        max_conf = 0 # for threshold = None


        for i in range(n_layers):

            classifier = classifiers["class_"+str(i)]
            conv_layer = conv_layers["conv_"+str(i)]
            decider = confidences["confidence_"+str(i)]

            main_outs = conv_layer(main_outs)
            flattened = torch.flatten(main_outs, start_dim=1)  

            if threshold is not None:

                confidence = decider(flattened).detach()[0].item()

                sum_confidences[i] += confidence

                # print("layer:",i,"confidence:",confidence,"threshold:",threshold,"accepted:",confidence >= threshold)

                if (confidence >= threshold):

                    # Add 1 to count for early exit on EE'd layer
                    exit_counts[i] += 1

                    scores = list(classifier(flattened).detach())
                    scores = list(scores[0].detach())
                    prediction = scores.index(max(scores))

                    if (prediction == label):
                        correct[i] += 1
                    # else:
                        # print("chose", prediction,"not",label)
                        # print("had confidence:", confidence)

                    break
                
                elif(i == n_layers-1):
                    # Adds Count for exit on final layer (none had enough confidence)
                    exit_counts[i] += 1

                    scores = list(classifier(flattened).detach())
                    scores = list(scores[0].detach())
                    prediction = scores.index(max(scores))

                    if (prediction == label):
                        correct[i] += 1
 
            else:

                conf = decider(flattened).detach()[0]
                if  conf > max_conf:

                    max_conf = conf
                    max_conf_ind = i
                    scores = list(classifier(flattened).detach())
                    scores = list(scores[0].detach())
                    
                    max_conf_pred = scores.index(max(scores))
                
        if threshold is None:

            # Add 1 to count for early exit on EE'd layer
            exit_counts[max_conf_ind] += 1

            prediction = max_conf_pred

            if (prediction == label):
                correct[max_conf_ind] += 1
            # else:
                # print("chose", prediction,"not",label)
                # print("had confidence:", max_conf)

    avg_confidences = np.zeros(n_layers)

    for i in range(n_layers):
        avg_confidences[i] = sum_confidences[i] / np.sum(exit_counts[i:])            

    total_accuracy = np.sum(correct)/len(test_gen)

    print("accuracy=", total_accuracy)

    accuracy = correct/exit_counts

    for i in range(n_layers):


        print("number of exits at layer",i,":",exit_counts[i])
        print("accuracy of exits at layer",i,":",accuracy[i])

    return exit_counts, accuracy, total_accuracy, avg_confidences

def conv_load_and_test_and_graph(channels=None, kernels=None, strides=None, pool_kernels=None, pool_strides=None, checkpoints=[], name='', threshold=0.8, classes=10):

    '''
    This function will load all of the pretrained networks give in the checkpoint param and will run a pass through the test data with it
    and will create a few figures on their test performance on a few metric. It will save these figures in the 'early_exit_results' folder

    Params: 

    - Layers: the layer sizes of our network in list form (e.g. [3072,400,10] would have (3072,400) as the size of convear layer 1) 

    - Checkpoints: a list of markers for checkpoints to use to get networks from the checkpoint folder (e.g. ['a_1','a_2'] would be the alternating net at 1 and 2 epochs )

    - Name: the name you'd like to save the figure under -- this will be the folder name 

    - Threshold: The threshold for acceptance at any layer compared to decider output

    '''

    paths = []

    # Creates a list of paths to the checkpoints 
    for cp in checkpoints:
        paths.append('../checkpoints/'+str(cp))

    # This is the path we will save results too, nested folder of threshold and the given name
    save_path = '../early_exit_results/threshold='+str(threshold)+'/'+str(name)+'/'

    # This will make the save path if it isn't made yet
    if not os.path.exists(save_path):
        os.makedirs(save_path, mode=0o777)
        os.chmod(save_path, mode=0o777)

    # Creating an empty dictionary for each type of layer
    conv_layers = {}
    classifiers = {}
    confidences = {}


    n_layers = len(kernels)
    classes = classes

    C,H,W = (3, 32, 32)

    for i in range(n_layers):

        #load all of the layers -- the main body, the classifiers, and the decision/confidence layer

        conv_layers["conv_"+str(i)] = Early_exit_conv_layer(in_channels=channels[i], out_channels=channels[i+1], kernel_size=kernels[i], stride=strides[i], padding=(0,0), pool_kernel=pool_kernels[i], pool_stride=pool_strides[i], pool_padding=(0,0))
        _,C,H,W = conv_layers["conv_"+str(i)].output_size(1,C,H,W)
        flattened = C*H*W

        classifiers["class_"+str(i)] = Early_exit_classifier(int(flattened), classes)
        confidences["confidence_"+str(i)] = Early_exit_confidence_layer(int(flattened))

    n_paths = len(paths)

    exit_counts = np.zeros((n_paths, n_layers))
    layer_accuracies = np.zeros((n_paths, n_layers))
    total_accuracies = np.zeros((n_paths))
    avg_confidences = np.zeros((n_paths, n_layers))


    for n, path in enumerate(paths):

        for i in range(n_layers):

            layer = torch.load(path + '/layer'+str(i)+'.pt')

            conv_layers["conv_"+str(i)].load_state_dict(layer["conv_"+str(i)])
            classifiers["class_"+str(i)].load_state_dict(layer["class_"+str(i)])
            confidences["confidence_"+str(i)].load_state_dict(layer["confidence_"+str(i)]) 
        
        exit_counts[n], layer_accuracies[n], total_accuracies[n], avg_confidences[n] = conv_run_test(conv_layers, classifiers, confidences, threshold=threshold)

    width = 1/ n_paths

    ## PLOTTING VARIOUS THINGS AND SAVING THEM IN 'early_exit_results' ################


    # Plotting the number of early exits by layer #####################################

    fig = plt.figure()

    for n in range(n_paths):
        plt.bar(np.arange(n_layers) + n*width, exit_counts[n], width=width, label=checkpoints[n])

    plt.legend()
    plt.title("early exit counts by layer -- threshold =" +str(threshold))
    plt.xlabel('layer number')
    plt.ylabel('number of early exits')

    fig.savefig(save_path+'early_exits.png')
    plt.close()

    # Plotting the avg accuracy of exited data by layer ###############################

    fig = plt.figure()

    for n in range(n_paths):
        plt.bar(np.arange(n_layers) + n*width, layer_accuracies[n], width=width, label=checkpoints[n])

    plt.legend()
    plt.title("Accuracy of layers output when chosen for EE -- threshold =" +str(threshold))
    plt.xlabel('layer number')
    plt.ylabel('avg accuracy of layer when chosen')

    fig.savefig(save_path+'early_exits_accuracies.png')
    plt.close()

    # Plotting the avg confidence by layer #############################################

    fig = plt.figure()

    for n in range(n_paths):
        plt.bar(np.arange(n_layers) + n*width, avg_confidences[n], width=width, label=checkpoints[n])

    plt.legend()
    plt.title("Average Confidences by layer with threshold=" +str(threshold))
    plt.xlabel('layer number')
    plt.ylabel('avg confidence of layer on seen data')

    fig.savefig(save_path+'avg_confidences.png')
    plt.close()

    # Plotting the total accuracy of the model ###########################################

    fig = plt.figure()

    for n in range(n_paths):
        plt.bar(np.arange(1) + n*width, total_accuracies[n], width=width, label=checkpoints[n])

    plt.legend()
    plt.title("total test accuracy with early exit -- threshold =" +str(threshold))
    plt.xlabel('layer number')
    plt.ylabel('test accuracy')

    fig.savefig(save_path+'total_accuracy.png')
    plt.close()
