from ultralytics import YOLO
import json
import os

def convert_predictions(gt_json_path, pred_json_path, output_json_path):
    # Load GT data to create filename to ID mapping
    with open(gt_json_path, 'r') as f:
        gt_data = json.load(f)

    filename_to_id = {
        img['file_name'].split("/")[-1].split(".")[0]: img['id']
        for img in gt_data['images']
    }

    # Load predictions
    with open(pred_json_path, 'r') as f:
        pred_data = json.load(f)

    # Convert predictions
    converted = []
    for ann in pred_data:
        file_name = ann['image_id']  # YOLO uses filename here
        if file_name in filename_to_id:
            ann['image_id'] = filename_to_id[file_name]
            converted.append(ann)

    # Save converted predictions
    with open(output_json_path, 'w') as f:
        json.dump(converted, f, indent=2)

def evaluate_model(model_path, output_name):
    print(f"\nEvaluating model: {model_path}")
    
    # Load model
    model = YOLO(model_path)
    
    # Run validation/test with JSON export
    metrics = model.val(data="data/forestpersons.yaml", device=7, save_json=True)
        
    # Print metrics
    print(f"Results for {output_name}:")
    print(f"mAP50-95: {metrics.box.map:.4f}")
    print(f"mAP50: {metrics.box.map50:.4f}")
    print(f"mAP75: {metrics.box.map75:.4f}")

# def main():
#     # Model paths
#     models = {
#         'standing': 'forestpersons_standing/yolov11n/weights/best.pt',
#         'summer': 'forestpersons_summer/yolov11n3/weights/best.pt',
#         'winter': 'forestpersons_winter/yolov11n/weights/best.pt'
#     }
    
#     # Evaluate each model
#     for model_type, model_path in models.items():
#         if os.path.exists(model_path):
#             evaluate_model(model_path, model_type)
#         else:
#             print(f"\nWarning: Model not found at {model_path}")

def main():
    # Convert predictions to required format
    pred_json_paths = ["runs/detect/val22/predictions.json",
                       "runs/detect/val23/predictions.json",
                       "runs/detect/val24/predictions.json",]# Default YOLO output path
    gt_json_path = "data/test.json"
    output_paths = ["converted_standing_predictions.json",
                    "converted_summer_predictions.json",
                    "converted_winter_predictions.json"]
    
    for pred_json_path, output_path in zip(pred_json_paths, output_paths):
        convert_predictions(gt_json_path, pred_json_path, output_path)

if __name__ == "__main__":
    main()