\name{extreme_deconvolution}
\alias{extreme_deconvolution}

\title{Density estimation using Gaussian mixtures in the presence of
  noisy, heterogeneous and incomplete data}

\description{We present a general algorithm to infer a d-dimensional
distribution function given a set of heterogeneous, noisy observations
or samples. This algorithm reconstructs the error-deconvolved or
'underlying' distribution function common to all samples, even when the
individual samples have unique error and missing-data properties. The
underlying distribution is modeled as a mixture of Gaussians, which is
completely general. Model parameters are chosen to optimize a justified,
scalar objective function: the logarithm of the probability of the data
under the error-convolved model, where the error convolution is
different for each data point. Optimization is performed by an
Expectation Maximization (EM) algorithm, extended by a regularization
technique and 'split-and-merge' procedure. These extensions mitigate
problems with singularities and local maxima, which are often
encountered when using the EM algorithm to estimate Gaussian density
mixtures.}

\usage{
extreme_deconvolution(ydata,ycovar,xamp,xmean,xcovar,
                      projection=NULL,weight=NULL,
                      fixamp=NULL,fixmean=NULL,fixcovar=NULL,
                      tol=1.e-6,maxiter=1e9,w=0,logfile=NULL,
                      splitnmerge=0,maxsnm=FALSE,likeonly=FALSE,
                      logweight=FALSE)
}

\arguments{
  \item{ydata}{[ndata,dy] matrix of observed quantities}
  \item{ycovar}{[ndata,dy] / [ndata,dy,dy] / [dy,dy,ndata] matrix, list or 3D array
    of observational error covariances
    (if [ndata,dy] then the error correlations are assumed to vanish)}
  \item{xamp}{[ngauss] array of initial amplitudes (*not* [1,ngauss])}
  \item{xmean}{[ngauss,dx] matrix of initial means}
  \item{xcovar}{[ngauss,dx,dx] list of matrices of initial covariances}
  \item{projection}{[ndata,dy,dx] list of projection matrices}
  \item{weight}{[ndata] array of weights to be applied to the data points}
  \item{logweight}{(bool, default=False) if True, weight is actually
                   log(weight)}
  \item{fixamp}{(default=None) None, True/False, or list of bools}
  \item{fixmean}{(default=None) None, True/False, or list of bools}
  \item{fixcovar}{(default=None) None, True/False, or list of bools}
  \item{tol}{(double, default=1.e-6) tolerance for convergence}
  \item{maxiter}{(long, default= 10**9) maximum number of iterations to perform}
  \item{w}{(double, default=0.) covariance regularization parameter
            (of the conjugate prior)}
  \item{logfile}{basename for several logfiles (_c.log has output from
                 the c-routine; _loglike.log has the log likelihood path of
                 all the accepted routes, i.e. only parts which increase
                 the likelihood are included, during splitnmerge)}
  \item{splitnmerge}{(int, default=0) depth to go down the splitnmerge path}
  \item{maxsnm}{(Bool, default=False) use the maximum number of split 'n'
                 merge steps, K*(K-1)*(K-2)/2}
  \item{likeonly}{(Bool, default=False) only compute the total log
                   likelihood of the data}
}

\value{
  \item{avgloglikedata}{avgloglikedata after convergence}
  \item{xamp}{updated xamp}
  \item{xmean}{updated xmean}
  \item{xcovar}{updated xcovar}
}

\details{
  ...
}

\author{Jo Bovy, David W. Hogg, & Sam T. Roweis}

\references{
Inferring complete distribution functions from
noisy, heterogeneous and incomplete observations Jo Bovy, David
W. Hogg, & Sam T. Roweis, Submitted to AOAS (2009) [arXiv/0905.2979]
}

\examples{
library(ExtremeDeconvolution)

ydata <-
c(2.62434536, 0.38824359, 0.47182825, -0.07296862, 1.86540763,
-1.30153870, 2.74481176, 0.23879310, 1.31903910, 0.75062962,
2.46210794, -1.06014071, 0.67758280, 0.61594565, 2.13376944,
-0.09989127, 0.82757179, 0.12214158, 1.04221375, 1.58281521,
-0.10061918, 2.14472371, 1.90159072, 1.50249434, 1.90085595,
0.31627214, 0.87710977, 0.06423057, 0.73211192, 1.53035547,
0.30833925, 0.60324647, 0.31282730, 0.15479436, 0.32875387,
0.98733540, -0.11731035, 1.23441570, 2.65980218, 1.74204416,
0.80816445, 0.11237104, 0.25284171, 2.69245460, 1.05080775,
0.36300435, 1.19091548, 3.10025514, 1.12015895, 1.61720311,
1.30017032, 0.64775015, -0.14251820, 0.65065728, 0.79110577,
1.58662319, 1.83898341, 1.93110208, 1.28558733, 1.88514116,
0.24560206, 2.25286816, 1.51292982, 0.70190717, 1.48851815,
0.92442829, 2.13162939, 2.51981682, 3.18557541, -0.39649633,
-0.44411380, 0.49553414, 1.16003707, 1.87616892, 1.31563495,
-1.02220122, 0.69379599, 1.82797464, 1.23009474, 1.76201118,
0.77767186, 0.79924193, 1.18656139, 1.41005165, 1.19829972,
1.11900865, 0.32933771, 1.37756379, 1.12182127, 2.12948391,
2.19891788, 1.18515642, 0.62471505, 0.36126959, 1.42349435,
1.07734007, 0.65614632, 1.04359686, 0.37999916, 1.69803203,
0.55287144, 2.22450770, 1.40349164, 1.59357852, -0.09491185,
1.16938243, 1.74055645, 0.04629940, 0.73378149, 1.03261455,
-0.37311732, 1.31515939, 1.84616065, 0.14048406, 1.35054598,
-0.31228341, 0.96130449, -0.61577236, 2.12141771, 1.40890054,
0.97538304, 0.22483838, 2.27375593, 2.96710175, -0.85798186,
2.23616403, 2.62765075, 1.33801170, -0.19926803, 1.86334532,
0.81907970, 0.39607937, -0.23005814, 1.55053750, 1.79280687,
0.37646927, 1.52057634, -0.14434139, 1.80186103, 1.04656730,
0.81343023, 0.89825413, 1.86888616, 1.75041164, 1.52946532,
1.13770121, 1.07782113, 1.61838026, 1.23249456, 1.68255141,
0.68988323, -1.43483776, 2.03882460, 3.18697965, 1.44136444,
0.89984477, 0.86355526, 0.88094581, 1.01740941, -0.12201873,
0.48290554, 0.00297317, 1.24879916, 0.70335885, 1.49521132,
0.82529684, 1.98633519, 1.21353390, 3.19069973, -0.89636092,
0.35308331, 1.90148689, 3.52832571, 0.75136522, 1.04366899,
0.77368576, 2.33145711, 0.71269214, 1.68006984, 0.68019840,
-0.27255875, 1.31354772, 1.50318481, 2.29322588, 0.88955297,
0.38263794, 1.56276110, 1.24073709, 1.28066508, 0.92688730,
2.16033857, 1.36949272, 2.90465871, 2.11105670, 1.65904980,
-0.62743834, 1.60231928, 1.42028220, 1.81095167, 2.04444209)

ydata  <- matrix(ydata,length(ydata),1)
N      <- dim(ydata)[1]
ycovar <- ydata*0 + 0.01
xamp   <- c(0.5,0.5)
xmean  <- matrix(c(0.86447943, 0.67078879, 0.322681, 0.45087394),2,2)
xcovar <-
  list(matrix(c(0.03821028, 0.04014796, 0.04108113, 0.03173839),2,2),
       matrix(c(0.06219194, 0.09738021, 0.04302473, 0.06778009),2,2))
projection <- list()
for (i in 1:N)
  projection[[i]] = matrix(c(i\%\%2,(i+1)\%\%2),1,2)
res <- extreme_deconvolution(ydata, ycovar, xamp, xmean, xcovar,
         projection=projection, logfile="ExDeconDemo")
print(res)
stopifnot((res$avgloglikedata - (-1.3114744655258121)) **2. < 10.**-8)
stopifnot((res$xmean[1,1]-2.30368235)**2. < 10.**-5)
stopifnot((res$xmean[1,2]-1.70701517)**2. < 10.**-5)
stopifnot((res$xmean[2,1]-1.08009397)**2. < 10.**-5)
stopifnot((res$xmean[2,2]-0.8888667)**2. < 10.**-5)
stopifnot((res$xcovar[[1]][1,1]-0.445645987259)**2. < 10.**-5.)
stopifnot((res$xamp[1]-0.11968415)**2. < 10.**-5)
stopifnot((res$xamp[2]-0.880315852981)**2. < 10.**-5)
}


LISP Interpreter Run

[[[[[

 Show that a formal axiomatic system (fas) can only prove
 that finitely many LISP expressions are elegant.
 (An expression is elegant if no smaller expression has
 the same value.)

 More precisely, show that a fas of LISP complexity N can't
 prove that a LISP expression X is elegant if X's size is
 greater than N + 356.

 (fas N) returns the theorem proved by the Nth proof
 (Nth S-expression) in the fas, or nil if the proof is
 invalid, or stop to stop everything.

]]]]]

[
 This expression searches for an elegant expression
 that is larger than it is and returns the value of
 that expression as its own value.
]

define expression  [Formal Axiomatic System #1]
       let (fas n) if = n 1 '(is-elegant x)
                   if = n 2  nil
                   if = n 3 '(is-elegant yyy)
                   [else]    stop

       let (loop n)
           let theorem [be] display (fas n)
           if = nil theorem [then] (loop + n 1)
           if = stop theorem [then] fas-has-stopped
           if = is-elegant car theorem
              if > display size cadr theorem
                   display + 356 size fas
                 [return] eval cadr theorem
              [else] (loop + n 1)
           [else] (loop + n 1)

       (loop 1)

define      expression
value       ((' (lambda (fas) ((' (lambda (loop) (loop 1))) ('
             (lambda (n) ((' (lambda (theorem) (if (= nil theo
            rem) (loop (+ n 1)) (if (= stop theorem) fas-has-s
            topped (if (= is-elegant (car theorem)) (if (> (di
            splay (size (car (cdr theorem)))) (display (+ 356
            (size fas)))) (eval (car (cdr theorem))) (loop (+
            n 1))) (loop (+ n 1))))))) (display (fas n))))))))
             (' (lambda (n) (if (= n 1) (' (is-elegant x)) (if
             (= n 2) nil (if (= n 3) (' (is-elegant yyy)) stop
            ))))))

[Show that this expression knows its own size.]

size expression

expression  (size expression)
value       456

[
 Run #1.

 Here it doesn't find an elegant expression
 larger than it is:
]

eval expression

expression  (eval expression)
display     (is-elegant x)
display     1
display     456
display     ()
display     (is-elegant yyy)
display     3
display     456
display     stop
value       fas-has-stopped

define expression  [Formal Axiomatic System #2]
       let (fas n) if = n 1 '(is-elegant x)
                   if = n 2  nil
                   if = n 3 '(is-elegant yyy)
                   if = n 4  cons is-elegant
                             cons ^ 10 509     [<=====]
                                  nil
                   [else]    stop

       let (loop n)
           let theorem [be] display (fas n)
           if = nil theorem [then] (loop + n 1)
           if = stop theorem [then] fas-has-stopped
           if = is-elegant car theorem
              if > display size cadr theorem
                   display + 356 size fas
                 [return] eval cadr theorem
              [else] (loop + n 1)
           [else] (loop + n 1)

       (loop 1)

define      expression
value       ((' (lambda (fas) ((' (lambda (loop) (loop 1))) ('
             (lambda (n) ((' (lambda (theorem) (if (= nil theo
            rem) (loop (+ n 1)) (if (= stop theorem) fas-has-s
            topped (if (= is-elegant (car theorem)) (if (> (di
            splay (size (car (cdr theorem)))) (display (+ 356
            (size fas)))) (eval (car (cdr theorem))) (loop (+
            n 1))) (loop (+ n 1))))))) (display (fas n))))))))
             (' (lambda (n) (if (= n 1) (' (is-elegant x)) (if
             (= n 2) nil (if (= n 3) (' (is-elegant yyy)) (if
            (= n 4) (cons is-elegant (cons (^ 10 509) nil)) st
            op)))))))

[Show that this expression knows its own size.]

size expression

expression  (size expression)
value       509

[
 Run #2.

 Here it finds an elegant expression
 exactly one character larger than it is:
]

eval expression

expression  (eval expression)
display     (is-elegant x)
display     1
display     509
display     ()
display     (is-elegant yyy)
display     3
display     509
display     (is-elegant 10000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            0000000000000000000000)
display     510
display     509
value       10000000000000000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            0000000000

