"""
Neural Network Designer

Interactive neural network designer with AI-powered architecture recommendations,
support for various layer types, and automatic configuration suggestions.
"""

import numpy as np
import pandas as pd
from typing import Dict, List, Any, Optional, Tuple, Union
import json
from dataclasses import dataclass, asdict
from enum import Enum
import warnings


class LayerType(Enum):
    """Supported layer types"""
    DENSE = "dense"
    CONV2D = "conv2d"
    MAXPOOL2D = "maxpool2d"
    LSTM = "lstm"
    GRU = "gru"
    DROPOUT = "dropout"
    BATCH_NORM = "batch_norm"
    FLATTEN = "flatten"
    ACTIVATION = "activation"


class ActivationType(Enum):
    """Supported activation functions"""
    RELU = "relu"
    SIGMOID = "sigmoid"
    TANH = "tanh"
    SOFTMAX = "softmax"
    LEAKY_RELU = "leaky_relu"
    ELU = "elu"
    GELU = "gelu"


class OptimizerType(Enum):
    """Supported optimizers"""
    ADAM = "adam"
    SGD = "sgd"
    RMSPROP = "rmsprop"
    ADAGRAD = "adagrad"
    ADADELTA = "adadelta"


@dataclass
class LayerConfig:
    """Configuration for a neural network layer"""
    layer_type: LayerType
    name: str
    units: Optional[int] = None
    activation: Optional[ActivationType] = None
    input_shape: Optional[Tuple[int, ...]] = None
    filters: Optional[int] = None
    kernel_size: Optional[Tuple[int, int]] = None
    strides: Optional[Tuple[int, int]] = None
    padding: Optional[str] = None
    pool_size: Optional[Tuple[int, int]] = None
    rate: Optional[float] = None  # For dropout
    return_sequences: Optional[bool] = None  # For LSTM/GRU
    params: Optional[Dict[str, Any]] = None
    
    def __post_init__(self):
        if self.params is None:
            self.params = {}


