import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

SEED = 42
np.random.seed(SEED)

df = pd.read_csv("used_car_price_dataset_extended.csv")

display(df.head())
display(df.describe())
display(df.describe(include=['O']))
df.info()
display(df.shape)

df = df.fillna("None")

# Distributions of categorical variables
sns.set_palette("colorblind")

# service_history
ax = sns.catplot(data=df, x='service_history', kind='count')
ax.set(title='Distribution of Service History')
plt.xticks(rotation=45)
plt.show()

# fuel_type
ax = sns.catplot(data=df, x='fuel_type', kind='count')
ax.set(title='Distribution of Fuel Type')
plt.xticks(rotation=45)
plt.show()

# transmission
ax = sns.catplot(data=df, x='transmission', kind='count')
ax.set(title='Distribution of Transmission Type')
plt.xticks(rotation=45)
plt.show()

# insurance_valid
ax = sns.catplot(data=df, x='insurance_valid', kind='count')
ax.set(title='Distribution of Insurance')
plt.show()

# owner_count
# owner_count is numeric, but I treat it as categorical since a bar plot makes more sense
ax = sns.catplot(data=df, x='owner_count', kind='count')
ax.set(title = 'Number of Cars By Owner Count')
plt.show()

# Distributions of continuous variables

# make_year
ax = sns.displot(df, x='make_year', bins=29)  # One bin per make year (1995-2023)
ax.set(xlabel='Year', ylabel='Number of Cars', title='Histogram of Car Make Year')
plt.show()

# price_usd
ax = sns.displot(df, x='price_usd')
ax.set(xlabel='Price', ylabel='Count', title='Histogram of Car Prices')
plt.show()

# mileage_kmpl
ax = sns.displot(data=df, x='mileage_kmpl')
ax.set(xlabel='KM Per Liter', ylabel='Count', title='Histogram Of Mileage - KM Per Liter')
plt.show()

sns.relplot(x="mileage_kmpl", y="price_usd", data=df, height=6, aspect=1.6)
plt.title('Mileage (km/l) vs Price')
plt.xlabel('Mileage (KM/l)')
plt.ylabel('Price')
plt.show()

# Mean car price for each year
df_line_mean = df.groupby('make_year')['price_usd'].mean().reset_index()

# Line plot of mean car price by year
sns.relplot(x="make_year", y="price_usd", kind="line", data=df_line_mean)
plt.title('Make Year vs Mean Car Price', fontsize=14)
plt.xlabel('Make Year', fontsize=12)
plt.ylabel('Price', fontsize=14)
plt.xticks(rotation=45)
plt.show()

# Median car price for each year
df_line_median = df.groupby('make_year')['price_usd'].median().reset_index()

# Line plot of median car price by year
sns.relplot(x="make_year", y="price_usd", kind="line", data=df_line_median)
plt.title('Make Year vs Median Car Price', fontsize=14)
plt.xlabel('Make Year', fontsize=12)
plt.ylabel('Price', fontsize=14)
plt.xticks(rotation=45)
plt.show()

# Mean engine capacity for each year
df_line_mean = df.groupby('make_year')['engine_cc'].mean().reset_index()

# Line plot of mean engine capacity by year
sns.relplot(x="make_year", y="engine_cc", kind="line", data=df_line_mean)
plt.title('Engine Capacity Each Year using Mean', fontsize=14)
plt.xlabel('Make Year', fontsize=12)
plt.ylabel('Engine Capacity', fontsize=14)
plt.xticks(rotation=45)
plt.show()

# Median engine capacity for each year
df_line_median = df.groupby('make_year')['engine_cc'].median().reset_index()

# Line plot of median engine capacity by year
sns.relplot(x="make_year", y="engine_cc", kind="line", data=df_line_median)
plt.title('Engine Capacity Each Year using Median', fontsize=14)
plt.xlabel('Make Year', fontsize=12)
plt.ylabel('Engine Capacity', fontsize=14)
plt.xticks(rotation=45)
plt.show()

# Return the difference in averages of column_name between the two groups of grouping_var
def diff_of_avgs(df_in, column_name, grouping_var):
    grpby_var = df_in.groupby(grouping_var, observed=True)
    avgs = grpby_var[column_name].mean()
    # Assumes df_in contains exactly two groups; the result is (second group - first group)
    return avgs.iloc[1] - avgs.iloc[0]

