import matplotlib.pyplot as plt
import mplfinance as mpl 
import pandas as pd

td_consumer_key = 'YOUR-KEY-HERE'
# 美国航空股票
ticker = 'AAL'
##periodType - day, month, year, ytd
##period - number of periods to show
##frequencyTYpe - type of frequency for each candle - day, month, year, ytd
##frequency - the number of the frequency type in each candle - minute, daily, weekly
endpoint = 'https://api.tdameritrade.com/v1/marketdata/{stock_ticker}/pricehistory?periodType={periodType}&period={period}&frequencyType={frequencyType}&frequency={frequency}'

# 获取数据
full_url = endpoint.format(stock_ticker=ticker,periodType='year',period=10,frequencyType='daily',frequency=1)
page = requests.get(url=full_url,params={'apikey' : td_consumer_key})
content = json.loads(page.content)

# 转成pandas可处理格式
df = pd.json_normalize(content['candles'])

# 设置时间戳为索引
df['timestamp'] = pd.to_datetime(df.datetime, unit='ms')
df = df.set_index("timestamp")

# 绘制数据
plt.figure(figsize=(15, 6), dpi=80)
plt.plot(df['close'])
plt.legend(['Closing Price'])
plt.show()

# 存储前一天的数据
df["previous_close"] = df["close"].shift(1)
df = df.dropna() # 删除缺失值

# 存储
df.to_csv('../data/stock_'+ticker+'.csv', mode='w', index=True, header=True)

# 读取数据
ticker = 'AAL'
df = pd.read_csv("../data/stock_"+ticker+".csv")

# 设置索引
df['DateIndex'] = pd.to_datetime(df['timestamp'], format="%Y/%m/%d")
df = df.set_index('DateIndex')

# 幅度缩放
df2 = df
cols = ['close', 'volume', 'previous_close']
features = df2[cols]
scaler = MinMaxScaler(feature_range=(0, 1)).fit(features.values)
features = scaler.transform(features.values)
df2[cols] = features

# 收盘价设为目标字段
X = df2.drop(['close','timestamp'], axis =1)
y = df2['close']

import math
# 计算切分点（以80%的训练数据为例）
train_percentage = 0.8
split_point = math.floor(len(X) * train_percentage)

# 时序切分
train_x, train_y = X[:split_point], y[:split_point]
test_x, test_y = X[split_point:], y[split_point:]

# 构建滑窗数据
import numpy.lib
from numpy.lib.stride_tricks import sliding_window_view

def genWindows(X_in, y_in, window_size):
    X_out = []
    y_out = []
    length = X_in.shape[0]
    for i in range(window_size, length):
        X_out.append(X_in[i-window_size:i, 0:4])
        y_out.append(y_in[i-1])
    return np.array(X_out), np.array(y_out)

# 窗口大小为5
window_size = 5
X_train_win, y_train_win = genWindows(np.array(train_x), np.array(train_y), window_size)
X_test_win, y_test_win = genWindows(np.array(test_x), np.array(test_y), window_size)

from tensorflow.keras import callbacks

# 早停止 回调函数
callback_early_stopping = callbacks.EarlyStopping(
    monitor="loss",
    patience=10,#look at last 10 epochs
    min_delta=0.0001,#loss must improve by this amount
    restore_best_weights=True,
)


from tensorflow import keras
from tensorflow.keras import layers
from keras.models import Sequential

# 构建RNN模型，结构为 输入-RNN-RNN-连续值输出
input_shape=(X_train_win.shape[1],X_train_win.shape[2])
print(input_shape)
model = Sequential(
    [
        layers.Input(shape=input_shape),
        layers.SimpleRNN(units=128, return_sequences=True),
        layers.SimpleRNN(64, return_sequences=False),
        layers.Dense(1, activation="linear"),
    ]
)

# 优化器
optimizer = keras.optimizers.Nadam(learning_rate=0.0001)
model.compile(optimizer=optimizer, loss="mse")

# 模型结构总结
model.summary()

# 模型训练
batch_size = 20
epochs = 50
history = model.fit(X_train_win, y_train_win,
  batch_size=batch_size, epochs=epochs,
  callbacks=[
      callback_early_stopping
    ])

# 训练集预测
pred_train_y = model.predict(X_train_win)

# 绘图
plt.figure(figsize=(15, 6), dpi=80)
plt.plot(np.array(train_y))
plt.plot(pred_train_y)
plt.legend(['Actual', 'Predictions'])
plt.show()

# 测试集预测
pred_test_y = model.predict(X_test_win)

# 预测结果绘制
plt.figure(figsize=(15, 6), dpi=80)
plt.plot(np.array(test_y))
plt.plot(pred_test_y)
plt.legend(['Actual', 'Predictions'])
plt.show()