import pandas as pd
import numpy as np
from sklearn.preprocessing import MinMaxScaler
from keras.models import Sequential
from keras.layers import Dense, LSTM
from ta.volatility import BollingerBands
from ta.trend import MACD
from ta.trend import EMAIndicator
from ta.volume import MFIIndicator
from sqlalchemy import create_engine

# 加载数据
engine_tech = create_engine('mysql+pymysql://root:r6kEwqWU9!v3@localhost:3307/qmt_stocks_tech?charset=utf8')
df = pd.read_sql_table('600000.SH_1d', con=engine_tech)


df = df[['open_back', 'high_back', 'low_back', 'close_back', 'volume_back', 'close_back', 'macd', 'willr', 'volume_back', 'dif', 'dea', 'rsi_6', 'rsi_12', 'rsi_24']].values

# df = df[~np.isnan(df).any(axis=1), :]

# 归一化
scaler = MinMaxScaler(feature_range=(0, 1))
scaled_data = scaler.fit_transform(df)
print(len(df), len(scaled_data))

# 构造训练数据
lookback = 90
target_days = 30

x_train = []
y_train = []
for i in range(lookback, len(df) - target_days):
    x_train.append(scaled_data[i - lookback:i])
    y_train.append(scaled_data[i:i + target_days, 0])

x_train, y_train = np.array(x_train), np.array(y_train)
print('11111111111111111111', x_train.shape, y_train.shape)


# 构建模型
model = Sequential()
model.add(LSTM(units=50, return_sequences=True, input_shape=(x_train.shape[1], x_train.shape[2])))
model.add(LSTM(units=50, return_sequences=True))
model.add(LSTM(units=50))
model.add(Dense(units=30))

model.compile(optimizer='adam', loss='mean_squared_error')

model.fit(x_train, y_train, epochs=10, batch_size=32)

# 预测未来30天走势
test_data = scaled_data[-lookback:]
test_data = np.expand_dims(test_data, axis=0)
print(test_data.shape)
predictions = model.predict(test_data)

# 将归一化的预测数据转换为原始数据
predictions = scaler.inverse_transform(predictions)

# 打印预测结果
print(predictions)