# -*- coding: utf-8 -*-
"""Assignment-3_notebook.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1vM9Ot32fNYF4wctXXrDhL4jk-FdOdT4E

Import and setup some auxiliary functions
"""

import torch
from logistic_regression import logistic_regression
from torchvision import transforms, datasets
import numpy as np
import timeit
from collections import OrderedDict
from pprint import pformat
from torch.utils.data.sampler import *
from tqdm import tqdm

torch.multiprocessing.set_sharing_strategy('file_system')

def compute_score(acc, min_thres, max_thres):
    if acc <= min_thres:
        base_score = 0.0
    elif acc >= max_thres:
        base_score = 100.0
    else:
        base_score = float(acc - min_thres) / (max_thres - min_thres) \
                     * 100
    return base_score


def run(algorithm, dataset_name, filename):
    start = timeit.default_timer()
    predicted_test_labels, gt_labels = algorithm(dataset_name)
    if predicted_test_labels is None or gt_labels is None:
      return (0, 0, 0)
    stop = timeit.default_timer()
    run_time = stop - start
    
    np.savetxt(filename, np.asarray(predicted_test_labels))

    correct = 0
    total = 0
    for label, prediction in zip(gt_labels, predicted_test_labels):
      total += label.size(0)
      correct += (prediction.cpu().numpy() == label.cpu().numpy()).sum().item()   # assuming your model runs on GPU
      
    accuracy = float(correct) / total
    
    print('Accuracy of the network on the 10000 test images: %d %%' % (100 * correct / total))
    return (correct, accuracy, run_time)

"""Main loop. Run time and total score will be shown below."""

def run_on_dataset(dataset_name, filename):
    if dataset_name == "MNIST":
        min_thres = 0.82
        max_thres = 0.92

    elif dataset_name == "CIFAR10":
        min_thres = 0.28
        max_thres = 0.38

    correct_predict, accuracy, run_time = run(logistic_regression, dataset_name, filename)

    score = compute_score(accuracy, min_thres, max_thres)
    result = OrderedDict(correct_predict=correct_predict,
                         accuracy=accuracy, score=score,
                         run_time=run_time)
    return result, score


def main():
    filenames = { "MNIST": "predictions_mnist_YourName_IDNumber.txt", "CIFAR10": "predictions_cifar10_YourName_IDNumber.txt"}
    result_all = OrderedDict()
    score_weights = [0.5, 0.5]
    scores = []
    for dataset_name in ["MNIST","CIFAR10"]:
        result_all[dataset_name], this_score = run_on_dataset(dataset_name, filenames[dataset_name])
        scores.append(this_score)
    total_score = [score * weight for score, weight in zip(scores, score_weights)]
    total_score = np.asarray(total_score).sum().item()
    result_all['total_score'] = total_score
    with open('result.txt', 'w') as f:
        f.writelines(pformat(result_all, indent=4))
    print("\nResult:\n", pformat(result_all, indent=4))


main()
