"""
Ablation Study Framework for Brain Tumor Detection Models
"""

import numpy as np
import tensorflow as tf
import json
import os
import pandas as pd
from datetime import datetime
from pathlib import Path
import matplotlib.pyplot as plt
import seaborn as sns


class AblationStudy:
    """
    Framework for conducting ablation studies on brain tumor detection models
    """
    
    def __init__(self, base_config, results_dir='./ablation_results'):
        """
        Initialize ablation study
        
        Args:
            base_config: Base configuration dictionary
            results_dir: Directory to save results
        """
        self.base_config = base_config
        self.results_dir = results_dir
        self.results = {}
        self.study_metadata = {
            'start_time': datetime.now().isoformat(),
            'base_config': base_config,
            'experiments': []
        }
        
        os.makedirs(results_dir, exist_ok=True)
    
    def add_experiment(self, name, config_modification, description=""):
        """
        Add an experiment configuration
        
        Args:
            name: Name of the experiment
            config_modification: Dictionary of config modifications
            description: Description of what this experiment tests
        """
        config = self.base_config.copy()
        config.update(config_modification)
        
        self.study_metadata['experiments'].append({
            'name': name,
            'config': config,
            'description': description,
            'modifications': config_modification
        })
    
    def run_experiment(self, experiment_idx, model_builder, data, metrics_calculator):
        """
        Run a single ablation experiment
        
        Args:
            experiment_idx: Index of experiment in study_metadata['experiments']
            model_builder: Function to build model with given config
            data: Tuple of (X_train, y_train, X_val, y_val)
            metrics_calculator: Function to calculate metrics
            
        Returns:
            Results dictionary
        """
        experiment = self.study_metadata['experiments'][experiment_idx]
        config = experiment['config']
        name = experiment['name']
        
        print(f"\n{'='*60}")
        print(f"Running Ablation Experiment: {name}")
        print(f"{'='*60}")
        print(f"Description: {experiment['description']}")
        print(f"Modifications: {experiment['modifications']}")
        
        # Build model
        model = model_builder(config)
        
        # Train model
        history = self._train_model(model, data, config)
        
        # Evaluate
        metrics = metrics_calculator(model, data)
        
        # Store results
        self.results[name] = {
            'metrics': metrics,
            'history': history.history if hasattr(history, 'history') else history,
            'config': config
        }
        
        print(f"Experiment {name} completed. Metrics: {metrics}")
        
        return self.results[name]
    
    def _train_model(self, model, data, config):
        """
        Train model with given configuration
        
        Args:
            model: Model to train
            data: Training data tuple
            config: Training configuration
            
        Returns:
            Training history
        """
        X_train, y_train, X_val, y_val = data
        
        # Compile model
        model.compile(
            optimizer=tf.keras.optimizers.Adam(learning_rate=config.get('learning_rate', 1e-4)),
            loss=config.get('loss_fn', 'binary_crossentropy'),
            metrics=config.get('metrics', ['accuracy'])
        )
        
        # Callbacks
        callbacks = [
            tf.keras.callbacks.EarlyStopping(
                monitor='val_loss',
                patience=config.get('patience', 10),
                restore_best_weights=True
            ),
            tf.keras.callbacks.ReduceLROnPlateau(
                monitor='val_loss',
                factor=0.5,
                patience=5,
                min_lr=1e-7
            )
        ]
        
        # Train
        history = model.fit(
            X_train,
            y_train,
            validation_data=(X_val, y_val),
            epochs=config.get('epochs', 50),
            batch_size=config.get('batch_size', 32),
            callbacks=callbacks,
            verbose=1
        )
        
        return history
    
    def run_all_experiments(self, model_builder, data, metrics_calculator):
        """
        Run all experiments in the study
        
        Args:
            model_builder: Function to build model with given config
            data: Training data tuple
            metrics_calculator: Function to calculate metrics
        """
        for i in range(len(self.study_metadata['experiments'])):
            self.run_experiment(i, model_builder, data, metrics_calculator)
        
        # Save results
        self.save_results()
        
        return self.results
    
    def save_results(self):
        """Save ablation study results"""
        # Save summary
        summary_path = os.path.join(self.results_dir, 'ablation_summary.json')
        with open(summary_path, 'w') as f:
            json.dump({
                'study_metadata': self.study_metadata,
                'results': self.results
            }, f, indent=2)
        
        # Save detailed results as CSV
        results_data = []
        for name, result in self.results.items():
            row = {'experiment': name}
            row.update(result['metrics'])
            results_data.append(row)
        
        results_df = pd.DataFrame(results_data)
        results_df.to_csv(os.path.join(self.results_dir, 'ablation_results.csv'), index=False)
        
        # Save plots
        self.plot_results()
        
        print(f"Results saved to {self.results_dir}")
    
    def plot_results(self):
        """Plot ablation study results"""
        if not self.results:
            return
        
        # Extract metrics
        experiments = list(self.results.keys())
        metrics_names = list(list(self.results.values())[0]['metrics'].keys())
        
        # Create subplots for each metric
        fig, axes = plt.subplots(1, len(metrics_names), figsize=(6*len(metrics_names), 5))
        if len(metrics_names) == 1:
            axes = [axes]
        
        for ax, metric_name in zip(axes, metrics_names):
            values = [self.results[exp]['metrics'][metric_name] for exp in experiments]
            
            # Create bar plot
            bars = ax.bar(experiments, values, color=plt.cm.Set3(np.linspace(0, 1, len(experiments))))
            ax.set_title(metric_name.replace('_', ' ').title())
            ax.set_ylabel(metric_name)
            ax.tick_params(axis='x', rotation=45)
            
            # Add value labels on bars
            for bar, value in zip(bars, values):
                ax.text(bar.get_x() + bar.get_width()/2., bar.get_height() + 0.001,
                       f'{value:.4f}', ha='center', va='bottom', fontsize=9)
        
        plt.tight_layout()
        plt.savefig(os.path.join(self.results_dir, 'ablation_results.png'), dpi=300, bbox_inches='tight')
        plt.close()
        
        # Plot training histories
        if len(experiments) > 0 and 'history' in self.results[experiments[0]]:
            self._plot_training_histories(experiments)
    
    def _plot_training_histories(self, experiments):
        """Plot training histories for all experiments"""
        fig, axes = plt.subplots(1, 2, figsize=(15, 5))
        
        for exp in experiments:
            history = self.results[exp]['history']
            if 'loss' in history:
                axes[0].plot(history['loss'], label=f'{exp} - train')
                if 'val_loss' in history:
                    axes[0].plot(history['val_loss'], label=f'{exp} - val', linestyle='--')
            
            if 'accuracy' in history:
                axes[1].plot(history['accuracy'], label=f'{exp} - train')
                if 'val_accuracy' in history:
                    axes[1].plot(history['val_accuracy'], label=f'{exp} - val', linestyle='--')
        
        axes[0].set_title('Loss')
        axes[0].set_xlabel('Epoch')
        axes[0].set_ylabel('Loss')
        axes[0].legend(bbox_to_anchor=(1.05, 1), loc='upper left')
        
        axes[1].set_title('Accuracy')
        axes[1].set_xlabel('Epoch')
        axes[1].set_ylabel('Accuracy')
        axes[1].legend(bbox_to_anchor=(1.05, 1), loc='upper left')
        
        plt.tight_layout()
        plt.savefig(os.path.join(self.results_dir, 'training_histories.png'), dpi=300, bbox_inches='tight')
        plt.close()
    
    def get_comparison_table(self):
        """
        Get comparison table of all experiments
        
        Returns:
            Pandas DataFrame with comparison results
        """
        if not self.results:
            raise ValueError("No results available. Run experiments first.")
        
        rows = []
        for name, result in self.results.items():
            row = {'Experiment': name}
            row.update(result['metrics'])
            row['Description'] = next(
                (exp['description'] for exp in self.study_metadata['experiments'] 
                 if exp['name'] == name), ''
            )
            rows.append(row)
        
        return pd.DataFrame(rows)


