# Label_v4.py
# ---------------------


import glob
import os
import random
import cv2

# --- CONFIGURATION ---
OUTPUT_DIRECTORY = "dataset_v4"
FRAMES_TO_EXTRACT = 20  # Number of random anchor frames to find targets
SCALE_FACTOR = 0.5  # Resize video to 1/2 size for easier viewing


# Defined classification options (YOLO maps these to IDs: 0, 1, 2, 3...)
CLASSIFICATION_OPTIONS = [
    "guardian",
    "interceptor",
    "sentinel",
    "rocket",
    "bomber",
]
# ---------------------


def get_dataset_counts():
    """Counts the number of existing bounding boxes by parsing YOLO text files."""
    counts = {option: 0 for option in CLASSIFICATION_OPTIONS}
    labels_dir = os.path.join(OUTPUT_DIRECTORY, "labels")

    if os.path.exists(labels_dir):
        # Scan all text annotation files in the YOLO labels directory
        for txt_file in glob.glob(os.path.join(labels_dir, "*.txt")):
            try:
                with open(txt_file, "r") as f:
                    for line in f:
                        parts = line.strip().split()
                        if parts:
                            class_id = int(parts[0])
                            if 0 <= class_id < len(CLASSIFICATION_OPTIONS):
                                option = CLASSIFICATION_OPTIONS[class_id]
                                counts[option] += 1
            except Exception:
                pass  # Skip corrupted or unreadable text files safely
    return counts

