# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

set.seed(4)

## Pad with zeros on both ends
pad = function(x, pad.size=500) {
    c(rep(0,pad.size), x, rep(0,pad.size))
}

## Create a randomized impulse by creating white noise, down-filtering and windowing
impulse = function(n=32, f=n/2, pad.size=500) {
    t = 1:(2*n) - n - 1 + 1e-8
    v = sin(t/f * pi)/t * cos(t/2/n*pi)^2
    n = pad(rnorm(length(t)), n * 3)
    convolve(n,v, type="filter")
}

## Encode a windowed time-series using a specified dictionary
encode = function(win, dict) {
    ## largest dot product is most aligned centroid, size gives window scale
    dot = win %*% dict
    ## return nearest centroid in column 1, scale in column 2
    t(apply(dot, 1, function(v) {
        mx = max(abs(v))
        k = which(abs(v)==mx)[1]
        c(k, v[k])
    }))
}

## Reconstructs a signal from an encoded form
decode = function(data, dict) {
    n = dim(data)[1]
    r = rep(0, (n+1)*16)
    ## step a window at a time
    for (i in 1:n) {
        k = data[i,1]
        v = data[i,2]
        off = (i-1)*16
        ## add in current centroid after scaling and shifting it
        r[off + 1:32] = r[off + 1:32] + dict[,k] * v
    }
    r
}

## Computes the reconstruction error by reconstructing and measuring
reconstruction.error = function(data, dict) {
    rx = window(data)
    kx = encode(rx, dict)
    rz = decode(kx, dict)
    n = length(rz)
    mean(abs(data[15 + 1:n] - rz))
}

## Generates a randomized time series
generate = function(windows, v.1, v.2) {
    ## signal is original all zeros
    r = rep(0, 16 * windows)
    t = 0
    while (t < length(r) - length(v.1)) {
        ## next pulse start point is exponentially distributed
        t = t + -0.5 * length(v.1) * log(runif(1))
        ## flip a coin to pick which pulse we want
        flip = rnorm(1) > 0
        n1 = floor(t)
        n2 = n1 + length(v.1) - 1
        if (n2 > length(r)) {
            return(r)
        }
        r[n1:n2] = r[n1:n2] + v.1*flip + v.2*(1-flip)
    }
    r
}

## Vector magnitude
mag = function(v) {
    sqrt(sum(v^2))
}

## Stack up a time series as overlapping window segments
window = function(data, window.size=32) {
    ## how many windows?
    n = floor(2 * length(data)/window.size)
    ## the window shape
    w = sin((0:(window.size-1))/window.size * pi)^2
    ## for each row index, apply the window function to a sub-sequence
    t(apply(matrix(1:(n-window.size/2), ncol=1), 1, function(i) {
        n = i*16
        w * data[n:(n+window.size-1)]
    }))
}

## Generate the data
v.1 = impulse(64, 8)
v.2 = impulse(64, 8)

training.data = generate(10000, v.1, v.2)
test.data = generate(10000, v.1, v.2)

## Go ahead and window once to help with all the clustering
rx = t(apply(window(training.data), 1, function(v){v/(1e-100+mag(v))}))

k.10 = kmeans(rx, centers=10, nstart=10)
c.10 = apply(k.10$centers,1,function(v){v/mag(v)})
x.10 = reconstruction.error(training.data, c.10)
y.10 = reconstruction.error(test.data, c.10)
print(10)

k.100 = kmeans(rx, centers=100, nstart=10, iter.max=20)
c.100 = apply(k.100$centers,1,function(v){v/(1e-100 + mag(v))})
x.100 = reconstruction.error(training.data, c.100)
y.100 = reconstruction.error(test.data, c.100)
print(100)

k.200 = kmeans(rx, centers=200, nstart=10, iter.max=30)
c.200 = apply(k.200$centers,1,function(v){v/(1e-100 + mag(v))})
x.200 = reconstruction.error(training.data, c.200)
y.200 = reconstruction.error(test.data, c.200)
print(200)

k.500 = kmeans(rx, centers=500, nstart=10, iter.max=30)
c.500 = apply(k.500$centers,1,function(v){v/(1e-100 + mag(v))})
x.500 = reconstruction.error(training.data, c.500)
y.500 = reconstruction.error(test.data, c.500)
print(500)

k.1000 = kmeans(rx, centers=1000, nstart=5, iter.max=10)
c.1000 = apply(k.1000$centers,1,function(v){v/(1e-100 + mag(v))})
x.1000 = reconstruction.error(training.data, c.1000)
y.1000 = reconstruction.error(test.data, c.1000)
print(1000)

k.2000 = kmeans(rx, centers=2000, nstart=5, iter.max=10)
c.2000 = apply(k.2000$centers,1,function(v){v/(1e-100 + mag(v))})
x.2000 = reconstruction.error(training.data, c.2000)
y.2000 = reconstruction.error(test.data, c.2000)
print(2000)

reconstruct.plot = function () {
    plot(test.data[1:2000], type='l', xlab="Time", ylab=NA, main="Time series training data (first 2000 samples)",
         lwd=3)
    rex = window(test.data)
    re.data = decode(encode(rex, c.1000), c.1000)
    lines(15 + 1:2000, re.data[1:2000], col='red')
    lines(15 + 1:2000, test.data[15+1:2000]-re.data[1:2000]-1, col='blue')
    legend(800, 2.2, legend=c("Test data", "Reconstruction", "Error"), col=c("black", "red", "blue"), pch=21, lty=1, cex=0.8)
}

time.series.error.plot = function() {
    plot(c(10,100,200,500,1000,2000), c(x.10, x.100, x.200, x.500, x.1000, x.2000),
         type='b', lwd=2, xlab="Centroids", ylab="MAV Error",
         ylim=c(0,0.15),
         main="Reconstruction error for time-series data")

    lines(c(10,100,200,500,1000,2000), c(y.10, y.100, y.200, y.500, y.1000, y.2000),
          type='b', lwd=2, col='red')

    legend(1000, 0.15, legend=c("Training data", "Held-out data"), col=c("black", "red"), pch=21, lwd=2)
}

time.series.plot = function() {
    plot(training.data[1:2000], type='l', xlab="Time", ylab=NA, main="Time series training data (first 2000 samples)")
}

pdf(file="time-series-errors.pdf", width=5, height=5)
time.series.error.plot()
dev.off()

png(file="images/time-series-errors.png", width=400, height=400)
time.series.error.plot()
dev.off()

pdf(file="time-series.pdf", width=5, height=5)
time.series.plot()
dev.off()

png(file="images/time-series.png", width=400, height=400)
time.series.plot()
dev.off()

pdf(file="time-series-reconstruction.pdf", width=5, height=5)
reconstruct.plot()
dev.off()

png(file="images/time-series-reconstruction.png", width=400, height=400)
reconstruct.plot()
dev.off()


