import pandas as pd

from sklearn.linear_model import LinearRegression

from sklearn.model_selection import cross_val_score, KFold

from sklearn.impute import SimpleImputer

import numpy as np

# Load the data

data = pd.read_csv("./input/data.csv")

# Select only numeric columns for imputation

numeric_columns = data.select_dtypes(include=[np.number]).columns

X = data[numeric_columns].drop(columns=["x_e_out [-]"])

y = data["x_e_out [-]"]

# Split the data into training and NaN sets

train_data = data.dropna(subset=["x_e_out [-]"])

nan_data = data[data["x_e_out [-]"].isna()]

# Impute missing values in features with mean

imputer = SimpleImputer(strategy="mean")

X_train_imputed = imputer.fit_transform(

train_data[numeric_columns].drop(columns=["x_e_out [-]"])

)

y_train = train_data["x_e_out [-]"]

# Train the linear regression model using cross-validation

model = LinearRegression()

kf = KFold(n_splits=10, shuffle=True, random_state=1)

rmse_scores = cross_val_score(

model, X_train_imputed, y_train, scoring="neg_root_mean_squared_error", cv=kf

)

print(f"10-fold CV RMSE: {-np.mean(rmse_scores):.4f} ( /- {np.std(rmse_scores):.4f})")

# Fit the model on the entire training set

model.fit(X_train_imputed, y_train)

# Load the test data

test_data = pd.read_csv("./input/sample_submission.csv")

# Prepare the test features

X_test = nan_data[numeric_columns].drop(columns=["x_e_out [-]"])

X_test_imputed = imputer.transform(X_test)

# Predict the missing values for the test set

nan_data["x_e_out [-]"] = model.predict(X_test_imputed)

# Merge predictions back into the test data

test_data = test_data.merge(nan_data[["id", "x_e_out [-]"]], on="id", how="left")

# Save the submission file

test_data.to_csv("./working/submission.csv", index=False)