## ----setup, include=FALSE-------------------------------------------
knitr::opts_chunk$set(
  collapse = TRUE,
  comment = "#>",
  fig.width = 7,
  fig.height = 4.5
)
suppressPackageStartupMessages(library(fitPS))

## ----zero-based-support---------------------------------------------
zeroBasedSurveySupport = function(n, type) {
  type = match.arg(type, c("P", "S"))
  if (identical(type, "P")) {
    return(n)
  }
  n - 1L
}

## ----poisson-constructor--------------------------------------------
externalPoissonModel = function() {
  psModel(
    model = "poisson",
    parameterNames = "lambda",
    subclass = "externalPoissonModel",
    supportedEngines = c("numerical", "mcmc"),
    mleStart = c(lambda = 1),
    mleLower = c(lambda = sqrt(.Machine$double.eps))
  )
}

## ----poisson-descriptor---------------------------------------------
poissonModel = externalPoissonModel()
class(poissonModel)
modelParameterNames(poissonModel)
supportedPosteriorEngines(poissonModel)

## ----poisson-observations-------------------------------------------
modelObservationData.externalPoissonModel = function(model, x, ...) {
  zeroBasedSurveySupport(x$data$n, x$type)
}

## ----poisson-probabilities------------------------------------------
modelProbabilities.externalPoissonModel = function(model,
                                                     parameters,
                                                     n,
                                                     type,
                                                     ...) {
  lambda = parameters[["lambda"]]
  support = zeroBasedSurveySupport(n, type)
  values = vapply(
    support,
    function(value) {
      dpois(value, lambda = lambda)
    },
    numeric(length(lambda))
  )

  if (length(lambda) == 1L) {
    values = matrix(values, nrow = 1L)
  }

  colnames(values) = paste0(type, n)
  values
}

## ----poisson-likelihood---------------------------------------------
modelLogLikelihood.externalPoissonModel = function(model,
                                                    parameters,
                                                    data,
                                                    ...) {
  lambda = parameters[["lambda"]]
  observations = modelObservationData(model, data)
  sum(data$data$rn * dpois(observations, lambda = lambda, log = TRUE))
}

## ----poisson-bayes-contract-----------------------------------------
modelLogPrior.externalPoissonModel = function(model, parameters, prior, ...) {
  lambda = parameters[["lambda"]]
  if (!is.list(prior) ||
      !all(c("shape", "rate") %in% names(prior)) ||
      any(!is.finite(c(prior$shape, prior$rate))) ||
      prior$shape <= 0 || prior$rate <= 0) {
    stop("Poisson prior must contain positive finite shape and rate values")
  }
  dgamma(lambda, shape = prior$shape, rate = prior$rate, log = TRUE)
}

modelBayesControl.externalPoissonModel = function(model, x, engine, prior, ...) {
  observations = modelObservationData(model, x)
  start = weighted.mean(observations, x$data$rn)
  if (!is.finite(start) || start <= 0) {
    start = 1
  }
  list(
    start = c(lambda = start),
    lower = c(lambda = 0),
    upper = c(lambda = Inf)
  )
}

modelToUnconstrained.externalPoissonModel = function(model, parameters, ...) {
  c(lambda = log(parameters[["lambda"]]))
}

modelFromUnconstrained.externalPoissonModel = function(model, unconstrained, ...) {
  c(lambda = exp(unconstrained[["lambda"]]))
}

modelLogJacobian.externalPoissonModel = function(model, unconstrained, ...) {
  unname(unconstrained[["lambda"]])
}

## ----register-model-methods-----------------------------------------
registerModelMethods = function(className, methods) {
  fitpsNamespace = asNamespace("fitPS")

  for (methodName in methods) {
    methodFunctionName = paste0(methodName, ".", className)
    methodFunction = get(
      methodFunctionName,
      envir = parent.frame(),
      inherits = TRUE
    )

    registerS3method(
      methodName,
      className,
      methodFunction,
      envir = fitpsNamespace
    )
  }
}

## ----poisson-register-----------------------------------------------
registerModelMethods(
  "externalPoissonModel",
  c(
    "modelObservationData",
    "modelProbabilities",
    "modelLogLikelihood",
    "modelLogPrior",
    "modelBayesControl",
    "modelToUnconstrained",
    "modelFromUnconstrained",
    "modelLogJacobian"
  )
)

## ----poisson-fit----------------------------------------------------
poissonData = makePSData(
  n = c(0, 1, 2, 3),
  count = c(30, 12, 5, 1),
  type = "P"
)

