
"""

This dataset was generated by:

# select a few seed records so there is some longer and shorter text, records with images and without, a few variations of each type
rm -rf data
python m4-ds-unpack.py --dataset_name_or_path /hf/m4-master/data/cm4/cm4-10000-v0.1 --ids 93-96,101,105,115,134,140,150 --target_path data

cd data

# shrink to 32x32 max, keeping ratio
mogrify -format jpg -resize 32x32\> */*jpg

# this is no longer needed found such existing record 150
# # adjust one record to have just images and no text (rec 134: 39 pairs with 3 images, 25 texts)
# cd 134
# rm texts*
# perl -le 'qx[touch texts_$_.null] for map {sprintf "%02d", $_} 0..38'
# cd -

cd ..

# create tarball
tar -cvzf data.tar.gz data

# prep dataset repo
https://huggingface.co/new-dataset => HuggingFaceM4/cm4-synthetic-testing
git clone https://huggingface.co/datasets/HuggingFaceM4/cm4-synthetic-testing

# complete the dataset repo
cp data.tar.gz cm4-synthetic-testing
cp cm4-synthetic-testing.py cm4-synthetic-testing
cp m4-ds-unpack.py cm4-synthetic-testing
echo "This dataset is designed to be used in testing. It's derived from cm4-10k dataset" >> cm4-synthetic-testing/README.md

# test dataset
datasets-cli test cm4-synthetic-testing/cm4-synthetic-testing.py --all_configs

# push the data
cd cm4-synthetic-testing
git add *
git commit -am "new dataset"
git push

# test
python -c 'from datasets import load_dataset; load_dataset("HuggingFaceM4/cm4-synthetic-testing")'

"""


from PIL import Image, ImageFile
from collections import defaultdict
from datasets import DatasetInfo
from pathlib import Path
from pprint import pprint
import datasets
import itertools
import json
import os

_CITATION = """\
@InProceedings{huggingface:dataset,
title = {Multimodal synthetic dataset for testing},
author={HuggingFace, Inc.},
year={2022}
}
"""

_DESCRIPTION = """This dataset is designed to be used in testing. It's derived from cm4-10k dataset"""
_HOMEPAGE = "https://huggingface.co/datasets/HuggingFaceM4/cm4-synthetic-testing"
_LICENSE = "bigscience-openrail-m"
_URL = "https://huggingface.co/datasets/HuggingFaceM4/cm4-synthetic-testing/resolve/main/data.tar.gz"
#_URL = "./data.tar.gz"

sizes = ["100", "300", "1k", "10k"]
types = ["unique", "repeat"]

