import os
from PIL import Image, ImageDraw, ImageFont
from .utils import image_label_2_color


def get_flattened_output(docs):
  print("Running Flattened Output")
  flattened_output = []
  annotation_key = 'output'
  for doc in docs:
    flattened_output_item = {annotation_key: []}
    doc_annotation = doc[annotation_key]
    for i, span in enumerate(doc_annotation):
      if len(span['words']) > 1:
        for span_chunk in span['words']:
          flattened_output_item[annotation_key].append(
              {
                  'label': span['label'],
                  'text': span_chunk['text'],
                  'words': [span_chunk]
              }
          )
          
      else:
        flattened_output_item[annotation_key].append(span)
    flattened_output.append(flattened_output_item)
  return flattened_output


def annotate_image(image_path, annotation_object):
  print("Annotating Images")
  img = None
  image = Image.open(image_path).convert('RGBA')
  tmp = image.copy()
  label2color = image_label_2_color(annotation_object)
  overlay = Image.new('RGBA', tmp.size, (0, 0, 0)+(0,))
  draw = ImageDraw.Draw(overlay)
  font = ImageFont.load_default()

  predictions = [span['label'] for span in annotation_object['output']]
  boxes = [span['words'][0]['box'] for span in annotation_object['output']]
  for prediction, box in zip(predictions, boxes):
      draw.rectangle(box, outline=label2color[prediction],
                     width=3, fill=label2color[prediction]+(int(255*0.33),))
      draw.text((box[0] + 10, box[1] - 10), text=prediction,
                fill=label2color[prediction], font=font)

  img = Image.alpha_composite(tmp, overlay)
  img = img.convert("RGB")

  image_name = os.path.basename(image_path)
  image_name = image_name[:image_name.find('.')]
  output_folder = 'inferenced/'
  os.makedirs(output_folder, exist_ok=True)

  img.save(os.path.join(output_folder, f'{image_name}_inference.jpg'))