class Dataset(object):
    """ Data iterator we can use to train or test a model on, responsible for both storing data
    and deciding how to batch it. """

    def get_epoch(self):
        """ Returns an iterator of batches/elements to train on, these elements are what will get
        passed to model.encode. Usually (but not necessarily) a list/batch of training examples """
        raise NotImplementedError(self.__class__)

    def get_batches(self, n_batches):
        if len(self) < n_batches:
            raise ValueError()
        return itertools.islice(self.get_epoch(), n_batches)

    def get_epochs(self, n_epochs: int):
        for _ in range(n_epochs):
            for batch in self.get_epoch():
                yield batch

    def get_samples(self, n_samples: int):
        """
        Sample for the data, be default we sample batches but subclasses can
        override this method to provide other kinds of sampling (like sampling individual elements).
        Must return both an iterator and the exact size of the iterator.
        """
        return self.get_batches(n_samples), n_samples

    def percent_filtered(self):
        # TODO nicer to just have "unfiltered_size"?
        """ If any filtering was done, the percent of examples that were filtered. Exposed so evaluators
         can compute percentages fairly even if some examples were removed during pre-processing """
        return None

    def __len__(self):
        """ Number of batches per an epoch """
        raise NotImplementedError(self.__class__) class TrainingData(Configurable):

    def get_train(self) -> Dataset:
        raise NotImplementedError()

    def get_eval(self) -> Dict[str, Dataset]:
        raise NotImplementedError()

    def get_train_corpus(self) -> object:
        """
        Return an object derived from the training data that will be passed to the model's initialization phase,
        what exactly is returned can be arbitrary, but will need to be compatible with
        the model's requirements. Example, return word counts to be used to decide what word vecs to train.
        """
        raise NotImplementedError()

    def get_resource_loader(self) -> ResourceLoader:
        return ResourceLoader() class ListDataset(Dataset):
    """ Dataset with a fixed list of elements """

    def __init__(self, data: List, batching: ListBatcher, unfiltered_len: Optional[int]=None):
        self.data = data
        self.batching = batching
        self.unfiltered_len = unfiltered_len

    def get_samples(self, n_examples) -> Iterator:
        n_batches = n_examples // self.batching.get_max_batch_size()
        return self.get_batches(n_batches), n_batches

    @property
    def batch_size(self):
        return self.batching.get_fixed_batch_size()

    def get_epoch(self):
        return self.batching.get_epoch(self.data)

    def percent_filtered(self):
        if self.unfiltered_len is None:
            return None
        return (self.unfiltered_len - len(self.data)) / self.unfiltered_len

    def get_n_examples(self):
        return len(self.data)

    def __len__(self):
        return self.batching.epoch_size(len(self.data)) class ListBatcher(Configurable):
    def get_epoch(self, data: List):
        raise NotImplementedError()

    def get_fixed_batch_size(self):
        """ Return the batch size if it is constant, else None """
        raise NotImplementedError()

    def get_max_batch_size(self):
        """ Return upper bound on the batch size """
        raise NotImplementedError()

    def epoch_size(self, n_elements):
        raise NotImplementedError() class ResourceLoader(object):
    """
    Abstraction for models the need access to external resources to setup, currently just
    for word-vectors.
    """

    def __init__(self, load_vec_fn=load_word_vectors):
        self.load_vec_fn = load_vec_fn

    def load_word_vec(self, vec_name, voc=None):
        return self.load_vec_fn(vec_name, voc) def flatten_iterable(listoflists: Iterable[Iterable[T]]) -> List[T]:
    return [item for sublist in listoflists for item in sublist] def max_or_none(a, b):
    if a is None or b is None:
        return None
    return max(a, b) class Configurable(object):
    """
    Configurable classes have names, versions, and a set of parameters that are either "simple" aka JSON serializable
    types or other Configurable objects. Configurable objects should also be serializable via pickle.
    Configurable classes are defined mainly to give us a human-readable way of reading of the `parameters`
    set for different objects and to attach version numbers to them.

    By default we follow the format sklearn uses for its `BaseEstimator` class, where parameters are automatically
    derived based on the constructor parameters.
    """

    @classmethod
    def _get_param_names(cls):
        # fetch the constructor or the original constructor before
        init = cls.__init__
        if init is object.__init__:
            # No explicit constructor to introspect
            return []

        init_signature = signature(init)
        parameters = [p for p in init_signature.parameters.values()
                      if p.name != 'self']
        if any(p.kind == p.VAR_POSITIONAL for p in parameters):
            raise RuntimeError()
        return sorted([p.name for p in parameters])

    @property
    def name(self):
        return self.__class__.__name__

    @property
    def version(self):
        return 0

    def get_params(self):
        out = {}
        for key in self._get_param_names():
            v = getattr(self, key, None)
            if isinstance(v, Configurable):
                out[key] = v.get_config()
            elif hasattr(v, "get_config"):  # for keras objects
                out[key] = {"name": v.__class__.__name__, "config": v.get_config()}
            else:
                out[key] = v
        return out

    def get_config(self) -> Configuration:
        params = {k: describe(v) for k,v in self.get_params().items()}
        return Configuration(self.name, self.version, params)

    def __getstate__(self):
        state = dict(self.__dict__)
        if "version" in state:
            if state["version"] != self.version:
                raise RuntimeError()
        else:
            state["version"] = self.version
        return state

    def __setstate__(self, state):
        if "version" not in state:
            raise RuntimeError("Version should be in state (%s)" % self.__class__.__name__)
        if state["version"] != self.version:
            warn(("%s loaded with version %s, but class " +
                 "version is %s") % (self.__class__.__name__, state["version"], self.version))

        if "state" in state:
            self.__dict__ = state["state"]
        else:
            del state["version"]
            self.__dict__ = state class DatasetBuilder(Configurable):

    def build_dataset(self, data, evidence) -> Dataset:
        """ Map the intermeidate format to a Dataset object """
        raise NotImplementedError()

    def build_stats(self, data) -> object:
        """ Map the intermeidate format to corpus statistic object, as will used in `TrainingData` """
        raise NotImplementedError() class FilteredData(object):
    def __init__(self, data: List, true_len: int):
        self.data = data
        self.true_len = true_len

    def __add__(self, other):
        return FilteredData(self.data + other.data, self.true_len + other.true_len)
