import pandas as pd
import numpy as np

from sklearn.neighbors import KNeighborsClassifier

# Load training data into a pandas dataframe
df_train = pd.read_csv("train_data.txt", delimiter=',', header=None)

# Specify column names
df_train.columns = ['age', 'workclass', 'education', 'marital-status', 'occupation', 'relationship', 'race', 'sex', 'capital-gain', 'capital-loss', 'hours-per-week', 'native-country']

# List of categorical variables
categorical_variables = ['workclass', 'education', 'marital-status', 'occupation', 'relationship', 'race', 'sex', 'native-country']

# Transform categorical variables into numerical by mapping each distinct value to an integer
for v in categorical_variables:
    df_train[v] = df_train[v].astype('category').cat.codes

# Convert the dataframe into a NumPy array
X_train = df_train.to_numpy()

# Load class labels of training samples
y_train = np.loadtxt('y_train.txt', dtype=int)

# Load test data into a pandas dataframe
df_test = pd.read_csv("test_data.txt", delimiter=',', header=None)

# Specify column names
df_test.columns = ['age', 'workclass', 'education', 'marital-status', 'occupation', 'relationship', 'race', 'sex', 'capital-gain', 'capital-loss', 'hours-per-week', 'native-country']

# Transform categorical variables into numerical by mapping each distinct value to an integer
for v in categorical_variables:
    df_test[v] = df_test[v].astype('category').cat.codes

# Convert the dataframe into a NumPy array
X_test = df_test.to_numpy()

# Initialize k-nn classifier
clf = KNeighborsClassifier(n_neighbors=5)
clf.fit(X_train, y_train)

# Use the classifier to make predictions
y_pred = clf.predict(X_test)

# Write predictions to a file
with open('sample_submission.csv', 'w') as f:
    f.write('ID,Class'+'\n')
    for i in range(y_pred.size):
        f.write(str(i)+','+str(y_pred[i])+'\n')