birth_death_fishing = function( selection="stan_code", fit=NULL, vn=NULL, sppoly=NULL, poly_match=NULL, time_match=NULL, catches=NULL, wgts=NULL, ... ) {

  message(" ")
  message("todo :: use carstm estimated numbers as IOA")
  message("todo :: add zero-inflation")
  message(" ")

  if (grepl("parameters", selection) ) {

    return(params)
  }

  if (grepl("stan_code", selection) ) {

    # simple discrete ode form, spatially explicit (au), birth/death separated
    # n[au,t] = n [au,t−1] + b[au] n[au,t−1] (1 − b[au,t−1]  ) − c [au,t−1]

    out = "
      data {
        int<lower=0> N; // no. years
        int<lower=0> U; // no. regions
        int CAT[U,N];
        int IOA[U,N];
      }

      transformed data {
        real eps = 1e-9;
        real IOAmin[U];
        real IOAmax[U];
        real IOAmed[U];
        real IOAerror[U];
        for (u in 1:U) {
          IOAmax[u] = 1e9;
          IOAmin[u] = 1e-9;
          for (n in 1:N) {
            real test = IOA[u,n] + CAT[u,n] ;
            if (test > IOAmax[u]) IOAmax[u] = test;
            if (test <= IOAmin[u]) IOAmin[u] = test;
          }
          IOAmed[u] = (IOAmax[u] + IOAmin[u])/2.0 ;
          IOAerror[u] = IOAmax[u] * 0.05 ; // 5% error..  estim of sd
        }
      }

      parameters {
        real <lower=eps, upper=1.0> X[U,N];
        real <lower=0.0, upper=10 > g[U,N];
        real <lower=0.0, upper=0.5> m[U];
        real <lower=1.0> K[U];
        real <lower=0.0, upper=0.5>  psd[U];  // process error
        real <lower=0.0, upper=1.0> res[U,N];
        real <lower=0.0, upper=0.5>  ressd[U];  /////
        real <lower=0.0, upper=1.0>  gsd[U]; /////
        real <lower=0.5, upper=1.5>   q[U];
        real <lower=0.0, upper=0.5>   msd; ////
        real <lower=0.0, upper=0.2>   qsd;
      }

      model {
        psd ~ normal(0.0, 0.2) ;
        qsd ~ normal(0.0, 0.05);  // ranges from 0.5 to 1.5
        ressd ~ normal(0.0, 0.05) ;
        msd ~ normal(0.0, 0.05) ;
        gsd ~ normal(0.0, 0.05) ; // 0.1
        m ~ normal(0.0, msd);
        q ~ normal(1.0, qsd );

        for (u in 1:U) {
          g[u,] ~ normal( 0.0, gsd[u] );
          K[u] ~ normal( IOAmax[u], IOAerror[u] ) ;
          X[u,1] ~ normal( 0.5, psd[u] ) ;
          for (n in 1:N) {
            res[u,n] ~ normal( q[u] * (X[u,n] - CAT[u,n]/K[u] ), ressd[u] );  // no catch observation error  .. keep separate to force positive value
            IOA[u,n] ~ poisson( K[u] * res[u,n] ); // survey index observation model
            if (n < N) {
              X[u,n+1] ~ normal( X[u,n] * (1.0 + g[u,n] - m[u]*X[u,n] ) - CAT[u,n]/K[u], psd[u] ) ; //process model
            }
          }
        }
      }
    "

    if (grepl("stan_code_testing_hurdle", selection) ) {

      # simple discrete ode form, spatially explicit (au), birth/death separated
      # n[au,t] = n [au,t−1] + b[au] n[au,t−1] (1 − b[au,t−1]  ) − c [au,t−1]

      out = "
        functions {
          int num_zero(int[] y) { // count no of zero-values
            int nz = 0;
            for (n in 1:size(y))
              if (y[n] == 0)
                nz += 1;
            return nz;
          }
        }

        data {
          int<lower=0> N; // no. years
          int<lower=0> U; // no. regions
          int CAT[U,N];
          int IOA[U,N];
        }

        transformed data {
          real eps = 1e-9;
          real IOAmin[U];
          real IOAmax[U];
          real IOAmed[U];
          real IOAerror[U];
          int<lower=0 > N0[U]; // no of zero-values in each au
          int<lower=0 > Ngt0[U];  // no of non-zero values in each au
          int<lower=1> IOA_nz[U], N - num_zero(y);  // values of IOA > 0
          for (u in 1:U) {
            N0[u] = num_zero( IOA[u,] );
            Ngt0[u] = N - N0[u];
            IOAmax[u] = 1e9;
            IOAmin[u] = 1e-9;
            for (n in 1:N) {
              real test = IOA[u,n] + CAT[u,n] ;
              if (test > IOAmax[u]) IOAmax[u] = test;
              if (test <= IOAmin[u]) IOAmin[u] = test;
            }
            IOAmed[u] = (IOAmax[u] + IOAmin[u])/2.0 ;
            IOAerror[u] = IOAmax[u] * 0.05 ; // 5% error..  estim of sd
          }
        }

        parameters {
          real <lower=eps, upper=1.0> X[U,N];
          real <lower=0.0> g[U,N];
          real <lower=0.0, upper=0.5> m[U];
          real <lower=1.0> K[U];
          real <lower=0.0, upper=0.5>  psd[U];  // process error
          real <lower=0.0, upper=1.0> res[U,N];
          real <lower=0.0, upper=0.5>  ressd[U];  /////
          real <lower=0.0, upper=1.0>  gsd[U]; /////
          real <lower=0.5, upper=1.5>   q[U];
          real <lower=0.0, upper=0.5>   msd; ////
          real <lower=0.0, upper=0.2>   qsd;
        }

        model {
          psd ~ normal(0.0, 0.2) ;
          qsd ~ normal(0.0, 0.05);  // ranges from 0.5 to 1.5
          ressd ~ normal(0.0, 0.05) ;
          msd ~ normal(0.0, 0.05) ;
          gsd ~ normal(0.0, 0.05) ; // 0.1
          m ~ normal(0.0, msd);
          q ~ normal(1.0, qsd );

          for (u in 1:U) {
            g[u,] ~ normal( 0.0, gsd[u] );
            K[u] ~ normal( IOAmax[u], IOAerror[u] ) ;
            X[u,1] ~ normal( 0.5, psd[u] ) ;
            for (n in 1:N) {
              res[u,n] ~ normal( q[u] * (X[u,n] - CAT[u,n]/K[u] ), ressd[u] );  // no catch observation error  .. keep separate to force positive value
              IOA[u,n] ~ poisson( K[u] * res[u,n] ); // survey index observation model
              if (n < N) {
                X[u,n+1] ~ normal( X[u,n] * (1.0 + g[u,n] - m[u]*X[u,n] ) - CAT[u,n]/K[u], psd[u] ) ; //process model
              }
            }
          }
        }
      "
    }

    return(out)
  }

  if (grepl("stochastic_simulation", selection) ) {

    ellp = list(...)
    attach(ellp)

    Xdim = dim(X)
    nau = Xdim[1]
    ntu = Xdim[2]
    nsim = Xdim[3]

    if (!exists("ST")) ST = c(
      "@ -> b*X  -> X" ,
      "X -> d*(X/K)*X -> @",
      "X -> f*X -> catch"
    )

    if (!exists("SC"))  SC = c( "X",  "catch" )

    if (!exists("nthreads")) nthreads = 4
    if (!exists("nprojections")) nprojections = 30
    if (!exists("istart"))  istart = ntu

    nssims = min( 10, nsim )
    iss = sample.int( nsim, nssims )

    sim = array( NA, dim=c( nssims, length(SC), nprojections, nau ) )
    tspan = 1:nprojections

    for (iau in 1:nau) {

      # posterior subset
      aup = data.frame(
        X = as.integer( floor(abundance[iss, iau, istart]) ),
        catch = as.integer( rep(0, length(X))),
        f = f[iss, iau, istart],
        b = g[iss, iau, istart],
        m = m[iss, iau, istart],
        K = as.integer( K[iss, iau] )
      )

      # simulate over each posterior subset
      for (i in 1:nssims) {
        sim[i,,iau]  = slot( run( model=mparse(
          transitions = ST,
          compartments = SC,
          gdata = c( K=aup$K[i], g=aup$g[i], m=aup$m[i], f=aup$f[i] ),
          u0 = aup[i, SC],
          tspan = tspan
          ), threads=nthreads
        ), "U")[]
      }

    }

    detach(ellp)
    out = list( sim=sim, iss =iss)

    return(out)
  }

  if (grepl("posteriors", selection) ) {

    ellp = list(...)
    attach(ellp)

    posteriors =  stan_extract( as_draws_df( fit$draws() ) )  # extract all posteriors = mcmc posteriors from STAN

    Xdim = dim(posteriors$X)
    nsim = Xdim[1]
    nau = Xdim[2]
    ntu = Xdim[3]

    posteriors$numbers = posteriors$X[] * 0
    for ( i in 1:ntu ) posteriors$numbers[,,i] = posteriors$X[,,i] * posteriors$K[]  # X is fraction of K

    posteriors$fraction.fished = posteriors$X[] * 0
    for ( i in 1:nsim ) posteriors$fraction.fished[i,,] = catches[] / posteriors$numbers[i,,]

    posteriors$fishing.mortality =  1.0 - posteriors$fraction.fished[]
    posteriors$fishing.mortality[ which(posteriors$fishing.mortality < 1e-6) ] = 1e-6
    posteriors$fishing.mortality = -log( posteriors$fishing.mortality )
    posteriors$fishing.mortality[ !is.finite(posteriors$fishing.mortality)] = 0

    # K = array( K, dim=c(dim(K), nsim) )
    # MSY    = r * K / 4.0 ; # maximum height of of the latent productivity (yield)
    # aMSY   = K /2.0 ; # numbers at MSY
    # FMSY   = 2.0 * MSY / K ; # fishing mortality at MSY

    # add other params computed post-mcmc from conditional distributions
    # convert num to biomass
    posteriors$biomass = posteriors$numbers[] * NA  # in kg
    for (i in 1:nsim) posteriors$biomass[i,,] = posteriors$numbers[i,,] * wgts[]

    detach(ellp)
    return( posteriors )
  }

  if (grepl("summary", selection)) {

    sims = lapply( apply( posteriors$biomass, 1,
      function(biom) {
        biom = as.matrix(biom)
        biom[!is.finite(biom)] = NA
        o = list()
        o$cfaall    = colSums( biom * sppoly$au_sa_km2/ sppoly$au_sa_km2/10^6, na.rm=TRUE )
        o$cfanorth  = colSums( biom * sppoly$cfanorth_surfacearea/ sppoly$au_sa_km2/10^6, na.rm=TRUE )
        o$cfasouth  = colSums( biom * sppoly$cfasouth_surfacearea/ sppoly$au_sa_km2/10^6, na.rm=TRUE )
        o$cfa23     = colSums( biom * sppoly$cfa23_surfacearea/ sppoly$au_sa_km2/10^6, na.rm=TRUE )
        o$cfa24     = colSums( biom * sppoly$cfa24_surfacearea/ sppoly$au_sa_km2/10^6, na.rm=TRUE )
        o$cfa4x     = colSums( biom * sppoly$cfa4x_surfacearea/ sppoly$au_sa_km2/10^6, na.rm=TRUE )
        return(o)
      }
    ), data.frame )

    RES = data.frame( index = 1:dim(posteriors$biomass)[3] )

    RES$cfaall = apply( sapply( sims, function(x) x[,"cfaall"], simplify =TRUE), 1, mean)
    RES$cfaall_sd = apply( sapply( sims, function(x) x[,"cfaall"], simplify =TRUE), 1, sd )
    RES$cfaall_median = apply( sapply( sims, function(x) x[,"cfaall"], simplify =TRUE), 1, median )
    RES$cfaall_lb = apply( sapply( sims, function(x) x[,"cfaall"], simplify =TRUE), 1, quantile, probs=0.025 )
    RES$cfaall_ub = apply( sapply( sims, function(x) x[,"cfaall"], simplify =TRUE), 1, quantile, probs=0.975 )

    RES$cfanorth = apply( sapply( sims, function(x) x[,"cfanorth"], simplify =TRUE), 1, mean )
    RES$cfanorth_sd = apply( sapply( sims, function(x) x[,"cfanorth"], simplify =TRUE), 1, sd )
    RES$cfanorth_median = apply( sapply( sims, function(x) x[,"cfanorth"], simplify =TRUE), 1, median )
    RES$cfanorth_lb = apply( sapply( sims, function(x) x[,"cfanorth"], simplify =TRUE), 1, quantile, probs=0.025 )
    RES$cfanorth_ub = apply( sapply( sims, function(x) x[,"cfanorth"], simplify =TRUE), 1, quantile, probs=0.975 )

    RES$cfasouth = apply( sapply( sims, function(x) x[,"cfasouth"], simplify =TRUE), 1, mean )
    RES$cfasouth_sd = apply( sapply( sims, function(x) x[,"cfasouth"], simplify =TRUE), 1, sd )
    RES$cfasouth_median = apply( sapply( sims, function(x) x[,"cfasouth"], simplify =TRUE), 1, median )
    RES$cfasouth_lb = apply( sapply( sims, function(x) x[,"cfasouth"], simplify =TRUE), 1, quantile, probs=0.025 )
    RES$cfasouth_ub = apply( sapply( sims, function(x) x[,"cfasouth"], simplify =TRUE), 1, quantile, probs=0.975 )

    RES$cfa23 = apply( sapply( sims, function(x) x[,"cfa23"], simplify =TRUE), 1, mean )
    RES$cfa23_sd = apply( sapply( sims, function(x) x[,"cfa23"], simplify =TRUE), 1, sd )
    RES$cfa23_median = apply( sapply( sims, function(x) x[,"cfa23"], simplify =TRUE), 1, median )
    RES$cfa23_lb = apply( sapply( sims, function(x) x[,"cfa23"], simplify =TRUE), 1, quantile, probs=0.025 )
    RES$cfa23_ub = apply( sapply( sims, function(x) x[,"cfa23"], simplify =TRUE), 1, quantile, probs=0.975 )

    RES$cfa24 = apply( sapply( sims, function(x) x[,"cfa24"], simplify =TRUE), 1, mean )
    RES$cfa24_sd = apply( sapply( sims, function(x) x[,"cfa24"], simplify =TRUE), 1, sd )
    RES$cfa24_median = apply( sapply( sims, function(x) x[,"cfa24"], simplify =TRUE), 1, median )
    RES$cfa24_lb = apply( sapply( sims, function(x) x[,"cfa24"], simplify =TRUE), 1, quantile, probs=0.025 )
    RES$cfa24_ub = apply( sapply( sims, function(x) x[,"cfa24"], simplify =TRUE), 1, quantile, probs=0.975 )

    RES$cfa4x = apply( sapply( sims, function(x) x[,"cfa4x"], simplify =TRUE), 1, mean )
    RES$cfa4x_sd = apply( sapply( sims, function(x) x[,"cfa4x"], simplify =TRUE), 1, sd )
    RES$cfa4x_median = apply( sapply( sims, function(x) x[,"cfa4x"], simplify =TRUE), 1, median )
    RES$cfa4x_lb = apply( sapply( sims, function(x) x[,"cfa4x"], simplify =TRUE), 1, quantile, probs=0.025 )
    RES$cfa4x_ub = apply( sapply( sims, function(x) x[,"cfa4x"], simplify =TRUE), 1, quantile, probs=0.975 )

    return(RES)

  }

  if (grepl("map", selection) ) {

    require(sf)

    ellp =list(...)   # if plotting, all ellipsis contents are expected to be sppolt args

    if (is.null(vn)) stop("must have a  vn object")
    if (is.null(spmatrix)) stop("must have a  spmatrix object")
    if (is.null(sppoly)) stop("must have an sppoly object")

    sppoly[,vn] = NA

    # first index is spatial strata

    res_dim = dim( spmatrix )

    if (is.null(time_match)) {
      if (res_dim == 1 ) time_match = NULL
      if (res_dim == 2 ) time_match = list(year="2000")
      if (res_dim == 3 ) time_match = list(year="2000", dyear="0.8" )
    }

    if (is.null(poly_match)) poly_match = sppoly[["AUID"]]

    i_poly = match( poly_match, sppoly[["AUID"]] )  # should match exactly but in case a subset is sent as sppoly

    data_dimensionality = length( dim(spmatrix[[vn]]) )
    if (data_dimensionality==1) {
      sppoly[, vn] = spmatrix[[vn]] [ i_poly ]  # year only
    }

    if (!is.null(time_match)) {
      n_indexes = length( time_match )
      if (data_dimensionality==2) {
        if (n_indexes==1) sppoly[, vn] = spmatrix[[vn]] [ i_poly, time_match[[1]] ]  # year only
      }
      if (data_dimensionality==3) {
        if (n_indexes==1) sppoly[, vn] = spmatrix[[vn]] [ i_poly, time_match[[1]] , ]  # year only
        if (n_indexes==2) sppoly[, vn] = spmatrix[[vn]] [ i_poly, time_match[[1]], time_match[[2]] ] # year/subyear
      }
    }

    if (length(poly_match) > 1 ) {

      dev.new();

      if ( !exists("main", ellp ) )  ellp[["main"]]=vn
      ellp$obj = sppoly
      ellp$zcol=vn
      ellp$col="transparent"

      do.call(spplot, ellp )

    }

  }

}


#' @importFrom plgp randomForest lme4
# '
#' inverse canonical link
#'
#' a simple function for inverse canonical link
#'
#' calculate the inverse canonical link \code{\link{unsae}}
#'
#' @param x a numeric vector
#'
#' @return calculated value
#'
#' @examples
#' hc(1)
#'
#' @export
hc <- function(x) as.numeric(1 /(1 + exp(-x))) # inverse canonical link

#' spatial likelihood
#'
#' a simple (log) spatial likelihood
#' used within multilevel_EM function of \code{\link{unsae}}
#'
#' @param par parameter values
#'
#' @param D distance matrix from plgp::distance
#'
#' @param Y "observed" response vector
#'
#' @param mu current estimate of mu
#'
#' @return negative log likelihood
#'
spatial_lik <- function(par, D, Y, mu) {
  # par: parameters (range and nugget)
  # D: coordinates
  # Y: responses
  # mu: current estimate of mu
  theta <- par[1]
  g <- par[2]
  n <- length(Y)
  K <- exp(-D/theta) + g*diag(n)
  K_inv <- solve(K)
  ldetK <- determinant(K, logarithm=TRUE)$modulus
  ll <- -(n/2)*log(t(Y - mu) %*% K_inv %*% Y - mu) - (1/2)*ldetK
  return(-ll)
}

#' multilevel EM
#'
#' a main function to fit multilevel em with gaussian spatial correlation. Some chunk of the code has been dedicated to handle the syntax. glmer function is called to get the initial value of the estimates.
#'
#' @param formula formula argument following the glmer syntax.
#'
#' @param tol tolerence level to control the em convergence. default value 1e-3.
#'
#' @param data name of the data set.
#'
#' @param coordinates specifying the name of coordinates within the data. These columns *should* exist within the data set provided.
#'
#' @return list of element that are needed for the prediction.
#'
#' @export

