# -*- coding: utf-8 -*-
"""
Created on Mon Oct 16 15:10:27 2017

@author: florent
"""

import numpy as np
import matplotlib.pyplot as plt
import scipy.stats as sps
import scipy.linalg as spl


#Ces deux lignes permettent de regler le nombre de decimales dans les "print" de tableaux np:
float_formatter = lambda x: "%.4f" % x
np.set_printoptions(formatter={'float_kind':float_formatter})

#Parametres fixes par l'utilisateur:
nbpas=int(5e1) # nombre de pas dans EM
rangeK=range(1,5) # valeurs de K testees pour le nombre de classes (dans la selection de modele)

#Creation du tableau de donnees "data" a partir du fichier:
#f=open("densitesOs.txt","r")
#datatxt=f.readlines()
#data=[]
#for line in datatxt:
#    line_of_data=[]
#    line = line.strip() #Returns a copy of the string line with possible leading and trailing whitespace removed
#    for number in line.split(): #(line.split() est la liste des str obtenus en scidant line aux espaces)
#        line_of_data.append(float(number))
#    data.append(np.array(line_of_data))
#Beaucoup, beaucoup plus simple :
data=np.loadtxt("densitesOs.txt")
n=len(data)
d=len(data[0])

#Cette fonction rend les parametres du modele, ainsi que les classes les plus probables des observations 
def EM(K,nbpas,x):  # x is an np.array of observations 
    alpha=np.random.rand(K)
    alpha/=np.sum(alpha)
    C=np.cov(x.T)
    mu_0=np.mean(x,axis=0)
    mu=[]
    Cov=[C]*K
    for k in xrange(K):
        mu.append(np.random.multivariate_normal(mean=mu_0,cov=C,size=1)[0])
    for t in xrange(nbpas):
        Probas=[]
        for xi in x:
            pp=np.array([alpha[k]*sps.multivariate_normal.pdf(xi, mean=mu[k], cov=Cov[k]) for k in xrange(K)])
            Probas.append(pp/np.sum(pp))
        Probas=np.array(Probas)
        alpha=np.mean(Probas,axis=0)
        for k in xrange(K):
            mu[k]=np.average(x,axis=0,weights=Probas[:,k])
            Cov[k]=np.cov(x.T,aweights=Probas[:,k])
    Classes=[]
    for xi in x:
        pp=[alpha[k]*sps.multivariate_normal.pdf(xi, mean=mu[k], cov=Cov[k]) for k in xrange(K)]
        Classes.append(np.argmax(pp))
    return (alpha,mu,Cov,Classes)

#Selection du modele (cad de "K") et estimation des parametres et classes pour le K choisi:
def logLikelyhood(K,alpha,mu,Cov,x):
    s=0
    for xi in x:
        s+=np.log(sum([alpha[k]*sps.multivariate_normal.pdf(xi, mean=mu[k], cov=Cov[k]) for k in xrange(K)]))
    return s   
    
def dim_model(K,d):
    return K*d+K*d*(d+1)*.5+K-1
    
estimations=[]   
model_sel_crit=[] 
Cl=[]
for K in rangeK:
    (alpha,mu,Cov,Classes)=EM(K,nbpas,data)
    model_sel_crit.append(-logLikelyhood(K,alpha,mu,Cov,data)+dim_model(K,d)*.5*np.log(n))
    estimations.append((alpha,mu,Cov))
    Cl.append(Classes)    
l=np.argmin(model_sel_crit)
Koptimal=rangeK[l]
(alpha_em,mu_em,covariance_em)=estimations[l]
Classes=np.array(Cl[l])

#Affichage des donnees et des resultats:

#Fonction utile pour materialiser la covariance et la moyenne d'un echantillon gaussien:
def coord_elliposide(mu=[0,0],C=np.eye(2),u=.5,m=1000):
    r=-2*np.log(u)
    A=spl.sqrtm(C)
    a,b,c=A[0,0],A[0,1],A[1,1]
    t=np.linspace(0,2*np.pi,m)
    co,si=r*np.cos(t),r*np.sin(t)
    x=mu[0]+a*co+b*si
    y=mu[1]+b*co+c*si
    return (x,y)

#Affichage des donnees:
plt.figure()
plt.plot(data[:,0],data[:,1],"r.")
plt.title("Les "+str(n)+" points du data set" )
plt.show()

#Affichage des donnees colorees en fonction de leur classe telle qu'estimee par la fct EM:
plt.figure()
for k in xrange(Koptimal):
    datak=np.array([data[i] for i in np.where(Classes == k)[0]])
    plt.plot(datak[:,0],datak[:,1],".")
    (x,y)=coord_elliposide(mu_em[k],covariance_em[k])
    plt.plot(x,y,"r")
    plt.plot(mu_em[k][0],mu_em[k][1],"r+")        
plt.title("Les "+str(n)+" points du data set clusterises par EM" )
plt.show()

#Affichage de la fonction dont le "K" optimal est l'argmin
plt.figure()
plt.plot(rangeK,model_sel_crit, "r")
plt.title("Model selection")
plt.show()

#Affichage des parametres du modele tels qu'estimes par la fct EM:
print("\nValeur de K :\n")
print(Koptimal)

print("\nEstimation de alpha :\n")
print("%.2f, "*Koptimal % tuple(alpha_em))

print("\nEstimation de mu ")
for k in xrange(Koptimal):
    print("\nClasse "+str(k)+" :")
    print("%.2f, "*d % tuple(mu_em[k]))

print("\nEstimation de la covariance :")
for k in xrange(Koptimal):
    print("\nClasse "+str(k)+" :")
    print(covariance_em[k])