from ultralytics import YOLO
import os
import shutil
from PIL import Image

# --- Configuration ---
# Load your trained model
model = YOLO('./best.pt') # Adjust 'train8' if needed

# Directory with new, unannotated images (your input)
unlabeled_image_dir = './untrained_images'

# Base directory for YOLOv8's output. YOLOv8 will create 'predictX' folders inside this.
# Example: /TempReview/predict, /TempReview/predict1, etc.
yolov8_project_dir = './.yolo_yemp' # This is where your 'predictX' folders are generated

# The final, flat directory where images and labels will be moved for LabelImg
flat_output_dir = './for_labelimg_review'

# --- Ensure input directory exists ---
os.makedirs(unlabeled_image_dir, exist_ok=True) # Make sure this exists if you haven't uploaded images yet

# Get all image files from the unlabeled directory
image_files = [f for f in os.listdir(unlabeled_image_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png'))]

if not image_files:
    print(f"No image files found in '{unlabeled_image_dir}'. Please upload images there.")
else:
    print(f"--- Step 1: Running YOLOv8 Inference on new images ---")
    print(f"Running inference on {len(image_files)} new images...")

    # Run inference on all images in the directory
    # YOLOv8 will automatically create a new 'predictX' folder (e.g., predict, predict1, predict2)
    # inside the yolov8_project_dir for each unique run.
    results_list = model(unlabeled_image_dir,
                         #save=True,          # Save images with plotted boxes
                         save_txt=True,      # Save the YOLO .txt label files
                         conf=0.6,           # Confidence threshold for predictions (adjust as needed)
                         project=yolov8_project_dir # This is your 'TempReview'
                        )

    # Find the most recently created 'predictX' folder
    # This assumes the latest run will be the one with the highest number
    # or the one most recently modified.
    predict_dirs = [d for d in os.listdir(yolov8_project_dir) if d.startswith('predict') and os.path.isdir(os.path.join(yolov8_project_dir, d))]
    
    if not predict_dirs:
        print(f"Error: No 'predictX' folders found in '{yolov8_project_dir}'. Inference might have failed.")
    else:
        # Sort by creation time (most recent first) or by name (highest number)
        # Sorting by name (predict, predict1, predict10, predict2...) doesn't always work numerically.
        # Sorting by modification time is safer.
        latest_predict_dir_name = max(predict_dirs, key=lambda d: os.path.getmtime(os.path.join(yolov8_project_dir, d)))
        
        yolov8_run_path = os.path.join(yolov8_project_dir, latest_predict_dir_name)
        yolov8_images_path = yolov8_run_path
        yolov8_labels_path = os.path.join(yolov8_run_path, 'labels')

        print(f"YOLOv8 results saved to: '{yolov8_run_path}'")

        print(f"\n--- Step 2: Flattening output structure for LabelImg ---")
        os.makedirs(flat_output_dir, exist_ok=True)

        # Move images
        if os.path.exists(yolov8_images_path):
            for img_file in os.listdir(yolov8_images_path):
                # Only move files that are original image files (e.g., .jpg, .png)
                if img_file.lower().endswith(('.jpg', '.jpeg', '.png')):
                    shutil.move(os.path.join(yolov8_images_path, img_file),
                                os.path.join(flat_output_dir, img_file))
            print(f"Moved images to '{flat_output_dir}'")
        else:
            print(f"Warning: No images found at '{yolov8_images_path}'.")

        # Process and move labels
        if os.path.exists(yolov8_labels_path):
            for label_file in os.listdir(yolov8_labels_path):
                if label_file.lower().endswith('.txt'):
                    label_path_src = os.path.join(yolov8_labels_path, label_file)
                    
                    # Read lines, parse class_id, and sort
                    with open(label_path_src, 'r') as f:
                        lines = f.readlines()
                    
                    # Sort lines based on the first element (class_id) as an integer
                    # Handle potential errors if a line is malformed, though unlikely from YOLO.
                    try:
                        sorted_lines = sorted(lines, key=lambda line: int(line.strip().split(' ')[0]))
                    except ValueError as e:
                        print(f"Warning: Could not sort lines in {label_file} due to format error: {e}. Skipping sort for this file.")
                        sorted_lines = lines # Fallback to original order
                    
                    # Write sorted lines to the destination file
                    label_path_dest = os.path.join(flat_output_dir, label_file)
                    with open(label_path_dest, 'w') as f:
                        f.writelines(sorted_lines)
                    
                    # Remove the original label file after processing (optional, but keeps source clean)
                    os.remove(label_path_src)
            print(f"Moved and sorted labels to '{flat_output_dir}'")
        else:
            print(f"Warning: No labels found at '{yolov8_labels_path}'.")

        # Clean up the intermediate YOLOv8 output directory (optional, but recommended for Colab space)
        # Be careful with shutil.rmtree - it deletes recursively!
        if os.path.exists(yolov8_run_path):
            print(f"Cleaning up temporary YOLOv8 output: '{yolov8_run_path}'")
            shutil.rmtree(yolov8_run_path)
            
        print(f"\n--- Process Complete! ---")
        print(f"Your images and predicted .txt labels are now in a flat structure in: '{flat_output_dir}'")
        print("You can now open LabelImg and point it to this directory to begin review and correction.")