import torch
import numpy as np
from torch.utils.data import Dataset
from sklearn.metrics import r2_score, mean_squared_error, mean_absolute_error

# 定义数据集类
class PharmacyDataset(Dataset):
    def __init__(self, data_X, data_Y):
        self.data_X = data_X
        self.data_Y = data_Y
        
    def __getitem__(self, index):
        return self.data_X[index], self.data_Y[index]
    
    def __len__(self):
        return len(self.data_X)

# 定义用于时间序列预测的数据处理函数
def create_dataset(datasetX, datasetY, look_back=1, predict_steps=1):
    dataX, dataY = [], []
    for i in range(0, len(datasetX) - look_back - predict_steps + 1):
        x = datasetX[i:(i + look_back), :]
        if predict_steps == 1:
            y = datasetY[i + look_back]
        else:
            y = datasetY[i + look_back:i + look_back + predict_steps, 0]  # 仅取销量列
        dataX.append(x)
        dataY.append(y)
    return np.array(dataX), np.array(dataY)

# 评估函数
def evaluate_model(y_true, y_pred):
    mse = mean_squared_error(y_true, y_pred)
    rmse = np.sqrt(mse)
    mae = mean_absolute_error(y_true, y_pred)
    r2 = r2_score(y_true, y_pred)
    
    # MAPE计算时避免除以0
    mask = y_true != 0
    y_true_masked = y_true[mask]
    y_pred_masked = y_pred[mask]
    mape = np.mean(np.abs((y_true_masked - y_pred_masked) / y_true_masked)) * 100
    
    return {
        'MSE': mse,
        'RMSE': rmse,
        'MAE': mae,
        'R²': r2,
        'MAPE(%)': mape
    }