from constants import DATASET_DIR, TRAINING_DIR, VALIDATION_DIR, TESTING_DIR, ORIGINAL_DATASET_FILES
from clean_data import clean_dataset
from helper import csv_to_parquet
from split import split_dataset, split_dataset_random
import nltk
import os
import shutil
import pandas as pd

def setup() -> None:
    # make sure nltk can be used later.
    nltk.download("all")

    for dataset_file in ORIGINAL_DATASET_FILES:
        if not os.path.exists(f"{DATASET_DIR}/{dataset_file}"):
            raise Exception(f"Please add {dataset_file} to {DATASET_DIR}")

        name = os.path.splitext(dataset_file)[0]
        if not os.path.exists(f"{DATASET_DIR}/{name}.parquet"):
            csv_to_parquet(f"{DATASET_DIR}/{dataset_file}", f"{DATASET_DIR}/{name}.parquet")
            print(f"finished converting {dataset_file} to parquet")
            clean_dataset(f"{name}.parquet")
            print(f"cleaned {name}.parquet")
            split_dataset_random(f"{name}.parquet")
            print(f"split {name}.parquet into traning, validation and test")

        # LIAR
        for dataset, destination in [("train.tsv", TRAINING_DIR), ("valid.tsv", VALIDATION_DIR), ("test.tsv", TESTING_DIR)]:
            if os.path.exists(f"{DATASET_DIR}/{dataset}"):
                df = pd.read_csv(f"{DATASET_DIR}/{dataset}", sep='\t', header=None)
                df = df.rename(columns={
                    1: "type",
                    2: "content"
                })
                name = os.path.splitext(dataset)[0]
                df.to_parquet(f"{DATASET_DIR}/{name}.parquet")
                clean_dataset(f"{name}.parquet")

                shutil.move(f"{DATASET_DIR}/{name}.parquet", f"{destination}/LIAR.parquet")

if __name__ == "__main__":
    setup()
