# -*- coding: utf-8 -*-

import numpy as np
from sklearn import cluster

def index(x, y, ncol):
    return x + y * ncol


def spectral_clustering(fvs, nstates):
    nrow, ncol, ncha = fvs.shape
    nsamples = nrow * ncol
    print("SP: Computing affinity")
    affinity = np.zeros((nsamples, nsamples), dtype=np.float32)
    for xi in range(ncol):
        for yi in range(nrow):
            for xj in range(max(0, xi - 1), min(xi + 1, ncol)):
                for yj in range(max(0, yi - 1), min(yi + 1, nrow)):
                    i = index(xi, yi, ncol)
                    j = index(xj, yj, ncol)
                    affinity[i, j] = np.float32((np.dot(fvs[yi, xi], fvs[yj, xj]) + 1.0) / 2.0)

    y = cluster.spectral_clustering(affinity, n_clusters=nstates)

    return y.reshape(nrow, ncol)
