from .lib.grnexus_core import *
from .lib.grnexus_activations import *
from .lib.grnexus_layers import *
from .lib.grnexus_normalization import *
from .lib.grnexus_numeric_proccessing import *
from .lib.grnexus_callbacks import *
from .lib.grnexus_machine_learning import *
import json
import time
import zlib
import random
from typing import List, Dict, Any, Tuple, Optional

class NeuralNetwork:
    """Complete Neural Network implementation with training, prediction, and model persistence."""
    
    def __init__(self, loss='mse', optimizer='sgd', learning_rate=0.01, name='model'):
        self.layers = []
        self.loss_function = loss
        self.optimizer = optimizer
        self.learning_rate = learning_rate
        self.name = name
        self.history = {'loss': [], 'accuracy': [], 'val_loss': [], 'val_accuracy': [], 'lr': []}
        self.optimizer_state = {}
        self.callbacks = []
        self.stop_training = False
    
    def add(self, layer):
        """Add a layer to the network."""
        self.layers.append(layer)
    
    def compile(self, loss=None, optimizer=None, learning_rate=None, metrics=None):
        """Configure the model for training."""
        if loss:
            self.loss_function = loss
        if optimizer:
            self.optimizer = optimizer
        if learning_rate:
            self.learning_rate = learning_rate
        self.metrics = metrics or []
    
    def forward(self, input_data, training=True):
        """Forward propagation through all layers."""
        output = input_data
        for layer in self.layers:
            if isinstance(layer, (DropoutLayer, BatchNormLayer)):
                output = layer.forward(output, training=training)
            else:
                output = layer.forward(output)
        return output
    
    def backward(self, gradient):
        """Backward propagation through all layers."""
        grad = gradient
        for layer in reversed(self.layers):
            if layer.trainable:
                grad = layer.backward(grad)
                # Update parameters if layer has them
                if hasattr(layer, 'weight_gradient') and layer.weight_gradient is not None:
                    layer.update_parameters(self.learning_rate)
            else:
                grad = layer.backward(grad)
        return grad
    
    def train(self, x_train, y_train, epochs=10, batch_size=32, verbose=True,
              validation_data=None, callbacks=None, shuffle=True):
        """Train the neural network."""
        self.callbacks = callbacks or []
        self.stop_training = False
        
        num_samples = len(x_train)
        num_batches = (num_samples + batch_size - 1) // batch_size
        
        # Validation data
        x_val, y_val = validation_data if validation_data else (None, None)
        
        # Callbacks - on_train_begin
        callback_logs = {'epochs': epochs, 'steps': num_batches, 'model': self}
        for cb in self.callbacks:
            cb.on_train_begin(callback_logs)
        
        for epoch in range(epochs):
            if self.stop_training:
                break
            
            # Callbacks - on_epoch_begin
            epoch_logs = {'epoch': epoch, 'model': self}
            for cb in self.callbacks:
                cb.on_epoch_begin(epoch, epoch_logs)
            
            epoch_loss = 0.0
            correct = 0
            total = 0
            
            # Shuffle data
            indices = list(range(num_samples))
            if shuffle:
                random.shuffle(indices)
            
            for batch_idx in range(num_batches):
                batch_start = batch_idx * batch_size
                batch_end = min(batch_start + batch_size, num_samples)
                batch_indices = indices[batch_start:batch_end]
                
                x_batch = [x_train[i] for i in batch_indices]
                y_batch = [y_train[i] for i in batch_indices]
                
                # Callbacks - on_batch_begin
                batch_logs = {'batch': batch_idx, 'size': len(batch_indices)}
                for cb in self.callbacks:
                    cb.on_batch_begin(batch_idx, batch_logs)
                
                # Forward pass
                predictions = self.forward(x_batch, training=True)
                
                # Calculate loss
                loss = self._calculate_loss(predictions, y_batch)
                epoch_loss += loss
                
                # Calculate accuracy (for classification)
                if self.loss_function == 'cross_entropy':
                    correct += self._calculate_accuracy(predictions, y_batch)
                    total += len(batch_indices)
                
                # Backward pass
                gradient = self._calculate_gradient(predictions, y_batch)
                self.backward(gradient)
                
                # Callbacks - on_batch_end
                batch_logs['loss'] = loss
                for cb in self.callbacks:
                    cb.on_batch_end(batch_idx, batch_logs)
            
            avg_loss = epoch_loss / num_batches
            self.history['loss'].append(avg_loss)
            self.history['lr'].append(self.learning_rate)
            
            # Calculate validation metrics if validation data provided
            if validation_data:
                val_results = self.evaluate(x_val, y_val)
                self.history['val_loss'].append(val_results['loss'])
                if 'accuracy' in val_results:
                    self.history['val_accuracy'].append(val_results['accuracy'])
            
            # Prepare epoch end logs
            epoch_end_logs = {
                'epoch': epoch,
                'loss': avg_loss,
                'lr': self.learning_rate,
                'model': self
            }
            
            if self.loss_function == 'cross_entropy':
                accuracy = (correct / total * 100) if total > 0 else 0
                self.history['accuracy'].append(accuracy)
                epoch_end_logs['accuracy'] = accuracy
            
            if validation_data:
                epoch_end_logs['val_loss'] = self.history['val_loss'][-1]
                if self.history['val_accuracy']:
                    epoch_end_logs['val_accuracy'] = self.history['val_accuracy'][-1]
            
            # Callbacks - on_epoch_end
            for cb in self.callbacks:
                cb.on_epoch_end(epoch, epoch_end_logs)
            
            # Check for early stopping
            for cb in self.callbacks:
                if hasattr(cb, 'stop_training') and cb.stop_training:
                    self.stop_training = True
                    break
            
            # Verbose output
            if verbose:
                output_str = f"Epoch {epoch + 1}/{epochs} - Loss: {avg_loss:.6f}"
                if self.loss_function == 'cross_entropy':
                    output_str += f" - Accuracy: {accuracy:.2f}%"
                if validation_data:
                    output_str += f" - Val Loss: {self.history['val_loss'][-1]:.6f}"
                    if self.history['val_accuracy']:
                        output_str += f" - Val Accuracy: {self.history['val_accuracy'][-1]:.2f}%"
                if any(isinstance(cb, (LearningRateScheduler, ReduceLROnPlateau)) for cb in self.callbacks):
                    output_str += f" - LR: {self.learning_rate}"
                print(output_str)
        
        # Callbacks - on_train_end
        for cb in self.callbacks:
            cb.on_train_end({'model': self})
        
        return self.history
    
    def predict(self, input_data):
        """Make predictions on input data."""
        return self.forward(input_data, training=False)
    
    def evaluate(self, x_test, y_test):
        """Evaluate the model on test data."""
        predictions = self.predict(x_test)
        loss = self._calculate_loss(predictions, y_test)
        
        result = {'loss': loss}
        
        if self.loss_function == 'cross_entropy':
            correct = self._calculate_accuracy(predictions, y_test)
            accuracy = (correct / len(y_test) * 100) if len(y_test) > 0 else 0
            result['accuracy'] = accuracy
        
        return result
    
    def save(self, filepath):
        """Save the model to a .nexus file."""
        import os
        
        # Expandir ruta para manejar rutas relativas y absolutas
        expanded_path = os.path.abspath(filepath)
        
        # Crear directorio si no existe
        directory = os.path.dirname(expanded_path)
        os.makedirs(directory, exist_ok=True)
        
        model_data = {
            'version': '2.0',
            'framework': 'GRNexus',
            'language': 'Python',
            'name': self.name,
            'created_at': time.strftime('%Y-%m-%dT%H:%M:%S'),
            'architecture': self._serialize_architecture(),
            'loss_function': self.loss_function,
            'optimizer': self.optimizer,
            'learning_rate': self.learning_rate,
            'history': self.history,
            'metadata': {
                'total_params': self.count_params()['total'],
                'trainable_params': self.count_params()['trainable'],
                'layers_count': len(self.layers)
            }
        }
        
        json_data = json.dumps(model_data)
        compressed_data = zlib.compress(json_data.encode('utf-8'))
        
        with open(expanded_path, 'wb') as f:
            f.write(compressed_data)
        
        print(f"Model saved to {expanded_path}")
    
    @staticmethod
    def load(filepath):
        """Load a model from a .nexus file."""
        import os
        
        # Expandir ruta para manejar rutas relativas y absolutas
        expanded_path = os.path.abspath(filepath)
        
        if not os.path.exists(expanded_path):
            raise FileNotFoundError(f"Model file not found: {expanded_path}")
        
        with open(expanded_path, 'rb') as f:
            compressed_data = f.read()
        
        json_data = zlib.decompress(compressed_data).decode('utf-8')
        model_data = json.loads(json_data)
        
        # Verificar compatibilidad
        version = model_data.get('version', '1.0')
        framework = model_data.get('framework', 'GRNexus')
        source_lang = model_data.get('language', 'Unknown')
        
        print(f"Loading model: {framework} v{version} (created in {source_lang})")
        
        model = NeuralNetwork(
            loss=model_data['loss_function'],
            optimizer=model_data['optimizer'],
            learning_rate=model_data['learning_rate'],
            name=model_data.get('name', 'model')
        )
        
        model._deserialize_architecture(model_data['architecture'])
        model.history = model_data['history']
        
        print(f"Model loaded from {expanded_path}")
        if 'metadata' in model_data:
            print(f"  Total params: {model_data['metadata']['total_params']}")
            print(f"  Layers: {model_data['metadata']['layers_count']}")
        return model
    
    @staticmethod
    def inspect_model(filepath):
        """Inspecciona un modelo guardado sin cargarlo completamente"""
        import os
        
        # Expandir ruta para manejar rutas relativas y absolutas
        expanded_path = os.path.abspath(filepath)
        
        if not os.path.exists(expanded_path):
            print(f"Error: Model file not found: {expanded_path}")
            return
        
        with open(expanded_path, 'rb') as f:
            compressed_data = f.read()
        
        json_data = zlib.decompress(compressed_data).decode('utf-8')
        model_data = json.loads(json_data)
        
        print("\n" + "=" * 80)
        print(f"MODEL INSPECTION: {filepath}")
        print("=" * 80)
        print(f"Framework: {model_data.get('framework', 'GRNexus')}")
        print(f"Version: {model_data.get('version', '1.0')}")
        print(f"Language: {model_data.get('language', 'Unknown')}")
        print(f"Name: {model_data['name']}")
        print(f"Created: {model_data['created_at']}")
        print(f"Loss Function: {model_data['loss_function']}")
        print(f"Optimizer: {model_data['optimizer']}")
        print(f"Learning Rate: {model_data['learning_rate']}")
        
        if 'metadata' in model_data:
            print("\nMetadata:")
            print(f"  Total Parameters: {model_data['metadata']['total_params']}")
            print(f"  Trainable Parameters: {model_data['metadata']['trainable_params']}")
            print(f"  Layers Count: {model_data['metadata']['layers_count']}")
        
        print("\nArchitecture:")
        print("-" * 80)
        for idx, layer_data in enumerate(model_data['architecture']):
            layer_type = layer_data['type'].split('.')[-1]
            print(f"  Layer {idx + 1}: {layer_type}")
            if 'units' in layer_data:
                print(f"    Units: {layer_data['units']}")
            if 'activation' in layer_data and layer_data['activation']:
                print(f"    Activation: {layer_data['activation'].split('.')[-1]}")
            print(f"    Trainable: {layer_data['trainable']}")
        
        if 'history' in model_data and 'loss' in model_data['history']:
            print("\nTraining History:")
            print(f"  Epochs trained: {len(model_data['history']['loss'])}")
            print(f"  Final loss: {model_data['history']['loss'][-1]:.6f}")
            if 'accuracy' in model_data['history'] and model_data['history']['accuracy']:
                print(f"  Final accuracy: {model_data['history']['accuracy'][-1]:.2f}%")
        
        print("=" * 80 + "\n")
    
    def summary(self, line_length=80):
        """Print a summary of the model architecture."""
        print("\n" + "=" * line_length)
        print(f"Model: {self.name}")
        print("=" * line_length)
        print(f"{' ' * 20}Output Shape{' ' * 10}Param #")
        print("-" * line_length)
        
        total_params = 0
        trainable_params = 0
        
        for idx, layer in enumerate(self.layers):
            layer_name = layer.__class__.__name__
            if hasattr(layer, 'activation') and layer.activation:
                layer_name = f"{layer_name} ({layer.activation.__class__.__name__})"
            
            # Calculate output shape
            output_shape = self._get_layer_output_shape(layer, idx)
            
            # Calculate parameters
            params = self._count_layer_params(layer)
            total_params += params
            trainable_params += params if layer.trainable else 0
            
            # Format output
            layer_info = f"{layer_name} ({idx + 1})"
            layer_info = layer_info.ljust(30)
            layer_info += output_shape.ljust(25)
            layer_info += str(params)
            
            print(layer_info)
        
        print("=" * line_length)
        print(f"Total params: {total_params}")
        print(f"Trainable params: {trainable_params}")
        print(f"Non-trainable params: {total_params - trainable_params}")
        print("=" * line_length + "\n")
    
    def to_json_architecture(self):
        """Export model architecture to JSON."""
        return json.dumps({
            'name': self.name,
            'layers': [
                {
                    'index': idx,
                    'type': layer.__class__.__name__,
                    'config': self._layer_config(layer),
                    'trainable': layer.trainable
                }
                for idx, layer in enumerate(self.layers)
            ]
        }, indent=2)
    
    def plot_history(self, metrics=None):
        """Plot training history (text-based)."""
        metrics = metrics or ['loss', 'accuracy']
        
        print("\n" + "=" * 80)
        print("Training History")
        print("=" * 80)
        
        for metric in metrics:
            if metric not in self.history or not self.history[metric]:
                continue
            
            print(f"\n{metric.capitalize()}:")
            values = self.history[metric]
            min_val = min(values)
            max_val = max(values)
            range_val = max_val - min_val
            
            for epoch, val in enumerate(values):
                normalized = int(((val - min_val) / range_val * 50)) if range_val > 0 else 25
                bar = '█' * normalized
                print(f"Epoch {epoch + 1}: {bar} {val:.4f}")
        
        print("=" * 80 + "\n")
    
    def count_params(self):
        """Count total, trainable, and non-trainable parameters."""
        total = 0
        trainable = 0
        
        for layer in self.layers:
            params = self._count_layer_params(layer)
            total += params
            trainable += params if layer.trainable else 0
        
        return {
            'total': total,
            'trainable': trainable,
            'non_trainable': total - trainable
        }
    
    # Private methods
    
    def _calculate_loss(self, predictions, targets):
        """Calculate loss based on the loss function."""
        if self.loss_function == 'mse':
            return self._mse_loss(predictions, targets)
        elif self.loss_function == 'cross_entropy':
            return self._cross_entropy_loss(predictions, targets)
        else:
            raise ValueError(f"Unknown loss function: {self.loss_function}")
    
    def _calculate_gradient(self, predictions, targets):
        """Calculate gradient based on the loss function."""
        if self.loss_function == 'mse':
            return self._mse_gradient(predictions, targets)
        elif self.loss_function == 'cross_entropy':
            return self._cross_entropy_gradient(predictions, targets)
        else:
            raise ValueError(f"Unknown loss function: {self.loss_function}")
    
    def _mse_loss(self, predictions, targets):
        """Mean Squared Error loss."""
        total_loss = 0.0
        for pred, target in zip(predictions, targets):
            for p, t in zip(pred, target):
                diff = p - t
                total_loss += diff * diff
        return total_loss / (len(predictions) * len(predictions[0]))
    
    def _mse_gradient(self, predictions, targets):
        """MSE gradient."""
        gradients = []
        n = len(predictions) * len(predictions[0])
        for pred, target in zip(predictions, targets):
            grad = [2.0 * (p - t) / n for p, t in zip(pred, target)]
            gradients.append(grad)
        return gradients
    
    def _cross_entropy_loss(self, predictions, targets):
        """Cross-entropy loss."""
        import math
        total_loss = 0.0
        epsilon = 1e-15
        
        for pred, target in zip(predictions, targets):
            for p, t in zip(pred, target):
                p_clipped = max(min(p, 1.0 - epsilon), epsilon)
                total_loss -= t * math.log(p_clipped)
        
        return total_loss / len(predictions)
    
    def _cross_entropy_gradient(self, predictions, targets):
        """Cross-entropy gradient (for softmax + cross-entropy)."""
        gradients = []
        for pred, target in zip(predictions, targets):
            grad = [p - t for p, t in zip(pred, target)]
            gradients.append(grad)
        return gradients
    
    def _calculate_accuracy(self, predictions, targets):
        """Calculate classification accuracy."""
        correct = 0
        for pred, target in zip(predictions, targets):
            pred_class = pred.index(max(pred))
            true_class = target.index(max(target))
            if pred_class == true_class:
                correct += 1
        return correct
    
    def _serialize_architecture(self):
        """Serialize the model architecture to a dictionary."""
        architecture = []
        
        for layer in self.layers:
            layer_data = {
                'type': f"{layer.__class__.__module__}.{layer.__class__.__name__}",
                'trainable': layer.trainable
            }
            
            # Serialize layer-specific parameters
            if hasattr(layer, 'weights') and layer.weights is not None:
                layer_data['weights'] = layer.weights
            
            if hasattr(layer, 'biases') and layer.biases is not None:
                layer_data['biases'] = layer.biases
            
            if hasattr(layer, 'units'):
                layer_data['units'] = layer.units
                if hasattr(layer, 'input_dim'):
                    layer_data['input_dim'] = layer.input_dim
            
            if hasattr(layer, 'activation') and layer.activation:
                layer_data['activation'] = f"{layer.activation.__class__.__module__}.{layer.activation.__class__.__name__}"
            
            if hasattr(layer, 'rate'):
                layer_data['rate'] = layer.rate
            
            if hasattr(layer, 'epsilon'):
                layer_data['epsilon'] = layer.epsilon
                layer_data['momentum'] = layer.momentum
                layer_data['gamma'] = layer.gamma
                layer_data['beta'] = layer.beta
                # CRITICAL: Save running statistics for BatchNorm inference
                if hasattr(layer, 'running_mean') and layer.running_mean is not None:
                    layer_data['running_mean'] = layer.running_mean
                    layer_data['running_var'] = layer.running_var
            
            architecture.append(layer_data)
        
        return architecture
    
    def _deserialize_architecture(self, architecture_data):
        """Deserialize the model architecture from a dictionary."""
        for layer_data in architecture_data:
            layer_class_name = layer_data['type'].split('.')[-1].split('::')[-1]
            
            # Reconstruct layer based on type
            if layer_class_name == 'DenseLayer':
                activation = None
                if 'activation' in layer_data and layer_data['activation']:
                    activation_class_name = layer_data['activation'].split('.')[-1].split('::')[-1]
                    activation = globals()[activation_class_name]()
                
                layer = DenseLayer(
                    units=layer_data['units'],
                    input_dim=layer_data['input_dim'],
                    activation=activation
                )
                if 'weights' in layer_data:
                    layer.weights = layer_data['weights']
                if 'biases' in layer_data:
                    layer.biases = layer_data['biases']
            
            elif layer_class_name == 'ActivationLayer':
                activation_class_name = layer_data['activation'].split('.')[-1].split('::')[-1]
                activation = globals()[activation_class_name]()
                layer = ActivationLayer(activation)
            
            elif layer_class_name == 'DropoutLayer':
                layer = DropoutLayer(rate=layer_data.get('rate', 0.5))
            
            elif layer_class_name == 'BatchNormLayer':
                layer = BatchNormLayer(
                    epsilon=layer_data.get('epsilon', 1e-5),
                    momentum=layer_data.get('momentum', 0.1)
                )
                layer.gamma = layer_data.get('gamma', 1.0)
                layer.beta = layer_data.get('beta', 0.0)
                # CRITICAL: Restore running statistics for BatchNorm inference
                if 'running_mean' in layer_data and 'running_var' in layer_data:
                    layer.running_mean = layer_data['running_mean']
                    layer.running_var = layer_data['running_var']
                else:
                    # Inicializar running stats si no existen en el archivo
                    layer.running_mean = None
                    layer.running_var = None
            
            else:
                # For other layers, try to instantiate with default parameters
                try:
                    layer = globals()[layer_class_name]()
                except:
                    print(f"Warning: Could not instantiate layer {layer_class_name}, skipping")
                    continue
            
            self.layers.append(layer)
    
    def _count_layer_params(self, layer):
        """Count parameters in a layer."""
        params = 0
        
        if hasattr(layer, 'weights') and layer.weights:
            params += sum(len(row) for row in layer.weights)
        
        if hasattr(layer, 'biases') and layer.biases:
            params += len(layer.biases)
        
        if hasattr(layer, 'gamma'):
            params += 2  # gamma and beta
        
        return params
    
    def _get_layer_output_shape(self, layer, idx):
        """Get output shape of a layer."""
        if hasattr(layer, 'units'):
            return f"(None, {layer.units})"
        elif hasattr(layer, 'filters'):
            return f"(None, ?, ?, {layer.filters})"
        elif layer.__class__.__name__ in ['DropoutLayer', 'BatchNormLayer']:
            # Dropout y BatchNorm mantienen la misma forma que la última capa con units
            # Buscar hacia atrás la última capa con units
            for i in range(idx - 1, -1, -1):
                if hasattr(self.layers[i], 'units'):
                    return f"(None, {self.layers[i].units})"
            return "(None, ?)"
        else:
            return "(None, ?)"
    
    def _layer_config(self, layer):
        """Get layer configuration."""
        config = {}
        
        if hasattr(layer, 'units'):
            config['units'] = layer.units
            if hasattr(layer, 'input_dim'):
                config['input_dim'] = layer.input_dim
        
        if hasattr(layer, 'activation') and layer.activation:
            config['activation'] = layer.activation.__class__.__name__
        
        if hasattr(layer, 'rate'):
            config['rate'] = layer.rate
        
        return config
