"""

PCB_metrics_func.py

"""
import numpy as np
from mrcnn import model as modellib, utils

class MetFunc():
  def __init__(self):
    pass

  # Function computes average metrics for the dataset (without details)
  def compute_metrics(self, dataset, in_model, print_enabled=True, img_skip=0):
    self.dataset = dataset
    self.in_model = in_model
    self.img_skip = img_skip

    APs_IoU_25 = []
    APs_IoU_50 = []
    APs_IoU_75 = []
    F1s = []
    F2s = []

    counter = 0

    # Loop through validation dataset
    for image_id in dataset.image_ids.copy():
            
            counter += 1

            # Performance - compute metrics only for every Xth image of the dataset
            if counter%self.img_skip != 0 and self.img_skip != 0:
              continue

            if print_enabled:
              print("Computing metrics for " + str(counter) + ". image...")

            # Load image
            image, image_meta, gt_class_id, gt_bbox, gt_mask = \
                modellib.load_image_gt(dataset,
                              in_model.config,
                              image_id,
                              use_mini_mask=False)

            molded_images = np.expand_dims(modellib.mold_image(image, in_model.config), 0)

            # Run detection
            results = in_model.detect(molded_images, verbose=0)
            r = results[0]

            # Determine APs for different IoUs
            # Compute IoU 0.25

            ap, precisions, recalls, overlaps = \
                utils.compute_ap(gt_bbox, gt_class_id, gt_mask,
                                r["rois"], r["class_ids"], r["scores"], r['masks'], iou_threshold=0.25)
            if print_enabled:
              print("AP IoU25:\t" + str(round(ap, 2)))
            APs_IoU_25.append(ap)

            # Compute IoU 0.75

            ap, precisions, recalls, overlaps = \
                utils.compute_ap(gt_bbox, gt_class_id, gt_mask,
                                r["rois"], r["class_ids"], r["scores"], r['masks'], iou_threshold=0.75)
            if print_enabled:
              print("AP IoU75:\t"+str(round(ap, 2)))
            APs_IoU_75.append(ap)

            # Compute IoU 0.50

            ap, precisions, recalls, overlaps = \
                utils.compute_ap(gt_bbox, gt_class_id, gt_mask,
                                r["rois"], r["class_ids"], r["scores"], r['masks'], iou_threshold=0.5)
            if print_enabled:
              print("AP IoU50:\t" + str(round(ap, 2)))
            APs_IoU_50.append(ap)

            # Determine F1 and F2 score ( using IoU 0.5 )

            F1 = np.mean((2 * precisions * recalls) / (precisions + recalls))
            F2 = np.mean((5 * (precisions * recalls)) / (5 * precisions + recalls))
            
            # Fix nans
            if np.isnan(F1):
              F1 = 0
            if np.isnan(F2):
              F2 = 0

            if print_enabled:
              print("Image F1 score:\t", round(F1, 2))
              print("Image F2 score:\t", round(F2, 2))
            F1s.append(F1)
            F2s.append(F2)


    # Average value of all images

    mAP_IoU_25  = np.mean(APs_IoU_25)
    mAP_IoU_50  = np.mean(APs_IoU_50)
    mAP_IoU_75  = np.mean(APs_IoU_75)
    mF1         = np.mean(F1s)
    mF2         = np.mean(F2s)

    # Print metrics information

    if print_enabled:
      print("Mean MAP (IoU 25):\t", round(mAP_IoU_25, 2))
      print("Mean MAP (IoU 50):\t", round(mAP_IoU_50, 2))
      print("Mean MAP (IoU 75):\t", round(mAP_IoU_75, 2))
      print("Mean F1 score:\t", round(mF1, 2))
      print("Mean F2 score:\t", round(mF2, 2))

    return round(mAP_IoU_25, 2),round(mAP_IoU_50, 2),round(mAP_IoU_75, 2),round(mF1, 2),round(mF2, 2)

  # Function computes average metrics for the dataset (without details)
  def compute_metrics_details(self, rfile, dataset, in_model, print_enabled=True, img_skip=0):
    self.rf = rfile
    self.dataset = dataset
    self.in_model = in_model
    self.img_skip = img_skip

    APs_IoU_25 = []
    APs_IoU_50 = []
    APs_IoU_75 = []
    F1s = []
    F2s = []

    counter = 0

    # Loop through validation dataset
    for image_id in dataset.image_ids.copy():
            
            counter += 1

            # Performance - compute metrics only for every Xth image of the dataset
            if self.img_skip !=0:
              if counter%self.img_skip != 0:
                continue

            if print_enabled:
              print("Computing metrics for " + str(counter) + ". image...")

            # Load image
            image, image_meta, gt_class_id, gt_bbox, gt_mask = \
                modellib.load_image_gt(dataset,
                              in_model.config,
                              image_id,
                              use_mini_mask=False)

            molded_images = np.expand_dims(modellib.mold_image(image, in_model.config), 0)

            # Run detection
            results = in_model.detect(molded_images, verbose=0)
            r = results[0]

            # Determine APs for different IoUs
            # Compute IoU 0.25

            ap, precisions, recalls, overlaps = \
                utils.compute_ap(gt_bbox, gt_class_id, gt_mask,
                                r["rois"], r["class_ids"], r["scores"], r['masks'], iou_threshold=0.25)
            if print_enabled:
              print("AP IoU25:\t" + str(round(ap, 2)))
            APs_IoU_25.append(ap)
            ap_25 = ap

            # Compute IoU 0.75

            ap, precisions, recalls, overlaps = \
                utils.compute_ap(gt_bbox, gt_class_id, gt_mask,
                                r["rois"], r["class_ids"], r["scores"], r['masks'], iou_threshold=0.75)
            if print_enabled:
              print("AP IoU75:\t"+str(round(ap, 2)))
            APs_IoU_75.append(ap)
            ap_75 = ap

            # Compute IoU 0.50

            ap, precisions, recalls, overlaps = \
                utils.compute_ap(gt_bbox, gt_class_id, gt_mask,
                                r["rois"], r["class_ids"], r["scores"], r['masks'], iou_threshold=0.5)
            if print_enabled:
              print("AP IoU50:\t" + str(round(ap, 2)))
            APs_IoU_50.append(ap)
            ap_50 = ap

            # Determine F1 and F2 score ( using IoU 0.5 )

            F1 = np.mean((2 * precisions * recalls) / (precisions + recalls))
            F2 = np.mean((5 * (precisions * recalls)) / (5 * precisions + recalls))
            
            # Fix nans
            if np.isnan(F1):
              F1 = 0
            if np.isnan(F2):
              F2 = 0

            if print_enabled:
              print("Image F1 score:\t", round(F1, 2))
              print("Image F2 score:\t", round(F2, 2))
            F1s.append(F1)
            F2s.append(F2)

            self.rf.add_to_metrics_table(path=str(image_id+1),
                                         AP_25=round(ap_25, 2),
                                         AP_50=round(ap_50, 2),
                                         AP_75=round(ap_75, 2),
                                         F1=round(F1, 2),
                                         F2=round(F2, 2))
            self.rf.add_to_labels(str(image_id+1))


    # Average value of all images

    mAP_IoU_25  = np.mean(APs_IoU_25)
    mAP_IoU_50  = np.mean(APs_IoU_50)
    mAP_IoU_75  = np.mean(APs_IoU_75)
    mF1         = np.mean(F1s)
    mF2         = np.mean(F2s)

    # Print metrics information

    if print_enabled:
      print("Mean MAP (IoU 25):\t", round(mAP_IoU_25, 2))
      print("Mean MAP (IoU 50):\t", round(mAP_IoU_50, 2))
      print("Mean MAP (IoU 75):\t", round(mAP_IoU_75, 2))
      print("Mean F1 score:\t", round(mF1, 2))
      print("Mean F2 score:\t", round(mF2, 2))

    return round(mAP_IoU_25, 2),round(mAP_IoU_50, 2),round(mAP_IoU_75, 2),round(mF1, 2),round(mF2, 2)
