text: import json
  import logging
  import os
  from collections import defaultdict, namedtuple
  from typing import *
  
  from allennlp.data.dataset_readers.dataset_reader import DatasetReader
  from allennlp.data.instance import Instance
  
  from .span_reader import SpanReader
  from ..utils import Span
  
  # logging.basicConfig(level=logging.DEBUG)
  
  # for v in logging.Logger.manager.loggerDict.values():
  # v.disabled = True
  
  logger = logging.getLogger(__name__)
  
  SpanTuple = namedtuple('Span', ['start', 'end'])
  
  
  @DatasetReader.register('better')
  class BetterDatasetReader(SpanReader):
      def __init__(
              self,
              eval_type,
              consolidation_strategy='first',
              span_set_type='single',
              max_argument_ss_size=1,
              use_ref_events=False,
              **extra
      ):
          super().__init__(**extra)
          self.eval_type = eval_type
          assert self.eval_type in ['abstract', 'basic']
  
          self.consolidation_strategy = consolidation_strategy
          self.unitary_spans = span_set_type == 'single'
          # event anchors are always singleton spans
          self.max_arg_spans = max_argument_ss_size
          self.use_ref_events = use_ref_events
  
          self.n_overlap_arg = 0
          self.n_overlap_trigger = 0
          self.n_skip = 0
          self.n_too_long = 0
  
      @staticmethod
      def post_process_basic_span(predicted_span, basic_entry):
          # Convert token offsets back to characters, also get the text spans as a sanity check
  
          # !!!!!
          # SF outputs inclusive idxs
          # char offsets are inc-exc
          # token offsets are inc-inc
          # !!!!!
  
          start_idx = predicted_span['start_idx']  # inc
          end_idx = predicted_span['end_idx']  # inc
  
          char_start_idx = basic_entry['tok2char'][predicted_span['start_idx']][0]  # inc
          char_end_idx = basic_entry['tok2char'][predicted_span['end_idx']][-1] + 1  # exc
  
          span_text = basic_entry['segment-text'][char_start_idx:char_end_idx]  # inc exc
          span_text_tok = basic_entry['segment-text-tok'][start_idx:end_idx + 1]  # inc exc
  
          span = {'string': span_text,
                  'start': char_start_idx,
                  'end': char_end_idx,
                  'start-token': start_idx,
                  'end-token': end_idx,
                  'string-tok': span_text_tok,
                  'label': predicted_span['label'],
                  'predicted': True}
          return span
  
      @staticmethod
      def _get_shortest_span(spans):
          # shortest_span_length = float('inf')
          # shortest_span = None
          # for span in spans:
          # span_tokens = span['string-tok']
          # span_length = len(span_tokens)
          # if span_length < shortest_span_length:
          # shortest_span_length = span_length
          # shortest_span = span
  
          # return shortest_span
          return [s[-1] for s in sorted([(len(span['string']), ix, span) for ix, span in enumerate(spans)])]
  
      @staticmethod
      def _get_first_span(spans):
          spans = [(span['start'], -len(span['string']), ix, span) for ix, span in enumerate(spans)]
          try:
              return [s[-1] for s in sorted(spans)]
          except:
              breakpoint()
  
      @staticmethod
      def _get_longest_span(spans):
          return [s[-1] for s in sorted([(len(span['string']), ix, span) for ix, span in enumerate(spans)], reverse=True)]
  
      @staticmethod
      def _subfinder(text, pattern):
          # https://stackoverflow.com/a/12576755
          matches = []
          pattern_length = len(pattern)
          for i, token in enumerate(text):
              try:
                  if token == pattern[0] and text[i:i + pattern_length] == pattern:
                      matches.append(SpanTuple(start=i, end=i + pattern_length - 1))  # inclusive boundaries
              except:
                  continue
          return matches
  
      def consolidate_span_set(self, spans):
          if self.consolidation_strategy == 'first':
              spans = BetterDatasetReader._get_first_span(spans)
          elif self.consolidation_strategy == 'shortest':
              spans = BetterDatasetReader._get_shortest_span(spans)
          elif self.consolidation_strategy == 'longest':
              spans = BetterDatasetReader._get_longest_span(spans)
          else:
              raise NotImplementedError(f"{self.consolidation_strategy} does not exist")
  
          if self.unitary_spans:
              spans = [spans[0]]
          else:
              spans = spans[:self.max_arg_spans]
  
          # TODO add some sanity checks here
  
          return spans
  
      def get_mention_spans(self, text: List[str], span_sets: Dict):
          mention_spans = defaultdict(list)
          for span_set_id in span_sets.keys():
              spans = span_sets[span_set_id]['spans']
              # span = BetterDatasetReader._get_shortest_span(spans)
              # span = BetterDatasetReader._get_earliest_span(spans)
              consolidated_spans = self.consolidate_span_set(spans)
              # if len(spans) > 1:
              # logging.info(f"Truncated a spanset from {len(spans)} spans to 1")
  
              if self.eval_type == 'abstract':
                  span = consolidated_spans[0]
                  span_tokens = span['string-tok']
  
                  span_indices = BetterDatasetReader._subfinder(text=text, pattern=span_tokens)
  
                  if len(span_indices) > 1:
                      pass
  
                  if len(span_indices) == 0:
                      continue
  
                  mention_spans[span_set_id] = span_indices[0]
              else:
                  # in basic, we already have token offsets in the right form
  
                  # if not span['string-tok'] == text[span['start-token']:span['end-token'] + 1]:
                  # print(span, text[span['start-token']:span['end-token'] + 1])
  
                  # we should use these token offsets only!
                  for span in consolidated_spans:
                      mention_spans[span_set_id].append(SpanTuple(start=span['start-token'], end=span['end-token']))
  
          return mention_spans
  
      def _read_single_file(self, file_path):
          with open(file_path) as fp:
              json_content = json.load(fp)
          if 'entries' in json_content:
              for doc_name, entry in json_content['entries'].items():
                  instance = self.text_to_instance(entry, 'train' in file_path)
                  yield instance
          else:  # TODO why is this split in 2 cases?
              for doc_name, entry in json_content.items():
                  instance = self.text_to_instance(entry, True)
                  yield instance
  
          logger.warning(f'{self.n_overlap_arg} overlapped args detected!')
          logger.warning(f'{self.n_overlap_trigger} overlapped triggers detected!')
          logger.warning(f'{self.n_skip} skipped detected!')
          logger.warning(f'{self.n_too_long} were skipped because they are too long!')
          self.n_overlap_arg = self.n_skip = self.n_too_long = self.n_overlap_trigger = 0
  
      def _read(self, file_path: str) -> Iterable[Instance]:
  
          if os.path.isdir(file_path):
              for fn in os.listdir(file_path):
                  if not fn.endswith('.json'):
                      logger.info(f'Skipping {fn}')
                      continue
                  logger.info(f'Loading from {fn}')
                  yield from self._read_single_file(os.path.join(file_path, fn))
          else:
              yield from self._read_single_file(file_path)
  
      def text_to_instance(self, entry, is_training=False):
          word_tokens = entry['segment-text-tok']
  
          # span sets have been trimmed to the earliest span mention
          spans = self.get_mention_spans(
              word_tokens, entry['annotation-sets'][f'{self.eval_type}-events']['span-sets']
          )
  
          # idx of every token that is a part of an event trigger/anchor span
          all_trigger_idxs = set()
  
          # actual inputs to the model
          input_spans = []
  
          self._local_child_overlap = 0
          self._local_child_total = 0
  
          better_events = entry['annotation-sets'][f'{self.eval_type}-events']['events']
  
          skipped_events = set()
          # check for events that overlap other event's anchors, skip them later
          for event_id, event in better_events.items():
              assert event['anchors'] in spans
  
              # take the first consolidated span for anchors
              anchor_start, anchor_end = spans[event['anchors']][0]
  
              if any(ix in all_trigger_idxs for ix in range(anchor_start, anchor_end + 1)):
                  logger.warning(
                      f"Skipped {event_id} with anchor span {event['anchors']}, overlaps a previously found event trigger/anchor")
                  self.n_overlap_trigger += 1
                  skipped_events.add(event_id)
                  continue
  
              all_trigger_idxs.update(range(anchor_start, anchor_end + 1))  # record the trigger
  
          for event_id, event in better_events.items():
              if event_id in skipped_events:
                  continue
  
              # arguments for just this event
              local_arg_idxs = set()
              # take the first consolidated span for anchors
              anchor_start, anchor_end = spans[event['anchors']][0]
  
              event_span = Span(anchor_start, anchor_end, event['event-type'], True)
              input_spans.append(event_span)
  
              def add_a_child(span_id, label):
                  # TODO this is a bad way to do this
                  assert span_id in spans
                  for child_span in spans[span_id]:
                      self._local_child_total += 1
                      arg_start, arg_end = child_span
  
                      if any(ix in local_arg_idxs for ix in range(arg_start, arg_end + 1)):
                          # logger.warn(f"Skipped argument {span_id}, overlaps a previously found argument")
                          # print(entry['annotation-sets'][f'{self.eval_type}-events']['span-sets'][span_id])
                          self.n_overlap_arg += 1
                          self._local_child_overlap += 1
                          continue
  
                      local_arg_idxs.update(range(arg_start, arg_end + 1))
                      event_span.add_child(Span(arg_start, arg_end, label, False))
  
              for agent in event['agents']:
                  add_a_child(agent, 'agent')
              for patient in event['patients']:
                  add_a_child(patient, 'patient')
  
              if self.use_ref_events:
                  for ref_event in event['ref-events']:
                      if ref_event in skipped_events:
                          continue
                      ref_event_anchor_id = better_events[ref_event]['anchors']
                      add_a_child(ref_event_anchor_id, 'ref-event')
  
              # if len(event['ref-events']) > 0:
              # breakpoint()
  
          fields = self.prepare_inputs(word_tokens, spans=input_spans)
          if self._local_child_overlap > 0:
              logging.warning(
                  f"Skipped {self._local_child_overlap} / {self._local_child_total} argument spans due to overlaps")
          return Instance(fields)
  
  