#Generate some data

set.seed(234)
n<-1000
K<-3
omega.vec<-c(0.20,0.5,0.3)
mu.vec<-c(10,50,70)
sig.vec<-c(10,10,10)

X<-sample(c(1:3),size=n,prob=omega.vec,rep=T)
Y<-sig.vec[X]*rnorm(n)+(X==1)*mu.vec[1]+(X==2)*mu.vec[2]+(X==3)*mu.vec[3]

#Starting values
omega.r<-omega.r1<-rep(1/3,3)
mu.vec.r<-mu.vec.r1<-quantile(Y,prob=c(0.25,0.5,0.75))
sig.vec.r<-sig.vec.r1<-rep(sqrt(var(Y)),3)

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

xvec<-seq(-35,120,length=1001)
yvec<-xvec*0
for(k in 1:K){
	yvec<-yvec+omega.vec[k]*dnorm(xvec,mu.vec[k],sig.vec[k])
}

hist(Y,breaks=seq(-35,120,by=5),ylim=range(0,n/5),col="gray")
points(Y,Y*0,pch=19)
lines(xvec,yvec*n*5,col="red",lwd=2)

plot(xvec,yvec,type="l",col="red",lwd=2,ylim=range(0,0.05))
points(Y,Y*0,pch=19)

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)
	}
	print(c(rsteps,Ly))

	if(rsteps %% 10 == 0){
		xvec<-seq(-20,120,length=1001)
		yvec<-xvec*0
		for(i in 1:1001){
			for(k in 1:K){
				yvec[i]<-yvec[i]+omega.r[k]*dnorm(xvec[i],mu.vec.r[k],sig.vec.r[k])
			}
		}
		lines(xvec,yvec)
	}
}

#Final density estimate
xtrue<-seq(-20,120,length=1001)
ytrue<-xtrue*0
for(k in 1:K){
	ytrue<-ytrue+omega.vec[k]*dnorm(xtrue,mu.vec[k],sig.vec[k])
}
plot(xtrue,ytrue,type="l",col="red",lwd=2,ylim=range(0,0.04))
points(Y,Y*0,pch=19)
lines(xvec,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(0,50,70,log(10),log(10),log(10),0,0)
f0<-optim(xstart,fn=mix.like,yv=Y,Kv=3,control=list(fnscale=-1,trace=10,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=3,control=list(fnscale=-1,trace=10,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)

f2<-optim(xstart,fn=mix.like,yv=Y,Kv=3,method="BFGS",control=list(fnscale=-1,trace=10,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)

y.opt1<-y.opt2<-xtrue*0
for(k in 1:K){
	y.opt1<-y.opt1+om.opt1[k]*dnorm(xtrue,mu.opt1[k],sig.opt1[k])
	y.opt2<-y.opt2+om.opt2[k]*dnorm(xtrue,mu.opt2[k],sig.opt2[k])
}
lines(xtrue,y.opt1,col="blue")
lines(xtrue,y.opt2,col="green")