import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import confusion_matrix

names = ['Married', 'Race', 'Gender', 'Age', 'Country', 'Income']

df = pd.read_csv('census.csv', header=1, names=names)
X = np.array(df.iloc[:, 0:4])
y = np.array(df['Income'])

# split into train and test
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.30)

knn = KNeighborsClassifier(n_neighbors=3)
knn.fit(X_train, y_train)
pred = knn.predict(X_test)

print ('\nModel accuracy score: ', accuracy_score(y_test, pred))
print(confusion_matrix(y_test, pred))
    
