# Path: lopq/lopq/utils.py
# def iterate_splits(x, splits):
#     """
#     A helper to iterate subvectors.
# 
#     :param ndarray x:
#         a vector to iterate over
#     :param int splits:
#         the number of subvectors
#     :returns (np.array, int):
#         subvector, split index pairs
#     """
#     split_size = len(x) / splits
#     for split in xrange(splits):
#         start = split * split_size
#         yield x[start:start + split_size], split
# 
# def predict_cluster(x, centroids):
#     """
#     Given a vector of dimension D and a matrix of centroids of dimension VxD,
#     return the id of the closest cluster
# 
#     :params np.array x:
#         the data to assign
#     :params np.array centroids:
#         a matrix of cluster centroids
#     :returns int:
#         cluster assignment
#     """
#     #return ((x - centroids) ** 2).sum(axis=1).argmin(axis=0)
# 
#     cid = ((x - centroids) ** 2).sum(axis=1).argmin(axis=0)
#     if centroids.shape[0] <= 256:
#         return np.uint8(cid)
#     elif centroids.shape[0] <= 65536:
#         return np.uint16(cid)
#     else:
#         return np.uint32(cid)
# 
#     # if centroids.shape[0] <= 256:
#     #     ocid = np.uint8(cid)
#     # elif centroids.shape[0] <= 65536:
#     #     ocid = np.uint16(cid)
#     # else:
#     #     ocid = np.uint32(cid)
#     # return ocid


import numpy as np
import logging
import sys
from sklearn.cluster import KMeans, MiniBatchKMeans
from collections import namedtuple
from .utils import iterate_splits, predict_cluster
        from scipy.io import savemat
        from .utils import concat_new_first
        from scipy.io import loadmat
        from .lopq_model_pb2 import LOPQModelParams
        from itertools import chain
        from .lopq_model_pb2 import LOPQModelParams
        from .utils import concat_new_first
        :returns tuple:
            a tuple of coarse codes
        """
        return tuple([predict_cluster(cx, self.Cs[split]) for cx, split in iterate_splits(x, self.num_coarse_splits)])

    def predict_fine(self, x, coarse_codes=None):
        """
        Compute the fine codes for a datapoint.

        :param ndarray x:
            the point to code
        :param ndarray coarse_codes:
            the coarse codes for the point
            if they are already computed

        :returns tuple:
            a tuple of fine codes
        """
        if coarse_codes is None:
            coarse_codes = self.predict_coarse(x)

        px = self.project(x, coarse_codes)

        fine_codes = []
        for cx, split in iterate_splits(px, self.num_coarse_splits):

            # Get product quantizer parameters for this split
            _, _, _, subC = self.get_split_parameters(split)

            # Compute subquantizer codes

# Path: lopq/lopq/utils.py
# def iterate_splits(x, splits):
#     """
#     A helper to iterate subvectors.
# 
#     :param ndarray x:
#         a vector to iterate over
#     :param int splits:
#         the number of subvectors
#     :returns (np.array, int):
#         subvector, split index pairs
#     """
#     split_size = len(x) / splits
#     for split in xrange(splits):
#         start = split * split_size
#         yield x[start:start + split_size], split
# 
# def predict_cluster(x, centroids):
#     """
#     Given a vector of dimension D and a matrix of centroids of dimension VxD,
#     return the id of the closest cluster
# 
#     :params np.array x:
#         the data to assign
#     :params np.array centroids:
#         a matrix of cluster centroids
#     :returns int:
#         cluster assignment
#     """
#     #return ((x - centroids) ** 2).sum(axis=1).argmin(axis=0)
# 
#     cid = ((x - centroids) ** 2).sum(axis=1).argmin(axis=0)
#     if centroids.shape[0] <= 256:
#         return np.uint8(cid)
#     elif centroids.shape[0] <= 65536:
#         return np.uint16(cid)
#     else:
#         return np.uint32(cid)
# 
#     # if centroids.shape[0] <= 256:
#     #     ocid = np.uint8(cid)
#     # elif centroids.shape[0] <= 65536:
#     #     ocid = np.uint16(cid)
#     # else:
#     #     ocid = np.uint32(cid)
#     # return ocid


# Path: lopq/lopq/model.py
import numpy as np
import logging
import sys
from sklearn.cluster import KMeans, MiniBatchKMeans
from collections import namedtuple
from .utils import iterate_splits, predict_cluster
        from scipy.io import savemat
        from .utils import concat_new_first
        from scipy.io import loadmat
        from .lopq_model_pb2 import LOPQModelParams
        from itertools import chain
        from .lopq_model_pb2 import LOPQModelParams
        from .utils import concat_new_first

        :returns tuple:
            a tuple of coarse codes
        """
        return tuple([predict_cluster(cx, self.Cs[split]) for cx, split in iterate_splits(x, self.num_coarse_splits)])

    def predict_fine(self, x, coarse_codes=None):
        """
        Compute the fine codes for a datapoint.

        :param ndarray x:
            the point to code
        :param ndarray coarse_codes:
            the coarse codes for the point
            if they are already computed

        :returns tuple:
            a tuple of fine codes
        """
        if coarse_codes is None:
            coarse_codes = self.predict_coarse(x)

        px = self.project(x, coarse_codes)

        fine_codes = []
        for cx, split in iterate_splits(px, self.num_coarse_splits):

            # Get product quantizer parameters for this split
            _, _, _, subC = self.get_split_parameters(split)

            # Compute subquantizer codes

            fine_codes += [predict_cluster(fx, subC[sub_split]) for fx, sub_split in iterate_splits(cx, self.num_fine_splits)]