import shutil
from typing import Callable

from PIL.Image import Image
from coco_eval import CocoEvaluator
from tqdm import tqdm

from yolo_dataset import YoloDataset, MaterialYoloDataset
from yolo_model import YoloModel

image_loader = Callable[[str], Image]


def coco_evaluate(model: YoloModel, dataset: YoloDataset, confidence_threshold=0.6):
    coco_gt = dataset.to_coco()
    # initialize evaluator with ground truth (gt)
    evaluator = CocoEvaluator(coco_gt=coco_gt, iou_types=["bbox"])

    print("Running evaluation...")
    coco_anns = []

    for image_id, annotations in tqdm(coco_gt.imgToAnns.items()):
        # get the inputs
        image = coco_gt.imgs[image_id]
        results = model.model(source=dataset.load_image(image["file_name"]))
        for result in results:
            coco_anns.extend(yolo_boxes_to_coco_annotations(image_id, result.boxes,
                                                            confidence_threshold=confidence_threshold))

    evaluator.update(coco_anns)
    evaluator.synchronize_between_processes()
    evaluator.accumulate()
    evaluator.summarize()


def yolo_evaluator(model: YoloModel, dataset: YoloDataset):
    return Metrics(model=model, material=dataset.to_material())


class Metrics:
    def __init__(self, model: YoloModel, material: MaterialYoloDataset):
        self.model = model
        self.material = material
        self.val = None

    def __enter__(self):
        if self.val is None:
            self.material.__enter__()
            self.val = self.model.model.val(data=self.material.yaml)
        return self.val

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.material.__exit__(exc_type, exc_val, exc_tb)
        shutil.rmtree(self.val.save_dir)


def yolo_boxes_to_coco_annotations(image_id: int, yolo_boxes, confidence_threshold=0.6):
    return [
        {
            "image_id": int(image_id),
            "category_id": int(box.cls.tolist()[0]),
            "bbox": box.xywh.tolist()[0],
            "score": box.conf.tolist()[0],
        }
        for box in yolo_boxes if box.conf.tolist()[0] > confidence_threshold
    ]


def main():
    yolo_dataset = YoloDataset.from_path('tests/coco8.zip')
    coco_gt = yolo_dataset.to_coco()
    model = YoloModel("ultralyticsplus/yolov8s", "yolov8s.pt")
    # model = YoloModel("SHOU-ISD/fire-and-smoke", "yolov8n.pt")
    # evaluate(model=model, coco_gt=coco_gt, loader=yolo_dataset.load_image, confidence_threshold=0.1)
    # Validate the model
    with yolo_dataset.to_material() as material:
        metrics = model.model.val(data=material.yaml)
        print(metrics.box.map)  # map50-95
        print(metrics.box.map50)  # map50
        print(metrics.box.map75)  # map75
        print(metrics.box.maps)  # a list contains map50-95 of each category


if __name__ == '__main__':
    main()