import numpy as np
import argparse
import os
import json

import sys
sys.path.append(os.path.join('..', '..'))
import bat_detect.train.train_utils as tu


def print_dataset_stats(data, split_name, classes_to_ignore):

    print('\nSplit:', split_name)
    print('Num files:', len(data))

    class_cnts = {}
    for dd in data:
        for aa in dd['annotation']:
            if aa['class'] not in classes_to_ignore:
                if aa['class'] in class_cnts:
                    class_cnts[aa['class']] += 1
                else:
                    class_cnts[aa['class']] = 1

    if len(class_cnts) == 0:
        class_names = []
    else:
        class_names = np.sort([*class_cnts]).tolist()
        print('Class count:')
        str_len = np.max([len(cc) for cc in class_names]) + 5

        for ii, cc in enumerate(class_names):
            print(str(ii).ljust(5) + cc.ljust(str_len) + str(class_cnts[cc]))

    return class_names


def load_file_names(file_name):

    if os.path.isfile(file_name):
        with open(file_name) as da:
            files = [line.rstrip() for line in da.readlines()]
        for ff in files:
            if ff.lower()[-3:] != 'wav':
                print('Error: Filenames need to end in .wav - ', ff)
                assert(False)
    else:
        print('Error: Input file not found - ', file_name)
        assert(False)

    return files


if __name__ == "__main__":

    info_str = '\nBatDetect - Prepare Data for Finetuning\n'

    print(info_str)
    parser = argparse.ArgumentParser()
    parser.add_argument('dataset_name', type=str, help='Name to call your dataset')
    parser.add_argument('audio_dir', type=str, help='Input directory for audio')
    parser.add_argument('ann_dir', type=str, help='Input directory for where the audio annotations are stored')
    parser.add_argument('op_dir', type=str, help='Path where the train and test splits will be stored')
    parser.add_argument('--percent_val', type=float, default=0.20,
                         help='Hold out this much data for validation. Should be number between 0 and 1')
    parser.add_argument('--rand_seed', type=int, default=2001,
                         help='Random seed used for creating the validation split')
    parser.add_argument('--train_file', type=str, default='',
                         help='Text file where each line is a wav file in train split')
    parser.add_argument('--test_file', type=str, default='',
                         help='Text file where each line is a wav file in test split')
    parser.add_argument('--input_class_names', type=str, default='',
                         help='Specify names of classes that you want to change. Separate with ";"')
    parser.add_argument('--output_class_names', type=str, default='',
                         help='New class names to use instead. One to one mapping with "--input_class_names". \
                         Separate with ";"')
    args = vars(parser.parse_args())


    np.random.seed(args['rand_seed'])

    classes_to_ignore  = ['', ' ', 'Unknown', 'Not Bat']
    generic_class      = ['Bat']
    events_of_interest = ['Echolocation']

    if args['input_class_names'] != '' and args['output_class_names'] != '':
        # change the names of the classes
        ip_names = args['input_class_names'].split(';')
        op_names = args['output_class_names'].split(';')
        name_dict = dict(zip(ip_names, op_names))
    else:
        name_dict = False

    # load annotations
    data_all, _, _ = tu.load_set_of_anns({'ann_path': args['ann_dir'], 'wav_path': args['audio_dir']},
                                          classes_to_ignore, events_of_interest, False, False,
                                          list_of_anns=True, filter_issues=True, name_replace=name_dict)

    print('Dataset name:         ' + args['dataset_name'])
    print('Audio directory:      ' + args['audio_dir'])
    print('Annotation directory: ' + args['ann_dir'])
    print('Ouput directory:      ' + args['op_dir'])
    print('Num annotated files:  ' + str(len(data_all)))

    if args['train_file'] != '' and args['test_file'] != '':
        # user has specifed the train / test split
        train_files = load_file_names(args['train_file'])
        test_files = load_file_names(args['test_file'])
        file_names_all = [dd['id'] for dd in data_all]
        train_inds = [file_names_all.index(ff) for ff in train_files if ff in file_names_all]
        test_inds = [file_names_all.index(ff) for ff in test_files if ff in file_names_all]

    else:
        # split the data into train and test at the file level
        num_exs = len(data_all)
        test_inds = np.random.choice(np.arange(num_exs), int(num_exs*args['percent_val']), replace=False)
        test_inds = np.sort(test_inds)
        train_inds = np.setdiff1d(np.arange(num_exs), test_inds)

    data_train = [data_all[ii] for ii in train_inds]
    data_test = [data_all[ii] for ii in test_inds]

    if not os.path.isdir(args['op_dir']):
        os.makedirs(args['op_dir'])
    op_name = os.path.join(args['op_dir'], args['dataset_name'])
    op_name_train = op_name + '_TRAIN.json'
    op_name_test  = op_name + '_TEST.json'

    class_un_train = print_dataset_stats(data_train, 'Train', classes_to_ignore)
    class_un_test = print_dataset_stats(data_test, 'Test', classes_to_ignore)

    if len(data_train) > 0 and len(data_test) > 0:
        if class_un_train != class_un_test:
            print('\nError: some classes are not in both the training and test sets.\
                   \nTry a different random seed "--rand_seed".')
            assert False

    print('\n')
    if len(data_train) == 0:
        print('No train annotations to save')
    else:
        print('Saving: ', op_name_train)
        with open(op_name_train, 'w') as da:
            json.dump(data_train, da, indent=2)

    if len(data_test) == 0:
        print('No test annotations to save')
    else:
        print('Saving: ', op_name_test)
        with open(op_name_test, 'w') as da:
            json.dump(data_test, da, indent=2)