dgp = function(y,mu,alpha){
  (mu/(1+alpha*mu))^y*(1+alpha*y)^(y-1)/factorial(y)*exp(-mu*(1+alpha*y)/(1+alpha*mu))
}

dzigp = function(y,mu,alpha,phi){
  ifelse(y==0,phi+(1-phi)*dgp(0,mu,alpha),(1-phi)*dgp(y,mu,alpha))
}

dzigp(y,3,0,0)
mu = 3
par(mfrow=c(2,4))
for (alpha in c(0,0.2,0.5,1)){
  ys    = 0:100
  prob  = dgp(ys,mu,alpha)
  p0    = round(prob[1],2)
  y.m   = round(sum(ys*prob),2)
  y.v   = round(sum((ys-y.m)^2*prob),2)
  plot(ys,prob,type="h",xlab="Count",ylab="Probability",xlim=c(0,50))
  title(paste("mu=",mu," - alpha=",alpha,"\n Mean=",y.m," - Var=",y.v," - Pr(0)=",p0,sep=""))
}


mu    = 3
alpha = 0.2
for (phi in c(0,0.2,0.5,0.75)){
  ys    = 0:100
  prob  = dzigp(ys,mu,alpha,phi)
  p0    = round(prob[1],2)
  y.m   = round(sum(ys*prob),2)
  y.v   = round(sum((ys-y.m)^2*prob),2)
  plot(ys,prob,type="h",xlab="Count",ylab="Probability",xlim=c(0,50))
  title(paste("mu=",mu," - alpha=",alpha," - phi=",phi,"\n Mean=",y.m," - Var=",y.v," - Pr(0)=",p0,sep=""))
}



# simulating some data
set.seed(2718282)
n     = 200
mu    = 5
alpha = 0.2
phi   = 0.1
true  = c(mu,alpha,phi)
ys    = 0:100
prob  = dzigp(ys,mu,alpha,phi)
p0    = round(dzigp(0,mu,alpha,phi),3)
E     = (1-phi)*mu
V     = E*((1+alpha*mu)^2+phi*mu)
y     = sample(ys,size=n,replace=TRUE,prob=prob)

cbind(c(E,V,p0),c(mean(y),var(y),mean(y==0)))

pdf(file="zigp-simdata.pdf",width=12,height=6)
par(mfrow=c(1,2))
plot(ys,prob,type="h",xlim=c(0,max(y)),lwd=2,ylab="Probability",xlab="Count")
plot(y,type="b",ylab="Count",xlab="Observation")
title(paste("mu=",mu," - alpha=",alpha," - phi=",phi,"\n Mean=",E," - Var=",V," - Pr(0)=",p0,sep=""))
dev.off()

# MCMC in action
mu     = mean(y)
alpha  = 0.01
burnin = 100000
M      = 10000
skip   = 100
niter  = burnin+skip*M
draws  = matrix(0,niter,3)
set.seed(12345)
for (iter in 1:niter){
  phi1   = rnorm(1,phi,0.01)
  if (phi1>0){
  lognum = sum(log(dzigp(y,mu,alpha,phi1)))
  loden  = sum(log(dzigp(y,mu,alpha,phi)))
  logaccept = min(0,lognum-logden)
  if (log(runif(1))<logaccept){
    phi = phi1
  }
  }
  alpha1    = exp(rnorm(1,log(alpha),0.1))
  mu1       = exp(rnorm(1,log(mu),0.1))
  lognum    = sum(log(dgp(y,mu1,alpha1)))-log(1+mu1^2)-log(1+alpha1^2)+log(mu)+log(alpha)
  logden    = sum(log(dgp(y,mu,alpha)))-log(1+mu^2)-log(1+alpha^2)+log(mu1)+log(alpha)
  logaccept = min(0,lognum-logden)
  if (log(runif(1))<logaccept){
    alpha = alpha1
    mu    = mu1
  }
  draws[iter,] = c(mu,alpha,phi)
}
draws = draws[seq(burnin+1,niter,by=skip),]

pairs(draws,labels=names)

pdf(file="zigp-posterior.pdf",width=9,height=8)
names = c(expression(mu),expression(alpha),expression(phi))
par(mfrow=c(3,3),mai=c(0.3,0.6,0.2,0.25),cex.lab=1.5)
for (i in 1:3){
  ts.plot(draws[,i],ylim=range(draws[,i],true[i]),xlab="Iteration",ylab=names[i])
  abline(h=true[i],col=2,lwd=2)
  acf(draws[,i],main="")
  hist(draws[,i],xlim=range(draws[,i],true[i]),prob=TRUE,main="",xlab="")
  abline(v=true[i],col=2,lwd=2)
}
dev.off()