#!/bin/bash

# Simple ConvNeXt training script for flower classification
# This script provides an easy way to train a flower classification model

echo "🌸 Flowerfy Simple Training Script"
echo "=================================="

# Check if training data exists
if [ ! -d "training_data/images" ]; then
    echo "❌ Training data directory not found!"
    echo "Please create 'training_data/images/' and organize your images by flower type."
    echo ""
    echo "Example structure:"
    echo "  training_data/images/roses/"
    echo "  training_data/images/tulips/"
    echo "  training_data/images/lilies/"
    echo "  training_data/images/orchids/"
    exit 1
fi

# Count training images
total_images=0
echo "Found flower types:"
for dir in training_data/images/*/; do
    if [ -d "$dir" ]; then
        flower_type=$(basename "$dir")
        count=$(find "$dir" -type f \( -iname "*.jpg" -o -iname "*.jpeg" -o -iname "*.png" -o -iname "*.webp" \) | wc -l)
        if [ "$count" -gt 0 ]; then
            echo "  - $flower_type: $count images"
            total_images=$((total_images + count))
        fi
    fi
done

if [ "$total_images" -lt 10 ]; then
    echo "❌ Insufficient training data. Found $total_images images."
    echo "You need at least 10 images to train the model."
    exit 1
fi

echo ""
echo "Total images: $total_images"
echo ""
echo "Training Configuration:"
echo "  - Method: Simple training (fast, lightweight)"
echo "  - Epochs: 3 (default)"
echo "  - Batch size: 4 (default)"
echo "  - Learning rate: 1e-5 (default)"
echo ""
echo "Starting training..."
echo ""

# Run the training
cd training
uv run python simple_trainer.py "$@"

echo ""
echo "Training completed! Check the output above for results."
echo "Your trained model will be in: training_data/trained_models/simple_trained/"