define expression  [Formal Axiomatic System #3]
       let (fas n) if = n 1 '(is-elegant x)
                   if = n 2  nil
                   if = n 3 '(is-elegant yyy)
                   if = n 4  cons is-elegant
                             cons ^ 10 508     [<=====]
                                  nil
                   [else]    stop

       let (loop n)
           let theorem [be] display (fas n)
           if = nil theorem [then] (loop + n 1)
           if = stop theorem [then] fas-has-stopped
           if = is-elegant car theorem
              if > display size cadr theorem
                   display + 356 size fas
                 [return] eval cadr theorem
              [else] (loop + n 1)
           [else] (loop + n 1)

       (loop 1)

define      expression
value       ((' (lambda (fas) ((' (lambda (loop) (loop 1))) ('
             (lambda (n) ((' (lambda (theorem) (if (= nil theo
            rem) (loop (+ n 1)) (if (= stop theorem) fas-has-s
            topped (if (= is-elegant (car theorem)) (if (> (di
            splay (size (car (cdr theorem)))) (display (+ 356
            (size fas)))) (eval (car (cdr theorem))) (loop (+
            n 1))) (loop (+ n 1))))))) (display (fas n))))))))
             (' (lambda (n) (if (= n 1) (' (is-elegant x)) (if
             (= n 2) nil (if (= n 3) (' (is-elegant yyy)) (if
            (= n 4) (cons is-elegant (cons (^ 10 508) nil)) st
            op)))))))

[Show that this expression knows its own size.]

size expression

expression  (size expression)
value       509

[
 Run #3.

 Here it finds an elegant expression
 exactly the same size as it is:
]

eval expression

expression  (eval expression)
display     (is-elegant x)
display     1
display     509
display     ()
display     (is-elegant yyy)
display     3
display     509
display     (is-elegant 10000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            000000000000000000000)
display     509
display     509
display     stop
value       fas-has-stopped

define expression  [Formal Axiomatic System #4]
       let (fas n) if = n 1 '(is-elegant x)
                   if = n 2  nil
                   if = n 3 '(is-elegant yyy)
                   if = n 4  cons is-elegant
                             cons cons "-
                                  cons ^ 10 600  [<=====]
                                  cons 1
                                       nil
                                  nil
                   [else]    stop

       let (loop n)
           let theorem [be] display (fas n)
           if = nil theorem [then] (loop + n 1)
           if = stop theorem [then] fas-has-stopped
           if = is-elegant car theorem
              if > display size cadr theorem
                   display + 356 size fas
                 [return] eval cadr theorem
              [else] (loop + n 1)
           [else] (loop + n 1)

       (loop 1)

define      expression
value       ((' (lambda (fas) ((' (lambda (loop) (loop 1))) ('
             (lambda (n) ((' (lambda (theorem) (if (= nil theo
            rem) (loop (+ n 1)) (if (= stop theorem) fas-has-s
            topped (if (= is-elegant (car theorem)) (if (> (di
            splay (size (car (cdr theorem)))) (display (+ 356
            (size fas)))) (eval (car (cdr theorem))) (loop (+
            n 1))) (loop (+ n 1))))))) (display (fas n))))))))
             (' (lambda (n) (if (= n 1) (' (is-elegant x)) (if
             (= n 2) nil (if (= n 3) (' (is-elegant yyy)) (if
            (= n 4) (cons is-elegant (cons (cons - (cons (^ 10
             600) (cons 1 nil))) nil)) stop)))))))

[Show that this expression knows its own size.]

size expression

expression  (size expression)
value       538

[
 Run #4.

 Here it finds an elegant expression
 much larger than it is, and evaluates it:
]

eval expression

expression  (eval expression)
display     (is-elegant x)
display     1
display     538
display     ()
display     (is-elegant yyy)
display     3
display     538
display     (is-elegant (- 10000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            00000000000000000000000000000000000000000000000000
            0000000000000000 1))
display     607
display     538
value       99999999999999999999999999999999999999999999999999
            99999999999999999999999999999999999999999999999999
            99999999999999999999999999999999999999999999999999
            99999999999999999999999999999999999999999999999999
            99999999999999999999999999999999999999999999999999
            99999999999999999999999999999999999999999999999999
            99999999999999999999999999999999999999999999999999
            99999999999999999999999999999999999999999999999999
            99999999999999999999999999999999999999999999999999
            99999999999999999999999999999999999999999999999999
            99999999999999999999999999999999999999999999999999
            99999999999999999999999999999999999999999999999999

End of LISP Run

Elapsed time is 0 seconds.


#######################################################################################################
#   Parameter estimation using R's optim function
#   Minimal error checking is done.  You should run is.marssMLE() before calling this.
#   Q and R are not allowed to be time-varying
#   Likelihood computation is via  Kalman filter
#######################################################################################################
MARSSoptim <- function(MLEobj) {
  # This function does not check if user specified a legal MLE object.
  ##
  # Define needed negLogLik function with chol transformed variances
  #
  neglogLik <- function(x, MLEobj = NULL) { # NULL assignment needed for optim call syntax
    # MLEobj is tmp.MLEobj so has altered free and fixed
    # x is the paramvector

    # update the MLEobj by putting the estimated pars from optim in
    MLEobj <- MARSSvectorizeparam(MLEobj, x)
    free <- MLEobj$marss$free
    pars <- MLEobj$par
    par.dims <- attr(MLEobj[["marss"]], "model.dims")
    for (elem in c("Q", "R", "V0")) {
      if (!is.fixed(free[[elem]])) # recompute par if needed since par in parlist is transformed
        {
          d <- sub3D(free[[elem]], t = 1) # this will be the one with the upper tri zero-ed out
          par.dim <- par.dims[[elem]][1:2]
          # t=1 since D not allowed to be time-varying; since code 4 lines down won't work otherwise
          L <- unvec(d %*% pars[[elem]], dim = par.dim) # this by def will have 0 row/col at the fixed values
          the.par <- tcrossprod(L) # L%*%t(L)
          # from f+Dm=M and if f!=0, D==0 so can leave off f
          MLEobj$par[[elem]] <- solve(crossprod(d)) %*% t(d) %*% vec(the.par)
          # solve(t(d)%*%d)%*%t(d)%*%vec(the.par)
        }
    } # end for over elem
    # This function is passed a special MLEobj with a marss.original element
    MLEobj$marss$fixed <- MLEobj$fixed.original
    MLEobj$marss$free <- MLEobj$free.original

    # kfsel selects the Kalman filter / smoother function based on MLEobj$fun.kf
    negLL <- MARSSkf(MLEobj, only.logLik = TRUE, return.lag.one = FALSE)$logLik

    -1 * negLL
  }

  if (!inherits(MLEobj, "marssMLE")) {
    stop("Stopped in MARSSoptim(). Object of class marssMLE is required.\n", call. = FALSE)
  }
  for (elem in c("Q", "R")) {
    if (dim(MLEobj$model$free[[elem]])[3] > 1) {
      stop(paste("Stopped in MARSSoptim() because this function does not allow estimated part of ", elem, " to be time-varying.\n", sep = ""), call. = FALSE)
    }
  }
  # the is.marssMODEL call is.validvarcov() which tests that the blocks are diagonal or unconstrained in the varcov matrices

  ## attach would be risky here since user might have one of these variables in their workspace
  MODELobj <- MLEobj[["marss"]]
  y <- MODELobj$data # must have time going across columns
  free <- MODELobj$free
  fixed <- MODELobj$fixed
  tmp.inits <- MLEobj$start
  control <- MLEobj$control
  par.dims <- attr(MODELobj, "model.dims")
  m <- par.dims[["x"]][1]
  n <- par.dims[["y"]][1]

  ## Set up the control list for optim; only pass in optim control elements
  control.names <- c("trace", "fnscale", "parscale", "ndeps", "maxit", "abstol", "reltol", "alpha", "beta", "gamma", "REPORT", "type", "lmm", "factr", "pgtol", "temp", "tmax")
  optim.control <- list()
  for (elem in control.names) {
    if (!is.null(control[[elem]])) optim.control[[elem]] <- control[[elem]]
  }
  if (is.null(control[["lower"]])) {
    lower <- -Inf
  } else {
    lower <- control[["lower"]]
  }
  if (is.null(control[["upper"]])) {
    upper <- Inf
  } else {
    upper <- control$upper
  }
  if (control$trace == -1) optim.control$trace <- 0

  # The code is used to set things up to use MARSSvectorizeparam to just select inits for the estimated parameters
  # Q=t(chol(Q)%*%chol(Q)); t(chol(Q)) has 0 in upper triangle
  tmp.MLEobj <- MLEobj
  # This is needed for the likelihood calculation
  tmp.MLEobj$fixed.original <- tmp.MLEobj$marss$fixed
  tmp.MLEobj$free.original <- tmp.MLEobj$marss$free

  tmp.MLEobj$par <- tmp.inits # set initial conditions for estimated parameters
  for (elem in c("Q", "R", "V0")) { # need the chol for these
    d <- sub3D(free[[elem]], t = 1) # free[[elem]] is required to be time constant
    f <- sub3D(fixed[[elem]], t = 1) # placeholder.  Need structure not actual values
    the.par <- unvec(f + d %*% tmp.inits[[elem]], dim = par.dims[[elem]][1:2])
    is.zero <- diag(the.par) == 0 # where the 0s on diagonal are
    if (any(is.zero)) diag(the.par)[is.zero] <- 1 # so the chol doesn't fail if there are zeros on the diagonal
    the.par <- t(chol(the.par)) # convert to transpose of chol
    if (any(is.zero)) diag(the.par)[is.zero] <- 0 # set back to 0
    if (!is.fixed(free[[elem]])) {
      # from f+Dm=M so m = solve(crossprod(d))%*%t(d)%*%(vec(the.par)-f)
      # but if d!=0,then f==0. if f!-0, then d==0.
      # Thus crossprod(d))%*%t(d) has 0 cols where fs appear in the.par and f is not needed
      tmp.MLEobj$par[[elem]] <- solve(crossprod(d)) %*% t(d) %*% vec(the.par)
    } else {
      tmp.MLEobj$par[[elem]] <- matrix(0, 0, 1)
    }
    # when being passed to optim, pars for var-cov mat is the chol, so need to reset free so we can get the L=t(chol) matrix
    # we don't need to reset fixed because it won't be used;
    # step 1, compute the D matrix corresponding to upper.tri=0 in t(chol)
    # note this only works because it is required that
    # a) if f!=0, d=0. so 1+a never appears in var-cov mat  b) a+b never appears in a var-cov mat, c) for BFGS, var-cov mat is time-invariant
    tmp.list.mat <- fixed.free.to.formula(f, d, par.dims[[elem]][1:2])
    tmp.list.mat[upper.tri(tmp.list.mat)] <- 0 # set upper tri to zero
    tmp.MLEobj$marss$free[[elem]] <- convert.model.mat(tmp.list.mat)$free
  }
  # will return the inits only for the estimated parameters
  pars <- MARSSvectorizeparam(tmp.MLEobj)

  if (substr(tmp.MLEobj$method, 1, 4) == "BFGS") {
    optim.method <- "BFGS"
  } else {
    optim.method <- "something wrong"
  }

  kf.function <- MLEobj$fun.kf # used for printing
  optim.output <- try(optim(pars, neglogLik, MLEobj = tmp.MLEobj, method = optim.method, lower = lower, upper = upper, control = optim.control, hessian = FALSE), silent = TRUE)

  if (inherits(optim.output, "try-error")) { # try MARSSkfss if the user did not use it
    if (MLEobj$fun.kf != "MARSSkfss") { # if user did not request MARSSkf
      cat("MARSSkfas returned error.  Trying MARSSkfss.\n")
      tmp.MLEobj$fun.kf <- "MARSSkfss"
      kf.function <- "MARSSkfss" # used for printing
      optim.output <- try(optim(pars, neglogLik, MLEobj = tmp.MLEobj, method = optim.method, lower = lower, upper = upper, control = optim.control, hessian = FALSE), silent = TRUE)
    }
  }

  # error returned
  if (inherits(optim.output, "try-error")) {
    optim.output <- list(convergence = 53, message = c("MARSSkfas and MARSSkfss tried to compute log likelihood and encountered numerical problems.\n", sep = ""))
  }

  MLEobj.return <- MLEobj
  MLEobj.return$iter.record <- optim.output$message
  #   MLEobj.return$control=MLEobj$control
  #   MLEobj.return$model=MLEobj$model
  MLEobj.return$start <- tmp.inits # set to what was used here
  MLEobj.return$convergence <- optim.output$convergence
  if (optim.output$convergence %in% c(1, 0)) {
    if ((!control$silent || control$silent == 2) && optim.output$convergence == 0) cat(paste("Success! Converged in ", optim.output$counts[1], " iterations.\n", "Function ", kf.function, " used for likelihood calculation.\n", sep = ""))
    if ((!control$silent || control$silent == 2) && optim.output$convergence == 1) cat(paste("Warning! Max iterations of ", control$maxit, " reached before convergence.\n", "Function ", kf.function, " used for likelihood calculation.\n", sep = ""))

    tmp.MLEobj <- MARSSvectorizeparam(tmp.MLEobj, optim.output$par)
    # par has the fixed and estimated values using t chol of Q and R

    # back transform Q, R and V0 if needed from chol form to usual form
    for (elem in c("Q", "R", "V0")) { # this works because by def fixed and free blocks of var-cov mats are independent
      if (!is.fixed(MODELobj$free[[elem]])) # get a new par if needed
        {
          d <- sub3D(tmp.MLEobj$marss$free[[elem]], t = 1) # this will be the one with the upper tri zero-ed out but ok since symmetric
          par.dim <- par.dims[[elem]][1:2]
          L <- unvec(tmp.MLEobj$marss$free[[elem]][, , 1] %*% tmp.MLEobj$par[[elem]], dim = par.dim) # this by def will have 0 row/col at the fixed values
          the.par <- tcrossprod(L) # L%*%t(L)
          tmp.MLEobj$par[[elem]] <- solve(crossprod(d)) %*% t(d) %*% vec(the.par)
        }
    } # end for

    pars <- MARSSvectorizeparam(tmp.MLEobj) # now the pars values have been adjusted back to normal scaling
    # now put the estimated values back into the original MLEobj; fixed and free matrices as in original
    MLEobj.return <- MARSSvectorizeparam(MLEobj.return, pars)
    kf.out <- try(MARSSkf(MLEobj.return), silent = TRUE)

    if (inherits(kf.out, "try-error")) {
      MLEobj.return$numIter <- optim.output$counts[1]
      MLEobj.return$logLik <- -1 * optim.output$value
      MLEobj.return$errors <- c(paste0("\nWARNING: optim() successfully fit the model but ", kf.function, " returned an error with the fitted model. Try MARSSinfo('optimerror54') for insight.", sep = ""), "\nError: ", kf.out[1])
      MLEobj.return$convergence <- 54
      MLEobj.return <- MARSSaic(MLEobj.return)
      kf.out <- NULL
    }
  } else {
    if (optim.output$convergence == 10) optim.output$message <- c("degeneracy of the Nelder-Mead simplex\n", paste("Function ", kf.function, " used for likelihood calculation.\n", sep = ""), optim.output$message)
    optim.output$counts <- NULL
    if (!control$silent) cat("MARSSoptim() stopped with errors. No parameter estimates returned.\n")
    if (control$silent == 2) cat("MARSSoptim() stopped with errors. No parameter estimates returned. See $errors in output for details.\n")

    MLEobj.return$par <- NULL
    MLEobj.return$errors <- optim.output$message
    kf.out <- NULL
  }

  if (!is.null(kf.out)) {
    if (control$trace > 0) MLEobj.return$kf <- kf.out
    MLEobj.return$states <- kf.out$xtT
    MLEobj.return$numIter <- optim.output$counts[1]
    MLEobj.return$logLik <- kf.out$logLik
  }
  MLEobj.return$method <- MLEobj$method

  ## Add AIC and AICc to the object
  if (!is.null(kf.out)) MLEobj.return <- MARSSaic(MLEobj.return)

  return(MLEobj.return)
}


#' brms model for estimating variant transmissibility
#'
#' @description Fit an additive mixture of a baseline and variant reproduction
#' number (Rt) to estimated non-parametric Rt values. The variant Rt is
#' described by a multiplication of the baseline Rt. The baseline Rt is
#' described by any covariate combination supported by `brms`.
#'
#' Uncertainty is included in the observed Rt estimates  by assuming that they
#' have student T distribution with standard deviation derived as a combination
#' of a modelled estimate and the observed standard error of the estimates.
#' Uncertainty is also included by modelling variant samples (`positive_prop`) as a
#' latent variable informed by the number of overall samples available and the
#' number of samples which are postive for the variant of interest.
#'
#' @param log_rt Formula describing the covariates  predicting the log of the
#' baseline variants time-varying reproduction number.
#' @param data Data frame that must contain the responses: `rt_mean`,
#' `rt_sd`, `positive_samples`, `samples` and a dummy variable `positive_prop` which
#' is `NA_real`. If wanting to not model uncertainty in the variant proportion
#' fill this variable with the point estimate of the proportion.
#' @param brm_fn Function from `brms` to apply. Common options are
#' `brm()` for model fitting,`get_prior()` to see the priors (not may not
#' be accurate in all cases), and `make_stancode()` to see the generated
#' stan code.
#' @param ... Additional arguments to pass to `brm_fn`.
#' @import brms
variant_rt <- function(log_rt = ~ 1, data, brm_fn = brm, ...) {

  # define Rt variant mixture
  rt_model <-
    bf(rt_mean | se(rt_sd, sigma = TRUE) ~ (1 + var) * exp(logRt),
       nl = TRUE) +
    lf(var ~ 0 + mi(positive_prop)) +
    lf(as.formula(paste0("logRt ", paste(log_rt, collapse = " ")))) +
    student()

  # define samples observation model
  samples_obs_model <- bf(positive_samples | trials(samples) ~ 0 + mi(positive_prop)) +
    binomial(link = "identity")

  # define model for latent variant proportion
  variant_model <- bf(positive_prop | mi() ~ 1, family = "beta")

  # set weak priors
  # set positive proportion to be latent and observed directly with no modifier
  priors <- c(
    prior(student_t(3, 0, 0.5), nlpar = "var", resp = "rtmean"),
    prior(student_t(3, 0, 0.5), nlpar = "logRt", resp = "rtmean"),
    prior(constant(1), class = "phi", resp = "positive_samples"),
    prior(constant(1), class = "Intercept", resp = "positive_samples"),
    prior(constant(1), resp = "positivesamples")
  )

# fit model
brm_out <- do.call(brm_fn, c(
  list(
    mvbf(samples_obs_model, variant_model, rt_model, rescor = FALSE),
    data = data,
    prior = priors),
  ...)
)

return(brm_out)
}


cat("insert number ")
a <- scan(nmax=1, quiet=TRUE)
cat("insert number ")
b <- scan(nmax=1, quiet=TRUE)
print(paste('a+b=', a+b))
print(paste('a-b=', a-b))
print(paste('a*b=', a*b))
print(paste('a%/%b=', a%/%b))
print(paste('a%%b=', a%%b))
print(paste('a^b=', a^b))


#' Create a bias-corrected distribution model
#'
#' This function implements a bayesLand model
#' @param x SpatialPolygonsDataFrame
#' @param effort Character, name of field in \code{x} indicating collection effort.
#' @param detect Character, name of field in \code{x} indicating number of detections of focal species.
#' @param stateProv Character, name of field with state/province.
#' @param county Character, name of field with county.
#' @param qGivenDetect Either \code{NULL} (default) or logical. If \code{TRUE} or \code{FALSE} then the model will try to estimate the rate of false detection (mistaken positive identifications). If \code{NULL}, then the model will not.  If \code{TRUE}, false detection rate \emph{q} reflects the probability of all collections in a location being of mistaken identification. If \code{FALSE} (default), then \emph{q} reflects the probability that for a county where the species does not truly occur, similar species (which could be mistaken for the focal species) do occur (and may or may not have been sampled and mis-identified).
#' @param covariate Either \code{NULL} (default) or a character or character vector. If not \code{NULL}, these are the name(s) field(s) with covariates.
#' @param niter Positive integer, number of MCMC iterations (including burn-in). The default is 2000, but this is often too low for most cases (i.e., try numbers in the range 10000, 100000, etc.).
#' @param nburnin Positive integer, number of burn-in samples (less than \code{niter}). The default is 1000, but this is often too low for most cases (i.e., try numbers in the range 10000, 100000, etc.).
#' @param nchains Positive integer, number of MCMC chains (default is 4).
#' @param thin Positive integer, number of MCMC samples by which to thin the results (default is 1; i.e., no thinning). To reduce memory requirements, you can use \code{thin} while increasing \code{niter} and/or \code{nburnin}.
#' @param na.rm Logical, if \code{TRUE} (default), then remove rows in \code{x} that have \code{NA} in any input field. If this is \code{FALSE} and \code{NA}s occur in any fields then an error may occur.
#' @param verbose Logical, if \code{TRUE} (default) display progress.
#' @param ... Arguments to pass to \code{\link[nimble]{configureMCMC}} and \code{\link[nimble]{runMCMC}}.
#' @return A list object with three elements, one with the MCMC chains, a second with the object \code{x} with model output appended to the data portion of the object, and a third with model information. The new columns in \code{x} represent:
#' \itemize{
#' 		\item{\code{cvScale}} (only if \code{covariate} is not \code{NULL}): Value of centered and scaled covariate
#' 		\item{\code{psi}} Probability of occurrence
#' 		\item{\code{psi95CI}} 95% credibility interval for probability of occurrence
#' 		\item{\code{p}} Probability of detecting the focal species assuming it is present
#' 		\item{\code{p95CI}} 95% credibility interval for probability of detection
#' }
#' @examples
#' @export