def bootstrap_mean_difference(original_sample, column_name, grouping_var, num_replications):
    '''This function returns an array of bootstrapped differences between two sample averages:
      original_sample: df containing the original sample
      column_name: name of column containing the variable to average
      grouping_var: name of variable according to which to group
      num_replications: number of bootstrap samples'''
    original_sample_size = original_sample.shape[0]  # Resample with the same sample size
    original_sample_cols_of_interest = original_sample[[column_name, grouping_var]]
    bstrap_mean_diffs = np.empty(num_replications)
    for i in range(num_replications):
        # Sample with replacement
        bootstrap_sample = original_sample_cols_of_interest.sample(original_sample_size, replace=True)
        resampled_mean_diff = diff_of_avgs(bootstrap_sample, column_name, grouping_var)
        bstrap_mean_diffs[i] = resampled_mean_diff
    return bstrap_mean_diffs

# Fuel type pairs to compare
pairs = [("Diesel", "Petrol"), ("Diesel", "Electric"), ("Electric", "Petrol")]

# Run the bootstrap for each pair
results = []  # Results for each pair
NUM_REPS = 5000

for g1, g2 in pairs:
    # Keep only the two groups in this pair
    sub_df = df[df['fuel_type'].isin([g1, g2])].copy()

    sub_df['fuel_type'] = pd.Categorical(sub_df['fuel_type'], categories=[g2, g1], ordered=True)

    # Observed difference in means
    obs_diff = diff_of_avgs(sub_df, 'price_usd', 'fuel_type')

    # Bootstrap distribution of the mean differences
    bstrap_diffs = bootstrap_mean_difference(
        original_sample=sub_df,
        column_name='price_usd',
        grouping_var='fuel_type',
        num_replications=NUM_REPS
    )

    # 95% confidence interval
    left_end = np.percentile(bstrap_diffs, 2.5, method='higher')
    right_end = np.percentile(bstrap_diffs, 97.5, method='higher')

    # P-value: the bootstrap distribution is centered at the observed difference, so shift it to be
    # centered at 0 (the null hypothesis) and count how often a difference at least as extreme appears
    null_diffs = bstrap_diffs - obs_diff
    p_value = np.mean(np.abs(null_diffs) >= np.abs(obs_diff))

    # Store results
    results.append({
        'Pair': f'{g1} vs {g2}',
        'Observed Diff in Means': f'{obs_diff:.2f}',
        '95% Confidence Interval': f'[{left_end:.2f}, {right_end:.2f}]',
        'p-value': p_value
    })

    # Plot the bootstrap distribution with the confidence interval
    ax = sns.displot(bstrap_diffs, height=4.5, aspect=1.5)
    plt.hlines(y=0, xmin=left_end, xmax=right_end,
               colors='orange', linestyles='solid', lw=7, clip_on=False)
    ax.set(
        xlabel=f'Difference in mean price ({g1} - {g2})',
        ylabel='frequency',
        title=f'Distribution of bootstrap estimates for the mean price difference\n{g1} vs {g2}'
    )

# Summary table
results_df = pd.DataFrame(results)
print(results_df.to_string(index=False))

ax = sns.catplot(data=df, x='fuel_type',y='price_usd',kind='bar',height=4,aspect=1)
ax.set(title='Average Car Price by Fuel Type', xlabel='Fuel Type', ylabel='Average Price')
plt.xticks(rotation=45)
plt.show()

# Average price of electric cars
mean_electric = df[df['fuel_type'] == 'Electric']['price_usd'].mean()

# Number of electric cars
sum_electric = df[df['fuel_type'] == 'Electric'].shape[0]

# Average price of non-electric cars (Petrol or Diesel)
mean_non_electric = df[df['fuel_type'].isin(['Petrol', 'Diesel'])]['price_usd'].mean()

print(f"Electric cars average price: {mean_electric:.2f}")
print(f"\nNon - Electric cars (Petrol and Diesel) average price: {mean_non_electric:.2f} ")
print(f"\nThere are {sum_electric} Electric cars out of 10000")

# Test settings
COLUMN_TO_TEST = 'price_usd'
GROUPING_VARIABLE = 'transmission'
GROUP_1 = 'Automatic'
GROUP_2 = 'Manual'
NUM_REPS = 5000

