# EE4 Selected Topics From Computer Vision Coursework
# Vasil Zlatanov, Nunzio Pucci

DATA_FILE = 'data.npz'
CLUSTER_CNT = 1337
KMEAN_PART = 33

import numpy as np
import matplotlib.pyplot as plt

from sklearn.cluster import KMeans

data = np.load(DATA_FILE)

train = data['train']

# Train part will contain 15 000 descriptors to generate KMeans
part_idx = np.random.random_integers(train.shape[1])
train_part = np.vstack(train[:][part_idx][300:1300])

kmeans = KMeans(n_clusters=CLUSTER_CNT, random_state=0).fit(train_part)


histogram = np.zeros((train.shape[0], train.shape[1],CLUSTER_CNT))

for i in range(train.shape[0])
    for j in range(train.shape[1])
        histogram[i][j] = np.bincount(kmeans.predict(train[i][j])

print(histogram.shape)