import tensorflow as tf
import tensorflow_addons as tfa
from matplotlib import pyplot as plt
import numpy as np
import optuna

densite = 10

start1 = 5
stop1 = 15
start2 = 35
stop2 = 45

start = -50
stop = 150

training = [np.concatenate((np.linspace(start1, stop1, (stop1-start1)*densite),
							np.linspace(start2, stop2, (stop2-start2)*densite)))]
training.append(np.sin(training[0])/3+0.5)

training[0] = training[0]/(stop-start)



validation = [np.linspace(start, stop, (stop-start)*densite)]
validation.append(np.sin(validation[0])/3+0.5)
validation[0] = validation[0]/(stop-start)


fig = plt.figure(1)
ax1 = fig.add_subplot(2, 1, 1)
ax2 = fig.add_subplot(2, 1, 2, sharex=ax1)

ax1.plot(*validation,'.', label="validation")
ax2.plot(*training,'.', label="apprentissage")

plt.legend()
plt.show()

bce = tf.keras.losses.BinaryCrossentropy()
mse = tf.keras.losses.MeanSquaredError()


def objectif(trial) :
	HIDDEN = trial.suggest_int('hidden',64,512)
	SIZE = trial.suggest_int('size',0,2)
	DROPOUT = trial.suggest_float('dropout', 0,0.3)
	model = tf.keras.Sequential()
	model.add(tf.keras.layers.Dense(HIDDEN, input_shape=(1,), activation='relu'))
	for i in range(SIZE) :
		model.add(tf.keras.layers.Dense(HIDDEN, activation='relu'))
		model.add(tf.keras.layers.Dropout(DROPOUT))
	model.add(tf.keras.layers.Dense(1, activation='sigmoid'))



	model.compile(optimizer='adam',
				  loss='bce',
				  metrics=['mse'])

	history = model.fit(x=training[0], y=training[1], batch_size=4,
		                epochs=3000, shuffle=True,
		                #validation_data=(validation[0], validation[1]),
		                verbose='auto')
	
	pred = model.predict(validation[0])

	plt.figure(2)
	plt.clf()
	plt.plot(*validation, label="validation")
	plt.plot(*training, label="apprentissage")
	plt.plot(validation[0], pred, label="prediction")
	plt.legend()
	rep = mse(pred-0.5,validation[1]-0.5).numpy()
	plt.savefig(f"images_relu/{rep}_{trial.number}_{HIDDEN}_{SIZE}_{DROPOUT}.png")
	return rep

study = optuna.create_study()
study.optimize(objectif, n_trials=25)