@dataclass
class NetworkConfig:
    """Complete neural network configuration"""
    name: str
    layers: List[LayerConfig]
    optimizer: OptimizerType
    learning_rate: float
    loss_function: str
    metrics: List[str]
    input_shape: Tuple[int, ...]
    output_shape: Tuple[int, ...]
    problem_type: str  # 'classification', 'regression', 'autoencoder'
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for serialization"""
        return {
            'name': self.name,
            'layers': [asdict(layer) for layer in self.layers],
            'optimizer': self.optimizer.value,
            'learning_rate': self.learning_rate,
            'loss_function': self.loss_function,
            'metrics': self.metrics,
            'input_shape': self.input_shape,
            'output_shape': self.output_shape,
            'problem_type': self.problem_type
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'NetworkConfig':
        """Create from dictionary"""
        layers = []
        for layer_data in data['layers']:
            layer_data['layer_type'] = LayerType(layer_data['layer_type'])
            if layer_data.get('activation'):
                layer_data['activation'] = ActivationType(layer_data['activation'])
            layers.append(LayerConfig(**layer_data))
        
        return cls(
            name=data['name'],
            layers=layers,
            optimizer=OptimizerType(data['optimizer']),
            learning_rate=data['learning_rate'],
            loss_function=data['loss_function'],
            metrics=data['metrics'],
            input_shape=tuple(data['input_shape']),
            output_shape=tuple(data['output_shape']),
            problem_type=data['problem_type']
        )


class NeuralNetworkDesigner:
    """
    Interactive neural network designer with AI-powered recommendations
    for architecture design, hyperparameters, and optimization.
    """
    
    def __init__(self):
        self.current_config = None
        self.design_history = []
        self.recommendations_cache = {}
        
    def create_network(self, name: str, input_shape: Tuple[int, ...], 
                      output_shape: Tuple[int, ...], problem_type: str = 'classification') -> NetworkConfig:
        """
        Create a new neural network configuration
        
        Args:
            name: Network name
            input_shape: Input tensor shape
            output_shape: Output tensor shape
            problem_type: Type of problem ('classification', 'regression', 'autoencoder')
            
        Returns:
            NetworkConfig object
        """
        config = NetworkConfig(
            name=name,
            layers=[],
            optimizer=OptimizerType.ADAM,
            learning_rate=0.001,
            loss_function=self._get_default_loss_function(problem_type),
            metrics=self._get_default_metrics(problem_type),
            input_shape=input_shape,
            output_shape=output_shape,
            problem_type=problem_type
        )
        
        self.current_config = config
        return config
    
    def add_layer(self, layer_config: LayerConfig) -> NetworkConfig:
        """Add a layer to the current network"""
        if not self.current_config:
            raise ValueError("No network configuration exists. Create one first.")
        
        self.current_config.layers.append(layer_config)
        return self.current_config
    
    def add_dense_layer(self, units: int, activation: ActivationType = ActivationType.RELU, 
                       name: Optional[str] = None) -> NetworkConfig:
        """Add a dense layer with sensible defaults"""
        if not name:
            name = f"dense_{len(self.current_config.layers)}"
        
        layer = LayerConfig(
            layer_type=LayerType.DENSE,
            name=name,
            units=units,
            activation=activation
        )
        
        return self.add_layer(layer)
    
    def add_conv2d_layer(self, filters: int, kernel_size: Tuple[int, int] = (3, 3),
                        activation: ActivationType = ActivationType.RELU,
                        strides: Tuple[int, int] = (1, 1),
                        padding: str = 'same', name: Optional[str] = None) -> NetworkConfig:
        """Add a 2D convolutional layer"""
        if not name:
            name = f"conv2d_{len(self.current_config.layers)}"
        
        layer = LayerConfig(
            layer_type=LayerType.CONV2D,
            name=name,
            filters=filters,
            kernel_size=kernel_size,
            activation=activation,
            strides=strides,
            padding=padding
        )
        
        return self.add_layer(layer)
    
    def add_lstm_layer(self, units: int, return_sequences: bool = False,
                      activation: ActivationType = ActivationType.TANH,
                      name: Optional[str] = None) -> NetworkConfig:
        """Add an LSTM layer"""
        if not name:
            name = f"lstm_{len(self.current_config.layers)}"
        
        layer = LayerConfig(
            layer_type=LayerType.LSTM,
            name=name,
            units=units,
            activation=activation,
            return_sequences=return_sequences
        )
        
        return self.add_layer(layer)
    
    def add_dropout_layer(self, rate: float = 0.5, name: Optional[str] = None) -> NetworkConfig:
        """Add a dropout layer"""
        if not name:
            name = f"dropout_{len(self.current_config.layers)}"
        
        layer = LayerConfig(
            layer_type=LayerType.DROPOUT,
            name=name,
            rate=rate
        )
        
        return self.add_layer(layer)
    
    def recommend_architecture(self, X: pd.DataFrame, y: pd.Series, 
                             data_type: str = 'tabular') -> NetworkConfig:
        """
        AI-powered architecture recommendation based on data characteristics
        
        Args:
            X: Feature matrix
            y: Target variable
            data_type: Type of data ('tabular', 'image', 'sequence', 'text')
            
        Returns:
            Recommended NetworkConfig
        """
        n_samples, n_features = X.shape
        n_classes = len(np.unique(y)) if len(np.unique(y)) < 20 else 1
        
        # Determine problem type
        if pd.api.types.is_numeric_dtype(y) and n_classes > 20:
            problem_type = 'regression'
        else:
            problem_type = 'classification'
        
        # Get input and output shapes
        if data_type == 'tabular':
            input_shape = (n_features,)
            output_shape = (n_classes,) if problem_type == 'classification' else (1,)
        elif data_type == 'image':
            # Assuming square images, need to be specified properly
            input_shape = (32, 32, 3)  # Default, should be adjusted
            output_shape = (n_classes,) if problem_type == 'classification' else (1,)
        elif data_type == 'sequence':
            # Assuming sequence data
            input_shape = (n_features, 1)
            output_shape = (n_classes,) if problem_type == 'classification' else (1,)
        else:
            input_shape = (n_features,)
            output_shape = (n_classes,) if problem_type == 'classification' else (1,)
        
        # Create recommended network
        config = self.create_network(
            name=f"recommended_{data_type}_network",
            input_shape=input_shape,
            output_shape=output_shape,
            problem_type=problem_type
        )
        
        # Add layers based on data type and problem characteristics
        if data_type == 'tabular':
            self._recommend_tabular_architecture(config, n_features, n_classes, n_samples)
        elif data_type == 'image':
            self._recommend_image_architecture(config, n_classes)
        elif data_type == 'sequence':
            self._recommend_sequence_architecture(config, n_features, n_classes)
        else:
            self._recommend_tabular_architecture(config, n_features, n_classes, n_samples)
        
        # Set optimizer and learning rate
        config.optimizer, config.learning_rate = self._recommend_optimizer(problem_type, n_samples)
        
        self.current_config = config
        return config
    
    def _recommend_tabular_architecture(self, config: NetworkConfig, n_features: int, 
                                       n_classes: int, n_samples: int):
        """Recommend architecture for tabular data"""
        
        # Input layer
        first_layer_units = min(max(n_features * 2, 64), 512)
        config.add_dense_layer(first_layer_units, ActivationType.RELU, "input_dense")
        config.add_dropout_layer(0.2, "input_dropout")
        
        # Hidden layers based on data size
        if n_samples > 10000:
            # Large dataset - deeper network
            config.add_dense_layer(first_layer_units // 2, ActivationType.RELU, "hidden_1")
            config.add_dropout_layer(0.3, "dropout_1")
            config.add_dense_layer(first_layer_units // 4, ActivationType.RELU, "hidden_2")
            config.add_dropout_layer(0.3, "dropout_2")
        elif n_samples > 1000:
            # Medium dataset - moderate network
            config.add_dense_layer(first_layer_units // 2, ActivationType.RELU, "hidden_1")
            config.add_dropout_layer(0.3, "dropout_1")
        else:
            # Small dataset - shallow network
            pass  # Only input layer
        
        # Output layer
        if config.problem_type == 'classification':
            if n_classes == 2:
                config.add_dense_layer(1, ActivationType.SIGMOID, "output")
            else:
                config.add_dense_layer(n_classes, ActivationType.SOFTMAX, "output")
        else:
            config.add_dense_layer(1, ActivationType.LINEAR if hasattr(ActivationType, 'LINEAR') else ActivationType.RELU, "output")
    
    def _recommend_image_architecture(self, config: NetworkConfig, n_classes: int):
        """Recommend architecture for image data"""
        
        # Convolutional layers
        config.add_conv2d_layer(32, (3, 3), ActivationType.RELU, name="conv1")
        config.add_conv2d_layer(64, (3, 3), ActivationType.RELU, name="conv2")
        
        # Add max pooling (would need to implement in actual framework)
        # config.add_layer(LayerConfig(layer_type=LayerType.MAXPOOL2D, name="pool1", pool_size=(2, 2)))
        
        config.add_dropout_layer(0.25, "dropout1")
        
        # Flatten layer (would need to implement)
        # config.add_layer(LayerConfig(layer_type=LayerType.FLATTEN, name="flatten"))
        
        # Dense layers
        config.add_dense_layer(128, ActivationType.RELU, "dense1")
        config.add_dropout_layer(0.5, "dropout2")
        
        # Output layer
        if config.problem_type == 'classification':
            if n_classes == 2:
                config.add_dense_layer(1, ActivationType.SIGMOID, "output")
            else:
                config.add_dense_layer(n_classes, ActivationType.SOFTMAX, "output")
        else:
            config.add_dense_layer(1, ActivationType.RELU, "output")
    
    def _recommend_sequence_architecture(self, config: NetworkConfig, n_features: int, n_classes: int):
        """Recommend architecture for sequence data"""
        
        # LSTM layers
        config.add_lstm_layer(64, return_sequences=True, name="lstm1")
        config.add_dropout_layer(0.2, "dropout1")
        config.add_lstm_layer(32, return_sequences=False, name="lstm2")
        config.add_dropout_layer(0.2, "dropout2")
        
        # Dense layers
        config.add_dense_layer(16, ActivationType.RELU, "dense1")
        
        # Output layer
        if config.problem_type == 'classification':
            if n_classes == 2:
                config.add_dense_layer(1, ActivationType.SIGMOID, "output")
            else:
                config.add_dense_layer(n_classes, ActivationType.SOFTMAX, "output")
        else:
            config.add_dense_layer(1, ActivationType.RELU, "output")
    
    def _recommend_optimizer(self, problem_type: str, n_samples: int) -> Tuple[OptimizerType, float]:
        """Recommend optimizer and learning rate"""
        
        if n_samples > 50000:
            # Large dataset - Adam with lower learning rate
            return OptimizerType.ADAM, 0.0001
        elif n_samples > 5000:
            # Medium dataset - Adam
            return OptimizerType.ADAM, 0.001
        else:
            # Small dataset - SGD with momentum
            return OptimizerType.SGD, 0.01
    
    def _get_default_loss_function(self, problem_type: str) -> str:
        """Get default loss function based on problem type"""
        if problem_type == 'classification':
            return 'categorical_crossentropy'
        elif problem_type == 'regression':
            return 'mse'
        elif problem_type == 'autoencoder':
            return 'mse'
        else:
            return 'mse'
    
    def _get_default_metrics(self, problem_type: str) -> List[str]:
        """Get default metrics based on problem type"""
        if problem_type == 'classification':
            return ['accuracy']
        elif problem_type == 'regression':
            return ['mae', 'mse']
        elif problem_type == 'autoencoder':
            return ['mse']
        else:
            return ['mse']
    
    def analyze_network_complexity(self, config: NetworkConfig) -> Dict[str, Any]:
        """Analyze the complexity of the network architecture"""
        
        total_params = 0
        layer_count = len(config.layers)
        dense_layers = sum(1 for layer in config.layers if layer.layer_type == LayerType.DENSE)
        conv_layers = sum(1 for layer in config.layers if layer.layer_type == LayerType.CONV2D)
        lstm_layers = sum(1 for layer in config.layers if layer.layer_type == LayerType.LSTM)
        dropout_layers = sum(1 for layer in config.layers if layer.layer_type == LayerType.DROPOUT)
        
        # Estimate parameters (simplified calculation)
        input_size = np.prod(config.input_shape)
        
        for i, layer in enumerate(config.layers):
            if layer.layer_type == LayerType.DENSE:
                if i == 0:
                    # First dense layer
                    params = input_size * layer.units + layer.units
                else:
                    # Subsequent dense layers
                    prev_units = config.layers[i-1].units if config.layers[i-1].layer_type == LayerType.DENSE else input_size
                    params = prev_units * layer.units + layer.units
                total_params += params
        
        complexity_score = layer_count * 10 + total_params / 1000
        
        return {
            'total_layers': layer_count,
            'dense_layers': dense_layers,
            'conv_layers': conv_layers,
            'lstm_layers': lstm_layers,
            'dropout_layers': dropout_layers,
            'estimated_parameters': total_params,
            'complexity_score': complexity_score,
            'complexity_level': self._get_complexity_level(complexity_score)
        }
    
    def _get_complexity_level(self, score: float) -> str:
        """Categorize network complexity"""
        if score < 50:
            return 'Simple'
        elif score < 200:
            return 'Moderate'
        elif score < 500:
            return 'Complex'
        else:
            return 'Very Complex'
    
    def get_optimization_recommendations(self, config: NetworkConfig) -> List[str]:
        """Get AI-powered optimization recommendations"""
        recommendations = []
        complexity = self.analyze_network_complexity(config)
        
        # Complexity-based recommendations
        if complexity['complexity_level'] == 'Very Complex':
            recommendations.append("Network is very complex. Consider simplifying to avoid overfitting")
            recommendations.append("Add more dropout layers or regularization")
        elif complexity['complexity_level'] == 'Simple':
            recommendations.append("Network is simple. Consider adding more layers for better representation")
        
        # Layer-specific recommendations
        if complexity['dropout_layers'] == 0 and complexity['total_layers'] > 3:
            recommendations.append("Consider adding dropout layers to prevent overfitting")
        
        if complexity['lstm_layers'] > 0 and complexity['dropout_layers'] < 2:
            recommendations.append("LSTM networks benefit from dropout layers between LSTM layers")
        
        # Optimizer recommendations
        if config.optimizer == OptimizerType.SGD and config.learning_rate > 0.01:
            recommendations.append("Consider reducing learning rate for SGD optimizer")
        
        if config.optimizer == OptimizerType.ADAM and config.learning_rate > 0.01:
            recommendations.append("Adam typically works better with learning rates <= 0.001")
        
        return recommendations
    
    def save_config(self, filepath: str, config: Optional[NetworkConfig] = None):
        """Save network configuration to file"""
        if config is None:
            config = self.current_config
        
        if not config:
            raise ValueError("No configuration to save")
        
        with open(filepath, 'w') as f:
            json.dump(config.to_dict(), f, indent=2)
    
    def load_config(self, filepath: str) -> NetworkConfig:
        """Load network configuration from file"""
        with open(filepath, 'r') as f:
            data = json.load(f)
        
        config = NetworkConfig.from_dict(data)
        self.current_config = config
        return config
    
    def export_to_keras(self, config: NetworkConfig) -> str:
        """Export configuration to Keras code"""
        code_lines = []
        code_lines.append("import tensorflow as tf")
        code_lines.append("from tensorflow import keras")
        code_lines.append("from tensorflow.keras import layers")
        code_lines.append("")
        code_lines.append(f"def create_{config.name.lower().replace(' ', '_')}():")
        code_lines.append("    model = keras.Sequential()")
        code_lines.append("")
        
        for layer in config.layers:
            if layer.layer_type == LayerType.DENSE:
                if layer.input_shape:
                    code_lines.append(f"    model.add(layers.Dense({layer.units}, activation='{layer.activation.value}', input_shape={layer.input_shape}, name='{layer.name}'))")
                else:
                    code_lines.append(f"    model.add(layers.Dense({layer.units}, activation='{layer.activation.value}', name='{layer.name}'))")
            elif layer.layer_type == LayerType.DROPOUT:
                code_lines.append(f"    model.add(layers.Dropout({layer.rate}, name='{layer.name}'))")
            elif layer.layer_type == LayerType.CONV2D:
                code_lines.append(f"    model.add(layers.Conv2D({layer.filters}, {layer.kernel_size}, activation='{layer.activation.value}', strides={layer.strides}, padding='{layer.padding}', name='{layer.name}'))")
            elif layer.layer_type == LayerType.LSTM:
                code_lines.append(f"    model.add(layers.LSTM({layer.units}, return_sequences={layer.return_sequences}, name='{layer.name}'))")
        
        code_lines.append("")
        code_lines.append(f"    model.compile(")
        code_lines.append(f"        optimizer='{config.optimizer.value}',")
        code_lines.append(f"        learning_rate={config.learning_rate},")
        code_lines.append(f"        loss='{config.loss_function}',")
        code_lines.append(f"        metrics={config.metrics}")
        code_lines.append(f"    )")
        code_lines.append("")
        code_lines.append("    return model")
        code_lines.append("")
        
        return "\n".join(code_lines)
    
    def get_network_summary(self, config: Optional[NetworkConfig] = None) -> Dict[str, Any]:
        """Get a comprehensive summary of the network"""
        if config is None:
            config = self.current_config
        
        if not config:
            raise ValueError("No network configuration available")
        
        summary = {
            'name': config.name,
            'problem_type': config.problem_type,
            'input_shape': config.input_shape,
            'output_shape': config.output_shape,
            'optimizer': config.optimizer.value,
            'learning_rate': config.learning_rate,
            'loss_function': config.loss_function,
            'metrics': config.metrics,
            'layers': []
        }
        
        for layer in config.layers:
            layer_info = {
                'name': layer.name,
                'type': layer.layer_type.value,
                'parameters': {}
            }
            
            if layer.units:
                layer_info['parameters']['units'] = layer.units
            if layer.activation:
                layer_info['parameters']['activation'] = layer.activation.value
            if layer.filters:
                layer_info['parameters']['filters'] = layer.filters
            if layer.kernel_size:
                layer_info['parameters']['kernel_size'] = layer.kernel_size
            if layer.rate:
                layer_info['parameters']['dropout_rate'] = layer.rate
            if layer.return_sequences is not None:
                layer_info['parameters']['return_sequences'] = layer.return_sequences
            
            summary['layers'].append(layer_info)
        
        # Add complexity analysis
        summary['complexity'] = self.analyze_network_complexity(config)
        
        return summary
