"""Program Synthesis dataset from dreamcoder. https://github.com/ellisk42/ec"""
from random import choice, shuffle
import datasets

from dreamcoder.domains.text.makeTextTasks import makeTasks as textMakeTasks
from dreamcoder.domains.list.main import main as listMakeTasks


_DESCRIPTION = """\
Generated program synthesis datasets used to train dreamcoder.
"""
_FEATURES = datasets.Features(
    {
        "description": datasets.Value("string"),
        "input": datasets.Value("string"),
        "output": datasets.Value("string"),
        "types": datasets.Value("string")
    }
)

_HOMEPAGE = "https://github.com/ellisk42/ec"

_LICENSE = "MIT License"

_MAX_STEPS = 10


class infIterator:
    def __init__(self, make_mthd):
        self.make_mthd = make_mthd
        self.i = None

    def reset(self):
        tasks = self.make_mthd()

        rows = []
        for task in tasks:
            base = {
                'types': str(task.request),
                "description": task.name,
            }
            for (inp, outp) in task.examples:
                rows.append(dict(input=str(inp), output=str(outp), **base))

        shuffle(rows)
        self.rows = rows
        self.i = 0

    def step(self):
        if self.i is None:
            self.reset()
        row = self.rows[self.i]
        self.i += 1
        if self.i >= len(self.rows):
            self.reset()
        return row


class ProgramSynthesis(datasets.GeneratorBasedBuilder):
    """Program Synthesis dataset from dreamcoder."""

    VERSION = datasets.Version("1.1.0")
    BUILDER_CONFIGS = [
        datasets.BuilderConfig(name="text", version=VERSION, description="Text tasks."),
        datasets.BuilderConfig(name="list", version=VERSION, description="List tasks."),
        datasets.BuilderConfig(name="all", version=VERSION, description="All tasks at once."),
    ]
    DEFAULT_CONFIG_NAME = "all"

    def _info(self):
        return datasets.DatasetInfo(
            description=_DESCRIPTION,
            features=_FEATURES,
            supervised_keys=("input", "output"),
            homepage=_HOMEPAGE,
            license=_LICENSE,
        )

    def _split_generators(self, dl_manager):
        return [
            datasets.SplitGenerator(
                name=datasets.Split.TRAIN,
            ),
        ]

    def _generate_examples(self):
        task_samples = {
            'text': infIterator(textMakeTasks),
            'list': infIterator(listMakeTasks)
        }
        for key in range(_MAX_STEPS):

            if self.config.name == 'all':
                dataset_type = choice(task_samples.keys())
            else:
                dataset_type = self.config.name

            yield key, task_samples[dataset_type].step()