"""
Utility Functions for Local Learning

Helper functions for visualization, analysis, and debugging.
"""

import torch
import numpy as np
import matplotlib.pyplot as plt
from typing import Dict, List, Optional


def apply_neuromodulation(gradients: Dict[str, torch.Tensor], 
                         modulators: Dict[str, float]) -> Dict[str, torch.Tensor]:
    """
    Apply neuromodulator scaling to gradients.
    
    Args:
        gradients: Dict mapping parameter names to gradients
        modulators: Dict of neuromodulator values
    
    Returns:
        Modulated gradients
    """
    modulated = {}
    
    for name, grad in gradients.items():
        if grad is not None:
            # Apply dopamine (reward-based scaling)
            scaled_grad = grad * modulators['dopamine']
            
            # Apply serotonin (stability)
            scaled_grad *= modulators['serotonin']
            
            # Apply norepinephrine (attention boost)
            if modulators['norepinephrine'] > 1.5:
                scaled_grad *= 1.2
            
            modulated[name] = scaled_grad
        else:
            modulated[name] = None
    
    return modulated


def compute_layer_importance(model: torch.nn.Module, 
                            inputs: torch.Tensor,
                            targets: torch.Tensor,
                            loss_fn) -> Dict[str, float]:
    """
    Compute importance of each layer based on gradient magnitude.
    
    Args:
        model: PyTorch model
        inputs: Input batch
        targets: Target labels
        loss_fn: Loss function
    
    Returns:
        Dict mapping layer names to importance scores
    """
    model.eval()
    
    # Forward pass
    outputs = model(inputs)
    loss = loss_fn(outputs, targets)
    
    # Backward pass
    model.zero_grad()
    loss.backward()
    
    # Compute gradient norms per layer
    importance = {}
    
    for name, param in model.named_parameters():
        if param.grad is not None:
            grad_norm = param.grad.norm().item()
            layer_name = name.split('.')[0]  # Get top-level layer name
            
            if layer_name in importance:
                importance[layer_name] += grad_norm
            else:
                importance[layer_name] = grad_norm
    
    # Normalize
    total = sum(importance.values())
    if total > 0:
        importance = {k: v/total for k, v in importance.items()}
    
    return importance


def visualize_neuromodulators(history: Dict[str, List[float]], 
                              save_path: Optional[str] = None):
    """
    Visualize neuromodulator levels over training.
    
    Args:
        history: Dict with lists of neuromodulator values
        save_path: Optional path to save figure
    """
    fig, axes = plt.subplots(2, 2, figsize=(14, 10))
    
    steps = list(range(len(history['dopamine'])))
    
    # Dopamine
    axes[0, 0].plot(steps, history['dopamine'], 'r-', linewidth=2)
    axes[0, 0].set_title('Dopamine (Reward)', fontsize=14, fontweight='bold')
    axes[0, 0].set_xlabel('Step')
    axes[0, 0].set_ylabel('Level')
    axes[0, 0].grid(True, alpha=0.3)
    axes[0, 0].axhline(y=1.5, color='g', linestyle='--', alpha=0.5, label='High')
    axes[0, 0].axhline(y=1.0, color='y', linestyle='--', alpha=0.5, label='Normal')
    axes[0, 0].legend()
    
    # Serotonin
    axes[0, 1].plot(steps, history['serotonin'], 'g-', linewidth=2)
    axes[0, 1].set_title('Serotonin (Stability)', fontsize=14, fontweight='bold')
    axes[0, 1].set_xlabel('Step')
    axes[0, 1].set_ylabel('Level')
    axes[0, 1].grid(True, alpha=0.3)
    
    # Norepinephrine
    axes[1, 0].plot(steps, history['norepinephrine'], 'b-', linewidth=2)
    axes[1, 0].set_title('Norepinephrine (Attention)', fontsize=14, fontweight='bold')
    axes[1, 0].set_xlabel('Step')
    axes[1, 0].set_ylabel('Level')
    axes[1, 0].grid(True, alpha=0.3)
    
    # Acetylcholine
    axes[1, 1].plot(steps, history['acetylcholine'], 'purple', linewidth=2)
    axes[1, 1].set_title('Acetylcholine (Plasticity)', fontsize=14, fontweight='bold')
    axes[1, 1].set_xlabel('Step')
    axes[1, 1].set_ylabel('Level')
    axes[1, 1].grid(True, alpha=0.3)
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=150, bbox_inches='tight')
        print(f"✅ Saved neuromodulator visualization to {save_path}")
    else:
        plt.show()
    
    plt.close()


