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 = 'static/temp/img_display/'
    os.makedirs(output_folder, exist_ok=True)

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