import argparse
import os
import glob
import pandas as pd
import h5py
import numpy as np
from cmu_mosei_std_folds import cmu_mosei_std_folds


def _parse():
    parser = argparse.ArgumentParser(description='create label file')
    parser.add_argument('label', type=str, help='CMU_MOSEI_Labels.csd')
    parser.add_argument('interval', type=str, help='cmumosei_interval.txt')
    parser.add_argument('outd', type=str, help='output label dir')
    parser.add_argument('--label-format', choices=['emo_class', 'sentiment_regress', 'sentiment_class'],
                        default='emo_class',
                        help='label format')
    args = parser.parse_args()
    return args


def get_interval_dict(mosei_interval):
    interval_dict = {}
    for l in open(mosei_interval):
        fid, segid, st, en = l.strip().split()
        if fid not in interval_dict:
            interval_dict[fid] = {}
        interval_dict[fid][(st, en)] = segid
    return interval_dict


def _main():
    args = _parse()
    os.makedirs(args.outd, exist_ok=True)

    interval_dict = get_interval_dict(args.interval)

    train_fold, valid_fold, test_fold = cmu_mosei_std_folds()
    label_dic = {'train':[], 'valid':[], 'test':[]}
    with h5py.File(args.label, 'r') as h:
        for fid in h['All Labels']['data']:
            if fid not in interval_dict:
                continue

            if fid in train_fold:
                dset = 'train'
            elif fid in valid_fold:
                dset = 'valid'
            elif fid in test_fold:
                dset = 'test'
            else:
                continue

            data = h['All Labels']['data'][fid]
            for f, i in zip(data['features'], data['intervals']):
                st, en = str(i[0]), str(i[1])
                segid = interval_dict[fid][(st, en)]
                clip_id = fid+'_'+segid
                if args.label_format == 'emo_class':
                    label = [str(v) for v in (f[1:] > 0).astype(np.int64)]
                elif args.label_format == 'sentiment_regress':
                    label = [str(f[0])]
                elif args.label_format == 'sentiment_class':
                    raise NotImplementedError
                label_dic[dset].append(' '.join([clip_id] + label))

    for fold, lines in label_dic.items():
        outf = os.path.join(args.outd, fold + '.txt')
        with open(outf, 'w') as fp:
            fp.write('\n'.join(lines)+'\n')


if __name__ == '__main__':
    _main()