def visualize_training_progress(history: Dict[str, List[float]],
                                save_path: Optional[str] = None):
    """
    Visualize training progress (loss, task loss, pred loss).
    
    Args:
        history: Dict with training metrics
        save_path: Optional path to save figure
    """
    fig, axes = plt.subplots(1, 2, figsize=(14, 5))
    
    steps = list(range(len(history['loss'])))
    
    # Total loss
    axes[0].plot(steps, history['loss'], 'b-', linewidth=2, label='Total Loss')
    axes[0].set_title('Training Loss', fontsize=14, fontweight='bold')
    axes[0].set_xlabel('Step')
    axes[0].set_ylabel('Loss')
    axes[0].grid(True, alpha=0.3)
    axes[0].legend()
    
    # Task vs Prediction loss
    axes[1].plot(steps, history['task_loss'], 'r-', linewidth=2, label='Task Loss')
    axes[1].plot(steps, history['pred_loss'], 'g-', linewidth=2, label='Prediction Loss')
    axes[1].set_title('Task vs Prediction Loss', fontsize=14, fontweight='bold')
    axes[1].set_xlabel('Step')
    axes[1].set_ylabel('Loss')
    axes[1].grid(True, alpha=0.3)
    axes[1].legend()
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=150, bbox_inches='tight')
        print(f"✅ Saved training progress to {save_path}")
    else:
        plt.show()
    
    plt.close()


def compare_models(model1_history: Dict, model2_history: Dict,
                   model1_name: str = "Model 1",
                   model2_name: str = "Model 2",
                   save_path: Optional[str] = None):
    """
    Compare two models' training histories.
    
    Args:
        model1_history: First model's history
        model2_history: Second model's history
        model1_name: Name for first model
        model2_name: Name for second model
        save_path: Optional path to save figure
    """
    fig, axes = plt.subplots(2, 2, figsize=(14, 10))
    
    steps1 = list(range(len(model1_history['loss'])))
    steps2 = list(range(len(model2_history['loss'])))
    
    # Loss comparison
    axes[0, 0].plot(steps1, model1_history['loss'], 'b-', linewidth=2, label=model1_name)
    axes[0, 0].plot(steps2, model2_history['loss'], 'r-', linewidth=2, label=model2_name)
    axes[0, 0].set_title('Loss Comparison', fontsize=14, fontweight='bold')
    axes[0, 0].set_xlabel('Step')
    axes[0, 0].set_ylabel('Loss')
    axes[0, 0].legend()
    axes[0, 0].grid(True, alpha=0.3)
    
    # Dopamine comparison
    if 'dopamine' in model1_history and 'dopamine' in model2_history:
        axes[0, 1].plot(steps1, model1_history['dopamine'], 'b-', linewidth=2, label=model1_name)
        axes[0, 1].plot(steps2, model2_history['dopamine'], 'r-', linewidth=2, label=model2_name)
        axes[0, 1].set_title('Dopamine Comparison', fontsize=14, fontweight='bold')
        axes[0, 1].set_xlabel('Step')
        axes[0, 1].set_ylabel('Dopamine Level')
        axes[0, 1].legend()
        axes[0, 1].grid(True, alpha=0.3)
    
    # Task loss comparison
    if 'task_loss' in model1_history and 'task_loss' in model2_history:
        axes[1, 0].plot(steps1, model1_history['task_loss'], 'b-', linewidth=2, label=model1_name)
        axes[1, 0].plot(steps2, model2_history['task_loss'], 'r-', linewidth=2, label=model2_name)
        axes[1, 0].set_title('Task Loss Comparison', fontsize=14, fontweight='bold')
        axes[1, 0].set_xlabel('Step')
        axes[1, 0].set_ylabel('Task Loss')
        axes[1, 0].legend()
        axes[1, 0].grid(True, alpha=0.3)
    
    # Summary stats
    axes[1, 1].axis('off')
    
    summary = f"📊 COMPARISON SUMMARY\n\n"
    summary += f"{model1_name}:\n"
    summary += f"  Initial Loss: {model1_history['loss'][0]:.4f}\n"
    summary += f"  Final Loss: {model1_history['loss'][-1]:.4f}\n"
    summary += f"  Improvement: {(1 - model1_history['loss'][-1]/model1_history['loss'][0])*100:.1f}%\n\n"
    
    summary += f"{model2_name}:\n"
    summary += f"  Initial Loss: {model2_history['loss'][0]:.4f}\n"
    summary += f"  Final Loss: {model2_history['loss'][-1]:.4f}\n"
    summary += f"  Improvement: {(1 - model2_history['loss'][-1]/model2_history['loss'][0])*100:.1f}%\n\n"
    
    winner = model1_name if model1_history['loss'][-1] < model2_history['loss'][-1] else model2_name
    summary += f"🏆 Winner: {winner}"
    
    axes[1, 1].text(0.1, 0.5, summary, fontsize=11,
                   verticalalignment='center', family='monospace',
                   bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=150, bbox_inches='tight')
        print(f"✅ Saved comparison to {save_path}")
    else:
        plt.show()
    
    plt.close()