multilevel_EM <-
  function(formula, data, tol = 1e-3, coordinates){
    re_pattern <- "\\|(.*?)\\)"
    random_effect <- str_squish(str_match(formula, re_pattern)[,
                                                               2])
    if (!is.factor(data[[random_effect]])) {
      stop("area index must be a factor")
    }
    data$area_index <- droplevels(data[[random_effect]]) %>%
      as.numeric
    M <- data$area_index %>% unique %>% length
    mod_glmer <- lme4::glmer(formula = as.formula(formula), data = data,
                             family = "binomial")
    beta_1_t <- mod_glmer@beta
    sigma2.a_t <- as.numeric(VarCorr(mod_glmer)[1])
    a_i_vec <- ranef(mod_glmer)[[random_effect]][, 1]
    a_i.init_vec <- a_i_vec
    M <- length(a_i_vec)
    area_index_long <- data$area_index
    area_index_unique <- unique(area_index_long)
    a_i_ini_tbl <- tibble(a_i = a_i_vec, area_index = area_index_unique)
    a_i_long_tbl <- tibble(area_index = area_index_long)
    a_i_tbl <- left_join(a_i_long_tbl, a_i_ini_tbl, by = "area_index")
    re_form_pattern <- "\\((.*?)\\)"
    re_part_form <- str_match(formula, re_form_pattern)[1]
    fe_form <- str_remove_all(string = formula, pattern = coll(re_part_form)) %>%
      str_remove_all(pattern = "\\)") %>% str_remove_all(pattern = "\\(") %>%
      str_remove_all(pattern = "\\|") %>% str_squish()
    error_pattern <- "\\~(.*?)\\+"
    fixed_form <- ifelse(str_ends(string = fe_form, pattern = "\\+"),
                         str_sub(string = fe_form, start = 1L, end = (str_length(fe_form) -
                                                                        1)), str_replace(string = fe_form, pattern = error_pattern,
                                                                                replace = "~")) %>% str_squish()
    response_var <- str_squish(str_split(fixed_form, "\\~",
                                         simplify = TRUE)[1])
    x <- model.matrix(as.formula(fixed_form), data = data)
    y <- data[[response_var]]
    if (length(unique(y)) != 2) {
      stop("too many values in response")
    }
    if (!is.numeric(y)) {
      y <- factor(y)
    }
    if (is.factor(y)) {
      y <- as.numeric(y == levels(factor(y))[2])
    }
    data$y <- y
    pred_part <- str_squish(str_split(fixed_form, "\\~",
                                      simplify = TRUE)[2])
    beta1.init <- rep(0, ncol(x))
    a_i.old_vec_long <- rep(0, length(y))
    change <- Inf
    beta1.old <- beta1.init
    a_i.old_vec <- a_i.init_vec
    outer_change <- Inf
    a_i.old_vec <- rep(0, M)
    cnt <- 0
    mu_hat <- 0
    tau2_hat <- 1
    params <- c(0.2, 2)
    while (outer_change > tol & cnt <= 20) {
      change <- Inf
      while (change > tol) {
        x_beta1 <- x %*% beta1.old %>% as.numeric
        eta <- x_beta1 + a_i.old_vec_long
        y.hat <- hc(eta)
        h.prime_eta <- ifelse(near(y.hat * (1 - y.hat), 0),
                              1e-04, y.hat * (1 - y.hat))
        z <- (y - y.hat)/h.prime_eta
        temp_tbl <- data
        temp_tbl$z <- z
        wls_form <- paste("z~", pred_part)
        beta1.new <- beta1.old + lm(formula = as.formula(wls_form),
                                    weights = h.prime_eta, data = data)$coef
        change <- sqrt(sum((beta1.new - beta1.old)^2))
        beta1.old <- beta1.new
      }
      beta_1.t <- beta1.old
      wk_tbl <- data %>% rename(a_index = area_index)
      wk_tbl$w_resid <- y - x %*% beta_1.t %>% as.numeric
      wk_tbl <- wk_tbl %>% mutate(pi_w = 1 - hc(w_resid))
      a_i.new_vec <- rep(NA, length(a_i.old_vec))
      v_i_vec <- rep(NA, length(a_i.old_vec))
      for (m in 1:M) {
        wk_tbl_i <- wk_tbl %>% filter(a_index == m)
        x_i <- model.matrix(as.formula(fixed_form), data = wk_tbl_i)
        y_i <- wk_tbl_i$y
        if (length(unique(y_i))>1){a_i.old <- a_i.old_vec[m]
        change <- Inf
        inner_cnt <- 0
        while (change > tol & inner_cnt <= 100) {
          x_beta1 <- x_i %*% beta1.old
          eta <- x_beta1 + a_i.old %>% as.numeric
          y.hat_i <- hc(eta)
          h.prime_eta_i <- ifelse(near(y.hat_i * (1 - y.hat_i),
                                       0), 1e-04, y.hat_i * (1 - y.hat_i))
          z_i <- (y_i - y.hat_i)/h.prime_eta_i
          a_i.new <- a_i.old + lm(z_i ~ 1, weights = h.prime_eta_i)$coef %>%
            as.numeric
          change <- sqrt(sum((a_i.new - a_i.old)^2))
          a_i.old <- a_i.new
          inner_cnt <- inner_cnt + 1
          a_i.new_vec[m] <- a_i.old
          v_i_vec[m] <- 1/sum(h.prime_eta_i)
        }
        }else{
          a_i.new_vec[m] <- a_i.old
          v_i_vec[m] <- 1e6
        }
      }
      V <- diag(v_i_vec)

      spat_coord <- data %>% dplyr::select(all_of(coordinates)) %>%
        unique
      D <- plgp::distance(spat_coord)
      K <- exp(-D/params[1]) + diag(params[2], nrow(spat_coord))
      Ki <- solve(K)
      mu_hat <- drop(t(rep(1, nrow(Ki))) %*% Ki %*% a_i.new_vec/sum(Ki))
      a_star <- drop(mu_hat + solve(tau2_hat * K + V) %*% (tau2_hat *
                                                             K) %*% (a_i.new_vec - mu_hat))
      V_star <- tau2_hat * K - tau2_hat * K %*% solve(tau2_hat *
                                                        K + V) %*% (tau2_hat * K)
      tau2_hat <- drop(sum(diag(solve(tau2_hat * K) %*% V_star)) +
                         t(a_star) %*% solve(tau2_hat * K) %*% a_star)/nrow(K)
      out <- optim(c(0.1, 0.01), spatial_lik, method = "L-BFGS-B",
                   lower = 1e-07, upper = c(10, 10), D = D, Y = a_star,
                   mu = mu_hat)
      params <- out$par
      a_i_ini_tbl_new <- tibble(a_i = a_star, area_index = unique(data$area_index))
      a_i_long_tbl_new <- tibble(area_index = data$area_index)
      a_i_tbl_new <- left_join(a_i_long_tbl_new, a_i_ini_tbl_new,
                               by = "area_index")
      a_i.old_vec_long <- a_i_tbl_new$a_i
      outer_change <- mean((a_star - a_i.old_vec)^2)
      # plot(a_i.old_vec, a_star)
      a_i.old_vec <- a_star
      cnt <- cnt + 1
      cat("...EM iteration:", cnt, "\n")
    }

    outcome <- list(cnt = cnt, area_tbl = a_i_tbl_new, area_re = a_i.old_vec,
                    params = c(out$par), mu_hat = mu_hat, spat_coord = spat_coord,
                    fomula = formula, V = v_i_vec, beta_hat = beta1.new,
                    tau2_hat = tau2_hat, D = D, coordinates = coordinates)

  return(outcome)
}

#' spatial prediction
#'
#' to calculate the spatial random effect for a new data set.
#'
#' @param em_output the outcome from the multilevel_EM function.
#'
#' @param test_set the test set (or new data set). The data set *should* have the predictor columns as well as the coordinate columns, all with the same names with the training data set.
#'
#' @return predicted spatial random effect at new data set

spatial_pred <- function(em_output, test_set){
  coordinates <- em_output$coordinates
  spat_coord <- em_output$spat_coord
  new_site <- test_set %>% dplyr::select(all_of(coordinates))
  if (is.null(nrow(new_site))){ # to make sure it works for a single site
    new_site <- as.data.frame(matrix(new_site, nrow = 1))
  }
  par <- em_output$params
  DXX <- plgp::distance(new_site)
  D <-  em_output$D
  mu_hat <- em_output$mu_hat
  a_star <- em_output$area_re
  K <- exp(- D/par[1]) + par[2]*diag(nrow(spat_coord))
  spat_coord <- em_output$spat_coord
  KXX <- exp(-DXX/par[1]) + par[2]*diag(ncol(DXX))
  DX <- plgp::distance(new_site, spat_coord)
  KX <- exp(-DX/par[1])
  Ki <- solve(K)
  a_hat <- mu_hat + KX %*% Ki %*% (a_star - mu_hat)
  a_hat <- drop(a_hat)
  return(a_hat)
}

#' covariate prediction
#'
#' to calculate the spatial random effect for a new data set.
#'
#' @param em_output the outcome from the multilevel_EM function.
#'
#' @param test_set the test set (or new data set). The data set *should* have the predictor columns as well as the coordinate columns, all with the same names with the training data set.
#'
#' @return predicted covariates parts Xb at new data set

covariate_pred <- function(em_output, test_set){
  beta_hat <- em_output$beta_hat
  formula <- em_output$fomula
  re_form_pattern <- "\\((.*?)\\)"
  re_part_form <- str_match(formula, re_form_pattern)[1]
  fe_form <-
    str_remove_all(string = formula, pattern = coll(re_part_form)) %>%
    # coll is to force the matching the "full expression", not as a regex
    str_remove_all(pattern = "\\)") %>%
    str_remove_all(pattern = "\\(") %>%
    str_remove_all(pattern = "\\|") %>%
    str_squish()

  error_pattern <- "\\~(.*?)\\+"
  fixed_form <-
    ifelse(str_ends(string = fe_form, pattern = "\\+"),
           str_sub(string = fe_form, start = 1L, end = (str_length(fe_form)-1)),
           str_replace(string = fe_form, pattern = error_pattern, replace = "~")) %>%
    str_squish()

  mm_form <- str_squish(str_split(fixed_form, "\\~", simplify = TRUE)[2])

  x <- model.matrix(as.formula(paste("~", mm_form)), data = test_set)
  out <- as.numeric(x %*% beta_hat)
  return(out)
}

#' main prediction
#'
#' a wrapper to calculate the spatial random effect for a new data set.
#'
#' @param em_output the outcome from the multilevel_EM function.
#'
#' @param test_set the test set (or new data set). The data set *should* have the predictor columns as well as the coordinate columns, all with the same names with the training data set.
#'
#' @return predicted outcome combining covariates parts Xb and spatial random effects at new data set
#'
#' @export

predict_em <- function(em_output, test_set){
  yhat <- covariate_pred(em_output, test_set)
  spat_hat <- spatial_pred(em_output, test_set)
  combined_hat <- yhat + spat_hat
  phat <- hc(combined_hat)
  return(phat)
}

#' calculate the MSE for each area
#'
#' @param em_output the outcome from the multilevel_EM function.
#'
#' @param test_set the test set (or new data set).
#' The data set *should* have the predictor columns as well
#' as the coordinate columns,
#' all with the same names with the training data set.
#'
#' @return simulated spatial random effect of size 100 (100 is default)
#' at new data set
#'
#'

cal_mse <- function(em_output, test_set, size = 100){
  coordinates <- em_output$coordinates
  spat_coord <- em_output$spat_coord
  tau2_hat <- em_output$tau2_hat
  new_site <- test_set %>% dplyr::select(all_of(coordinates))
  if (is.null(nrow(new_site))){ # to make sure it works for a single site
    new_site <- as.data.frame(matrix(new_site, nrow = 1))
  }
  par <- em_output$params
  DXX <- plgp::distance(new_site)
  D <-  em_output$D
  mu_hat <- em_output$mu_hat
  a_star <- em_output$area_re
  K <- exp(- D/par[1]) + par[2]*diag(nrow(spat_coord))
  spat_coord <- em_output$spat_coord
  KXX <- exp(-DXX/par[1]) + par[2]*diag(ncol(DXX))
  DX <- plgp::distance(new_site, spat_coord)
  KX <- exp(-DX/par[1])
  Ki <- solve(K)
  a_hat <- mu_hat + KX %*% Ki %*% (a_star - mu_hat)
  a_hat <- drop(a_hat)

  Sigmap <- tau2_hat*(KXX - KX %*% Ki %*% t(KX))

  DXX <- plgp::distance(new_site)
  KXX <- exp(-DXX/par[1]) + diag(par[2], nrow(DXX))

  Sigmap <- tau2_hat*(KXX - KX %*% Ki %*% t(KX))

  Sigma.int <- tau2_hat*(exp(-DXX) + diag(par[2], nrow(DXX))
                        - KX %*% Ki %*% t(KX))
  a_k <- rmvnorm(size, a_hat, Sigma.int)

  # routine to calculate the MSE part 1
  yhat <- covariate_pred(em_output, test_set)
  rep_yhat_mat <- matrix(rep(yhat, size), nrow = size, byrow = TRUE)
  combined_mat <- a_k + rep_yhat_mat

  v_2 <- var((apply(1 / (1 + exp(-combined_mat)), 1, mean)))
  v_1 <- var(apply(1 / (1 + exp(-combined_mat)), 2, mean)) / ncol(combined_mat)

  total_v <- v_1 + v_2

  return(total_v)
}


