from collections import defaultdict
import os
from pathlib import Path

import joblib
import numpy as np
import pandas as pd
import torch

from .transforms.util import get_transforms
from .engine.tester import Tester
from .datasets import Searchstims, VOCDetection
from .utils.general import make_save_path


def test(csv_file,
         dataset_type,
         net_name,
         number_nets_to_train,
         epochs_list,
         batch_size,
         restore_path,
         test_results_save_path,
         configfile,
         random_seed,
         root=None,
         num_classes=2,
         pad_size=500,
         embedding_n_out=512,
         loss_func='CE',
         method='transfer',
         mode='classify',
         num_workers=4,
         data_parallel=False):
    """measure accuracy of trained convolutional neural networks on test set of visual search stimuli

    Parameters
    ----------
    csv_file : str
        name of .csv file containing prepared data sets.
        Generated by searchnets.data.split function from a csv created by the searchstims library.
    dataset_type : str
        one of {'searchstims', 'VSD'}. Specifies whether dataset is images generated by searchstims package, or
        images from Pascal-VOC dataset that were used to create the Visual Search Difficulty 1.0 dataset.
    net_name : str
        name of convolutional neural net architecture to train.
        One of {'alexnet', 'VGG16'}
    number_nets_to_train : int
        number of training "replicates"
    epochs_list : list
        of training epochs. Replicates will be trained for each
        value in this list. Can also just be one value, but a list
        is useful if you want to test whether effects depend on
        number of training epochs.
    batch_size : int
        number of samples in a batch of training data
    restore_path : str
        path to directory where checkpoints and train models were saved
    test_results_save_path : str
        path to directory where results from measuring accuracy on test set should be saved
    configfile : str
        filename of config.ini file. Used (without .ini extension) as name for output file
        that is saved in test_results_save_path.
    random_seed : int
        to seed random number generator
    root : str
        path to dataset root. Used with VOCDetection dataset to specify where VOC data was downloaded to.
    num_classes : int
        number of classes. Default is 2 (target present, target absent).
    pad_size : int
        size to which images in PascalVOC / Visual Search Difficulty dataset should be padded.
        Images are padded by making an array of zeros and randomly placing the image within it
        so that the entire image is still within the boundaries of (pad size x pad size).
        Default value is specified by searchnets.transforms.functional.VSD_PAD_SIZE.
        Argument has no effect if the dataset_type is not 'VOC'.
        Used to determine transforms to use at test time.
    loss_func : str
        type of loss function to use. One of {'CE', 'BCE'}. Default is 'CE',
        the standard cross-entropy loss. Used to determine transforms to use at test time.
    num_workers : int
        number of workers used by torch.DataLoaders. Default is 4.
    data_parallel : bool
        if True, use torch.nn.dataparallel to train network on multiple GPUs. Default is False.
    method : str
        training method. One of {'initialize', 'transfer'}.
        'initialize' means randomly initialize all weights and train the
        networks "from scratch".
        'transfer' means perform transfer learning, using weights pre-trained
        on imagenet.
        Default is 'transfer'.
    mode : str
        training mode. One of {'classify', 'detect'}.
        'classify' is standard image classification.
        'detect' trains to detect whether specified target is present or absent.
        Default is 'classify'.
    embedding_n_out : int
        for DetectNet, number of output features from input embedding.
        I.e., the output size of the linear layer that accepts the
        one hot vector querying whether a specific class is present as input.
        Default is 512.

    Returns
    -------
    None

    saves .npz output file with following keys:
        acc_per_set_size_per_model : numpy.ndarray
            where rows are models and columns are set size
        acc_per_set_size_model_dict : dict
            where keys are paths to model and values are accuracy by set size.
            The actual set sizes are in the .npz file saved by data, under
            the key 'set_size_vec_test'.
        predictions_per_model_dict : dict
            where keys are paths to model and values are array
            of predictions made by that model for test set
    """
    if mode == 'detect' and loss_func != 'BCE':
        print(
            f"when mode is 'detect', loss_func must be 'BCE', but was {loss_func}. Setting to 'BCE."
        )
        loss_func = 'BCE'

    if random_seed:
        np.random.seed(random_seed)  # for shuffling in batch_generator
        torch.manual_seed(random_seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False

    if torch.cuda.is_available():
        device = torch.device('cuda')
    else:
        device = torch.device('cpu')

    df_dataset = pd.read_csv(csv_file)

    for epochs in epochs_list:
        print(f'measuring accuracy on test set for {net_name} model trained for {epochs} epochs')

        # ------ initialize variables to hold outputs from all training replicates -------------------------------------
        if dataset_type == 'searchstims':
            # ---- for searchstims save just a results.gz file with predictions per model plus accuracies per model
            stim_types = df_dataset['stimulus'].unique()
            set_sizes_by_stim_type = {}
            for stim_type in stim_types:
                set_sizes_by_stim_type[stim_type] = df_dataset[df_dataset['stimulus'] == stim_type]['set_size'].unique()

            set_sizes = []
            for stim_type, set_sizes_this_stim in set_sizes_by_stim_type.items():
                if set_sizes == []:
                    set_sizes = set_sizes_this_stim
                else:
                    if not np.all(np.equal(set_sizes_this_stim, set_sizes)):
                        raise ValueError('set sizes are not the same across visual search stimuli')
                    else:
                        continue

            y_test = df_dataset[df_dataset['split'] == 'test']['target_condition']
            y_test = np.asarray(
                [1 if el == 'present' else 0 for el in y_test]
            )

            acc_per_set_size_per_model = []
            acc_per_set_size_model_dict = {}

        elif dataset_type == 'VSD':
            # ---- for VSD save results.gz **and** a .csv, because we have multiple metrics,
            # and because csv files are better anyway
            test_records = defaultdict(list)  # records gets turned into pandas DataFrame, then saved as .csv
            img_names_per_model_dict = {}

        predictions_per_model_dict = {}

        for net_number in range(1, number_nets_to_train + 1):
            transform, target_transform = get_transforms(dataset_type, loss_func, pad_size)

            if dataset_type == 'VSD':
                testset = VOCDetection(root=root,
                                       csv_file=csv_file,
                                       image_set='trainval',
                                       split='test',
                                       download=True,
                                       transform=transform,
                                       target_transform=target_transform)

            elif dataset_type == 'searchstims':
                testset = Searchstims(csv_file=csv_file,
                                      split='test',
                                      transform=transform,
                                      target_transform=target_transform)

            restore_path_this_net = make_save_path(restore_path, net_name, net_number, epochs)

            print(f'Loading model from {restore_path_this_net}')
            tester = Tester.from_config(net_name=net_name,
                                        num_classes=num_classes,
                                        loss_func=loss_func,
                                        testset=testset,
                                        mode=mode,
                                        embedding_n_out=embedding_n_out,
                                        restore_path=restore_path_this_net,
                                        batch_size=batch_size,
                                        device=device,
                                        num_workers=num_workers,
                                        data_parallel=data_parallel)

            test_results = tester.test()

            if dataset_type == 'searchstims':
                acc, y_pred = test_results['acc'], test_results['pred']
                set_size_vec_test = df_dataset[df_dataset['split'] == 'test']['set_size']
                if mode == 'detect':
                    y_pred = np.squeeze(y_pred)  # because DetectNet adds a dimension to output. My bad.
                acc_per_set_size = []
                for set_size in set_sizes:
                    # in line below, [0] at end because np.where returns a tuple
                    inds = np.where(set_size_vec_test == set_size)[0]
                    acc_this_set_size = (np.sum(y_pred[inds] == y_test[inds]) / y_test[inds].shape)
                    acc_per_set_size.append(acc_this_set_size)

                acc_set_size_str = ''
                acc_set_size_zip = zip(set_sizes, acc_per_set_size)
                for set_size, acc in acc_set_size_zip:
                    acc_set_size_str += f'set size {set_size}: {acc}. '
                print(acc_set_size_str)

                # append to list of lists which we convert into matrix
                acc_per_set_size_per_model.append(acc_per_set_size)
                # and insert into dictionary where model name is key
                # and list of accuracies per set size is the "value"
                acc_per_set_size_model_dict[restore_path_this_net] = acc_per_set_size

            elif dataset_type == 'VSD':
                test_records['net_name'].append(net_name)
                test_records['replicate'].append(net_number)
                test_records['method'].append(method)
                test_records['loss_func'].append(loss_func)
                test_records['restore_path'] = restore_path_this_net
                if mode == 'classify':
                    for metric in ['f1', 'acc_largest', 'acc_random']:
                        test_records[metric].append(test_results[metric])
                    results_str = ', '.join(
                        [f'{key}: {test_results[key]:7.3f}'
                         for key in ['loss', 'f1', 'acc_largest', 'acc_random']]
                    )
                elif mode == 'detect':
                    test_records['acc'].append(test_results['acc'])
                    results_str = f"acc: {test_results['acc']:7.3f}"
                print(f'test results: {results_str}')

                y_pred, img_names = test_results['pred'], test_results['img_names']
                img_names_per_model_dict[restore_path_this_net] = img_names

            # regardless of dataset type we store predictions in this dict
            predictions_per_model_dict[restore_path_this_net] = y_pred

        if dataset_type == 'searchstims':
            acc_per_set_size_per_model = np.asarray(acc_per_set_size_per_model)
            acc_per_set_size_per_model = np.squeeze(acc_per_set_size_per_model)

        # ---- create results dict, save to results.gz file
        if not os.path.isdir(test_results_save_path):
            os.makedirs(test_results_save_path)
        results_fname_stem = str(Path(configfile).stem)  # remove .ini extension

        results_fname = os.path.join(test_results_save_path,
                                     f'{results_fname_stem}_trained_{epochs}_epochs_test_results.gz')
        results_dict = dict(predictions_per_model_dict=predictions_per_model_dict)
        if dataset_type == 'searchstims':
            results_dict.update(dict(acc_per_set_size_per_model=acc_per_set_size_per_model,
                                     acc_per_set_size_model_dict=acc_per_set_size_model_dict,
                                     set_sizes=set_sizes)
                                )
        elif dataset_type == 'VSD':
            results_dict.update(dict(img_names_per_model_dict=img_names_per_model_dict))

        joblib.dump(results_dict, results_fname)

        # ---- finally for VSD dataset, create .csv
        if dataset_type == 'VSD':
            results_csv_fname = os.path.join(test_results_save_path,
                                             f'{results_fname_stem}_trained_{epochs}_epochs_test_results.csv')
            results_df = pd.DataFrame.from_records(test_records)
            results_df.to_csv(results_csv_fname, index=False)