class SegmentationAblationStudy(AblationStudy):
    """
    Ablation study framework specifically for segmentation models
    """
    
    def __init__(self, base_config, results_dir='./segmentation_ablation_results'):
        super().__init__(base_config, results_dir)
    
    def add_segmentation_experiment(self, name, model_config, training_config, description=""):
        """
        Add a segmentation experiment
        
        Args:
            name: Experiment name
            model_config: Model configuration modifications
            training_config: Training configuration modifications
            description: Description of the experiment
        """
        config = {
            **self.base_config,
            **model_config,
            **training_config
        }
        
        self.study_metadata['experiments'].append({
            'name': name,
            'config': config,
            'description': description,
            'model_modifications': model_config,
            'training_modifications': training_config
        })
    
    def run_segmentation_experiment(self, experiment_idx, model_builder, data, metrics_calculator):
        """
        Run a segmentation experiment
        
        Args:
            experiment_idx: Index of experiment
            model_builder: Function to build model
            data: Tuple of (X_train, y_train, X_val, y_val) where y are masks
            metrics_calculator: Function to calculate segmentation metrics
            
        Returns:
            Results dictionary
        """
        experiment = self.study_metadata['experiments'][experiment_idx]
        config = experiment['config']
        name = experiment['name']
        
        print(f"\n{'='*60}")
        print(f"Running Segmentation Ablation: {name}")
        print(f"{'='*60}")
        
        # Build model
        model = model_builder(config)
        
        # Compile
        model.compile(
            optimizer=tf.keras.optimizers.Adam(learning_rate=config.get('learning_rate', 1e-4)),
            loss=config.get('loss_fn', 'binary_crossentropy'),
            metrics=[
                tf.keras.metrics.MeanIoU(num_classes=2),
                'accuracy'
            ]
        )
        
        # Train
        X_train, y_train, X_val, y_val = data
        
        callbacks = [
            tf.keras.callbacks.EarlyStopping(
                monitor='val_loss',
                patience=config.get('patience', 10),
                restore_best_weights=True
            )
        ]
        
        history = model.fit(
            X_train,
            y_train,
            validation_data=(X_val, y_val),
            epochs=config.get('epochs', 100),
            batch_size=config.get('batch_size', 16),
            callbacks=callbacks
        )
        
        # Evaluate with custom metrics
        metrics = metrics_calculator(model, (X_val, y_val))
        
        # Store results
        self.results[name] = {
            'metrics': metrics,
            'history': history.history,
            'config': config
        }
        
        print(f"Segmentation ablation {name} completed. Metrics: {metrics}")
        
        return self.results[name]