poissonFit = fit(
  poissonData,
  model = externalPoissonModel(),
  nterms = 5
)

poissonFit$lambda

## ----poisson-fit-contract-------------------------------------------
class(poissonFit)
class(poissonFit$modelObject)
fitted(poissonFit)
predict(poissonFit, newdata = 0:4, interval = "none")

## ----poisson-comparison---------------------------------------------
logLik(poissonFit)
deviance(poissonFit)
AIC(poissonFit)
BIC(poissonFit)

## ----poisson-bayes-fit----------------------------------------------
poissonPrior = list(shape = 2, rate = 1)
poissonBayesFit = fit(
  poissonData,
  model = externalPoissonModel(),
  method = "bayes",
  prior = poissonPrior,
  nterms = 5
)

summary(poissonBayesFit)
posteriorProbs(poissonBayesFit, n = 5)

## ----poisson-s-support----------------------------------------------
poissonSData = makePSData(
  n = c(1, 2, 3, 4),
  count = c(30, 12, 5, 1),
  type = "S"
)

poissonSFit = fit(
  poissonSData,
  model = externalPoissonModel(),
  nterms = 5
)

c(P = poissonFit$lambda, S = poissonSFit$lambda)
unname(fitted(poissonFit))
unname(fitted(poissonSFit))

## ----poisson-normal-constructor-------------------------------------
externalPoissonNormalModel = function() {
  psModel(
    model = "poissonNormal",
    parameterNames = c("mu", "sigma"),
    subclass = "externalPoissonNormalModel",
    supportedEngines = c("numerical", "mcmc")
  )
}

## ----poisson-normal-control-----------------------------------------
modelObservationData.externalPoissonNormalModel = function(model, x, ...) {
  zeroBasedSurveySupport(x$data$n, x$type)
}

modelMleControl.externalPoissonNormalModel = function(model, x, ...) {
  observations = modelObservationData(model, x)
  weights = x$data$rn
  meanStart = weighted.mean(observations, weights)
  varianceStart = weighted.mean((observations - meanStart)^2, weights)

  if (!is.finite(meanStart) || meanStart <= 0) {
    meanStart = 1
  }
  if (!is.finite(varianceStart)) {
    varianceStart = meanStart
  }

  extraVariance = max(
    varianceStart - meanStart,
    meanStart^2 * 1e-6
  )
  sigmaSquaredStart = log1p(extraVariance / meanStart^2)
  sigmaStart = sqrt(max(sigmaSquaredStart, 1e-6))
  muStart = log(meanStart) - sigmaSquaredStart / 2

  list(
    start = c(mu = muStart, sigma = sigmaStart),
    lower = c(mu = -20, sigma = sqrt(.Machine$double.eps)),
    upper = c(mu = 20, sigma = 5)
  )
}

## ----poisson-normal-pmf---------------------------------------------
poissonNormalProbability = function(n, mu, sigma) {
  if (!is.finite(mu) || !is.finite(sigma) || sigma <= 0) {
    return(rep(NaN, length(n)))
  }

  vapply(n, function(value) {
    if (!is.finite(value) || value < 0 || value != floor(value)) {
      return(0)
    }

    integrand = function(u) {
      z = mu + sigma * u
      dpois(value, lambda = exp(z)) * dnorm(u)
    }

    result = integrate(
      integrand,
      lower = -Inf,
      upper = Inf,
      rel.tol = 1e-8,
      subdivisions = 200L,
      stop.on.error = FALSE
    )

    if (!identical(result$message, "OK")) {
      return(NaN)
    }

    result$value
  }, numeric(1L))
}

## ----poisson-normal-methods-----------------------------------------
modelProbabilities.externalPoissonNormalModel = function(model,
                                                          parameters,
                                                          n,
                                                          type,
                                                          ...) {
  parameterFrame = as.data.frame(parameters)
  if (!all(c("mu", "sigma") %in% names(parameterFrame))) {
    stop("parameters must contain mu and sigma")
  }
  support = zeroBasedSurveySupport(n, type)
  values = vapply(seq_len(nrow(parameterFrame)), function(row) {
    poissonNormalProbability(
      support,
      mu = parameterFrame$mu[row],
      sigma = parameterFrame$sigma[row]
    )
  }, numeric(length(support)))
  values = t(values)
  colnames(values) = paste0(type, n)
  values
}

modelLogLikelihood.externalPoissonNormalModel = function(model,
                                                           parameters,
                                                           data,
                                                           ...) {
  mu = parameters[["mu"]]
  sigma = parameters[["sigma"]]
  observations = modelObservationData(model, data)
  probabilities = poissonNormalProbability(
    observations,
    mu = mu,
    sigma = sigma
  )

  if (any(!is.finite(probabilities)) || any(probabilities <= 0)) {
    return(-Inf)
  }

  sum(data$data$rn * log(probabilities))
}

