import sys
from tqdm import tqdm
from typing import List, Tuple


def remove_large_sentences(src_path: str, tgt_path: str) -> Tuple[int, List[str], List[str]]:
    """
    Removes large sentences from a parallel dataset of source and target data.

    Args:
        src_path (str): path to the file containing the source language data.
        tgt_path (str): path to the file containing the target language data.

    Returns:
        Tuple[int, List[str], List[str]]: a tuple of
            - an integer representing the number of sentences removed
            - a list of strings containing the source language data after removing large sentences
            - a list of strings containing the target language data after removing large sentences
    """
    count = 0
    new_src_lines, new_tgt_lines = [], []

    src_num_lines = sum(1 for line in open(src_path, "r", encoding="utf-8"))
    tgt_num_lines = sum(1 for line in open(tgt_path, "r", encoding="utf-8"))
    assert src_num_lines == tgt_num_lines

    with open(src_path, encoding="utf-8") as f1, open(tgt_path, encoding="utf-8") as f2:
        for src_line, tgt_line in tqdm(zip(f1, f2), total=src_num_lines):
            src_tokens = src_line.strip().split(" ")
            tgt_tokens = tgt_line.strip().split(" ")

            if len(src_tokens) > 200 or len(tgt_tokens) > 200:
                count += 1
                continue

            new_src_lines.append(src_line)
            new_tgt_lines.append(tgt_line)

    return count, new_src_lines, new_tgt_lines


def create_txt(out_file: str, lines: List[str]):
    """
    Creates a text file and writes the given list of lines to file.

    Args:
        out_file (str): path to the output file to be created.
        lines (List[str]): a list of strings to be written to the output file.
    """
    add_newline = not "\n" in lines[0]
    outfile = open("{}".format(out_file), "w", encoding="utf-8")
    for line in lines:
        if add_newline:
            outfile.write(line + "\n")
        else:
            outfile.write(line)
    outfile.close()


if __name__ == "__main__":

    src_path = sys.argv[1]
    tgt_path = sys.argv[2]
    new_src_path = sys.argv[3]
    new_tgt_path = sys.argv[4]

    count, new_src_lines, new_tgt_lines = remove_large_sentences(src_path, tgt_path)
    print(f"{count} lines removed due to seq_len > 200")
    create_txt(new_src_path, new_src_lines)
    create_txt(new_tgt_path, new_tgt_lines)