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.voc_assayer import VOCAssayer
from .datasets import VOCDetection
from .utils.general import make_save_path


def assay(csv_file,
          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):
    """assay behavior of models trained with Pascal VOC Detection set

    Parameters
    ----------
    csv_file : str
        name of .csv file containing prepared data sets.
        Generated by searchnets.data.split function.
    net_name : str
        name of convolutional neural net architecture to train.
        One of {'alexnet', 'VGG16', 'CORnet_Z', 'CORnet_S'}
    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:
        arrays_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')

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

        # ------ initialize variables to hold outputs from all training replicates ------------------------------------
        # ---- for VSD save results.gz **and** a .csv, because we have multiple metrics,
        # and because csv files are better anyway
        assay_records = defaultdict(list)  # records gets turned into pandas DataFrame, then saved as .csv
        # these will be lists of DataFrames, one for each training replicate
        df_lists = defaultdict(list)
        arrays_per_model = {}  # inputs/outputs of model, where key is restore path, and value is dict of arrays

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

            testset = VOCDetection(root=root,
                                   csv_file=csv_file,
                                   image_set='trainval',
                                   split='test',
                                   download=True,
                                   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}')
            assayer = VOCAssayer.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)

            results = assayer.assay()

            # --- add columns to image + trial dataframes before appending to list
            for key in ('images_df', 'trials_df'):
                df = results[key]
                df['net_name'] = net_name
                df['replicate'] = net_number
                df['mode'] = mode
                df['method'] = method
                df['loss_func'] = loss_func
                df['restore_path'] = restore_path_this_net
                df_lists[key].append(df)

            # ---- add columns + metrics to our 'results across replicates' records for that data frame
            assay_records['net_name'].append(net_name)
            assay_records['replicate'].append(net_number)
            assay_records['mode'].append(mode)
            assay_records['method'].append(method)
            assay_records['loss_func'].append(loss_func)
            assay_records['restore_path'] = restore_path_this_net

            for metric in ['acc', 'd_prime']:
                assay_records[metric].append(results[metric])
            results_str = ', '.join(
                [f'{key}: {results[key]:7.3f}'
                 for key in ['acc', 'd_prime']]
            )
            print(f'assay results: {results_str}')

            arrays_per_model[restore_path_this_net] = results['arrays']

        # ---- 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

        arrays_fname = os.path.join(test_results_save_path,
                                    f'{results_fname_stem}_trained_{epochs}_epochs_assay_arrays.gz')
        joblib.dump(arrays_per_model, arrays_fname)

        summary_csv_fname = os.path.join(test_results_save_path,
                                         f'{results_fname_stem}_trained_{epochs}_epochs_assay_results.csv')
        results_df = pd.DataFrame.from_records(assay_records)
        results_df.to_csv(summary_csv_fname, index=False)

        for key, df_list in df_lists.items():
            csv_fname = f"{results_fname_stem}_trained_{epochs}_epochs_assay_{key.replace('_df', '')}.csv"
            csv_path = os.path.join(test_results_save_path, csv_fname)
            df = pd.concat(df_list)
            df.to_csv(csv_path, index=False)
