"""
Deep Learning Trainer

Advanced training system for neural networks with real-time monitoring,
early stopping, learning rate scheduling, and comprehensive logging.
"""

import numpy as np
import pandas as pd
from typing import Dict, List, Any, Optional, Callable, Tuple
import time
import json
from dataclasses import dataclass, asdict
import warnings
from abc import ABC, abstractmethod

# Try to import deep learning frameworks
try:
    import tensorflow as tf
    from tensorflow import keras
    TENSORFLOW_AVAILABLE = True
except ImportError:
    TENSORFLOW_AVAILABLE = False

try:
    import torch
    import torch.nn as nn
    import torch.optim as optim
    from torch.utils.data import DataLoader, TensorDataset
    PYTORCH_AVAILABLE = True
except ImportError:
    PYTORCH_AVAILABLE = False


@dataclass
class TrainingConfig:
    """Configuration for training parameters"""
    epochs: int = 100
    batch_size: int = 32
    validation_split: float = 0.2
    early_stopping_patience: int = 10
    reduce_lr_patience: int = 5
    reduce_lr_factor: float = 0.5
    min_lr: float = 1e-7
    shuffle: bool = True
    verbose: int = 1
    save_best_model: bool = True
    monitor_metric: str = 'val_loss'
    monitor_mode: str = 'min'  # 'min' or 'max'


@dataclass
class TrainingHistory:
    """Training history and metrics"""
    epoch: int
    train_loss: float
    val_loss: Optional[float] = None
    train_metrics: Dict[str, float] = None
    val_metrics: Dict[str, float] = None
    learning_rate: Optional[float] = None
    timestamp: float = None
    
    def __post_init__(self):
        if self.train_metrics is None:
            self.train_metrics = {}
        if self.val_metrics is None:
            self.val_metrics = {}
        if self.timestamp is None:
            self.timestamp = time.time()


class Callback(ABC):
    """Abstract base class for training callbacks"""
    
    @abstractmethod
    def on_epoch_begin(self, epoch: int, logs: Dict[str, Any]):
        pass
    
    @abstractmethod
    def on_epoch_end(self, epoch: int, logs: Dict[str, Any]):
        pass
    
    @abstractmethod
    def on_train_begin(self, logs: Dict[str, Any]):
        pass
    
    @abstractmethod
    def on_train_end(self, logs: Dict[str, Any]):
        pass


class EarlyStopping(Callback):
    """Early stopping callback"""
    
    def __init__(self, monitor: str = 'val_loss', patience: int = 10, 
                 mode: str = 'min', restore_best_weights: bool = True):
        self.monitor = monitor
        self.patience = patience
        self.mode = mode
        self.restore_best_weights = restore_best_weights
        self.best_value = None
        self.wait = 0
        self.stopped_epoch = 0
        self.best_weights = None
        
    def on_train_begin(self, logs: Dict[str, Any]):
        self.wait = 0
        self.stopped_epoch = 0
        self.best_value = np.inf if self.mode == 'min' else -np.inf
        
    def on_epoch_end(self, epoch: int, logs: Dict[str, Any]):
        current = logs.get(self.monitor)
        if current is None:
            return
        
        if (self.mode == 'min' and current < self.best_value) or \
           (self.mode == 'max' and current > self.best_value):
            self.best_value = current
            self.wait = 0
            # Store best weights (simplified)
            self.best_weights = True
        else:
            self.wait += 1
            if self.wait >= self.patience:
                self.stopped_epoch = epoch
                return True  # Stop training
        
        return False


class ReduceLROnPlateau(Callback):
    """Reduce learning rate when metric plateaus"""
    
    def __init__(self, monitor: str = 'val_loss', factor: float = 0.5, 
                 patience: int = 5, min_lr: float = 1e-7, mode: str = 'min'):
        self.monitor = monitor
        self.factor = factor
        self.patience = patience
        self.min_lr = min_lr
        self.mode = mode
        self.best = np.inf if mode == 'min' else -np.inf
        self.wait = 0
        
    def on_epoch_end(self, epoch: int, logs: Dict[str, Any]):
        current = logs.get(self.monitor)
        if current is None:
            return
        
        if (self.mode == 'min' and current < self.best) or \
           (self.mode == 'max' and current > self.best):
            self.best = current
            self.wait = 0
        else:
            self.wait += 1
            if self.wait >= self.patience:
                # Reduce learning rate (simplified - would need actual optimizer access)
                new_lr = logs.get('learning_rate', 0.001) * self.factor
                new_lr = max(new_lr, self.min_lr)
                logs['learning_rate'] = new_lr
                self.wait = 0


