import os
import torch
import pandas as pd
from rfdetr import RFDETRBase


def run_inference(model, image_path, conf_threshold, save_path):

    test_images = sorted([
        f for f in os.listdir(image_path)
        if f.lower().endswith((".jpg", ".jpeg", ".png"))
    ])

    bboxes = []
    category_ids = []
    test_images_names = []

    for image_name in test_images:
        test_images_names.append(image_name)

        image_file = os.path.join(image_path, image_name)

        bbox = []
        category_id = []

        preds = model.predict(image_file)

        if preds is not None and preds.xyxy is not None and len(preds.xyxy) > 0:
            for box, score, label in zip(
                preds.xyxy,
                preds.confidence,
                preds.class_id
            ):
                score = float(score)
                if score >= conf_threshold:
                    xmin, ymin, xmax, ymax = map(float, box)

                    width = xmax - xmin
                    height = ymax - ymin

                    bbox.append([xmin, ymin, width, height])
                    category_id.append(int(label))

        bboxes.append(bbox)
        category_ids.append(category_id)

    df_predictions = pd.DataFrame(columns=["file_name", "bbox", "category_id"])

    for i in range(len(test_images_names)):
        new_row = pd.DataFrame({
            "file_name": test_images_names[i],
            "bbox": str(bboxes[i]),
            "category_id": str(category_ids[i])
        }, index=[0])

        df_predictions = pd.concat([df_predictions, new_row], ignore_index=True)

    df_predictions.to_csv(save_path, index=False)


if __name__ == "__main__":

    TEST_IMAGE_PATH = r"rf-detr\dataset\test"
    SUBMISSION_SAVE_PATH = "submission.csv"
    CONF_THRESHOLD = 0.30

    model = RFDETRBase(
        checkpoint_path="checkpoint_best_total.pth",
        device="cuda" if torch.cuda.is_available() else "cpu"
    )

    run_inference(model, TEST_IMAGE_PATH, CONF_THRESHOLD, SUBMISSION_SAVE_PATH)