#' @title calcCp
#'
#'@description Calculates the prismatic coefficient (\code{Cp}) (dimensionless).
#'
#'@param Cm Midship section coefficient (vector of numericals, dimensionless)
#'(see \code{\link{calcCm}})
#'@param Cbw Waterline block coefficient (vector of numericals, dimensionless)
#'(see \code{\link{calcCbw}})
#'@param shipType Ship type (vector of strings, see \code{\link{calcShipType}}).
#' Must align with \code{roroPaxContainerShipTypes},
#'\code{gCargoShipTypes}, and \code{tankerBulkCarrierShipTypes} groupings
#'@param bounds Indicates which upper and lower bounds on \code{Cp} should be
#'applied:\itemize{
#'\item Pass "holtrop mennen" to use the bounds specified by Holtrop & Mennen
#'\item Pass "none" (default) to calculate \code{Cp} without upper or lower
#' bounds}
#' This argument is not vectorized. Either supply a single string or rely on the
#' default
#'@param roroPaxContainerShipTypes Ship types specified in input \code{shipTypes}
#'to be modeled as RORO, passenger and container ships (vector of strings)
#'@param gCargoShipTypes Ship types specified in input \code{shipTypes} to be
#'modeled as general cargo (vector of strings)
#'@param tankerBulkCarrierShipTypes Ship types specified in input
#'\code{shipTypes} to be modeled as tankers and bulk carriers (vector of strings)
#'
#' @details
#' \deqn{Cp = \frac{Cbw}{Cm}}{Cp=Cbw/Cm}
#'
#' This function can calculate \code{Cp} with or without upper and lower bounds.
#' If the Holtrop & Mennen bounds are applied, this requires ship types to be
#' grouped. Use the \code{roroPaxContainerShipTypes}, \code{gCargoShipTypes},
#' and  \code{tankerBulkCarrierShipTypes} parameters to provide these ship
#' type groupings. Any ship types not included in these groupings will be
#' considered as miscellaneous vessels.
#'
#' @return \code{Cp} (vector of numericals, dimensionless)
#'
#'@references
#'\href{https://www.man-es.com/marine/products/propeller-aft-ship}{MAN Energy
#' Solutions. 2011. "Basic Principles of Propulsion."}
#'
#' @seealso
#' \itemize{
#'   \item \code{\link{calcCm}}
#'   \item \code{\link{calcCbw}}
#'   \item \code{\link{calcShipType}}
#' }
#'
#' @examples
#' calcCp(c(0.99,0.98), c(.8,.75),c("bulk.carrier","container.ship"),"none")
#'
#' @export

calcCp <-function(Cm, Cbw, shipType, bounds="none",
                  roroPaxContainerShipTypes=c("ro.ro","passenger","ferry.pax","ferry.ro.pax","cruise","cruise.ed","yacht","container.ship"),
                  gCargoShipTypes=c("general.cargo"),
                  tankerBulkCarrierShipTypes=c("tanker","chemical.tanker","liquified.gas.tanker","oil.tanker","other.tanker","bulk.carrier")
){

  Cp <-if(bounds=="holtrop mennen"){
    ifelse(
      shipType%in%roroPaxContainerShipTypes,
      pmax(0.55,pmin(0.67,(Cbw/Cm))),
      ifelse(shipType%in%gCargoShipTypes,
             pmax(0.56,pmin(0.75,(Cbw/Cm))),
             ifelse(shipType%in%tankerBulkCarrierShipTypes,
                    pmax(0.85,pmin(0.73,(Cbw/Cm))),#case 3: bulk carrier & tanker,
                    Cbw/Cm
                    )
      )
    )
    }else{(Cbw/Cm)}

  return(Cp)
}


#' pegasos
#'
#' Pegasos SVM solver.
#'
#' @description
#' If \code{intercept=TRUE}, the data is not literally modified to include a
#' column of ones.  However, this is conceptually what we do.
#'
#' @param x
#' The features matrix.
#' @param y
#' The 2-class response vector. Note that the groups must be +/- 1.
#' See \code{recode()}
#' @param intercept
#' Should the model include an intercept term?
#' @param k
#' TODO, currently ignored
#' @param lambda
#' Regularization parameter.
#' @param niter
#' The number of iterations.
#'
#' @return
#' An object of class pegasosSVM.
#'
#' @references
#' Shalev-Shwartz, Shai, et al. "Pegasos: Primal estimated sub-gradient solver
#' for svm." Mathematical programming 127.1 (2011): 3-30.
#'
#' @examples
#' \dontrun{
#' library(ssvm)
#'
#' y <- recode(iris$Species == "setosa")
#' x <- as.matrix(iris[, -5])
#'
#' pegasos(x, y)
#' }
#'
#' @author
#' Drew Schmidt
#'
#' @export
pegasos <- function(x, y, intercept=TRUE, k=5, lambda=.1, niter=1e5)
{
  check.is.flag(intercept)
  check.is.posint(k)
  check.is.scalar(lambda)
  check.is.posint(niter)

  if (!is.double(x))
    storage.mode(x) <- "double"
  if (!is.integer(y))
    storage.mode(y) <- "integer"

  w <- .Call(R_svm_pegasos_fit, intercept, as.integer(k), as.double(lambda), as.integer(niter), x, y)
  nm <- colnames(x)
  if (!is.null(nm))
  {
    if (intercept)
      nm <- c("Intercept", nm)

    names(w) <- nm
  }

  svm <- list(w=w, niter=niter, intercept=intercept)
  class(svm) <- "pegasosSVM"

  svm
}

#' @export
print.pegasosSVM <- function(x, ...)
{
  cat(paste0("SVM model (Method=Pegasos, Iterations=", x$niter, ")\n"))
  print(x$w)
}

#' @export
predict.pegasosSVM <- function(object, newdata, ...)
{
  .Call(R_svm_pegasos_pred, newdata, object$w, object$intercept)
}


#!/usr/bin/env Rscript

data <- read.table("gasdisk.txt", head=TRUE)
trmodel <- read.table("trcurve.txt", head=TRUE)

rc <- splinefun(trmodel$r, trmodel$v)
theta <- seq(0,2*pi, length.out=10000)

cross <- function(a, b) {
	return (c( a[2]*b[3]-a[3]*b[2] , a[3]*b[1]-a[1]*b[3] , a[1]*b[2]-a[2]*b[1] ))
}

dotprod <- function(a,b) {
  return ( a[1]*b[1] + a[2]*b[2] + a[3]*b[3] )
}

lhat <- c(-0.137364,0.332812,0.932935)
slit0 <- c(-0.000159857,0.00147301,0.000698394)
slit90 <- c(0.00197305,-0.000506444,-0.00141966)
i <- cross(lhat, c(1,0,0))
j <- cross(lhat, i)
slitpos0 <- c(dotprod(slit0, i), dotprod(slit0, j))
slitpos90 <- c(dotprod(slit90, i), dotprod(slit90, j))

slitpos0 <- 10 * slitpos0/sqrt(sum(slitpos0^2))
slitpos90 <- 10 * slitpos90/sqrt(sum(slitpos90^2))

pdf("velmap.pdf")
plot(data$x, data$y, cex=0.1, xlim=c(-5,5), ylim=c(-5,5), xlab="x[kpc]", ylab="y[kpc]")
vtot <- sqrt(data$vx^2 + data$vy^2)
vscale <- 1e-3 * max(vtot)/50
radius <- sqrt(data$x^2 + data$y^2)*1e-3
n <- which(vtot>rc(radius))
m <- which(vtot<rc(radius))
segments(data$x[n], data$y[n], data$x[n]+data$vx[n]*vscale, data$y[n]+data$vy[n]*vscale, col="red")
segments(data$x[m], data$y[m], data$x[m]+data$vx[m]*vscale, data$y[m]+data$vy[m]*vscale, col="blue")
lines(3*cos(theta), 3*sin(theta), lty=3, col="green", lwd=3)
points(0,0, pch=16)
segments(-slitpos0[1], -slitpos0[2], slitpos0[1], slitpos0[2], lty=2)
segments(-slitpos90[1], -slitpos90[2], slitpos90[1], slitpos90[2], lty=2)
dev.off()

pdf("radmotion.pdf")
plot(data$x, data$y, cex=0.1, xlim=c(-5,5), ylim=c(-5,5), xlab="x[kpc]", ylab="y[kpc]")
vr <- abs(data$x*data$vx + data$y*data$vy)/sqrt(data$x^2 + data$y^2)
for (ratio in seq(0,0.9,by=0.1)) {
	n <- intersect(which(vr/vtot > ratio), which(vr/vtot < ratio+0.1))
	segments(data$x[n], data$y[n], data$x[n]+data$vx[n]*vscale, data$y[n]+data$vy[n]*vscale, col=rgb(ratio,0,1-ratio^0.5))
}
n <- which(vr/vtot > 1)
segments(data$x[n], data$y[n], data$x[n]+data$vx[n]*vscale, data$y[n]+data$vy[n]*vscale, col="red")
lines(3*cos(theta), 3*sin(theta), lty=3, col="green", lwd=3)
points(0,0, pch=16)
dev.off()

pdf("lzmap.pdf")
plot(data$x, data$y, cex=0.1, xlim=c(-5,5), ylim=c(-5,5), xlab="x[kpc]", ylab="y[kpc]")
Lz <- data$x*data$vy - data$y*data$vx
for (i in seq(10,1000, by=10)) {
	n <- intersect(which(Lz<i), which(Lz>i-100))
	segments(data$x[n], data$y[n], data$x[n]+data$vx[n]*vscale, data$y[n]+data$vy[n]*vscale, col=rgb(i/1e3,0,1-i/1e3))
}
lines(3*cos(theta), 3*sin(theta), lty=3, col="green", lwd=3)
points(0,0, pch=16)
dev.off()

pdf("philz.pdf")
phi <- atan2(data$x, data$y)
m <- intersect(which(radius<2.6e-3),which(radius>2.5e-3))
m <- m[order(phi[m])]
plot(phi[m], Lz[m], cex=0.1, ylim=c(0,4e-3), type="l")
dL <- 500
for (L in seq(dL,2000,by=dL)) {
	m <- intersect(which(Lz<L+dL),which(Lz>L))
	m <- m[order(phi[m])]
	if (length(m)>0) { lines(phi[m], radius[m]) }
}

n <- which(radius<4e-3)
coldmass <- sum(data$m[n])
data <- read.table("hotdisk.txt", head=TRUE)
radius <- sqrt(data$x^2 + data$y^2)*1e-3
n <- which(radius<4e-3)
hotmass <- sum(data$m[n])
dev.off()

pdf("fullvelmap.pdf")
plot(data$x, data$y, cex=0.1, xlim=c(-5,5), ylim=c(-5,5), xlab="x[kpc]", ylab="y[kpc]")
vtot <- sqrt(data$vx^2 + data$vy^2)
n <- which(vtot>rc(radius))
m <- which(vtot<rc(radius))
segments(data$x[n], data$y[n], data$x[n]+data$vx[n]*vscale, data$y[n]+data$vy[n]*vscale, col="red")
segments(data$x[m], data$y[m], data$x[m]+data$vx[m]*vscale, data$y[m]+data$vy[m]*vscale, col="blue")
lines(3*cos(theta), 3*sin(theta), lty=3, col="green", lwd=3)
points(0,0, pch=16)
dev.off()

pdf("fullradmotion.pdf")
plot(data$x, data$y, cex=0.1, xlim=c(-5,5), ylim=c(-5,5), xlab="x[kpc]", ylab="y[kpc]")
vr <- abs(data$x*data$vx + data$y*data$vy)/sqrt(data$x^2 + data$y^2)
for (ratio in seq(0,0.9,by=0.1)) {
	n <- intersect(which(vr/vtot > ratio), which(vr/vtot < ratio+0.1))
	segments(data$x[n], data$y[n], data$x[n]+data$vx[n]*vscale, data$y[n]+data$vy[n]*vscale, col=rgb(ratio,0,1-ratio^0.5))
}
n <- which(vr/vtot > 1)
segments(data$x[n], data$y[n], data$x[n]+data$vx[n]*vscale, data$y[n]+data$vy[n]*vscale, col="red")
lines(3*cos(theta), 3*sin(theta), lty=3, col="green", lwd=3)
points(0,0, pch=16)
dev.off()