from collections import Counter
from typing import List, Union, Optional, Set, Dict
from docqa.dataset import Dataset, TrainingData, ListDataset, ListBatcher
from docqa.utils import ResourceLoader, flatten_iterable, max_or_none
from docqa.configurable import Configurable
from docqa.data_processing.preprocessed_corpus import DatasetBuilder, FilteredData
import numpy as np
        voc.update(point.question)
        if point.answer is not None:
            voc.update(point.answer.get_vocab())
        voc.update(point.get_context())
    return voc


class ParagraphQuestionFilter(Configurable):
    def keep(self, data_point: ContextAndQuestion) -> bool:
        raise NotImplemented()


class QuestionFilter(ParagraphQuestionFilter):
    def __init__(self, ques_size_th: int):
        self.ques_size_th = ques_size_th

    def keep(self, data_point: ContextAndQuestion):
        return len(data_point.question) <= self.ques_size_th


class AnswerWord(ParagraphQuestionFilter):
    def __init__(self, para_size_th: int):
        self.para_size_th = para_size_th

    def keep(self, data_point: ContextAndQuestion):
        return all(ans.para_word_end < self.para_size_th for ans in data_point.answer)


class AnySplitAnswerFilter(ParagraphQuestionFilter):
    def keep(self, data_point: SentencesAndQuestion):
        for answer in data_point.answer:
            if answer.sent_start != answer.sent_end:
                return False
        return True


class AnswerSentence(ParagraphQuestionFilter):
    def __init__(self, sent_size_th: Union[int, None] = None,
                 num_sent_th: Union[int, None]=None):
        self.num_sent_th = num_sent_th
        self.sent_size_th = sent_size_th

    def keep(self, data_point: SentencesAndQuestion):
        for answer in data_point.answer:
            if self.num_sent_th is not None:
                if answer.sent_end >= self.num_sent_th:
                    return False
            if self.sent_size_th is not None:
                if answer.word_start >= self.sent_size_th:
                    return False
        return True


def apply_filters(data: List, data_filters: List[ParagraphQuestionFilter], name: str):
    if len(data) == 0:
        raise ValueError()
    if len(data_filters) == 0:
        return data
    else:
        pruned = []
        removed = np.zeros(len(data_filters), dtype=np.int32)
        for x in data:
            keep = True
            for i,f in enumerate(data_filters):
                if not f.keep(x):
                    keep = False
                    removed[i] += 1
                    break
            if keep:
                pruned.append(x)
        for i,x in enumerate(data_filters):
            print("\t%s filtered %d(%.5f) from %s" % (x.__class__.__name__, removed[i], removed[i]/len(data), name))
        n_removed = len(data)-len(pruned)
        print("Pruned a total of %d/%d (%.3f) for %s" % (n_removed, len(data), n_removed/len(data), name))
        return pruned


def build_spec(batch_size: int,
               max_batch_size: int,
               data: List[ContextAndQuestion]) -> ParagraphAndQuestionSpec:
    max_ques_size = 0
    max_word_size = 0
    max_para_size = 0
    for data_point in data:
        context = data_point.get_context()
        max_word_size = max(max_word_size, max(len(word) for word in context))
        max_para_size = max(max_para_size, len(context))
        if data_point.question is not None:
            max_ques_size = max(max_ques_size, len(data_point.question))
            max_word_size = max(max_word_size, max(len(word) for word in data_point.question))
    return ParagraphAndQuestionSpec(batch_size, max_ques_size, max_para_size,
                                    max_word_size, max_batch_size)


class ParagraphAndQuestionDataset(ListDataset):
    """ Dataset with q/a pairs and that exposes some meta-data about its elements """
    def get_spec(self) -> ParagraphAndQuestionSpec:
        return build_spec(self.batching.get_fixed_batch_size(),
                          self.batching.get_max_batch_size(),
                          self.data)

    def get_vocab(self) -> Set[str]:
        return compute_voc(self.data)


class ParagraphAndQuestionsBuilder(DatasetBuilder):
    """ For use with the preprocesed_corpus framework """
    def __init__(self, batching: ListBatcher, sample=None, sample_seed=None):
        if sample_seed is not None and sample is None:
            raise ValueError("Seed set, but sampling not requested")
        self.batching = batching
        self.sample_seed = sample_seed
        self.sample = sample

    def build_stats(self, data):
        if isinstance(data, FilteredData):
            return QaCorpusLazyStats(data.data)
        else:
            return QaCorpusLazyStats(data)


    def build_dataset(self, data, evidence) -> Dataset: