import numpy as np
from sklearn.datasets import fetch_openml
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import OneHotEncoder
import matplotlib.pyplot as plt

def preprocess_data():
    """
    Fetches the MNIST dataset, normalizes the features, encodes the labels, 
    and splits the data into training, validation, and test sets.

    Returns:
        X_train (ndarray): Training feature data.
        X_val (ndarray): Validation feature data.
        X_test (ndarray): Test feature data.
        y_train (ndarray): Training labels.
        y_val (ndarray): Validation labels.
        y_test (ndarray): Test labels.
    """
    # Fetch the MNIST dataset
    mnist = fetch_openml('mnist_784', version=1)
    X = mnist["data"].to_numpy()
    y = mnist["target"].astype(int).to_numpy()

    # Normalize the feature data
    X = X / 255.0

    # One-hot encode the target labels
    encoder = OneHotEncoder()
    y = encoder.fit_transform(y.reshape(-1, 1)).toarray()

    # Split data into training and temporary sets
    X_train, X_temp, y_train, y_temp = train_test_split(X, y, test_size=0.3, random_state=42)

    # Split the temporary set into validation and test sets
    X_val, X_test, y_val, y_test = train_test_split(X_temp, y_temp, test_size=0.5, random_state=42)

    return X_train, X_val, X_test, y_train, y_val, y_test

# Usage Example
X_train, X_val, X_test, y_train, y_val, y_test = preprocess_data()
print(f"Training set size:\nX_train {X_train.shape}\ny_train {y_train.shape}")
print(f"\nValidation set size:\nX_val {X_val.shape}\ny_val {y_val.shape}")
print(f"\nTest set size:\nX_test {X_test.shape}\ny_test {y_test.shape}")

class SimpleNNWithDropout:
    def __init__(self, input_size: int, hidden_size: int, output_size: int, 
                 dropout_rate: float = 0.5, learning_rate: float = 0.01, use_dropout: bool = True):
        """
        Initializes the neural network with specified parameters.

        Args:
            input_size (int): Number of input features.
            hidden_size (int): Number of neurons in the hidden layer.
            output_size (int): Number of output classes.
            dropout_rate (float, optional): Dropout rate for regularization. Defaults to 0.5.
            learning_rate (float, optional): Learning rate for gradient descent. Defaults to 0.01.
            use_dropout (bool, optional): Flag to apply dropout during training. Defaults to True.
        """
        np.random.seed(42)  # For reproducibility
        self.W1 = np.random.randn(input_size, hidden_size) * np.sqrt(2. / input_size)
        self.b1 = np.zeros((1, hidden_size))
        self.W2 = np.random.randn(hidden_size, output_size) * np.sqrt(2. / hidden_size)
        self.b2 = np.zeros((1, output_size))
        self.dropout_rate = dropout_rate
        self.learning_rate = learning_rate
        self.use_dropout = use_dropout

    # Define the ReLU activation function to apply to input data.
    def relu(self, x: np.ndarray) -> np.ndarray:
        """Applies the ReLU activation function."""
        return np.maximum(0, x)

    # Define the function to compute the derivative of the ReLU activation function.
    def relu_derivative(self, x: np.ndarray) -> np.ndarray:
        """Computes the derivative of the ReLU activation function."""
        return (x > 0).astype(float)

    # Define the softmax activation function to apply to input data, converting logits to probabilities.
    def softmax(self, x: np.ndarray) -> np.ndarray:
        """Applies the softmax activation function."""
        exp_x = np.exp(x - np.max(x, axis=1, keepdims=True))
        return exp_x / np.sum(exp_x, axis=1, keepdims=True)

# Define the forward pass for applying dropout, generating a dropout mask and applying it to the input data.
def dropout_forward(self, X: np.ndarray) -> (np.ndarray, np.ndarray):
    """
    Applies dropout during the forward pass.

    Args:
        X (np.ndarray): Input data.

    Returns:
        np.ndarray: Data after applying dropout.
        np.ndarray: Dropout mask.
    """
    mask = np.random.rand(*X.shape) > self.dropout_rate
    mask = mask / (1.0 - self.dropout_rate)
    X_dropout = X * mask
    return X_dropout, mask

# Define the backward pass for applying the dropout mask to the gradient during backpropagation.
def dropout_backward(self, dout: np.ndarray, mask: np.ndarray) -> np.ndarray:
    """Applies the dropout mask during the backward pass."""
    return dout * mask