class CM4Synthetic(datasets.GeneratorBasedBuilder):

    VERSION = datasets.Version("1.1.2")

    # splits = [f"{s}.{t}" for s in sizes for t in types]
    # BUILDER_CONFIGS = [] # can't use list comprehension and access VERSION due to python scoping design
    # for split in splits:
    #     BUILDER_CONFIGS.append(datasets.BuilderConfig(name=split, version=VERSION, description=f"{split} items split"))
    DEFAULT_CONFIG_NAME = "100.unique"

    def _info(self):
        # script_dir = os.path.abspath(os.path.dirname(__file__))
        # path = os.path.join(script_dir, "dataset_info.json")
        # ds_info = DatasetInfo.from_directory(path)
        # pprint(ds_info)
        # return ds_info

        # XXX: automate
        return datasets.DatasetInfo(
            description=_DESCRIPTION,
            citation=_CITATION,
            homepage=_HOMEPAGE,
            license=_LICENSE,
            features={
                "texts": {
                    "feature": {"dtype": "string", "id": None, "_type": "Value"},
                    "length": -1,
                    "id": None,
                    "_type": "Sequence",
                },
                "images": {
                    "feature": {"decode": True, "id": None, "_type": "Image"},
                    "length": -1,
                    "id": None,
                    "_type": "Sequence",
                },
            },
        )

    def _split_generators(self, dl_manager):
        url = _URL
        data_dir = dl_manager.download_and_extract(url)

        return [
            datasets.SplitGenerator(
                name=self.config.name,
                # These kwargs will be passed to _generate_examples
                gen_kwargs={
                    "data_path": os.path.join(data_dir, "data"),
                },
            )
        ]

    def _generate_examples(self, data_path):
        # the split name acts as the designator of how many rows to generate

        size, type = self.config.name.split(".")

        print(f"Generating {size}-long {type} records split")

        # for now handling 100, 10k - can add m
        total_examples = int(size.replace("k", "000"))

        def pack_example(path):
            """ put the directory with texts and images into a single datasets record """
            images = []
            for img_file in sorted(path.glob("images_*")):
                if str(img_file).endswith(".null"):
                    images.append(None)
                else:
                    images.append(Image.open(img_file))

            texts = []
            for txt_file in sorted(path.glob("texts_*")):
                if str(txt_file).endswith(".null"):
                    texts.append(None)
                else:
                    text = "".join([l for l in open(txt_file)])
                    texts.append(text)

            return dict(images=images, texts=texts)

        def dump_example_shapes(idx, row):
            """ dump the row stats """

            imgs = defaultdict(int)
            for img in row["images"]:
                if img is None:
                    imgs["0"] += 1
                else:
                    shape = "x".join(map(str, img.size))
                    imgs[shape] += 1
            imgs_summary = ", ".join([f"{v} {k}" for k,v in sorted(imgs.items(), key=lambda x: int(x[0].split("x")[0]))])

            txts = defaultdict(int)
            for txt in row["texts"]:
                if txt is None:
                    txts[0] += 1
                else:
                    shape = len(txt)
                    txts[shape] += 1
            txts_summary = ", ".join([f"{v} {k}" for k,v in sorted(txts.items(), key=lambda x: int(x[0]))])

            print(f"\nrec{idx}: {len(row['images'])} pairs with {len(row['images'])-imgs['0']} images, {len(row['texts'])-txts[0]} texts")
            print(f"- img: {imgs_summary}")
            print(f"- txt: {txts_summary}")

        print()
        rows = [pack_example(subdir) for subdir in sorted(Path(data_path).glob("[0-9]*"))]
        num_rows = len(rows)
        if num_rows == 0:
            raise ValueError(f"can't find any data - check {data_path}")

        print(f"\nStats for {len(rows)} unique records used:")
        for i, row in enumerate(rows): dump_example_shapes(i, row)

        one_zero_texts = 0
        one_none_texts = 0
        def gen_unique_rec(idx, row):
            nonlocal one_zero_texts, one_none_texts
            """ insert idx as a string at the end of the first non-None entry, or create a new one if all are
            None. The replacement will overwrite the last few characters of the previous string. This ensures
            that each record will be unique """

            texts = row["texts"]
            num_texts = len(texts)
            uniq_text = str(idx)
            if num_texts == 0:
                # keep one record that has no text entries and no None's (which is still unique)
                if one_zero_texts == 0:
                    one_zero_texts = 1
                else:
                    texts.append(uniq_text)
            else:
                for i in range(num_texts):
                    if texts[i] is not None:
                        texts[i] = texts[i][:-len(uniq_text)] + uniq_text
                        break
                else:
                    # keep one record that has only None's for texts (which is still unique)
                    if one_none_texts == 0:
                        one_none_texts = 1
                    else:
                        texts[i] = uniq_text
            return row

        # this being a synthetic dataset we rotate the 1 or more available rows until we generate enough records.
        # in the case of unique type we tweak one text record to be unique
        for i in range(total_examples):
            idx = i % num_rows
            if type == "repeat":
                yield i, rows[idx]
            elif type == "unique":
                yield i, gen_unique_rec(i, rows[idx])
