import os
import codecs
import random
import numpy as np
from tqdm import tqdm
from collections import Counter
from nltk.tokenize import word_tokenize
from util.data_util import load_json, load_lines, load_pickle, save_pickle, time_to_index

PAD, UNK = "<PAD>", "<UNK>"

class CharadesProcessor:
    def __init__(self):
        super(CharadesProcessor, self).__init__()
        self.idx_counter = 0

    def reset_idx_counter(self):
        self.idx_counter = 0

    def process_data(self, data, charades, scope):
        results = []
        for line in tqdm(data, total=len(data), desc='process charades-sta {}'.format(scope)):
            line = line.lstrip().rstrip()
            if len(line) == 0:
                continue
            video_info, sentence = line.split('##')
            vid, start_time, end_time = video_info.split(' ')
            duration = float(charades[vid]['duration'])
            start_time = max(0.0, float(start_time))
            end_time = min(float(end_time), duration)
            words = word_tokenize(sentence.strip().lower(), language="english")
            record = {'sample_id': self.idx_counter, 'vid': str(vid), 's_time': start_time, 'e_time': end_time,
                      'duration': duration, 'words': words}
            results.append(record)
            self.idx_counter += 1
        return results

    def convert(self, data_dir):
        self.reset_idx_counter()
        if not os.path.exists(data_dir):
            raise ValueError('data dir {} does not exist'.format(data_dir))
        # load raw data
        charades = load_json(os.path.join(data_dir, 'charades.json'))
        train_data = load_lines(os.path.join(data_dir, 'charades_sta_train.txt'))
        # train_data = load_lines(os.path.join(data_dir, 'charades_sta_train_pseudo.txt'))
        test_data = load_lines(os.path.join(data_dir, 'charades_sta_test.txt'))
        # process data
        train_set = self.process_data(train_data, charades, scope='train')
        test_set = self.process_data(test_data, charades, scope='test')
        return train_set, None, test_set  # train/val/test

def gen_or_load_dataset(configs):
    if not os.path.exists(configs.save_dir):
        os.makedirs(configs.save_dir)
    data_dir = os.path.join('data', 'dataset', configs.task)
    feature_dir = os.path.join('data', 'features', configs.task, configs.fv)
    if configs.suffix is None:
        save_path = os.path.join(configs.save_dir, '_'.join([configs.task, 
            configs.fv, str(configs.max_pos_len)]) + '.pkl')
    else:
        save_path = os.path.join(configs.save_dir, '_'.join([configs.task, 
            configs.fv, str(configs.max_pos_len), 
            configs.suffix]) + '.pkl')
    if os.path.exists(save_path):
        dataset = load_pickle(save_path)
        return dataset
    feat_len_path = os.path.join(feature_dir, 'feature_shapes.json')    
    emb_path = os.path.join('data', 'features', 'glove.840B.300d.txt') 
    # load video feature length
    vfeat_lens = load_json(feat_len_path)
    for vid, vfeat_len in vfeat_lens.items():
        vfeat_lens[vid] = min(configs.max_pos_len, vfeat_len)
    # load data
    if configs.task == 'charades':
        processor = CharadesProcessor()
    elif configs.task == 'activitynet':
        processor = ActivityNetProcessor()
    elif configs.task == 'tacos':
        processor = TACoSProcessor()
    elif configs.task == 'charades_cd':
        processor = CharadesCDProcessor()
    elif configs.task == 'charades_ov':
        processor = CharadesOVProcessor()
    elif configs.task == 'activitynet_cd':
        processor = ActivityNetCDProcessor()
    elif configs.task == 'activitynet_ov':
        processor = ActivityNetOVProcessor()
    elif configs.task == 'charades_anet':
        processor = CharadesAnetProcessor()
    else:
        raise ValueError('Unknown task {}!!!'.format(configs.task))
    train_data, val_data, test_data = processor.convert(data_dir)
    # generate dataset
    data_list = [train_data, test_data] if val_data is None else [train_data, val_data, test_data]
    # data_list = [train_data, val_data]
    word_dict, char_dict, vectors = vocab_emb_gen(data_list, emb_path)
    train_set = dataset_gen(train_data, vfeat_lens, word_dict, char_dict, configs.max_pos_len, 'train')
    val_set = None if val_data is None else dataset_gen(val_data, vfeat_lens, word_dict, char_dict,
                                                        configs.max_pos_len, 'val')
    test_set = dataset_gen(test_data, vfeat_lens, word_dict, char_dict, configs.max_pos_len, 'test')
    # save dataset
    n_val = 0 if val_set is None else len(val_set)
    dataset = {'train_set': train_set, 'val_set': val_set, 'test_set': test_set, 'word_dict': word_dict,
               'char_dict': char_dict, 'word_vector': vectors, 'n_train': len(train_set), 'n_val': n_val,
               'n_test': len(test_set), 'n_words': len(word_dict), 'n_chars': len(char_dict)}
    save_pickle(dataset, save_path)
    return dataset