#!/usr/bin/env python3.11

import numpy as np
from sklearn.neighbors import KNeighborsRegressor
#--------------------------------------------------------------------------------------------------
#----The last column is the class
def get_data_from_csv(file_name):

    data = np.genfromtxt(file_name, delimiter=',', names=True, filling_values=0, dtype=None, \
ndmin=1)
    # print(f"Extracting from {file_name} and got {data}")
#----Extract all except the last column name as feature names
    feature_names = list(data.dtype.names[:-1])
#----Extract last column as class name
    class_names = [data.dtype.names[-1]]
    # print(f"Feature names {feature_names} Class name {class_names}")

#----Extract all except the last column as features
    features = np.column_stack([data[name] for name in feature_names])
    # print(f"Features are {features}")
#----Extract last column as classes
    classes = data[data.dtype.names[-1]]
    # print(f"Classes {classes}")

    return features,classes,feature_names,class_names
#--------------------------------------------------------------------------------------------------
def get_data():

#----Give meaningful names to features and classes
    feature_names = ["square footage"]
    class_names = ["price"]
    # print(f"Feature names {feature_names} Class name {class_names}")

#----Square Footage (must be 2D array)
    features_train = np.array([[800], [1200], [1500], [1800], [2200], [2500]])
    # print(f"Features are {features_train}")
#----Price in Dollars
    classes_train = np.array([150000, 230000, 280000, 310000, 400000, 450000])
    # print(f"Classes {classes_train}")

#----Test houses at 2000 sqft
    features_test = np.array([[1250],[2000]])
    # print(f"Features test {features_test}")

    return features_train,classes_train,features_test,feature_names,class_names
#--------------------------------------------------------------------------------------------------
def get_predictions(K,features_train,classes_train,features_test):

#----Initialize the K-NN Regressor
    knn_regressor = KNeighborsRegressor(n_neighbors=K)
#----Train the model
    knn_regressor.fit(features_train, classes_train)
#----Predict 
    predictions = knn_regressor.predict(features_test)

    return predictions
#--------------------------------------------------------------------------------------------------
def print_predictions(predictions,features_test,class_names):

#----Output the result for each test data
    for index, value in enumerate(features_test):
        print(f"The predicted {class_names[0]} for {value} is {predictions[index]:.2f}")
#--------------------------------------------------------------------------------------------------
def main():

#----Get the training and test data
    # features_train,classes_train,features_test,feature_names,class_names = get_data()
    features_train,classes_train,feature_names,class_names = get_data_from_csv("houses_train.csv")
    features_test, *_ = get_data_from_csv("houses_test.csv")

#----We'll look at the 3 nearest neighbors (k=3)
    predictions = get_predictions(3,features_train,classes_train,features_test)

#----Report results
    print_predictions(predictions,features_test,class_names)
#--------------------------------------------------------------------------------------------------
if __name__ == "__main__":
    main()
#--------------------------------------------------------------------------------------------------