import json
import os
from uuid import uuid4


def export_one_prediction_to_label_studio(prediction, label_studio_relative_path: str, output_folder: str,
                                          model_name: str, from_name: str,
                                          threshold_to_export: float = 0.1):
    image_name = label_studio_relative_path.split('/')[-1]
    output_file = image_name.replace('.jpg', ".json").replace('.jpeg', ".json")
    output_path = os.path.join(output_folder, output_file)

    original_width, original_height = prediction.orig_shape
    boxes_filtered = filter(lambda x: float(x.conf[0]) > threshold_to_export, prediction.boxes)
    json_content = {
        "data": {
            "img": label_studio_relative_path,
        },
        "predictions": [
            {
                "model_version": model_name,
                "result": [
                    {
                        "original_width": original_width,
                        "original_height": original_height,
                        "score": float(box.conf[0]),
                        "value": {
                            # nightmareeeeeee
                            "x": float(box.xywhn[0][0] - box.xywhn[0][2] / 2) * 100,
                            "y": float(box.xywhn[0][1] - box.xywhn[0][3] / 2) * 100,
                            "width": float(box.xywhn[0][2]) * 100,
                            "height": float(box.xywhn[0][3]) * 100,
                            "rotation": 0
                        },
                        "id": str(uuid4()),
                        "from_name": from_name,
                        "to_name": "image",
                        "type": "rectangle"
                    } for box in boxes_filtered],
            }]
    }

    with open(output_path, 'w') as fp:
        json.dump(json_content, fp)

    return json_content


def export_one_yolo_output_to_label_studio_rectangles(yolo_prediction,
                                                      model_version: str,
                                                      label_studio_relative_path: str,
                                                      output_folder: str,
                                                      label_studio_to_name: str = 'image',
                                                      label_studio_from_name: str = 'label',
                                                      confidence_threshold: float = 0.1,
                                                      overwrite_existing_files: bool = False):
    os.makedirs(output_folder, exist_ok=True)

    image_name = label_studio_relative_path.split('/')[-1]
    output_file = image_name.replace('.jpg', ".json").replace('.jpeg', ".json")
    output_path = os.path.join(output_folder, output_file)

    original_width, original_height = yolo_prediction.orig_shape
    class_names = yolo_prediction.names
    boxes_filtered = filter(lambda x: float(x.conf[0]) >= confidence_threshold
                            , yolo_prediction.boxes)

    if not os.path.exists(output_path) or overwrite_existing_files:
        json_content = {
            "data": {
                "image": label_studio_relative_path,
            },
            "predictions": [
                {
                    "model_version": model_version,
                    "result": [
                        {
                            "original_width": original_width,
                            "original_height": original_height,
                            "score": float(box.conf[0]),
                            "value": {
                                # label studio x, y, width and height are in percent multiplied by 100
                                # x and y are the position of the box top left corner
                                "x": float(box.xywhn[0][0] - box.xywhn[0][2] / 2) * 100,
                                "y": float(box.xywhn[0][1] - box.xywhn[0][3] / 2) * 100,
                                "width": float(box.xywhn[0][2]) * 100,
                                "height": float(box.xywhn[0][3]) * 100,
                                "rotation": 0,
                                "rectanglelabels": [class_names[int(box.cls)]]
                            },
                            "id": str(uuid4()),
                            "from_name": label_studio_from_name,
                            "to_name": label_studio_to_name,
                            "type": "rectanglelabels",
                        } for box in boxes_filtered],
                }]
        }

        with open(output_path, 'w') as fp:
            json.dump(json_content, fp)

        return json_content