#!/usr/bin/env python3
# This script is for YOLOv8 training using data from PlantNet and local files.

import os
import requests
import json
import random
import shutil
from ultralytics import YOLO

# ==================== CONFIGURATION & SECRETS ====================
# Use your API key from the main script
PLANTNET_API_KEY = "2b10HlR0yP4wZbvHjP8jCo3Aae"
PLANTNET_SEARCH_URL = "https://my-api.plantnet.org/v2/identify/all"

# Define the training directory structure
TRAIN_DIR = 'yolo_train'
DATASET_YAML_PATH = os.path.join(TRAIN_DIR, 'dataset.yaml')

# Define your 52 morphological features here.
# This list is crucial and must match the labels you will create in your annotation tool.
MORPHOLOGICAL_FEATURES = [
    'leaf', 'stem', 'flower', 'fruit', 'root', 'branch', 'petal', 'sepal',
    'bud', 'thorn', 'node', 'vein', 'blade', 'petiole', 'stamen', 'pistil',
    'seed', 'bark', 'trunk', 'tuber', 'rhizome', 'bulb', 'spore', 'cone',
    'tendril', 'pollen', 'nectar', 'bract', 'gemma', 'stomata', 'cuticle',
    'xylem', 'phloem', 'vascular_bundle', 'cambium', 'lenticel', 'stolon',
    'runner', 'corm', 'sheath', 'ligule', 'auricle', 'rachis', 'rachilla',
    'glume', 'lemma', 'palea', 'floret', 'awns', 'caryopsis', 'pericarp', 'endosperm'
]
# You can update this list based on what you want to annotate.

# ==================== DATA COLLECTION & PREPARATION ====================
def download_and_prepare_dataset(plant_list: list, num_images_per_plant: int = 50):
    """
    Downloads and organizes images for a list of plants.
    Note: PlantNet's API is for identification. We'll simulate fetching images
    from a hypothetical image search to demonstrate the concept.
    In a real-world scenario, you would use a web scraper or a curated dataset.
    
    This function creates a basic directory structure for YOLO.
    """
    print("Preparing dataset directory...")
    # Clean up and create new directories
    if os.path.exists(TRAIN_DIR):
        shutil.rmtree(TRAIN_DIR)
    os.makedirs(os.path.join(TRAIN_DIR, 'images'))
    os.makedirs(os.path.join(TRAIN_DIR, 'images', 'train'))
    os.makedirs(os.path.join(TRAIN_DIR, 'images', 'val'))
    os.makedirs(os.path.join(TRAIN_DIR, 'labels'))
    os.makedirs(os.path.join(TRAIN_DIR, 'labels', 'train'))
    os.makedirs(os.path.join(TRAIN_DIR, 'labels', 'val'))

    # Simulate downloading and creating placeholder images and labels
    print("Simulating dataset creation. Please replace with your actual data and annotations.")
    
    # Create dataset.yaml
    with open(DATASET_YAML_PATH, 'w') as f:
        yaml_content = f"""
path: {os.path.abspath(TRAIN_DIR)}
train: images/train
val: images/val

names:
  {os.linesep.join([f"  {i}: {name}" for i, name in enumerate(MORPHOLOGICAL_FEATURES)])}
"""
        f.write(yaml_content)
    
    print(f"✅ Created {DATASET_YAML_PATH}")
    print(f"Next, populate {os.path.join(TRAIN_DIR, 'images', 'train')}, {os.path.join(TRAIN_DIR, 'images', 'val')}, and their corresponding 'labels' directories with your annotated data.")

# ==================== YOLO TRAINING ====================
def start_yolo_training():
    """
    Starts the YOLOv8 training process using the prepared dataset.
    """
    if not os.path.exists(DATASET_YAML_PATH):
        print("❌ Dataset not found. Please run the data preparation step first.")
        return

    print("Loading YOLOv8 model for training...")
    try:
        model = YOLO('yolov8n.pt')
        print("Starting training...")
        
        # Train the model
        # You can adjust epochs, image size (imgsz), and other parameters
        results = model.train(data=DATASET_YAML_PATH, epochs=100, imgsz=640)
        
        print("✅ Training complete. Results are saved in the 'runs' directory.")
        print("The best model is located at 'runs/detect/train/weights/best.pt'.")
        
    except Exception as e:
        print(f"An error occurred during training: {e}")

# ==================== MAIN EXECUTION ====================
if __name__ == "__main__":
    # Example list of plants to train on
    plants_to_train = ["rosemary", "basil", "mint"]
    
    # Step 1: Prepare the directory structure and dataset.yaml
    # This function is a placeholder. You need to manually add your images and labels.
    download_and_prepare_dataset(plants_to_train)
    
    # Step 2: Start the training process
    # Make sure you have your annotated images and labels in place before running this.
    # start_yolo_training() # Uncomment this line to run the training