# Perform the forward pass through the network, applying ReLU and softmax activation functions.
# Apply dropout during training if the dropout flag is set.
def forward_pass(self, X: np.ndarray, training: bool = True) -> np.ndarray:
    """
    Performs the forward pass of the network.

    Args:
        X (np.ndarray): Input data.
        training (bool, optional): Flag to indicate if the model is in training mode. Defaults to True.

    Returns:
        np.ndarray: Output probabilities after the forward pass.
    """
    # First layer
    self.z1 = X.dot(self.W1) + self.b1
    self.a1 = self.relu(self.z1)

    # Apply dropout during training
    if training and self.use_dropout:
        self.a1_dropout, self.mask = self.dropout_forward(self.a1)
    else:
        self.a1_dropout = self.a1

    # Second layer
    self.z2 = self.a1_dropout.dot(self.W2) + self.b2
    self.a2 = self.softmax(self.z2)
    return self.a2

# Perform the backward pass to compute gradients for the weights and biases.
# Update the network parameters using gradient descent. 
def backward_pass(self, X: np.ndarray, y: np.ndarray):
    """
    Performs the backward pass and updates the network parameters.

    Args:
        X (np.ndarray): Input data.
        y (np.ndarray): True labels.
    """
    m = X.shape[0]

    # Gradients for output layer
    dz2 = self.a2 - y
    dW2 = self.a1_dropout.T.dot(dz2) / m
    db2 = np.sum(dz2, axis=0, keepdims=True) / m

    # Gradient for hidden layer with dropout
    da1_dropout = dz2.dot(self.W2.T)
    if self.use_dropout:
        da1 = self.dropout_backward(da1_dropout, self.mask)
    else:
        da1 = da1_dropout

    # Gradients for hidden layer
    dz1 = da1 * self.relu_derivative(self.z1)
    dW1 = X.T.dot(dz1) / m
    db1 = np.sum(dz1, axis=0, keepdims=True) / m

    # Update weights and biases
    self.W1 -= self.learning_rate * dW1
    self.b1 -= self.learning_rate * db1
    self.W2 -= self.learning_rate * dW2
    self.b2 -= self.learning_rate * db2

# Compute the cross-entropy loss between the true labels and predicted probabilities.
def compute_loss(self, y: np.ndarray, y_hat: np.ndarray) -> float:
    """
    Computes the cross-entropy loss.

    Args:
        y (np.ndarray): True labels.
        y_hat (np.ndarray): Predicted probabilities.

    Returns:
        float: Cross-entropy loss.
    """
    m = y.shape[0]
    return -np.sum(y * np.log(y_hat + 1e-8)) / m

# Train the neural network over a specified number of epochs.
# Perform forward and backward passes, update parameters, and optionally print training progress.
def train(self, X: np.ndarray, y: np.ndarray, epochs: int = 100, verbose: bool = True):
    """
    Trains the neural network.

    Args:
        X (np.ndarray): Training feature data.
        y (np.ndarray): Training labels.
        epochs (int, optional): Number of training epochs. Defaults to 100.
        verbose (bool, optional): Flag to print training progress. Defaults to True.
    """
    for epoch in range(epochs):
        y_hat = self.forward_pass(X, training=True)
        loss = self.compute_loss(y, y_hat)
        self.backward_pass(X, y)
        if verbose and epoch % 10 == 0:
            print(f"Epoch {epoch}/{epochs} - Loss: {loss:.4f}")

# Predict labels for the given input data by performing a forward pass without dropout.
def predict(self, X: np.ndarray) -> np.ndarray:
    """
    Predicts labels for the given input data.

    Args:
        X (np.ndarray): Input data.

    Returns:
        np.ndarray: Predicted labels.
    """
    y_hat = self.forward_pass(X, training=False)
    return np.argmax(y_hat, axis=1)

