import numpy as np
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from scipy.spatial.distance import wasserstein_distance
from sklearn.preprocessing import MinMaxScaler
from pyESN import ESN
from tensorflow_probability import distributions as tfd
from bayes_opt import BayesianOptimization

# Load your data here
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# Scale the data to the range of -1 to 1
scaler = MinMaxScaler(feature_range=(-1, 1))
X_train = scaler.fit_transform(X_train)
X_test = scaler.transform(X_test)

def esn_error(esn_params):
    # Define the ESN object
    esn = ESN(**esn_params)
    
    # Train the ESN on the training data
    train_len = int(0.7*len(X_train))
    esn.fit(X_train[:train_len], y_train[:train_len])
    
    # Generate output for the remaining data
    y_pred = esn.predict(X_train[train_len:])
    
    # Calculate the Wasserstein distance between the predicted and true outputs
    distance = wasserstein_distance(y_train[train_len:].flatten(), y_pred.flatten())
    
    return distance

def sample_hyperparameters(num_samples):
    # Define the distribution over hyperparameters
    lstm_units = tfd.TruncatedNormal(loc=128, scale=32, low=32, high=256)
    spectral_radius = tfd.TruncatedNormal(loc=0.9, scale=0.1, low=0.1, high=1.0)
    sparsity = tfd.TruncatedNormal(loc=0.1, scale=0.05, low=0.01, high=0.5)
    noise = tfd.TruncatedNormal(loc=0.1, scale=0.05, low=0.01, high=0.5)
    
    # Sample hyperparameters from the distribution
    lstm_units_samples = lstm_units.sample(num_samples).numpy().astype(int)
    spectral_radius_samples = spectral_radius.sample(num_samples).numpy()
    sparsity_samples = sparsity.sample(num_samples).numpy()
    noise_samples = noise.sample(num_samples).numpy()
    
    # Create a list of dictionaries with the hyperparameters
    hyperparameters = [{'n_reservoir': lstm_units_samples[i],
                        'spectral_radius': spectral_radius_samples[i],
                        'sparsity': sparsity_samples[i],
                        'noise': noise_samples[i]} for i in range(num_samples)]
    
    return hyperparameters

def evaluate_hyperparameters(num_samples):
    # Sample hyperparameters from the distribution
    hyperparameters = sample_hyperparameters(num_samples)
    
    # Evaluate the ESN on the hyperparameters
    errors = []
    for params in hyperparameters:
        error = esn_error(params)
        errors.append(error)
    
    # Calculate the mean and standard deviation of the errors
    mean_error = np.mean(errors)
    std_error = np.std(errors)
    
    return mean_error, std_error

pbounds = {'num_samples': (10, 100)}

optimizer = BayesianOptimization(
    f=evaluate_hyperparameters,
    pbounds=pbounds,
    random_state=42,
)

optimizer.maximize(init_points=5, n_iter=10)

print(optimizer.max)
