from math import ceil
from pathlib import Path
import random

import numpy as np
import pandas as pd


def _samples_per_set_size(split, n_samples, n_set_sizes):
    """evenly divide samples for dataset split across visual search set sizes"""
    # below, divide by two for target present / absent
    samples_per_set_size = (n_samples / n_set_sizes) / 2
    if samples_per_set_size.is_integer():
        samples_per_set_size = int(samples_per_set_size)
        samples_per_set_size = [samples_per_set_size for _ in range(n_set_sizes)]
        return samples_per_set_size
    else:
        raise TypeError(f'{split}_size_per_set_size, {samples_per_set_size}, is is not a whole number.\n'
                        f'It is calculated as: ({split}_size_per_stim_type / len(set_sizes)) / 2\n'
                        '(2 is for target present or absent).\n'
                        'Adjust total number of samples, or number of set sizes.')


def split(csv_file_in,
          train_size,
          dataset_type='searchstims',
          csv_file_out=None,
          stim_types=None,
          val_size=None,
          test_size=None,
          train_size_per_set_size=None,
          val_size_per_set_size=None,
          test_size_per_set_size=None,
          shard_train=False,
          shard_size=None):
    """split a dataset into training, validation, and test sets.

    Takes a .csv file representing the dataset, and adds a 'split' field
    with string values that specify which split each sample belongs to, one of {'train', 'val', 'test'}.

    Parameters
    ----------
    csv_file_in : str
        path to .csv file generated by searchstims package. Typically in root of train_dir.
    train_size : int
        number of samples to put in training data set.
        Note that if there is more than on stimulus type in the source data set, then the number of stimuli from
        each type will be train_size / number of types.
    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.
    csv_file_out : str
        name of .csv file that this function will save. Default is None, in which case the suffix '_split' is added
        to csv_file_in.
    stim_types : list
        of strings; specifies which visual search stimulus types to use when creating dataset. Strings must be keys in
        .json file within train_dir. Default is None, in which case all types found in .csv file will be used.
    val_size : int
        number of samples to put in validation data set.
        Default is None, in which case there will be no validation set.
        Note that if there is more than on stimulus type in the source data set, then the number of stimuli from
        each type will be val_size / number of types.
    test_size : int
        number of samples to put in test data set.
        Default is None, in which case all samples not used in training and validation sets are used for test set.
        Note that if there is more than on stimulus type in the source data set, then the number of stimuli from
        each type will be test_size / number of types.
    train_size_per_set_size, val_size_per_set_size, test_size_per_set_size : list
        number of samples in split per visual search set size. Default is None, in which case the total number of
        samples for each stimulus type will be divided by the number of set sizes, so that an equal number is used
        for each set size.
    shard_train : bool
        if True, split the training set into shards. This will add an additional column to the .csv generated, 'shard'
        of integer values representing the shard that the sample in each row of the .csv belongs to.
    shard_size : int
        number of samples per shard

    Returns
    -------
    None
    """
    if csv_file_out is None:
        # if no name declared for csv_file_out, then add '_split' to csv_file_in and use that for name
        csv_file_in_stem = Path(csv_file_in).stem
        csv_file_out = csv_file_in_stem + '_split.csv'
        csv_file_out = Path(csv_file_in).parent.joinpath(csv_file_out)

    df = pd.read_csv(csv_file_in)

    if dataset_type == 'VSD':
        if any([not type(size) == float for size in [train_size, val_size, test_size]]):
            raise ValueError(
                'when using Visual Difficulty Score dataset, sizes of splits (train_size, val_size, and test_size) '
                'must be specified as proportions of dataset (e.g. 0.5, 0.25, 0.25).'
            )
        df.columns = ['img', 'difficulty_score']
        inds = df.index.values

        n_samples = inds.shape[0]
        train_size = np.round(n_samples * train_size).astype(int)
        val_size = np.round(n_samples * val_size).astype(int)
        test_size = np.round(n_samples * test_size).astype(int)

        # shuffle then grab indices
        inds = np.random.permutation(inds)
        train_inds = inds[0: train_size]
        val_inds = inds[train_size: train_size + val_size]
        test_inds = inds[train_size + val_size: train_size + val_size + test_size]
        leftover_inds = inds[train_size + val_size + test_size:]
        split_col = np.asarray([''] * inds.shape[0], dtype='object')
        split_col[train_inds] = 'train'
        split_col[val_inds] = 'val'
        split_col[test_inds] = 'test'
        if leftover_inds:
            split_col[leftover_inds] = 'None'
        df['split'] = split_col

    elif dataset_type == 'searchstims':
        if any([not type(size) == int for size in [train_size, val_size, test_size]]):
            raise ValueError(
                'when using searchstims dataset, sizes of splits (train_size, val_size, and test_size) '
                'must be specified as number of samples from dataset (e.g. 6400, 1200, 800).'
            )

        df_stim_types = df['stimulus'].unique().tolist()
        if stim_types:
            if type(stim_types) != list or not all([type(stim_type) == str for stim_type in stim_types]):
                raise TypeError('stim_types must be a list of strings')
            # check that all stim_types are actually in the DataFrame

            if not all([stim_type in df_stim_types
                        for stim_type in stim_types]):
                not_in_fnames_dict = [stim_type for stim_type in stim_types if stim_type not in df_stim_types]
                raise ValueError(
                    f'the following stimulus types were not found in {csv_file_in}: {not_in_fnames_dict}'
                )

            # if all stim_types are valid, filter by them
            df = df[df['stimulus'].isin(stim_types)]
            # reset index to zero, don't keep the original indices
            # this is necessary if we don't use all stimuli, because old
            # indices if used for indexing might not correspond to 
            # new rows of the shorter DataFrame
            df = df.reset_index(drop=True)
        else:
            stim_types = df_stim_types
        num_stim_types = len(stim_types)

        # split is the column we will add to the dataframe. The value in each row will be the split that row belongs to,
        # i.e., it will be one of {'train', 'val', 'test'}
        split_col = np.asarray(['' for _ in range(len(df))], dtype='object')
        if shard_train:
            # shard is the column we will add to the dataframe that indicates which shard each sample belongs to. I.e.,
            # the value in each row will be an integer 0, 1, ..., n-1 where n is the number of shards
            shard_col = np.zeros(shape=(len(df),), dtype=int)

        # ----- (1) divide the split size by the number of visual search stimulus types --------------------------------
        train_size_per_stim_type = train_size / num_stim_types
        if train_size_per_stim_type.is_integer():
            train_size_per_stim_type = int(train_size_per_stim_type)
        else:
            raise TypeError(f'train_size_per_stim_type, {train_size_per_stim_type}, is is not a whole number.\n'
                            'It is calculated as: (train_size / number of visual search stimulus types))\n'
                            'Adjust total number of samples, or number of stimulus types.')

        if train_size_per_set_size:
            total_train_size_from_per_set_size = sum(train_size_per_set_size)
            if total_train_size_from_per_set_size != train_size_per_stim_type:
                raise ValueError(
                    f'total number of training samples specified in '
                    f'train_size_per_set_size, {total_train_size_from_per_set_size} does not equal number determined '
                    f'by dividing train_size up by number of stim_types: {train_size_per_stim_type}'
                )

        if val_size:
            val_size_per_stim_type = val_size / num_stim_types
            if val_size_per_stim_type.is_integer():
                val_size_per_stim_type = int(val_size_per_stim_type)
            else:
                raise TypeError('val_size_per_set_size is not a whole number, adjust '
                                'total number of samples, or number of set sizes.')
        else:
            val_size_per_stim_type = 0

        if test_size:
            test_size_per_stim_type = test_size / num_stim_types
            if test_size_per_stim_type.is_integer():
                test_size_per_stim_type=int(test_size_per_stim_type)
            else:
                raise TypeError('test_size_per_set_size is not a whole number, adjust '
                                'total number of samples, or number of set sizes.')
        else:
            # "-1" means "use the remaining samples for the test set"
            test_size_per_stim_type = -1

        # ----- (2) determine the number of samples in each split *for each set visual search set size* ----------------
        for stim_type in stim_types:
            # and this will be set sizes declared by user for this stimulus (could be diff't for each stimulus type).
            # First have to convert set size from char to int
            set_sizes = df[df['stimulus'] == stim_type]['set_size'].unique()
            n_set_sizes = len(set_sizes)

            # if user did not specify train_size_per_set_size, divide evenly among set sizes
            if train_size_per_set_size is None:
                train_size_per_set_size_this_stim = _samples_per_set_size('train',
                                                                          train_size_per_stim_type,
                                                                          n_set_sizes)
            else:
                # if train_size_per_set_size is not None, divide each element in two (for target present / absent)
                train_size_per_set_size_this_stim = [item // 2 if item % 2 == 0 else item / 2
                                                     for item in train_size_per_set_size]

            if val_size:
                # if user did not specify val_size_per_set_size, divide evenly among set sizes
                if val_size_per_set_size is None:
                    val_size_per_set_size_this_stim = _samples_per_set_size('val',
                                                                            val_size_per_stim_type,
                                                                            n_set_sizes)
                else:
                    # if val_size_per_set_size is not None, divide each element in two (for target present / absent)
                    val_size_per_set_size_this_stim = [item // 2 if item % 2 == 0 else item / 2
                                                       for item in val_size_per_set_size]
            else:
                val_size_per_set_size_this_stim = [0 for _ in set_sizes]

            if test_size is None or test_size == -1:
                # "-1" means "use the remaining samples for the test set"
                test_size_per_set_size_this_stim = [-1 for _ in set_sizes]
            elif test_size > 0:
                if test_size_per_set_size is None:
                    test_size_per_set_size_this_stim = _samples_per_set_size('test',
                                                                             test_size_per_stim_type,
                                                                             n_set_sizes)
                else:
                    # if test_size_per_set_size is not None, divide each element in two (for target present / absent)
                    test_size_per_set_size_this_stim = [item // 2 if item % 2 == 0 else item / 2
                                                        for item in test_size_per_set_size]
            else:
                raise ValueError(f'invalid test size: {test_size}')

            # do some extra juggling to make sure we have equal number of target present
            # and target absent stimuli for each "set size", in training and test datasets
            for (set_size,
                 train_size,
                 val_size,
                 test_size) in zip(set_sizes,
                                   train_size_per_set_size_this_stim,
                                   val_size_per_set_size_this_stim,
                                   test_size_per_set_size_this_stim):
                df_stim_set = df[
                    (df['stimulus'] == stim_type) &
                    (df['set_size'] == set_size)
                ]
                target_conditions = df_stim_set['target_condition'].unique()
                if len(target_conditions) != 2:
                    raise ValueError(f'incorrect number of target conditions: {target_conditions}')

                for target_condition in target_conditions:
                    df_target = df_stim_set[df_stim_set['target_condition'] == target_condition]
                    total_size = sum([size
                                      for size in (train_size, val_size, test_size)
                                      if size is not 0 and size is not -1])
                    if total_size > len(df_target):
                        raise ValueError(
                            f'number of samples for training + validation set, {total_size}, '
                            f'is larger than number of samples in data set, {len(df_target)}'
                        )
                    inds = df_target.index.values.tolist()
                    random.shuffle(inds)
                    train_inds = np.asarray(
                        [inds.pop() for _ in range(train_size)]
                    )

                    if val_size > 0:
                        val_inds = np.asarray(
                            [inds.pop() for _ in range(val_size)]
                        )
                    else:
                        val_inds = None

                    if test_size > 0:
                        test_inds = np.asarray([inds.pop() for _ in range(test_size)])
                    elif test_size == -1:
                        # "-1" means "use the remaining samples for the test set"
                        test_inds = np.asarray([inds.pop() for _ in range(len(inds))])

                    split_zip = zip(
                        ['train', 'val', 'test'],
                        [train_inds, val_inds, test_inds]
                    )
                    for split_name, split_inds in split_zip:
                        if split_inds is not None:
                            split_col[split_inds] = split_name
                    if inds:  # is not empty
                        leftover_inds = np.asarray(inds)
                        split_col[leftover_inds] = 'None'

        df['split'] = split_col

        if shard_train:
            inds = np.nonzero(split_col == 'train')[0]
            np.random.shuffle(inds)
            num_shards = int(ceil(len(inds) / shard_size))

            for shard_num in range(num_shards):
                start = shard_size * shard_num
                stop = shard_size * (shard_num + 1)
                inds_this_shard = inds[start:stop]
                shard_col[inds_this_shard] = shard_num

            df['shard'] = shard_col

    # make parent directory for csv_file_out, if needed
    csv_parent = Path(csv_file_out).parent
    if not csv_parent.is_dir():
        csv_parent.mkdir(parents=True)

    df.to_csv(csv_file_out)
