import os

from scipy.cluster.vq import whiten
from scipy.cluster.vq import kmeans
import numpy as np

import data_loader as dl
import my_img_process as mip


def cluster(src_dir, k_or_guess_book_dir, tgt_dir, regular_shape):
    src_dir += "/"
    tgt_dir += "/"

    if not os.path.exists(tgt_dir):
        os.mkdir(tgt_dir)

    # load data_set
    data_set = dl.load_patterns(src_dir)[0]
    whitened_data_set = whiten(data_set)
    stds = np.std(data_set, axis=0)
    stds[stds == 0] = 1

    if isinstance(k_or_guess_book_dir, int):
        k_or_whitend_guessed_code_books = k_or_guess_book_dir
    else:
        k_or_guess_book_dir + '/'
        guessed_code_books = dl.load_patterns(k_or_guess_book_dir)[0]
        std_mat = np.tile(stds, (len(guessed_code_books), 1))
        k_or_whitend_guessed_code_books = guessed_code_books / std_mat

    whitened_centroids = kmeans(
        whitened_data_set,
        k_or_guess=k_or_whitend_guessed_code_books,
        check_finite=False
    )[0]

    centroids = np.round(whitened_centroids, decimals=0)
    centroids[centroids > 0] = 1
    centroids[centroids != 1] = 0
    mats = []
    for centroid in centroids:
        mat = centroid.reshape(regular_shape)
        mats.append(mat)

    mip.save_mats_as_imgs(mats, tgt_dir, "centroid_", 1)


def test():
    src_dir = "F:/shenzhen_si_chars/"
    k_or_guess_book_dir = "G:/abcd/"
#     k_or_guess_book_dir = 11
    tgt_dir = "g:/test_kmeans_zz4/"
    regular_shape = (15, 10)
    cluster(src_dir, k_or_guess_book_dir, tgt_dir, regular_shape)

if __name__ == '__main__':
    test()