from numpy import *
import time
import numpy
import pylab as p
import mpl3d.mplot3d as p3
from pylab import *

import datamind.ml
import datamind.data

def ring_gaussians_data(big_r=10, nbseeds=5, cov=numpy.diag([1,1,1]), n=100):
    '''
    Return a matrix representing several gaussians dispatced along a ring.

    big_r :     big ring radius.
    nbseeds :   number of gaussians dispatched along the ring.
    cov :       covariance matrix of all gaussians.
    n :         number of generated dots in each gaussian'''
    v = numpy.linspace(0, 2 * numpy.pi, nbseeds, endpoint=False)
    C = big_r * numpy.column_stack((numpy.cos(v), numpy.sin(v), numpy.cos(v+1.09)))
    t = [numpy.random.multivariate_normal(c, cov, n) for c in C]
    return numpy.vstack(t)

k = ring_gaussians_data(7, 15, n = 50)

fig=p.figure()
ax = p3.Axes3D(fig)
ax.plot3D(k[:,0], k[:,1], k[:,2], 'bo')
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')
fig.add_axes(ax)
p.show()

fig2 = p.figure(2)

import fff.clustering as FC
from datamind.ml.classifier import ofunc
from datamind.ml.classifier import optimizers

grid = range(1, 30, 2)
energy = []

class EvalCmeans(ofunc.OFunc):
    def __init__(self, X):
        self._X = X
    def eval(self, params = None, eval_data = None):
        nb_clusters = params['clusters number']
        centers, labels, j = FC.cmeans(self._X, nb_clusters)
        c = centers
        ax.clear()
        ax.scatter3D(k[:,0], k[:,1], k[:,2], c=labels, cmap=cm.prism)
        ax.plot3D(c[:,0], c[:,1], c[:,2], 'wD')
        val = j * nb_clusters
        energy.append(val)
        plot(grid[0:len(energy)], energy, 'r')
        p.show()
        return {'value' : val, 'centers' : centers, 'labels' : labels}


def clusterize(X):
    ofunc = EvalCmeans(X)
    opt = optimizers.Grid({"clusters number" : grid})
    res = opt.optimize(ofunc)
    centers = res['ofunc res']['centers']
    return centers


c = clusterize(k)