class ParagraphAndQuestionSpec(object):
    """ Bound on the size of `ParagraphAndQuestion` objects """
    def __init__(self, batch_size, max_question_words=None,
                 max_num_context_words=None, max_word_size=None, max_batch_size=None):
        if batch_size is not None:
            if max_batch_size is None:
                max_batch_size = batch_size
            elif max_batch_size != batch_size:
                raise ValueError()
        self.batch_size = batch_size
        self.max_num_quesiton_words = max_question_words
        self.max_num_context_words = max_num_context_words
        self.max_word_size = max_word_size
        self.max_batch_size = max_batch_size

    def __add__(self, o):
        return ParagraphAndQuestionSpec(
            max_or_none(self.batch_size, o.batch_size),
            max_or_none(self.max_num_quesiton_words, o.max_num_quesiton_words),
            max_or_none(self.max_num_context_words, o.max_num_context_words),
            max_or_none(self.max_word_size, o.max_word_size),
            max_or_none(self.max_batch_size, o.max_batch_size)
        ) class WordCounts(object):
    def __init__(self, word_counts):
        self.word_counts = word_counts

    def get_word_counts(self):
        return self.word_counts class ParagraphAndQuestion(ContextAndQuestion):

    def __init__(self, context: List[str], question: List[str],
                 answer: Optional[Answer], question_id: object, doc_id=None):
        super().__init__(question, answer, question_id, doc_id)
        self.context = context

    @property
    def n_context_words(self):
        return len(self.context)

    def get_context(self):
        return self.context class ContextAndQuestion(object):
    """
    This is our standard unit of training data, context and a single question/answer.
    The answer type is unspecified and depends on the application.
    """

    def __init__(self, question: List[str], answer: Optional[Answer],
                 question_id: object, doc_id=None):
        self.question = question
        self.answer = answer
        self.question_id = question_id
        self.doc_id = doc_id

    @property
    def n_context_words(self) -> int:
        raise NotImplementedError()

    def get_context(self) -> List[str]:
        raise NotImplementedError() 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 TokenSpans(Answer):
    __slots__ = ["answer_text", "answer_spans"]

    def __init__(self, answer_text: List[str], answer_spans: np.ndarray):
        """
        :param answer_text: list of text answers
        :param answer_spans: (n, 2) array of inclusive (start, end) occurrences of an answer
        """
        self.answer_text = answer_text
        self.answer_spans = answer_spans

    def get_vocab(self):
        return [] 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 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 ClusteredBatcher(ListBatcher):
    def __init__(self,
                 batch_size: int,
                 clustering: Callable,
                 shuffle_buckets=False,
                 truncate_batches=False):
        self.batch_size = batch_size
        self.clustering = clustering
        self.shuffle_buckets = shuffle_buckets
        self.truncate_batches = truncate_batches

    def get_fixed_batch_size(self):
        return None if self.truncate_batches else self.batch_size

    def get_max_batch_size(self):
        return self.batch_size

    def get_epoch(self, data: List):
        data = sorted(data, key=self.clustering)
        if self.shuffle_buckets:
            shuffle_list_buckets(data, self.clustering, np.random)
        n_batches = len(data) // self.batch_size
        intervals = [(i*self.batch_size, (i + 1)*self.batch_size) for i in range(0, n_batches)]
        remainder = len(data) % self.batch_size
        if self.truncate_batches and remainder > 0:
            intervals.append((len(data) - remainder, len(data)))
        np.random.shuffle(intervals)
        for i, j in intervals:
            yield data[i:j]

    def epoch_size(self, n_elements):
        size = n_elements // self.batch_size
        if self.truncate_batches and (n_elements % self.batch_size) > 0:
            size += 1
        return size 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
from docqa.data_processing.qa_training_data import ParagraphAndQuestionSpec, WordCounts, ParagraphAndQuestion, \
    ContextAndQuestion, ParagraphAndQuestionDataset