sub_df = df.copy()

sub_df[GROUPING_VARIABLE] = pd.Categorical(sub_df[GROUPING_VARIABLE], categories=[GROUP_2, GROUP_1], ordered=True)

# Observed difference in means
obs_diff = diff_of_avgs(sub_df, COLUMN_TO_TEST, GROUPING_VARIABLE)

# Bootstrap distribution of the mean differences
bstrap_diffs = bootstrap_mean_difference(
    original_sample=sub_df,
    column_name=COLUMN_TO_TEST,
    grouping_var=GROUPING_VARIABLE,
    num_replications=NUM_REPS
)

# 95% confidence interval
left_end = np.percentile(bstrap_diffs, 2.5)
right_end = np.percentile(bstrap_diffs, 97.5)

# P-value: the bootstrap distribution is centered at the observed difference, so shift it to be
# centered at 0 (the null hypothesis) and count how often a difference at least as extreme appears
null_diffs = bstrap_diffs - obs_diff
p_value = np.mean(np.abs(null_diffs) >= np.abs(obs_diff))

# Summary table
results = {
    'Comparison': f'{GROUP_1} vs {GROUP_2}',
    'Observed Diff in Means (Automatic - Manual)': f'${obs_diff:,.2f}',
    '95% Confidence Interval': f'[${left_end:,.2f}, ${right_end:,.2f}]',
    'P-value': f'{p_value:.4f}'
}
results_df = pd.DataFrame([results])
print("Hypothesis Test Results:")
print(results_df.to_string(index=False))

# Plot the bootstrap distribution with the confidence interval
ax = sns.displot(bstrap_diffs, height=4.5, aspect=1.5)
plt.hlines(y=0, xmin=left_end, xmax=right_end,
            colors='orange', linestyles='solid', lw=10, label='95% Confidence Interval')

plt.title(f'Bootstrap Distribution of Mean Price Difference\n({GROUP_1} vs. {GROUP_2})', fontsize=16)
plt.xlabel(f'Difference in Mean Price (USD)', fontsize=12)
plt.ylabel('Frequency', fontsize=12)
plt.show()

ax = sns.catplot(data=df, x='transmission',y='price_usd',kind='bar',height=4,aspect=1)
ax.set(title='Average Car Price by Transmission Type', xlabel='Transmission Type', ylabel='Average Price')
plt.xticks(rotation=45)
plt.show()

# Average price of manual cars
mean_manual = df[df['transmission'] == 'Manual']['price_usd'].mean()

# Number of manual cars
sum_manual = df[df['transmission'] == 'Manual'].shape[0]

# Average price of automatic cars
mean_auto = df[df['transmission'] == 'Automatic']['price_usd'].mean()

print(f"Manual cars average price: {mean_manual:.2f}")
print(f"\nAutomatic cars average price: {mean_auto:.2f} ")
print(f"\nThere are {sum_manual} Manual cars out of 10000")

# Split the continuous price into 3 equal-sized bins, stored as integers (as KNN expects):
#   0 - Budget (bottom third)
#   1 - Mid-Range (middle third)
#   2 - Expensive (top third)
df['price_bracket'] = pd.qcut(df['price_usd'], q=3, labels=[0, 1, 2]).astype(int)

# Price range and number of cars in each bin
ranges = pd.qcut(df['price_usd'], q=3).cat.categories
counts = df['price_bracket'].value_counts().sort_index()

print(f"Budget price range [{ranges[0].left:.0f}, {ranges[0].right:.0f}] - {counts[0]} cars")
print(f"Mid-Range price range [{ranges[1].left:.0f}, {ranges[1].right:.0f}] - {counts[1]} cars")
print(f"Expensive price range [{ranges[2].left:.0f}, {ranges[2].right:.0f}] - {counts[2]} cars")


# Remove price_usd to prevent leakage
df = df.drop(columns=['price_usd'])
target = 'price_bracket'

df.info()

# One-hot encoding
encoded_df = pd.get_dummies(df,
                                columns=['fuel_type', 'brand', 'transmission', 'color', 'insurance_valid'], drop_first=True)
encoded_df.head()

# Ordinal encoding of service_history
encoded_df['service_history'] = encoded_df['service_history'].map({'None': 1, 'Partial': 2, 'Full': 3})
encoded_df["service_history"] = encoded_df["service_history"].astype(int)
encoded_df.head()

