import argparse
import os
import pickle
from bisect import bisect_right
from collections import defaultdict
from multiprocessing import cpu_count

from datasets import load_from_disk
from tqdm import tqdm


def get_pairs(byterange):
    """
    Returns pairs generated by
    https://github.com/google-research/deduplicate-text-datasets#collecting-the-duplicates-together
    """
    print("Getting pairs")
    pairs = []
    with open(byterange, "r") as f:
        save = False
        for line in tqdm(f):
            if line.strip() == "out":
                save = True
                continue
            if save:
                left, right = line.strip().split()
                pairs.append((int(left), int(right)))
    print("num pairs", len(pairs))
    return pairs


def get_bytes(pairs, data):
    """
    Return bytes constituring the duplicated substring. There seems to be something off here, see:
    https://github.com/google-research/deduplicate-text-datasets/issues/24
    """
    print("Getting bytes")
    byte_array = []
    for left, right in tqdm(pairs):
        byte_array.append(data[left:right])

    print("byte_array size", len(byte_array))
    return byte_array


def get_doc_id(pos, pos2id, pos2id_list):
    """
    Gets id of the datapoint at position.
    """
    pos = bisect_right(pos2id_list, pos)
    doc_id = pos2id[pos2id_list[pos - 1]]
    return doc_id


def get_url(row, dataset_name):
    if dataset_name == "oscar":
        return row["meta"]["warc_headers"]["warc-target-uri"]
    if dataset_name == "the_pile" or dataset_name == "roots_en":
        return None
    return None


if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--dataset_name",
        "-d",
        type=str,
        required=True,
    )
    parser.add_argument(
        "--base_dir",
        "-b",
        type=str,
        required=True,
    )
    args = parser.parse_args()
    dataset_name = args.dataset_name
    base_dir = args.base_dir

    HUGGINGFACE_TOKEN = os.environ.get("HUGGINGFACE_TOKEN")

    print("Loading", dataset_name)
    dataset = load_from_disk("/home/piktus_huggingface_co/lumi/preprocessed_data/{}".format(dataset_name))["train"]
    dataset_raw = load_from_disk("/home/piktus_huggingface_co/lumi/{}".format(dataset_name))["train"]
    print("Done loading", dataset_name)

    data_path = base_dir + "/{ds}.train".format(ds=dataset_name)
    byterange = base_dir + "/{ds}.train.byterange".format(ds=dataset_name)
    pairs = get_pairs(byterange)
    data = open(data_path, "rb").read()
    pos2id = pickle.load(open(base_dir + "/{ds}.train.pos2id.pkl".format(ds=dataset_name), "rb"))
    pos2id_list = sorted(pos2id.keys())

    doc_pairs = defaultdict(list)
    dup_len = defaultdict(int)
    doc_bytes = defaultdict(list)
    repetitions = defaultdict(set)
    clusters = defaultdict(set)
    included_docs = set()

    print("Getting bytes")
    byte_array = []
    for left, right in tqdm(pairs):
        byte_array.append(data[left:right])

    print("byte_array size", len(byte_array))

    print("Calculating repetitions")
    for (l, r), b in tqdm(zip(pairs, byte_array)):
        i = get_doc_id(l, pos2id, pos2id_list)
        doc_pairs[i].append((l, r))
        dup_len[i] += (r - l)
        doc_bytes[i].append(b)
        repetitions[b].add(i)

    print(
        "Singleton pairs:",
        sum(1 for i in doc_pairs.keys() if len(doc_pairs[i]) == 1)
    )

    for idx in pos2id.values():
        included_docs.add(idx)
    print("Num included", len(included_docs))

    print("Calculating clusters")
    for i, _ in tqdm(doc_pairs.items()):
        c = set()
        for rep in doc_bytes[i]:
            c |= repetitions[rep]
        clusters[i] = c

    # cleanup
    del byterange
    del data
    del pos2id
    del pos2id_list
    del pairs

    def add_duplication_info(example, idx):
        example["text"] = dataset_raw[idx]["text"]
        text_length = len(example["text"])
        example["text_length"] = text_length
        example["url"] = get_url(example, dataset_name)
        example["domain"] = example["url"].split("/")[2] if example["url"] is not None else None
        example["dup_ratio"] = dup_len[idx] / text_length
        example["pairs"] = doc_pairs[idx]
        example["repetitions"] = doc_bytes[idx]
        example["included_in_dedup"] = idx in included_docs
        example["cluster"] = clusters[idx]
        return example

    print("Mapping")
    dataset = dataset.map(add_duplication_info, num_proc=max(1, cpu_count() - 32), with_indices=True, new_fingerprint="13")
    dataset.save_to_disk("/home/piktus_huggingface_co/lumi/preprocessed_data/{}-dedup".format(dataset_name))
    dataset.push_to_hub(
        "datablations/{}-dedup".format(dataset_name),
        private=True,
        token=os.environ.get("HUGGINGFACE_TOKEN"),
    )