def redraw_frame(param):
    """Regenerates the display frame using the current list of annotations."""
    # Start with a fresh, clean copy of the frame
    display_frame = param["frame"].copy()
    
    # Draw all active boxes
    for ann in param["annotations"]:
        class_id, label, x1, y1, x2, y2 = ann
        cv2.rectangle(display_frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cv2.putText(
            display_frame,
            label,
            (x1, max(15, y1 - 5)),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.4,
            (0, 255, 0),
            1,
        )
    
    # Update state and window
    param["display_frame"] = display_frame
    cv2.imshow(param["window_name"], display_frame)

def save_annotations_to_file(param):
    """Overwrites the YOLO text file and updates the image file on disk."""
    video_base = os.path.splitext(os.path.basename(param["video_name"]))[0]
    filename_base = f"{video_base}_frame_{param['frame_idx']}"

    image_save_path = os.path.join(param["images_dir"], f"{filename_base}.jpg")
    label_save_path = os.path.join(param["labels_dir"], f"{filename_base}.txt")

    # If all annotations were undone, delete the files
    if len(param["annotations"]) == 0:
        if os.path.exists(label_save_path):
            os.remove(label_save_path)
        if os.path.exists(image_save_path):
            os.remove(image_save_path)
        print(f" Removed empty annotation files for {filename_base}")
        param["box_drawn"] = False
    else:
        # Save the frame image
        cv2.imwrite(image_save_path, param["frame"])
        
        # Write/Overwrite the text file
        img_h, img_w = param["frame"].shape[0], param["frame"].shape[1]
        with open(label_save_path, "w") as f:
            for ann in param["annotations"]:
                class_id, _, x1, y1, x2, y2 = ann
                
                # YOLO Normalized Coordinate calculations
                true_center_x = (x1 + x2) / 2.0
                true_center_y = (y1 + y2) / 2.0
                width = x2 - x1
                height = y2 - y1

                x_center_norm = true_center_x / img_w
                y_center_norm = true_center_y / img_h
                width_norm = width / img_w
                height_norm = height / img_h

                f.write(
                    f"{class_id} {x_center_norm:.6f} {y_center_norm:.6f} {width_norm:.6f} {height_norm:.6f}\n"
                )
        print(f" Updated dataset files for {filename_base}")
        param["box_drawn"] = True

def mouse_click_callback(event, x, y, flags, param):
    """Handles click-and-drag bounding box creation and right-click undo actions."""
    if param["session_label"] == "background":
        if event == cv2.EVENT_LBUTTONDOWN:
            print(
                " Info: In background mode, do not click targets. Just press SPACE to save frame."
            )
        return

    frame = param["frame"]
    display_frame = param["display_frame"]
    window_name = param["window_name"]

    # --- Right Click: Undo Last Box ---
    if event == cv2.EVENT_RBUTTONDOWN:
        if param["annotations"]:
            removed = param["annotations"].pop()
            print(f" Undid last box ({removed[1]})")
            save_annotations_to_file(param)
            redraw_frame(param)
        else:
            print(" No boxes left to undo on this frame!")
        return

    # --- Left Click Down: Start Dragging ---
    if event == cv2.EVENT_LBUTTONDOWN:
        param["drawing"] = True
        param["ix"], param["iy"] = x, y

    # --- Mouse Move: Live Preview ---
    elif event == cv2.EVENT_MOUSEMOVE:
        if param["drawing"]:
            # Wipe frame clean to draw fresh temporary drag lines
            temp_frame = display_frame.copy()
            cv2.rectangle(temp_frame, (param["ix"], param["iy"]), (x, y), (0, 255, 0), 2)
            cv2.imshow(window_name, temp_frame)

    # --- Left Click Release: Lock Box & Save ---
    elif event == cv2.EVENT_LBUTTONUP:
        param["drawing"] = False
        ix, iy = param["ix"], param["iy"]

        # Handle drawing in any physical direction
        x1, x2 = min(ix, x), max(ix, x)
        y1, y2 = min(iy, y), max(iy, y)

        # Ignore tiny accidental micro-clicks
        width = x2 - x1
        height = y2 - y1
        if width < 5 or height < 5:
            cv2.imshow(window_name, display_frame)
            return

        # Keep box coordinates within image boundaries
        img_h, img_w = frame.shape[0], frame.shape[1]
        x1 = max(0, min(x1, img_w - 1))
        y1 = max(0, min(y1, img_h - 1))
        x2 = max(0, min(x2, img_w - 1))
        y2 = max(0, min(y2, img_h - 1))

        # Store the coordinates in our dynamic box list
        param["annotations"].append((
            param["class_id"],
            param["session_label"],
            x1, y1, x2, y2
        ))

        # Save updates to disk and redraw
        save_annotations_to_file(param)
        redraw_frame(param)

def label_video_session():
    print("=========================================")
    print("         AI YOLO DATASET LABELER          ")
    print("=========================================\n")

    # 1. Automatically find and select a random MP4 video
    mp4_files = glob.glob("*.mp4")
    if not mp4_files:
        print(
            "Error: No .mp4 files found in the current directory.\n"
            "Please place this script in the same folder as your video files."
        )
        return

    selected_video = random.choice(mp4_files)
    print(f"-> Selected Video File: '{selected_video}'")

    # 2. Fetch and display current dataset balance dashboard
    counts = get_dataset_counts()
    print("\n-----------------------------------------")
    print("      CURRENT DATASET BALANCE STATUS     ")
    print("-----------------------------------------")
    for option, count in counts.items():
        print(f"  • {option.capitalize():<12} : {count} boxes logged")
    print("-----------------------------------------")
    print(
        "💡 Tip: Try to choose options with lower counts to keep data balanced!"
    )

    # 3. Force selection of a valid label
    print("\nAvailable Classification Targets:")
    for idx, option in enumerate(CLASSIFICATION_OPTIONS, start=1):
        print(f"  [{idx}] {option} (Current Boxes: {counts[option]})")

    while True:
        try:
            choice = input(
                f"\nEnter the number (1-{len(CLASSIFICATION_OPTIONS)}) of what you are labeling: "
            ).strip()
            choice_idx = int(choice) - 1
            if 0 <= choice_idx < len(CLASSIFICATION_OPTIONS):
                session_label = CLASSIFICATION_OPTIONS[choice_idx]
                session_class_id = choice_idx
                break
            else:
                print("Invalid selection. Please choose a number from the list.")
        except ValueError:
            print("Invalid input. Please enter a number.")

    # Open video and check validity
    cap = cv2.VideoCapture(selected_video)
    if not cap.isOpened():
        print(f"Error: Could not open video file {selected_video}")
        return

    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    if total_frames == 0:
        print("Error: Video has 0 frames.")
        return

    # Select random unique anchor frames to jump between
    num_frames = min(FRAMES_TO_EXTRACT, total_frames)
    frame_indices = random.sample(range(total_frames), num_frames)
    frame_indices.sort()

    # Set up standard YOLO dataset directory branches
    images_dir = os.path.join(OUTPUT_DIRECTORY, "images")
    labels_dir = os.path.join(OUTPUT_DIRECTORY, "labels")
    os.makedirs(images_dir, exist_ok=True)
    os.makedirs(labels_dir, exist_ok=True)

    # Comprehensive User Instructions
    print("\n" + "=" * 50)
    print("                HOW TO LABEL")
    print("=" * 50)
    print(f"• TARGET OBJECT: {session_label.upper()}")
    print(f"• SESSION SCOPE: {num_frames} random search zones in '{selected_video}'")
    print("-" * 50)
    print("► CONTROLS:")
    print("  [Click & Drag]   - Draw a custom box around the target (auto-saves).")
    print("  [Right-Click]    - Undo / delete the last drawn box on the current frame.")
    print("  [Right Arrow]    - Move FORWARD 10 frames (sequential tracking, no save).")
    print("  [Left Arrow]     - Move BACKWARD 10 frames (sequential tracking, no save).")
    print("  [SPACE / ENTER]  - Jump completely to the NEXT random video search spot.")
    print("  [Q Key]          - Quit and save all progress up to this point.")
    print("=" * 50)
    input("\nPress ENTER when you are ready to begin...")

    # Cross-platform arrow key mappings for waitKeyEx()
    LEFT_KEYS = [81, 2, 2424832, 65361, 63234]
    RIGHT_KEYS = [83, 3, 2490368, 65363, 63235]

    i = 0
    current_frame = frame_indices[i]

    while i < len(frame_indices):
        cap.set(cv2.CAP_PROP_POS_FRAMES, current_frame)
        ret, frame = cap.read()
        if not ret:
            print(f" Warning: Could not read frame {current_frame}. Moving to next random anchor.")
            i += 1
            if i < len(frame_indices):
                current_frame = frame_indices[i]
            continue

        # Resize the frame by SCALE_FACTOR for laptop viewability
        frame_resized = cv2.resize(
            frame, (0, 0), fx=SCALE_FACTOR, fy=SCALE_FACTOR
        )
        frame_display = frame_resized.copy()

        window_name = f"Labeling: {session_label} | Spot {i+1}/{num_frames} (Frame {current_frame})"
        cv2.namedWindow(window_name)

        # Expanded state tracking for multi-box lists
        session_state = {
            "frame": frame_resized,
            "display_frame": frame_display,
            "images_dir": images_dir,
            "labels_dir": labels_dir,
            "frame_idx": current_frame,
            "window_name": window_name,
            "video_name": selected_video,
            "class_id": session_class_id,
            "session_label": session_label,
            "box_drawn": False,
            "drawing": False,
            "ix": -1,
            "iy": -1,
            "annotations": [],  # Holds tuple info of drawn boxes for live Undo action
        }

        cv2.setMouseCallback(window_name, mouse_click_callback, session_state)

        # Keep window open until user acts
        cv2.imshow(window_name, frame_display)
        while True:
            cv2.moveWindow(window_name, 0, 0)
            key = cv2.waitKeyEx(1)
            
            # Enter (13) or Space (32) -> Jump to the next random anchor spot
            if key == 13 or key == 32:
                if session_label == "background":
                    video_base = os.path.splitext(
                        os.path.basename(selected_video)
                    )[0]
                    filename_base = f"{video_base}_frame_{current_frame}"
                    image_save_path = os.path.join(
                        images_dir, f"{filename_base}.jpg"
                    )
                    label_save_path = os.path.join(
                        labels_dir, f"{filename_base}.txt"
                    )

                    cv2.imwrite(image_save_path, frame_resized)
                    open(label_save_path, "a").close()  # Creates empty file
                    print(f" Registered negative background frame: {filename_base}")
                
                i += 1
                if i < len(frame_indices):
                    current_frame = frame_indices[i]
                break

            # Right Arrow -> Move forward
            elif key in RIGHT_KEYS:
                if current_frame < total_frames - 10:
                    current_frame += 10
                else:
                    print(" Already at the very last frame of the video!")
                    continue
                break

            # Left Arrow -> Move backward
            elif key in LEFT_KEYS:
                if current_frame > 0:
                    current_frame -= 10
                else:
                    print(" Already at the first frame of the video!")
                    continue
                break

            # Quit 
            elif key == ord("q") or key == ord("Q") or (key & 0xFF) in [ord("q"), ord("Q")]:
                print("\nQuitting session early... Saving progress.")
                cap.release()
                cv2.destroyAllWindows()
                return

        cv2.destroyWindow(window_name)

    cap.release()
    cv2.destroyAllWindows()
    print(f"\nSession finished! Dataset updated successfully inside '{OUTPUT_DIRECTORY}'.")

if __name__ == "__main__":
    label_video_session()