from sklearn.model_selection import train_test_split

# Shuffle the dataframe
encoded_df = encoded_df.sample(frac=1, random_state=SEED)

# Make sure price_bracket exists
assert 'price_bracket' in encoded_df.columns

Y = encoded_df['price_bracket']
X = encoded_df.drop(columns=['price_bracket'])

X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size=0.20, random_state=SEED)

# Combine the training features and target to compute correlations
train_data = pd.concat([X_train, Y_train], axis=1)
train_data.corr()

# Correlation between each pair of variables
correlations = train_data.corr()

# Plot the heatmap
plt.figure(figsize=(10,10))
sns.heatmap(correlations, annot=True)
plt.show()

# Drop color and repeat the encoding, split and correlation steps
encoded_df = df.drop(columns=['color'])

encoded_df = pd.get_dummies(encoded_df,
                                columns=['fuel_type', 'brand', 'transmission', 'insurance_valid'], drop_first=True)

encoded_df['service_history'] = encoded_df['service_history'].map({'None': 1, 'Partial': 2, 'Full': 3})
encoded_df["service_history"] = encoded_df["service_history"].astype(int)

# Shuffle the dataframe
encoded_df = encoded_df.sample(frac=1, random_state=SEED)

# Make sure price_bracket exists
assert 'price_bracket' in encoded_df.columns

Y = encoded_df['price_bracket']
X = encoded_df.drop(columns=['price_bracket'])

X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size=0.20, random_state=SEED)



train_data = pd.concat([X_train, Y_train], axis=1)


# Correlation between each pair of variables
correlations = train_data.corr()

# Plot the heatmap
plt.figure(figsize=(15,15))
sns.heatmap(correlations, annot=True, cmap="seismic")
plt.show()

correlations.reindex(correlations["price_bracket"].abs().sort_values().index)["price_bracket"]  # Sorted by absolute correlation

# Keep only the selected features (in both the train and test sets)
selected_features = ['engine_cc', 'make_year', 'owner_count', 'mileage_kmpl', 'fuel_type_Electric',
                     'fuel_type_Petrol']

X_train = X_train[selected_features]
display(X_train)

X_test = X_test[selected_features]
X_test

# Standardize the features using z-scores. KNN is a distance-based algorithm, so without standardization,
# features with larger numeric ranges would dominate the distance metric. Giving each feature a mean of 0
# and a standard deviation of 1 makes all features contribute equally to the neighbor search
from sklearn.preprocessing import StandardScaler
df_columns = X_train.columns
scaler = StandardScaler()
scaled_X_train = scaler.fit_transform(X_train)

scaled_X_test = scaler.transform(X_test)

scaled_df = pd.DataFrame(scaled_X_train, columns=df_columns)
scaled_df.describe()

# 10-fold cross-validation on the training data only, to find the best k
from sklearn.neighbors import KNeighborsClassifier
from sklearn.model_selection import cross_val_score

mean_cv_scores = []
k_list = range(1, 51)
for nn in k_list:
    knn_cv = KNeighborsClassifier(n_neighbors=nn)
    cv_scores = cross_val_score(knn_cv, scaled_X_train, Y_train, cv=10)
    mean_cv_scores.append(cv_scores.mean())

# Output results
best_k = k_list[int(np.argmax(mean_cv_scores))]  # The k with the highest mean CV accuracy
print('Highest accuracy is obtained for k =', best_k, 'and equals', max(mean_cv_scores))
plt.plot(k_list, mean_cv_scores, '-o')
plt.xlabel('Number of neighbors (k)')
plt.ylabel('Accuracy');

# Retrain KNN with the best k on the whole training set and evaluate it on the test set
# The confusion matrix, precision and recall are adjusted for 3 classes instead of 2
from sklearn.metrics import confusion_matrix
from sklearn.metrics import precision_score
from sklearn.metrics import recall_score

knn_classifier = KNeighborsClassifier(n_neighbors=best_k)
knn_classifier.fit(scaled_X_train, Y_train)
print('accuracy of the classifier is', knn_classifier.score(scaled_X_test, Y_test))

# Confusion matrix
predictions = knn_classifier.predict(X=scaled_X_test)  # The classifier's predictions
print('confusion matrix: \n', confusion_matrix(y_true=Y_test, y_pred=predictions, labels=[0, 1, 2]))  # Rows are true values, columns are predicted values