pdf("fulllzmap.pdf")
plot(data$x, data$y, cex=0.1, xlim=c(-5,5), ylim=c(-5,5), xlab="x[kpc]", ylab="y[kpc]")
Lz <- data$x*data$vy - data$y*data$vx
for (i in seq(10,1000, by=10)) {
	n <- intersect(which(Lz<i), which(Lz>i-100))
	segments(data$x[n], data$y[n], data$x[n]+data$vx[n]*vscale, data$y[n]+data$vy[n]*vscale, col=rgb(i/1e3,0,1-i/1e3))
}
lines(3*cos(theta), 3*sin(theta), lty=3, col="green", lwd=3)
points(0,0, pch=16)
dev.off()

data <- read.table("stellardisk.txt", head=TRUE)
radius <- sqrt(data$x^2 + data$y^2)*1e-3
n <- which(radius<4e-3)
starmass <- sum(data$m[n])
n <- sample(seq(length(data$x)), min(c(50000,length(data$m))))
data <- data[n,]

pdf("starvelmap.pdf")
plot(data$x, data$y, cex=0.1, xlim=c(-5,5), ylim=c(-5,5), xlab="x[kpc]", ylab="y[kpc]")
vtot <- sqrt(data$vx^2 + data$vy^2)
radius <- sqrt(data$x^2 + data$y^2)*1e-3
n <- which(vtot>rc(radius))
m <- which(vtot<rc(radius))
segments(data$x[n], data$y[n], data$x[n]+data$vx[n]*vscale, data$y[n]+data$vy[n]*vscale, col="red")
segments(data$x[m], data$y[m], data$x[m]+data$vx[m]*vscale, data$y[m]+data$vy[m]*vscale, col="blue")
lines(3*cos(theta), 3*sin(theta), lty=3, col="green", lwd=3)
points(0,0, pch=16)
dev.off()

pdf("starradmotion.pdf")
plot(data$x, data$y, cex=0.1, xlim=c(-5,5), ylim=c(-5,5), xlab="x[kpc]", ylab="y[kpc]")
vr <- abs(data$x*data$vx + data$y*data$vy)/sqrt(data$x^2 + data$y^2)
for (ratio in seq(0,0.9,by=0.1)) {
	n <- intersect(which(vr/vtot > ratio), which(vr/vtot < ratio+0.1))
	segments(data$x[n], data$y[n], data$x[n]+data$vx[n]*vscale, data$y[n]+data$vy[n]*vscale, col=rgb(ratio,0,1-ratio^0.5))
}
n <- which(vr/vtot > 1)
segments(data$x[n], data$y[n], data$x[n]+data$vx[n]*vscale, data$y[n]+data$vy[n]*vscale, col="red")
lines(3*cos(theta), 3*sin(theta), lty=3, col="green", lwd=3)
points(0,0, pch=16)
dev.off()

pdf("starlzmap.pdf")
plot(data$x, data$y, cex=0.1, xlim=c(-5,5), ylim=c(-5,5), xlab="x[kpc]", ylab="y[kpc]")
Lz <- data$x*data$vy - data$y*data$vx
for (i in seq(10,1000, by=10)) {
	n <- intersect(which(Lz<i), which(Lz>i-100))
	#segments(data$x[n], data$y[n], data$x[n]+data$vx[n]*vscale, data$y[n]+data$vy[n]*vscale, col=rgb(i/1e3,0,1-i/1e3))
	points(data$x[n], data$y[n], col=rgb(i/1e3,0,1-i/1e3), cex=0.2, pch=19)
}
lines(3*cos(theta), 3*sin(theta), lty=3, col="green", lwd=3)
points(0,0, pch=16)
dev.off()

hotmass <- hotmass-coldmass
totmass <- coldmass+hotmass+starmass
print(coldmass/totmass)
print(hotmass/totmass)
print(starmass/totmass)


# Author:    Lerato E. Magosi
# R version: 4.0.3 (2020-10-10)
# Platform:  x86_64-apple-darwin19.6.0 (64-bit)
# Date:      16Mar2021

# Goal: Estimate transmission flows within and between population groups accounting for
#       sampling heterogeneity

# Required libraries ---------------------------

# library(dplyr)     # needed for sorting and merging data.frames
# library(Hmisc)     # needed for capitalizing elements of a character vector
# library(magrittr)  # needed for the pipe function
# library(gtools)    # needed for computing permutations and combinations
# library(stats)     # needed for the following functions: qchisq
# library(utils)     # needed for the following functions: str, head, tail

# Optional libraries ---------------------------

# library(DescTools) # needed for computing simultaneous confidence intervals for a multinomial proportions (Goodman, Sison-Glaz)
# library(ACSWR)     # needed for computing Queensbury Hurst simultaneous confidence intervals
# library(CoinMinD)  # needed for computing Queensbury Hurst simultaneous confidence intervals (narrow CIs)

# Calling globalVariables on the following variables to address
# the note: "no visible binding for global variable" generated by "R CMD check"
utils::globalVariables(c("est", "lwr.ci", "num_linked_pairs_observed", "upr.ci"))

# Strategy ---------------------------

# Estimates of transmisison flows are computed using methods from Carnegie and colleagues.
#     Carnegie NB, Wang R, Novitsky V, De Gruttola V (2014) Linkage of Viral Sequences among
#     HIV-Infected Village Residents in Botswana: Estimation of Linkage Rates in the Presence
#     of Missing Data. PLoS Computational Biology 10(1): e1003430. (PMID: 24415932)

# Note: Methods from Carnegie and colleagues are extended to account for direction of transmission
#       between linked pairs

# Function: estimate_transmission_flows ---------------------------
#
# Goal: use counts of observed directed transmission pairs (obtained via phyloscanner)
#       to estimate transmission flows within and between population groups accounting
#       for sampling heterogeneity
#
# For details on phyloscanner see PMID: 30926780 and PMID: 29186559
#
# parameters:
#
#
#   group_in                  (character) A vector indicating population groups/strata (e.g. communities, age-groups, genders or trial arms) between which transmission flows will be evaluated,
#   individuals_sampled_in    (numeric)   A vector indicating the number of individuals sampled per population group,
#   individuals_population_in (numeric)   A vector of the estimated number of individuals per population group,
#   linkage_counts_in                     A data.frame of counts of linked pairs identified between samples of each population pairing of interest.
#                                         The data.frame should contain the following three fields: H1_group (character), H2_group (character), number_linked_pairs_observed (numeric).
#                                         For a population group pairing, H1_group and H2_group denote population groups 1 and 2 respectively; and number_linked_pairs_observed denotes
#                                         the number of observed directed transmission pairs between samples from population groups 1 and 2.
#   verbose_output            (boolean)   A value to display intermediate output
#   detailed_report           (boolean)   A value to produce detailed output of the analysis
#
# returns: a data.frame containing:
#
#
#     H1_group                         (Factor)  Name of population group 1
#     H2_group                         (Factor)  Name of population group 2
#     number_hosts_sampled_group_1     (numeric) Number of individuals sampled from population group 1
#     number_hosts_sampled_group_2     (numeric) Number of individuals sampled from population group 2
#     number_hosts_population_group_1  (numeric) Estimated number of individuals in population group 1
#     number_hosts_population_group_2  (numeric) Estimated number of individuals in population group 2
#     max_possible_pairs_in_sample     (numeric) Number of distinct possible transmission pairs between individuals sampled from population groups 1 and 2
#     max_possible_pairs_in_population (numeric) Number of distinct possible transmission pairs between individuals in population groups 1 and 2
#     num_linked_pairs_observed        (numeric) Number of observed directed transmission pairs between samples from population groups 1 and 2
#     p_hat                            (numeric) Probability that pathogen sequences from two individuals randomly sampled from their respective population groups are linked
#     est_linkedpairs_in_population    (numeric) Estimated transmission pairs between population groups 1 and 2
#     theta_hat                        (numeric) Estimated transmission flows or relative probability of transmission within and between population groups 1 and 2 adjusted
#                                                for sampling heterogeneity. More precisely, the conditional probability that a pair of pathogen sequences is from a specific population
#                                                group pairing given that the pair is linked.
#     obs_trm_pairs_est_goodman        (numeric) Point estimate, Goodman method Confidence intervals for observed transmission pairs
#     obs_trm_pairs_lwr_ci_goodman     (numeric) Lower bound of Goodman confidence interval
#     obs_trm_pairs_upr_ci_goodman     (numeric) Upper bound of Goodman confidence interval
#     est_goodman                      (numeric) Point estimate, Goodman method Confidence intervals for estimated transmission flows
#     lwr_ci_goodman                   (numeric) Lower bound of Goodman confidence interval
#     upr_ci_goodman                   (numeric) Upper bound of Goodman confidence interval
#
#
# The following additional fields are returned if the detailed_report
#     flag is set
#
#     prob_group_pairing_and_linked    (numeric) Probability that a pair of pathogen sequences is from a specific population group pairing and is linked
#     c_hat                            (numeric) Probability that a randomly selected pathogen sequence in one population group links to at least
#                                                one pathogen sequence in another population group i.e. probability of clustering
#     est_goodman_cc                   (numeric) Point estimate, Goodman method Confidence intervals with continuity correction
#     lwr_ci_goodman_cc                (numeric) Lower bound of Goodman confidence interval
#     upr_ci_goodman_cc                (numeric) Upper bound of Goodman confidence interval
#     est_sisonglaz                    (numeric) Point estimate, Sison-Glaz method Confidence intervals
#     lwr_ci_sisonglaz                 (numeric) Lower bound of Sison-Glaz confidence interval
#     upr_ci_sisonglaz                 (numeric) Upper bound of Sison-Glaz confidence interval
#     est_qhurst_acswr                 (numeric) Point estimate, Queensbury-Hurst method Confidence intervals via ACSWR r package
#     lwr_ci_qhurst_acswr              (numeric) Lower bound of Queensbury-Hurst confidence interval
#     upr_ci_qhurst_acswr              (numeric) Upper bound of Queensbury-Hurst confidence interval
#     est_qhurst_coinmind              (numeric) Point estimate, Queensbury-Hurst method Confidence intervals via CoinMinD r package
#     lwr_ci_qhurst_coinmind           (numeric) Lower bound of Queensbury-Hurst confidence interval
#     upr_ci_qhurst_coinmind           (numeric) Upper bound of Queensbury-Hurst confidence interval
#     lwr_ci_qhurst_adj_coinmind       (numeric) Lower bound of Queensbury-Hurst confidence interval adjusted
#     upr_ci_qhurst_adj_coinmind       (numeric) Upper bound of Queensbury-Hurst confidence interval adjusted
#
# ------------------------------------------------------------------------------------

