#                          R function plot.fitPoisNegb
#                          (Version 2.0; Enzmann 2007)
#
# Plot of proportion of observed counts of variable x and expected probabilities
# of poisson and negative binomial distributed counts of x showing counts from 0
# to counts with cumulated proportion (= cumprop; default cumprop=.95) or counts
# (= count) as specified
#
# Note: Parameter "cumprob" can only be specified if "count=0". Parameter
#       "count" overrides "cumprop" if "count" is specified to be > 0.
#
# Options:
#
# normal=T : superimpose normal curve (default: normal=F)
# tables=T : show table of (cumulated) probabilities of counts(default: table=F)
# xlab     : label of x-axis (default: name of count variable)
# ylab     : label of y-axis (default: ylab='proportion & probability')
# main     : main title
#            (default: main='Fit of Poisson and Negative Binomial Distribution')
# legend=F : suppress legend (default: legend=T)
# statistics=F: suppress statistics on x-axis (default: statistics=T)
# sleep    : delay in seconds for plotting lines (default: sleep=0)
#
# Example:
#
# x = rnbinom(500,size=1/10.0,mu=1.65)
# plot.fitPoisNegb(x,normal=T,sleep=5)  # wait 5 seconds per line!

library(MASS)

plot.fitPoisNegb = function(x, cumprop=.95, count=0, normal=F, tables=F,
                            xlab='', ylab='', main='', legend=T, statistics=T,
                            sleep=0)
{
# -------------------------
# Function fill.counts:
fill.counts = function(x)
{
  tabdat = table(x)
  d = as.numeric(dimnames(tabdat)$x)
  counts = NULL
  k = 0
  for (i in 1:(d[length(d)]+1))
  {
    k = k + 1
    if (d[i] > (i-1))
    {
      d=append(d,i-1,i-1)
      counts=c(counts,0)
      k = k - 1
    }
    else
    {
      counts=c(counts,tabdat[[k]])
    }
  }
#  print(cbind(0:d[length(d)],counts))
  counts
}

# -------------------------
# Function pnorm.counts (normal curve probabilities for counts):
pnorm.counts = function(x1,x2,mean=0,sd=1)
{
  pnorm.int = function(x1,x2,mean=0,sd=1)
  {
    pnorm(max(x1,x2),mean=mean,sd=sd)-pnorm(min(x1,x2),mean=mean,sd=sd)
  }
  x1=round(x1)
  x2=round(x2)
  p = NULL
  for (i in x1:x2)
  {
    p = c(p,pnorm.int(i-.5,i+.5,mean=mean,sd=sd))
  }
  p
}

# -------------------------
  if (xlab=='') { xlab=deparse(substitute(x)) }

  x = x[!is.na(x)]
  obs=fill.counts(x)/length(x)
  if (count==0)
  {
    max.c = length(obs[cumprop > cumsum(obs)])
  }
  else
  {
    max.c = count
  }
  obs=obs[1:(max.c+1)]

  pois.exp = dpois(0:max.c,mean(x))
  fit=fitdistr(x,densfun="negative binomial")
  nbin.exp = dnbinom(0:max.c,size=fit$estimate[1],mu=fit$estimate[2])

  if (tables==T)
  {
  cat('Poisson  Probabilities for lambda =',round(mean(x),6),'\n\n')
  print(cbind(k=0:max.c,pprob=round(pois.exp,7),pcum=ppois(0:max.c,mean(x))))
  cat('\nNegative Binomial Probabilities\nfor mean =',round(mean(x),6),
      '& overdispersion =',round(1/fit$estimate[1],6),'\n\n')
  print(cbind(k=0:max.c,nbprob=round(nbin.exp,7),nbcum=pnbinom(0:max.c,
        size=fit$estimate[1],mu=fit$estimate[2])))
  if (normal==T)
    {
    cat('\nDiscrete Normal Curve Probabilities\nfor mean =',round(mean(x),6),
        '& sd =',round(sd(x),6),'\n\n')
    tempX = pnorm.counts(0,max.c,mean=mean(x),sd=sd(x))
    print(cbind(k=0:max.c,nprob=round(tempX,7),ncum=cumsum(tempX)))
    }
  }

  if (main=='') { main='Fit of Poisson and Negative Binomial Distribution' }
  if (ylab=='') { ylab='proportion & probability' }
  if (statistics==T)
  {
    sub=paste('(mean = ',round(mean(x),3),', overdispersion = ',
              round(1/fit$estimate[1],3),')',sep="")
  }
  else { sub='' }

  plot(obs~c(0:max.c),ylim=range(0,obs,pois.exp,nbin.exp),
       xlab=xlab, ylab=ylab, sub=sub, main=main,type='o')
  u=par("usr")
  ux = u[2]-(u[2]-u[1])/35
  uy = u[4]-(u[4]-u[3])/25
  if (normal==T)
  {
    if (legend==T)
    {
    legend(ux,uy,xjust=1,yjust=1,c('observed','normal','poisson',
                                   'negative binomial'),
           lty=c(1,2,2,2),col=c('black','green','blue','red'))
    }
    Sys.sleep(sleep)
    lines(c(0:max.c),pnorm.counts(0,max.c,mean=mean(x),sd=sd(x)),col='green',
            lty=2,type='o')
  }
  else if (legend==T)
  {
    legend(ux,uy,xjust=1,yjust=1,c('observed','poisson','negative binomial'),
           lty=c(1,2,2),col=c('black','blue','red'))
  }
  Sys.sleep(sleep)
  lines(pois.exp~c(0:max.c),type='o',col='blue',lty=2)
  Sys.sleep(sleep)
  lines(nbin.exp~c(0:max.c),type='o',col='red',lty=2)
}