## ----poisson-normal-bayes-contract----------------------------------
modelLogPrior.externalPoissonNormalModel = function(model, parameters, prior, ...) {
  mu = parameters[["mu"]]
  sigma = parameters[["sigma"]]
  required = c("muMean", "muSd", "sigmaScale")
  if (!is.list(prior) || !all(required %in% names(prior)) ||
      any(!is.finite(unlist(prior[required], use.names = FALSE))) ||
      prior$muSd <= 0 || prior$sigmaScale <= 0) {
    stop("Poisson-normal prior must contain muMean, positive muSd, and positive sigmaScale")
  }
  if (!is.finite(sigma) || sigma <= 0) {
    return(-Inf)
  }
  dnorm(mu, mean = prior$muMean, sd = prior$muSd, log = TRUE) +
    log(2) + dnorm(sigma, mean = 0, sd = prior$sigmaScale, log = TRUE)
}

modelBayesControl.externalPoissonNormalModel = function(model, x, engine, prior, ...) {
  modelMleControl(model, x)
}

modelToUnconstrained.externalPoissonNormalModel = function(model, parameters, ...) {
  c(mu = parameters[["mu"]], sigma = log(parameters[["sigma"]]))
}

modelFromUnconstrained.externalPoissonNormalModel = function(model, unconstrained, ...) {
  c(mu = unconstrained[["mu"]], sigma = exp(unconstrained[["sigma"]]))
}

modelLogJacobian.externalPoissonNormalModel = function(model, unconstrained, ...) {
  unname(unconstrained[["sigma"]])
}

## ----poisson-normal-register----------------------------------------
registerModelMethods(
  "externalPoissonNormalModel",
  c(
    "modelObservationData",
    "modelMleControl",
    "modelProbabilities",
    "modelLogLikelihood",
    "modelLogPrior",
    "modelBayesControl",
    "modelToUnconstrained",
    "modelFromUnconstrained",
    "modelLogJacobian"
  )
)

## ----poisson-normal-fit---------------------------------------------
poissonNormalData = makePSData(
  n = 0:8,
  count = c(3032, 3240, 2035, 997, 426, 168, 64, 24, 9),
  type = "P"
)

poissonNormalFit = fit(
  poissonNormalData,
  model = externalPoissonNormalModel(),
  nterms = 9
)

c(mu = poissonNormalFit$mu, sigma = poissonNormalFit$sigma)

## ----poisson-normal-generic-apis------------------------------------
fitted(poissonNormalFit)
predict(poissonNormalFit, newdata = 0:8, interval = "none")
logLik(poissonNormalFit)
deviance(poissonNormalFit)
AIC(poissonNormalFit)
BIC(poissonNormalFit)

## ----poisson-normal-s-support---------------------------------------
poissonNormalP = modelProbabilities(
  externalPoissonNormalModel(),
  parameters = list(mu = poissonNormalFit$mu, sigma = poissonNormalFit$sigma),
  n = 0:8,
  type = "P"
)
poissonNormalS = modelProbabilities(
  externalPoissonNormalModel(),
  parameters = list(mu = poissonNormalFit$mu, sigma = poissonNormalFit$sigma),
  n = 1:9,
  type = "S"
)

all.equal(
  as.numeric(poissonNormalP),
  as.numeric(poissonNormalS)
)

## ----poisson-normal-bayes-fit---------------------------------------
poissonNormalBayesData = makePSData(
  n = 0:2,
  count = c(30, 12, 4),
  type = "P"
)
poissonNormalPrior = list(muMean = 0, muSd = 2, sigmaScale = 1)

poissonNormalBayesFit = fit(
  poissonNormalBayesData,
  model = externalPoissonNormalModel(),
  method = "bayes",
  prior = poissonNormalPrior,
  nterms = 4
)

summary(poissonNormalBayesFit)
posteriorProbs(poissonNormalBayesFit, n = 4)

## ----retained-model-------------------------------------------------
class(poissonFit$modelObject)
class(poissonBayesFit$modelObject)
class(poissonNormalFit$modelObject)
class(poissonNormalBayesFit$modelObject)

## ----serialization--------------------------------------------------
path = tempfile(fileext = ".rds")
saveRDS(poissonNormalBayesFit, path)
restoredFit = readRDS(path)
class(restoredFit$modelObject)
summary(restoredFit)