KNN

import numpy as np
import pandas as pd
from sklearn.model_selection import cross_val_score
from sklearn.neighbors import KNeighborsRegressor
from sklearn.metrics import r2_score
from typing import Any
from sklearn.preprocessing import MinMaxScaler
from sklearn.model_selection import KFold
from matplotlib import pyplot as plt
import time
import model_evaluation

#using lab2 as a template for KNN regression
# From lab2 
# This file can be executed to train a KNN regression model and print its evaluation metrics.

# This function helps to visualize cross-validation results with the help of a graph
def visualize_cv_results(
    k_values: list[int],
    cv_scores: list[float], 
    cv_std: list[float],
    optimal_k: int
) -> None:
    """
    Create visualization of cross-validation results across k values.
    
    Plots CV scores with error bars and highlights the optimal k value.
    Adds annotations for overfitting and underfitting regions. 
    
    Args:
        k_values: List of k values tested
        cv_scores: Mean CV scores for each k
        cv_std: Standard deviation of CV scores  
        optimal_k: The k value with best CV performance
    """
    plt.figure(figsize=(10, 6))
    plt.errorbar(k_values,cv_scores) 
    # This helps to specify while K value is the optimal one and how it compares to other values
    plt.axvline(x=optimal_k, color='red', linestyle='--', alpha=0.7,
                label=f'Optimal k={optimal_k}') 
    plt.xlabel('Number of Neighbors (k)')
    plt.ylabel('Cross-validation RMSE')
    plt.title('Cross-validation Performance vs Number of Neighbors')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.show()


def train_knn_model(
    X_train: np.ndarray, 
    y_train: np.ndarray,
    k_values: list[int],
    cv_folds: int | KFold | None = 5 
) -> tuple[int, list[float], dict[str, Any]]:
    """
    Train KNN model with cross-validation to find optimal k.
    
    Args:
        X_train: Training features
        y_train: Training target
        k_values: List of k values to test
        cv_folds: Number of cross-validation folds of a KFold object
        
    Returns:
        Optimal k value, CV scores, and training results.
    """    
    cv_scores = []
    cv_std = []
    # Perfoming cross validation for each k value to find the best one
    for k in k_values:
        knn = KNeighborsRegressor(n_neighbors=k, weights='distance') #Create KNN model with distance weighting 
        scores = cross_val_score(knn,X_train,y_train,cv=cv_folds,scoring='neg_root_mean_squared_error') 
        rmse_scores = -scores 
        cv_scores.append(rmse_scores.mean())
        cv_std.append(rmse_scores.std())
        print(f"{k}\t{rmse_scores.mean():.4f}\t\t{rmse_scores.std():.4f}")
    optimal_k = k_values[np.argmin(cv_scores)] 
    print(f"\nOptimal k: {optimal_k}")
    print(f"Best CV RMSE: {np.min(cv_scores)}") 
    final_model = KNeighborsRegressor(n_neighbors=optimal_k, weights='distance')
    final_model.fit(X_train,y_train) 
    training_results = {
        'final_model': final_model,
        'cv_std': cv_std,
        'cv_scores': cv_scores,
    }
    return optimal_k, cv_scores, training_results

# Main function to execute train the KNN models and evaluate the performance of this model.
# The evalutation metrics are printed using the print_evaluation in model_evaluation.
def main():
    train_set = pd.read_csv("train.csv")
    test_set = pd.read_csv("test.csv")
    features = ['OCC_YEAR','OCC_MONTH','PREMISES','HOOD','LONG','LAT','WEIGHT','NSI']
    target = 'TARGET'
    X_train = train_set[features].values
    y_train = train_set[target].values
    X_test = test_set[features].values
    y_test = test_set[target].values
    scaler = MinMaxScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_test_scaled = scaler.transform(X_test)
    kfold = KFold(n_splits=5, shuffle=True, random_state=42)
    range1=range(1,10)
    range2=range(10,300,20)
    k_values = list(range1) + list(range2)
    start= time.time()
    optimal_k, cv_scores, training_results = train_knn_model(X_train_scaled, y_train, k_values, cv_folds=kfold)
    visualize_cv_results(k_values, cv_scores, training_results['cv_std'], optimal_k)
    end= time.time()
    print(f"Training Time: {end-start:.2f} seconds")
    final_model = training_results['final_model']
    y_pred = final_model.predict(X_test_scaled)
    r2 = r2_score(y_test, y_pred)
    model_evaluation.print_evaluation(y_test, y_pred)
    print("END")

if __name__ == "__main__":
    main()