# LSTM Implementation
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import mean_squared_error
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, LSTM, Dropout
from tensorflow.keras.callbacks import EarlyStopping
import yfinance as yf

# 1. Load the Dataset
stock_data = yf.download('AAPL', start='2012-01-01', end='2022-01-01')
df = stock_data[['Close']]

# Apply a moving average to smooth the data
df['Close'] = df['Close'].rolling(window=5).mean()
df.dropna(inplace=True)  # Drop NaN values created by the moving average

# 2. Data Preprocessing
scaler = MinMaxScaler(feature_range=(0, 1))
scaled_data = scaler.fit_transform(df)
# Create dataset for LSTM
def create_dataset(data, time_step=1):
    X, Y = [], []
    for i in range(len(data)-time_step-1):
        X.append(data[i:(i+time_step), 0])
        Y.append(data[i + time_step, 0])
    return np.array(X), np.array(Y)
    
# Prepare the dataset with a time step of 60 (60 days of data)
time_step = 60
X, Y = create_dataset(scaled_data, time_step)

# Split the data into training and test sets
train_size = int(len(X) * 0.8)
X_train, X_test = X[:train_size], X[train_size:]
Y_train, Y_test = Y[:train_size], Y[train_size:]

# Reshape the input to be [samples, time steps, features] for LSTM
X_train = X_train.reshape(X_train.shape[0], X_train.shape[1], 1)
X_test = X_test.reshape(X_test.shape[0], X_test.shape[1], 1)

# 3. Build and Train the LSTM Model
model = Sequential()
model.add(LSTM(50, return_sequences=True, input_shape=(time_step, 1)))
model.add(Dropout(0.2))  # Adding dropout for regularization
model.add(LSTM(50, return_sequences=False))
model.add(Dropout(0.2))  # Adding another dropout layer
model.add(Dense(25))
model.add(Dense(1))
model.compile(optimizer='adam', loss='mean_squared_error')

# Add early stopping
early_stop = EarlyStopping(monitor='val_loss', patience=10, restore_best_weights=True)

# Train the model
history = model.fit(X_train, Y_train, epochs=100, batch_size=32, validation_data=(X_test, Y_test), callbacks=[early_stop], verbose=1)

# 4. Prediction and Evaluation
train_predict = model.predict(X_train)
test_predict = model.predict(X_test)

# Inverse transform to get the actual values
train_predict = scaler.inverse_transform(train_predict)
test_predict = scaler.inverse_transform(test_predict)
Y_train = scaler.inverse_transform([Y_train])
Y_test = scaler.inverse_transform([Y_test])

# Calculate RMSE
train_rmse = np.sqrt(mean_squared_error(Y_train[0], train_predict[:,0]))
test_rmse = np.sqrt(mean_squared_error(Y_test[0], test_predict[:,0]))

print(f'Train RMSE: {train_rmse}')
print(f'Test RMSE: {test_rmse}')

# 5. Future Predictions using LSTM
n_future = 100  # Predicting 30 days into the future
last_sequence = scaled_data[-time_step:]
predictions = []

for _ in range(n_future):
    next_prediction = model.predict(last_sequence.reshape(1, time_step, 1))
    predictions.append(next_prediction[0, 0])
    last_sequence = np.append(last_sequence[1:], next_prediction, axis=0)
predictions = np.array(predictions)
future_predictions = scaler.inverse_transform(predictions.reshape(-1, 1))

# 6. Visualizations
# LSTM Predictions vs. Actual Values
plt.figure(figsize=(12,6))
# Plot the actual data
plt.plot(df.index, df['Close'], label='Actual', color='blue')

# Plot the training predictions
plt.plot(df.index[time_step:time_step + len(train_predict)], train_predict, color='green', label='Train Predictions')

# Plot the testing predictions
plt.plot(df.index[-len(test_predict):], test_predict, color='red', label='Test Predictions')

# Plot the future predictions
future_dates = pd.date_range(df.index[-1], periods=n_future, freq='D')
plt.plot(future_dates, future_predictions, color='orange', label='LSTM Future Predictions')

plt.legend()
plt.title('LSTM Predictions vs. Actual')
plt.xlabel('Time')
plt.ylabel('Stock Price')
plt.show()