#!/usr/bin/env python
# coding: utf-8

# In[2]:


import pandas as pd
from tqdm import tqdm
import os
#先分训练集和测试集
folder_path = "/home/lxm/1vs1/total"  # 替换为你的文件夹路径
# # 获取文件夹中的所有文件
file_list = os.listdir(folder_path)
data_train=pd.DataFrame()
data_test=pd.DataFrame()
count=0
for file_name in tqdm(file_list, desc="Processing files"):
    if file_name.endswith(".csv"):
        count+=1
        if count<=4:
            file_path = os.path.join(folder_path, file_name)
            # 解析CSV文件
            df = pd.read_csv(file_path)  # 读取csv文件到dataframe中
            data =df.iloc[:,1:] 
            data1=pd.DataFrame(data,columns=['tar_x','tar_y','tar_z','tarEli','position_change','v','v_change','r','eli_change',
                                            'h_change','id'])
            data_train = pd.concat([data_train, data1], axis=0, ignore_index=True)  # 将文件内容合并到data_all中
#             print("train",file_path)
        else:
            file_path = os.path.join(folder_path, file_name)
            # 解析CSV文件
            df = pd.read_csv(file_path)  # 读取csv文件到dataframe中
            data =df.iloc[:,1:] 
            data1=pd.DataFrame(data,columns=['tar_x','tar_y','tar_z','tarEli','position_change','v','v_change','r','eli_change',
                                            'h_change','id'])
            data_test = pd.concat([data_test, data1], axis=0, ignore_index=True)  # 将文件内容合并到data_all中
#             print("test",file_path)



# In[33]:


import numpy as np
#分别在测试集和训练集中分x和y
def create_dataset(dataset,window,offset,predict_length):
    dataX, dataY = [], []
    for i in range(0,len(dataset)-window-1,offset):
        if (i + window+predict_length)<len(dataset) :
            id1=dataset[i,-1]
            id2=dataset[i+window-1,-1]
            id3=dataset[i+window+predict_length-1,-1]
            if id1==id2 and id1==id3:
                a = dataset[i:(i+window),:-1]
                dataX.append(a)
                dataY.append(dataset[i + window:(i + window+predict_length),0:3])
    return np.array(dataX), np.array(dataY)


train=data_train.values
test=data_test.values
from sklearn.preprocessing import MinMaxScaler
scaler = MinMaxScaler()
scaler.fit(train) ## 生成规则
data_tr = scaler.transform(train) ## 将规则应用于训练集
data_te = scaler.transform(test) ## 将规则应用于测试集


train_X,train_Y=create_dataset(data_tr,200,100,100)
test_X,test_Y=create_dataset(data_te,200,100,100)




# In[70]:


from keras.models import load_model
# load model from single file
model1 = load_model('/home/lxm/model/lstm_model3.h5')
yhat = model1.predict(test_X)



# In[83]:


from numpy import concatenate
from numpy import sqrt
from sklearn.metrics import mean_squared_error

temp = test_X.reshape(test_X.shape[0]*test_X.shape[1], test_X.shape[2])
temp=temp[:373700,2:]
iny_yhat=yhat.reshape(yhat.shape[0]*yhat.shape[1],yhat.shape[2])
#预测数据逆缩放 invert scaling for forecast
iny = concatenate((iny_yhat, temp), axis=1)
inverse_y = scaler.inverse_transform(iny)
inv_y = inverse_y[:, 0:3]
# inv_yhat = np.array(inv_yhat)
inv_y_reshape=inv_y.reshape(int(inv_y.shape[0]/100),100,3)

#真实数据
tempy=test_Y.reshape(test_Y.shape[0]*test_Y.shape[1], test_Y.shape[2])
iny_temp= concatenate((tempy, temp), axis=1)#水平方向合并，为了适配scaler
iny_temp1=scaler.inverse_transform(iny_temp)
iny_real=iny_temp1[:,0:3]
iny_real_reshape=iny_real.reshape(int(iny_real.shape[0]/100),100,3)

# # # calculate RMSE
rmse = sqrt(mean_squared_error(iny_real_reshape[3],inv_y_reshape[3]))
print('Test RMSE: %.3f' % rmse)


# In[78]:


import matplotlib.pyplot as plt
#画出真实数据和预测数据

predict_x,predict_y,predict_z=[],[],[]
true_x,true_y,true_z=[],[],[]
for i in range(len(inv_y_reshape[1])):
    predict_x.append(inv_y_reshape[1][i][0])
    predict_y.append(inv_y_reshape[1][i][1])
    predict_z.append(inv_y_reshape[1][i][2])
    true_x.append(iny_real_reshape[1][i][0])
    true_y.append(iny_real_reshape[1][i][1])
    true_z.append(iny_real_reshape[1][i][2])
fig = plt.figure()
ax = fig.gca(projection='3d')
ax.scatter(true_x, true_y, true_z, s=2,label='true')
ax.scatter(predict_x, predict_y, predict_z, s=2,label='prediction')
plt.legend()
plt.show()


# In[80]:




# In[ ]:




