"""
Construct tfrecord files of WLASL dataset
"""
import tensorflow as tf
import numpy as np

from random import shuffle

from modules import wlasl_utils, dataset_utils
from modules.config import CFG


def wlasl():
    # read annotation
    data_dict = wlasl_utils.read_annotation(CFG.WLASL_DATASET_PATHS['orig_annotations'])
    # shuffle annotation keys
    keys = list(data_dict.keys())
    shuffle(keys)
    # load class map
    class_dict = dataset_utils.load_json(CFG.WLASL_DATASET_PATHS['class_map'])
    # counter of samples
    dataset_classes = {'train': {100: 0, 300: 0, 1000: 0, 2000: 0},
                       'val': {100: 0, 300: 0, 1000: 0, 2000: 0},
                       'test': {100: 0, 300: 0, 1000: 0, 2000: 0}}
    # train, val and test subsets for WLASL100, 300, 1000, and 2000
    with tf.io.TFRecordWriter(CFG.WLASL_TFRECORDS_PATH[100]['train']) as writer_train_100, \
         tf.io.TFRecordWriter(CFG.WLASL_TFRECORDS_PATH[100]['val']) as writer_val_100, \
         tf.io.TFRecordWriter(CFG.WLASL_TFRECORDS_PATH[100]['test']) as writer_test_100, \
         tf.io.TFRecordWriter(CFG.WLASL_TFRECORDS_PATH[300]['train']) as writer_train_300, \
         tf.io.TFRecordWriter(CFG.WLASL_TFRECORDS_PATH[300]['val']) as writer_val_300, \
         tf.io.TFRecordWriter(CFG.WLASL_TFRECORDS_PATH[300]['test']) as writer_test_300, \
         tf.io.TFRecordWriter(CFG.WLASL_TFRECORDS_PATH[1000]['train']) as writer_train_1000, \
         tf.io.TFRecordWriter(CFG.WLASL_TFRECORDS_PATH[1000]['val']) as writer_val_1000, \
         tf.io.TFRecordWriter(CFG.WLASL_TFRECORDS_PATH[1000]['test']) as writer_test_1000, \
         tf.io.TFRecordWriter(CFG.WLASL_TFRECORDS_PATH[2000]['train']) as writer_train_2000, \
         tf.io.TFRecordWriter(CFG.WLASL_TFRECORDS_PATH[2000]['val']) as writer_val_2000, \
         tf.io.TFRecordWriter(CFG.WLASL_TFRECORDS_PATH[2000]['test']) as writer_test_2000:
        # tfrecord writers
        writers = {
            'train': ((100, writer_train_100), (300, writer_train_300), (1000, writer_train_1000), (2000, writer_train_2000)),
            'val': ((100, writer_val_100), (300, writer_val_300), (1000, writer_val_1000), (2000, writer_val_2000)),
            'test': ((100, writer_test_100), (300, writer_test_300), (1000, writer_test_1000), (2000, writer_test_2000))
            }
        # for every annotation entry
        for key in keys:
            annot = data_dict[key]
            # read annotation data
            gloss, _, split, _, _, fused_keypoints = wlasl_utils.data_from_annotation(
                annot, CFG.WLASL_DATASET_PATHS['videos'], CFG.WLASL_DATASET_PATHS['keypoints_annotations'], clip_frames=0)
            # if annotation does not contain keypoint annotations, skip
            if fused_keypoints is None:
                continue
            # transform keypoints to input space
            hand_keypoints, face_keypoints = dataset_utils.keypoints_to_input_space(*fused_keypoints)
            # convert everything to lists
            hand_keypoints_list = np.array(hand_keypoints).reshape(-1).tolist()
            face_keypoints_list = np.array(face_keypoints).reshape(-1).tolist()
            cls = class_dict[gloss]
            # construct example
            example = tf.train.Example(features=tf.train.Features(feature={
                    'hand_keypoints':
                        tf.train.Feature(float_list=tf.train.FloatList(value=hand_keypoints_list)),
                    'face_keypoints':
                        tf.train.Feature(float_list=tf.train.FloatList(value=face_keypoints_list)),
                    'class':
                        tf.train.Feature(int64_list=tf.train.Int64List(value=[int(cls)])),
                }))
            # write example to associated tfrecords
            # the first 100 words are in all datasets (WLASL100, 300, 1000, and 2000),
            # 300 only in WLASL300, 1000, and 2000 and so on..
            for n_classes, writer in writers[split]:
                if cls < n_classes:
                    writer.write(example.SerializeToString())
                    # update counters
                    dataset_classes[split][n_classes] += 1
    print(dataset_classes)


if __name__ == '__main__':
    wlasl()