class SimpleNNWithDropout:
    def __init__(self, input_size: int, hidden_size: int, output_size: int, 
                 dropout_rate: float = 0.5, learning_rate: float = 0.01, use_dropout: bool = True):
        """
        Initializes the neural network with specified parameters.

        Args:
            input_size (int): Number of input features.
            hidden_size (int): Number of neurons in the hidden layer.
            output_size (int): Number of output classes.
            dropout_rate (float, optional): Dropout rate for regularization. Defaults to 0.5.
            learning_rate (float, optional): Learning rate for gradient descent. Defaults to 0.01.
            use_dropout (bool, optional): Flag to apply dropout during training. Defaults

 to True.
        """
        np.random.seed(42)  # For reproducibility
        self.W1 = np.random.randn(input_size, hidden_size) * np.sqrt(2. / input_size) # He Initialization (aka He Normal Initialization)
        self.b1 = np.zeros((1, hidden_size))
        self.W2 = np.random.randn(hidden_size, output_size) * np.sqrt(2. / hidden_size)
        self.b2 = np.zeros((1, output_size))
        self.dropout_rate = dropout_rate
        self.learning_rate = learning_rate
        self.use_dropout = use_dropout

    # Define the ReLU activation function to apply to input data.
    def relu(self, x: np.ndarray) -> np.ndarray:
        """Applies the ReLU activation function."""
        return np.maximum(0, x)

    # Define the function to compute the derivative of the ReLU activation function.
    def relu_derivative(self, x: np.ndarray) -> np.ndarray:
        """Computes the derivative of the ReLU activation function."""
        return (x > 0).astype(float)

    # Define the softmax activation function to apply to input data, converting logits to probabilities.
    def softmax(self, x: np.ndarray) -> np.ndarray:
        """Applies the softmax activation function."""
        exp_x = np.exp(x - np.max(x, axis=1, keepdims=True))
        return exp_x / np.sum(exp_x, axis=1, keepdims=True)

    # Define the forward pass for applying dropout, generating a dropout mask and applying it to the input data.
    def dropout_forward(self, X: np.ndarray) -> (np.ndarray, np.ndarray):
        """
        Applies dropout during the forward pass.

        Args:
            X (np.ndarray): Input data.

        Returns:
            np.ndarray: Data after applying dropout.
            np.ndarray: Dropout mask.
        """
        mask = np.random.rand(*X.shape) > self.dropout_rate
        mask = mask / (1.0 - self.dropout_rate)
        X_dropout = X * mask
        return X_dropout, mask

    # Define the backward pass for applying the dropout mask to the gradient during backpropagation.
    def dropout_backward(self, dout: np.ndarray, mask: np.ndarray) -> np.ndarray:
        """Applies the dropout mask during the backward pass."""
        return dout * mask

    # Perform the forward pass through the network, applying ReLU and softmax activation functions.
    # Apply dropout during training if the dropout flag is set.
    def forward_pass(self, X: np.ndarray, training: bool = True) -> np.ndarray:
        """
        Performs the forward pass of the network.

        Args:
            X (np.ndarray): Input data.
            training (bool, optional): Flag to indicate if the model is in training mode. Defaults to True.

        Returns:
            np.ndarray: Output probabilities after the forward pass.
        """
        # First layer
        self.z1 = X.dot(self.W1) + self.b1
        self.a1 = self.relu(self.z1)

        # Apply dropout during training
        if training and self.use_dropout:
            self.a1_dropout, self.mask = self.dropout_forward(self.a1)
        else:
            self.a1_dropout = self.a1

        # Second layer
        self.z2 = self.a1_dropout.dot(self.W2) + self.b2
        self.a2 = self.softmax(self.z2)
        return self.a2

    # Perform the backward pass to compute gradients for the weights and biases.
    # Update the network parameters using gradient descent. 
    def backward_pass(self, X: np.ndarray, y: np.ndarray):
        """
    Performs the backward pass and updates the network parameters.

    Args:
        X (np.ndarray): Input data.
        y (np.ndarray): True labels.
    """
    m = X.shape[0]

    # Gradients for output layer
    dz2 = self.a2 - y
    dW2 = self.a1_dropout.T.dot(dz2) / m
    db2 = np.sum(dz2, axis=0, keepdims=True) / m

    # Gradient for hidden layer with dropout
    da1_dropout = dz2.dot(self.W2.T)
    if self.use_dropout:
        da1 = self.dropout_backward(da1_dropout, self.mask)
    else:
        da1 = da1_dropout

    # Gradients for hidden layer
    dz1 = da1 * self.relu_derivative(self.z1)
    dW1 = X.T.dot(dz1) / m
    db1 = np.sum(dz1, axis=0, keepdims=True) / m

    # Update weights and biases
    self.W1 -= self.learning_rate * dW1
    self.b1 -= self.learning_rate * db1
    self.W2 -= self.learning_rate * dW2
    self.b2 -= self.learning_rate * db2

