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

# Load the arrow head dataset
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)

# Ensure the target variables are floats (if needed)
y_train = y_train.astype("float")
y_test = y_test.astype("float")

# Initialize and train the KNeighborsTimeSeriesRegressor
regressor = KNeighborsTimeSeriesRegressor()
regressor.fit(X_train, y_train)

# Predict and evaluate
y_pred = regressor.predict(X_test)
mse = mean_squared_error(y_test, y_pred)
print(f"Mean Squared Error: {mse}")