import time
from datetime import datetime
from tensorflow.keras.layers import Dense, Dropout, LSTM
from tensorflow.keras.models import Sequential
from tensorflow.keras.callbacks import CSVLogger

current_time = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")

# Define the LSTM model architecture
model = Sequential()
model.add(LSTM(units=160, activation='relu', return_sequences=True, input_shape=(x_train.shape[1], 1)))
model.add(Dropout(0.2))

model.add(LSTM(units=64, activation='relu', return_sequences=True))
model.add(Dropout(0.2))

model.add(LSTM(units=64, activation='relu', return_sequences=True))
model.add(Dropout(0.1))

model.add(LSTM(units=64, activation='relu'))
model.add(Dropout(0.4))

model.add(Dense(units=1))

model.summary()

import tensorflow as tf

# Compile the model
model.compile(optimizer='adam', loss='mean_squared_error', metrics=[tf.keras.metrics.MeanAbsoluteError()])

start_time = time.time()
# Record the training end time
csv_logger = CSVLogger('drive/MyDrive/StockLog/amgn_training_logs_{}.csv'.format(current_time))

# Train the model
model.fit(x_train, y_train, epochs=100, callbacks=[csv_logger])

end_time = time.time()
# Calculate training time
training_time = end_time - start_time

# Save the trained model weights
model.save("drive/MyDrive/StockModelWeight/amgn_model_weights_{}.h5".format(current_time))
print("Model training time: {} seconds".format(training_time))

# Prepare data for prediction
past_100_days = pd.DataFrame(train_close[-100:])
print(past_100_days)

test_df = pd.DataFrame(test_close)

final_df = past_100_days.append(test_df, ignore_index=True)
input_data = scaler.fit_transform(final_df)

x_test = []
y_test = []
for i in range(100, input_data.shape[0]):
    x_test.append(input_data[i-100: i])
    y_test.append(input_data[i, 0])

x_test, y_test = np.array(x_test), np.array(y_test)
print(x_test.shape)
print(y_test.shape)

# Make predictions
y_pred = model.predict(x_test)

print(scaler.scale_)

scale_factor = 1/scaler.scale_
y_pred = y_pred * scale_factor
y_test = y_test * scale_factor

# Save predictions to a CSV file
df = pd.DataFrame({'y_pred': y_pred.flatten(), 'y_test': y_test})
df.to_csv('drive/MyDrive/StockResultData/predictions_amgn.csv', index=False)

import matplotlib.pyplot as plt

# Visualize the predictions
plt.figure(figsize=(12,6))
plt.plot(y_test, 'b', label="Original Price")
plt.plot(y_pred, 'r', label="Predicted Price")
plt.xlabel('Time')
plt.ylabel('Price')
plt.legend()
plt.grid(True)
plt.show()