"""
Training script for fine-tuning BLIP-2 on agricultural photos
"""

import os
import sys
import argparse
import json
from datetime import datetime

# Add src to path
sys.path.append(os.path.dirname(__file__))

from data.training_data_processor import TrainingDataProcessor
from model.fine_tuner import AgriculturalBLIPFineTuner

def main():
    parser = argparse.ArgumentParser(description='Train agricultural keyword generation model')
    
    # Data arguments
    parser.add_argument('--data-dir', type=str, default='data/training',
                       help='Directory containing training images')
    parser.add_argument('--metadata-file', type=str, default='data/training/metadata.csv',
                       help='CSV file with image filenames and keywords')
    parser.add_argument('--create-sample', action='store_true',
                       help='Create sample metadata for testing')
    
    # Training arguments
    parser.add_argument('--output-dir', type=str, default='models/agricultural_blip',
                       help='Directory to save trained model')
    parser.add_argument('--epochs', type=int, default=5,
                       help='Number of training epochs')
    parser.add_argument('--batch-size', type=int, default=8,
                       help='Training batch size')
    parser.add_argument('--learning-rate', type=float, default=5e-5,
                       help='Learning rate')
    parser.add_argument('--val-split', type=float, default=0.2,
                       help='Validation split ratio')
    
    # Model arguments
    parser.add_argument('--model-name', type=str, default='Salesforce/blip-image-captioning-base',
                       help='Pre-trained model name')
    parser.add_argument('--resume-from', type=str, default=None,
                       help='Resume training from checkpoint')
    
    # Hardware arguments
    parser.add_argument('--num-workers', type=int, default=4,
                       help='Number of data loader workers')
    
    args = parser.parse_args()
    
    print("🚜 Agricultural Photo Keyword Training")
    print("=" * 50)
    
    # Create sample metadata if requested
    if args.create_sample:
        print("Creating sample metadata for testing...")
        processor = TrainingDataProcessor(args.data_dir)
        os.makedirs(args.data_dir, exist_ok=True)
        processor.create_sample_metadata(args.metadata_file, num_samples=100)
        print(f"Sample metadata created: {args.metadata_file}")
        return
    
    # Check if metadata file exists
    if not os.path.exists(args.metadata_file):
        print(f"❌ Metadata file not found: {args.metadata_file}")
        print("Use --create-sample to create sample data for testing")
        return
    
    try:
        # Initialize components
        print("Initializing training components...")
        data_processor = TrainingDataProcessor(args.data_dir)
        fine_tuner = AgriculturalBLIPFineTuner(args.model_name, args.output_dir)
        
        # Load model
        print("Loading pre-trained model...")
        fine_tuner.load_model()
        
        # Prepare training data
        print("Preparing training data...")
        image_paths, keyword_lists = data_processor.prepare_training_data(args.metadata_file)
        
        if len(image_paths) == 0:
            print("❌ No valid training data found!")
            return
        
        print(f"Found {len(image_paths)} training examples")
        
        # Analyze training data
        analysis = data_processor.analyze_training_data(keyword_lists)
        print(f"Training data analysis:")
        print(f"  - Total images: {analysis['total_images']}")
        print(f"  - Unique keywords: {analysis['unique_keywords']}")
        print(f"  - Avg keywords per image: {analysis['avg_keywords_per_image']:.1f}")
        
        # Create train/val split
        print("Creating train/validation split...")
        train_paths, val_paths, train_keywords, val_keywords = data_processor.create_train_val_split(
            image_paths, keyword_lists, val_size=args.val_split
        )
        
        print(f"Training set: {len(train_paths)} images")
        print(f"Validation set: {len(val_paths)} images")
        
        # Create data loaders
        print("Creating data loaders...")
        train_loader, val_loader = data_processor.create_dataloaders(
            train_paths, train_keywords, val_paths, val_keywords,
            fine_tuner.processor, batch_size=args.batch_size, num_workers=args.num_workers
        )
        
        # Setup training
        print("Setting up training...")
        fine_tuner.setup_training(train_loader, val_loader, learning_rate=args.learning_rate)
        
        # Resume from checkpoint if specified
        if args.resume_from:
            print(f"Resuming from checkpoint: {args.resume_from}")
            fine_tuner.load_checkpoint(args.resume_from)
        
        # Save training configuration
        config = {
            'model_name': args.model_name,
            'data_dir': args.data_dir,
            'metadata_file': args.metadata_file,
            'epochs': args.epochs,
            'batch_size': args.batch_size,
            'learning_rate': args.learning_rate,
            'val_split': args.val_split,
            'training_data_analysis': analysis,
            'timestamp': datetime.now().isoformat()
        }
        
        config_path = os.path.join(args.output_dir, 'training_config.json')
        data_processor.save_training_config(config, config_path)
        
        # Start training
        print(f"\n🚀 Starting training for {args.epochs} epochs...")
        print(f"Output directory: {args.output_dir}")
        
        training_history = fine_tuner.train(
            train_loader, val_loader,
            num_epochs=args.epochs,
            save_every=1,
            early_stopping_patience=3
        )
        
        # Training summary
        print("\n✅ Training completed!")
        print(f"Best validation loss: {fine_tuner.best_val_loss:.4f}")
        print(f"Total epochs: {len(training_history)}")
        print(f"Model saved to: {args.output_dir}")
        
        # Test the trained model
        print("\n🧪 Testing trained model...")
        test_model(fine_tuner, train_paths[:3])  # Test on first 3 training images
        
    except Exception as e:
        print(f"\n❌ Training failed: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)

def test_model(fine_tuner, test_image_paths):
    """Test the trained model on sample images"""
    print("Testing keyword generation on sample images:")
    print("-" * 50)
    
    for image_path in test_image_paths:
        try:
            keywords = fine_tuner.generate_keywords(image_path)
            filename = os.path.basename(image_path)
            print(f"Image: {filename}")
            print(f"Keywords: {', '.join(keywords)}")
            print("-" * 50)
        except Exception as e:
            print(f"Error testing {image_path}: {e}")

if __name__ == "__main__":
    main()