trainBayesODM <- function(
	x,
	effort,
	detect,
	stateProv,
	county,
	qGivenDetect = NULL,
	covariate = NULL,
	niter = 2000,
	nburnin = 1000,
	nchains = 4,
	thin = 1,
	na.rm=TRUE,
	verbose=TRUE,
	...
) {

	### prepare data
	################

		if (verbose) omnibus::say('Preparing data...')

		# remove missings
		if (na.rm) {
			ok <- complete.cases(x@data[ , c(effort, detect, stateProv, county, covariate)])
			x <- x[ok, ]
		}

		### CAR setup for counties
		##########################

			neighs <- spdep::poly2nb(x, queen=TRUE)

			hasNeighs <- !sapply(neighs, function(x) { length(x == 1) && (x == 0) })

			neighsList <- spdep::nb2WB(neighs)
			carAdjCounty <- neighsList$adj
			carWeightCounty <- neighsList$weights
			carNumCounty <- neighsList$num

		### inputs
		##########

			### data
			data <- list(
				y = x@data[ , detect],
				N = x@data[ , effort],
				carAdjCounty = carAdjCounty,
				carWeightCounty = carWeightCounty,
				carNumCounty = carNumCounty
			)

			# add covariates
			if (!is.null(covariate)) {
				predVals <- as.numeric(scale(x@data[ , covariate]))
				predList <- list(predVals)
				names(predList) <- 'predictor'
				data <- c(data, predList)
			}

			### constants
			constants <- list(
				numCounties = nrow(x),
				numStates = length(unique(x@data[ , stateProv])),
				state = as.numeric(as.factor(x@data[ , stateProv])),
				lengthAdjCounties = length(carAdjCounty),
				isOcc = as.numeric(data$y > 0)
			)

			### initializations
			inits <- list(

				z = as.numeric(data$y > 0),

				q = 0.01,

				p = runif(constants$numCounties, 0.1, 0.3),
				logit_p = logit(p),
				p_stateMean = rnorm(constants$numStates, 0, 0.5),

				psi = runif(constants$numCounties, 0.4, 0.6),

				psi_island = rnorm(constants$numCounties, 0, 0.1),
				psi_islandMean = -0.1,

				psi_tau = 0.1,
				psi_car = rnorm(constants$numCounties)

			)

			if (!is.null(covariate)) {

				inits <- c(
					inits,
					list(psi_beta = 0.1)
				)

			}

			### monitors
			monitors <- names(inits)
			monitors <- monitors[!(monitors %in% c('logit_p'))]
			monitors <- monitors[!(monitors %in% c('z'))]

	### model setup
	###############

		if (verbose) omnibus::say('Model setup...')

		### model specification
		if (is.null(covariate)) {

			### WITHOUT covariate
			code <- nimble::nimbleCode({

				## likelihood
				for (i in 1:numCounties) {

					# detection
					y[i] ~ dbin(pstar[i], N[i])
					pstar[i] <- psi[i] * p[i] + (1 - psi[i]) * isOcc[i] * (p[i] * q) # viable but gives large q if isOcc is {0, 1}
					logit(p[i]) ~ dnorm(p_stateMean[state[i]], tau=0.001)

					# occupancy
					logit(psi[i]) <- hasNeighs[i] * psi_car[i] + (1 - hasNeighs[i]) * psi_island[i]
					psi_island[i] ~ dnorm(psi_islandMean, tau=0.001)

				}

				## priors
				q ~ dbeta(10, 1)

				# state effects
				for (j in 1:numStates) {
					p_stateMean[j] ~ dnorm(0, tau=0.001)
				}

				psi_islandMean ~ dnorm(0, tau=0.001)

				# county occupancy CAR
				psi_tau ~ dgamma(0.001, 0.001)
				psi_car[1:numCounties] ~ dcar_normal(adj=carAdjCounty[1:lengthAdjCounties], weights=carWeightCounty[1:lengthAdjCounties], num=carNumCounty[1:numCounties], tau=psi_tau, c=3, zero_mean=0)

			})

		} else {

			### WITH covariate
			code <- nimble::nimbleCode({

				## likelihood
				for (i in 1:numCounties) {

					# detection
					y[i] ~ dbin(pstar[i], N[i])
					pstar[i] <- psi[i] * p[i] + (1 - psi[i]) * isOcc[i] * (p[i] * q) # viable but gives large q if isOcc is {0, 1}
					logit(p[i]) ~ dnorm(p_stateMean[state[i]], tau=0.001)

					# occupancy
					logit(psi[i]) <- hasNeighs[i] * psi_car[i] + (1 - hasNeighs[i]) * psi_island[i] + psi_beta * predictor[i]
					psi_island[i] ~ dnorm(psi_islandMean, tau=0.001)

				}

				## priors
				q ~ dbeta(10, 1)

				# state effects
				for (j in 1:numStates) {
					p_stateMean[j] ~ dnorm(0, tau=0.001)
				}

				psi_islandMean ~ dnorm(0, tau=0.001)

				beta ~ dnorm(0, tau=0.001)

				# county occupancy CAR
				psi_tau ~ dgamma(0.001, 0.001)
				psi_car[1:numCounties] ~ dcar_normal(adj=carAdjCounty[1:lengthAdjCounties], weights=carWeightCounty[1:lengthAdjCounties], num=carNumCounty[1:numCounties], tau=psi_tau, c=3, zero_mean=0)

			})

		} # if covariate

		### construct and compile model
		model <- nimble::nimbleModel(code=code, constants=constants, data=data, inits=inits, check=TRUE)
		flush.console()

		# conf <- nimble::configureMCMC(model, monitors = monitors, thin = thin, print = verbose, ...)
		conf <- nimble::configureMCMC(model, monitors = monitors, thin = thin, print = verbose)
		flush.console()

		## modify samplers
		conf$removeSamplers('psi_tau')
		conf$addSampler(target='psi_tau', type='slice')

		conf$removeSamplers('q')
		conf$addSampler(target='q', type='slice')

		if (!is.null(covariate)) {
			conf$removeSamplers('psi_beta')
			conf$addSampler(target='psi_beta', type='slice')
		}

		for (i in 1:constants$numCounties) {
			conf$removeSamplers(paste0('logit_p[', i, ']'))
			conf$addSampler(target=paste0('logit_p[', i, ']'), type='slice')

		}

		for (i in 1:constants$numCounties) {
			conf$removeSamplers(paste0('psi_island[', i, ']'))
			conf$addSampler(target=paste0('psi_island[', i, ']'), type='slice')

		}

		confBuild <- nimble::buildMCMC(conf)
		flush.console()
		compiled <- nimble::compileNimble(model, confBuild)
		flush.console()

	### MCMC
	########

		if (verbose) omnibus::say('Modeling...')
		# mcmc <- runMCMC(compiled$confBuild, niter = niter, nburnin = nburnin, nchains = nchains, inits = inits, progressBar = verbose, samplesAsCodaMCMC = TRUE, summary = FALSE, ...)
		mcmc <- runMCMC(compiled$confBuild, niter = niter, nburnin = nburnin, nchains = nchains, inits = inits, progressBar = verbose, samplesAsCodaMCMC = TRUE, summary = FALSE)
		flush.console()

		# rm(model, compiled, confBuild)
		# gc()

	### process model output
	########################

		omnibus::say('Processing model output...')

		### remove unneeded stuff
		#########################

			mcmc <- list(samples = mcmc)

			# remove unneeded
			ignores <- c('logit_p[[]', 'z[[]')
			for (ignore in ignores) {
				bads <- grepl(colnames(mcmc$samples$chain1), pattern=ignore)
				if (any(bads)) {
					for (chain in 1:nchains) {
						mcmc$samples[[chain]] <- mcmc$samples[[chain]][ , !bads]
					}
				}
			}

			# remove island coefficients for non-islands and CAR coefficients for islands
			if (any(!hasNeighs)) {

				# psi_car for islands
				index <- which(!hasNeighs)
				bads <- paste0('psi_car[', index, ']')
				for (chain in 1:nchains) {
					remove <- which(colnames(mcmc$samples[[chain]]) %in% bads)
					mcmc$samples[[chain]] <- mcmc$samples[[chain]][ , -remove]
				}

				# psi_island for non-islands
				index <- which(hasNeighs)
				bads <- paste0('psi_island[', index, ']')
				for (chain in 1:nchains) {
					remove <- which(colnames(mcmc$samples[[chain]]) %in% bads)
					mcmc$samples[[chain]] <- mcmc$samples[[chain]][ , -remove]
				}

			# no islands
			} else {

				# remove psi_island and psi_islandMean for all
				index <- which(hasNeighs)
				bads <- c(paste0('psi_island[', index, ']'), 'psi_islandMean')
				for (chain in 1:nchains) {
					remove <- which(colnames(mcmc$samples[[chain]]) %in% bads)
					mcmc$samples[[chain]] <- mcmc$samples[[chain]][ , -remove]
				}

			}

						# conv <- rhatStats(mcmc$samples, rhatThresh=1.1, minConv=minUnconverged)
						# print(str(conv))
						# w <- as.Bwiqid(mcmc$samples)

						# caterplot(mcmc$samples, regex='p[[]')
						# caterplot(mcmc$samples, regex='psi[[]')

						# diagPlot(w, params='p')
						# diagPlot(w, params='psi')
						# diagPlot(w, params='q')
						# diagPlot(w, params='psi_tau')

		### calculate summary
		#####################

			mcmc$summary$all.chains <- matrix(NA, nrow=ncol(mcmc$samples[[1]]), ncol=6)
			rownames(mcmc$summary$all.chains) <- colnames(mcmc$samples[[1]])
			colnames(mcmc$summary$all.chains) <- c('Mean', 'Median', 'St.Dev', '95%CI_low', '95%CI_upp', 'rhat')

			samps <- mcmc$samples[[1]]
			if (nchains > 1) {
				for (chain in 2:nchains) {
					samps <- rbind(samps, mcmc$samples[[chain]])
				}
			}

			mcmc$summary$all.chains[ , 'Mean'] <- colMeans(samps)
			mcmc$summary$all.chains[ , 'Median'] <- apply(samps, 2, median)
			mcmc$summary$all.chains[ , 'St.Dev'] <- apply(samps, 2, sd)
			mcmc$summary$all.chains[ , '95%CI_low'] <- apply(samps, 2, quantile, probs=0.025)
			mcmc$summary$all.chains[ , '95%CI_upp'] <- apply(samps, 2, quantile, probs=0.975)
			mcmc$summary$all.chains[ , 'rhat'] <- wiqid::simpleRhat(mcmc$samples, n.chains=nchains)

			rm(samps); gc()

		### update shapefile
		####################

			psi <- mcmc$summary$all.chains[grepl(rownames(mcmc$summary$all.chains), pattern='psi[[]'), 'Mean']
			p <- mcmc$summary$all.chains[grepl(rownames(mcmc$summary$all.chains), pattern='p[[]'), 'Mean']

			uncer <- mcmc$summary$all.chains[ , '95%CI_upp'] - mcmc$summary$all.chains[ , '95%CI_low']
			names(uncer) <- rownames(mcmc$summary$all.chains)
			occUncer <- uncer[grepl(names(uncer), pattern='psi[[]')]
			pUncer <- uncer[grepl(names(uncer), pattern='p[[]')]

			x@data$psi <- psi
			x@data$occUncer <- occUncer
			x@data$p <- p
			x@data$pUncer <- pUncer

			names(x@data)[(ncol(x@data) - 3):ncol(x@data)] <- c('psi', 'psi95CI', 'p', 'p95CI')

		### collate
		###########

			meta <- list(
				niter=niter,
				nburnin=nburnin,
				nchains=nchains,
				thin=thin,
				effort=effort,
				detect=detect,
				stateProv=stateProv,
				county=county,
				covariate=covariate,
				hasIslands=any(!hasNeighs),
				qGivenDetect=qGivenDetect,
				na.rm=na.rm
			)

			attr(x, 'meta') <- meta

	### output
	list(mcmc=mcmc, code=code, shape=x, meta=meta)

}


# Glossary ----

# n.b. the following employs a mix of snake_case and camelCase that is sure to
#  vex some, but represents the author's best attempt to balance to the competing
#  aims of clarity & brevity.

# Y: observed outcome
# nY: number of observed outcomes
# X: predictor/contrast matrix
# nX: number of predictors (columns in the contrast matrix)
# rX: number of rows in the contrast matrix X
# (i)ndividual: a unit of observation within which correlated measurements may take place
# (c)ondition: a labelled set of observations within an individual that share some feature/predictor or conjunction of features/predictors
# Xc: condition-level contrast matrix
# nXc: number of predictors in the condition-level contrast matrix
# rXc: number of rows in the condition-level contrast matrix
# yXc: for each observation in y, an index indicating the associated row in Xc corresponding to that observation's individual/condition combo
# (g)roup: a collection of individuals that share some feature/predictor
# Xg: group-level contrast matrix
# nXg: number of predictors in the group-level contrast matrix
# rXg: number of rows in the group-level contrast matrix
# Z: matrix of coefficient row-vectors to be dot-product'd with a contrast matrix
# iZc: matrix of coefficient row-vectors associated with each individual

# Preamble (options, installs, imports & custom functions) ----

options(warn=1) # really should be default in R
`%!in%` = Negate(`%in%`) # should be in base R!

# specify the packages used:
required_packages = c(
	'github.com/rmcelreath/rethinking' # for rlkjcorr & rmvrnom2
	, 'github.com/stan-dev/cmdstanr' # for Stan stuff
	, 'github.com/mike-lawrence/aria/aria' # for aria
	, 'tidyverse' # for all that is good and holy
)

# load the helper functions:
for(file in fs::dir_ls('r')){
	cat('Loading function: ',fs::path_ext_remove(fs::path_file(file)),'()\n',sep='')
	source(file)
}

# install any required packages not already present
install_if_missing(required_packages)

# load tidyverse & aria
library(tidyverse)
library(aria)

# Simulate data ----
set.seed(1) # change this to make different data

# setting the data simulation parameters

# parameters you can play with
nG_vars = 1 # number of 2-level variables manipulated as crossed and across individuals, must be an integer >0
nC_vars = 2 # number of 2-level variables manipulated as crossed and within each individual, must be an integer >0
nI_per_group = 3 # number of individuals, must be an integer >1
nY_per_ic = 4 # number of observations per individual/condition combo, must be an integer >1
# the latter two combine to determine whether centered or non-centered will sample better

# the rest of these you shouldn't touch
nXc = 2^(nC_vars)
nXg = 2^(nG_vars)
Z = matrix(rnorm(nXg*nXc),nrow=nXg,ncol=nXc)
iZc_sd = rweibull(nXc,2,1)
iZc_r_mat = rethinking::rlkjcorr(1,nXc,eta=1)
Y_sd = rweibull(1,2,1)
iZc_r_vec = iZc_r_mat[lower.tri(iZc_r_mat)]

# compute Xg
Xg = sim_contrasts(nG_vars,'G')

# compute gZc
gZc = matrix(NA,nrow=nXg,ncol=nXc)
for(this_nXc in 1:nXc){
	for(this_nXg in 1:nXg){
		gZc[this_nXg,this_nXc] = Xg[this_nXg,] %*% Z[,this_nXc]
	}
}

# compute iZc
(
	1:nXg
	%>% map_dfr(
		.f = function(this_nXg){(
			rethinking::rmvnorm2(
				n = nI_per_group
				, Mu = gZc[this_nXg,]
				, sigma = iZc_sd
				, Rho = iZc_r_mat
			)
			%>% as_tibble(
				.name_repair = function(x) paste0('iZc[.,',1:length(x),']')
			)
			%>% mutate(
				g = this_nXg
				, individual = paste(g,1:n(),sep='_') # important that individuals have unique identifiers across groups
			)
		)}
	)
) ->
	iZc

# compute Xc
uXc = sim_contrasts(nC_vars,'C')

# compute iZc_dot_Xc
(
	iZc
	%>% group_by(g,individual)
	%>% summarise(
		{function(x){
			x = t(as.matrix(select(x,starts_with('iZc')),))
			this_dot = rep(NA,nXc)
			for(this_nXc in 1:nXc){
				this_dot[this_nXc] = uXc[this_nXc,] %*% x
			}
			tibble(c=1:nXc,value=this_dot)
		}}(cur_data())
		, .groups = 'keep'
	)
) ->
	iZc_dot_Xc

# compute Y
(
	iZc_dot_Xc
	%>% group_by(c,.add=T)
	%>% summarise(
		value = rnorm(nY_per_ic,value,Y_sd)
		, .groups = 'keep'
	)
) ->
	Y

# join with group & condition columns
(
	Y
	%>% left_join((
		attr(Xg,'data')
		%>% mutate(g=1:n())
	))
	%>% left_join((
		attr(uXc,'data')
		%>% mutate(c=1:n())
	))
	%>% ungroup()
	%>% select(everything(),-g,-c,-value,value)
) ->
	dat

# dat is now what one would typically have as collected data
dat

# if there's no missing data (as in this synthetic example), we could proceed with modelling,
#  but to demonstrate how to handle missing data, we'll do a bit of extra work:

# Uncomment next section to induce data-missingness (causes some Ss to be missing whole conditions)
# (
# 	dat
# 	# select all but value
# 	%>% select(-value)
# 	# get distinct rows (no more individual trials)
# 	%>% distinct()
# 	# group by individual
# 	%>% group_by(individual)
# 	# toss one condition
# 	%>% slice_sample(n=(nXc-1))
# 	# semi-join to keep all in dat that have been kept above
# 	%>% semi_join(
# 		x = dat
# 		, y = .
# 	)
# ) ->
# 	dat

nrow(dat)