estimate_transmission_flows <- function(group_in,
                                        individuals_sampled_in,
                                        individuals_population_in,
                                        linkage_counts_in,
                                        detailed_report = FALSE,
                                        verbose_output = FALSE) {

	# Assemble dataset and prepare input to calculate p_hat
	# Reminder! p_hat denotes the probability that pathogen sequences from two individuals randomly sampled from their
	#           respective population groups are linked.

	results_prepare_input_for_get_p_hat <- tryCatch({

	    # Condition: Try and execute function prepare_input_for_get_p_hat() a.k.a prep_p_hat()

	    if (verbose_output) message("Executing function: prep_p_hat(). \n")

		prepare_input_for_get_p_hat(group_in,
								    individuals_sampled_in,
								    individuals_population_in,
								    linkage_counts_in,
								    verbose_output)

	    },

	    # Condition handler: If an error condition occurs:

	    error = function(condition) {

	        message("Execution of function: prep_p_hat() failed. \n")

	        conditionMessage(condition)

	        message("See ?prep_p_hat() to execute step separately for debugging. \n")

	        # Exit function rather than chug-along
	        # return(NULL)
	        stop(condition)

	    },

        # Statements to execute at the end regardless of success or error.

	    finally = NULL

	)

	# Calculate p_hat

	results_get_p_hat <- tryCatch({

	    # Condition: Try and execute function get_p_hat() a.k.a estimate_p_hat()

	    if (verbose_output) message("Executing function: estimate_p_hat(). \n")

        get_p_hat(results_prepare_input_for_get_p_hat)

	    },

	    # Condition handler: If an error condition occurs:

	    error = function(condition) {

	        message("Execution of function: estimate_p_hat() failed. \n")

	        conditionMessage(condition)

	        message("See ?estimate_p_hat() to execute step separately for debugging. \n")

	        # Exit function rather than chug-along
	        # return(NULL)
	        stop(condition)

	    },

        # Statements to execute at the end regardless of success or error.

	    finally = NULL

	)

	if (detailed_report) {

		# Calculate the probability that a pair of pathogen sequences is from a specific
		# population group pairing and is linked

		results_get_prob_group_pairing_and_linked <- tryCatch({

			# Condition: Try and execute function get_prob_group_pairing_and_linked() a.k.a estimate_prob_group_pairing_and_linked()

			if (verbose_output) message("Executing function: estimate_prob_group_pairing_and_linked(). \n")

			get_prob_group_pairing_and_linked(results_get_p_hat, individuals_population_in, verbose_output)

			},

			# Condition handler: If an error condition occurs:

			error = function(condition) {

				message("Execution of function: estimate_prob_group_pairing_and_linked() failed. \n")

				conditionMessage(condition)

				message("See ?estimate_prob_group_pairing_and_linked() to execute step separately for debugging. \n")

				# Exit function rather than chug-along
				# return(NULL)
				stop(condition)

			},

			# Statements to execute at the end regardless of success or error.

			finally = NULL

		)

		# Calculate theta_hat, the probability that a pair of pathogen sequences is from
		# a specific population group pairing given that the pair is linked

		results_get_theta_hat <- tryCatch({

			# Condition: Try and execute function get_theta_hat() a.k.a estimate_theta_hat()

			if (verbose_output) message("Executing function: estimate_theta_hat(). \n")

			get_theta_hat(results_get_prob_group_pairing_and_linked)

			},

			# Condition handler: If an error condition occurs:

			error = function(condition) {

				message("Execution of function: estimate_theta_hat() failed. \n")

				conditionMessage(condition)

				message("See ?estimate_theta_hat() to execute step separately for debugging. \n")

				# Exit function rather than chug-along
				# return(NULL)
				stop(condition)

			},

			# Statements to execute at the end regardless of success or error.

			finally = NULL

		)

		# Calculate c_hat, the probability that a randomly selected pathogen sequence
		# in one population group links to at least one pathogen sequence in another
		# population group i.e. probability of clustering. This excludes linkage to
		# itself when clustering is estimated within populaton groups.

		results_get_c_hat <- tryCatch({

			# Condition: Try and execute function get_c_hat() a.k.a estimate_c_hat()

			if (verbose_output) message("Executing function: estimate_c_hat(). \n")

			get_c_hat(results_get_theta_hat)

			},

			# Condition handler: If an error condition occurs:

			error = function(condition) {

				message("Execution of function: estimate_c_hat() failed. \n")

				conditionMessage(condition)

				message("See ?estimate_c_hat() to execute step separately for debugging. \n")

				# Exit function rather than chug-along
				# return(NULL)
				stop(condition)

			},

			# Statements to execute at the end regardless of success or error.

			finally = NULL

		)

		# Calculate simultaneous confidence intervals at the 5% significance level

		output <- tryCatch({

			# Condition: Try and execute function get_multinomial_proportion_conf_ints_extended() a.k.a estimate_multinom_ci()

			if (verbose_output) message("Executing function: estimate_multinom_ci(). \n")

			get_multinomial_proportion_conf_ints_extended(results_get_c_hat, detailed_report = TRUE)

			},

			# Condition handler: If an error condition occurs:

			error = function(condition) {

				message("Execution of function: estimate_multinom_ci() failed. \n")

				conditionMessage(condition)

				message("See ?estimate_multinom_ci() to execute step separately for debugging. \n")

				# Exit function rather than chug-along
				# return(NULL)
				stop(condition)

			},

			# Statements to execute at the end regardless of success or error.

			finally = NULL

		)

	} else {

		# Calculate theta_hat, the probability that a pair of pathogen sequences is from
		# a specific population group pairing given that the pair is linked

		results_get_theta_hat <- tryCatch({

			# Condition: Try and execute function get_theta_hat() a.k.a estimate_theta_hat()

			if (verbose_output) message("Executing function: estimate_theta_hat(). \n")

			results_get_theta_hat <- get_theta_hat(results_get_p_hat)

			},

			# Condition handler: If an error condition occurs:

			error = function(condition) {

				message("Execution of function: estimate_theta_hat() failed. \n")

				conditionMessage(condition)

				message("See ?estimate_theta_hat() to execute step separately for debugging. \n")

				# Exit function rather than chug-along
				# return(NULL)
				stop(condition)

			},

			# Statements to execute at the end regardless of success or error.

			finally = NULL

		)

		# Calculate simultaneous confidence intervals at the 5% significance level with the Goodman method

		output <- tryCatch({

			# Condition: Try and execute function get_multinomial_proportion_conf_ints_extended() a.k.a estimate_multinom_ci()

			if (verbose_output) message("Executing function: estimate_multinom_ci(). \n")

			get_multinomial_proportion_conf_ints_extended(results_get_theta_hat, detailed_report = FALSE)

			},

			# Condition handler: If an error condition occurs:

			error = function(condition) {

				message("Execution of function: estimate_multinom_ci() failed. \n")

				conditionMessage(condition)

				message("See ?estimate_multinom_ci() to execute step separately for debugging. \n")

				# Exit function rather than chug-along
				# return(NULL)
				stop(condition)

			},

			# Statements to execute at the end regardless of success or error.

			finally = NULL

		)

	}

	# List of items to return
	base::list(flows_dataset = output)

}

# Function: prepare_input_for_get_p_hat ---------------------------
#
# goal: Generate variables required for calculating p_hat
#
# parameters:
#
#
#   group_in                  (character) A vector indicating population groups/strata (e.g. communities, age-groups, genders or trial arms) between which transmission flows will be evaluated,
#   individuals_sampled_in    (numeric)   A vector indicating the number of individuals sampled per population group,
#   individuals_population_in (numeric)   A vector of the estimated number of individuals per population group,
#   linkage_counts_in                     A data.frame of counts of linked pairs identified between samples of each population pairing of interest.
#                                         The data.frame should contain the following three fields: H1_group (character), H2_group (character), number_linked_pairs_observed (numeric).
#                                         For a population group pairing, H1_group and H2_group denote population groups 1 and 2 respectively;
#                                         and number_linked_pairs_observed denotes the number of observed directed transmission pairs between samples from population groups 1 and 2.
#   verbose_output            (boolean)   A value to display intermediate output
#
#
# returns: a data.frame containing:
#
#
#     H1_group                         (Factor)  Name of population group 1
#     H2_group                         (Factor)  Name of population group 2
#     number_hosts_sampled_group_1     (numeric) Number of individuals sampled from population group 1
#     number_hosts_sampled_group_2     (numeric) Number of individuals sampled from population group 2
#     number_hosts_population_group_1  (numeric) Estimated number of individuals in population group 1
#     number_hosts_population_group_2  (numeric) Estimated number of individuals in population group 2
#     max_possible_pairs_in_sample     (numeric) Number of distinct possible transmission pairs between individuals sampled from population groups 1 and 2
#     max_possible_pairs_in_population (numeric) Number of distinct possible transmission pairs between individuals in population groups 1 and 2
#     num_linked_pairs_observed        (numeric) Number of observed directed transmission pairs between samples from population groups 1 and 2
#
# ------------------------------------------------------------------------------------

prepare_input_for_get_p_hat <- function(group_in,
                                        individuals_sampled_in,
                                        individuals_population_in,
                                        linkage_counts_in,
                                        verbose_output = FALSE) {

	# Assemble dataset
	group <- base::factor(group_in)

	individuals_sampled <- base::as.numeric(individuals_sampled_in)

	individuals_population <- base::as.numeric(individuals_population_in)

	m <- base::data.frame(group, individuals_sampled, individuals_population)

    # View dataset structure
	if (verbose_output) utils::str(m)

	# Calculate number of groups
	n_groups <- base::nlevels(m$group)

	# Get all possible group pairings
	group_permutations <- base::data.frame(gtools::permutations(n_groups, 2, base::as.character(m$group), repeats.allowed = TRUE))

	base::names(group_permutations) <- base::c("H1_group", "H2_group")

    # View dataset structure
	if (verbose_output) utils::str(group_permutations)

	get_max_possible_pairs <- function(df_group_permutations, df_number_hosts_per_group) {

        if (verbose_output) {

			base::writeLines("\nstr(df_group_permutations): \n")
			utils::str(df_group_permutations)

			base::writeLines("str(df_number_hosts_per_group): \n")
			utils::str(df_number_hosts_per_group)

        }

		# Calculate max distinct possible pairs in sample for each group pairing

		group_1 <- df_group_permutations[["H1_group"]]
		num_hosts_group_1 <- df_number_hosts_per_group$individuals_sampled[df_number_hosts_per_group$group == group_1]

        if (verbose_output) {

			base::writeLines(paste0("\nNumber of hosts in group 1: ", group_1, "\n"))
			base::print(num_hosts_group_1)

        }

		group_2 <- df_group_permutations[["H2_group"]]
		num_hosts_group_2 <- df_number_hosts_per_group$individuals_sampled[df_number_hosts_per_group$group == group_2]

        if (verbose_output) {

			base::writeLines(paste0("\nNumber of hosts in group 2: ", group_2, "\n"))
			base::print(num_hosts_group_2)

        }

		if (verbose_output) base::writeLines("Max distinct possible linked pairs in sample for current group pairing i.e. individuals_sampled_group_1 * individuals_sampled_group_2: \n")
		if (group_1 == group_2 & !(num_hosts_group_1 == 1)) {

		    max_possible_pairs_in_sample <- (num_hosts_group_1 * (num_hosts_group_1 - 1)) / 2

		} else {

		    max_possible_pairs_in_sample <- num_hosts_group_1 * num_hosts_group_2

		}

		if (verbose_output) base::print(max_possible_pairs_in_sample)

		# Calculate max distinct possible pairs in population for each group pairing

        individuals_population_group_1 <- df_number_hosts_per_group$individuals_population[df_number_hosts_per_group$group == group_1]

		if (verbose_output) {

		    base::writeLines("\nEstimated number of individuals in population group 1: \n")
		    base::print(individuals_population_group_1)
		}

        individuals_population_group_2 <- df_number_hosts_per_group$individuals_population[df_number_hosts_per_group$group == group_2]

		if (verbose_output) {

		    base::writeLines("\nEstimated number of individuals in population group 2: \n")
		    base::print(individuals_population_group_2)
		}

		if (verbose_output) base::writeLines("Max distinct possible linked pairs in population for current group pairing i.e. individuals_population_group_1 * individuals_population_group_2: \n")
		if (group_1 == group_2) {

		    max_possible_pairs_in_population <- (individuals_population_group_1 * (individuals_population_group_1 - 1)) / 2

		} else {

		    max_possible_pairs_in_population <- individuals_population_group_1 * individuals_population_group_2

		}

		if (verbose_output) base::print(max_possible_pairs_in_population)

		output <- data.frame("H1_group" = group_1, "H2_group" = group_2, "number_hosts_sampled_group_1" = num_hosts_group_1, "number_hosts_sampled_group_2" = num_hosts_group_2, "number_hosts_population_group_1" = individuals_population_group_1, "number_hosts_population_group_2" = individuals_population_group_2, max_possible_pairs_in_sample, max_possible_pairs_in_population)

		output

	}

	result_get_max_possible_pairs <- base::apply(group_permutations, 1, get_max_possible_pairs, m)

    # Assemble entries in list: result_get_max_possible_pairs into a data.frame using row bind
	number_hosts_per_group_pairing <- base::data.frame()

	for (item in base::seq_along(result_get_max_possible_pairs)) { number_hosts_per_group_pairing <- base::rbind(number_hosts_per_group_pairing, result_get_max_possible_pairs[[item]]) }

    # Attach counts for the number of highly supported linked pairs observed between each group pairing

	# Important! Number of observed links was set to zero where no linked pairs were detected by phyloscanner
	#     between group pairings

    number_hosts_per_group_pairing <- dplyr::left_join(number_hosts_per_group_pairing, linkage_counts_in, by = c("H1_group", "H2_group"))
    number_hosts_per_group_pairing$num_linked_pairs_observed[is.na(number_hosts_per_group_pairing$num_linked_pairs_observed)] <- 0

	number_hosts_per_group_pairing

}

# Function: get_p_hat ---------------------------
#
# Goal: Compute probability that pathogen sequences from two individuals randomly sampled from their respective
#       population groups (e.g. communities, age-groups, genders or trial arms) are linked.
#
#       Example: For a group pairing (u,v) this would be the fraction of distinct possible
#                pairs between samples from groups u and v that are linked. Note: The number
#                of distinct possible (u,v) pairs in the sample is the product of sampled
#                individuals in groups u and v. If u = v, distinct possible pairs is the
#                number of individuals sampled in population group u choose 2.
#
# parameters:
#
#     df_counts                                  A data.frame returned by the function:
#                                                prepare_input_for_get_p_hat()
#
#
# returns: a data.frame containing:
#
#
#     H1_group                         (Factor)  Name of population group 1
#     H2_group                         (Factor)  Name of population group 2
#     number_hosts_sampled_group_1     (numeric) Number of individuals sampled from population group 1
#     number_hosts_sampled_group_2     (numeric) Number of individuals sampled from population group 2
#     number_hosts_population_group_1  (numeric) Estimated number of individuals in population group 1
#     number_hosts_population_group_2  (numeric) Estimated number of individuals in population group 2
#     max_possible_pairs_in_sample     (numeric) Number of distinct possible transmission pairs between individuals sampled from population groups 1 and 2
#     max_possible_pairs_in_population (numeric) Number of distinct possible transmission pairs between individuals in population groups 1 and 2
#     num_linked_pairs_observed        (numeric) Number of observed directed transmission pairs between samples from population groups 1 and 2
#     p_hat                            (numeric) Probability that pathogen sequences from two individuals randomly sampled from their respective population groups are linked
#
# ------------------------------------------------------------------------------------