# Compute the cross-entropy loss between the true labels and predicted probabilities.
def compute_loss(self, y: np.ndarray, y_hat: np.ndarray) -> float:
    """
    Computes the cross-entropy loss.

    Args:
        y (np.ndarray): True labels.
        y_hat (np.ndarray): Predicted probabilities.

    Returns:
        float: Cross-entropy loss.
    """
    m = y.shape[0]
    return -np.sum(y * np.log(y_hat + 1e-8)) / m

# Train the neural network over a specified number of epochs.
# Perform forward and backward passes, update parameters, and optionally print training progress.
def train(self, X: np.ndarray, y: np.ndarray, epochs: int = 100, verbose: bool = True):
    """
    Trains the neural network.

    Args:
        X (np.ndarray): Training feature data.
        y (np.ndarray): Training labels.
        epochs (int, optional): Number of training epochs. Defaults to 100.
        verbose (bool, optional): Flag to print training progress. Defaults to True.
    """
    for epoch in range(epochs):
        y_hat = self.forward_pass(X, training=True)
        loss = self.compute_loss(y, y_hat)
        self.backward_pass(X, y)
        if verbose and epoch % 10 == 0:
            print(f"Epoch {epoch}/{epochs} - Loss: {loss:.4f}")

# Predict labels for the given input data by performing a forward pass without dropout.
def predict(self, X: np.ndarray) -> np.ndarray:
    """
    Predicts labels for the given input data.

    Args:
        X (np.ndarray): Input data.

    Returns:
        np.ndarray: Predicted labels.
    """
    y_hat = self.forward_pass(X, training=False)
    return np.argmax(y_hat, axis=1)

def compute_accuracy(y_true: np.ndarray, y_pred: np.ndarray) -> float:
    """
    Computes the accuracy of the model predictions.

    Args:
        y_true (np.ndarray): True labels.
        y_pred (np.ndarray): Predicted probabilities.

    Returns:
        float: Accuracy of the predictions.
    """
    predictions = np.argmax(y_pred, axis=1)
    actuals = np.argmax(y_true, axis=1)
    return np.mean(predictions == actuals)

def train_model(model, X_train: np.ndarray, y_train: np.ndarray, 
                X_val: np.ndarray, y_val: np.ndarray, num_epochs: int):
    """
    Trains the neural network model and tracks the training and validation losses and accuracies.

    Args:
        model: Instance of the neural network model.
        X_train (np.ndarray): Training feature data.
        y_train (np.ndarray): Training labels.
        X_val (np.ndarray): Validation feature data.
        y_val (np.ndarray): Validation labels.
        num_epochs (int): Number of training epochs.

    Returns:
        tuple: Lists of training and validation losses and accuracies.
    """
    train_losses, val_losses = [], []
    train_accuracies, val_accuracies = [], []

    for epoch in range(num_epochs):
        # Forward pass and loss computation for training data
        y_hat_train = model.forward_pass(X_train, training=True)
        train_loss = model.compute_loss(y_train, y_hat_train)

        # Backward pass and parameter update
        model.backward_pass(X_train, y_train)

        # Forward pass and loss computation for validation data
        y_hat_val = model.forward_pass(X_val, training=False)
        val_loss = model.compute_loss(y_val, y_hat_val)

        # Compute accuracies
        train_accuracy = compute_accuracy(y_train, y_hat_train)
        val_accuracy = compute_accuracy(y_val, y_hat_val)

        # Record losses and accuracies
        train_losses.append(train_loss)
        val_losses.append(val_loss)
        train_accuracies.append(train_accuracy)
        val_accuracies.append(val_accuracy)

        # Print progress
        print(f'Epoch {epoch + 1}/{num_epochs}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}, '
              f'Train Acc: {train_accuracy:.4f}, Val Acc: {val_accuracy:.4f}')

    return train_losses, val_losses, train_accuracies, val_accuracies

def evaluate_model(model, X_test: np.ndarray, y_test: np.ndarray):
    """
    Evaluates the neural network model on the test dataset.

    Args:
        model: Instance of the neural network model.
        X_test (np.ndarray): Test feature data.
        y_test (np.ndarray): Test labels.

    Returns:
        tuple: Predicted

 labels, actual labels, test loss, and test accuracy.
    """
    y_hat_test = model.forward_pass(X_test, training=False)
    test_loss = model.compute_loss(y_test, y_hat_test)
    test_accuracy = compute_accuracy(y_test, y_hat_test)

    print(f'Test Loss: {test_loss:.4f}, Test Accuracy: {test_accuracy:.4f}')

    predictions = np.argmax(y_hat_test, axis=1)
    actuals = np.argmax(y_test, axis=1)

    return predictions, actuals, test_loss, test_accuracy