class ModelCheckpoint(Callback):
    """Model checkpoint callback"""
    
    def __init__(self, filepath: str, monitor: str = 'val_loss', 
                 save_best_only: bool = True, mode: str = 'min'):
        self.filepath = filepath
        self.monitor = monitor
        self.save_best_only = save_best_only
        self.mode = mode
        self.best = np.inf if mode == 'min' else -np.inf
        
    def on_epoch_end(self, epoch: int, logs: Dict[str, Any]):
        current = logs.get(self.monitor)
        if current is None:
            return
        
        if self.save_best_only:
            if (self.mode == 'min' and current < self.best) or \
               (self.mode == 'max' and current > self.best):
                self.best = current
                # Save model (simplified)
                self._save_model(epoch, logs)
        else:
            self._save_model(epoch, logs)
    
    def _save_model(self, epoch: int, logs: Dict[str, Any]):
        """Save model to file"""
        # This would implement actual model saving
        pass


class DeepLearningTrainer:
    """
    Advanced deep learning trainer supporting multiple frameworks
    with comprehensive monitoring and optimization features.
    """
    
    def __init__(self, framework: str = 'tensorflow'):
        """
        Initialize trainer
        
        Args:
            framework: 'tensorflow' or 'pytorch'
        """
        self.framework = framework.lower()
        self.model = None
        self.history = []
        self.callbacks = []
        self.training_config = None
        
        # Check framework availability
        if self.framework == 'tensorflow' and not TENSORFLOW_AVAILABLE:
            raise ImportError("TensorFlow is not available. Install with: pip install tensorflow")
        elif self.framework == 'pytorch' and not PYTORCH_AVAILABLE:
            raise ImportError("PyTorch is not available. Install with: pip install torch")
    
    def set_model(self, model):
        """Set the model to be trained"""
        self.model = model
    
    def compile_model(self, optimizer: str = 'adam', learning_rate: float = 0.001,
                     loss: str = 'mse', metrics: List[str] = None):
        """Compile the model with specified parameters"""
        if metrics is None:
            metrics = []
        
        if self.framework == 'tensorflow':
            if optimizer == 'adam':
                opt = keras.optimizers.Adam(learning_rate=learning_rate)
            elif optimizer == 'sgd':
                opt = keras.optimizers.SGD(learning_rate=learning_rate)
            elif optimizer == 'rmsprop':
                opt = keras.optimizers.RMSprop(learning_rate=learning_rate)
            else:
                opt = optimizer
            
            self.model.compile(optimizer=opt, loss=loss, metrics=metrics)
        
        elif self.framework == 'pytorch':
            # PyTorch compilation is different - optimizer is set separately
            pass
    
    def fit(self, X_train: np.ndarray, y_train: np.ndarray,
            X_val: Optional[np.ndarray] = None, y_val: Optional[np.ndarray] = None,
            config: Optional[TrainingConfig] = None) -> List[TrainingHistory]:
        """
        Train the model with specified configuration
        
        Args:
            X_train: Training features
            y_train: Training targets
            X_val: Validation features (optional)
            y_val: Validation targets (optional)
            config: Training configuration
            
        Returns:
            Training history
        """
        if config is None:
            config = TrainingConfig()
        
        self.training_config = config
        
        # Prepare validation data
        if X_val is None and config.validation_split > 0:
            # Split training data
            split_idx = int(len(X_train) * (1 - config.validation_split))
            X_val = X_train[split_idx:]
            y_val = y_train[split_idx:]
            X_train = X_train[:split_idx]
            y_train = y_train[:split_idx]
        
        # Initialize callbacks
        self._setup_callbacks(config)
        
        # Train based on framework
        if self.framework == 'tensorflow':
            return self._train_tensorflow(X_train, y_train, X_val, y_val, config)
        elif self.framework == 'pytorch':
            return self._train_pytorch(X_train, y_train, X_val, y_val, config)
    
    def _train_tensorflow(self, X_train: np.ndarray, y_train: np.ndarray,
                         X_val: Optional[np.ndarray], y_val: Optional[np.ndarray],
                         config: TrainingConfig) -> List[TrainingHistory]:
        """Train using TensorFlow/Keras"""
        
        # Prepare callbacks for Keras
        keras_callbacks = []
        
        if config.early_stopping_patience > 0:
            early_stop = keras.callbacks.EarlyStopping(
                monitor=config.monitor_metric,
                patience=config.early_stopping_patience,
                mode=config.monitor_mode,
                restore_best_weights=True
            )
            keras_callbacks.append(early_stop)
        
        if config.reduce_lr_patience > 0:
            reduce_lr = keras.callbacks.ReduceLROnPlateau(
                monitor=config.monitor_metric,
                factor=config.reduce_lr_factor,
                patience=config.reduce_lr_patience,
                min_lr=config.min_lr,
                mode=config.monitor_mode
            )
            keras_callbacks.append(reduce_lr)
        
        # Prepare validation data
        validation_data = (X_val, y_val) if X_val is not None else None
        
        # Train the model
        history = self.model.fit(
            X_train, y_train,
            epochs=config.epochs,
            batch_size=config.batch_size,
            validation_data=validation_data,
            callbacks=keras_callbacks,
            verbose=config.verbose,
            shuffle=config.shuffle
        )
        
        # Convert to TrainingHistory objects
        training_history = []
        for epoch in range(len(history.history['loss'])):
            epoch_history = TrainingHistory(
                epoch=epoch + 1,
                train_loss=history.history['loss'][epoch],
                val_loss=history.history.get('val_loss', [None])[epoch],
                train_metrics={k: v[epoch] for k, v in history.history.items() 
                             if k not in ['loss', 'val_loss']},
                val_metrics={k.replace('val_', ''): v[epoch] for k, v in history.history.items() 
                           if k.startswith('val_') and k != 'val_loss'}
            )
            training_history.append(epoch_history)
        
        self.history = training_history
        return training_history
    
    def _train_pytorch(self, X_train: np.ndarray, y_train: np.ndarray,
                      X_val: Optional[np.ndarray], y_val: Optional[np.ndarray],
                      config: TrainingConfig) -> List[TrainingHistory]:
        """Train using PyTorch"""
        
        # Convert to tensors
        X_train_tensor = torch.FloatTensor(X_train)
        y_train_tensor = torch.FloatTensor(y_train) if y_train.dtype.kind == 'f' else torch.LongTensor(y_train)
        
        train_dataset = TensorDataset(X_train_tensor, y_train_tensor)
        train_loader = DataLoader(train_dataset, batch_size=config.batch_size, shuffle=config.shuffle)
        
        if X_val is not None:
            X_val_tensor = torch.FloatTensor(X_val)
            y_val_tensor = torch.FloatTensor(y_val) if y_val.dtype.kind == 'f' else torch.LongTensor(y_val)
            val_dataset = TensorDataset(X_val_tensor, y_val_tensor)
            val_loader = DataLoader(val_dataset, batch_size=config.batch_size, shuffle=False)
        else:
            val_loader = None
        
        # Setup optimizer and loss
        optimizer = optim.Adam(self.model.parameters(), lr=0.001)  # Default to Adam
        criterion = nn.MSELoss()  # Default to MSE
        
        training_history = []
        best_val_loss = float('inf')
        patience_counter = 0
        
        for epoch in range(config.epochs):
            # Training phase
            self.model.train()
            train_loss = 0.0
            train_metrics = {}
            
            for batch_X, batch_y in train_loader:
                optimizer.zero_grad()
                outputs = self.model(batch_X)
                loss = criterion(outputs, batch_y)
                loss.backward()
                optimizer.step()
                train_loss += loss.item()
            
            train_loss /= len(train_loader)
            
            # Validation phase
            val_loss = None
            val_metrics = {}
            if val_loader:
                self.model.eval()
                val_loss = 0.0
                with torch.no_grad():
                    for batch_X, batch_y in val_loader:
                        outputs = self.model(batch_X)
                        loss = criterion(outputs, batch_y)
                        val_loss += loss.item()
                val_loss /= len(val_loader)
            
            # Early stopping
            if val_loss is not None:
                if config.monitor_mode == 'min':
                    if val_loss < best_val_loss:
                        best_val_loss = val_loss
                        patience_counter = 0
                    else:
                        patience_counter += 1
                else:
                    if val_loss > best_val_loss:
                        best_val_loss = val_loss
                        patience_counter = 0
                    else:
                        patience_counter += 1
                
                if patience_counter >= config.early_stopping_patience:
                    print(f"Early stopping at epoch {epoch + 1}")
                    break
            
            # Record history
            epoch_history = TrainingHistory(
                epoch=epoch + 1,
                train_loss=train_loss,
                val_loss=val_loss,
                train_metrics=train_metrics,
                val_metrics=val_metrics,
                learning_rate=optimizer.param_groups[0]['lr']
            )
            training_history.append(epoch_history)
            
            if config.verbose > 0:
                print(f"Epoch {epoch + 1}/{config.epochs} - "
                      f"loss: {train_loss:.4f}" + 
                      (f" - val_loss: {val_loss:.4f}" if val_loss is not None else ""))
        
        self.history = training_history
        return training_history
    
    def _setup_callbacks(self, config: TrainingConfig):
        """Setup training callbacks"""
        self.callbacks = []
        
        if config.early_stopping_patience > 0:
            self.callbacks.append(
                EarlyStopping(
                    monitor=config.monitor_metric,
                    patience=config.early_stopping_patience,
                    mode=config.monitor_mode
                )
            )
        
        if config.reduce_lr_patience > 0:
            self.callbacks.append(
                ReduceLROnPlateau(
                    monitor=config.monitor_metric,
                    factor=config.reduce_lr_factor,
                    patience=config.reduce_lr_patience,
                    min_lr=config.min_lr,
                    mode=config.monitor_mode
                )
            )
    
    def predict(self, X: np.ndarray) -> np.ndarray:
        """Make predictions"""
        if self.model is None:
            raise ValueError("No model has been trained or loaded")
        
        if self.framework == 'tensorflow':
            return self.model.predict(X)
        elif self.framework == 'pytorch':
            self.model.eval()
            with torch.no_grad():
                X_tensor = torch.FloatTensor(X)
                predictions = self.model(X_tensor)
                return predictions.numpy()
    
    def evaluate(self, X: np.ndarray, y: np.ndarray) -> Dict[str, float]:
        """Evaluate the model"""
        if self.model is None:
            raise ValueError("No model has been trained or loaded")
        
        if self.framework == 'tensorflow':
            results = self.model.evaluate(X, y, verbose=0)
            if isinstance(results, list):
                return dict(zip(self.model.metrics_names, results))
            else:
                return {'loss': results}
        elif self.framework == 'pytorch':
            self.model.eval()
            criterion = nn.MSELoss()  # Default
            total_loss = 0.0
            with torch.no_grad():
                X_tensor = torch.FloatTensor(X)
                y_tensor = torch.FloatTensor(y) if y.dtype.kind == 'f' else torch.LongTensor(y)
                outputs = self.model(X_tensor)
                loss = criterion(outputs, y_tensor)
                total_loss = loss.item()
            
            return {'loss': total_loss}
    
    def get_training_history(self) -> List[TrainingHistory]:
        """Get the complete training history"""
        return self.history
    
    def plot_training_history(self, save_path: Optional[str] = None):
        """Plot training history"""
        try:
            import matplotlib.pyplot as plt
            
            if not self.history:
                print("No training history available")
                return
            
            epochs = [h.epoch for h in self.history]
            train_losses = [h.train_loss for h in self.history]
            val_losses = [h.val_loss for h in self.history if h.val_loss is not None]
            
            fig, axes = plt.subplots(2, 2, figsize=(15, 10))
            
            # Loss plot
            axes[0, 0].plot(epochs, train_losses, label='Training Loss')
            if val_losses:
                val_epochs = epochs[:len(val_losses)]
                axes[0, 0].plot(val_epochs, val_losses, label='Validation Loss')
            axes[0, 0].set_title('Training and Validation Loss')
            axes[0, 0].set_xlabel('Epoch')
            axes[0, 0].set_ylabel('Loss')
            axes[0, 0].legend()
            
            # Learning rate plot (if available)
            lrs = [h.learning_rate for h in self.history if h.learning_rate is not None]
            if lrs:
                lr_epochs = epochs[:len(lrs)]
                axes[0, 1].plot(lr_epochs, lrs)
                axes[0, 1].set_title('Learning Rate Schedule')
                axes[0, 1].set_xlabel('Epoch')
                axes[0, 1].set_ylabel('Learning Rate')
            
            # Metrics plot (if available)
            if self.history[0].train_metrics:
                metric_names = list(self.history[0].train_metrics.keys())
                if metric_names:
                    for metric in metric_names[:2]:  # Plot first 2 metrics
                        metric_values = [h.train_metrics.get(metric) for h in self.history]
                        if metric_values and all(v is not None for v in metric_values):
                            axes[1, 0].plot(epochs, metric_values, label=f'Training {metric}')
                            
                            # Validation metrics
                            val_metric_values = [h.val_metrics.get(metric) for h in self.history 
                                               if h.val_metrics and metric in h.val_metrics]
                            if val_metric_values:
                                val_epochs = epochs[:len(val_metric_values)]
                                axes[1, 0].plot(val_epochs, val_metric_values, 
                                               label=f'Validation {metric}', linestyle='--')
                    
                    axes[1, 0].set_title('Training Metrics')
                    axes[1, 0].set_xlabel('Epoch')
                    axes[1, 0].set_ylabel('Metric Value')
                    axes[1, 0].legend()
            
            # Loss difference plot
            if val_losses:
                loss_diff = np.array(train_losses[:len(val_losses)]) - np.array(val_losses)
                axes[1, 1].plot(val_epochs, loss_diff)
                axes[1, 1].set_title('Training vs Validation Loss Difference')
                axes[1, 1].set_xlabel('Epoch')
            axes[1, 1].set_ylabel('Loss Difference')
            
            plt.tight_layout()
            if save_path:
                plt.savefig(save_path)
            plt.show()
            
        except ImportError:
            print("Matplotlib is required for plotting. Install with: pip install matplotlib")
    
    def save_model(self, filepath: str):
        """Save the trained model"""
        if self.model is None:
            raise ValueError("No model to save")
        
        if self.framework == 'tensorflow':
            self.model.save(filepath)
        elif self.framework == 'pytorch':
            torch.save(self.model.state_dict(), filepath)
    
    def load_model(self, filepath: str):
        """Load a trained model"""
        if self.framework == 'tensorflow':
            self.model = keras.models.load_model(filepath)
        elif self.framework == 'pytorch':
            # Note: Model architecture must be defined first
            if self.model is not None:
                self.model.load_state_dict(torch.load(filepath))
            else:
                raise ValueError("Model architecture must be defined before loading weights")
    
    def get_model_summary(self) -> str:
        """Get model summary"""
        if self.model is None:
            return "No model loaded"
        
        if self.framework == 'tensorflow':
            import io
            import sys
            
            # Capture model summary
            old_stdout = sys.stdout
            sys.stdout = buffer = io.StringIO()
            self.model.summary()
            sys.stdout = old_stdout
            return buffer.getvalue()
        elif self.framework == 'pytorch':
            return str(self.model)
    
    def analyze_training_stability(self) -> Dict[str, Any]:
        """Analyze training stability and provide insights"""
        if not self.history:
            return {"error": "No training history available"}
        
        train_losses = [h.train_loss for h in self.history]
        val_losses = [h.val_loss for h in self.history if h.val_loss is not None]
        
        analysis = {
            'total_epochs': len(self.history),
            'final_train_loss': train_losses[-1],
            'best_train_loss': min(train_losses),
            'train_loss_variance': np.var(train_losses),
            'convergence_epoch': None
        }
        
        if val_losses:
            analysis.update({
                'final_val_loss': val_losses[-1],
                'best_val_loss': min(val_losses),
                'val_loss_variance': np.var(val_losses),
                'overfitting_detected': val_losses[-1] > min(val_losses) * 1.1
            })
            
            # Find convergence epoch (when validation loss stopped improving significantly)
            for i in range(1, len(val_losses)):
                if abs(val_losses[i] - val_losses[i-1]) / val_losses[i-1] < 0.001:
                    analysis['convergence_epoch'] = i + 1
                    break
        
        # Stability assessment
        if analysis['train_loss_variance'] < 0.001:
            analysis['stability'] = 'Very Stable'
        elif analysis['train_loss_variance'] < 0.01:
            analysis['stability'] = 'Stable'
        elif analysis['train_loss_variance'] < 0.1:
            analysis['stability'] = 'Moderately Stable'
        else:
            analysis['stability'] = 'Unstable'
        
        return analysis
