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

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

# Build the LSTM model
model = Sequential()
model.add(LSTM(units=224, activation='relu', return_sequences=True, input_shape=(x_train.shape[1], 1)))
model.add(Dropout(0.3))

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

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

model.add(LSTM(units=160, activation='relu'))
model.add(Dropout(0.5))

model.add(Dense(units=1))

import tensorflow as tf

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

# Record the start time of training
start_time = time.time()

# Log the training process
csv_logger = CSVLogger('drive/MyDrive/StockLog/aapl_adx_training_logs_{}.csv'.format(current_time))

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

# Record the end time of training
end_time = time.time()

# Calculate the training time
training_time = end_time - start_time

# Save the trained model weights
model.save_weights("drive/MyDrive/StockModelWeight/aapl_adx_training_model_weights_{}.h5".format(current_time))

print("Training time: {} seconds".format(training_time))

# Prepare test data
past_100_days = pd.DataFrame(train_close[-100:])
test_df = pd.DataFrame(test_close)
final_df = past_100_days.append(test_df, ignore_index=True)

# Scale the input data
input_data = scaler.fit_transform(final_df)

# Prepare test inputs and outputs
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)

# Make predictions
y_pred = model.predict(x_test)

# Scale back the predictions and true values
scale_factor = 1/scaler.scale_
y_pred = y_pred * scale_factor
y_test = y_test * scale_factor

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

import matplotlib.pyplot as plt

# Plot the predicted and true prices
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()