"""
Architectures Module - Neural Network Implementations for CPD Learning
======================================================================

This module contains all neural network architectures for learning CPDs:
- LSTM: Long Short-Term Memory networks
- BNN: Bayesian Neural Networks  
- VAE: Variational Autoencoders
- Autoencoder: Standard autoencoders
- Normalizing Flow: Normalizing flows
- Simple/Advanced/Ultra/Mega: Traditional architectures
"""

import numpy as np
import pandas as pd
from abc import ABC, abstractmethod
import warnings
warnings.filterwarnings('ignore')

# TensorFlow setup
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow import keras
from tensorflow.keras import layers, Model, Sequential
from tensorflow.keras.optimizers import Adam, AdamW
from tensorflow.keras.callbacks import EarlyStopping, ReduceLROnPlateau
from tensorflow.keras.layers import (
    Dense, Dropout, LSTM, GRU, Conv1D, Flatten, 
    BatchNormalization, LayerNormalization, Input,
    MultiHeadAttention, Concatenate, Add
)

# Scikit-learn
from sklearn.neural_network import MLPRegressor, MLPClassifier
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split

# PGMPy
from pgmpy.factors.discrete import TabularCPD


class BaseCPDLearner(ABC):
    """Base class for all CPD learning architectures."""
    
    def __init__(self, node, data, true_model, learnt_bn_structure, num_parameters,
                 epochs=100, batch_size=32, learning_rate=0.001, validation_split=0.2,
                 early_stopping=True, verbose=True, random_state=42, **kwargs):
        self.node = node
        self.data = data
        self.true_model = true_model
        self.learnt_bn_structure = learnt_bn_structure
        self.num_parameters = num_parameters
        self.epochs = epochs
        self.batch_size = batch_size
        self.learning_rate = learning_rate
        self.validation_split = validation_split
        self.early_stopping = early_stopping
        self.verbose = verbose
        self.random_state = random_state
        self.kwargs = kwargs
        
        # Set random seeds
        np.random.seed(random_state)
        tf.random.set_seed(random_state)
        
        # Initialize model components
        self.model = None
        self.scaler = StandardScaler()
        self.is_fitted = False
        
        # Prepare data
        self._prepare_data()
    
    def _prepare_data(self):
        """Prepare training data from the input DataFrame."""
        # Get parents of the target node
        parents = list(self.learnt_bn_structure.predecessors(self.node))
        
        if len(parents) == 0:
            # Root node - no parents
            self.X = np.ones((len(self.data), 1))  # Dummy input
            self.y = self.data[self.node].values
        else:
            # Node with parents
            self.X = self.data[parents].values
            self.y = self.data[self.node].values
        
        # Store variable cardinalities
        self.node_card = len(self.data[self.node].unique())
        self.parent_cards = [len(self.data[parent].unique()) for parent in parents] if parents else [1]
        self.parents = parents
        
        if self.verbose:
            print("Data prepared for node '{}':".format(self.node))
            print(f"   Parents: {parents if parents else 'None (root node)'}")
            print(f"   Input shape: {self.X.shape}")
            print(f"   Output shape: {self.y.shape}")
            print(f"   Node cardinality: {self.node_card}")
    
    @abstractmethod
    def _build_model(self):
        """Build the neural network model (to be implemented by subclasses)."""
        pass
    
    @abstractmethod
    def fit(self, data):
        """Fit the model to the data (to be implemented by subclasses)."""
        pass
    
    def get_cpd(self):
        """Convert model predictions to TabularCPD format."""
        if not self.is_fitted:
            raise ValueError("Model must be fitted before getting CPD")
        
        # Generate all possible parent combinations
        if len(self.parents) == 0:
            # Root node
            if hasattr(self.model, 'predict_proba'):
                # Sklearn classifier with probability prediction
                pred = self.model.predict_proba(np.ones((1, 1)))[0]
            else:
                # TensorFlow model or sklearn regressor
                try:
                    pred = self.model.predict(np.ones((1, 1)), verbose=0)[0]
                except TypeError:
                    # Sklearn model without verbose parameter
                    pred = self.model.predict(np.ones((1, 1)))[0]
            
            if len(pred) == 1 or not hasattr(pred, '__len__'):
                # Binary case or single value
                if hasattr(pred, '__len__'):
                    prob = pred[0] if len(pred) == 1 else pred
                else:
                    prob = pred
                cpd_values = np.array([[1-prob], [prob]])
            else:
                cpd_values = pred.reshape(-1, 1)
        else:
            # Node with parents
            parent_combinations = []
            for cards in self.parent_cards:
                parent_combinations.append(list(range(cards)))
            
            import itertools
            all_combinations = list(itertools.product(*parent_combinations))
            
            X_test = np.array(all_combinations)
            
            # Get predictions
            if hasattr(self.model, 'predict_proba'):
                # Sklearn classifier with probability prediction
                predictions = self.model.predict_proba(X_test)
            else:
                # TensorFlow model or sklearn regressor
                try:
                    predictions = self.model.predict(X_test, verbose=0)
                except TypeError:
                    # Sklearn model without verbose parameter
                    predictions = self.model.predict(X_test)
            
            # Reshape for CPD format
            if self.node_card == 2:
                # Binary target variable
                cpd_values = np.zeros((2, len(all_combinations)))
                for i, pred in enumerate(predictions):
                    if hasattr(pred, '__len__') and len(pred) > 1:
                        cpd_values[0, i] = pred[0]  # P(node=0|parents)
                        cpd_values[1, i] = pred[1]  # P(node=1|parents)
                    else:
                        cpd_values[0, i] = 1 - pred  # P(node=0|parents)
                        cpd_values[1, i] = pred      # P(node=1|parents)
            else:
                # Multi-class case
                cpd_values = predictions.T
        
        # Normalize to ensure probabilities sum to 1
        cpd_values = cpd_values / cpd_values.sum(axis=0, keepdims=True)
        
        # Create TabularCPD
        if len(self.parents) == 0:
            cpd = TabularCPD(
                variable=self.node,
                variable_card=self.node_card,
                values=cpd_values
            )
        else:
            cpd = TabularCPD(
                variable=self.node,
                variable_card=self.node_card,
                values=cpd_values,
                evidence=self.parents,
                evidence_card=self.parent_cards
            )
        
        return cpd


