import torch
import numpy as np
import fiftyone as fo
# from box_metrics import BoxMetrics
from seametrics.fo_utils.utils import fo_to_payload 
from tqdm import tqdm
import evaluate

tags = []
dataset_name = "SENTRY_VIDEOS_DATASET_QA"
model = "cerulean-level-17_11_2023_RL_SPLIT_ep147_CNN"
det_gt_field = "ground_truth_det"

dataset = fo.load_dataset(dataset_name)
dataset_view = fo.load_dataset(dataset_name).match_tags(tags) if tags else fo.load_dataset(dataset_name)
sequences = dataset_view.distinct("sequence")

# bbox_metric = BoxMetrics(max_iou=0.01)
bbox_metric = evaluate.load("SEA-AI/box-metrics")

for sequence in tqdm(sequences):
    try:
        payload = fo_to_payload(dataset = dataset_name,
                                gt_field = det_gt_field,
                                models = [model],
                                tracking_mode = False,
                                sequence_list = [sequence],
                                group_slices = ["thermal_wide"],
                                excluded_classes = ["SUN_REFLECTIONS"],
                                )
        bbox_metric.add_payload(payload)
    except Exception as e:
        print(f"Error in {sequence}: {e}")

result = bbox_metric.compute()
summary = bbox_metric.summarize(result)


for sequence in result:
    print(f"Sequence: {sequence}")
    for model in result[sequence]:
        print(f"Model: {model}")
        for metric in result[sequence][model]:
            print(f"{metric}: {result[sequence][model][metric]}")
    print("\n")