def print_comparison_table(v1_results: Dict[str, float], 
                           v2_results: Dict[str, float]):
    """
    Print a comparison table of two training runs.
    
    Args:
        v1_results: Results from version 1
        v2_results: Results from version 2
    """
    print("\n" + "="*80)
    print("📊 TRAINING COMPARISON TABLE")
    print("="*80)
    
    print(f"\n{'Metric':<30} {'v1.0':<20} {'v2.0':<20} {'Improvement':<10}")
    print("-"*80)
    
    # Loss metrics
    print(f"{'Initial Loss':<30} {v1_results.get('initial_loss', 0):<20.4f} {v2_results.get('initial_loss', 0):<20.4f} {'-':<10}")
    print(f"{'Final Loss':<30} {v1_results.get('final_loss', 0):<20.4f} {v2_results.get('final_loss', 0):<20.4f} {((v1_results.get('final_loss', 0) - v2_results.get('final_loss', 0))/v1_results.get('final_loss', 1)*100):<10.1f}%")
    print(f"{'Total Improvement %':<30} {v1_results.get('improvement_pct', 0):<20.1f}% {v2_results.get('improvement_pct', 0):<20.1f}% {(v2_results.get('improvement_pct', 0) - v1_results.get('improvement_pct', 0)):<10.1f}%")
    
    # Neuromodulators
    if 'avg_dopamine' in v1_results and 'avg_dopamine' in v2_results:
        print(f"\n{'Avg Dopamine':<30} {v1_results.get('avg_dopamine', 0):<20.3f} {v2_results.get('avg_dopamine', 0):<20.3f} {'-':<10}")
        print(f"{'Avg Serotonin':<30} {v1_results.get('avg_serotonin', 0):<20.3f} {v2_results.get('avg_serotonin', 0):<20.3f} {'-':<10}")
    
    # Features
    print(f"\n{'Predictive Coding':<30} {v1_results.get('predictive_coding', 'No'):<20} {v2_results.get('predictive_coding', 'No'):<20} {'-':<10}")
    print(f"{'Communication Strength':<30} {v1_results.get('comm_strength', 0):<20.3f} {v2_results.get('comm_strength', 0):<20.3f} {'-':<10}")
    print(f"{'R_signal':<30} {v1_results.get('r_signal', 0):<20} {v2_results.get('r_signal', 0):<20} {'-':<10}")
    
    print("\n" + "="*80)
    
    # Determine winner
    if v2_results.get('final_loss', float('inf')) < v1_results.get('final_loss', float('inf')):
        print("🏆 WINNER: v2.0 (Better final loss)")
    else:
        print("🏆 WINNER: v1.0 (Better final loss)")
    
    print("="*80 + "\n")


if __name__ == "__main__":
    print("🧪 Testing Utility Functions\n")
    
    # Create dummy history
    import numpy as np
    
    history = {
        'loss': list(np.linspace(3.0, 1.5, 100)),
        'task_loss': list(np.linspace(2.8, 1.3, 100)),
        'pred_loss': list(np.linspace(0.5, 0.2, 100)),
        'dopamine': list(np.random.uniform(0.8, 1.8, 100)),
        'serotonin': list(np.random.uniform(1.0, 1.5, 100)),
        'norepinephrine': list(np.random.uniform(0.8, 2.0, 100)),
        'acetylcholine': list(np.random.uniform(0.7, 1.5, 100)),
    }
    
    print("✅ Visualizing neuromodulators...")
    visualize_neuromodulators(history, save_path='test_neuromodulators.png')
    
    print("✅ Visualizing training progress...")
    visualize_training_progress(history, save_path='test_progress.png')
    
    # Test comparison table
    v1_results = {
        'initial_loss': 7.18,
        'final_loss': 3.49,
        'improvement_pct': 51.4,
        'avg_dopamine': 1.004,
        'avg_serotonin': 1.500,
        'predictive_coding': 'No',
        'comm_strength': 0.167,
        'r_signal': 64,
    }
    
    v2_results = {
        'initial_loss': 5.22,
        'final_loss': 2.33,
        'improvement_pct': 55.3,
        'avg_dopamine': 0.928,
        'avg_serotonin': 1.500,
        'predictive_coding': 'Yes',
        'comm_strength': 0.167,
        'r_signal': 96,
    }
    
    print_comparison_table(v1_results, v2_results)
    
    print("\n✅ All utility tests passed!")