import json
import os
import random

DATA_DIR = f"/home/{os.environ['USER']}/data/wit/all_jsons"
SEED = 0
PROPORTION_TRAIN = 0.98
PROPORTION_VALID = 0.01

random.seed(SEED)

all_files = [
    f"{DATA_DIR}/{file_}" for file_ in os.listdir(DATA_DIR) if ("all" not in file_)
]

print(all_files)

examples = []
for file_ in all_files:
    print(file_)
    with open(file_) as f:
        file_examples = [
            json.dumps(json.loads(line), ensure_ascii=False) for line in f.readlines()
        ]
    print(len(file_examples))
    examples.extend(file_examples)

print(f"Before dedup: {len(examples)}")
examples = list(set(examples))
print(f"After dedup: {len(examples)}")

print(examples[0])
# Shuffle examples
random.shuffle(examples)
print(examples[0])

split_dataset = {}
split_dataset["train"] = examples[: int(len(examples) * PROPORTION_TRAIN)]
split_dataset["valid"] = examples[
    int(len(examples) * PROPORTION_TRAIN) : int(
        len(examples) * (PROPORTION_TRAIN + PROPORTION_VALID)
    )
]
split_dataset["test"] = examples[
    int(len(examples) * (PROPORTION_TRAIN + PROPORTION_VALID)) :
]


for split in ["train", "valid", "test"]:
    print("-----")
    print(len(split_dataset[split]))
    print("-----")
    with open(
        f"/home/{os.environ['USER']}/data/wit/all_jsons/{split}_dataset_all_98_1_1_split.json",
        "w",
    ) as f:
        f.write("\n".join(split_dataset[split]))