from PIL import Image
import matplotlib.pyplot as plt
import io


# COCO classes
CLASSES = [
    'N/A', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
    'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A',
    'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse',
    'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack',
    'umbrella', 'N/A', 'N/A', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis',
    'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove',
    'skateboard', 'surfboard', 'tennis racket', 'bottle', 'N/A', 'wine glass',
    'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich',
    'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake',
    'chair', 'couch', 'potted plant', 'bed', 'N/A', 'dining table', 'N/A',
    'N/A', 'toilet', 'N/A', 'tv', 'laptop', 'mouse', 'remote', 'keyboard',
    'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'N/A',
    'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier',
    'toothbrush'
]
COLORS = [
    [0.000, 0.447, 0.741],
    [0.850, 0.325, 0.098],
    [0.929, 0.694, 0.125],
    [0.494, 0.184, 0.556],
    [0.466, 0.674, 0.188],
    [0.301, 0.745, 0.933],
]


# Update JSON dictionary with rounded values and classes
def generate_output_json(json_dict):
    json_dict['scores'] = [round(score, 3) for score in json_dict['scores']]
    json_dict['boxes'] = [[round(coord, 3) for coord in box] for box in json_dict['boxes']]
    json_dict['labels'] = [CLASSES[label] for label in json_dict['labels']]
    return json_dict


# Generate matplotlib figure from prediction scores and boxes
def generate_output_figure(image_path, results, threshold):
    pil_img = Image.open(image_path)

    plt.figure(figsize=(16, 10))
    plt.imshow(pil_img)
    ax = plt.gca()
    colors = COLORS * 100

    print("\t Detailed information...")
    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
        #box = [round(i, 2) for i in box]
        print(
            f"\t\t Detected {label} with confidence "
            f"{score} at location {box}"
        )

        if score > threshold:
            c = COLORS[hash(label) % len(COLORS)]
            ax.add_patch(
                plt.Rectangle((box[0], box[1]), box[2] - box[0], box[3] - box[1], fill=False, color=c, linewidth=3)
            )
            text = f"{label}: {score:0.2f}"
            ax.text(box[0], box[1], text, fontsize=15, bbox=dict(facecolor="yellow", alpha=0.5))
    plt.axis("off")

    return plt.gcf()


#  Generate PIL image from matplotlib figure
def generate_output_image(output_figure):
    # Convert matplotlib figure to PIL image
    #output_figure = plt.gcf()
    buf = io.BytesIO()
    output_figure.savefig(buf, bbox_inches="tight")
    buf.seek(0)
    output_pil_img = Image.open(buf)

    return output_pil_img


def generate_gradio_outputs(image_path, response_dict, threshold):
    output_json = generate_output_json(response_dict)
    output_figure = generate_output_figure(image_path, output_json, threshold)
    output_pil_img = generate_output_image(output_figure)
    return output_json, output_pil_img