import os
import pandas as pd
from ultralytics import YOLO

def run_inference(model, image_path, conf_threshold, save_path):
    test_images = os.listdir(image_path)
    test_images.sort()
    
    bboxes = []
    category_ids = []
    test_images_names = []
    
    # Iterate through images for inference
    for image_name in test_images:
        # Skip any non-image files if they exist in the directory
        if not image_name.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp')):
            continue

        full_image_path = os.path.join(image_path, image_name)
        
        current_image_bboxes = []
        current_image_category_ids = []
        
        # Perform inference
        results = model(full_image_path)
        
        # Process results
        # results is a list of Results objects, one for each image
        # Since we pass one image at a time, results[0] is the relevant object
        for pred in results[0].boxes:
            # Bounding box in xyxy format, confidence, class_id
            xmin, ymin, xmax, ymax = pred.xyxy[0].tolist()
            conf = pred.conf.item()
            class_id = int(pred.cls.item())

            if conf >= conf_threshold:
                width = xmax - xmin
                height = ymax - ymin

                current_image_bboxes.append([xmin, ymin, width, height])
                current_image_category_ids.append(class_id)
            
        test_images_names.append(image_name)
        bboxes.append(current_image_bboxes)
        category_ids.append(current_image_category_ids)
    
    # Create DataFrame for predictions
    df_predictions = pd.DataFrame(columns=["file_name", "bbox", "category_id"])
    
    for i in range(len(test_images_names)):
        file_name = test_images_names[i]
        new_row = pd.DataFrame({"file_name": file_name,
                                "bbox": str(bboxes[i]),
                                "category_id": str(category_ids[i]),
                                }, index=[0])
        df_predictions = pd.concat([df_predictions, new_row], ignore_index=True)
        
    # Ensure the save directory exists
    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    df_predictions.to_csv(save_path, index=False)
    print(f"Inference results saved to: {save_path}")


if __name__ == "__main__":
    # Define paths
    # You might need to change TEST_IMAGE_PATH to where your actual test images are stored
    current_directory = os.path.dirname(os.path.abspath(__file__))
    TEST_IMAGE_PATH = "/tmp/data/test_images"
    SUBMISSION_SAVE_PATH = os.path.join(current_directory, "submission.csv") 
    
    # Path to your trained model weights
    MODEL_WEIGHTS_PATH = os.path.join(current_directory, "best.pt")
    CONF_THRESHOLD = 0.30 # Confidence threshold for predictions
    
    # Load the YOLO model
    model = YOLO(MODEL_WEIGHTS_PATH) # Using ultralytics.YOLO for loading

    # Run inference
    run_inference(model, TEST_IMAGE_PATH, CONF_THRESHOLD, SUBMISSION_SAVE_PATH)