"""
Visualization tool - search sign videos based on sign names
"""
import os
import cv2

from collections import defaultdict

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


def make_vocabulary():
    """Construct vocabulary object, mapping dataset samples to sign names and save it to file
    """
    # read annotation and prepare vocabulary map
    annotation = wlasl_utils.read_annotation(CFG.WLASL_DATASET_PATHS['orig_annotations'], shuffle_glosses=False)
    vocab_dict = defaultdict(lambda: list())
    # parse annotation and save the entries to the vocabulary map
    for annot in annotation.values():
        vocab_dict[annot['gloss']].append(annot['video_file_name'])
    # save to file
    dump_json(CFG.WLASL_DATASET_PATHS['vocabulary'], vocab_dict)


def main():
    # if vocabulary file does not exist, construct it
    if not os.path.isfile(CFG.WLASL_DATASET_PATHS['vocabulary']):
        make_vocabulary()
    # load vocabulary
    vocab = load_json(CFG.WLASL_DATASET_PATHS['vocabulary'])
    print('Pass empty string to exit...')
    while True:
        # wait for input sign name
        gloss = input('Enter sign to display: ')
        # terminate on empty string
        if gloss == '':
            break
        try:
            # get files of received sign name
            files = vocab[gloss]
            # for each video file
            for file in files:
                path = os.path.join(CFG.WLASL_DATASET_PATHS['videos'], file)
                # repeat videos until keypress
                repeat = True
                while repeat:
                    # read video file
                    capture = cv2.VideoCapture(path)
                    while capture.isOpened():
                        # read frame and viz it
                        ret, frame = capture.read()
                        if not ret:
                            break
                        cv2.imshow('vocabulary', frame)
                        key = cv2.waitKey(40)
                        # escape - exit
                        if key == 27:
                            capture.release()
                            raise KeyboardInterrupt
                        # Q - next video
                        elif key == 113:
                            repeat = False
                            break
                    capture.release()
        except KeyError:
            print('Sign not found in dataset')
        except KeyboardInterrupt:
            pass
    cv2.destroyAllWindows()


if __name__ == '__main__':
    main()
