# evaluation/vizwiz_localization_eval.py
"""
VizWiz Object Localization Evaluation Module
Adapted from EvalAI vqaEval.py for VizWiz Object Localization Challenge
"""

import json
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval


class Localization:
    """
    Original Localization class from EvalAI
    """
    def __init__(self, eval_bbox=True, eval_segm=True):
        self.eval_bbox = eval_bbox
        self.eval_segm = eval_segm
        print("Created Localization Object")
    
    def compute_coco_metrics(self, gt_annotations, sub_annotations):
        """
        Compute COCO metrics for object detection/segmentation
        
        Args:
            gt_annotations (str): Path to ground truth annotations JSON
            sub_annotations (str): Path to submission annotations JSON
            
        Returns:
            tuple: (bbox_mAP, bbox_AP50, segm_mAP, segm_AP50)
        """
        print("In compute_coco_metrics")
        # Load ground truth and submission annotations
        coco_gt = COCO(gt_annotations)
        coco_sub = coco_gt.loadRes(sub_annotations)
        
        bbox_mAP = 0.0
        bbox_AP50 = 0.0
        segm_mAP = 0.0
        segm_AP50 = 0.0
        
        if self.eval_bbox:
            # Object Detection Evaluation
            coco_eval_bbox = COCOeval(coco_gt, coco_sub, iouType='bbox')
            coco_eval_bbox.evaluate()
            coco_eval_bbox.accumulate()
            coco_eval_bbox.summarize()
            
            # CORRECTED: stats[0] is mAP, stats[1] is AP50
            bbox_mAP = coco_eval_bbox.stats[0]
            bbox_AP50 = coco_eval_bbox.stats[1]
        
        if self.eval_segm:
            # Instance Segmentation Evaluation
            coco_eval_mask = COCOeval(coco_gt, coco_sub, iouType='segm')
            coco_eval_mask.evaluate()
            coco_eval_mask.accumulate()
            coco_eval_mask.summarize()
            
            segm_mAP = coco_eval_mask.stats[0]
            segm_AP50 = coco_eval_mask.stats[1]
        
        return bbox_mAP, bbox_AP50, segm_mAP, segm_AP50


def evaluate_submission(ground_truth_path, predictions_path):
    """
    Evaluate a submission and return score compatible with UI code
    
    Args:
        ground_truth_path (str): Path to ground truth JSON
        predictions_path (str): Path to predictions JSON
        
    Returns:
        float or None: Score as percentage (0-100), or None if error
    """
    try:
        # Basic validation
        with open(predictions_path, 'r') as f:
            predictions = json.load(f)
        
        print("Loaded Predictions file")
        
        if not isinstance(predictions, list) or len(predictions) == 0:
            return None
        
        # Check required fields
        required_fields = ['image_id', 'category_id', 'bbox', 'score']
        if not all(field in predictions[0] for field in required_fields):
            return None
        print(f"Required Fields Checked")
        
        # Run evaluation
        evaluator = Localization()
        print("Evaluation Object Instantiated")
        bbox_mAP, bbox_AP50, segm_mAP, segm_AP50 = evaluator.compute_coco_metrics(
            ground_truth_path,
            predictions_path
        )
        print(f"bbox score calculated : {bbox_mAP}")
        
        return float(bbox_mAP), float(bbox_AP50), segm_mAP, segm_AP50
        
    except Exception as e:
        print(f"Evaluation error: {e}")
        return None