import copy


def extract_data(dataset, queries):
    """
    From [dataset], extract antecedent annotations for [queries].
    Returns (new_data, num_queries) where:
     - [new_data] is a subset of [dataset] which contains annotations only for spans in [queries]
     - [num_queries] is number of annotated spans in [new_data]
    """
    stop = len(queries)
    new_data = []
    num_queries = 0

    # iterate over [dataset] with stopping condition for faster search
    for doc in dataset:

        if stop <= 0:
            # already queried labels for all sampled spans
            break

        doc_key = doc['doc_key']
        if doc_key in queries:
            # at least one span in doc was queried
            gold_map = doc['antecedent_map']
            # get most recent antecedent for queried spans
            sub_map = {}
            for span in queries.get(doc_key, []):
                if span in gold_map:
                    ante = gold_map[span]
                    if type(ante) is list:
                        # grab most recent antecedent if not '0' cluster
                        ante = [max(ante)]
                    sub_map[span] = ante
                else:
                    sub_map[span] = None
                num_queries += 1

            if len(sub_map) > 0:
                # create new copy of doc with new sub_map
                new_doc = copy.copy(doc)
                new_doc['antecedent_map'] = sub_map
                new_data.append(new_doc)
            stop -= 1
    return new_data, num_queries