# Compute inputs to Stan model ----

# Some data prep
(
	dat
	# ungroup (necessary)
	%>% ungroup()
	# ensure individual is a sequential numeric
	%>% mutate(
		individual = as.numeric(factor(individual))
	)
	# arrange rows by individual
	%>% arrange(individual)
) ->
	dat

# compute Xg from distinct combinations of groups
(
	dat
	# select down to any G vars
	%>% select(starts_with('G'))
	# collapse to distinct set of rows (combinations of grouping variables)
	%>% distinct()
	# arrange (not really necessary, but why not)
	%>% arrange()
	# add the contrast matrix columns
	%>% mutate(
		contrasts = get_contrast_matrix_rows_as_list(
			data = .
			# the following complicated specification of the formula is a by-product of making this example
			# work for any number of variables =  normally you would do something like this
			# (for 2 variables for example):
			# formula = ~ G1*G2
			, formula = as.formula(paste0('~',paste0('G',1:nG_vars,collapse='*')))
			# half-sum contrasts are nice for 2-level variables bc they yield parameters whose value
			# is the difference between conditions
			, contrast_kind = halfsum_contrasts
		)
	)
) ->
	Xg_with_vars

# show contrasts
(
	Xg_with_vars
	%>% unnest(contrasts)
)

# join Xg with dat to label individuals with corresponding row from Xg
(
	Xg_with_vars
	# add row identifier
	%>% mutate(Xg_row=1:n())
	# join with dat, collapsed to 1-row per individual with their group info
	%>% right_join((# right-join to apply the row order from dat
		dat
		%>% select(individual,starts_with('G'))
		%>% distinct()

	))
	%>% arrange(individual)
	%>% pull(Xg_row)
) ->
	iXg

# compute complete_Xc from distinct combinations of individuals & conditions
(
	dat
	# select individual & any condition-defining columns
	%>% select(individual,starts_with('C'))
	# collapse down to distinct rows (1 per individual/conditions combo)
	%>% distinct()
	# expand (in case there's any missing data; doesn't hurt if not)
	# %>% exec(expand_grid,!!!.) #should work, but doesn't, so need the next three instead
	%>% as.list()
	%>% map(unique)
	%>% cross_df()
	# arrange (not really necessary, but why not)
	%>% arrange()
	# add the contrast matrix columns
	%>% mutate(
		contrasts = get_contrast_matrix_rows_as_list(
			data = .
			# the following complicated specification of the formula is a by-product of making this example
			# work for any number of variables =  normally you would do something like this
			# (for 2 variables for example):
			# formula = ~ C1*C2
			, formula = as.formula(paste0('~',paste0('C',1:nC_vars,collapse='*')))
			# half-sum contrasts are nice for 2-level variables bc they yield parameters whose value
			# is the difference between conditions
			, contrast_kind = halfsum_contrasts
		)
	)
) ->
	complete_Xc_with_vars

# show the unique contrasts
(
	complete_Xc_with_vars
	%>% select(-individual)
	%>% distinct()
	%>% unnest(contrasts)
)

# subset down to just those individual-condition combos actually present in the data
#  it's ok if there's no missing data and nrow(complete_Xc_with_vars)==nrow(Xc_with_vars)
(
	complete_Xc_with_vars
	%>% semi_join(dat)
) ->
	Xc_with_vars

# join Xc with dat to label observations with corresponding row from Xc
(
	Xc_with_vars
	# add row identifier
	%>% mutate(Xc_row=1:n())
	# right-join with dat to preserve dat's row order
	%>% right_join(
		mutate(dat,dat_row=1:n())
	)
	%>% arrange(dat_row)
	# pull the Xc row identifier
	%>% pull(Xc_row)
) ->
	yXc

# package for stan & sample ----

data_for_stan = lst( # lst permits later entries to refer to earlier entries

	# # # #
	# Entries we need to specify ourselves
	# # # #

	# Xg: group-level predictor matrix
	Xg = (
		Xg_with_vars
		%>% select(contrasts)
		%>% unnest(contrasts)
		%>% as.matrix()
	)

	# iXg: which group each individual is associated with
	, iXg = iXg

	# Xc: condition-level predictor matrix
	, Xc = (
		Xc_with_vars
		%>% select(contrasts)
		%>% unnest(contrasts)
		%>% as.matrix()
	)

	# iXc: which individual is associated with each row in Xc
	, iXc = as.numeric(factor(Xc_with_vars$individual))

	# Y: observations
	, Y = dat$value

	# yXc: which row in Xc is associated with each observation in Y
	, yXc = yXc

	# # # #
	# Entries computable from the above
	# # # #

	# nXg: number of cols in the group-level predictor matrix
	, nXg = ncol(Xg)

	# rXg: number of rows in the group-level predictor matrix
	, rXg = nrow(Xg)

	# nI: number of individuals
	, nI = max(iXc)

	# nXc: number of cols in the condition-level predictor matrix
	, nXc = ncol(Xc)

	# rXc: number of rows in the condition-level predictor matrix
	, rXc = nrow(Xc)

	# nY: num entries in the observation vectors
	, nY = length(Y)

)

# double-check:
glimpse(data_for_stan)

# set the model path (automated bc in this repo there's only one)
mod_path = fs::dir_ls(
	path = 'stan'
	, glob = '*.stan'
)

# set the model centered/non-centeredness
#  generally, if *either* nI_per_group *or* num_Y_per_ic is small, non-centered will sample better than centered
data_for_stan$centered = F

# conversion to 1/0 for stan
data_for_stan$centered = as.numeric(data_for_stan$centered)

# set the posterior path (automated but you could do your own if you had multiple models)
(
	mod_path
	%>% fs::path_file()
	%>% fs::path_ext_remove()
	%>% paste0(
		ifelse(data_for_stan$centered,'_c','_nc')
	)
	%>% fs::path(
		'posteriors'
		, .
		, ext = 'netcdf4'
	)
) -> post_path

# ensure model is compiled
aria:::check_and_compile(mod_path,block=T)

# compose
aria::compose(
	data = data_for_stan
	, code_path = mod_path
	, out_path = post_path
	, overwrite = T
	, block = T
)

# how long did it take?
aria::marginalia()$time

# check posterior diagnostics ----
post = aria::coda(post_path)

# Check treedepth, divergences, & rebfmi
(
	post$draws(group='sample_stats')
	%>% posterior::as_draws_df()
	%>% group_by(.chain)
	%>% summarise(
		max_treedepth = max(treedepth)
		, num_divergent = sum(divergent)
		, rebfmi = var(energy)/(sum(diff(energy)^2)/n()) # n.b. reciprocal of typical EBFMI, so bigger=bad, like rhat
	)
)

# gather summary for core parameters (inc. r̂ & ess)
(
	post$draws(group='parameters')
	%>% posterior::summarise_draws(.cores=parallel::detectCores())
) ->
	par_summary

# show the ranges of r̂/ess's
(
	par_summary
	%>% select(rhat,contains('ess'))
	%>% summary()
)

# View those with suspect r̂
(
	par_summary
	%>% filter(rhat>1.01)
	%>% (function(suspects){
		if(nrow(suspects)>1){
			View(suspects)
		}
		return(paste('# suspect parameters:',nrow(suspects)))
	})()
)

# Viz recovery of non-correlation parameters ----
(
	post$draws(variables=c('Y_sd','Z','iZc_sd'))
	%>% posterior::as_draws_df()
	%>% select(-.draw)
	%>% pivot_longer(
		cols = -c(.chain,.iteration)
		, names_to = 'variable'
	)
	%>% group_by(variable)
	%>% arrange(variable,.chain,.iteration)
	%>% summarise(
		rhat = 1.01<posterior::rhat(matrix(value,ncol=length(unique(.chain))))
		, ess_bulk = 100>posterior::ess_bulk(matrix(value,ncol=length(unique(.chain))))
		, ess_tail = 100>posterior::ess_tail(matrix(value,ncol=length(unique(.chain))))
		, as_tibble(t(posterior::quantile2(value,c(.05,.25,.5,.75,.95))))
	)
	#add true values
	%>% left_join(
		bind_rows(
			tibble(
				true = Y_sd
				, variable = 'Y_sd'
			)
			, tibble(
				true = as.vector(Z)
				, variable = paste0(
					'Z['
					,rep(1:nXg,times=nXc)
					,','
					,rep(1:nXc,each=nXg)
					,']'
				)
			)
			, tibble(
				true = iZc_sd
				, variable = paste0('iZc_sd[',1:length(true),']')
			)
		)
	)
	%>% ggplot()
	+ geom_hline(yintercept = 0)
	+ geom_linerange(
		mapping = aes(
			x = variable
			, ymin = q5
			, ymax = q95
			, colour = ess_tail
		)
		, alpha = .5
	)
	+ geom_linerange(
		mapping = aes(
			x = variable
			, ymin = q25
			, ymax = q75
			, colour = ess_bulk
		)
		, size = 3
		, alpha = .5
	)
	+ geom_point(
		mapping = aes(
			x = variable
			, y = q50
			, fill = rhat
		)
		, shape = 21
		, size = 3
	)
	+ geom_point(
		mapping = aes(
			x = variable
			, y = true
		)
		, size = 4
		, shape = 4
		, colour = 'blue'
	)
	+ coord_flip()
	+ scale_color_manual(
		values = lst(`TRUE`='red',`FALSE`='black')
		, labels = lst(`TRUE`='<100',`FALSE`='>=100')
	)
	+ scale_fill_manual(
		values = lst(`TRUE`='red',`FALSE`='white')
		, labels = lst(`TRUE`='>1.01',`FALSE`='<=1.01')
	)
	+ labs(
		y = 'True & Posterior Value'
		, x = 'Variable'
		, colour = 'ESS'
		, fill = 'Rhat'
	)
)

# Viz recovery of correlations ----
(
	post$draws(variables='iZc_r_vec')
	%>% posterior::as_draws_df()
	%>% select(-.draw)
	%>% pivot_longer(
		cols = -c(.chain,.iteration)
		, names_to = 'variable'
	)
	%>% group_by(variable)
	%>% arrange(variable,.chain,.iteration)
	%>% summarise(
		rhat = 1.01<posterior::rhat(matrix(value,ncol=length(unique(.chain))))
		, ess_bulk = 100>posterior::ess_bulk(matrix(value,ncol=length(unique(.chain))))
		, ess_tail = 100>posterior::ess_tail(matrix(value,ncol=length(unique(.chain))))
		, as_tibble(t(posterior::quantile2(value,c(.05,.25,.5,.75,.95))))
	)
	#add true values
	%>% left_join(
		tibble(
			true = iZc_r_vec
			, variable = case_when(
				length(true)==1 ~ 'r'
				, T ~ paste0('iZc_r_vec[',1:length(true),']')
			)
		)
	)
	%>% mutate(variable = factor_1d(variable))
	%>% ggplot()
	+ geom_hline(yintercept = 0)
	+ geom_linerange(
		mapping = aes(
			x = variable
			, ymin = q5
			, ymax = q95
			, colour = ess_tail
		)
		, alpha = .5
	)
	+ geom_linerange(
		mapping = aes(
			x = variable
			, ymin = q25
			, ymax = q75
			, colour = ess_bulk
		)
		, size = 3
		, alpha = .5
	)
	+ geom_point(
		mapping = aes(
			x = variable
			, y = q50
			, fill = rhat
		)
		, shape = 21
		, size = 3
	)
	+ geom_point(
		mapping = aes(
			x = variable
			, y = true
		)
		, size = 4
		, shape = 4
		, colour = 'blue'
	)
	+ coord_flip()
	+ scale_color_manual(
		values = lst(`TRUE`='red',`FALSE`='black')
		, labels = lst(`TRUE`='<100',`FALSE`='>=100')
	)
	+ scale_fill_manual(
		values = lst(`TRUE`='red',`FALSE`='white')
		, labels = lst(`TRUE`='>1.01',`FALSE`='<=1.01')
	)
	+ labs(
		y = 'True & Posterior Value'
		, x = 'Variable'
		, colour = 'ESS'
		, fill = 'Rhat'
	)
)

# fit with brms to compare sampling speed ----
library(brms)
brms_data = mutate(dat,id = factor(individual))
brms_formula = bf(
	formula = value ~ G1*G2*C1*C2 + ( C1*C2 | id )
	, family = gaussian()
)
brms_prior = c(
	prior(normal(0, 1), class = 'Intercept')
	, prior(normal(0, 1), class = 'b')
	, prior(weibull(2, 1), class = 'sigma')
	, prior(lkj(1), class = 'cor')
	, prior(normal(0, 1), class = 'sd')
)
brms_standata = make_standata(
	data = brms_data
	, formula = brms_formula
)
attr(brms_standata,'class') = NULL
brms_stancode = make_stancode(
	data = brms_data
	, formula = brms_formula
	, prior = brms_prior
	, normalize = F
)
brms_code_path = 'brms/brms_hmg.stan'
# write( brms_stancode, file = brms_code_path ) #commented-out bc after initial write I edited the file a bit

# ensure model is compiled
aria:::check_and_compile(
	brms_code_path
	, compile_debug = F
	, run_debug = F
	, block = T
)

