import numpy as np
import matplotlib.pyplot as plt
from sklearn.neural_network import MLPRegressor
from sklearn.preprocessing import StandardScaler

solover = 'lbfgs' # 'sgd' 'adam'
max_iter = 1000

def func(x):
    return np.sin(0.2 * x ** 2)

def train_mlp(n_hidden, activation_function):
    """
    Trains an MLP with one hidden layer and the given number of hidden nodes and activation function.
    Fits the model to 250 randomly sampled (x, y) points in the range [-10, 10].
    Returns the mean sum of squared error.
    """
    # Generate the dataset
    X = np.random.uniform(-10, 10, size=(250, 1))
    y = func(X.ravel())#np.sin(0.2 * X.ravel() ** 2)
    # Scale the input data
    scaler = StandardScaler()
    scaler.fit(X)
    X_scaled = scaler.transform(X)

    # Train the MLP
    mlp = MLPRegressor(
        hidden_layer_sizes=(n_hidden,),
        activation=activation_function,
        solver=solover,
        random_state=42,
        max_iter=max_iter,
    )
    mlp.fit(X_scaled, y)

    # Compute the mean sum of squared error
    y_pred = mlp.predict(X_scaled)
    mse = np.mean((y - y_pred) ** 2)
    return mse

def get_mean_error(n_hidden, activation_function, num_runs=10):
    """
    Calls the train_mlp function 10 times and calculates the mean error over the runs.
    """
    errors = [train_mlp(n_hidden, activation_function) for _ in range(num_runs)]
    return np.mean(errors)

# Call the get_mean_error function for different numbers of hidden nodes and activation functions
n_hidden_nodes = [3, 5, 7, 9, 11, 13, 15, 19]
mse_logistic = [get_mean_error(n, 'logistic') for n in n_hidden_nodes]
mse_relu = [get_mean_error(n, 'relu') for n in n_hidden_nodes]

# Plot the mean square error as the number of hidden nodes increases
plt.figure(figsize=(10, 6))
plt.plot(n_hidden_nodes, mse_logistic, label='Logistic activation')
plt.plot(n_hidden_nodes, mse_relu, label='ReLU activation')
plt.xlabel('Number of hidden nodes')
plt.ylabel('Mean Squared Error')
plt.title('MLP Approximation of sin(0.2x^2)')
plt.legend()

X = np.random.uniform(-10, 10, size=(250, 1))
y = func(X.ravel())
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

# Plot the best fitting curve with 19 hidden nodes (Logistic activation)
mlp_logistic = MLPRegressor(hidden_layer_sizes=(19,),
                            activation='logistic',
                            solver=solover,
                            random_state=42,
                            max_iter=max_iter,)
mlp_logistic.fit(X_scaled, y)
y_pred_logistic = mlp_logistic.predict(X_scaled)
plt.figure(figsize=(10, 6))
plt.plot(np.linspace(-10,10,100), func(np.linspace(-10,10,100)), label='True function')
plt.scatter(X, y_pred_logistic, label='MLP with Logistic activation',c='k')
plt.xlabel('x')
plt.ylabel('f(x)')
plt.title('Best Fitting Curve (Logistic, 19 hidden nodes)')
plt.legend()

# # Plot the best fitting curve with 19 hidden nodes (ReLU activation)
mlp_relu = MLPRegressor(hidden_layer_sizes=(19,),
                        activation='relu',
                        solver=solover,
                        random_state=42,
                        max_iter=max_iter,)
mlp_relu.fit(X_scaled, y)
y_pred_relu = mlp_relu.predict(X_scaled)
plt.figure(figsize=(10, 6))
plt.plot(np.linspace(-10,10,100), func(np.linspace(-10,10,100)), label='True function')
plt.scatter(X, y_pred_relu, label='MLP with ReLU activation',c='k')
plt.xlabel('x')
plt.ylabel('f(x)')
plt.title('Best Fitting Curve (ReLU, 19 hidden nodes)')
plt.legend()
plt.show()