from sktime.datasets import load_arrow_head
from sktime.regression.distance_based import KNeighborsTimeSeriesRegressor
from sklearn.metrics import mean_squared_error

# Tải tập dữ liệu arrow head
X_train, y_train = load_arrow_head(split="train", return_X_y=True)
X_test, y_test = load_arrow_head(split="test", return_X_y=True)

# Đảm bảo các biến mục tiêu là số thực (nếu cần)
y_train = y_train.astype("float")
y_test = y_test.astype("float")

# Khởi tạo và huấn luyện KNeighborsTimeSeriesRegressor
regressor = KNeighborsTimeSeriesRegressor()
regressor.fit(X_train, y_train)

# Dự đoán và đánh giá
y_pred = regressor.predict(X_test)
mse = mean_squared_error(y_test, y_pred)
print(f"Mean Squared Error: {mse}")