import os
import cv2
import numpy as np
from glob import glob
from tqdm import tqdm

# 1. INPUT: Your seed data (Original 31 images)
INPUT_FOLDER = r"C:\Users\charu\Desktop\seed_data"

# 2. OUTPUT: Where to save the checked images
OUTPUT_FOLDER = r"C:\Users\charu\Desktop\seed_verification_output"

def verify_and_save():
    if not os.path.exists(OUTPUT_FOLDER):
        os.makedirs(OUTPUT_FOLDER)
        
    print(f"🕵️ Checking data in: {INPUT_FOLDER}")
    txt_files = glob(os.path.join(INPUT_FOLDER, "*.txt"))
    
    if not txt_files:
        print("❌ No text files found! Check your path.")
        return

    print(f"   -> Found {len(txt_files)} files. Drawing polygons...")

    for txt_path in tqdm(txt_files):
        # Load Image
        filename = os.path.basename(txt_path).replace('.txt', '')
        img_path_jpg = os.path.join(INPUT_FOLDER, filename + ".jpg")
        img_path_png = os.path.join(INPUT_FOLDER, filename + ".png")
        img_path_jpeg = os.path.join(INPUT_FOLDER, filename + ".jpeg")
        
        if os.path.exists(img_path_jpg): img_path = img_path_jpg
        elif os.path.exists(img_path_png): img_path = img_path_png
        elif os.path.exists(img_path_jpeg): img_path = img_path_jpeg
        else:
            print(f"⚠️ Image missing for {filename}")
            continue

        img = cv2.imread(img_path)
        if img is None: continue
        h, w, _ = img.shape
        
        # Load Labels
        with open(txt_path, 'r') as f:
            lines = f.readlines()

        for line in lines:
            parts = line.strip().split()
            if not parts: continue
            
            try:
                cls_id = int(parts[0])
                
                # Parse Polygons
                coords = [float(x) for x in parts[1:]]
                
                # Check format (Is it Box or Polygon?)
                if len(coords) == 4:
                    # It is a BOX (x, y, w, h) -> Draw Rectangle
                    cx, cy, bw, bh = coords
                    x1 = int((cx - bw/2) * w)
                    y1 = int((cy - bh/2) * h)
                    x2 = int((cx + bw/2) * w)
                    y2 = int((cy + bh/2) * h)
                    cv2.rectangle(img, (x1, y1), (x2, y2), (0, 0, 255), 3)
                    label_text = "BOX"
                    pt = (x1, y1 - 10)
                else:
                    # It is a POLYGON -> Draw Shape
                    points = []
                    for i in range(0, len(coords), 2):
                        px = int(coords[i] * w)
                        py = int(coords[i+1] * h)
                        points.append([px, py])
                    
                    pts = np.array(points, np.int32).reshape((-1, 1, 2))
                    cv2.polylines(img, [pts], isClosed=True, color=(0, 255, 0), thickness=2)
                    label_text = "POLY"
                    pt = tuple(points[0])

                # Add Text
                cv2.putText(img, f"{cls_id}-{label_text}", pt, cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
            except Exception as e:
                print(f"Error parsing {filename}: {e}")

        # Save to Output Folder
        save_path = os.path.join(OUTPUT_FOLDER, filename + "_check.jpg")
        cv2.imwrite(save_path, img)

    print(f"\n✅ Done! Go open this folder and look at the images:")
    print(f"   {OUTPUT_FOLDER}")

if __name__ == "__main__":
    verify_and_save()