#!/usr/bin/env python3

import glob
import os

import random
from tqdm import tqdm

from datasets import Dataset, DatasetDict, load_dataset


def convert(list_of_dicts):
    res = {}
    for d in list_of_dicts:
        for k, v in d.items():
            res.setdefault(k, []).append(v)
    return res


v10 = load_dataset("fever", "v1.0")
name_lst = ["train", "labelled_dev"]

old_to_new_label_map = {"SUPPORTS": "supported", "REFUTES": "refuted"}

data_map = {}

for name in name_lst:
    instance_lst = []

    for entry in tqdm(v10[name]):
        id_ = entry["id"]
        label = entry["label"]
        claim = entry["claim"]

        evidence_id = entry["evidence_id"]
        evidence_wiki_url = entry["evidence_wiki_url"]

        if evidence_id != -1:
            assert label in {"SUPPORTS", "REFUTES"}

            instance = {"id": id_, "label": old_to_new_label_map[label], "claim": claim}
            instance_lst.append(instance)

    key = "dev" if name in {"labelled_dev"} else name

    instance_lst = sorted([dict(t) for t in {tuple(d.items()) for d in instance_lst}], key=lambda d: d["claim"])

    label_to_instance_lst = {}
    for e in instance_lst:
        if e["label"] not in label_to_instance_lst:
            label_to_instance_lst[e["label"]] = []
        label_to_instance_lst[e["label"]].append(e)

    min_len = min(len(v) for k, v in label_to_instance_lst.items())

    new_instance_lst = []
    for k in sorted(label_to_instance_lst.keys()):
        new_instance_lst += label_to_instance_lst[k][:min_len]

    random.Random(42).shuffle(new_instance_lst)
    data_map[key] = new_instance_lst

ds_path = "pminervini/hl-fever"

task_to_ds_map = {k: Dataset.from_dict(convert(v)) for k, v in data_map.items()}
ds_dict = DatasetDict(task_to_ds_map)

ds_dict.push_to_hub(ds_path, "v1.0")

# breakpoint()