import re
from pathlib import Path

DATA_DIR = "data/imdb-rating"
FILE_PRE = DATA_DIR.rsplit("/", maxsplit=1)[1].replace("-", "_")


def clean_text(text):
    # remove whitespace characters
    text = re.sub(r"\s+", " ", text, flags=re.MULTILINE | re.DOTALL)

    return text.strip()


def main():
    # load raw data
    with open(Path(DATA_DIR) / "texts.txt", encoding="utf-8") as f:
        texts = f.readlines()
    with open(Path(DATA_DIR) / "score.txt", encoding="utf-8") as f:
        labels = f.readlines()

    # clean raw data
    texts = list(map(clean_text, texts))
    labels = list(map(lambda x: x.strip(), labels))

    assert len(texts) == len(labels)

    # generate full-size dataset
    output = ["".join([y, "\t", x, "\n"]) for x, y in zip(texts, labels)]
    output[-1] = output[-1].strip()
    with open(Path(DATA_DIR) / f"{FILE_PRE}_raw_texts.txt", "w", encoding="utf-8") as f:
        f.writelines(output)


if __name__ == "__main__":
    """Instructions for preparing the imdb-rating dataset.

    1. Download the original data at https://zenodo.org/record/5257310.
        - texts: "texts.txt"
        - labels: "score.txt"
    2. Run this script with appropriate file path defined in DATA_DIR.
    3. Compress the output files using "bzip2 -fkv *raw_texts*.txt".
    """
    main()