"""
Fine-tuning module for agricultural keyword generation using BLIP-2
"""

import os
import torch
import torch.nn as nn
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR
from transformers import BlipProcessor, BlipForConditionalGeneration
from transformers import get_linear_schedule_with_warmup
import logging
from typing import Dict, List, Optional, Tuple
import json
from tqdm import tqdm
import numpy as np
from datetime import datetime

class AgriculturalBLIPFineTuner:
    """Fine-tune BLIP-2 model for agricultural keyword generation"""
    
    def __init__(self, model_name: str = "Salesforce/blip-image-captioning-base",
                 output_dir: str = "models/agricultural_blip"):
        """
        Initialize fine-tuner
        
        Args:
            model_name: Pre-trained BLIP model name
            output_dir: Directory to save fine-tuned model
        """
        self.model_name = model_name
        self.output_dir = output_dir
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        # Create output directory
        os.makedirs(output_dir, exist_ok=True)
        
        # Setup logging
        self.setup_logging()
        
        # Initialize model and processor
        self.processor = None
        self.model = None
        self.optimizer = None
        self.scheduler = None
        
        # Training state
        self.current_epoch = 0
        self.best_val_loss = float('inf')
        self.training_history = []
    
    def setup_logging(self):
        """Setup logging for training"""
        log_file = os.path.join(self.output_dir, 'training.log')
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(levelname)s - %(message)s',
            handlers=[
                logging.FileHandler(log_file),
                logging.StreamHandler()
            ]
        )
        self.logger = logging.getLogger(__name__)
    
    def load_model(self):
        """Load pre-trained BLIP model and processor"""
        self.logger.info(f"Loading model: {self.model_name}")
        
        self.processor = BlipProcessor.from_pretrained(self.model_name)
        self.model = BlipForConditionalGeneration.from_pretrained(self.model_name)
        
        # Move model to device
        self.model.to(self.device)
        
        self.logger.info(f"Model loaded on device: {self.device}")
        
        # Print model info
        total_params = sum(p.numel() for p in self.model.parameters())
        trainable_params = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
        
        self.logger.info(f"Total parameters: {total_params:,}")
        self.logger.info(f"Trainable parameters: {trainable_params:,}")
    
    def setup_training(self, train_loader, val_loader, learning_rate: float = 5e-5,
                      weight_decay: float = 0.01, warmup_steps: int = 500):
        """
        Setup training components
        
        Args:
            train_loader: Training data loader
            val_loader: Validation data loader
            learning_rate: Learning rate for optimizer
            weight_decay: Weight decay for regularization
            warmup_steps: Number of warmup steps for scheduler
        """
        # Setup optimizer
        self.optimizer = AdamW(
            self.model.parameters(),
            lr=learning_rate,
            weight_decay=weight_decay,
            betas=(0.9, 0.999),
            eps=1e-8
        )
        
        # Calculate total training steps
        total_steps = len(train_loader) * 10  # Assuming 10 epochs max
        
        # Setup scheduler
        self.scheduler = get_linear_schedule_with_warmup(
            self.optimizer,
            num_warmup_steps=warmup_steps,
            num_training_steps=total_steps
        )
        
        self.logger.info(f"Training setup complete:")
        self.logger.info(f"  - Learning rate: {learning_rate}")
        self.logger.info(f"  - Weight decay: {weight_decay}")
        self.logger.info(f"  - Warmup steps: {warmup_steps}")
        self.logger.info(f"  - Total steps: {total_steps}")
    
    def train_epoch(self, train_loader) -> Dict[str, float]:
        """Train for one epoch"""
        self.model.train()
        total_loss = 0.0
        num_batches = len(train_loader)
        
        progress_bar = tqdm(train_loader, desc=f"Epoch {self.current_epoch + 1}")
        
        for batch_idx, batch in enumerate(progress_bar):
            # Move batch to device
            batch = {k: v.to(self.device) for k, v in batch.items()}
            
            # Forward pass
            outputs = self.model(
                pixel_values=batch['pixel_values'],
                input_ids=batch['input_ids'],
                attention_mask=batch['attention_mask'],
                labels=batch['labels']
            )
            
            loss = outputs.loss
            
            # Backward pass
            self.optimizer.zero_grad()
            loss.backward()
            
            # Gradient clipping
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
            
            # Update weights
            self.optimizer.step()
            self.scheduler.step()
            
            # Update metrics
            total_loss += loss.item()
            avg_loss = total_loss / (batch_idx + 1)
            
            # Update progress bar
            progress_bar.set_postfix({
                'loss': f'{loss.item():.4f}',
                'avg_loss': f'{avg_loss:.4f}',
                'lr': f'{self.scheduler.get_last_lr()[0]:.2e}'
            })
        
        return {'train_loss': total_loss / num_batches}
    
    def validate_epoch(self, val_loader) -> Dict[str, float]:
        """Validate for one epoch"""
        self.model.eval()
        total_loss = 0.0
        num_batches = len(val_loader)
        
        with torch.no_grad():
            for batch in tqdm(val_loader, desc="Validation"):
                # Move batch to device
                batch = {k: v.to(self.device) for k, v in batch.items()}
                
                # Forward pass
                outputs = self.model(
                    pixel_values=batch['pixel_values'],
                    input_ids=batch['input_ids'],
                    attention_mask=batch['attention_mask'],
                    labels=batch['labels']
                )
                
                total_loss += outputs.loss.item()
        
        return {'val_loss': total_loss / num_batches}
    
    def train(self, train_loader, val_loader, num_epochs: int = 5,
              save_every: int = 1, early_stopping_patience: int = 3) -> Dict:
        """
        Main training loop
        
        Args:
            train_loader: Training data loader
            val_loader: Validation data loader
            num_epochs: Number of epochs to train
            save_every: Save model every N epochs
            early_stopping_patience: Stop if no improvement for N epochs
            
        Returns:
            Training history dictionary
        """
        self.logger.info(f"Starting training for {num_epochs} epochs")
        
        patience_counter = 0
        
        for epoch in range(num_epochs):
            self.current_epoch = epoch
            
            # Train epoch
            train_metrics = self.train_epoch(train_loader)
            
            # Validate epoch
            val_metrics = self.validate_epoch(val_loader)
            
            # Combine metrics
            epoch_metrics = {**train_metrics, **val_metrics, 'epoch': epoch + 1}
            self.training_history.append(epoch_metrics)
            
            # Log metrics
            self.logger.info(
                f"Epoch {epoch + 1}/{num_epochs} - "
                f"Train Loss: {train_metrics['train_loss']:.4f}, "
                f"Val Loss: {val_metrics['val_loss']:.4f}"
            )
            
            # Save model if improved
            if val_metrics['val_loss'] < self.best_val_loss:
                self.best_val_loss = val_metrics['val_loss']
                self.save_model('best_model')
                patience_counter = 0
                self.logger.info(f"New best model saved with val_loss: {self.best_val_loss:.4f}")
            else:
                patience_counter += 1
            
            # Save checkpoint
            if (epoch + 1) % save_every == 0:
                self.save_model(f'checkpoint_epoch_{epoch + 1}')
            
            # Early stopping
            if patience_counter >= early_stopping_patience:
                self.logger.info(f"Early stopping triggered after {epoch + 1} epochs")
                break
        
        # Save final model
        self.save_model('final_model')
        
        # Save training history
        self.save_training_history()
        
        self.logger.info("Training completed!")
        return self.training_history
    
    def save_model(self, checkpoint_name: str):
        """Save model checkpoint"""
        checkpoint_dir = os.path.join(self.output_dir, checkpoint_name)
        os.makedirs(checkpoint_dir, exist_ok=True)
        
        # Save model and processor
        self.model.save_pretrained(checkpoint_dir)
        self.processor.save_pretrained(checkpoint_dir)
        
        # Save training state
        state = {
            'epoch': self.current_epoch,
            'best_val_loss': self.best_val_loss,
            'model_name': self.model_name,
            'training_history': self.training_history
        }
        
        torch.save(state, os.path.join(checkpoint_dir, 'training_state.pt'))
        
        self.logger.info(f"Model saved: {checkpoint_dir}")
    
    def load_checkpoint(self, checkpoint_path: str):
        """Load model from checkpoint"""
        self.logger.info(f"Loading checkpoint: {checkpoint_path}")
        
        # Load model and processor
        self.processor = BlipProcessor.from_pretrained(checkpoint_path)
        self.model = BlipForConditionalGeneration.from_pretrained(checkpoint_path)
        self.model.to(self.device)
        
        # Load training state if available
        state_path = os.path.join(checkpoint_path, 'training_state.pt')
        if os.path.exists(state_path):
            state = torch.load(state_path, map_location=self.device)
            self.current_epoch = state.get('epoch', 0)
            self.best_val_loss = state.get('best_val_loss', float('inf'))
            self.training_history = state.get('training_history', [])
        
        self.logger.info("Checkpoint loaded successfully")
    
    def save_training_history(self):
        """Save training history to JSON"""
        history_path = os.path.join(self.output_dir, 'training_history.json')
        with open(history_path, 'w') as f:
            json.dump(self.training_history, f, indent=2)
        
        self.logger.info(f"Training history saved: {history_path}")
    
    def generate_keywords(self, image_path: str, max_length: int = 50) -> List[str]:
        """
        Generate keywords for a single image using fine-tuned model
        
        Args:
            image_path: Path to image file
            max_length: Maximum generation length
            
        Returns:
            List of generated keywords
        """
        if self.model is None or self.processor is None:
            raise ValueError("Model not loaded. Call load_model() or load_checkpoint() first.")
        
        self.model.eval()
        
        with torch.no_grad():
            # Load and process image
            from PIL import Image
            image = Image.open(image_path).convert('RGB')
            
            # Process image
            inputs = self.processor(image, return_tensors="pt")
            inputs = {k: v.to(self.device) for k, v in inputs.items()}
            
            # Generate
            outputs = self.model.generate(
                **inputs,
                max_length=max_length,
                num_beams=5,
                temperature=0.7,
                do_sample=True,
                early_stopping=True
            )
            
            # Decode
            generated_text = self.processor.decode(outputs[0], skip_special_tokens=True)
            
            # Parse keywords
            keywords = [kw.strip() for kw in generated_text.split(',')]
            keywords = [kw for kw in keywords if kw and len(kw) > 1]
            
            return keywords[:10]  # Limit to 10 keywords
