#!/home/calum/anaconda3/bin/python

from sklearn import decomposition
import skfuzzy as fuzz
from sklearn import preprocessing
from sklearn.mixture import GMM

from scipy.stats import norm
from scipy import linalg

import numpy as np
import matplotlib.pyplot as plt
import plotting_tools as pt
from matplotlib.patches import Ellipse

def fclustering(data, nclusters, n_pca='mle', hard_cluster='standard', hc_prob = 0.7, error=1e-5, init=None):
    
	# normalising data
    data_scale = preprocessing.scale(data,axis=0)
    
    plt.figure(figsize=(20,10))
    for i in range(len(data[0])):
        pt.plot_hist([row[i] for row in data_scale],200,str(i))
    plt.legend()
    plt.xlim(-10,10)
    plt.show()
    
    pca = decomposition.PCA(n_components=n_pca)
    data_pca = pca.fit_transform(data_scale)
    
    print(data_pca.shape)
    
    plt.figure(figsize=(20,10))
    for i in range(len(data_pca[0])):
        pt.plot_hist([row[i] for row in data_pca],200,str(i))
    plt.legend()
    plt.xlim(-10,10)
    plt.show()
    
    #data_pca = np.array(data)
    
    print('Conducting cluster analysis...')
    cntr, nu ,u0, d, jm, p, fpc = fuzz.cluster.cmeans(data_pca.T,nclusters,2,error=error,maxiter=500,init=init)
    
    print('Number of iterations:',p,'FPC',fpc)
    
    if hard_cluster == 'standard':
        cluster_membership = np.argmax(nu,axis=0)
    elif hard_cluster == 'special':
        cluster_membership = []    
        for row in nu.T:
            if np.max(row) > hc_prob:
                cluster_membership.append(np.argmax(row))
            else:
                cluster_membership.append(-1)
    
    cluster_membership = np.array(cluster_membership)
    
    print('cluster membership',cluster_membership.shape)
    
    print('u shape',nu.shape)
    
    ypts = np.array([row[0] for row in data])
       
    plt.figure(figsize=(20,10))
    pt.plot_hist(ypts,200,'one',normed=False,set_bin_size=True,bin_size=0.1)
    for i in range(nclusters):
        pt.plot_hist(ypts[cluster_membership==i],200,str(i),
		     normed=False,set_bin_size=True,bin_size=0.1)
        print('Length of cluster',str(i),':',len(ypts[cluster_membership==i]))
    plt.legend()
    plt.show()
    
    if len(data[0]) > 1:
        xpts = np.array([row[1] for row in data]) 
            
        plt.figure(figsize=(20,10))
        pt.plot_contour(xpts,ypts,fill=True,bins=70,lines=20,colours='ocean_r')
        for i in range(nclusters):
            pt.plot_contour(xpts[cluster_membership == i],
			    ypts[cluster_membership== i],bins=25,lines=5,colours='k')
            print('')   

        plt.xlim(-24,-16)
        plt.ylim(1,4)  
        
        plt.figure(figsize=(20,10))
        for i in range(nclusters):
            plt.plot(xpts[cluster_membership == i],ypts[cluster_membership== i],'.')
            
        # Mark the center of each fuzzy cluster
        for point in cntr:
            plt.plot(point[0], point[1], 'rs')
            
        plt.xlim(-24,-16)
        plt.ylim(1,4)

        
        plt.show()
        
    return u0  

def gmm_clustering(data,scaled_data,n_clusters,n_iter,params,covariance_type):

	gmm = GMM(n_clusters,n_iter=n_iter,params=params,covariance_type=covariance_type)
	gmm.fit(scaled_data)
	
    # visualise the data
	fig, ax = plt.subplots(figsize=(20,10))

	# scaled data
	scaled_x = np.array([row[0] for row in scaled_data])
	scaled_y = np.array([row[1] for row in scaled_data])

	# raw data
	x = np.array([row[0] for row in data])
	y = np.array([row[1] for row in data])

	pt.plot_contour(scaled_x,scaled_y,lines=15)

	for i in range(gmm.n_components):
		mean = gmm.means_[i][:2]
		if covariance_type == 'diag':
		    covar = [[gmm.covars_[i][0],0],[0,gmm.covars_[i][1]]]
		elif covariance_type == 'full':
		    covar = gmm.covars_[i]
		else:
			print('Invalid covariance type input')
			quit()
		v,w = linalg.eig(covar)
		v = 2.* np.sqrt(2.)*np.sqrt(v)
		u = w[0]/linalg.norm(w[0])
		angle = np.arctan(u[1] / u[0])
		angle = 180. * angle / np.pi  # convert to degrees
		ax.plot(gmm.means_[i][0],gmm.means_[i][1],'rs')
		ax.add_artist(Ellipse(mean,2*v[0],2*v[1],180.+angle,fill=False))
		ax.add_artist(Ellipse(mean,v[0],v[1],180.+angle,fill=False))

	cluster_probs = gmm.predict_proba(scaled_data)
	cluster_membership = np.argmax(cluster_probs.T, axis=0)

	for i in range(gmm.n_components):
		plt.figure(figsize=(20,10))
		pt.plot_contour(x,y,lines=15)
		plt.plot(x[cluster_membership==i],y[cluster_membership==i],'.')

	#plt.xlim(-22,-16)
	#plt.ylim(1,4)

	plt.figure(figsize=(20,10))

	pt.plot_hist(y,200,'l',set_bin_size=True,bin_size=0.1,normed=False)

	for i in range(3):
		pt.plot_hist(y[cluster_membership==i],200,str(i), set_bin_size=True, bin_size=0.1, normed=False)
	plt.legend()   

	cluster_membership = np.argmax(cluster_probs.T[:2], axis=0)
	
	plt.figure(figsize=(20,10))

	for i in range(gmm.n_components):
		plt.plot(x[cluster_membership==i],y[cluster_membership==i],'.')

	#plt.xlim(-22,-16)
	#plt.ylim(1,4)

	plt.figure(figsize=(20,10))

	pt.plot_hist(y,200,'l',set_bin_size=True,bin_size=0.1,normed=False)

	for i in range(2):
		pt.plot_hist(y[cluster_membership==i],200,str(i), set_bin_size=True, bin_size=0.1, normed=False)
	plt.legend()  

	return cluster_probs 


	






