"""Calculate the sensitivity of the plugin network to incorrect evidence.

This script exhaustively searches all potential evidence from the SUN397
fine-grained classification task with partial evidence, as discussed by
Koperski et al. in "Plugin Networks for Inference under Partial Evidence."

Example:
    Similar to the instructions in the readme, sensitivity.py should be run
    from the workspace/sun397 folder. The command is::

        $ python ../../pluginnet/sensitivity.py output/[] test
"""
import argparse
import csv
import os
import json
import torch
import numpy as np
import matplotlib.pyplot as plt
from pluginnet.sun397.testing import create_mode_pe as sun397_create_mode_pe
from pluginnet.sun397.testing import create_mode_base as sun397_create_mode_base

# pylint: disable=C0103,R0914
tasks = {'sun397_pe': sun397_create_mode_pe, 'sun397': sun397_create_mode_base}

def load_conf(conf_file):
    """Loads the JSON file.

    Args:
        conf_file: The path to the JSON config file.

    Returns:
        A dict containing the information from the JSON config file.
    """
    with open(conf_file, 'r') as fd:
        conf = json.load(fd)
        return conf

def main():
    """Runs the sensitivity analysis.

    Args:
        None

    Returns:
        None
    """
    parser = argparse.ArgumentParser(
        description='Find the sensitivity of our model to incorrect evidence.')
    parser.add_argument('model_dir')
    parser.add_argument('split', choices=['train', 'val', 'test'])
    parser.add_argument('-w', '--workers', default=4, type=int, metavar='N',
                        help='number of data loading workers (default: 4)')
    parser.add_argument('-b', '--batch-size', default=64, type=int,
                        metavar='N', help='mini-batch size')

    args = parser.parse_args()

    # Setup
    # This is entirely from the original codebase.
    conf = load_conf(os.path.join(args.model_dir, 'conf.json'))
    conf['split'] = args.split
    conf['base_model_file'] = os.path.join(
        args.model_dir, 'model_best_state_dict.pth.tar')
    test_set, net, _, _ = tasks[conf['task']](conf)
    test_loader = torch.utils.data.DataLoader(
        test_set, batch_size=1, shuffle=False, num_workers=args.workers, \
        pin_memory=torch.cuda.is_available(), drop_last=False)

    # Set the device.
    device = 'cpu'
    if torch.cuda.is_available():
        device = 'cuda'

    # Move the network to the device.
    net = net.eval().to(device)

    # Put all the potential pieces of evidence in one tensor,
    # for one pass through the network.
    all_evidence = torch.zeros((0, 3))
    for i in range(0, 2):
        for j in range(0, 2):
            for k in range(0, 2):
                if i == 0 and j == 0 and k == 0:
                    continue
                this_row = torch.tensor(np.array([[i, j, k]])).float()
                all_evidence = torch.cat((all_evidence, this_row), dim=0)

    all_evidence = all_evidence.to(device)

    # How many pieces of incorrect evidence cause additional error.
    causes_additional_error = 0
    # How many pieces of incorrect evidence there are.
    sample_count = 0
    # How many images respond to zero through 7 pieces of evidence.
    ae_histogram = torch.zeros(7)

    # Same as above, but only for examples where the classification is correct.
    correct_samples = 0 # Not quite the same as above I guess...
    causes_additional_error_and_correct = 0
    ae_correct_histogram = torch.zeros(7)

    write_list = []
    # Loop through all the test samples, no gradients.

    with torch.no_grad():
        for dataset_index, (input_data, target) in enumerate(test_loader):

            #image_to_save = to_pil(reverse_normalize(input_data[1]).squeeze())
            #image_to_save.save("out_images/"+str(dataset_index)+".jpg")
            # Move information to the device.
            evidence = input_data[0].to(device)
            image = input_data[1].to(device)
            target = target.to(device)

            write_dict = {'index': dataset_index}
            write_dict['evidence'] = str(evidence.cpu().int().numpy())

            # Find the prediction for every piece of evidence.
            guesses = net((all_evidence, image.repeat(7, 1, 1, 1))).argmax(
                dim=1)

            # What is the correct answer?
            correct_answer = target.argmax()
            write_dict['correct_answer'] = correct_answer.item()
            for guess_idx in range(guesses.shape[0]):
                write_dict[str(all_evidence[guess_idx].int().cpu().numpy())] =\
                        guesses[guess_idx].item()

            # And do the guesses match the correct answer?
            all_correct = (correct_answer.repeat(7) == guesses).float()

            solve_as_loaded = net((evidence, image)).argmax(dim=1)

            # Is the real (i.e., with the correct evidence) answer correct?
            correct_evidence_correct = (
                solve_as_loaded == correct_answer).float()

            # Calculate additional error for all examples.
            # If the correct guess is incorrect, there will never be
            # additional error.
            # Also, the correct answer will never contribute.
            additional_error = torch.nn.functional.relu(
                correct_evidence_correct-all_correct)

            # Count the pieces of incorrect evidence that cause
            # additional error.
            causes_additional_error += additional_error.sum().item()
            # Remember how many pieces of evidence in this image cause
            # additional error.
            ae_histogram[additional_error.sum().int().item()] += 1
            # And track the total number of incorrect evidences.
            sample_count += 7

            # Keep track specifically of instances where
            # the correct answer is correct.
            # The variables are mostly the same as above.
            if correct_evidence_correct:
                causes_additional_error_and_correct += \
                        additional_error.sum().item()
                ae_correct_histogram[additional_error.sum().int().item()] += 1
                correct_samples += 1

            write_list.append(write_dict)

    # save specific output
    with open("evidence_results.csv", 'w') as csvfile:
        writer = csv.DictWriter(csvfile, fieldnames=write_list[0].keys())
        writer.writeheader()
        for row in write_list:
            writer.writerow(row)
    # Print the results.
    print("There are " + str(sample_count) + " incorrect pieces of evidence.")
    print("Of these, " + str(causes_additional_error) + \
          " cause additional error")
    print("The histogram is: " + str(ae_histogram))
    print("There are a total of " + str(correct_samples) + \
          " correctly classified samples.")
    print("Of which " + str(causes_additional_error_and_correct) + \
          " incorrect pieces of evidence cause additional error.")
    print("The histogram is: " + str(ae_correct_histogram))

    from IPython import embed
    embed()
    # and create plots
    plt.bar(np.arange(7), ae_histogram/ae_histogram.sum()*100)
    plt.xlabel("# labels which cause additional error")
    plt.ylabel("% of Images")
    plt.savefig("error-dist-all-"+str(args.split)+".eps", bbox_inches="tight")

    plt.cla()
    plt.bar(np.arange(7), ae_correct_histogram/ae_correct_histogram.sum()*100)
    plt.xlabel("# labels which cause additional error")
    plt.ylabel("% of Images")
    plt.savefig(f"error-dist-correctonly-{args.split}.eps", bbox_inches="tight")

if __name__ == '__main__':
    main()