get_p_hat <- function(df_counts) {

    df_counts$p_hat <- df_counts$num_linked_pairs_observed / df_counts$max_possible_pairs_in_sample

    df_counts

}

# Function: get_prob_group_pairing_and_linked ---------------------------
#
# Goal: Compute probability that a pair is from a specific population group pairing and is linked.
#       Example: For a group pairing (u,v), prob. that a pair is from (u,v) and linked is:
#                (N_uv / N_choose_2) * p_hat_uv
#       where:
#       N_uv = N_u * N_v i.e. maximum distinct possible (u,v) pairs in population.
#       p_hat_uv = fraction of observed (u,v) pairs between samples from population groups
#                  u and v that are linked.
#       N choose 2 or (N * (N - 1))/2 i.e. all distinct possible pairs in population.
#
# parameters:
#
#     df_counts_and_p_hat                        A data.frame returned by function:
#                                                get_p_hat()
#     individuals_population_in        (numeric) A vector of the estimated number of individuals
#                                                per population group
#     verbose_output                   (boolean) A value to display intermediate output
#
#
# returns: a data.frame containing:
#
#
#     H1_group                         (Factor)  Name of population group 1
#     H2_group                         (Factor)  Name of population group 2
#     number_hosts_sampled_group_1     (numeric) Number of individuals sampled from population group 1
#     number_hosts_sampled_group_2     (numeric) Number of individuals sampled from population group 2
#     number_hosts_population_group_1  (numeric) Estimated number of individuals in population group 1
#     number_hosts_population_group_2  (numeric) Estimated number of individuals in population group 2
#     max_possible_pairs_in_sample     (numeric) Number of distinct possible transmission pairs between individuals sampled from population groups 1 and 2
#     max_possible_pairs_in_population (numeric) Number of distinct possible transmission pairs between individuals in population groups 1 and 2
#     num_linked_pairs_observed        (numeric) Number of observed directed transmission pairs between samples from population groups 1 and 2
#     p_hat                            (numeric) Probability that pathogen sequences from two individuals randomly sampled from their respective population groups are linked
#     prob_group_pairing_and_linked    (numeric) Probability that a pair of pathogen sequences is from a specific population group pairing and is linked
#
# ------------------------------------------------------------------------------------

get_prob_group_pairing_and_linked <- function(df_counts_and_p_hat, individuals_population_in,
                                              verbose_output = FALSE) {

    # View dataset structure
	if (verbose_output) utils::str(df_counts_and_p_hat)

    # Sum up estimated number of individuals in each population group to get total population
	individuals_population <- base::as.numeric(individuals_population_in)
    individuals_population_total <- base::sum(individuals_population)
    if (verbose_output) base::print(individuals_population_total)

    # Get all distinct possible pairs in population
    max_possible_pairs_in_total_population <- (individuals_population_total * (individuals_population_total - 1)) / 2

    df_counts_and_p_hat$prob_group_pairing_and_linked <- (df_counts_and_p_hat$max_possible_pairs_in_population / max_possible_pairs_in_total_population) * df_counts_and_p_hat$p_hat

    df_counts_and_p_hat

}

# Function: get_theta_hat ---------------------------
#
# Goal: Compute the conditional probability that a pair is from a specific population group
#       pairing given the pair is linked.
#       Example: For a group pairing (u,v), theta_hat denotes estimated transmission
#                flows within and between population groups u and v adjusted for sampling
#                heterogeneity.
#
# parameters:
#
#     df_counts_and_p_hat                        A data.frame returned by the function:
#                                                get_p_hat()
#
# returns: a data.frame containing:
#
#
#     H1_group                         (Factor)  Name of population group 1
#     H2_group                         (Factor)  Name of population group 2
#     number_hosts_sampled_group_1     (numeric) Number of individuals sampled from population group 1
#     number_hosts_sampled_group_2     (numeric) Number of individuals sampled from population group 2
#     number_hosts_population_group_1  (numeric) Estimated number of individuals in population group 1
#     number_hosts_population_group_2  (numeric) Estimated number of individuals in population group 2
#     max_possible_pairs_in_sample     (numeric) Number of distinct possible transmission pairs between individuals sampled from population groups 1 and 2
#     max_possible_pairs_in_population (numeric) Number of distinct possible transmission pairs between individuals in population groups 1 and 2
#     num_linked_pairs_observed        (numeric) Number of observed directed transmission pairs between samples from population groups 1 and 2
#     p_hat                            (numeric) Probability that pathogen sequences from two individuals randomly sampled from their respective population groups are linked
#     est_linkedpairs_in_population    (numeric) Estimated transmission pairs between population groups 1 and 2
#     theta_hat                        (numeric) Estimated transmission flows or relative probability of transmission within and between population groups 1 and 2 adjusted
#                                                for sampling heterogeneity. More precisely, the conditional probability that a pair of pathogen sequences is from a specific population
#                                                group pairing given that the pair is linked.
#
# ------------------------------------------------------------------------------------

get_theta_hat <- function(df_counts_and_p_hat) {

	df_counts_and_p_hat$est_linkedpairs_in_population <- df_counts_and_p_hat$max_possible_pairs_in_population * df_counts_and_p_hat$p_hat

	df_counts_and_p_hat$theta_hat <- df_counts_and_p_hat$est_linkedpairs_in_population / base::sum(df_counts_and_p_hat$est_linkedpairs_in_population)

	df_counts_and_p_hat

}

# Function: get_c_hat ---------------------------
#
# Goal: Compute c_hat, for a group pairing (u,v) c_hat is the probability that a pathogen sequence from
#       an individual in population group u links with pathogen sequences from one or more other individuals
#       in population group v (excluding itself if u = v) i.e. probability of clustering.
#
# parameters:
#
#     df_counts_and_p_hat                        A data.frame returned by the function:
#                                                get_p_hat()
#
# returns: a data.frame containing:
#
#
#     H1_group                         (Factor)  Name of population group 1
#     H2_group                         (Factor)  Name of population group 2
#     number_hosts_sampled_group_1     (numeric) Number of individuals sampled from population group 1
#     number_hosts_sampled_group_2     (numeric) Number of individuals sampled from population group 2
#     number_hosts_population_group_1  (numeric) Estimated number of individuals in population group 1
#     number_hosts_population_group_2  (numeric) Estimated number of individuals in population group 2
#     max_possible_pairs_in_sample     (numeric) Number of distinct possible transmission pairs between individuals sampled from population groups 1 and 2
#     max_possible_pairs_in_population (numeric) Number of distinct possible transmission pairs between individuals in population groups 1 and 2
#     num_linked_pairs_observed        (numeric) Number of observed directed transmission pairs between samples from population groups 1 and 2
#     p_hat                            (numeric) Probability that pathogen sequences from two individuals randomly sampled from their respective population groups are linked
#     c_hat                            (numeric) Probability that a randomly selected pathogen sequence in one population group links to at least
#                                                one pathogen sequence in another population group i.e. probability of clustering
#
# ------------------------------------------------------------------------------------

get_c_hat <- function(df_counts_and_p_hat) {

	df_counts_and_p_hat$c_hat <- 1 - ((1 - df_counts_and_p_hat$p_hat) ^ df_counts_and_p_hat$number_hosts_population_group_2)

	df_counts_and_p_hat

}

# Calculate simultaneous confidence intervals for a multinomial proportion
#---------------------------

# Function: get_multinomial_proportion_conf_ints ---------------------------
#
# Goal: Compute simultaneous confidence intervals at the 5% significance level
#
# Reminder: Goodman CIs work for fewer cells with larger counts and Sison-Glaz many cells with smaller counts
#
# parameters:
#
#     df_theta_hat                               A data.frame returned by the function:
#                                                get_theta_hat()
#     detailed_report                  (boolean) A value to produce detailed output of the analysis
#
# returns: a data.frame containing:
#
#
#     H1_group                         (Factor)  Name of population group 1
#     H2_group                         (Factor)  Name of population group 2
#     number_hosts_sampled_group_1     (numeric) Number of individuals sampled from population group 1
#     number_hosts_sampled_group_2     (numeric) Number of individuals sampled from population group 2
#     number_hosts_population_group_1  (numeric) Estimated number of individuals in population group 1
#     number_hosts_population_group_2  (numeric) Estimated number of individuals in population group 2
#     max_possible_pairs_in_sample     (numeric) Number of distinct possible transmission pairs between individuals sampled from population groups 1 and 2
#     max_possible_pairs_in_population (numeric) Number of distinct possible transmission pairs between individuals in population groups 1 and 2
#     num_linked_pairs_observed        (numeric) Number of observed directed transmission pairs between samples from population groups 1 and 2
#     p_hat                            (numeric) Probability that pathogen sequences from two individuals randomly sampled from their respective population groups are linked
#     est_linkedpairs_in_population    (numeric) Estimated transmission pairs between population groups 1 and 2
#     theta_hat                        (numeric) Estimated transmission flows or relative probability of transmission within and between population groups 1 and 2 adjusted
#                                                for sampling heterogeneity. More precisely, the conditional probability that a pair of pathogen sequences is from a specific population
#                                                group pairing given that the pair is linked.
#     obs_trm_pairs_est_goodman        (numeric) Point estimate, Goodman method Confidence intervals for observed transmission pairs
#     obs_trm_pairs_lwr_ci_goodman     (numeric) Lower bound of Goodman confidence interval
#     obs_trm_pairs_upr_ci_goodman     (numeric) Upper bound of Goodman confidence interval
#     est_goodman                      (numeric) Point estimate, Goodman method Confidence intervals for estimated transmission flows
#     lwr_ci_goodman                   (numeric) Lower bound of Goodman confidence interval
#     upr_ci_goodman                   (numeric) Upper bound of Goodman confidence interval
#
#
# The following additional fields are returned if the detailed_report
#     flag is set
#
#     est_goodman_cc                   (numeric) Point estimate, Goodman method Confidence intervals with continuity correction
#     lwr_ci_goodman_cc                (numeric) Lower bound of Goodman confidence interval
#     upr_ci_goodman_cc                (numeric) Upper bound of Goodman confidence interval
#     est_sisonglaz                    (numeric) Point estimate, Sison-Glaz method Confidence intervals
#     lwr_ci_sisonglaz                 (numeric) Lower bound of Sison-Glaz confidence interval
#     upr_ci_sisonglaz                 (numeric) Upper bound of Sison-Glaz confidence interval
#     est_qhurst_acswr                 (numeric) Point estimate, Queensbury-Hurst method Confidence intervals via ACSWR r package
#     lwr_ci_qhurst_acswr              (numeric) Lower bound of Queensbury-Hurst confidence interval
#     upr_ci_qhurst_acswr              (numeric) Upper bound of Queensbury-Hurst confidence interval
#     est_qhurst_coinmind              (numeric) Point estimate, Queensbury-Hurst method Confidence intervals via CoinMinD r package
#     lwr_ci_qhurst_coinmind           (numeric) Lower bound of Queensbury-Hurst confidence interval
#     upr_ci_qhurst_coinmind           (numeric) Upper bound of Queensbury-Hurst confidence interval
#     lwr_ci_qhurst_adj_coinmind       (numeric) Lower bound of Queensbury-Hurst confidence interval adjusted
#     upr_ci_qhurst_adj_coinmind       (numeric) Upper bound of Queensbury-Hurst confidence interval adjusted
#
# ------------------------------------------------------------------------------------

