predict

#!/usr/bin/env python3

import os
import sys
import numpy as np
from sklearn.neighbors import KNeighborsClassifier

#print(sk.__path__)
#print(sk.__version__)

k = int(sys.argv[1])

size_f64 = np.dtype(np.float64).itemsize
size_i32 = np.dtype(np.int32).itemsize

file = sys.argv[2]

with open(file, 'rb') as f:
	N = int.from_bytes(f.read(4), byteorder='little')
	M = int.from_bytes(f.read(4), byteorder='little')
	D = int.from_bytes(f.read(4), byteorder='little')
	_ = int.from_bytes(f.read(4), byteorder='little')
	
	train = np.frombuffer(f.read(N*D*size_f64), dtype=np.float64).reshape((N,D))
	train_y = np.frombuffer(f.read(N*size_i32), dtype=np.uint32)

	test = np.frombuffer(f.read(M*D*size_f64), dtype=np.float64).reshape((M,D))

test_y = np.array(KNeighborsClassifier(k).fit(train, train_y).predict(test), dtype=np.int32)


predict_name = os.path.splitext(file)

predict_name = '{}-k{}{}'.format(predict_name[0], k, predict_name[1])

with open(sys.argv[3], 'wb') as f:
	f.write(np.int32(M))
	f.write(test_y);