class LSTMCPDLearner(BaseCPDLearner):
    """LSTM-based CPD learner for sequential dependencies."""
    
    def _build_model(self):
        """Build LSTM model for CPD learning."""
        input_dim = self.X.shape[1]
        
        # Create sequence-based input (treat each sample as a sequence)
        inputs = Input(shape=(1, input_dim), name='lstm_input')
        
        # LSTM layers
        lstm1 = LSTM(64, return_sequences=True, dropout=0.2, recurrent_dropout=0.2)(inputs)
        lstm2 = LSTM(32, dropout=0.2, recurrent_dropout=0.2)(lstm1)
        
        # Dense layers
        dense1 = Dense(64, activation='relu')(lstm2)
        dense1 = Dropout(0.3)(dense1)
        dense2 = Dense(32, activation='relu')(dense1)
        dense2 = Dropout(0.2)(dense2)
        
        # Output layer
        if self.node_card == 2:
            outputs = Dense(1, activation='sigmoid', name='cpd_output')(dense2)
        else:
            outputs = Dense(self.node_card, activation='softmax', name='cpd_output')(dense2)
        
        model = Model(inputs=inputs, outputs=outputs, name=f'LSTM_CPD_{self.node}')
        
        # Compile model
        if self.node_card == 2:
            loss = 'binary_crossentropy'
            metrics = ['accuracy']
        else:
            loss = 'sparse_categorical_crossentropy'
            metrics = ['sparse_categorical_accuracy']
            
        model.compile(
            optimizer=Adam(learning_rate=self.learning_rate),
            loss=loss,
            metrics=metrics
        )
        
        return model
    
    def fit(self, data):
        """Fit LSTM model to the data."""
        self.model = self._build_model()
        
        # Reshape data for LSTM (add sequence dimension)
        X_lstm = self.X.reshape(self.X.shape[0], 1, self.X.shape[1])
        
        # Prepare callbacks
        callbacks = []
        if self.early_stopping:
            callbacks.append(EarlyStopping(patience=10, restore_best_weights=True))
            callbacks.append(ReduceLROnPlateau(patience=5, factor=0.5))
        
        # Train model
        if self.verbose:
            print("Training LSTM model for node '{}'...".format(self.node))
        
        history = self.model.fit(
            X_lstm, self.y,
            epochs=self.epochs,
            batch_size=self.batch_size,
            validation_split=self.validation_split,
            callbacks=callbacks,
            verbose=1 if self.verbose else 0
        )
        
        self.is_fitted = True
        if self.verbose:
            print("LSTM training completed!")
        
        return history