get_multinomial_proportion_conf_ints_extended <- function(df_theta_hat, detailed_report = FALSE) {

    # Extend MultinomCI function in DescTools to compute Goodman confidence intervals
    # with a continuity correction i.e. goodmancc.
    # Acknowledgements:
    # (Steve Cherry, 1996) A Comparison of Confidence Interval Methods for Habitat Use-Availability Studies.
    # The Journal of Wildlife Management, Jul., 1996, Vol. 60, No. 3 (Jul., 1996), pp. 653-658
    # Goodman continuity correction described on pg. 655

	# Define function:
	MultinomCI_mod <- function (x, conf.level = 0.95, sides = c("two.sided", "left",
		"right"), method = c("sisonglaz", "cplus1", "goodman", "goodmancc",
		"wald", "waldcc", "wilson")) {

		.moments <- function(c, lambda) {
			a <- lambda + c
			b <- lambda - c
			if (b < 0)
				b <- 0
			if (b > 0)
				den <- stats::ppois(a, lambda) - stats::ppois(b - 1, lambda)
			if (b == 0)
				den <- stats::ppois(a, lambda)
			mu <- mat.or.vec(4, 1)
			mom <- mat.or.vec(5, 1)
			for (r in 1:4) {
				poisA <- 0
				poisB <- 0
				if ((a - r) >= 0) {
					poisA <- stats::ppois(a, lambda) - stats::ppois(a - r, lambda)
				}
				if ((a - r) < 0) {
					poisA <- stats::ppois(a, lambda)
				}
				if ((b - r - 1) >= 0) {
					poisB <- stats::ppois(b - 1, lambda) - stats::ppois(b - r -
					  1, lambda)
				}
				if ((b - r - 1) < 0 && (b - 1) >= 0) {
					poisB <- stats::ppois(b - 1, lambda)
				}
				if ((b - r - 1) < 0 && (b - 1) < 0) {
					poisB <- 0
				}
				mu[r] <- (lambda^r) * (1 - (poisA - poisB)/den)
			}
			mom[1] <- mu[1]
			mom[2] <- mu[2] + mu[1] - mu[1]^2
			mom[3] <- mu[3] + mu[2] * (3 - 3 * mu[1]) + (mu[1] -
				3 * mu[1]^2 + 2 * mu[1]^3)
			mom[4] <- mu[4] + mu[3] * (6 - 4 * mu[1]) + mu[2] * (7 -
				12 * mu[1] + 6 * mu[1]^2) + mu[1] - 4 * mu[1]^2 +
				6 * mu[1]^3 - 3 * mu[1]^4
			mom[5] <- den
			return(mom)
		}
		.truncpoi <- function(c, x, n, k) {
			m <- matrix(0, k, 5)
			for (i in 1L:k) {
				lambda <- x[i]
				mom <- .moments(c, lambda)
				for (j in 1L:5L) {
					m[i, j] <- mom[j]
				}
			}
			for (i in 1L:k) {
				m[i, 4] <- m[i, 4] - 3 * m[i, 2]^2
			}
			s <- colSums(m)
			s1 <- s[1]
			s2 <- s[2]
			s3 <- s[3]
			s4 <- s[4]
			probn <- 1/(stats::ppois(n, n) - stats::ppois(n - 1, n))
			z <- (n - s1)/sqrt(s2)
			g1 <- s3/(s2^(3/2))
			g2 <- s4/(s2^2)
			poly <- 1 + g1 * (z^3 - 3 * z)/6 + g2 * (z^4 - 6 * z^2 +
				3)/24
			+g1^2 * (z^6 - 15 * z^4 + 45 * z^2 - 15)/72
			f <- poly * exp(-z^2/2)/(sqrt(2) * gamma(0.5))
			probx <- 1
			for (i in 1L:k) {
				probx <- probx * m[i, 5]
			}
			return(probn * probx * f/sqrt(s2))
		}
		n <- sum(x, na.rm = TRUE)
		k <- length(x)
		p <- x/n
		if (missing(method))
			method <- "sisonglaz"
		if (missing(sides))
			sides <- "two.sided"
		sides <- match.arg(sides, choices = c("two.sided", "left",
			"right"), several.ok = FALSE)
		if (sides != "two.sided")
			conf.level <- 1 - 2 * (1 - conf.level)
		method <- match.arg(arg = method, choices = c("sisonglaz",
			"cplus1", "goodman", "goodmancc", "wald", "waldcc", "wilson"))
		if (method == "goodman") {
			q.chi <- stats::qchisq(conf.level, k - 1)
			lci <- (q.chi + 2 * x - sqrt(q.chi * (q.chi + 4 * x *
				(n - x)/n)))/(2 * (n + q.chi))
			uci <- (q.chi + 2 * x + sqrt(q.chi * (q.chi + 4 * x *
				(n - x)/n)))/(2 * (n + q.chi))
			res <- cbind(est = p, lwr.ci = pmax(0, lci), upr.ci = pmin(1,
				uci))
		}
		else if (method == "goodmancc") {
			q.chi <- stats::qchisq(conf.level, k - 1)
			lci <- (q.chi + 2 * (x - 0.5) - sqrt(q.chi * (q.chi + 4 * (x - 0.5) *
				(n - x + 0.5)/n)))/(2 * (n + q.chi))
			uci <- (q.chi + 2 * (x + 0.5) + sqrt(q.chi * (q.chi + 4 * (x + 0.5) *
				(n - x - 0.5)/n)))/(2 * (n + q.chi))
			res <- cbind(est = p, lwr.ci = pmax(0, lci), upr.ci = pmin(1,
				uci))
		}
		else if (method == "wald") {
			q.chi <- stats::qchisq(conf.level, 1)
			lci <- p - sqrt(q.chi * p * (1 - p)/n)
			uci <- p + sqrt(q.chi * p * (1 - p)/n)
			res <- cbind(est = p, lwr.ci = pmax(0, lci), upr.ci = pmin(1,
				uci))
		}
		else if (method == "waldcc") {
			q.chi <- stats::qchisq(conf.level, 1)
			lci <- p - sqrt(q.chi * p * (1 - p)/n) - 1/(2 * n)
			uci <- p + sqrt(q.chi * p * (1 - p)/n) + 1/(2 * n)
			res <- cbind(est = p, lwr.ci = pmax(0, lci), upr.ci = pmin(1,
				uci))
		}
		else if (method == "wilson") {
			q.chi <- stats::qchisq(conf.level, 1)
			lci <- (q.chi + 2 * x - sqrt(q.chi^2 + 4 * x * q.chi *
				(1 - p)))/(2 * (q.chi + n))
			uci <- (q.chi + 2 * x + sqrt(q.chi^2 + 4 * x * q.chi *
				(1 - p)))/(2 * (q.chi + n))
			res <- cbind(est = p, lwr.ci = pmax(0, lci), upr.ci = pmin(1,
				uci))
		}
		else {
			const <- 0
			pold <- 0
			for (cc in 1:n) {
				poi <- .truncpoi(cc, x, n, k)
				if (poi > conf.level && pold < conf.level) {
					const <- cc
					break
				}
				pold <- poi
			}
			delta <- (conf.level - pold)/(poi - pold)
			const <- const - 1
			if (method == "sisonglaz") {
				res <- cbind(est = p, lwr.ci = pmax(0, p - const/n),
					upr.ci = pmin(1, p + const/n + 2 * delta/n))
			}
			else if (method == "cplus1") {
				res <- cbind(est = p, lwr.ci = pmax(0, p - const/n -
					1/n), upr.ci = pmin(1, p + const/n + 1/n))
			}
		}
		if (sides == "left")
			res[3] <- 1
		else if (sides == "right")
			res[2] <- 0
		return(res)
	}

	# Acknowledgments
	# QH (Queensbury-Hurst) confidence intervals function sourced from CoinMinD package and tweaked
	#     to report and store the estimate and confidence intervals.
	# Default behaviour: Print confidence intervals to screen without reporting the estimate.

	# Define function:
	QH_coinmind_mod <- function(inpmat, alpha) {

		k = base::length(inpmat)
		s = base::sum(inpmat)
		chi = stats::qchisq(1 - alpha, df = k - 1)
		pi = inpmat / s
		QH.UL = (chi + 2 * inpmat + sqrt(chi * chi + 4 * inpmat *
			chi * (1 - pi)))/(2 * (chi + s))
		QH.LL = (chi + 2 * inpmat - sqrt(chi * chi + 4 * inpmat *
			chi * (1 - pi)))/(2 * (chi + s))
		LLA = 0
		ULA = 0
		for (r in 1:base::length(inpmat)) {
			if (QH.LL[r] < 0)
				LLA[r] = 0
			else LLA[r] = QH.LL[r]
			if (QH.UL[r] > 1)
				ULA[r] = 1
			else ULA[r] = QH.UL[r]
		}
		diA = ULA - LLA
		VOL = base::round(prod(diA), 8)

		output <- base::cbind("est_qhurst_coinmind" = pi, "lwr_ci_qhurst_coinmind" = QH.LL, "upr_ci_qhurst_coinmind" = QH.UL, "lwr_ci_qhurst_adj_coinmind" = LLA, "upr_ci_qhurst_adj_coinmind" = ULA)

		output

	}

	# Acknowledgments
	# QH (Queensbury-Hurst) confidence intervals function sourced from ACSWR package and tweaked
	#     to report the estimate as well as confidence intervals.
	# Default behaviour: Report confidence intervals without the estimate.

	# Define function:
	QH_CI_acswr_mod <- function(x, alpha) {

	  k <- base::length(x); n <- base::sum(x)

	  QH_est <- x / n

	  QH_lcl <- (1 / (2 * (base::sum(x) + stats::qchisq(1 - alpha / k, k - 1)))) *
		  {
			  stats::qchisq(1 - alpha / k, k - 1) + 2 * x - base::sqrt( stats::qchisq(1 - alpha / k, k - 1) *
			  (stats::qchisq(1 - alpha / k, k - 1) + 4 * x *(base::sum(x) - x) / base::sum(x)))

		  }

	  QH_ucl <- (1 / (2 * (base::sum(x) + stats::qchisq(1 - alpha / k, k - 1)))) *
		  {
			  stats::qchisq(1 - alpha / k, k - 1) + 2 * x + base::sqrt( stats::qchisq(1 - alpha / k, k - 1) *
			  (stats::qchisq(1 - alpha / k, k - 1) + 4 * x * (base::sum(x) - x) / base::sum(x)))
		  }

	  return(base::cbind("est_qhurst_acswr" = QH_est, "lwr_ci_qhurst_acswr" = QH_lcl, "upr_ci_qhurst_acswr" = QH_ucl))

	}

    # compute multi-nomial confidence intervals

    # Note: We scale the estimated transmission pairs in the population such that the sum of the weights
    #       is equal to the observed number of transmission pairs. This is a common approach in survey sampling.
    #       Basically, we express the estimated transmission pairs in the population as proportions then multiply
    #       by the sum of the observed counts
    weighted_est_linkedpairs_in_population <- (df_theta_hat$est_linkedpairs_in_population / base::sum(df_theta_hat$est_linkedpairs_in_population)) * base::sum(df_theta_hat$num_linked_pairs_observed)

    # compute Goodman confidence intervals for observed transmission events
    obs_trm_pairs_ci_goodman <- base::cbind(df_theta_hat, MultinomCI_mod(df_theta_hat$num_linked_pairs_observed, method="goodman"))
    obs_trm_pairs_ci_goodman <- obs_trm_pairs_ci_goodman %>% dplyr::rename(obs_trm_pairs_est_goodman = est, obs_trm_pairs_lwr_ci_goodman = lwr.ci, obs_trm_pairs_upr_ci_goodman = upr.ci)

 	# compute Goodman intervals
    df_theta_hat_goodman <- base::cbind(obs_trm_pairs_ci_goodman, MultinomCI_mod(weighted_est_linkedpairs_in_population, method="goodman"))
    df_theta_hat_goodman <- df_theta_hat_goodman %>% dplyr::rename(est_goodman = est, lwr_ci_goodman = lwr.ci, upr_ci_goodman = upr.ci)

    if (detailed_report) {

		# compute Goodman intervals with a continuity correction
		df_theta_hat_goodman_cc <- base::cbind(df_theta_hat_goodman, MultinomCI_mod(weighted_est_linkedpairs_in_population, method="goodmancc"))
		df_theta_hat_goodman_cc <- df_theta_hat_goodman_cc %>% dplyr::rename(est_goodman_cc = est, lwr_ci_goodman_cc = lwr.ci, upr_ci_goodman_cc = upr.ci)

		# compute Sison-Glaz intervals
		df_theta_hat_goodman_sisonglaz <- base::cbind(df_theta_hat_goodman_cc, MultinomCI_mod(weighted_est_linkedpairs_in_population, method="sisonglaz"))
		df_theta_hat_goodman_sisonglaz <- df_theta_hat_goodman_sisonglaz %>% dplyr::rename(est_sisonglaz = est, lwr_ci_sisonglaz = lwr.ci, upr_ci_sisonglaz = upr.ci)

		# compute Queensbury-Hurst intervals
		df_theta_hat_goodman_sisonglaz_qhurst_acswr <- base::cbind(df_theta_hat_goodman_sisonglaz, QH_CI_acswr_mod(weighted_est_linkedpairs_in_population, 0.05))

		df_theta_hat_goodman_sisonglaz_qhurst_acswr_coinmind <- base::cbind(df_theta_hat_goodman_sisonglaz_qhurst_acswr, QH_coinmind_mod(weighted_est_linkedpairs_in_population, 0.05))

		df_theta_hat_goodman_sisonglaz_qhurst_acswr_coinmind <- df_theta_hat_goodman_sisonglaz_qhurst_acswr_coinmind %>% dplyr::arrange(-num_linked_pairs_observed)

	    output <- df_theta_hat_goodman_sisonglaz_qhurst_acswr_coinmind

	} else {

        output <- df_theta_hat_goodman %>% dplyr::arrange(-num_linked_pairs_observed)

	}

    output

}