import matplotlib.pyplot as plt

def plot_learning_curves(train_losses: list, val_losses: list, 
                         train_accuracies: list, val_accuracies: list, 
                         title_suffix: str = ""):
    """
    Plots the learning curves for training and validation loss and accuracy.

    Args:
        train_losses (list): List of training losses over epochs.
        val_losses (list): List of validation losses over epochs.
        train_accuracies (list): List of training accuracies over epochs.
        val_accuracies (list): List of validation accuracies over epochs.
        title_suffix (str, optional): Suffix for the plot titles. Defaults to "".
    """
    plt.figure(figsize=(12, 5))

    # Plotting loss curves
    plt.subplot(1, 2, 1)
    plt.plot(train_losses, label='Train Loss')
    plt.plot(val_losses, label='Validation Loss')
    plt.xlabel('Epochs')
    plt.ylabel('Loss')
    plt.legend()
    plt.title(f'Loss vs Epochs {title_suffix}')

    # Plotting accuracy curves
    plt.subplot(1, 2, 2)
    plt.plot(train_accuracies, label='Train Accuracy')
    plt.plot(val_accuracies, label='Validation Accuracy')
    plt.xlabel('Epochs')
    plt.ylabel('Accuracy')
    plt.legend()
    plt.title(f'Accuracy vs Epochs {title_suffix}')

    plt.show()

def visualize_predictions(predictions: np.ndarray, actuals: np.ndarray, 
                          X_test: np.ndarray, title: str):
    """
    Visualizes the predictions of the model on test data.

    Args:
        predictions (np.ndarray): Predicted labels.
        actuals (np.ndarray): Actual labels.
        X_test (np.ndarray): Test feature data.
        title (str): Title for the plot.
    """
    fig, axes = plt.subplots(2, 5, figsize=(10, 5))

    for i, ax in enumerate(axes.flat):
        ax.imshow(X_test[i].reshape(28, 28), cmap='gray')
        ax.set_title(f'Pred: {predictions[i]}, True: {actuals[i]}')
        ax.axis('off')

    plt.suptitle(title)
    plt.show()

def run_experiment(params: dict, use_dropout: bool, num_epochs: int = 100):
    """
    Trains, evaluates, and visualizes results for a model with or without dropout using the given parameters.

    Args:
        params (dict): Parameters for initializing the neural network.
        use_dropout (bool): Flag to indicate if dropout should be used.
        num_epochs (int, optional): Number of training epochs. Defaults to 100.
    """
    # Initialize the model with given parameters
    model = SimpleNNWithDropout(**params, use_dropout=use_dropout)

    # Train the model and capture training and validation metrics
    train_losses, val_losses, train_accs, val_accs = train_model(
        model, X_train, y_train, X_val, y_val, num_epochs=num_epochs
    )

    # Evaluate the model on the test set
    predictions, actuals, test_loss, test_acc = evaluate_model(model, X_test, y_test)

    # Visualize predictions and learning curves
    title = f"Model {'with' if use_dropout else 'without'} Dropout"
    visualize_predictions(predictions, actuals, X_test, title=title)
    plot_learning_curves(train_losses, val_losses, train_accs, val_accs, title_suffix=title)

# Preprocess data
X_train, X_val, X_test, y_train, y_val, y_test = preprocess_data()

default_params = {
    "input_size": X_train.shape[1],
    "hidden_size": 128,
    "output_size": y_train.shape[1],
    "dropout_rate": 0.5,
    "learning_rate": 0.01,
}

# Run an experiment with dropout
run_experiment(default_params, use_dropout=True, num_epochs=1000)

# Run an experiment without dropout
run_experiment(default_params, use_dropout=False, num_epochs=1000)

import torch.nn as nn

