# -*- coding: utf-8 -*-
"""dgs_document_split.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/19pHmLuIEAKFn4BqVr7cwNRaVxQHWKNI7
"""

# ! pip install sign-language-datasets==0.0.12

import json

import numpy as np

import tensorflow_datasets as tfds
import sign_language_datasets.datasets
from sign_language_datasets.datasets.config import SignDatasetConfig
from sign_language_datasets.datasets.dgs_corpus.dgs_utils import get_elan_sentences

from typing import Optional, Tuple

np.random.seed(1)

# Videos 1177918 and 1432043 have 25 fps, start and end frame won't match

INCORRECT_FRAMERATE = ["1432043", "1177918"]


def get_split_indexes(total_size: int, dev_size: int, test_size) -> Tuple[np.array, np.array, np.array]:
    train_indexes = np.arange(total_size, dtype=np.int32)

    np.random.shuffle(train_indexes)

    # high inclusive

    dev_indexes = np.random.choice(train_indexes, size=(dev_size,), replace=False)

    remaining_train_indexes = np.asarray([i for i in train_indexes if i not in dev_indexes])

    test_indexes = np.random.choice(remaining_train_indexes, size=(test_size,), replace=False)

    remaining_train_indexes = np.asarray([i for i in remaining_train_indexes if i not in test_indexes])

    return remaining_train_indexes, dev_indexes, test_indexes


config = SignDatasetConfig(name="only-annotations", version="1.0.0", include_video=False, include_pose=None)
dgs_corpus = tfds.load("dgs_corpus", builder_kwargs=dict(config=config))


def get_split(dev_size: int, test_size: int):
    ids = np.array([datum["id"].numpy().decode("utf-8") for datum in dgs_corpus["train"] if datum["id"] not in INCORRECT_FRAMERATE])

    train_indexes, dev_indexes, test_indexes = get_split_indexes(len(ids), dev_size=dev_size, test_size=test_size)

    print("Number of entire files in each split:")
    print(str({"train": len(train_indexes), "dev": len(dev_indexes), "test": len(test_indexes)}))

    return {
        "dgs_corpus_version": "3.0.0",
        "train": list(ids[train_indexes]),
        "dev": list(ids[dev_indexes]),
        "test": list(ids[test_indexes]),
    }


split = get_split(dev_size=10, test_size=10)

with open("split.json", "w") as outfile:
    json.dump(split, outfile, indent=4)

# ! cat split.json

"""## compute sentence statistics for this split"""


def get_split_name_from_id(_id: str) -> str:
    for key in split.keys():
        if _id in split[key]:
            return key

    return "none"


sentences_found = {"train": 0, "dev": 0, "test": 0, "none": 0}

for datum in dgs_corpus["train"]:

    _id = datum["id"].numpy().decode("utf-8")

    split_name = get_split_name_from_id(_id)

    elan_path = datum["paths"]["eaf"].numpy().decode("utf-8")
    sentences = get_elan_sentences(elan_path)

    for sentence in sentences:
        gloss_sequence = " ".join([s["gloss"] for s in sentence["glosses"]])
        german_sentence = sentence["german"]

        if gloss_sequence != "" and german_sentence != "":
            sentences_found[split_name] += 1

print(sentences_found)
