import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.datasets import fetch_california_housing
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
import numpy as np

# Load and preprocess the dataset
data = fetch_california_housing()
X, y = data.data, data.target
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# Standardize the features
scaler = StandardScaler()
X_train = scaler.fit_transform(X_train)
X_test = scaler.transform(X_test)

# Convert data to PyTorch tensors
X_train = torch.tensor(X_train, dtype=torch.float32)
y_train = torch.tensor(y_train, dtype=torch.float32)
X_test = torch.tensor(X_test, dtype=torch.float32)
y_test = torch.tensor(y_test, dtype=torch.float32)

# Define a simple feedforward neural network model
class SimpleNN(nn.Module):
    def __init__(self, input_dim):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(input_dim, 64)
        self.fc2 = nn.Linear(64, 64)
        self.fc3 = nn.Linear(64, 1)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# Train a single model
def train_model(model, X_train, y_train, epochs=100, lr=0.001):
    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=lr)
    model.train()

    for epoch in range(epochs):
        optimizer.zero_grad()
        outputs = model(X_train)
        loss = criterion(outputs.squeeze(), y_train)
        loss.backward()
        optimizer.step()
        
    return model

# Define the number of models in the ensemble
ensemble_size = 10
input_dim = X_train.shape[1]

# Create an ensemble of models
ensemble_models = []
for i in range(ensemble_size):
    model = SimpleNN(input_dim)
    trained_model = train_model(model, X_train, y_train, epochs=100, lr=0.001)
    ensemble_models.append(trained_model)

# Make predictions using each model in the ensemble and calculate the mean and variance
def ensemble_predict(models, X):
    predictions = []
    for model in models:
        model.eval()
        with torch.no_grad():
            predictions.append(model(X).squeeze().numpy())
    
    predictions = np.array(predictions)
    mean_prediction = np.mean(predictions, axis=0)
    variance_prediction = np.var(predictions, axis=0)  # Uncertainty estimate
    
    return mean_prediction, variance_prediction

# Get predictions on the test set
mean_pred, var_pred = ensemble_predict(ensemble_models, X_test)

# Print the results
print("Mean Predictions:", mean_pred)
print("Uncertainty (Variance of Predictions):", var_pred)