import json
import numpy as np
from pycocotools.coco import COCO
from pydef main():
    import os
    
    # Path to your prediction file - using absolute path for safety
    pred_file = os.path.abspath("outputs_standing/inference/coco_instances_results.json")
    # Path to your ground truth COCO annotation file
    gt_file = os.path.abspath("data/test/annotations.json")  # Update this path
    
    # Check if files exist
    if not os.path.exists(pred_file):
        print(f"Error: Prediction file not found: {pred_file}")
        return
    if not os.path.exists(gt_file):
        print(f"Error: Ground truth file not found: {gt_file}")
        return
        
    print(f"Loading predictions from: {pred_file}")
    print(f"Loading ground truth from: {gt_file}")
    
    # Create analysis directory
    analysis_dir = os.path.abspath("analysis_results")
    os.makedirs(analysis_dir, exist_ok=True)
    
    # Load predictions and ground truth
    try:
        predictions = load_predictions(pred_file)
        print(f"Successfully loaded {len(predictions)} predictions")
        coco_gt = COCO(gt_file)
        print(f"Successfully loaded ground truth annotations").cocoeval import COCOeval
import pandas as pd
from collections import defaultdict

def load_predictions(pred_file):
    with open(pred_file, 'r') as f:
        predictions = json.load(f)
    return predictions

def filter_annotations_by_attributes(coco_gt, annotations, attribute_filter):
    """
    Filter annotations based on specific attributes
    attribute_filter: dict with keys as attribute names and values as required values
    """
    filtered_anns = []
    for ann in annotations:
        img_info = coco_gt.imgs[ann['image_id']]
        matches_all = True
        for attr, value in attribute_filter.items():
            if attr not in img_info or img_info[attr] != value:
                matches_all = False
                break
        if matches_all:
            filtered_anns.append(ann)
    return filtered_anns

def evaluate_subset(coco_gt, predictions, attribute_filter):
    """
    Evaluate predictions on a subset of data with specific attributes
    """
    # Create a new COCO object with filtered annotations
    filtered_anns = filter_annotations_by_attributes(coco_gt, coco_gt.dataset['annotations'], attribute_filter)
    if not filtered_anns:
        return None
    
    # Create a subset of the dataset with filtered annotations
    subset_dataset = {
        'info': coco_gt.dataset['info'],
        'licenses': coco_gt.dataset['licenses'],
        'images': coco_gt.dataset['images'],
        'categories': coco_gt.dataset['categories'],
        'annotations': filtered_anns
    }
    
    # Create temporary file for filtered ground truth
    coco_subset = COCO()
    coco_subset.dataset = subset_dataset
    coco_subset.createIndex()
    
    # Filter predictions to only include images in the subset
    valid_img_ids = set(img['id'] for img in subset_dataset['images'])
    filtered_preds = [p for p in predictions if p['image_id'] in valid_img_ids]
    
    if not filtered_preds:
        return None
    
    # Run evaluation
    coco_eval = COCOeval(coco_subset, coco_subset.loadRes(filtered_preds), 'bbox')
    coco_eval.evaluate()
    coco_eval.accumulate()
    coco_eval.summarize()
    
    # Return AP values
    return {
        'AP': coco_eval.stats[0],  # AP @ IoU=0.50:0.95
        'AP50': coco_eval.stats[1],  # AP @ IoU=0.50
        'AP75': coco_eval.stats[2],  # AP @ IoU=0.75
    }

def main():
    # Path to your prediction file
    pred_file = "outputs_standing/inference/coco_instances_results.json"
    # Path to your ground truth COCO annotation file
    gt_file = "/mnt/home2/annonymous/ForestPersons_v3/annotations/test.json"  # Update this path
    
    # Load predictions and ground truth
    predictions = load_predictions(pred_file)
    coco_gt = COCO(gt_file)
    
    # Define attribute combinations to evaluate
    pose_types = ['standing', 'sitting', 'lying']
    seasons = ['summer', 'fall', 'winter']
    visibility_levels = [20, 40, 70, 100]
    
    results = []
    
    # Evaluate by pose type
    for pose in pose_types:
        result = evaluate_subset(coco_gt, predictions, {'pose': pose})
        if result:
            results.append({
                'Category': pose,
                'Type': 'Pose',
                **result
            })
    
    # Evaluate by season
    for season in seasons:
        result = evaluate_subset(coco_gt, predictions, {'season': season})
        if result:
            results.append({
                'Category': season,
                'Type': 'Season',
                **result
            })
    
    # Evaluate by pose and visibility
    for pose in pose_types:
        for vis in visibility_levels:
            result = evaluate_subset(coco_gt, predictions, {
                'pose': pose,
                'visibility': vis
            })
            if result:
                results.append({
                    'Category': f'{pose}_vis{vis}',
                    'Type': 'Pose+Visibility',
                    **result
                })
    
    # Convert results to DataFrame and save
    df = pd.DataFrame(results)
    print("\nEvaluation Results:")
    print(df.to_string(index=False))
    
    # Save results to CSV
    output_file = pred_file.replace('.json', '_analysis.csv')
    df.to_csv(output_file, index=False)
    print(f"\nResults saved to: {output_file}")

if __name__ == "__main__":
    main()