def create_attention_ablation_study(base_config, results_dir='./attention_ablation'):
    """
    Create an ablation study for attention mechanisms
    
    Args:
        base_config: Base configuration
        results_dir: Directory to save results
        
    Returns:
        AblationStudy instance with experiments added
    """
    study = AblationStudy(base_config, results_dir)
    
    # Baseline without attention
    study.add_experiment(
        name='baseline_no_attention',
        config_modification={'use_attention': False},
        description='Baseline model without any attention mechanisms'
    )
    
    # With attention gates in skip connections
    study.add_experiment(
        name='attention_skip_connections',
        config_modification={'use_attention': True},
        description='Model with attention gates in U-Net skip connections'
    )
    
    # With channel attention
    study.add_experiment(
        name='channel_attention',
        config_modification={'attention_type': 'channel'},
        description='Model with channel-wise attention mechanism'
    )
    
    # With spatial attention
    study.add_experiment(
        name='spatial_attention',
        config_modification={'attention_type': 'spatial'},
        description='Model with spatial attention mechanism'
    )
    
    # With both channel and spatial attention
    study.add_experiment(
        name='cbam_attention',
        config_modification={'attention_type': 'cbam'},
        description='Model with combined channel and spatial attention (CBAM)'
    )
    
    return study


def create_architecture_ablation_study(base_config, results_dir='./architecture_ablation'):
    """
    Create an ablation study for architecture variations
    
    Args:
        base_config: Base configuration
        results_dir: Directory to save results
        
    Returns:
        AblationStudy instance with experiments added
    """
    study = AblationStudy(base_config, results_dir)
    
    # Baseline
    study.add_experiment(
        name='baseline',
        config_modification={},
        description='Baseline architecture'
    )
    
    # Different depths
    study.add_experiment(
        name='shallow_network',
        config_modification={'num_layers': 3},
        description='Shallower network with fewer layers'
    )
    
    study.add_experiment(
        name='deep_network',
        config_modification={'num_layers': 6},
        description='Deeper network with more layers'
    )
    
    # Different filter sizes
    study.add_experiment(
        name='smaller_filters',
        config_modification={'base_filters': 32},
        description='Network with smaller base number of filters'
    )
    
    study.add_experiment(
        name='larger_filters',
        config_modification={'base_filters': 128},
        description='Network with larger base number of filters'
    )
    
    # With residual connections
    study.add_experiment(
        name='residual_connections',
        config_modification={'use_residual': True},
        description='Network with residual connections'
    )
    
    # With dense connections
    study.add_experiment(
        name='dense_connections',
        config_modification={'use_dense': True},
        description='Network with dense connections'
    )
    
    return study


