from keras.models import Sequential
import matplotlib.pyplot as plt
from keras.layers import Dense
import math
import numpy as np


def f1(x):
    return x * math.sin(6*math.pi*x)*math.exp(math.pow(-x, 2))


model = Sequential()
model.add(Dense(200, input_dim=1, activation='sigmoid', weights=[100*np.random.randn(1, 200), 10*np.random.randn(200)]))
model.add(Dense(1, activation='sigmoid'))

model.compile(optimizer='sgd', loss='mse')

draw_x = np.linspace(-1, 1, 10000)
draw_y = np.array([f1(x) for x in draw_x])

data = np.random.random(10000)
data = np.array([-x if i % 2 == 0 else x for i, x in enumerate(data)])
results = np.array([f1(x) for x in data])
model.fit(data[:7000], results[:7000], epochs=200, batch_size=32)
print(model.evaluate(data[7000:], results[7000:]))
predicted_results = model.predict(draw_x)

plt.plot(draw_x, draw_y)
plt.plot(draw_x, predicted_results)
plt.show()