# Precision and recall, macro-averaged over the 3 classes
print('precision: ', precision_score(y_true=Y_test, y_pred=predictions, average='macro'))
print('recall: ', recall_score(y_true=Y_test, y_pred=predictions, average='macro'))

# Keep only the 4 features with the highest correlation (in both the train and test sets)
selected_features = ['engine_cc', 'make_year', 'owner_count', 'fuel_type_Electric']

X_train = X_train[selected_features]
display(X_train)

X_test = X_test[selected_features]
X_test

# Standardize the features using z-scores (same as above)
from sklearn.preprocessing import StandardScaler
df_columns = X_train.columns
scaler = StandardScaler()
scaled_X_train = scaler.fit_transform(X_train)

scaled_X_test = scaler.transform(X_test)

scaled_df = pd.DataFrame(scaled_X_train, columns=df_columns)
scaled_df.describe()

# 10-fold cross-validation on the training data only, to find the best k
from sklearn.neighbors import KNeighborsClassifier
from sklearn.model_selection import cross_val_score

mean_cv_scores = []
k_list = range(1, 51)
for nn in k_list:
    knn_cv = KNeighborsClassifier(n_neighbors=nn)
    cv_scores = cross_val_score(knn_cv, scaled_X_train, Y_train, cv=10)
    mean_cv_scores.append(cv_scores.mean())

# Output results
best_k = k_list[int(np.argmax(mean_cv_scores))]  # The k with the highest mean CV accuracy
print('Highest accuracy is obtained for k =', best_k, 'and equals', max(mean_cv_scores))
plt.plot(k_list, mean_cv_scores, '-o')
plt.xlabel('Number of neighbors (k)')
plt.ylabel('Accuracy');

# Retrain KNN with the best k on the whole training set and evaluate it on the test set
# The confusion matrix, precision and recall are adjusted for 3 classes instead of 2
from sklearn.metrics import confusion_matrix
from sklearn.metrics import precision_score
from sklearn.metrics import recall_score

knn_classifier = KNeighborsClassifier(n_neighbors=best_k)
knn_classifier.fit(scaled_X_train, Y_train)
print('accuracy of the classifier is', knn_classifier.score(scaled_X_test, Y_test))

# Confusion matrix
predictions = knn_classifier.predict(X=scaled_X_test)  # The classifier's predictions
print('confusion matrix: \n', confusion_matrix(y_true=Y_test, y_pred=predictions, labels=[0, 1, 2]))  # Rows are true values, columns are predicted values

# Precision and recall, macro-averaged over the 3 classes
print('precision: ', precision_score(y_true=Y_test, y_pred=predictions, average='macro'))
print('recall: ', recall_score(y_true=Y_test, y_pred=predictions, average='macro'))

# Reload the dataset, since price_usd was dropped in the classification section
reg_df = pd.read_csv("used_car_price_dataset_extended.csv")
reg_df = reg_df.fillna("None")

# Drop color, as in the classification section
reg_df = reg_df.drop(columns=['color'])

# One-hot encoding of the nominal features
reg_df = pd.get_dummies(reg_df, columns=['fuel_type', 'brand', 'transmission', 'insurance_valid'], drop_first=True)

# Ordinal encoding of service_history
reg_df['service_history'] = reg_df['service_history'].map({'None': 1, 'Partial': 2, 'Full': 3})
reg_df['service_history'] = reg_df['service_history'].astype(int)

Y_reg = reg_df['price_usd']
X_reg = reg_df.drop(columns=['price_usd'])

X_train_reg, X_test_reg, Y_train_reg, Y_test_reg = train_test_split(X_reg, Y_reg, test_size=0.20, random_state=SEED)
print('Train set:', X_train_reg.shape, '| Test set:', X_test_reg.shape)

# Correlation of each feature with price_usd, using the training set only
train_data_reg = pd.concat([X_train_reg, Y_train_reg], axis=1)
reg_correlations = train_data_reg.corr()['price_usd']
reg_correlations.reindex(reg_correlations.abs().sort_values().index)  # Sorted by absolute correlation

from sklearn.metrics import r2_score, mean_squared_error, mean_absolute_error