# compose
brms_post_path = 'posteriors/brms.netcdf4'
aria::compose(
	data = brms_standata
	, code_path = brms_code_path
	, out_path = brms_post_path
	, overwrite = T
	, block = T
)

# how long did it take?
aria::marginalia()$time

# check posterior diagnostics ----
post = aria::coda(brms_post_path)

# Check treedepth, divergences, & rebfmi
(
	post$draws(group='sample_stats')
	%>% posterior::as_draws_df()
	%>% group_by(.chain)
	%>% summarise(
		max_treedepth = max(treedepth)
		, num_divergent = sum(divergent)
		, rebfmi = var(energy)/(sum(diff(energy)^2)/n()) # n.b. reciprocal of typical EBFMI, so bigger=bad, like rhat
	)
)

# gather summary for core parameters (inc. r̂ & ess)
(
	post$draws(group='parameters')
	%>% posterior::summarise_draws(.cores=parallel::detectCores())
) ->
	par_summary

# show the ranges of r̂/ess's
(
	par_summary
	%>% select(rhat,contains('ess'))
	%>% summary()
)

# View those with suspect r̂
(
	par_summary
	%>% filter(rhat>1.01)
	%>% (function(suspects){
		if(nrow(suspects)>1){
			View(suspects)
		}
		return(paste('# suspect parameters:',nrow(suspects)))
	})()
)


### Function to simulate traits with causative markers for GWAS

# need genomic data . A sample of 22k random SNPs from the Arabidopsis 1001G Projekt is available as example for Arabidopsis traits
# a file with the Information of all 2029 accessions for which genotype data exist is also provided (lat_long_2029.csv). This information has been downloaded from the AraPheno database

#A<-read.csv('~/git/GWAS/data/lat_long_2029.csv')
#If you use own data ensure that the id should be in a column called 'accession_id'
#load('~/git/GWAS/data/X_Mix.rda')
#rownames(X) needs to be the ids of the accessions
# for models with multiple causative markers, it is recommended to have them orderd by the amount of variance explained (ve)

# the simulation create random, but reproducible (seed) phenotypes.
# to simulate only in specific accession provide an integer of accessions ids in sp_acc (length need to be >1).
# to siumulate distinct markers, the script can be modified manually.

sim_Y<-function(n=100,acc=A,sp_acc=0,no_acc=200,fix_acc=TRUE,SNPs=X,no_snps=1,ve=.1,mac=5,h2=0.7,seed=42) {
  stopifnot(no_snps>0)
  stopifnot(length(ve)==no_snps)
  set.seed(seed)

  Sim<-list()
  Caus<-list()
if (length(sp_acc)>1){
  a<-sp_acc
  no_acc=length(a)
}else {
a<-sample(acc$accession_id,no_acc) }

X_<-SNPs[rownames(SNPs)%in%a,]

af<-apply(X_,2,sum)
X_ok<-X_[,which(af>mac&af<(no_acc-mac))]

u<-1
# set the seed
set.seed(seed+u)

caus<-X_ok[,sample(1:ncol(X_ok),(no_snps+1))]

#generating polygenic background

X3<-X_ok[,!colnames(X_ok)%in%colnames(caus)]

back<-X3[,sample(1:ncol(X3),1000)]

betas<-rnorm(1000,mean=0,sd=0.1)
first<- back %*% betas
### adding genetic background to data
sim<-data.frame(ecot_id=as.integer(rownames(back)),value=first)
### set heritability
dat<-var(sim[,2])

h_2<-dat/h2-dat
fix1<-rnorm(nrow(back),0,sqrt(h_2))
sim_<-data.frame(ecot_id=as.integer(rownames(back)),value=first+fix1)

for ( t in 1:length(ve)) {
beta<-sqrt((ve[t]/(1-ve[t]))*(var(sim_[,2])/var(caus[,t])))

cand<-beta*caus[,t]
sim_$value<-sim_$value+cand

}
Sim[[u]]<-sim_
Caus[[u]]<-colnames(caus)[1:t]

if (fix_acc==FALSE) {
  for ( u in 2:n ) {
    set.seed(seed+u)
  a<-sample(A$accession_id,no_acc)
  X_<-subset(X,rownames(X)%in%a)

  af<-apply(X_,2,sum)
  X_ok<-X_[,which(af>mac&af<(no_acc-mac))]

  caus<-X_ok[,sample(1:ncol(X_ok),(no_snps+1))]

  #generating polygenic background

  X3<-X_ok[,!colnames(X_ok)%in%colnames(caus)]

  back<-X3[,sample(1:ncol(X3),1000)]

  betas<-rnorm(1000,mean=0,sd=0.1)
  first<- back %*% betas
  ### adding genetic background to data
  sim<-data.frame(ecot_id=as.integer(rownames(back)),value=first)
  ### set heritability
  dat<-var(sim[,2])

  h_2<-dat/h2-dat
  fix1<-rnorm(nrow(back),0,sqrt(h_2))
  sim_<-data.frame(ecot_id=as.integer(rownames(back)),value=first+fix1)

  for ( t in 1:length(ve)) {
    beta<-sqrt((ve[t]/(1-ve[t]))*(var(sim_[,2])/var(caus[,t])))

    cand<-beta*caus[,t]
    sim_$value<-sim_$value+cand

  }
  Sim[[u]]<-sim_
  Caus[[u]]<-colnames(caus)[1:t]
}
  }else {
  for ( u in 2:n ) {
    set.seed(seed+u)

  caus<-X_ok[,sample(1:ncol(X_ok),(no_snps+1))]

  #generating polygenic background

  X3<-X_ok[,!colnames(X_ok)%in%colnames(caus)]

  back<-X3[,sample(1:ncol(X3),1000)]

  betas<-rnorm(1000,mean=0,sd=0.1)
  first<- back %*% betas
  ### adding genetic background to data
  sim<-data.frame(ecot_id=as.integer(rownames(back)),value=first)
  ### set heritability
  dat<-var(sim[,2])

  h_2<-dat/h2-dat
  fix1<-rnorm(nrow(back),0,sqrt(h_2))
  sim_<-data.frame(ecot_id=as.integer(rownames(back)),value=first+fix1)

  for ( t in 1:length(ve)) {
    beta<-sqrt((ve[t]/(1-ve[t]))*(var(sim_[,2])/var(caus[,t])))

    cand<-beta*caus[,t]
    sim_$value<-sim_$value+cand

  }
  Sim[[u]]<-sim_
  Caus[[u]]<-colnames(caus)[1:t]
}
}

return(list(Y=Sim,Caus=Caus))
}


#' ms2knots
#'
#' Conversion speed meter per second in knot per second.
#'
#' @param numeric   ms Speed in meter per second [m/s].
#' @return
#'
#'
#' @author  Istituto di Biometeorologia Firenze Italy  Alfonso Crisci \email{a.crisci@@ibimet.cnr.it}
#' @keywords  ms2knots
#'
#' @export
#'
#'
#'
#'

ms2knots<-function(ms)
{
  return(ms * 1.94384);
}


rm(list = ls())
gc()
set.seed(1954)

# Adjust to your settings.
.libPaths("~/Rlib/")
setwd("~/Code/ising_model/scripts/")

library(ggplot2)
library(posterior)
source("tools.r")
source("algorithms.r")
source("algorithms_potts.r")

## Tuning parameters for experiments
# For Potts model with q = 4.
n_states <- 4
# beta_range <- c(0.25, 0.45, 0.5493062, 0.65, 0.85)  # grid graph
# beta_range <- c(1.35, 1.55, 1.647918, 1.75, 1.95)  # complete graph (q = 4)
# beta_range <- c(0.5, 0.8, 1, 1.2, 2)  # gaussian graph
# beta_range <- c(0.25, 0.44, 0.63)
# beta_range <- c(0.5, 1, 1.5) # complete graph for q = 2
beta_range <- rep(1, 14)  # hopfield model

n_beta <- length(beta_range)
# Options for graph types: "complete", "grid", "gaussian", "hopfield"
# m_hopfield parameter is only relevant for Hopfield model
m_hopfield <- c(1, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 150, 200, 256)
  # c(1) # c(150, 200) # c(1, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100)
graph_types <- rep("hopfield", n_beta)
anti_corr <- FALSE
graph_sizes <- rep(256, n_beta)
anti_corr <- rep(FALSE, n_beta)

n_iter <- 1e4
n_chains <- 4  # need multiple chains for diagnostics
total_sample <- n_iter * n_chains

# For Gaussian graphs, Wolff and SW cannot be used.
# algo_names <- c("AG_potts", "HB", "Wolff", "SW", "AG_potts_lowrank")  # complete
# algo_names <- c("AG_potts", "HB")  # Gaussian
algo_names <- c("AG_potts", "HB", "AG_potts_lowrank")
# algo_names <- c("AG_potts", "HB", "Wolff", "SW", "AG_potts_lowrank")
n_algorithms <- length(algo_names)

# Number of generated quantities, based on samples
# (for now only do one)
n_gen_quant <- 2
n_gen_sample <- n_iter / 2

burn_in <- 1e3

length_graph_size <- length(graph_sizes)
if (FALSE) {
  beta_range[1] <- 1.60
  length_graph_size <- 1  # only run one loop.
}

for (k in 1:length_graph_size) {
  graph <- graph_types[k]
  n_part <- graph_sizes[k]
  beta <- beta_range[k]
  print(paste0("graph size: ", n_part, ", type: ", graph_types[k],
               ", beta: ", beta_range[k], ", anti_corr: ", anti_corr[k]))

  # Build adjacency graph
  A <- adjacency_graph(n_part, type = graph, anti_corr = anti_corr[k],
                       m_hopfield = m_hopfield[k])

  init <- matrix(NA, nrow = n_chains, ncol = n_part)
  for (c in 1:n_chains) init[c, ] <- sample(1:n_states, n_part, replace = T)

  samples <- array(NA, c(n_iter, n_part, n_algorithms, n_chains))
  time <- rep(NA, n_algorithms)

  print("Doing AG sampling")
  time[1] <- system.time(
    for (c in 1:n_chains) {
      samples[, , 1, c] <- ag_potts_simple(A = A, beta = beta, init = init[c, ],
                                           n_iter = n_iter, n_states = n_states)
    })[3]

  print("Doing HB sampling")
  time[2] <- system.time(
    for (c in 1:n_chains) {
      samples[, , 2, c] <- hb_sampler(A = A, beta = beta, init = init[c, ],
                                      n_iter = n_iter, is_potts = TRUE,
                                      n_states = n_states,
                                      sub_sample = ceiling(n_part / 10))
    })[3]

  if (graph != "gaussian" & graph != "hopfield") {
    print("Doing Wolff sampling")
    time[3] <- system.time(
      for (c in 1:n_chains) {
        samples[, , 3, c] <- wolff_sampler(A = A, beta = beta, init = init[c, ],
                                           n_iter = n_iter, is_potts = TRUE,
                                           n_states = n_states)
      })[3]

    print("Doing SW sampling")
    time[4] <- system.time(
      for (c in 1:n_chains) {
        samples[, , 4, c] <- sw_sampler(A = A, beta = beta, init = init[c, ],
                                        n_iter = n_iter, is_potts = TRUE,
                                        n_states = n_states)
      })[3]
  }

  if (graph == "complete" || graph == "hopfield") {
    print("Doing AG lowrank sampling")
    if (graph == "complete") index = 5
    if (graph == "hopfield") index = 3

    # REMARK: default epsilon = 1e-12 doesn't remove approximatively 0 eigenvalues.
    time[index] <- system.time(
      for (c in 1:n_chains) {
        samples[, , index, c] <- ag_potts_lowrank(A = A, beta = beta, init = init[c, ],
                                              n_states = n_states, n_iter = n_iter,
                                              alpha_modif = 0,
                                              epsilon = 1e-5)
      })[3]
  }

  all_samples <- array(NA, dim = c(total_sample, n_part, n_algorithms))
  for (c in 1:n_chains) {
    iteration_index <- ((c - 1) * n_iter + 1):(c * n_iter)
    all_samples[iteration_index, , ] <- samples[, , , c]
  }

  gen_quant <- array(NA, c(n_gen_sample, n_algorithms, n_chains, n_gen_quant))
  index_start <- n_iter - n_gen_sample + 1
  print("Computing generated quantities.")
  for (j in 1:n_algorithms) {
    for (c in 1:n_chains) {
      for (i in 1:n_gen_sample) {
        gen_quant[i, j, c, 1] <- potts_log_kernel(beta, A,
                                  samples[index_start - 1 + i, , j, c])
        gen_quant[i, j, c, 2] <- gen_quant[i, j, c, 1]
      }
    }
  }

  # Compute the effective samples size and results
  mean <- array(NA, dim = c(n_gen_quant, n_algorithms))
  mcse <- array(NA, dim = c(n_gen_quant, n_algorithms))
  rhat <- array(NA, dim = c(n_gen_quant, n_algorithms))
  ess <- array(NA, dim = c(n_gen_quant, n_algorithms))
  eff <- array(NA, dim = c(n_gen_quant, n_algorithms))

  for (j in 1:n_algorithms) {
    draws_formatted <- as_draws_df(gen_quant[, j, , ])
    names(draws_formatted)[1] <- c("log_kernel")
    summary <- summarize_draws(draws_formatted,
                               measure = c("mean", "mcse_mean", "ess_bulk", "rhat"))
    names(summary) <- c("variable", "mean", "mcse_mean", "ess_bulk", "rhat")
    mean[, j] <- summary$mean
    mcse[, j] <- summary$mcse_mean
    rhat[, j] <- summary$rhat
    ess[, j] <- summary$ess_bulk   # / total_sample * 1e3
    eff[, j] <- ess[, j] / time[j]   # (time[j] * 1e3 / total_sample)
    # Note: returns ESS / time, and corrects for the fact ess is the number
    # of effectively independent samples per 1000 iterations.
  }

  eff_data <- data.frame(cbind(mean, mcse, rhat, ess, eff))
  names(eff_data) <- c(paste0("mean_", algo_names), paste0("mcse_", algo_names),
                       paste0("rhat_", algo_names),
                       paste0("ess_", algo_names), paste0("eff_", algo_names))

  output_name <- paste0("deliv/measurements/potts/", n_part, graph,
                        "_", beta, "_", n_states, "_", m_hopfield[k], "_efficiency")
  write.csv(eff_data, paste0(output_name, ".csv"))
}


