import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.preprocessing import MinMaxScaler
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, LSTM, Dropout
from transformers import TimeSeriesTransformerModel

# Load the dataset
df = pd.read_csv("GOOG.csv", date_parser=True)

# Select the "Close" column as the target variable
data = df.filter(["Close"]).values

# Split the data into train and test sets
train_size = int(len(data) * 0.8)
train_data = data[:train_size, :]
test_data = data[train_size:, :]

# Scale the data to be between 0 and 1
scaler = MinMaxScaler(feature_range=(0, 1))
train_data = scaler.fit_transform(train_data)
test_data = scaler.transform(test_data)

# Create the time series transformer
transformer = TimeSeriesTransformerModel(
    input_size=1,
    output_size=1,
    seq_len=60,
    n_layers=2,
    n_heads=2,
    d_model=64,
    d_ff=128,
    dropout=0.2,
    activation="relu",
)

# Define the model
model = Sequential()
model.add(transformer)
model.add(LSTM(units=64, return_sequences=True))
model.add(Dropout(0.2))
model.add(LSTM(units=64, return_sequences=True))
model.add(Dropout(0.2))
model.add(LSTM(units=64))
model.add(Dropout(0.2))
model.add(Dense(units=1))

# Compile the model
model.compile(optimizer="adam", loss="mean_squared_error")

# Train the model
model.fit(
    train_data[:-1], train_data[1:], epochs=100, batch_size=32, verbose=2
)

# Make predictions on the test data
predictions = model.predict(test_data[:-1])
predictions = scaler.inverse_transform(predictions)

# Plot the predictions against the actual values
plt.plot(test_data[1:], label="Actual")
plt.plot(predictions, label="Predicted")
plt.legend()
plt.show()