class BayesianNeuralNetworkCPDLearner(BaseCPDLearner):
    """Bayesian Neural Network for uncertainty quantification in CPD learning."""
    
    def _build_model(self):
        """Build Bayesian Neural Network using TensorFlow Probability."""
        input_dim = self.X.shape[1]
        
        # Bayesian neural network with variational layers
        model = Sequential([
            Input(shape=(input_dim,)),
            
            # Bayesian dense layers with uncertainty
            tfp.layers.DenseVariational(
                64, 
                make_posterior_fn=tfp.layers.default_mean_field_normal_fn(),
                make_prior_fn=tfp.layers.default_multivariate_normal_fn,
                kl_weight=1/len(self.X),
                activation='relu'
            ),
            Dropout(0.3),
            
            tfp.layers.DenseVariational(
                32,
                make_posterior_fn=tfp.layers.default_mean_field_normal_fn(),
                make_prior_fn=tfp.layers.default_multivariate_normal_fn,
                kl_weight=1/len(self.X),
                activation='relu'
            ),
            Dropout(0.2),
            
            # Output layer
            tfp.layers.DenseVariational(
                self.node_card if self.node_card > 2 else 1,
                make_posterior_fn=tfp.layers.default_mean_field_normal_fn(),
                make_prior_fn=tfp.layers.default_multivariate_normal_fn,
                kl_weight=1/len(self.X),
                activation='softmax' if self.node_card > 2 else 'sigmoid'
            )
        ], name=f'BNN_CPD_{self.node}')
        
        # Compile with ELBO loss
        if self.node_card == 2:
            loss = lambda y, p_y: -p_y.log_prob(y)
            metrics = ['accuracy']
        else:
            loss = lambda y, p_y: -p_y.log_prob(y)
            metrics = ['sparse_categorical_accuracy']
        
        model.compile(
            optimizer=Adam(learning_rate=self.learning_rate),
            loss=loss,
            metrics=metrics
        )
        
        return model
    
    def fit(self, data):
        """Fit Bayesian Neural Network to the data."""
        self.model = self._build_model()
        
        # Prepare callbacks
        callbacks = []
        if self.early_stopping:
            callbacks.append(EarlyStopping(patience=15, restore_best_weights=True))
            callbacks.append(ReduceLROnPlateau(patience=7, factor=0.5))
        
        # Train model
        if self.verbose:
            print("Training Bayesian Neural Network for node '{}'...".format(self.node))
        
        history = self.model.fit(
            self.X, self.y,
            epochs=self.epochs,
            batch_size=self.batch_size,
            validation_split=self.validation_split,
            callbacks=callbacks,
            verbose=1 if self.verbose else 0
        )
        
        self.is_fitted = True
        if self.verbose:
            print("BNN training completed!")
        
        return history