from docqa.data_processing.span_data import TokenSpans
from docqa.dataset import Dataset, ListBatcher, ClusteredBatcher
from docqa.data_processing.preprocessed_corpus import DatasetBuilder, FilteredData
import numpy as np
                    not (force_answer and len(sample1) == 1)):  # don't bother if there only is one answer
                ix = 0
                for i, p in enumerate(q.paragraphs):
                    if len(p.answer_spans) > 0:
                        sample1 += [i] * self.overample_first_answer[ix]
                        ix += 1
                        if ix >= len(self.overample_first_answer):
                            break

            permutations = []
            for i in sample1:
                for j in range(len(q.paragraphs)):
                    if j != i:
                        permutations.append((i, j))
            permutations = np.array(permutations, dtype=np.int32)
            np.random.shuffle(permutations)
            self._order.append(permutations)

    def get_vocab(self):
        voc = set()
        for q in self.questions:
            voc.update(q.question)
            for para in q.paragraphs:
                voc.update(para.text)
        return voc

    def get_spec(self):
        max_q_len = max(len(q.question) for q in self.questions)
        max_c_len = max(max(len(p.text) for p in q.paragraphs) for q in self.questions)
        return ParagraphAndQuestionSpec(None, max_q_len, max_c_len, None)

    def get_epoch(self):
        return self._build_expanded_batches(self.questions)

    def _build_expanded_batches(self, questions):
        out = []
        # Decide what paragraphs to use for each question
        for i, q in enumerate(questions):
            order = self._order[i]
            out.append(ParagraphSelection(q, order[self._on[i]]))
            self._on[i] += 1
            if self._on[i] == len(order):
                self._on[i] = 0
                np.random.shuffle(order)

        # Sort by context length
        out.sort(key=lambda x: x.n_context_words)

        # Yield the correct batches
        group = 0
        for selection_batch in self.batcher.get_epoch(out):
            batch = []
            for selected in selection_batch:
                q = selected.question
                if self.merge:
                    paras = [q.paragraphs[i] for i in selected.selection]
                    # Sort paragraph by reading order, not rank order
                    paras.sort(key=lambda x: x.get_order())
                    answer_spans = []
                    text = []
                    for para in paras:
                        answer_spans.append(len(text) + para.answer_spans)
                        text += para.text
                    batch.append(ParagraphAndQuestion(text, q.question,
                                                      TokenSpans(q.answer_text, np.concatenate(answer_spans)),
                                                      q.question_id))
                else:
                    for i in selected.selection:
                        para = q.paragraphs[i]
                        batch.append(para.build_qa_pair(q.question, q.question_id, q.answer_text, group))
                    group += 1
            yield batch

    def get_samples(self, n_examples):
        n_batches = self.batcher.epoch_size(n_examples)
        return self._build_expanded_batches(np.random.choice(self.questions, n_examples, replace=False)), n_batches

    def percent_filtered(self):
        return (self.true_len - len(self.questions)) / self.true_len

    def __len__(self):
        return self.batcher.epoch_size(len(self.questions))


def multi_paragraph_word_counts(data):
    wc = Counter()
    for point in data:
        wc.update(point.question)
        for para in point.paragraphs:
            wc.update(para.get_context())
    return WordCounts(wc)


class IndividualParagraphBuilder(DatasetBuilder):
    """ Treat each paragraph as its own training point """

    def __init__(self, batcher: ListBatcher, force_answer: float):
        self.batcher = batcher
        self.force_answer = force_answer

    def build_stats(self, data: FilteredData):
        return multi_paragraph_word_counts(data.data)

    def build_dataset(self, data: FilteredData, corpus) -> Dataset:
        flattened = []
        for point in data.data:
            for para in point.paragraphs:
                flattened.append(ParagraphAndQuestion(para.text, point.question,
                                            TokenSpans(point.answer_text, para.answer_spans),
                                            point.question_id))
        return ParagraphAndQuestionDataset(flattened, self.batcher)


class RandomParagraphsBuilder(DatasetBuilder):
    def __init__(self, batching: ListBatcher, force_answer: float, n_to_sample=1):
        self.batching = batching
        self.force_answer = force_answer
        self.n_to_sample = n_to_sample

    def build_stats(self, data: Union[FilteredData, List]):

        if isinstance(data, FilteredData):