#!/usr/bin/Rscript

# friedman.post.hoc.simple(t(amg2$mean))$ranks
# friedman.post.hoc.simple(t(r$succRatio[1:4,]),minFlag=F)
# friedman.post.hoc.simple(t(r$mean[1:4,]))

#"
#Useful commands:
#source("friedman.r"); r <- friedman.post.hoc.meth.bench("intsp.csv", minFlag=FALSE)
#source("friedman.r"); r <- friedman.post.hoc.meth.bench("strp.csv", minFlag=FALSE)
#sort(r$ranks)
#r$p.value
#r$cmp.matrix
#r$cmp.p.values
#r$cmp.p.values
#r$p.value
#r$cmp.matrix
#"

friedman.post.hoc.meth.bench <- function(fname, minFlag=TRUE)
{
# KK
# reads a csv file of the form: methods x benchmarks
# transforms a table into format expected by friedman.post and calls friedman.post
# set minFlag to true if the variable is minimized, otherwise to false

    options(width=10000) # prevent the wrapping of output
    #options(width=300)

    options(digits=2)
	  b <- read.csv(file=fname, header=TRUE, sep=";", row.names=1, check.names=FALSE)
    b <- b[rownames(b) != "All", ]
    b <- t(b)
	print(summary(b))
	r <- friedman.post.hoc.simple(b,minFlag)
    print(sort(r$ranks))
    r #sort(r.ranks)
}

## Does Friedman from dataframe
friedman.post.hoc.simple <- function(df, minFlag=T,...)  {
	res <- data.frame(var=rep(NA, 0), alg=rep("", 0),
		block=rep("", 0), stringsAsFactors=FALSE, check.names=FALSE)
    i = 1
   print(df)
	for(r in 1:nrow(df))
		for(c in 1:ncol(df))
		{
			res[i,1] = df[r,c]
			res[i,2] = row.names(df)[r]
			res[i,3] = colnames(df)[c]
			i=i+1
		}
	friedman.post.hoc (var ~ alg | block, data=res,minimize=minFlag, ...)
}

friedman.post.hoc <- function(formu, data, minimize=TRUE, p.signif=0.05, post.signif=p.signif, ctrl=NULL, method=NULL)
{
	#
	# formu is a formula of the shape: 	Y ~ X | block   i.e.:  Value ~ Algorithm | Problemf
	# data is a long data.frame with three columns:	[[ Y (numeric), X (factor), block (factor) ]]

	# get the names out of the formula
	formu.names <- all.vars(formu)
	Y.name <- formu.names[1]
	X.name <- formu.names[2]
	block.name <- formu.names[3]

	X <- factor(data[,X.name ])
	Y <- data[,Y.name]
	block <- factor(data[,block.name ])

	n <- nlevels(block)
	k <- nlevels(X)

	if(sum(is.na(Y)) > 0) stop("Function stopped: This function doesn't handle NA's. In case of NA in Y in one of the blocks, then that entire block should be removed.")
	if(k == 2) { warning(paste("'",X.name,"'", "has only two levels. Consider using paired wilcox.test instead of friedman test"))}

	# convert data.frame to matrix
	values <- matrix(nrow=n, ncol=k, dimnames=list(levels(block), levels(X)))
	values[cbind(block, X)] <- Y

	ranks.in.blocks <- t(apply(if (minimize) values else -values, 1, rank))
	rmean <- mean(ranks.in.blocks)
	ranks <- apply(ranks.in.blocks, 2, mean)

	Ff <- n*(k-1)*(n*sum((ranks-rmean)^2))/(sum((ranks.in.blocks-rmean)^2))
	p.value <- pchisq(Ff, k-1, lower.tail=F)

	if (p.value < p.signif)
	{
		if (is.null(ctrl)) {
			# NxN
			#stop("NxN posthoc not implemented yet")
			method <- match.arg(method, c('shaffer', p.adjust.methods))

			#z <- outer(1:k,1:k, function(a,b) abs(ranks[a]-ranks[b]) / sqrt((k*(k+1))/(6*n)))
			rowidx <- unlist(lapply(2:k, function(x) x:k))
			colidx <- rep(1:(k-1), (k-1):1)
			z <- abs(ranks[rowidx] - ranks[colidx]) / sqrt((k*(k+1))/(6*n))
			p <- 2*pnorm(-z)

			cmp.p.values <- switch(method,
					shaffer = {
						ShafferH <- function(k) {
							if (k <= 1) return(0)
							S <- vector("list", k)
							S[[1]] <- 0

							for (i in 2:k) {
								res <- choose(i, 2)
								for (j in (i-1):1) res <- union(res, choose(j, 2) + S[[i-j]])
								S[[i]] <- res
							}
							return(sort(S[[k]]))
						}

						i <- seq_len(length(p))
						o <- order(p)
						ro <- order(o)
						m <- k*(k-1)/2

						val <- matrix(nrow=k,ncol=k)
						#val[lower.tri(val)] <- p.adjust(p) # Homel
						#val[lower.tri(val)] <- pmin(1, cummax( (m - i + 1L) * p[o] ))[ro] # Homel

						H <- ShafferH(k)
						val[lower.tri(val)] <- pmin(1, cummax( unlist(lapply(0:(m-1), function(x) max(H[H<=(m-x)]))) * p[o] ))[ro] # Shaffer's static S1
						val
					},
					{
						val <- matrix(nrow=k,ncol=k)
						val[lower.tri(val)] <- p.adjust(p)
						val
					}
				)

			cmp.p.values[upper.tri(cmp.p.values)] <- t(cmp.p.values)[upper.tri(cmp.p.values)]
			cmp.matrix <- matrix(ifelse(cmp.p.values < p.signif, outer(1:k,1:k, function(a,b) sign(ranks[b]-ranks[a])), 0), nrow=k,ncol=k)
			dimnames(cmp.p.values) <- dimnames(cmp.matrix) <- list(levels(X), levels(X))
		} else {
			# 1xN
			if (is.character(ctrl))
				ctrl <- match(ctrl, names(ranks))
			z <- abs(ranks[ctrl]-ranks[-ctrl]) / sqrt((k*(k+1))/(6*n))
			p <- 2*pnorm(-z)

			if (is.null(method)) method <- 'holm';
			cmp.p.values <- matrix(p.adjust(p, method=method), nrow=1, dimnames=list(names(ranks)[ctrl], names(p)))
			cmp.matrix <- ifelse(cmp.p.values < p.signif, sign(ranks[-ctrl] - ranks[ctrl]), 0)

		}

		return(list(statistic=Ff, parameter=k-1, p.value=p.value, ranks=ranks, cmp.p.values=cmp.p.values, cmp.matrix=cmp.matrix, cmp.method=method))
	} else return(list(statistic=Ff, parameter=k-1, p.value=p.value, ranks=ranks))
}
# friedman.post.hoc(successes ~ method | problem, sstab, minimize=FALSE, ctrl=c('X_A0.2','X_A0.3'))
#
# friedman.post.hoc(successes ~ method | problem, sstab, minimize=FALSE)
#
# friedman.post.hoc(successes ~ method | problem, sstab, minimize=FALSE, ctrl='X_A0.2')
#
# friedman.post.hoc(successes ~ method | problem, sstab, minimize=FALSE, ctrl=c('X_A0.2','X_A0.3'))
#
# friedman.post.hoc(successes ~ method | problem, sstab)

#################################

args <- commandArgs(trailingOnly = TRUE)
friedman.post.hoc.meth.bench(args[1], minFlag=args[2])

#friedman.post.hoc.meth.bench("tmp.csv", minFlag=FALSE)

#r <- friedman.post.hoc.meth.bench("table-success-small.csv", minFlag=FALSE)
#sort(r$ranks)
#r$p.value
#r$cmp.matrix
#r$cmp.p.values
#source("friedman.r"); r <- friedman.post.hoc.meth.bench("strp.csv", minFlag=FALSE)


#!/usr/bin/env Rscript

# ypar
# reference for main analysis: https://github.com/chr1swallace/coloc

args = commandArgs(trailingOnly=TRUE)
print(args)

library(dplyr)
library(data.table)
library(coloc)

Var.data <- function(f, N) {
  1 / (2 * N * f * (1 - f))
}

Var.data.cc <- function(f, N, s) {
  1 / (2 * N * f * (1 - f) * s * (1 - s))
}

logsum <- function(x) {
  my.max <- max(x)                              ##take out the maximum value in log form
  my.res <- my.max + log(sum(exp(x - my.max )))
  return(my.res)
}

logdiff <- function(x,y) {
  my.max <- max(x,y)                              ##take out the maximum value in log form
  my.res <- my.max + log(exp(x - my.max ) - exp(y-my.max))
  return(my.res)
}

approx.bf.p <- function(p,f,type, N, s, suffix=NULL) {
  if(type=="quant") {
    sd.prior <- 0.15
    V <- Var.data(f, N)
  } else {
    sd.prior <- 0.2
    V <- Var.data.cc(f, N, s)
  }
  z <- qnorm(0.5 * p, lower.tail = FALSE)
  ## Shrinkage factor: ratio of the prior variance to the total variance
  r <- sd.prior^2 / (sd.prior^2 + V)
  ## Approximate BF  # I want ln scale to compare in log natural scale with LR diff
  lABF = 0.5 * (log(1-r) + (r * z^2))
  ret <- data.frame(V,z,r,lABF)
  if(!is.null(suffix))
    colnames(ret) <- paste(colnames(ret), suffix, sep=".")
  return(ret)
}

approx.bf.estimates <- function (z, V, type, suffix=NULL, sdY=1) {
  sd.prior <- if (type == "quant") { 0.15*sdY } else { 0.2 }
  r <- sd.prior^2/(sd.prior^2 + V)
  lABF = 0.5 * (log(1 - r) + (r * z^2))
  ret <- data.frame(V, z, r, lABF)
  if(!is.null(suffix))
    colnames(ret) <- paste(colnames(ret), suffix, sep = ".")
  return(ret)
}

combine.abf <- function(l1, l2, p1, p2, p12) {
  lsum <- l1 + l2
  lH0.abf <- 0
  lH1.abf <- log(p1) + logsum(l1)
  lH2.abf <- log(p2) + logsum(l2)
  lH3.abf <- log(p1) + log(p2) + logdiff(logsum(l1) + logsum(l2), logsum(lsum))
  lH4.abf <- log(p12) + logsum(lsum)

  all.abf <- c(lH0.abf, lH1.abf, lH2.abf, lH3.abf, lH4.abf)
  my.denom.log.abf <- logsum(all.abf)
  pp.abf <- exp(all.abf - my.denom.log.abf)
  names(pp.abf) <- paste("PP.H", (1:length(pp.abf)) - 1, ".abf", sep = "")
  print(signif(pp.abf,3))
  print(paste("PP abf for shared variant: ", signif(pp.abf["PP.H4.abf"],3)*100 , '%', sep=''))
  return(pp.abf)
}

sdY.est <- function(vbeta, maf, n) {
    warning("estimating sdY from maf and varbeta, please directly supply sdY if known")
    oneover <- 1/vbeta
    nvx <- 2 * n * maf * (1-maf)
    m <- lm(nvx ~ oneover - 1)
    cf <- coef(m)[['oneover']]
    if(cf < 0)
        stop("estimated sdY is negative - this can happen with small datasets, or those with errors.  A reasonable estimate of sdY is required to continue.")
    return(sqrt(cf))
}