def create_loss_ablation_study(base_config, results_dir='./loss_ablation'):
    """
    Create an ablation study for different loss functions
    
    Args:
        base_config: Base configuration
        results_dir: Directory to save results
        
    Returns:
        AblationStudy instance with experiments added
    """
    study = AblationStudy(base_config, results_dir)
    
    # Baseline cross-entropy
    study.add_experiment(
        name='cross_entropy',
        config_modification={'loss_fn': 'binary_crossentropy'},
        description='Standard binary cross-entropy loss'
    )
    
    # Dice loss
    study.add_experiment(
        name='dice_loss',
        config_modification={'loss_fn': 'dice_loss'},
        description='Dice loss for better handling of class imbalance'
    )
    
    # Combined loss
    study.add_experiment(
        name='combined_dice_bce',
        config_modification={'loss_fn': 'combined_dice_bce', 'loss_weights': [0.5, 0.5]},
        description='Combined Dice and BCE loss'
    )
    
    # Focal loss
    study.add_experiment(
        name='focal_loss',
        config_modification={'loss_fn': 'focal_loss', 'focal_gamma': 2.0},
        description='Focal loss for hard example mining'
    )
    
    # Tversky loss
    study.add_experiment(
        name='tversky_loss',
        config_modification={'loss_fn': 'tversky_loss', 'alpha': 0.5, 'beta': 0.5},
        description='Tversky loss for imbalanced segmentation'
    )
    
    return study


def create_data_augmentation_ablation_study(base_config, results_dir='./augmentation_ablation'):
    """
    Create an ablation study for data augmentation strategies
    
    Args:
        base_config: Base configuration
        results_dir: Directory to save results
        
    Returns:
        AblationStudy instance with experiments added
    """
    study = AblationStudy(base_config, results_dir)
    
    # No augmentation
    study.add_experiment(
        name='no_augmentation',
        config_modification={'use_augmentation': False},
        description='No data augmentation'
    )
    
    # Basic augmentation
    study.add_experiment(
        name='basic_augmentation',
        config_modification={
            'use_augmentation': True,
            'augmentation': ['flip', 'rotation']
        },
        description='Basic augmentation: flips and rotations'
    )
    
    # Advanced augmentation
    study.add_experiment(
        name='advanced_augmentation',
        config_modification={
            'use_augmentation': True,
            'augmentation': ['flip', 'rotation', 'zoom', 'contrast', 'brightness']
        },
        description='Advanced augmentation with multiple transformations'
    )
    
    # With MixUp
    study.add_experiment(
        name='mixup_augmentation',
        config_modification={
            'use_augmentation': True,
            'augmentation': ['flip', 'rotation', 'mixup'],
            'mixup_alpha': 0.2
        },
        description='Augmentation with MixUp strategy'
    )
    
    # With CutMix
    study.add_experiment(
        name='cutmix_augmentation',
        config_modification={
            'use_augmentation': True,
            'augmentation': ['flip', 'rotation', 'cutmix'],
            'cutmix_alpha': 1.0
        },
        description='Augmentation with CutMix strategy'
    )
    
    return study


def calculate_segmentation_metrics(model, data, thresholds=None):
    """
    Calculate comprehensive segmentation metrics
    
    Args:
        model: Trained segmentation model
        data: Tuple of (X_val, y_val)
        thresholds: List of thresholds for binary classification
        
    Returns:
        Dictionary of metrics
    """
    from sklearn.metrics import jaccard_score, confusion_matrix

    def _dice_score(y_true, y_pred, smooth=1e-6):
        y_true = np.asarray(y_true).ravel().astype(np.float32)
        y_pred = np.asarray(y_pred).ravel().astype(np.float32)
        intersection = float(np.sum(y_true * y_pred))
        return (2.0 * intersection + smooth) / (float(np.sum(y_true) + np.sum(y_pred)) + smooth)

    X_val, y_val = data

    # Predict
    y_pred = model.predict(X_val)

    # Use threshold of 0.5 by default
    if thresholds is None:
        thresholds = [0.5]

    metrics = {}

    for threshold in thresholds:
        y_pred_binary = (y_pred >= threshold).astype(int)

        # Flatten for metric calculation
        y_val_flat = y_val.flatten()
        y_pred_flat = y_pred_binary.flatten()

        # Dice coefficient (local implementation; sklearn has no dice_score)
        dice = _dice_score(y_val_flat, y_pred_flat)
        
        # IoU (Jaccard index)
        iou = jaccard_score(y_val_flat, y_pred_flat)
        
        # Precision, Recall, F1
        tn, fp, fn, tp = confusion_matrix(y_val_flat, y_pred_flat).ravel()
        precision = tp / (tp + fp) if (tp + fp) > 0 else 0
        recall = tp / (tp + fn) if (tp + fn) > 0 else 0
        f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0
        
        # Specificity
        specificity = tn / (tn + fp) if (tn + fp) > 0 else 0
        
        metrics[f'dice_t{threshold}'] = float(dice)
        metrics[f'iou_t{threshold}'] = float(iou)
        metrics[f'precision_t{threshold}'] = float(precision)
        metrics[f'recall_t{threshold}'] = float(recall)
        metrics[f'f1_t{threshold}'] = float(f1)
        metrics[f'specificity_t{threshold}'] = float(specificity)
    
    return metrics