class VAECPDLearner(BaseCPDLearner):
    """Variational Autoencoder for probabilistic CPD learning."""
    
    def _build_model(self):
        """Build VAE model for CPD learning."""
        input_dim = self.X.shape[1]
        latent_dim = min(16, max(2, input_dim // 2))
        
        # Encoder
        encoder_inputs = Input(shape=(input_dim,), name='encoder_input')
        h = Dense(64, activation='relu')(encoder_inputs)
        h = Dropout(0.3)(h)
        h = Dense(32, activation='relu')(h)
        
        z_mean = Dense(latent_dim, name='z_mean')(h)
        z_log_var = Dense(latent_dim, name='z_log_var')(h)
        
        # Sampling function
        def sampling(args):
            z_mean, z_log_var = args
            epsilon = tf.random.normal(shape=tf.shape(z_mean))
            return z_mean + tf.exp(0.5 * z_log_var) * epsilon
        
        z = layers.Lambda(sampling, output_shape=(latent_dim,), name='z')([z_mean, z_log_var])
        
        # Decoder
        decoder_h = Dense(32, activation='relu')(z)
        decoder_h = Dropout(0.2)(decoder_h)
        decoder_h = Dense(64, activation='relu')(decoder_h)
        
        # CPD output
        if self.node_card == 2:
            cpd_output = Dense(1, activation='sigmoid', name='cpd_output')(decoder_h)
        else:
            cpd_output = Dense(self.node_card, activation='softmax', name='cpd_output')(decoder_h)
        
        # Reconstruction output
        reconstruction_output = Dense(input_dim, activation='sigmoid', name='reconstruction')(decoder_h)
        
        # VAE model
        vae = Model(encoder_inputs, [cpd_output, reconstruction_output], name=f'VAE_CPD_{self.node}')
        
        # Custom VAE loss
        def vae_loss(y_true, y_pred):
            cpd_pred, recon_pred = y_pred
            
            # CPD loss
            if self.node_card == 2:
                cpd_loss = tf.keras.losses.binary_crossentropy(y_true[0], cpd_pred)
            else:
                cpd_loss = tf.keras.losses.sparse_categorical_crossentropy(y_true[0], cpd_pred)
            
            # Reconstruction loss
            recon_loss = tf.keras.losses.mse(y_true[1], recon_pred)
            
            # KL divergence
            kl_loss = -0.5 * tf.reduce_mean(1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var))
            
            return cpd_loss + 0.1 * recon_loss + 0.01 * kl_loss
        
        vae.compile(
            optimizer=Adam(learning_rate=self.learning_rate),
            loss=vae_loss,
            metrics=['accuracy']
        )
        
        return vae
    
    def fit(self, data):
        """Fit VAE model to the data."""
        self.model = self._build_model()
        
        # Prepare targets (CPD target + reconstruction target)
        y_cpd = self.y
        y_recon = self.X  # Reconstruct input
        
        # Prepare callbacks
        callbacks = []
        if self.early_stopping:
            callbacks.append(EarlyStopping(patience=15, restore_best_weights=True))
            callbacks.append(ReduceLROnPlateau(patience=7, factor=0.5))
        
        # Train model
        if self.verbose:
            print("Training VAE for node '{}'...".format(self.node))
        
        history = self.model.fit(
            self.X, [y_cpd, y_recon],
            epochs=self.epochs,
            batch_size=self.batch_size,
            validation_split=self.validation_split,
            callbacks=callbacks,
            verbose=1 if self.verbose else 0
        )
        
        self.is_fitted = True
        if self.verbose:
            print("VAE training completed!")
        
        return history


class AutoencoderCPDLearner(BaseCPDLearner):
    """Standard Autoencoder for dimensionality reduction in CPD learning."""
    
    def _build_model(self):
        """Build Autoencoder model."""
        input_dim = self.X.shape[1]
        encoding_dim = max(2, input_dim // 2)
        
        # Encoder
        input_layer = Input(shape=(input_dim,))
        encoded = Dense(64, activation='relu')(input_layer)
        encoded = Dropout(0.3)(encoded)
        encoded = Dense(encoding_dim, activation='relu')(encoded)
        
        # Decoder for reconstruction
        decoded = Dense(64, activation='relu')(encoded)
        decoded = Dropout(0.2)(decoded)
        reconstruction = Dense(input_dim, activation='sigmoid', name='reconstruction')(decoded)
        
        # CPD prediction branch
        cpd_branch = Dense(32, activation='relu')(encoded)
        cpd_branch = Dropout(0.2)(cpd_branch)
        
        if self.node_card == 2:
            cpd_output = Dense(1, activation='sigmoid', name='cpd_output')(cpd_branch)
        else:
            cpd_output = Dense(self.node_card, activation='softmax', name='cpd_output')(cpd_branch)
        
        # Model
        autoencoder = Model(input_layer, [cpd_output, reconstruction], name=f'Autoencoder_CPD_{self.node}')
        
        # Compile
        autoencoder.compile(
            optimizer=Adam(learning_rate=self.learning_rate),
            loss={
                'cpd_output': 'binary_crossentropy' if self.node_card == 2 else 'sparse_categorical_crossentropy',
                'reconstruction': 'mse'
            },
            loss_weights={'cpd_output': 1.0, 'reconstruction': 0.1},
            metrics={'cpd_output': 'accuracy'}
        )
        
        return autoencoder
    
    def fit(self, data):
        """Fit Autoencoder model to the data."""
        self.model = self._build_model()
        
        # Prepare targets
        y_cpd = self.y
        y_recon = self.X
        
        # Prepare callbacks
        callbacks = []
        if self.early_stopping:
            callbacks.append(EarlyStopping(patience=10, restore_best_weights=True))
            callbacks.append(ReduceLROnPlateau(patience=5, factor=0.5))
        
        # Train model
        if self.verbose:
            print("Training Autoencoder for node '{}'...".format(self.node))
        
        history = self.model.fit(
            self.X, {'cpd_output': y_cpd, 'reconstruction': y_recon},
            epochs=self.epochs,
            batch_size=self.batch_size,
            validation_split=self.validation_split,
            callbacks=callbacks,
            verbose=1 if self.verbose else 0
        )
        
        self.is_fitted = True
        if self.verbose:
            print("Autoencoder training completed!")
        
        return history


class NormalizingFlowCPDLearner(BaseCPDLearner):
    """Normalizing Flow for exact probability modeling in CPD learning."""
    
    def _build_model(self):
        """Build Normalizing Flow model."""
        input_dim = self.X.shape[1]
        
        # Create normalizing flow
        bijectors = []
        for i in range(4):  # 4 flow layers
            bijectors.append(
                tfp.bijectors.RealNVP(
                    num_masked=input_dim // 2,
                    shift_and_log_scale_fn=tfp.bijectors.real_nvp_default_template(
                        hidden_layers=[64, 32]
                    )
                )
            )
            if i < 3:
                bijectors.append(tfp.bijectors.Permute(np.random.permutation(input_dim)))
        
        flow = tfp.bijectors.Chain(bijectors)
        
        # Base distribution
        base_distribution = tfp.distributions.MultivariateNormalDiag(
            loc=tf.zeros(input_dim), scale_diag=tf.ones(input_dim)
        )
        
        # Transformed distribution
        self.flow_distribution = tfp.distributions.TransformedDistribution(
            distribution=base_distribution,
            bijector=flow
        )
        
        # CPD prediction network
        inputs = Input(shape=(input_dim,))
        h = Dense(64, activation='relu')(inputs)
        h = Dropout(0.3)(h)
        h = Dense(32, activation='relu')(h)
        h = Dropout(0.2)(h)
        
        if self.node_card == 2:
            outputs = Dense(1, activation='sigmoid')(h)
        else:
            outputs = Dense(self.node_card, activation='softmax')(h)
        
        model = Model(inputs, outputs, name=f'NormalizingFlow_CPD_{self.node}')
        
        # Custom loss combining flow likelihood and CPD prediction
        def flow_cpd_loss(y_true, y_pred):
            # CPD prediction loss
            if self.node_card == 2:
                cpd_loss = tf.keras.losses.binary_crossentropy(y_true, y_pred)
            else:
                cpd_loss = tf.keras.losses.sparse_categorical_crossentropy(y_true, y_pred)
            
            # Flow likelihood (maximize likelihood = minimize negative log-likelihood)
            flow_loss = -tf.reduce_mean(self.flow_distribution.log_prob(inputs))
            
            return cpd_loss + 0.1 * flow_loss
        
        model.compile(
            optimizer=Adam(learning_rate=self.learning_rate),
            loss=flow_cpd_loss,
            metrics=['accuracy']
        )
        
        return model
    
    def fit(self, data):
        """Fit Normalizing Flow model to the data."""
        self.model = self._build_model()
        
        # Prepare callbacks
        callbacks = []
        if self.early_stopping:
            callbacks.append(EarlyStopping(patience=15, restore_best_weights=True))
            callbacks.append(ReduceLROnPlateau(patience=7, factor=0.5))
        
        # Train model
        if self.verbose:
            print("Training Normalizing Flow for node '{}'...".format(self.node))
        
        history = self.model.fit(
            self.X, self.y,
            epochs=self.epochs,
            batch_size=self.batch_size,
            validation_split=self.validation_split,
            callbacks=callbacks,
            verbose=1 if self.verbose else 0
        )
        
        self.is_fitted = True
        if self.verbose:
            print("Normalizing Flow training completed!")
        
        return history


# Traditional architectures (using sklearn MLPRegressor/MLPClassifier)
class SimpleCPDLearner(BaseCPDLearner):
    """Simple neural network using sklearn."""
    
    def _build_model(self):
        """Build simple sklearn neural network."""
        if self.node_card == 2:
            return MLPClassifier(
                hidden_layer_sizes=(32, 16),
                activation='relu',
                solver='adam',
                alpha=0.001,
                batch_size=self.batch_size,
                learning_rate='adaptive',
                max_iter=self.epochs,
                random_state=self.random_state
            )
        else:
            return MLPRegressor(
                hidden_layer_sizes=(32, 16),
                activation='relu', 
                solver='adam',
                alpha=0.001,
                batch_size=self.batch_size,
                learning_rate='adaptive',
                max_iter=self.epochs,
                random_state=self.random_state
            )
    
    def fit(self, data):
        """Fit simple model to the data."""
        self.model = self._build_model()
        
        if self.verbose:
            print("Training Simple Neural Network for node '{}'...".format(self.node))
        
        self.model.fit(self.X, self.y)
        self.is_fitted = True
        
        if self.verbose:
            print("Simple NN training completed!")


class AdvancedCPDLearner(SimpleCPDLearner):
    """Advanced neural network using sklearn."""
    
    def _build_model(self):
        """Build advanced sklearn neural network."""
        if self.node_card == 2:
            return MLPClassifier(
                hidden_layer_sizes=(64, 32, 16),
                activation='relu',
                solver='adam',
                alpha=0.0005,
                batch_size=self.batch_size,
                learning_rate='adaptive',
                max_iter=self.epochs,
                random_state=self.random_state
            )
        else:
            return MLPRegressor(
                hidden_layer_sizes=(64, 32, 16),
                activation='relu',
                solver='adam', 
                alpha=0.0005,
                batch_size=self.batch_size,
                learning_rate='adaptive',
                max_iter=self.epochs,
                random_state=self.random_state
            )


class UltraCPDLearner(SimpleCPDLearner):
    """Ultra neural network using sklearn."""
    
    def _build_model(self):
        """Build ultra sklearn neural network."""
        if self.node_card == 2:
            return MLPClassifier(
                hidden_layer_sizes=(128, 64, 32, 16),
                activation='relu',
                solver='adam',
                alpha=0.0001,
                batch_size=self.batch_size,
                learning_rate='adaptive',
                max_iter=self.epochs,
                random_state=self.random_state
            )
        else:
            return MLPRegressor(
                hidden_layer_sizes=(128, 64, 32, 16),
                activation='relu',
                solver='adam',
                alpha=0.0001,
                batch_size=self.batch_size,
                learning_rate='adaptive',
                max_iter=self.epochs,
                random_state=self.random_state
            )


class MegaCPDLearner(SimpleCPDLearner):
    """Mega neural network using sklearn."""
    
    def _build_model(self):
        """Build mega sklearn neural network."""
        if self.node_card == 2:
            return MLPClassifier(
                hidden_layer_sizes=(256, 128, 64, 32, 16),
                activation='relu',
                solver='adam',
                alpha=0.00005,
                batch_size=self.batch_size,
                learning_rate='adaptive',
                max_iter=self.epochs,
                random_state=self.random_state
            )
        else:
            return MLPRegressor(
                hidden_layer_sizes=(256, 128, 64, 32, 16),
                activation='relu',
                solver='adam',
                alpha=0.00005,
                batch_size=self.batch_size,
                learning_rate='adaptive',
                max_iter=self.epochs,
                random_state=self.random_state
            )