process.dataset <- function(d, suffix) {
  #message('Processing dataset')

  nd <- names(d)
  if (! 'type' %in% nd)
    stop("dataset ",suffix,": ",'The variable type must be set, otherwise the Bayes factors cannot be computed')

  if(!(d$type %in% c("quant","cc")))
      stop("dataset ",suffix,": ","type must be quant or cc")

  if(d$type=="cc") {
      if(! "s" %in% nd)
          stop("dataset ",suffix,": ","please give s, proportion of samples who are cases")
      if("pvalues" %in% nd && !( "MAF" %in% nd))
          stop("dataset ",suffix,": ","please give MAF if using p values")
      if(d$s<=0 || d$s>=1)
          stop("dataset ",suffix,": ","s must be between 0 and 1")
  }

  if(d$type=="quant") {
      if(!("sdY" %in% nd || ("MAF" %in% nd && "N" %in% nd )))
          stop("dataset ",suffix,": ","must give sdY for type quant, or, if sdY unknown, MAF and N so it can be estimated")
  }

  if("beta" %in% nd && "varbeta" %in% nd) {  ## use beta/varbeta.  sdY should be estimated by now for quant
    if(length(d$beta) != length(d$varbeta))
      stop("dataset ",suffix,": ","Length of the beta vectors and variance vectors must match")
    if(!("snp" %in% nd))
      d$snp <- sprintf("SNP.%s",1:length(d$beta))
    if(length(d$snp) != length(d$beta))
      stop("dataset ",suffix,": ","Length of snp names and beta vectors must match")

    if(d$type=="quant" && !('sdY' %in% nd))
          d$sdY <- sdY.est(d$varbeta, d$MAF, d$N)
   df <- approx.bf.estimates(z=d$beta/sqrt(d$varbeta),
                              V=d$varbeta, type=d$type, suffix=suffix, sdY=d$sdY)
    df$snp <- as.character(d$snp)
    return(df)
  }

  if("pvalues" %in% nd & "MAF" %in% nd & "N" %in% nd) { ## no beta/varbeta: use p value / MAF approximation
    if (length(d$pvalues) != length(d$MAF))
      stop('Length of the P-value vectors and MAF vector must match')
    if(!("snp" %in% nd))
      d$snp <- sprintf("SNP.%s",1:length(d$pvalues))
    df <- data.frame(pvalues = d$pvalues,
                     MAF = d$MAF,
                     snp=as.character(d$snp))
    colnames(df)[-3] <- paste(colnames(df)[-3], suffix, sep=".")
    df <- subset(df, df$MAF>0 & df$pvalues>0) # all p values and MAF > 0
    abf <- approx.bf.p(p=df$pvalues, f=df$MAF, type=d$type, N=d$N, s=d$s, suffix=suffix)
    df <- cbind(df, abf)
    return(df)
  }

  stop("Must give, as a minimum, one of:\n(beta, varbeta, type, sdY)\n(beta, varbeta, type, MAF)\n(pvalues, MAF, N, type)")
}

coloc.abf <- function(dataset1, dataset2, MAF=NULL,
                      p1=1e-4, p2=1e-4, p12=1e-5) {

  if(!is.list(dataset1) || !is.list(dataset2))
    stop("dataset1 and dataset2 must be lists.")
  if(!("MAF" %in% names(dataset1)) & !is.null(MAF))
    dataset1$MAF <- MAF
  if(!("MAF" %in% names(dataset2)) & !is.null(MAF))
    dataset2$MAF <- MAF

  df1 <- process.dataset(d=dataset1, suffix="df1")
  df2 <- process.dataset(d=dataset2, suffix="df2")
  # the na.omit is the fix here:
  merged.df <- na.omit(merge(df1,df2))

   if(!nrow(merged.df))
    stop("dataset1 and dataset2 should contain the same snps in the same order, or should contain snp names through which the common snps can be identified")

  merged.df$internal.sum.lABF <- with(merged.df, lABF.df1 + lABF.df2)
  ## add SNP.PP.H4 - post prob that each SNP is THE causal variant for a shared signal
  my.denom.log.abf <- logsum(merged.df$internal.sum.lABF)
  merged.df$SNP.PP.H4 <- exp(merged.df$internal.sum.lABF - my.denom.log.abf)

  pp.abf <- combine.abf(merged.df$lABF.df1, merged.df$lABF.df2, p1, p2, p12)
  common.snps <- nrow(merged.df)
  results <- c(nsnps=common.snps, pp.abf)

  output<-list(summary=results, results=merged.df)
  return(output)
}

############################## above here are general coloc package functions

# example file
# due to constraints in the computing environment,
# individual genes were tested first before parallelization

# head -2 input_per_gene/Whole_Blood/Coronary_Artery_Disease/ENSG00000000457.13_for_coloc.txt
# gene_id	eqtl_variant_id	tss_distance	eqtl_ma_samples	eqtl_ma_count	eqtl_maf	eqtl_pval_nominal	eqtl_slope	eqtl_slope_se	tissue	eqtl_sample_size	gwas_variant_id	chromosome	position	eqtl_effect_allele	eqtl_non_effect_allele	current_build	frequency	gwas_sample_size	gwas_zscore	gwas_pvalue	gwas_effect_size	gwas_standard_error	gwas_imputation_status	gwas_n_cases	trait
# ENSG00000000457.13	chr1_168894411_A_T_b38	-999856	30	30	0.0223881	0.4447680000000001	0.0509097	0.0665773	Whole_Blood	670	rs114383479	chr1	168894411	T	A	hg38	0.023088023088023088	184305	1.3586797714233398	0.17424808484624	NA	NA	imputed	60801	Coronary_Artery_Disease

# this script assumes input files contains per-gene or per-ldblock variants only

# inputfile = 'ENSG00000000457.13_for_coloc.txt'
inputfile = args[1]
tissueid = args[2]
traitid = args[3]

workdir = '/project/chrbrolab/gtex/coloc_v8/'
# inputdir = paste(workdir, 'input_per_gene/', sep='')
outputdir = paste(workdir, 'output_per_gene_eur/', tissueid, '/', traitid, '/', sep='')
# file names are changed to reflect new pheno labels
priorsfile = 'ypark_enloc_enrichment_and_prior_estimation_results.txt'

geneid = unlist(strsplit(inputfile, '_for_'))[1]

print(c(tissueid, traitid, geneid))

# before organizing
# inputdir = paste(workdir, 'input_per_gene/', sep='')
# after organizing
inputdir = paste(workdir, 'input_per_gene_eur/', tissueid, '/', traitid, '/', sep='')
outfile = paste(outputdir, geneid, '_coloc_output_file.txt', sep='')
output = c()

inputdata = tbl_df(fread(paste(inputdir, inputfile, sep='')))
inputdata = inputdata[inputdata$frequency > 0,]
inputdata$gwas_varbeta = 1

inputdata = inputdata %>% filter(complete.cases(gwas_zscore)) %>% mutate(max_gwas_zscore = max(abs(gwas_zscore)))
inputdata = inputdata %>% arrange(desc(abs(gwas_zscore)))

priors = tbl_df(fread(paste(workdir, priorsfile, sep='')))
priorsused = priors %>% filter(trait == traitid & tissue == tissueid)
p1used = priorsused$p1_trait_prior
p2used = priorsused$p2_tissue_prior
p12used = priorsused$p12

# eqtlinput = list(beta = inputdata$eqtl_slope, varbeta = (inputdata$eqtl_slope_se)**2, N = inputdata$eqtl_sample_size, type = 'quant', snp = inputdata$eqtl_variant_id, MAF = inputdata$eqtl_maf)
eqtlinput = list(pvalues = inputdata$eqtl_pval_nominal, N = inputdata$eqtl_sample_size, type = 'quant', snp = inputdata$eqtl_variant_id, MAF = inputdata$frequency)

caseproportion = inputdata$gwas_n_cases/inputdata$gwas_sample_size
datatype = ifelse(is.na(caseproportion[1]), 'quant', 'cc')

if(datatype == 'quant') {
  gwasinput = list(beta = inputdata$gwas_zscore, varbeta = inputdata$gwas_varbeta, N = inputdata$gwas_sample_size, type = datatype, snp = inputdata$eqtl_variant_id, MAF = inputdata$frequency)
} else {
  gwasinput = list(beta = inputdata$gwas_zscore, varbeta = inputdata$gwas_varbeta, N = inputdata$gwas_sample_size, type = datatype, snp = inputdata$eqtl_variant_id, MAF = inputdata$frequency, s = caseproportion)
}

# if(datatype == 'quant') {
#   gwasinput = list(pvalues = inputdata$gwas_pvalue, N = inputdata$gwas_sample_size, type = datatype, snp = inputdata$eqtl_variant_id, MAF = inputdata$frequency)
# } else {
#   gwasinput = list(pvalues = inputdata$gwas_pvalue, N = inputdata$gwas_sample_size, type = datatype, snp = inputdata$eqtl_variant_id, MAF = inputdata$frequency, s = caseproportion)
# }

# priors were stored in a table like this:
# tissue	trait	p1_trait_prior	p2_tissue_prior	p12	pm1	pm2	enloc_enrich_intercept_effect(inteffect)	inteffect_95ci_lower	inteffect_95ci_upper	enloc_enrich_log_odds_ratio_effect(logor)	logor_95ci_lower	logor_95ci_upper
# Adipose_Subcutaneous	Adiponectin	7.31680E-06	3.32503E-03	3.05385E-08	7.34734E-06	3.32506E-03	-11.822	-11.849	-11.795	0.224	-2.384	2.832
# Adipose_Subcutaneous	Alzheimers_Disease	1.32258E-05	3.32490E-03	1.60605E-07	1.33864E-05	3.32506E-03	-11.230	-11.409	-11.052	1.292	-2.956	5.540

runcoloc = coloc.abf(gwasinput, eqtlinput, p1=p1used, p2=p2used, p12=p12used)
output = rbind(output, c(unlist(inputdata[1,]), runcoloc$summary))
output = as.data.table(output)

output = mutate(output, PP.H1.abf = ifelse(PP.H0.abf == 1, 0, PP.H1.abf))
output = mutate(output, PP.H2.abf = ifelse(PP.H0.abf == 1, 0, PP.H2.abf))
output = mutate(output, PP.H3.abf = ifelse(PP.H0.abf == 1, 0, PP.H3.abf))
output = mutate(output, PP.H4.abf = ifelse(PP.H0.abf == 1, 0, PP.H4.abf))

outputdata = tbl_df(as.data.frame(output)) %>% mutate(PP.H4.abf = as.numeric(as.character(PP.H4.abf))) %>% arrange(desc(PP.H4.abf))
write.table(outputdata, file=outfile, quote=F, col.names=T, row.names=F, sep='\t')


###############################################################################
#  Copyright (C) 1994 - 2009, Performance Dynamics Company                    #
#                                                                             #
#  This software is licensed as described in the file COPYING, which          #
#  you should have received as part of this distribution. The terms           #
#  are also available at http://www.perfdynamics.com/Tools/copyright.html.    #
#                                                                             #
#  You may opt to use, copy, modify, merge, publish, distribute and/or sell   #
#  copies of the Software, and permit persons to whom the Software is         #
#  furnished to do so, under the terms of the COPYING file.                   #
#                                                                             #
#  This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY  #
#  KIND, either express or implied.                                           #
###############################################################################

# Load the library
library("pdq")
#
# T/put bounds for closed tandem QNM with increasing Z
# Created by NJG on Fri Feb 27, 2009
#
clients = 300
thinkZ  = 800 * 10^(-3) # ms as seconds
stime   = 20  * 10^(-3) # ms as seconds
node1   = "n1"
node2   = "n2"
node3   = "n3"
work    = "w"
for(z in 1:3) {
   xc<-0 # prudent to reset these
   yc<-0
   for (i in 1:clients) {
      Init("")
      if (z==1) {think<-thinkZ * 0}
      if (z==2) {think<-thinkZ * 2}
      if (z==3) {think<-thinkZ * 4}
      CreateClosed(work, TERM, as.double(i), think)
      CreateNode(node1, CEN, FCFS)
      CreateNode(node2, CEN, FCFS)
      CreateNode(node3, CEN, FCFS)
      SetDemand(node1, work, stime)
      SetDemand(node2, work, stime)
      SetDemand(node3, work, stime)
      Solve(APPROX)
      xc[i]<-as.double(i)
      yc[i]<-GetThruput(TERM, work)
      nopt<-(3*stime + think)/stime
     }
    if (z==1) {
    	# establish plot frame and first curve
        plot(xc, yc, type="l", ylim=c(0,1/stime), lwd=2, xlab="N Clients",
             ylab="Throughput X(N)")
        title("Increasing Z Values")
      	abline(0, 1/(nopt*stime), lty="dashed",  col = "blue")
      	abline(1/stime, 0, lty="dashed",  col = "red")
      	abline(v=nopt, lty="dashed", col="gray50")
        text(20+nopt, 35, paste("Z=", as.numeric(think)))
      	text(3+nopt, 6, paste("N*=", as.numeric(nopt)))
    } else {
    	# add the other curves
      	points(xc, yc, type="l", lwd=2)
      	abline(0, 1/(nopt*stime), lty="dashed",  col = "blue")
      	abline(v=nopt, lty="dashed", col="gray50")
      	format(think, scientific = TRUE)
      	text(4+nopt, 35, paste("Z=", as.numeric(think)))
      	text(3+nopt, 6, paste("N*=", as.numeric(nopt)))
    }
}