# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
# -*- coding: utf-8 -*-
from dataiku.doctor.deephub.object_detection_performance import (
    ObjectDetectionPerformanceComputer,
)
from dataiku.doctor.deephub.deephub_params import TargetRemapping

import json
import pandas as pd
import dataiku

df = dataiku.Dataset("test_scored").get_dataframe()
output_dataset = dataiku.Dataset("evaluation")

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
all_categories = set()
cat2id = {}
for i in range(len(df)):
    for label in json.loads(df["label"].iloc[i]):
        all_categories.add(label["category"])
all_categories = sorted(list(all_categories))
for index, value in enumerate(all_categories):
    cat2id[value] = index

target_remapping = TargetRemapping(
    [
        {"sourceValue": k, "mappedValue": cat2id[k], "sampleFreq": 100}
        for k in all_categories
    ]
)
metric_params = {
    "confidenceScoreThresholdOptimMetric": "f1",
    "evaluationMetric": "AVERAGE_PRECISION_IOU50",
}
origin_index = list(range(len(df)))

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
detections_df = df.drop(labels="label", axis=1)
detections_df["confidence"] = 0.9
detections_df["prediction"] = detections_df["prediction"].map(lambda x: json.loads(x))
detections_df.rename({"record_id": "image_id"}, axis=1, inplace=True)

for i in range(len(detections_df)):
    for j in range(len(detections_df.iloc[i]["prediction"])):
        detections_df.iloc[i]["prediction"][j]["category"] = cat2id[
            detections_df.iloc[i]["prediction"][j]["category"]
        ]

ground_truth_df = df.drop(labels="prediction", axis=1)
ground_truth_df.rename({"label": "target"}, axis=1, inplace=True)
ground_truth_df.rename({"record_id": "image_id"}, axis=1, inplace=True)
ground_truth_df["target"] = ground_truth_df["target"].map(lambda x: json.loads(x))

for i in range(len(ground_truth_df)):
    for j in range(len(ground_truth_df.iloc[i]["target"])):
        ground_truth_df.iloc[i]["target"][j]["category"] = cat2id[
            ground_truth_df.iloc[i]["target"][j]["category"]
        ]

metric_computer = ObjectDetectionPerformanceComputer(
    target_remapping, origin_index, ground_truth_df, detections_df, metric_params
)
result = metric_computer.compute_performance().to_dict()

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
short_label = {
    "averagePrecisionAllIOU": "Average Precision (All IoUs)",
    "averagePrecisionIOU50": "Average Precision (IoU=0.5)",
    "averagePrecisionIOU75": "Average Precision (IoU=0.75)",
}

result_dict = {"metric": [], "category": [], "value": []}

for metric in result["metrics"]:
    result_dict["metric"].append(short_label[metric])
    result_dict["category"].append("All classes")
    result_dict["value"].append(result["metrics"][metric])


for cat in range(len(result["perCategoryMetrics"])):
    for metric in result["perCategoryMetrics"][cat]:
        result_dict["metric"].append(short_label[metric])
        result_dict["category"].append(all_categories[cat])
        result_dict["value"].append(result["perCategoryMetrics"][cat][metric])

output_dataset.write_with_schema(
    pd.DataFrame.from_dict(result_dict),
)