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

# ================= CONFIGURATION =================
# 1. PATHS (Separated)
IMAGES_DIR = r"C:\Users\charu\Desktop\04-02-2026\images"
LABELS_DIR = r"C:\Users\charu\Desktop\04-02-2026\labels"

# 2. Target Count per class
TARGET_PER_CLASS = 300

# 3. Class Names
CLASS_NAMES = {0: "Blast", 1: "Brown Spot", 2: "Sheath Blight"}
# =================================================

def load_dataset():
    dataset = {0: [], 1: [], 2: []}
    
    # scan labels folder
    txt_files = glob(os.path.join(LABELS_DIR, "*.txt"))
    
    print(f"📂 Scanning Labels: {LABELS_DIR}")
    print(f"   -> Found {len(txt_files)} text files.")
    
    if len(txt_files) == 0:
        print(" Error: No text files found! Check the path.")
        return dataset

    for txt_path in txt_files:
        filename = os.path.basename(txt_path).replace('.txt', '')
        
        # Look for matching image in IMAGES_DIR
        img_jpg = os.path.join(IMAGES_DIR, filename + ".jpg")
        img_png = os.path.join(IMAGES_DIR, filename + ".png")
        img_jpeg = os.path.join(IMAGES_DIR, filename + ".jpeg")
        
        if os.path.exists(img_jpg): img_path = img_jpg
        elif os.path.exists(img_png): img_path = img_png
        elif os.path.exists(img_jpeg): img_path = img_jpeg
        else: 
            # If no image found for this label, skip it
            continue 

        with open(txt_path, 'r') as f:
            lines = f.readlines()
            
        if lines:
            try:
                # Read class ID
                class_id = int(lines[0].split()[0])
                if class_id in dataset:
                    dataset[class_id].append((img_path, lines))
            except:
                pass 
    return dataset

def augment_polygon(img_path, lines, new_filename):
    img = cv2.imread(img_path)
    if img is None: return
    
    action = random.choice(["h_flip", "v_flip", "bright", "noise"])
    new_lines = []
    
    if action == "h_flip":
        new_img = cv2.flip(img, 1)
        for line in lines:
            parts = line.strip().split()
            cls = parts[0]
            coords = [float(x) for x in parts[1:]]
            new_coords = []
            for i, val in enumerate(coords):
                if i % 2 == 0: new_coords.append(1.0 - val) # X
                else:          new_coords.append(val)       # Y
            new_lines.append(f"{cls} " + " ".join([f"{c:.6f}" for c in new_coords]) + "\n")

    elif action == "v_flip":
        new_img = cv2.flip(img, 0)
        for line in lines:
            parts = line.strip().split()
            cls = parts[0]
            coords = [float(x) for x in parts[1:]]
            new_coords = []
            for i, val in enumerate(coords):
                if i % 2 == 0: new_coords.append(val)       # X
                else:          new_coords.append(1.0 - val) # Y
            new_lines.append(f"{cls} " + " ".join([f"{c:.6f}" for c in new_coords]) + "\n")

    elif action == "bright":
        beta = random.randint(-30, 30)
        new_img = cv2.convertScaleAbs(img, alpha=1.0, beta=beta)
        new_lines = lines 

    elif action == "noise":
        noise = np.random.normal(0, 15, img.shape).astype(np.uint8)
        new_img = cv2.add(img, noise)
        new_lines = lines 
        
    else: 
        new_img = img
        new_lines = lines

    # SAVE TO SEPARATE FOLDERS
    cv2.imwrite(os.path.join(IMAGES_DIR, new_filename + ".jpg"), new_img)
    with open(os.path.join(LABELS_DIR, new_filename + ".txt"), 'w') as f:
        f.writelines(new_lines)

def main():
    print("🚀 Loading Dataset (Separated Folders)...")
    data_map = load_dataset()
    
    print("\n📊 Current Counts:")
    for cid in [0, 1, 2]:
        print(f"   - {CLASS_NAMES[cid]}: {len(data_map[cid])} images")
        
    print("\n🛠️ augmenting...")
    for cid in [0, 1, 2]:
        items = data_map[cid]
        needed = TARGET_PER_CLASS - len(items)
        
        if needed > 0 and items:
            print(f"   -> Generating {needed} images for {CLASS_NAMES[cid]}...")
            for i in tqdm(range(needed)):
                src_img, src_lines = random.choice(items)
                augment_polygon(src_img, src_lines, f"aug_{cid}_{i}")
        elif not items:
            print(f"⚠️ Warning: No images found for {CLASS_NAMES[cid]}!")

    print("\n✅ Done!")

if __name__ == "__main__":
    main()