#Galaxy data

set.seed(234)

Y<-MASS::galaxies
n<-length(Y)

#Starting values
omega.r<-omega.r1<-rep(1/K,K)
mu.vec.r<-mu.vec.r1<-quantile(Y,prob=c(1:K)/(K+1))
sig.vec.r<-sig.vec.r1<-rep(sqrt(var(Y)),K)

varomega.r<-matrix(0,nrow=n,ncol=K)

xvec<-seq(5000,35000,by=100)
yvec<-xvec*0
for(k in 1:K){
	yvec<-yvec+omega.r[k]*dnorm(xvec,mu.vec.r[k],sig.vec.r[k])
}


#Run the EM algorithm for a fixed number of R steps.
R<-10000

for(rsteps in 1:R){

	for(i in 1:n){
		for(k in 1:K){
			varomega.r[i,k]<-omega.r[k]*dnorm(Y[i],mu.vec.r[k],sig.vec.r[k])
		}
		varomega.r[i,]<-varomega.r[i,]/sum(varomega.r[i,])
	}

	for(k in 1:K){
		omega.r1[k]<-sum(varomega.r[,k])
	}
	omega.r1<-omega.r1/sum(omega.r1)

	for(k in 1:K){
		mu.vec.r1[k]<-sum(varomega.r[,k]*Y)/sum(varomega.r[,k])
		sig.vec.r1[k]<-sqrt(sum(varomega.r[,k]*(Y-mu.vec.r1[k])^2)/sum(varomega.r[,k]))
	}

	omega.r<-omega.r1
	mu.vec.r<-mu.vec.r1
	sig.vec.r<-sig.vec.r1

	Ly<-0
	for(i in 1:n){
		Lsum<-0
		for(k in 1:K){
			Lsum<-Lsum+omega.r[k]*dnorm(Y[i],mu.vec.r[k],sig.vec.r[k])
		}
		Ly<-Ly+log(Lsum)
	}

	if(rsteps %% 1000 == 0){
		print(c(rsteps,Ly))
		xvec<-seq(5000,35000,by=100)
		yvec<-xvec*0
		for(i in 1:length(xvec)){
			for(k in 1:K){
				yvec[i]<-yvec[i]+omega.r[k]*dnorm(xvec[i],mu.vec.r[k],sig.vec.r[k])
			}
		}
	}
}
yvec.em<-yvec

##################################################################
#Check using optim and direct optimization of 'incomplete' data likelihood.

mix.like<-function(xv,yv,Kv){

	muv<-xv[1:Kv]
	sigv<-exp(xv[(Kv+1):(2*Kv)])
	ov<-c(exp(xv[(2*Kv+1):(3*Kv-1)]),1)
	ov<-ov/sum(ov)
	dvec<-rep(0,length(yv))
	for(k in 1:Kv){
		dvec<-dvec+ov[k]*dnorm(yv,muv[k],sigv[k])
	}
	return(sum(log(dvec)))
}


np<-3*K-1
xstart<-c(quantile(Y,prob=c(0:(K-1))/(K-1)),log(rep(sqrt(var(Y)),K)),rep(0,K-1))

#Using Nelder-Mead
f0<-optim(xstart,fn=mix.like,yv=Y,Kv=K,control=list(fnscale=-1,trace=0,maxit=10000))
muv<-f0$par[1:K]
sigv<-exp(f0$par[(K+1):(2*K)])
ov<-c(exp(f0$par[(2*K+1):(3*K-1)]),1)
ov<-ov/sum(ov)
f1<-optim(f0$par,fn=mix.like,yv=Y,Kv=K,control=list(fnscale=-1,trace=0,maxit=10000))
mu.opt1<-f1$par[1:K]
sig.opt1<-exp(f1$par[(K+1):(2*K)])
om.opt1<-c(exp(f1$par[(2*K+1):(3*K-1)]),1)
om.opt1<-om.opt1/sum(om.opt1)

f1<-optim(f1$par,fn=mix.like,yv=Y,Kv=K,control=list(fnscale=-810,trace=0,maxit=10000))
mu.opt1<-f1$par[1:K]
sig.opt1<-exp(f1$par[(K+1):(2*K)])
om.opt1<-c(exp(f1$par[(2*K+1):(3*K-1)]),1)
om.opt1<-om.opt1/sum(om.opt1)

#Using BFGS Quasi-Newton
f2<-optim(xstart,fn=mix.like,yv=Y,Kv=K,method="BFGS",control=list(fnscale=-1,trace=0,maxit=10000))
mu.opt2<-f2$par[1:K]
sig.opt2<-exp(f2$par[(K+1):(2*K)])
om.opt2<-c(exp(f2$par[(2*K+1):(3*K-1)]),1)
om.opt2<-om.opt2/sum(om.opt2)

#f2 with BFGS does a better job

##############################

#For comparison of EM and optim (with BFGS)
hist(Y,breaks=seq(5000,35000,by=1000),col="gray",ylim=range(0,30),main='Galaxy data with fitted densities')
box()
y.opt1<-y.opt2<-xvec*0
for(k in 1:K){
	y.opt1<-y.opt1+om.opt1[k]*dnorm(xvec,mu.opt1[k],sig.opt1[k])
	y.opt2<-y.opt2+om.opt2[k]*dnorm(xvec,mu.opt2[k],sig.opt2[k])
}
lines(xvec,y.opt1*n*1000,col="blue",lwd=2)
lines(xvec,y.opt2*n*1000,col="green",lwd=2)
lines(xvec,yvec.em*n*1000,col="red",lwd=2)
legend(5000,30,c(paste('Nelder-Mead',round(f1$value,3)),paste('BFGS',round(f2$value,3)),paste('EM',round(Ly,3))),col=c('blue','green','red'),lwd=2)
text(30000,30,substitute(K==Kval,list(Kval=K)))