import os
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
from PyEMD import CEEMDAN
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MinMaxScaler


def ceemdan(x, num_trials):
    time = np.arange(len(x))  # 创建一个时间数组
    cemd = CEEMDAN()
    cemd.trials = num_trials
    imfs = cemd.emd(x, time)  # 在emd函数中传递时间数组
    return imfs


def prepare_data(imfs, seq_len):
    data = np.array(imfs).T
    data = MinMaxScaler().fit_transform(data)
    X = []
    for i in range(len(data) - seq_len):
        X.append(data[i:i + seq_len])
    X = np.array(X)
    X = torch.tensor(X, dtype=torch.float32)
    return X


class Transformer(nn.Module):
    def init(self, d_model, num_heads, hidden_units, num_layers, dropout_rate, seq_len):
        super(Transformer, self).init()
        self.d_model = d_model
        self.num_heads = num_heads
        self.hidden_units = hidden_units
        self.num_layers = num_layers
        self.dropout_rate = dropout_rate
        self.seq_len = seq_len

        self.encoder = nn.TransformerEncoderLayer(d_model, num_heads, hidden_units, dropout_rate, batch_first=True)
        self.transformer = nn.TransformerEncoder(self.encoder, num_layers)
        self.global_avg_pool = nn.AdaptiveAvgPool2d((1, 1))
        self.linear = nn.Linear(d_model, 1)

    def forward(self, x):
        x = self.transformer(x)
        x = self.global_avg_pool(x)
        x = x.view(-1, self.d_model)  # Reshape
        x = self.linear(x)
        return x


# 加载数据
os.chdir("C:\\Users\\Administrator\\Desktop")
data_reader = pd.read_csv("cleaned_data0.csv", encoding='GBK')
# data = data_reader.values.astype(float)
data = data_reader['锚链1拉力'].values.astype(np.float32)
# 分解成IMFs
imfs = []
for chunk in data:
    imfs_chunk = ceemdan(chunk, num_trials=10)  # 对每个分块进行CEEMDAN分解
    imfs.extend(imfs_chunk)

# 准备数据
seq_len = 10
X = prepare_data(imfs, seq_len)
y = data[seq_len:]
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 构建Transformer模型
d_model = 64
num_heads = 4
hidden_units = 64
num_layers = 4
dropout_rate = 0.2

model = Transformer(d_model, num_heads, hidden_units, num_layers, dropout_rate, seq_len)
optim = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.MSELoss()

# 训练模型
num_epochs = 10
batch_size = 32
total_steps = len(X_train) // batch_size

for epoch in range(num_epochs):
    model.train()  # 设置为训练模式
    for i in range(0, len(X_train), batch_size):
        inputs = X_train[i:i + batch_size]
        targets = y_train[i:i + batch_size]

        # 前向传播
        outputs = model(inputs)
        loss = criterion(outputs, targets)

        # 反向传播和优化
        optim.zero_grad()
        loss.backward()
        optim.step()

    print(f"Epoch {epoch + 1}/{num_epochs}, Loss: {loss.item():.4f}")

# 预测
model.eval()  # 设置为评估模式
y_pred = model(X_test).detach().numpy()

# 计算均方误差
mse = mean_squared_error(y_test, y_pred)
print("Mean Squared Error: ", mse)这段代码中有
Traceback (most recent call last):
  File "E:\python-tensorflow\tfgpu\ceemdanfenjie.py", line 63, in <module>
    imfs_chunk = ceemdan(chunk, num_trials=10)  # 对每个分块进行CEEMDAN分解
  File "E:\python-tensorflow\tfgpu\ceemdanfenjie.py", line 14, in ceemdan
    time = np.arange(len(x))  # 创建一个时间数组
TypeError: object of type 'numpy.float32' has no len()这个错误