class ConvNetWithDropout(nn.Module):
    def __init__(self):
        super(ConvNetWithDropout, self).__init__()

        self.conv1 = nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=1, padding=1)
        self.dropout2d = nn.Dropout2d(p=0.5)
        self.fc1 = nn.Linear(64 * 32 * 32, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = nn.ReLU()(x)
        x = self.dropout2d(x)
        x = x.view(x.size(0), -1)
        x = self.fc1(x)
        return x

import torch
import torch.nn as nn

class SimpleNNWithDropout(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, dropout_rate=0.5):
        """
        Initializes the SimpleNNWithDropout.

        Args:
            input_size (int): The number of input features.
            hidden_size (int): The number of neurons in the hidden layer.
            output_size (int): The number of output features (e.g., number of classes).
            dropout_rate (float): The probability of dropping a neuron during training.
        """
        super(SimpleNNWithDropout, self).__init__()
        # Define the first fully connected layer
        self.fc1 = nn.Linear(input_size, hidden_size)
        # Define the dropout layer
        self.dropout = nn.Dropout(dropout_rate)
        # Define the second fully connected layer
        self.fc2 = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        """
        Defines the forward pass of the network.

        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, input_size).

        Returns:
            torch.Tensor: Output tensor of shape (batch_size, output_size).
        """
        # Apply the first fully connected layer followed by ReLU activation
        x = torch.relu(self.fc1(x))
        # Apply dropout to the hidden layer's output during training
        x = self.dropout(x)
        # Apply the second fully connected layer
        x = self.fc2(x)
        return x

def mc_dropout_inference(model, X, num_samples=100):
    """
    Performs Monte Carlo Dropout inference on the given model.

    Args:
        model (nn.Module): The neural network model with dropout layers.
        X (torch.Tensor): The input data tensor.
        num_samples (int, optional): The number of forward passes to perform. Defaults to 100.

    Returns:
        torch.Tensor: The mean of the predictions.
        torch.Tensor: The variance of the predictions.
    """
    # Set the model to training mode to enable dropout during inference
    model.train()
    predictions = []

    # Perform multiple forward passes with dropout enabled
    with torch.no_grad():  # Disable gradient calculation for inference
        for _ in range(num_samples):
            preds = model(X)  # Perform a forward pass
            predictions.append(preds)  # Store the predictions

    # Stack the predictions to create a tensor of shape (num_samples, batch_size, output_size)
    predictions = torch.stack(predictions)

    # Compute the mean of the predictions along the first dimension (num_samples)
    prediction_mean = predictions.mean(dim=0)

    # Compute the variance of the predictions along the first dimension (num_samples)
    prediction_variance = predictions.var(dim=0)

    return prediction_mean, prediction_variance

def visualize_predictions(prediction_mean, prediction_variance):
    """
    Visualizes the prediction mean and variance to understand model confidence.

    Args:
        prediction_mean (torch.Tensor): The mean of the predictions.
        prediction_variance (torch.Tensor): The variance of the predictions.
    """
    prediction_mean = prediction_mean.numpy()
    prediction_variance = prediction_variance.numpy()

    x = np.arange(len(prediction_mean))

    plt.figure(figsize=(14, 7))

    plt.subplot(1, 2, 1)
    plt.plot(x, prediction_mean, 'o-', label='Prediction Mean')
    plt.fill_between(x, prediction_mean.squeeze() - prediction_variance.squeeze(), 
                     prediction_mean.squeeze() + prediction_variance.squeeze(), alpha=0.2, label='Variance')
    plt.xlabel('Sample Index')
    plt.ylabel('Prediction Mean')
    plt.title('Prediction Mean with Variance')
    plt.legend()

    plt.subplot(1, 2, 2)
    plt.plot(x, prediction_variance, 'o-', label='Prediction Variance')
    plt.xlabel('Sample Index')
    plt.ylabel('Prediction Variance')
    plt.title('Prediction Variance')
    plt.legend()

    plt.tight_layout()
    plt.show()

# Instantiate the model with specified input size, hidden size, output size, and dropout rate
input_size = 1000
hidden_size = 5
output_size = 1
dropout_rate = 0.5
model = SimpleNNWithDropout(input_size, hidden_size, output_size, dropout_rate)

# Example input data: Batch of 10 samples, each with 10 features
X = torch.randn(15, input_size)

# Perform Monte Carlo Dropout inference with 100 samples
prediction_mean, prediction_variance = mc_dropout_inference(model, X, num_samples=100)

# Print the mean and variance of the predictions
print("Prediction Mean:\n", prediction_mean)
print("Prediction Variance:\n", prediction_variance)

# Visualize the predictions to understand model confidence
visualize_predictions(prediction_mean, prediction_variance)