# Return the R², RMSE and MAE of the predictions
def regression_metrics(y_true, y_pred):
    return {
        'R²': r2_score(y_true, y_pred),
        'RMSE': np.sqrt(mean_squared_error(y_true, y_pred)),
        'MAE': mean_absolute_error(y_true, y_pred)
    }


# Baseline: predict the average training price for every car
baseline_predictions = np.full(len(Y_test_reg), Y_train_reg.mean())
baseline_metrics = regression_metrics(Y_test_reg, baseline_predictions)
print('Baseline:', {name: round(value, 3) for name, value in baseline_metrics.items()})

from sklearn.linear_model import LinearRegression

linear_model = LinearRegression()
linear_model.fit(X_train_reg, Y_train_reg)

# Evaluate on the test set
linear_predictions = linear_model.predict(X_test_reg)
linear_metrics = regression_metrics(Y_test_reg, linear_predictions)
print('Linear regression:', {name: round(value, 3) for name, value in linear_metrics.items()})

# Coefficients, sorted by absolute size
coefficients = pd.Series(linear_model.coef_, index=X_train_reg.columns)
coefficients.reindex(coefficients.abs().sort_values(ascending=False).index)

# Predicted vs. actual price, with the line of perfect predictions
plt.figure(figsize=(7, 6))
plt.scatter(Y_test_reg, linear_predictions, alpha=0.3, s=10)
plt.plot([Y_test_reg.min(), Y_test_reg.max()], [Y_test_reg.min(), Y_test_reg.max()], color='orange', lw=2)
plt.title('Linear Regression: Predicted vs. Actual Price')
plt.xlabel('Actual Price')
plt.ylabel('Predicted Price')
plt.show()

# Residuals vs. predicted price
residuals = Y_test_reg - linear_predictions
plt.figure(figsize=(7, 5))
plt.scatter(linear_predictions, residuals, alpha=0.3, s=10)
plt.axhline(y=0, color='orange', lw=2)
plt.title('Linear Regression: Residuals vs. Predicted Price')
plt.xlabel('Predicted Price')
plt.ylabel('Residual (Actual - Predicted)')
plt.show()

from sklearn.neighbors import KNeighborsRegressor

# Keep only the 6 features with the highest correlation with price (in both the train and test sets)
reg_selected_features = ['engine_cc', 'make_year', 'fuel_type_Electric', 'owner_count', 'mileage_kmpl', 'fuel_type_Petrol']

# Standardize the features using z-scores, fitted on the training set only
reg_scaler = StandardScaler()
scaled_X_train_reg = reg_scaler.fit_transform(X_train_reg[reg_selected_features])
scaled_X_test_reg = reg_scaler.transform(X_test_reg[reg_selected_features])

# 10-fold cross-validation on the training data only, to find the best k
mean_cv_rmse = []
k_list = range(1, 51)
for nn in k_list:
    knn_reg_cv = KNeighborsRegressor(n_neighbors=nn)
    cv_scores = cross_val_score(knn_reg_cv, scaled_X_train_reg, Y_train_reg, cv=10, scoring='neg_root_mean_squared_error')
    mean_cv_rmse.append(-cv_scores.mean())  # Scores are negative RMSE, so flip the sign

# Output results
best_k_reg = k_list[int(np.argmin(mean_cv_rmse))]  # The k with the lowest mean CV RMSE
print('Lowest RMSE is obtained for k =', best_k_reg, 'and equals', round(min(mean_cv_rmse), 2))
plt.plot(k_list, mean_cv_rmse, '-o')
plt.xlabel('Number of neighbors (k)')
plt.ylabel('Mean CV RMSE');

# Retrain KNN with the best k on the whole training set and evaluate it on the test set
knn_regressor = KNeighborsRegressor(n_neighbors=best_k_reg)
knn_regressor.fit(scaled_X_train_reg, Y_train_reg)

knn_predictions = knn_regressor.predict(scaled_X_test_reg)
knn_metrics = regression_metrics(Y_test_reg, knn_predictions)
print('KNN regression:', {name: round(value, 3) for name, value in knn_metrics.items()})

# Test set metrics of all the models
comparison = pd.DataFrame(
    [baseline_metrics, linear_metrics, knn_metrics],
    index=['Baseline (mean)', 'Linear regression', f'KNN regression (k={best_k_reg})']
)
comparison.round(3)