library(scales)

# set your folder
setwd("//Users//wayne//Documents//My research//0. Pop CSGLD//")

# total number of iterations
total_iters = 1e6
# learning rate for gradient Langevin dynamics
eta = 0.1
# temperature
T = 1
# hyperparameter
zeta = 0.9
# Total number of partitions
parts = 100
# Energy gap
div = 1
# Minimum energy for partition
min_energy = 0.15

# build the energy function and the gradient
energy = function(x) return(-log(0.4 * dnorm(x, -6, 1) + 0.6 * dnorm(x, 4, 1)))
grad_f = function(x) return(numDeriv::grad(energy, x) + 0.03 * rnorm(1, 0, 1))

# find the index for new samples
find_idx = function(sample) {
    stochastic_energy = energy(sample) + 0.03 * rnorm(1, 0, 1)
    return(min((as.integer((energy(sample) - min_energy) / div) + 1), parts-1))
}

random_field = function(theta, idx) {
    sub_random_field = -theta
    sub_random_field[idx] = sub_random_field[idx] + 1
    return(theta[idx] * sub_random_field)
}

# print process
print_progress = function(iter, total_iters) {
    if (iter %% (total_iters / 10) == 0) {
        print(paste('Iter ', scientific(iter, digits = 3), '/',  total_iters,' completed'))
    }
}

# Since we are converging to a biased equilibrium that approximates the ground truth, we will not compare it and will only
# focus on validate the reduction of variance
#############################################     Ground truth density of states  #########################################
parts_for_presentation = 11
# simulate samples from the Gaussian mixture distribution
real_samples = c(rnorm(1e6 * 0.4, -6, 1), rnorm(1e6 * 0.6, 4, 1))
# obtain the energy value for these samples
energy_each_sample = mapply(energy, real_samples)

# establish energy partition
grids = seq(1, parts_for_presentation) * div + min_energy
# compute density of states (ignore the first partition)
real_density_of_states = rep(0, (parts_for_presentation-1))
for (i in 1:(parts_for_presentation-1)) {
    real_density_of_states[i]  = (sum(energy_each_sample <= grids[i+1]) - sum(energy_each_sample <= grids[i])) / 1e6
}

chains = 10
#########################################    ICSGLD  (many short chains)    ##############################################
theta_ICSGLD = c()
candidate_seeds = seq(1, 10)
for (seeds in candidate_seeds) {
    set.seed(seeds)
    print(paste("Run ICSGLD (many short chains) based on the random seed", seeds))
    theta = rep(1, parts) / parts
    samples = rep(0, chains)
    indexs = rep(1, chains)
    for (iter in 1: total_iters) {
        decay = 1 / (iter^0.6 + 100)
        print_progress(iter, total_iters)
        # adaptive sampling part (parallelizable)
        for (cidx in 1: chains) {
            multiplier = 1 + zeta * T * (log(theta[indexs[cidx]+1]) - log(theta[indexs[cidx]])) / div
            samples[cidx] = samples[cidx] - eta * multiplier * grad_f(samples[cidx])  + sqrt(2 * eta * T) * rnorm(1, 0, 1)
            indexs[cidx] = find_idx(samples[cidx])
        }
        # interacting stochastic approximation
        interacting_random_field = 0
        for (cidx in 1: chains) {
            interacting_random_field = interacting_random_field + random_field(theta, indexs[cidx]) / chains
        }
        theta = theta + decay * interacting_random_field
    }
    # ignore the first empty partition where the minimum energy is around 1.42
    theta_ICSGLD = cbind(theta_ICSGLD, theta[2:parts_for_presentation])
}

#########################################    CSGLD  (a single long chain)    ##############################################
theta_CSGLD = c()
for (seeds in candidate_seeds) {
    set.seed(seeds)
    print(paste("Run CSGLD (a single long chain) based on the random seed", seeds))
    theta = rep(1, parts) / parts
    sample = 0
    idx = 1
    for (iter in 1: (total_iters*chains)) {
        decay = 1 / (iter^0.6 + 100)
        print_progress(iter, total_iters*chains)
        # adaptive sampling part
        multiplier = 1 + zeta * T * (log(theta[idx+1]) - log(theta[idx])) / div
        sample = sample - eta * multiplier * grad_f(sample)  + sqrt(2 * eta * T) * rnorm(1, 0, 1)
        idx = find_idx(sample)
        theta = theta + decay * random_field(theta, idx)
    }
    # ignore the first empty partition where the minimum energy is around 1.42
    theta_CSGLD = cbind(theta_CSGLD, theta[2:parts_for_presentation])
}

library(ggplot2)

ICSGLD_estimated_density_of_states = rowMeans(theta_ICSGLD)^zeta / sum(rowMeans(theta_ICSGLD)^zeta)
CSGLD_estimated_density_of_states = rowMeans(theta_CSGLD)^zeta / sum(rowMeans(theta_CSGLD)^zeta)

# we don't know the real equilibrium, we approximate it instead
biased_equilibrium = (ICSGLD_estimated_density_of_states + CSGLD_estimated_density_of_states) / 2

wdata = cbind(CSGLD_estimated_density_of_states, ICSGLD_estimated_density_of_states, biased_equilibrium)
ddf = as.data.frame(wdata)
ddf$X = seq(1, length(ICSGLD_estimated_density_of_states))
ddf$ICSGLD_theta_upper = ddf$ICSGLD_estimated_density_of_states + 2 * apply(theta_ICSGLD, 1, sd)
ddf$ICSGLD_theta_lower = ddf$ICSGLD_estimated_density_of_states - 2 * apply(theta_ICSGLD, 1, sd)
ddf$CSGLD_theta_upper = ddf$CSGLD_estimated_density_of_states + 2 * apply(theta_CSGLD, 1, sd)
ddf$CSGLD_theta_lower = ddf$CSGLD_estimated_density_of_states - 2 * apply(theta_CSGLD, 1, sd)

sub_p = ggplot(ddf, aes(X)) +
    geom_line(aes(y=ICSGLD_estimated_density_of_states, col="#7570B3"), size=0.5) +
    geom_line(aes(y=CSGLD_estimated_density_of_states, col="#FF9999"), size=0.5) +
    geom_line(aes(y=biased_equilibrium, col="black"), size=1) +
    geom_ribbon(aes(ymin=ICSGLD_theta_lower,ymax=ICSGLD_theta_upper, x=X),alpha=0.3, fill="#7570B3") +
    geom_ribbon(aes(ymin=CSGLD_theta_lower,ymax=CSGLD_theta_upper, x=X),alpha=0.3, fill="#FF9999") +
    scale_x_continuous(name="Energy") +
    scale_y_continuous(name="PDF in Energy (density of states)") +
    scale_color_identity(name = "Multipliers",
                       breaks = c("#7570B3", "#FF9999", "black"),
                       labels = c("ICSGLD xP10", "CSGLD xT10", bquote(hat(theta)["*"])),
                       guide = "legend")  +
    theme(legend.position=c(0.8, 0.2),
        legend.title = element_blank(),
        axis.title=element_text(size=20),
        legend.text = element_text(colour="grey15", size=24),
        legend.key.size = unit(2,"line"),
        legend.key = element_blank(),
        legend.background = element_rect(fill=alpha('grey', 0.1)),
        axis.text.y = element_text(size=16),
        axis.text.x = element_text(size=16))

print(sub_p)


LISP Interpreter Run

[[[
 First steps with my new construction for
 a self-delimiting universal Turing machine.
 We show that
    H(x,y) <= H(x) + H(y) + c
 and determine c.
 Consider a bit string x of length |x|.
 We also show that
    H(x) <= 2|x| + c
 and that
    H(x) <= |x| + H(the binary string for |x|) + c
 and determine both these c's.
]]]

[
 Here is the self-delimiting universal Turing machine!
]
define (U p) cadr try no-time-limit 'eval read-exp p

define      U
value       (lambda (p) (car (cdr (try no-time-limit (' (eval
            (read-exp))) p))))

(U bits 'cons x cons y cons z nil)

expression  (U (bits (' (cons x (cons y (cons z nil))))))
value       (x y z)

(U append bits 'cons a debug read-exp
          bits '(b c d)
)

expression  (U (append (bits (' (cons a (debug (read-exp)))))
            (bits (' (b c d)))))
debug       (b c d)
value       (a b c d)

[
 The length of alpha in bits is the
 constant c in H(x) <= 2|x| + 2 + c.
]
define alpha
let (loop) let x read-bit
           let y read-bit
           if = x y
              cons x (loop)
              nil
(loop)

define      alpha
value       ((' (lambda (loop) (loop))) (' (lambda () ((' (lam
            bda (x) ((' (lambda (y) (if (= x y) (cons x (loop)
            ) nil))) (read-bit)))) (read-bit)))))

length bits alpha

expression  (length (bits alpha))
value       1104

(U
 append
   bits alpha
   '(0 0 1 1 0 0 1 1 0 1)
)

expression  (U (append (bits alpha) (' (0 0 1 1 0 0 1 1 0 1)))
            )
value       (0 1 0 1)

(U
 append
   bits alpha
   '(0 0 1 1 0 0 1 1 0 0)
)

expression  (U (append (bits alpha) (' (0 0 1 1 0 0 1 1 0 0)))
            )
value       out-of-data

[
 The length of beta in bits is the
 constant c in H(x,y) <= H(x) + H(y) + c.
]
define beta
cons eval read-exp
cons eval read-exp
     nil

define      beta
value       (cons (eval (read-exp)) (cons (eval (read-exp)) ni
            l))

length bits beta

expression  (length (bits beta))
value       432

(U
 append
   bits  beta
 append
   bits 'cons a cons b cons c nil
   bits 'cons x cons y cons z nil
)

expression  (U (append (bits beta) (append (bits (' (cons a (c
            ons b (cons c nil))))) (bits (' (cons x (cons y (c
            ons z nil))))))))
value       ((a b c) (x y z))

(U
 append
   bits beta
 append
   append bits alpha '(0 0 1 1 0 0 1 1 0 1)
   append bits alpha '(1 1 0 0 1 1 0 0 1 0)
)

expression  (U (append (bits beta) (append (append (bits alpha
            ) (' (0 0 1 1 0 0 1 1 0 1))) (append (bits alpha)
            (' (1 1 0 0 1 1 0 0 1 0))))))
value       ((0 1 0 1) (1 0 1 0))

[
 The length of gamma in bits is the
 constant c in H(x) <= |x| + H(|x|) + c
]
define gamma
let (loop k)
   if = 0 k nil
   cons read-bit (loop - k 1)
(loop base2-to-10 eval read-exp)

define      gamma
value       ((' (lambda (loop) (loop (base2-to-10 (eval (read-
            exp)))))) (' (lambda (k) (if (= 0 k) nil (cons (re
            ad-bit) (loop (- k 1)))))))

length bits gamma

expression  (length (bits gamma))
value       1024

(U
 append
   bits gamma
 append
   [Arbitrary program for U to compute number of bits]
   bits' '(1 0 0 0)
   [That many bits of data]
   '(0 0 0 0  0 0 0 1)
)

expression  (U (append (bits gamma) (append (bits (' (' (1 0 0
             0)))) (' (0 0 0 0 0 0 0 1)))))
value       (0 0 0 0 0 0 0 1)

End of LISP Run

Elapsed time is 0 seconds.


#' @title calcKristWakeFrac
#'
#'@description Calculate wake fraction (\code{wakeFraction}) (dimensionless)
#'using the Kristensen method.
#'
#'@param shipType Ship type (vector of strings, see \code{\link{calcShipType}})
#'@param breadth Moulded breadth (vector of numericals, m)
#'@param lwl Waterline length (vector of numericals, m) (see
#'\code{\link{calclwl}})
#'@param Cbw Waterline block coefficient (vector of numericals, dimensionless)
#'(see \code{\link{calcCbw}})
#'@param propDiam Propeller diameter (vector of numericals, m) (see
#'\code{\link{calcPropDia}})
#'@param M Fineness/slenderness coefficient (vector of numericals,
#'dimensionless) (see \code{\link{calcShipM}})
#'@param nProp Number of propellers (vector of numericals, see
#'\code{\link{calcPropNum}})
#'@param tankerBulkCarrierShipTypes Ship types specified in input
#'\code{shipTypes} to be modeled as tankers and bulk carriers vessels (vector of
#' strings)
#'
#' @return \code{wakeFraction} (vector of numericals, dimensionless)
#'
#' @details
#' "The speed of advance of the propeller relative to the water in which it is working is
#'  lower than the observed speed of the vessel. This difference in speed,
#'  expressed as a percentage of the ship speed, is known as the wake fraction coefficient".
#'  \url{https://www.wartsila.com/encyclopedia/term/wake-fraction-coefficient}
#'
#'Wake fraction is a component of hull efficiency as well as a component of
#'propeller efficiency.
#'
#' This method this requires ship types to be grouped. Use the
#' \code{tankerBulkCarrierShipTypes} grouping  parameters to provide these ship
#' type groupings. Any ship types not included in this grouping will be considered as
#' miscellaneous vessels.
#'
#' @references
#'Kristensen, H. O. and Lutzen, M. 2013. "Prediction of Resistance and Propulsion
#'Power of Ships."
#'
#'\href{https://gitlab.gbar.dtu.dk/oceanwave3d/Ship-Desmo}{Kristensen, H. O.
#'"Ship-Desmo-Tool." https://gitlab.gbar.dtu.dk/oceanwave3d/Ship-Desmo}
#'
#'@seealso \itemize{
#'\item \code{\link{calclwl}}
#'\item \code{\link{calcCbw}}
#'\item \code{\link{calcPropDia}}
#'\item \code{\link{calcShipM}}
#'\item \code{\link{calcPropNum}}
#'\item \code{\link{calcShipType}}}
#'
#' @family Kristensen Calculations
#'
#' @examples
#' calcKristWakeFrac("bulk.carrier",32.25,218.75, 0.8, 6.7,5.2, 1)
#' calcKristWakeFrac("oil.tanker",32.25,218.75, 0.8, 6.7,5.2, 1)
#' calcKristWakeFrac(c("bulk.carrier","container.ship"),
#'                   c(32.25,49),
#'                   c(218.75,400),
#'                   c(0.8,0.75),
#'                   c(6.7,9.6),
#'                   c(5.2,6.7),
#'                   c(1,1))
#'
#' @export

calcKristWakeFrac <- function(shipType, breadth,lwl, Cbw, propDiam,M, nProp,
                              tankerBulkCarrierShipTypes=c("tanker","chemical.tanker","liquified.gas.tanker","oil.tanker","other.tanker","bulk.carrier")
){

  wakeFraction<-ifelse(nProp==1,
            #a
            ((((0.1*breadth)/
                 lwl)+0.149)+
               (
                 #b
                 (((0.05*breadth)/
                     lwl)+0.449)/
                   #c
                   (
                     (585-((5027*breadth)/
                             lwl)+
                        11700*((breadth/
                                  lwl)^2)
                     )*
                       ((0.98-Cbw)^3)+1)))+
              # Assume w2 = 0 because we assume F_a = 0 where
              #F_a = [-2,0,2] and F_a=0 represents a N-shaped hull form
              #take this assumption for lack of better information
              #w3
              pmin(0.1,-0.18+(0.00756/((propDiam/lwl)+
                                         0.002))),
            #nprop == 2 case
            1.133*(Cbw^2)-0.797*Cbw+0.215
            )

  #apply adjustments from Kristensen, 2016 (to correct Harvald's algorithms toward data in more recent model tests)
  wakeFraction<-ifelse( (nProp==1) & ( shipType %in% tankerBulkCarrierShipTypes), #End of case definition
             (0.7*wakeFraction-0.45+0.08*M),#second case= container ships
             0.7*wakeFraction
            )

  return(wakeFraction)

}# end of function: calcWakeFrac


#' Runs cross-validation of Random Forests classification
#'
#' @param x A matrix of numeric predictors
#' @param y A factor or logical of responses
#' @param nfolds number of cross-validation folds (nfolds==-1 means leave-one-out)
#' @param folds optional, if not \code{NULL} then uses these fold assignments (one fold index for each row in \code{x})
#' @param verbose Use verbose output
#' @param ... additional parameters for randomForest
#'
#' @return
#'  \item{y}{Observed values}
#'  \item{predicted}{CV predicted values}
#'  \item{probabilities}{CV predicted class probabilities (or NULL if unavailable)}
#'  \item{confusion.matrix}{Confusion matrix (true-by-predicted)}
#'  \item{nfolds}{Number of folds used (or -1 for leave-one-out)}
#'  \item{params}{List of additional parameters, if any}
#'  \item{importances}{Feature-by-fold matrix of importances of features (mean decrease in accuracy) across folds}
#'  \item{final.model}{Final random forests model trained on whole data set}
#' @examples
#' # generate fake data: 50 samples, 20 variables
#' x <- matrix(rnorm(2000),100)
#' rownames(x) <- sprintf('Sample%02d',1:100)
#' colnames(x) <- sprintf('Feature%02d',1:20)
#' # generate fake response variable, function of columns 1-5 plus noise
#' y <- rowSums(sweep(x[,1:5],2,runif(5),'*')) + rnorm(100,sd=.5)
#' # y is now binary response variable
#' y <- factor(y > median(y))
#' names(y) <- rownames(x)
#'
#' # 10-fold cross validation, returning predictions
#' res.rf <- rf.cross.validation(x,y2)
#'
#' # plot importance of top ten variables
#' sorted.importances <- sort(rowMeans(res.rf$importances), decreasing=TRUE)
#' barplot(rev(sorted.importances[1:10]),horiz=TRUE, xlab='Mean Decrease in Accuracy')
#'
#' # Plot classification ROC curve
#' roc <- probabilities.to.ROC(res.rf$probabilities[,2], y2,plot=TRUE)
#'
#' # Report ROC AUC
#' cat('ROC AUC was:',roc$auc,'\n')
#'

"rf.cross.validation" <- function(x, y, nfolds=10, folds=NULL, verbose=FALSE, ...){
    require('randomForest')
    if(nfolds==-1) nfolds <- length(y)
    if(is.null(folds)) folds <- balanced.folds(y,nfolds=nfolds)
    result <- list()
    result$y <- as.factor(y)
    result$predicted <- result$y
    result$probabilities <- matrix(0, nrow=length(result$y), ncol=length(levels(result$y)))
    rownames(result$probabilities) <- rownames(x)
    colnames(result$probabilities) <- levels(result$y)
    result$importances <- matrix(0,nrow=ncol(x),ncol=nfolds)
    result$errs <- numeric(length(unique(folds)))

    # K-fold cross-validation
    for(fold in sort(unique(folds))){
        if(verbose) cat(sprintf('Fold %d...\n',fold))
        foldix <- which(folds==fold)
        model <- randomForest(x[-foldix,], factor(result$y[-foldix]), importance=TRUE, do.trace=verbose, ...)
        newx <- x[foldix,,drop=F]
        if(length(foldix)==1) newx <- matrix(newx,nrow=1)
        result$predicted[foldix] <- predict(model, newx)
        probs <- predict(model, newx, type='prob')
        result$probabilities[foldix,colnames(probs)] <- probs
        result$errs[fold] <- mean(result$predicted[foldix] != result$y[foldix])
        result$importances[,fold] <- model$importance[,'MeanDecreaseAccuracy']
    }
    rownames(result$importances) <- colnames(x)
    result$err <- mean(result$predicted != result$y)
    result$nfolds <- nfolds
    result$params <- list(...)
    result$confusion.matrix <- t(sapply(levels(y), function(level) table(result$predicted[y==level])))

    # train one final model on the whole data set
    result$final.model <- randomForest(x,result$y)

    return(result)
}

#' Runs cross-validation of LARS regression
#'
#' @param x A matrix of numeric predictors
#' @param y A factor or logical of responses
#' @param nfolds number of cross-validation folds (nfolds==-1 means leave-one-out)
#' @param folds optional, if not \code{NULL} then uses these fold assignments (one fold index for each row in \code{x})
#' @param verbose Use verbose output
#' @param ... additional parameters for lars
#'
#' @return:
#'  \item{y}{Observed values}
#'  \item{predicted}{CV predicted values}
#'  \item{fractions}{Vector of LASSO fractions used across folds}
#'  \item{best.fraction}{Best LASSO fraction used}
#'  \item{MSE}{Mean squared error}
#'  \item{RMSE}{Root mean squared error}
#'  \item{nfolds}{Number of folds used (or -1 for leave-one-out)}
#'  \item{params}{List of additional parameters, if any}
#'  \item{coefficients}{Feature-by-fold matrix of coefficients of features across folds}
#'  \item{final.model}{Final lars model trained on whole data set}
#' @examples
#' # generate fake data: 50 samples, 20 variables
#' x <- matrix(rnorm(2000),100)
#' rownames(x) <- sprintf('Sample%02d',1:100)
#' colnames(x) <- sprintf('Feature%02d',1:20)
#' # generate fake response variable, function of columns 1-5 plus noise
#' y <- rowSums(sweep(x[,1:5],2,runif(5),'*')) + rnorm(100,sd=.5)
#' names(y) <- rownames(x)
#'
#' # 10-fold cross validation with predictions
#' res.lars <- lars.cross.validation(x,y)
#' # plot predictions
#' plot(res.lars$y, res.lars$predicted, xlab="Observed", ylab="Predicted (hold-out)")
#'
#' # correlation
#' cat('Pearson correlation of predictions with truth?\n')
#' print(cor.test(res.lars$y, res.lars$predicted))
#'
#' # plot fraction of times each variable was included in a fold
#' sorted.importances <- sort(rowMeans(res.lars$coefficients > 0), decreasing=TRUE)
#' # bar plot of top ten variables
#' par(mar=c(4.5,6,2,1),las=2) # larger margin on left, perpendicular axis labels
#' barplot(rev(sorted.importances[1:10]),horiz=TRUE, xlab='Frequency of selection')
#'

"lars.cross.validation" <- function(x, y, nfolds=10, folds=NULL, verbose=FALSE, ...){
    require('lars')
    if(nfolds==-1) {
        nfolds <- length(y)
        folds <- 1:length(y)
    } else{
        folds <- sample(rep(1:nfolds,times=ceiling(length(y)/nfolds))[1:length(y)])
    }
    result <- list()
    result$y <- as.numeric(y)
    result$predicted <- result$y
    result$coefficients <- matrix(0,nrow=ncol(x),ncol=nfolds)
    result$fractions <- numeric(length(y))

    # K-fold cross-validation
    for(fold in sort(unique(folds))){
        if(verbose) cat(sprintf('Fold %d...\n',fold))
        foldix <- which(folds==fold)

        # do cv.lars on the training set to choose the best penalty
        cv.res <- cv.lars(x[-foldix,], result$y[-foldix], plot.it=FALSE, ...)
        threshold <- cv.res$index[which.min(cv.res$cv)]
        res <- lars(x[-foldix,], result$y[-foldix])

        # predict on test set
        newx <- x[foldix,,drop=F]
        yhat <- predict(res, newx, mode='fraction',type='fit',s=threshold)$fit

        # extract results
        result$predicted[foldix] <- yhat
        result$coefficients[,fold] <- coef(res,mode='fraction',s=threshold)
        result$fractions[fold] <- threshold
    }
    rownames(result$coefficients) <- colnames(x)
    result$MSE <- mean((result$predicted - result$y)^2)
    result$RMSE <- sqrt(result$MSE)
    result$nfolds <- nfolds
    result$params <- list(...)
    result$best.fraction <- mean(result$fractions)

    # train one final model on the whole data set
    result$final.model <- lars(x,result$y)
    return(result)
}

#' Get balanced folds where each fold has close to overall class ratio. This only works for non-numeric inputs.
#'
#' @param y Vector of classes; should be a factor, logical, or character
#' @param nfolds Number of folds. -1 means leave-one-out
#' @return Numeric vector of fold indices, one for each element of \code{y}
"balanced.folds" <- function(y, nfolds=10){
    if(is.numeric(y)) stop("MWAS balanced.folds function only accepts non-numeric inputs.")
    folds = rep(0, length(y))
    y <- as.factor(y)
    classes = levels(y)
    # size of each class
    Nk = table(y)
    # -1 or nfolds = len(y) means leave-one-out
    if (nfolds == -1 || nfolds == length(y)){
        invisible(1:length(y))
    }
    else{
    # Can't have more folds than there are items per class
    nfolds = min(nfolds, max(Nk))
    # Assign folds evenly within each class, then shuffle within each class
        for (k in 1:length(classes)){
            ixs <- which(y==classes[k])
            folds_k <- rep(1:nfolds, ceiling(length(ixs) / nfolds))
            folds_k <- folds_k[1:length(ixs)]
            folds_k <- sample(folds_k)
            folds[ixs] = folds_k
        }
        invisible(folds)
    }
}

#' Performs ROC analysis (and plot) on probability vector
#'
#' @param p Vector or probabilities of TRUE
#' @param y Logical, factor, or character of TRUE/FALSE
#' @param step.size Increments at which to calculate false and true positive rates
#' @param plot If \code{TRUE} then make a simple ROC curve plot.
#' @param ... Additional parameters to pass to \code{plot} function
#' @return
#'  \item{fprs}{Vector of false positive rates}
#'  \item{tprs}{Vector of true positive rates}
#'  \item{auc}{Area under the curve}
"probabilities.to.ROC" <- function(p,y,step.size=0.01, plot=TRUE,...){
    require('flux')
    y <- as.logical(y)
    fprs <- sapply(seq(1,0,-abs(step.size)), function(xx) mean(p[!y] >= xx))
    tprs <- sapply(seq(1,0,-abs(step.size)), function(xx) mean(p[y] >= xx))
    tprs[fprs > tprs] <- fprs[fprs > tprs]
    auc.res <- auc(fprs, tprs)
    plot(fprs,tprs,type='l', xlab='False Positive Rate',ylab='True Positive Rate',...)

    return(list(fprs=fprs, tprs=tprs, auc=auc.res))
}


x <- c(-8, -3, 9)
pr_1 <- 4
pr_2 <- 2
pr_3 <- pr_1 + pr_2

# Svolgimento
library(MASS)
print(0)
print(fractions(1/pr_3))
a <- x[3]/-(x[1]*pr_1 + x[2]*pr_2 + x[3]*(-pr_3))
print(fractions(a))
print(x[1]^2*(a*pr_1) + x[2]^2*(a*pr_2) + x[3]^2*(1-(a*pr_3)))


# écriture de stochastic_model sans array
# avec output true

# endogènes -----------------
## qpib pib en volume
## vpib pib en valeur (mds euro)
## og écart de production
## ppib prix du pib
## inf  inflation
## imp_dep, imp_po, imp_int décompositon de l'impulsion budgétaire en trois composantes
##                           dépenses, recettes, et intérêt
## dette dette publique, en valeur (mds euro)
## tppdette dette publique en point de pib potentiel
## dettep dette publique en point de pib
## cha_int  charge d'intérêt en valeur (mds euro)
## tchaint  charge d'intérêt en point de pib potentiel
## r_app    taux d'intérêt souverain national apparent (1% = 0.01)
## r_inst   taux d'intérêt souverain instantané (moyen année) (1% = 0.01)
## gpot croissance potentielle en volume
## aide de odin https://mrc-ide.github.io/odin/
## tcho est le taux de chômage au sens du BIT
## tvnairu est le NAIRU de AMECO tandis que Nairu est le chômage d'équilibre

# exogènes --------------
## dep_prim dépenses primaires en valeur
## rec_po   recettes fiscales (et non fiscales) en valeur (mds euro)
## pibpot pib potentiel en volume

# écart de production -------------------
og <- lagog + dog
dog <- og_lag * (lagog - lag2og) +
  - og_mce * (lagog - 0) +
  og_dep * ib_dep +
  og_po * ib_po +
  og_cin * ib_cin +
  og_noise +
  - og_ddpot * ((pibpot/lagpibpot-1)-(lagpibpot/lag2pibpot-1)) +
  - pot_r*(r_inst-lagr_inst)
# on retarde de cette façon, cela permet de passer une condition initiale
update(lagog) <- og
update(lag2og) <- lagog
update(og_noise) <- ogn_ar * og_noise + rnorm(0, ogn_sigma)

# Taux de chômage -------------------
tcho <- lagtcho - tcho_okun * dog - tcho_mce * (lagtcho+lagog-lagtvnairu)
update(lagtcho) <- tcho
update(lag2tcho) <- lagtcho

# Dynamique du NAIRU -------------------
tvnairu <- lagtvnairu - tvnairu_mce * (lagtvnairu - nairu)
update(lagtvnairu) <- tvnairu

# prix du pib -------------------
txppib <- lagtxppib + inf_lag * (lagtxppib - lag2txppib) - inf_mce * (lagtxppib - infstar + inf_tcho * (lagtcho - lagtvnairu))
ppib <- (1 + txppib)*lagppib
update(lagppib) <- ppib
update(lag2ppib) <- lagppib
update(lag3ppib) <- lag2ppib
lagtxppib <- lagppib/lag2ppib-1
lag2txppib <- lag2ppib/lag3ppib-1

# pib volume&valeur&potentiel -----------
pibpot <- lagpibpot*(1+(gpot[step]-1*((tvnairu-lagtvnairu)/(1-(tvnairu-i_lagtvnairu)))))
update(lagpibpot) <- pibpot
update(lag2pibpot) <- lagpibpot
qpib <- (1 + og) * pibpot
vpib <- qpib * ppib
lagvpib <- lagqpib * lagppib
lag2vpib <- lag2qpib * lag2ppib
update(lagqpib) <- qpib
update(lag2qpib) <- lagqpib
gpib <- qpib/lagqpib-1
update(laggpib) <- gpib
gpot_v <- pibpot/lagpibpot-1

# dette ------------
# dette au 31 décembre
dette <- lagdette + ddette + dette_oneoff[step]
update(lagdette) <- dette
update(lag2dette) <- lagdette
lagdettep <- lagdette/lagvpib
lag2dettep <- lag2dette/lag2vpib
ddette <- - rec_po + dep_prim + cin - taut[step] * vpib

# taux d'intérêt ---------------
#r_inst <- gpot[step] + lagtxppib + ecart_c + r_dette * (lagdettep - dstarr)
r_inst <- gpot_v + infstar + ecart_c + r_dette * (lagdettep - dstarr)
r_app <- (1-1/r_mat) * lagr_app +  1/r_mat * r_inst
update(lagr_app) <- r_app
ecart_c <- lagecart_c + decart_c
decart_c <- -ec_mce * (lagecart_c - ecstar)
update(lagecart_c) <- ecart_c
update(lagr_inst) <- r_inst

# impulsion --------------
# l'impulsion sur les dépenses est définie comme la variation du taux des dépenses par rapport au potentiel
dep_prim <- tdeppp * ppib_dep * pibpot_dep

ppib_dep <- lagppib_dep*(1+ infdep_star*txppib + (1-infdep_star)*lagtxppib_dep - infdep_mce*(lagppib_dep/lagppib-1))
pibpot_dep <- lagpibpot_dep*(1+ potdep_star*gpot_v + (1-potdep_star)*lagtxpibpot_dep - potdep_mce*(lagpibpot_dep/lagpibpot-1))

#pibpot_dep <- lagpibpot_dep*(1+(-((lag2pibpot_dep-lagpibpot_dep)/lag2pibpot_dep)-potdep_mce*(lagpibpot_dep/lagpibpot-1)))
lagtxppib_dep <- (lagppib_dep-lag2ppib_dep)/lag2ppib_dep
lagtxpibpot_dep <- (lagpibpot_dep-lag2pibpot_dep)/lag2pibpot_dep

update(lagppib_dep) <- ppib_dep
update(lag2ppib_dep) <- lagppib_dep
update(lagpibpot_dep) <- pibpot_dep
update(lag2pibpot_dep) <- lagpibpot_dep
update(lagtdeppp) <- tdeppp
update(lag2tdeppp) <- lagtdeppp

tdepp_e <- dep_prim/pibpot/ppib
update(lagtdepp_e) <- tdepp_e

# l'impulsion sur les recettes est définie comme la variation du tpo
rec_po <- tpo * vpib
update(lagtpo) <- tpo

# l'impulsion sur les intérêts (le multiplicateur est suceptible d'être nul (og_int))
# charge d'intérêts nette de la part BC
cin  <- r_app * lagdette - dettebc * lagdette * r_bc
cinp <- cin / vpib
update(lagcinp) <- cinp
update(lag2cinp) <- lagcinp

ib_cin <- lagcinp - lag2cinp

# règle budgétaire --------------
spp <- (rec_po - dep_prim)/vpib + taut[step]
update(lagspp) <- spp
update(lag2spp) <- lagspp
spstar <- (ecart_c+r_dette*(dstar-dstarr))/(1+gpot_v+infstar) * dstar

ib_fr <- tpo_sstar * (lagspp - spstar) + tpo_1sstar * (lagspp - lag2spp) + tpo_sstar2 * (lagspp - spstar)^2 +
         tpo_dstar * (lagdettep - dstar) + tpo_1dstar * (lagdettep - lag2dettep) + tpo_dstar2 * (lagdettep - dstar)^2 +
         tpo_og*(lagog) + tpo_og2*(lagog)^2 + tpo_1og*(lagog - lag2og) +
         tpo_ec*(ecart_c - ecstar) +
         tpo_inf*(lagtxppib- infstar) +
         phi[step]

ib_sup <- max(min(ib_fr, ibcap), -ibcap)

ib_spont_dep <- tx_deriv_dep[step] * lagtdeppp +
  lagtdeppp * (lagppib_dep*lagpibpot_dep / lagpibpot / lagppib - lag2ppib_dep * lag2pibpot_dep/ lag2pibpot / lag2ppib)
ib_dep <- ib_spont_dep + (1-pcpo) * ib_sup
tdeppp <- lagtdeppp + ib_dep

ib_spont_po <-  - tx_deriv_po[step] * lagtpo
ib_po <- ib_spont_po + pcpo * ib_sup
tpo <- lagtpo - ib_po

# fonction de perte ----------------------
# utilisée pour la partie optimisation
update(loss_d_only) <- loss_d_only + loss_d * (step >= loss_t)*(lagdettep - dstar)^2
initial(loss_d_only) <- 0
update(loss_no_d) <- loss_no_d + og^2/(1+loss_df)^(step-1) + loss_dog * (og-lagog)^2 + loss_ib * (ib_sup)^2
initial(loss_no_d) <- 0

# sorties ------
# dans dust c'est un peu capilotracté
# sans output, on est obligé de faire un update()
# et de définir une CI nulle
# on pense ensuite à décaler les variables pour les utilisations
#
update(dettep_o) <- dette / vpib
update(og_o) <- og
update(txppib_o) <- ppib/lagppib-1
update(gpib_o) <- qpib/lagqpib-1
update(gpot_o) <- gpot_v
update(spp_o) <- spp
update(ci_o) <- cinp
update(tpo_o) <- tpo
update(tdep_o) <- tdepp_e * pibpot/qpib
update(tdeppp_o) <- tdepp_e
update(ib_dep_o) <- ib_dep
update(ib_po_o) <- ib_po
update(ib_o) <- ib_po + ib_dep
update(qpib_o) <- qpib
update(ppib_o) <- ppib
update(tcho_o) <- tcho
update(tvnairu_o) <- tvnairu
update(spstar_o) <- spstar
update(r_app_o) <- r_app
update(loss_o) <- loss_no_d + loss_d_only
update(loss_nd_o) <- loss_no_d

initial(txppib_o) <- 0
initial(gpib_o) <- 0
initial(gpot_o) <- 0
initial(spp_o) <- 0
initial(og_o) <- 0
initial(dettep_o) <- 0
initial(ci_o) <- 0
initial(tpo_o) <- 0
initial(tdep_o) <- 0
initial(ib_dep_o) <- 0
initial(ib_o) <- 0
initial(ib_po_o) <- 0
initial(qpib_o) <- 0
initial(ppib_o) <- 0
initial(tcho_o) <- 0
initial(tvnairu_o) <- 0
initial(tdeppp_o) <- 0
initial(spstar_o) <- 0
initial(r_app_o) <- 0
initial(loss_o) <- 0
initial(loss_nd_o) <- 0

# init ----------------
i_lagog <- user(0)
i_lagtcho <- user(0)
i_lag2tcho <- user(0)
i_lagtvnairu <- user(0)
i_lag2og <- user(0)
i_lagqpib <- user(1)
i_lag2qpib <- user(1)
i_lag3qpib <- user(1)
i_lagppib <- user(1)
i_lag2ppib <- user(1)
i_lag3ppib <- user(1)
i_lagpibpot <- user(1)
i_lag2pibpot <- user(1)
i_lagtpo <- user(0.5)
i_lagtdeppp <- user(0.5)
i_lagr_app <- user(0.04)
i_lagr_inst <- user(0.04)
i_lagdette <- user(1)
i_lagspp <- user(0.0)
i_lag2spp <- user(0.0)
i_lagecart_c <- user(0.0)
i_ogn <- user(0) # rnorm(0, ogn_sigma)+ogn*(rnorm(0, ogn_sigma) + ogn*(rnorm(0, ogn_sigma)+ogn*rnorm(0, ogn_sigma)))
i_lagcip <- user(0.02)
i_lag2cip <- user(0.02)
i_lag2dette <- user(1)
i_lag2tdeppp <- user(0.5)

## Initial conditions
# initial(og) <- init_og
initial(lagog) <- i_lagog
initial(lag2og) <- i_lag2og
initial(lagtcho) <- i_lagtcho
initial(lag2tcho) <- i_lag2tcho
initial(lagtvnairu) <- i_lagtvnairu
initial(og_noise) <- i_ogn
initial(lagppib) <- i_lagppib
initial(lag2ppib) <- i_lag2ppib
initial(lag3ppib) <- i_lag3ppib
initial(lagppib_dep) <- i_lagppib
initial(lag2ppib_dep) <- i_lag2ppib
initial(lagqpib) <- i_lagqpib
initial(lag2qpib) <- i_lag2qpib
initial(laggpib) <- i_lag2qpib/i_lag3qpib - 1
initial(lagpibpot) <- i_lagpibpot
initial(lagpibpot_dep) <- i_lagpibpot
initial(lag2pibpot_dep) <- i_lag2pibpot
initial(lag2pibpot) <- i_lag2pibpot
initial(lagtpo) <- i_lagtpo
initial(lagtdepp_e) <- i_lagtdeppp
initial(lagtdeppp) <- i_lagtdeppp
initial(lag2tdeppp) <- i_lag2tdeppp
initial(lagr_app) <- i_lagr_app
initial(lagr_inst) <- i_lagr_inst
initial(lagdette) <- i_lagdette
initial(lag2dette) <- i_lag2dette
initial(lagspp) <- i_lagspp
initial(lag2spp) <- i_lag2spp
initial(lagcinp) <- i_lagcip
initial(lag2cinp) <- i_lag2cip
initial(lagecart_c) <- i_lagecart_c

# params -------------
periods <- user(50)

## parameters
og_lag   <- user(0)
og_mce   <- user(0.29)
og_ddpot   <- user(0.29)
tvnairu_mce   <- user(0.13)
tcho_okun   <- user(0.33)
tcho_mce   <- user(0.27)
inf_tcho <- user(0.1)
inf_lag <- user(0.0)
inf_mce <- user(0.68)
infdep_mce <- user(0.5)
infdep_star <- user(0.5)
potdep_star <- user(0.5)
potdep_mce <- user(0.5)
pot_r   <- user(0.15)
og_dep  <- user(0.7)
og_po  <-  user(0.7)
og_cin  <- user(0)
infstar <- user(0.0175)
r_mat   <- user(8)
dstar <- user(0.6)
dstarr <- user(0.6)
nairu <- user(0.07)
tpo_sstar <- user(0.4)
tpo_dstar <- user(0.1)
tpo_sstar2 <- user(0)
tpo_1sstar <- user(0)
tpo_dstar2 <- user(0)
tpo_1dstar <- user(0)
tpo_1og <- user(0)
tpo_og <- user(0)
tpo_og2 <- user(0)
tpo_ec <- user(0)
tpo_inf <- user(0)
pcpo <- user(1)
loss_df <- user(0.02)
loss_ib <- user(0)
ogn_ar <- user(0)
ogn_sigma <- user(0.005)
r_dette <- user(0.01)
ibcap <- user(0.01)
ecstar <- user(-0.015)
ec_mce <- user(0.1)
dettebc <- user(0)
r_bc <- user(0)

# loss
loss_dog <- user(0)
loss_d <- user(0.5)
loss_t <- user(20)

# controls&exogenes -----------
dette_oneoff[] <- user(0)
phi[] <- user(0)
taut[] <- user(0)
gpot[] <- user(0)
tx_deriv_dep[] <- user(0)
tx_deriv_po[] <- user(0)

dim(dette_oneoff) <- periods
dim(phi) <- periods
dim(gpot) <- periods
dim(taut) <- periods
dim(tx_deriv_dep) <- periods
dim(tx_deriv_po) <- periods


#==============================================
#  functions
#----------------------------------------------
#  reworked function to speed up the bootstrap analysis

#----------------------------------------------------
#  Make Tables with specified categoriges
#  Enables fixed format output and zero entries
#----------------------------------------------------

Table <- function(x,Labels){
  #=c("D614","G614") ){
  xx <- rep(0, length(Labels) )
  names(xx) <- Labels
  xx.t <- table( x )
  xx[ names(xx.t) ] <- xx.t
  return(xx)
}

#=================================================================
#  Compute isotonic logistic regression, using the pava rountine
#  from Iso
#  Input:  yy     matrix of successes, failures, ordered in time
#          decreasing = F
#          theta  initial value, default is zero
#          err    L1 error on estimated probs.  Recommended 10e-5
#          eps    numerical stability for prob going to 0 or 1
IsoLog <- function(Y,decreasing=F,theta=0,err=1e-5,min.w=1e-6){
  G <- 1
  theta0 <- theta
  pp <- exp(theta0)/(1+exp(theta0))

  yy   <- Y[,1]
  nn   <- apply(Y,1,sum)

  while ( G > err ){
    ww <- nn*pmax( pp*(1-pp), min.w )  # ensure the weights are always larger than eps
    zz <- theta0 + (yy-nn*pp)/ww
    theta1 <- pava(zz,ww,decreasing=decreasing)
    p1 <- exp(theta1)/(1+exp(theta1))
    G <- sum( abs(p1-pp) )
    #
    theta0 <- theta1
    pp <- p1
  }
  return(theta1)
}

#===================================================================================
#  fit a Ashaped
Ashape <- function(Y,lmode,theta=0,err=1e-5,min.w=1e-6){
  G <- 1
  theta0 <- theta
  pp <- exp( theta0)/( 1+exp(theta0) )

  yy <- Y[,1]
  nn <- apply(Y,1,sum)
  iidx <- (1:length(yy)) < (lmode+1)

  while ( G > err ){
    ww <- nn*pmax( pp*(1-pp), min.w )  # ensure the weights are always larger than eps
    zz <- theta0 + (yy-nn*pp)/ww
    theta1 <- pava(zz[iidx],ww[iidx],decreasing=F)
    theta2 <- pava(zz[!iidx],ww[!iidx],decreasing=T)
    ttheta <- c(theta1,theta2)
    p1 <- exp(ttheta)/(1+exp(ttheta))
    G <- sum( abs(p1-pp) )
    #
    theta0 <- ttheta
    pp <- p1
  }
  loglik <- sum(yy*log(pp) + (nn-yy)*log(1-pp))
  return( list( theta=theta0,p=pp,loglik=loglik ) )
}

Ashape2 <- function(Y,Lmode=c(1,dim(Y)[1]),theta=0,err=1e-6,min.w=1e-6){
  lmode <- Lmode[1]:Lmode[2]
  Ffit <- vector("list",length = length(lmode))
  Ffit[[1]] <- Ashape(Y,lmode[1],theta)
  for ( k in 2:length(lmode) ){
    Ffit[[k]] <- Ashape(Y,lmode=lmode[k],Ffit[[k-1]]$theta)
  }
  iidx <- which.max(sapply(Ffit, FUN=function(x) x$loglik))
  return( list( theta=Ffit[[iidx]]$theta,
                p=Ffit[[iidx]]$p,
                loglik=Ffit[[iidx]]$loglik,
                mode = lmode[iidx] ) )
}

#===================================================================
#  Function to compute from one bootstrap sample
#  To be used in parallel
#  input: empty vector, needed to get it to run in parallel
#         y vector of 0-1
#         ddate: vector of dates
#         p0: probability under the null distribution
#         Decreasing: test decresing? T=yes, F=no
#         Labels: vector of labels to be used in Table
#===================================================================
Boot.Pval <- function(zz,yy,ddates,p0,Decreasing=F,Labels=c("M","W")){

    y.b <- sample(yy, replace = F)  # permutation of strain labels
    ddt.b <- split(y.b,ddates) # split by date
    BB.b <- t( sapply( ddt.b , Table, Labels=Labels ) )

    ffit.b <- IsoLog( as.matrix( BB.b ), decreasing=Decreasing)
    p.b <- exp(ffit.b)/(1+exp(ffit.b))

    # calculate relative loglikelihood
    LL <- sum( BB.b[,1]*log(p.b/p0) + BB.b[,2]*log((1-p.b)/(1-p0)) )

    return(LL)

  }


# Start point bias plots
    # Set overlap values
    overlaps <- seq(0,1.5,0.025)
    # Run RRW based on overlaps and bias
    # Unbiased start point
    unbiased <- getMomentsOfRRWoverlap(overlaps,boundary=10, startValue=0, loops = 500,noiseSD = 2)
    #Positive start point bias
    posbias1 <- getMomentsOfRRWoverlap(overlaps,boundary=10, startValue=0.25, loops = 500,noiseSD = 2)
    posbias2 <- getMomentsOfRRWoverlap(overlaps,boundary=10, startValue=0.50, loops = 500,noiseSD = 2)
    # Negative start point bias
    negbias1 <- getMomentsOfRRWoverlap(overlaps,boundary=10, startValue= -0.25, loops = 500,noiseSD = 2)
    negbias2 <- getMomentsOfRRWoverlap(overlaps,boundary=10, startValue= -0.50, loops = 500,noiseSD = 2)

    # Generate plot found in paper
     with(unbiased[unbiased$correct==T,],plot(pCross ~ overlap, pch=16, col="black",frame.plot=F,xlab="Overlap",ylab="pHO",ylim=c(0,1),las=1,cex=0.75,xaxt='n'))
       abline(0.5,0,lty=2)
       abline(v=1,lty=2)
       axis(1, at = seq(0, 1.5, by = .1), las=2)
       with(posbias1[posbias1$correct==T,],points(pCross ~ overlap, pch=16, col="sky blue",cex=0.75))
       with(posbias2[posbias2$correct==T,],points(pCross ~ overlap, pch=16, col="blue",cex=0.75))
       with(negbias1[negbias1$correct==T,],points(pCross ~ overlap, pch=16, col="light pink",cex=0.75))
       with(negbias2[negbias2$correct==T,],points(pCross ~ overlap, pch=16, col="red",cex=0.75))
       legend(0.2,0.4, legend=c("0.50","0.25","0 (unbiased)","-0.25","-0.50"),title="Start Point",col=c("blue", "skyblue","black","light pink","red"),pch=16,cex=0.75)
       # Write to a pdf
       dev.copy(pdf,"pHVOxOverlapVal.pdf")
       dev.off()


#########################################################
#install.packages("argparse")
library(argparse)
library(mdatools)
parser <- ArgumentParser()

parser$add_argument("--file", type="character", default= "no_input_return_error",
    help="file containing variant frequencies")
parser$add_argument("--plot_bool", type="logical", default= FALSE,
    help="determines whether pdf plot approx_plot.pdf is produced or not")
parser$add_argument("--var_calling_threshold", type="double", default= 0.03,
    help="variant calling threshold")
parser$add_argument("--Nb_min", type="integer", default= 1,
    help="Minimum bottleneck value considered")
parser$add_argument("--Nb_max", type="integer", default= 200,
    help="Maximum bottleneck value considered")
parser$add_argument("--confidence_level", type="double", default= .95,
    help="Confidence level (determines bounds of confidence interval)")
args <- parser$parse_args()
if (args$file == "no_input_return_error" ) { stop("file with lists of donor and recipient frequencies is a required argument.", call.=FALSE)}

plot_bool  <- args$plot_bool
var_calling_threshold  <- args$var_calling_threshold
Nb_min <- args$Nb_min # Minimum bottleneck size we consider.
Nb_max <-  args$Nb_max
confidence_level <- args$confidence_level
donor_and_recip_freqs_observed <- read.table(args$file) #donor_freqs_observed <- read.table(args[1])
donor_freqs_observed <- as.data.frame(donor_and_recip_freqs_observed[, 1]) # We read in and save the list of donor frequencies
recipient_freqs_observed <- as.data.frame(donor_and_recip_freqs_observed[, 2]) # We read in and save the list of recipient frequencies
n_variants <- nrow(donor_and_recip_freqs_observed)#nrow(donor_freqs_observed) # number of variants

#######################################################
#######################################################
generate_log_likelihood_function_approx <- function(donor_freqs_observed, recipient_freqs_observed, Nb_min, Nb_max, var_calling_threshold, confidence_level)
{

    num_NB_values <- Nb_max -Nb_min + 1
likelihood_matrix <- matrix( 0, n_variants, num_NB_values)#This matrix stores the likelihood values for each variant frequency and bottleneck size.
log_likelihood_matrix <- matrix( 0, n_variants, num_NB_values)#This matrix stores the log likelihood values for each variant frequency and bottleneck size.  We can sum the log likelihoods of all the variant alleles to get the total log likelihood.
log_likelihood_function <- matrix( 0, Nb_max )

for (i in 1:n_variants) {for (j in 1:num_NB_values) {
  Nb_val <- (j - 1 + Nb_min)
    nu_donor <- donor_freqs_observed[i, 1]
    nu_recipient <- recipient_freqs_observed[i, 1]
    if (recipient_freqs_observed[i, 1] >= var_calling_threshold)
        { # implement variant calling threshold

for (k in 0:Nb_val){
      likelihood_matrix[i, j] <- likelihood_matrix[i, j] +
  (dbeta(nu_recipient, k, (Nb_val - k))*dbinom(k, size=Nb_val, prob= nu_donor))
        }
  log_likelihood_matrix[i,j] = log(likelihood_matrix[i, j])
   }
 if (recipient_freqs_observed[i, 1] < var_calling_threshold)
        { # implement variant calling threshold

   likelihood_matrix[i, j] = 0
   log_likelihood_matrix[i,j] = 0
   for (k in 0:Nb_val){   likelihood_matrix[i, j] <- likelihood_matrix[i, j] +
               (pbeta(var_calling_threshold, k, (Nb_val - k))*dbinom(k, size=Nb_val, prob= nu_donor))
                      }
log_likelihood_matrix[i,j] = log(likelihood_matrix[i, j])
            }
# Now we sum over log likelihoods of the variants at different loci to get the total log likelihood for each value of Nb
log_likelihood_function[ Nb_val] <- log_likelihood_function[ Nb_val] + log_likelihood_matrix[i,j]
# Shifting entries of log_likelihood function to ensure plot begins at proper point on x axis
}}

return(log_likelihood_function)

}
#################################################################
#################################################################
restrict_log_likelihood <- function(log_likelihood_function, Nb_min, Nb_max) # restricts log likelihood to the interval of interst
{
for (h in 1:(Nb_min )){
    if(h< Nb_min)
    {log_likelihood_function[h] = - 999999999}        # kludge for ensuring that these values less than Nb_min don't interfere with our search for the max of log likelihood in the interval of Nb_min to Nb_max
    }

return(log_likelihood_function)
#print(erfinv(percent_confidence_interval)*sqrt(2))
}

return_bottleneck_size <- function(log_likelihood_function, Nb_min, Nb_max)
{ max_log_likelihood = which(log_likelihood_function == max(log_likelihood_function))

    return(max_log_likelihood)
}

return_CI_lower <- function(log_likelihood_function, Nb_min, Nb_max) ## returns lower bound of confidence interval
  { max_log_likelihood = which(log_likelihood_function == max(log_likelihood_function))  ## This is the point on the x-axis (bottleneck size) at which log likelihood is maximized
max_val =  max(log_likelihood_function)  ## This is the maximum value of the log likelihood function, found when the index is our bottleneck estimate
CI_height = max_val - erfinv(confidence_level)*sqrt(2)  # This value (  height on y axis) determines the confidence intervals using the likelihood ratio test

    CI_index_lower = Nb_min
  CI_index_upper = max_log_likelihood
for (h in 1:Nb_min){
    if(h< Nb_min)
    {log_likelihood_function[h] = NA}   #  Removing parameter values less than Nb_min from plot
        }
## above loop just enforces our minimum bottleneck cutoff
for (h in Nb_min:max_log_likelihood){
    test1 = (log_likelihood_function[CI_index_lower] - CI_height) * (log_likelihood_function[CI_index_lower] - CI_height)
    test2 = (log_likelihood_function[h] - CI_height) * (log_likelihood_function[h] - CI_height)
        if( test2 < test1){  CI_index_lower = h  }
}
if(  (log_likelihood_function[CI_index_lower] - CI_height) > 0  ){CI_index_lower = CI_index_lower - 1   }
# above loops use likelihood ratio test to find lower confidence interval
for (h in max_log_likelihood:Nb_max)
{       test1 = (log_likelihood_function[CI_index_upper] - CI_height) * (log_likelihood_function[CI_index_upper] - CI_height)
      test2 = (log_likelihood_function[h] - CI_height) * (log_likelihood_function[h] - CI_height)
        if( test2 < test1  ){CI_index_upper = h   }
}
if(  (log_likelihood_function[CI_index_upper] - CI_height) > 0  ){CI_index_upper = CI_index_upper + 1   }

  return(CI_index_lower)

  }

return_CI_upper <- function(log_likelihood_function,  Nb_min, Nb_max) ## returns upper bound of confidence interval
  { max_log_likelihood = which(log_likelihood_function == max(log_likelihood_function))  ## This is the point on the x-axis (bottleneck size) at which log likelihood is maximized
   max_val =  max(log_likelihood_function)  ## This is the maximum value of the log likelihood function, found when the index is our bottleneck estimate
   CI_height = max_val - erfinv(confidence_level)*sqrt(2)  # This value (  height on y axis) determines the confidence intervals using the likelihood ratio test

    CI_index_lower = Nb_min
  CI_index_upper = max_log_likelihood
for (h in 1:Nb_min){
    if(h< Nb_min)
    {log_likelihood_function[h] = NA}   #  Removing parameter values less than Nb_min from plot
        }
## above loop just enforces our minimum bottleneck cutoff
for (h in Nb_min:max_log_likelihood){
    test1 = (log_likelihood_function[CI_index_lower] - CI_height) * (log_likelihood_function[CI_index_lower] - CI_height)
    test2 = (log_likelihood_function[h] - CI_height) * (log_likelihood_function[h] - CI_height)
        if( test2 < test1){  CI_index_lower = h  }
}
if(  (log_likelihood_function[CI_index_lower] - CI_height) > 0  ){CI_index_lower = CI_index_lower - 1   }
# above loops use likelihood ratio test to find lower confidence interval
for (h in max_log_likelihood:Nb_max)
{       test1 = (log_likelihood_function[CI_index_upper] - CI_height) * (log_likelihood_function[CI_index_upper] - CI_height)
      test2 = (log_likelihood_function[h] - CI_height) * (log_likelihood_function[h] - CI_height)
        if( test2 < test1  ){CI_index_upper = h   }
}
if(  (log_likelihood_function[CI_index_upper] - CI_height) > 0  ){CI_index_upper = CI_index_upper + 1   }

  return(CI_index_upper)

  }

log_likelihood_function <- generate_log_likelihood_function_approx(donor_freqs_observed, recipient_freqs_observed, Nb_min, Nb_max, var_calling_threshold, confidence_level)
log_likelihood_function <- restrict_log_likelihood(log_likelihood_function, Nb_min, Nb_max)
bottleneck_size <- return_bottleneck_size(log_likelihood_function,  Nb_min, Nb_max)
CI_index_lower <- return_CI_lower(log_likelihood_function,  Nb_min, Nb_max)
CI_index_upper <- return_CI_upper(log_likelihood_function,  Nb_min, Nb_max)

        ##########################
##############################################################################################  ABOVE THIS LINE DETERMINES PEAK LOG LIKELIHOOD AND CONFIDENCE INTERVALS
 # Npw we plot the result
if(plot_bool == TRUE)
{pdf(file="approx_plot.pdf")
plot(log_likelihood_function)
abline(v = bottleneck_size, col="black" )  # Draws a verticle line at Nb value for which log likelihood is maximized
abline(v = CI_index_lower, col="green" ) # confidence intervals
abline(v = CI_index_upper, col="green" )
dev.off()
}

print("Bottleneck size")
print(bottleneck_size)
print("confidence interval left bound")
print(CI_index_lower)
print("confidence interval right bound")
print(CI_index_upper)


#' @title
#' Reduced-rank ridge regression with rank and tuning parameter selected by cross validation
#'
#' @description
#' This function performs reduced-rank ridge regression with the rank and the
#' tuning parameter selected by cross validation.
#'
#' @usage
#' RRS.cv(Y, X, nfold = 10, rankmax = min(dim(Y), dim(X)), nlam = 100,
#'         lambda = seq(0, 100, length = nlam), norder = NULL,
#'         nest.tune = FALSE, fold.drop = 0)
#'
#' @param Y response matrix.
#' @param X design matrix.
#' @param nfold the number of folds used in cross validation. Default is 10.
#' @param rankmax the maximum rank allowed.
#' @param nlam the number of tuning parameter candidates. Default is 100.
#' @param lambda the tuning sequence of length \code{nlam}.
#' @param norder a vector of length n that assigns samples to multiple folds
#'               for cross validation. Default is NULL and then \code{norder} will be
#'               generated randomly.
#' @param nest.tune a logical value to specify whether to tune the rank and lambda
#'                  in a nested way. Default is FALSE.
#' @param fold.drop the number of folds to drop. Default is 0.
#'
#'
#' @return The function returns a list:
#'   \item{cr_path}{a matrix displays the path of model selection.}
#'   \item{C}{the estimated low-rank coefficient matrix.}
#'   \item{rank}{the selected rank.}
#'   \item{lam}{the selected tuning parameter for ridge penalty.}
#'
#' @references Mukherjee, A., & Zhu, J. (2011).
#' Reduced Rank Ridge Regression and Its Kernel Extensions.
#' Statistical analysis and data mining,
#' 4(6), 612–622.
#'
# @importFrom rrpack rrs.fit
#' @export
RRS.cv<-function(Y, X, nfold = 10, rankmax = min(dim(Y), dim(X)), nlam = 100,
                     lambda = seq(0, 100, length = nlam), norder = NULL, nest.tune = FALSE,
                     fold.drop = 0)
{
  p <- ncol(X)
  q <- ncol(Y)
  n <- nrow(Y)
  ndel <- round(n/nfold)
  if (is.null(norder)) {
    norder <- sample(1:n, n)
  }
  cr_path <- array(dim = c(nlam, rankmax + 1, nfold), NA)
  for (f in 1:nfold) {
    if (f != nfold) {
      iddel <- norder[(1 + ndel * (f - 1)):(ndel * f)]
    }
    else {
      iddel <- norder[(1 + ndel * (f - 1)):n]
    }
    ndel <- length(iddel)
    nf <- n - ndel
    Xf <- X[-iddel, ]
    Xfdel <- X[iddel, ]
    Yf <- Y[-iddel, ]
    Yfdel <- Y[iddel, ]
    for (ll in 1:nlam) {
      ini <- rrs.fit(Yf, Xf, lambda = lambda[ll], nrank = rankmax)
      C_ls <- ini$coef.ls
      A <- ini$A
      tempFit <- Xfdel %*% C_ls
      tempC <- matrix(nrow = nrow(A), ncol = nrow(A), 0)
      for (i in 1:rankmax) {
        tempC <- tempC + A[, i] %*% t(A[, i])
        cr_path[ll, i + 1, f] <- sum((Yfdel - tempFit %*%
                                        tempC)^2)
      }
      cr_path[ll, 1, f] <- sum(Yfdel^2)
    }
  }
  index <- order(apply(cr_path, 3, sum))[1:ifelse(nfold > 5,
                                                  nfold - fold.drop, nfold)]
  crerr <- apply(cr_path[, , index], c(1, 2), sum)/length(index) *
    nfold
  minid <- which.min(crerr)
  minid2 <- ceiling(minid/nlam)
  rankest <- minid2 - 1
  minid1 <- minid - nlam * (minid2 - 1)
  if (nest.tune == TRUE) {
    lambda <- exp(seq(log(lambda[max(minid1 - 1, 1)] + 1e-04),
                      log(lambda[min(minid1 + 1, nlam)]), length = nlam))
    cr_path <- array(dim = c(nlam, rankmax + 1, nfold), NA)
    ndel <- round(n/nfold)
    for (f in 1:nfold) {
      if (f != nfold) {
        iddel <- norder[(1 + ndel * (f - 1)):(ndel *
                                                f)]
      }
      else {
        iddel <- norder[(1 + ndel * (f - 1)):n]
        ndel <- length(iddel)
      }
      nf <- n - ndel
      Xf <- X[-iddel, ]
      Xfdel <- X[iddel, ]
      Yf <- Y[-iddel, ]
      Yfdel <- Y[iddel, ]
      for (ll in 1:nlam) {
        ini <- rrs.fit(Yf, Xf, lambda = lambda[ll], nrank = rankmax)
        C_ls <- ini$coef.ls
        A <- ini$A
        tempFit <- Xfdel %*% C_ls
        tempC <- matrix(nrow = nrow(A), ncol = nrow(A),
                        0)
        for (i in 1:rankmax) {
          tempC <- tempC + A[, i] %*% t(A[, i])
          cr_path[ll, i + 1, f] <- sum((Yfdel - tempFit %*%
                                          tempC)^2)
        }
        cr_path[ll, 1, f] <- sum(Yfdel^2)
      }
    }
    index <- order(apply(cr_path, 3, sum))[1:ifelse(nfold >
                                                      5, nfold - fold.drop, nfold)]
    crerr <- apply(cr_path[, , index], c(1, 2), sum)/length(index) *
      nfold
    minid <- which.min(crerr)
    minid2 <- ceiling(minid/nlam)
    rankest <- minid2 - 1
    minid1 <- minid - nlam * (minid2 - 1)
  }
  if (rankest == 0) {
    list(cr_path = cr_path, CRE = crerr, norder = norder,
         C = matrix(nrow = p, ncol = q, 0), rank = 0)
  }
  else {
    fit <- rrs.fit(Y, X, lambda = lambda[minid1], nrank = rankmax)
    C_lslam <- fit$coef.ls
    A <- fit$A
    list(cr_path = cr_path, CRE = crerr, ID = c(minid1, minid2),
         norder = norder, C = C_lslam %*% A[, 1:rankest] %*%
           t(A[, 1:rankest]), rank = rankest, lam = lambda[minid1])
  }
}


# <-- encoding UTF-8 -->
# Empirical study on NHSS dataset (county level)
# -------------------------------------
## DOC STRING
#
#
# Tianhao Zhao (GitHub: Clpr)
# Dec 2018
# -------------------------------------

# -------------------------------------
## SECTION 0: ENVIRONMENT
library(sqldf)  # sql enquiry
library(openxlsx)  # easy xlsx IO
source("./src/mathtools.r")  # math tools, e.g. Lorenz, Theil, descriptive stats
source("./src/mapplots.r")  # tools of map plots, using ggplot2
# environment par dict
envNHSS$Output <- "./output/proc_2_NHSS/"  # alter output directory

# -------------------------------------
## SECTION 1: GENERAL DESCRIPTIVE STATISTICS
cat("\nGeneral descriptive statistics of NHSS dataset:\n")
cat("(please note: the income now is the residuals of income ~ edu)\n")
cat("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
cat("PLEASE NOTE: because we have normalized PCA components, we only do descriptive statistics for income & edu!!!!!!")
cat("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
# --------------
li_Descript_NHSS <- list()  # a list of descriptive statistics
for(tmpYname in envNHSS$Ynames){
    # 1.1 get a namelist of all numeric variables, both Y and Xcore
    tmp <- c( tmpYname, envNHSS$Xcore )
    li_Descript_NHSS[[tmpYname]] <- func_DescrStat( li_Dat_NHSS[[tmpYname]][,tmp] )
    # 1.2 print
    print(li_Descript_NHSS[[tmpYname]]); cat("-----------------------------------------\n")
}
# output to xlsx
openxlsx::write.xlsx(li_Descript_NHSS, paste(sep="",envNHSS$Output,"Descript_NHSS.xlsx")  )

# -------------------------------------
## SECTION 2: COUNTY-LEVEL LORENZ CURVE & GINI COEF
# NOTE: in this section, we use counties as units/individuals to plot Lorenz curves in every year (98,03,08),
#       meanwhile, Gini coefficients in the three years are calculated;
#       later, in following sections, we will add other inequlity indices

# 2.0 prepare a df for Gini coefs
df_Gini_NHSS <- data.frame(
    year = c(1998, 2003, 2008),   # years
    counties = c( sum(df_NHSS$year == 1998),sum(df_NHSS$year == 2003),sum(df_NHSS$year == 2008)  ),  # number of counties in every year's data
    illnessratio = NA, # prevalence
    illnessday = NA,
    chronicratio = NA
)

# 2.1 use functions in mathtools.r to compute Lorenz Curve & Gini coefficients
# NOTE: please refer to mathtools.r for how we calculated the both
#       and, if you want to see the plotting of Lorenz Curve, you may DIY (using tmp$LorenzX and tmp$LorenzY)
for(tmpYear in df_Gini_NHSS$year){
    for(tmpYname in envNHSS$Ynames){
        # 2.1.1 compute Lorenz Curve & Gini in specific year
        # eval(parse(text=paste(sep="",
        #     "tmp <- LorenzCurve( df_NHSS$",tmpYname,"[df_NHSS$",envNHSS$flagT," == ",tmpYear,"] )"
        # )))
        tmp <- li_Dat_NHSS[[tmpYname]]
        tmp <- LorenzCurve( tmp[ tmp[,envNHSS$flagT] == tmpYear ,tmpYname] )
        # 2.1.2 save Gini coef
        df_Gini_NHSS[,tmpYname][df_Gini_NHSS$year == tmpYear] <- tmp$GiniIdx
    }
}
# print & output to xlsx
print(df_Gini_NHSS)

# NOTE: the results will be output to a csv file with other kinds of indices, later

# -------------------------------------
## SECTION 3: COUNTY-LEVEL THEIL I & II INDEXES
# NOTE: in this section, we use counties as units/individuals to calculate Theil-I & Theil-II indexes.
#       meanwhile, indices in the three years are calculated;
#       later, in following sections, we will add other inequlity indices

# 3.0 prepare two df
df_Theil1_NHSS <- data.frame(
    year = c(1998, 2003, 2008),   # years
    counties = c( sum(df_NHSS$year == 1998),sum(df_NHSS$year == 2003),sum(df_NHSS$year == 2008)  ),  # number of counties in every year's data
    illnessratio = NA, # prevalence
    illnessday = NA,
    chronicratio = NA
)
df_Theil2_NHSS <- df_Theil1_NHSS

# 3.1 use functions in mathtools.r to do computing
for(tmpYear in df_Gini_NHSS$year){  # loop on years
    for(tmpYname in envNHSS$Ynames){  # loop on health outcomes
        # slice a temp dataset
        tmpdf <- li_Dat_NHSS[[tmpYname]]
        # 3.1.1 compute Theil-I & Theil-II in a specific year
        tmp <- Theil( tmpdf[,tmpYname][tmpdf[,envNHSS$flagT] == tmpYear], Type = "T" )
        df_Theil1_NHSS[,tmpYname][ df_Theil1_NHSS$year == tmpYear ] <- tmp
        tmp <- Theil( tmpdf[,tmpYname][tmpdf[,envNHSS$flagT] == tmpYear], Type = "L" )
        df_Theil2_NHSS[,tmpYname][ df_Theil2_NHSS$year == tmpYear ] <- tmp
    }
}

# 3.2 print info
cat("\nThe county-level Theil-I type index in every year of NHSS dataset are:\n")
print(df_Theil1_NHSS)
cat("\nThe county-level Theil-II type index in every year of NHSS dataset are:\n")
print(df_Theil2_NHSS)

# NOTE: the results will be output to a csv file with other kinds of indices, later

# -------------------------------------
## SECTION 4: COUNTY-LEVEL COEF OF VARIANCE (C.V.) & VARIANCE
# NOTE: in this section, we use counties as units/individuals to calculate C.V. & variances
#       meanwhile, indices in the three years are calculated;

# 4.0 prepare two df
df_CoefVar_NHSS <- data.frame(
    year = c(1998, 2003, 2008),   # years
    counties = c( sum(df_NHSS$year == 1998),sum(df_NHSS$year == 2003),sum(df_NHSS$year == 2008)  ),  # number of counties in every year's data
    illnessratio = NA, # prevalence
    illnessday = NA,
    chronicratio = NA
)
df_Variance_NHSS <- df_CoefVar_NHSS

# 4.1 calculation
for(tmpYear in df_Gini_NHSS$year){  # loop on years
    for(tmpYname in envNHSS$Ynames){  # loop on health outcomes
        # slice a temp dataset
        tmpdf <- li_Dat_NHSS[[tmpYname]]
        # 4.1.1 get data vector
        tmp <- tmpdf[,tmpYname][ tmpdf[,envNHSS$flagT] == tmpYear ]
        df_Variance_NHSS[,tmpYname][df_Variance_NHSS$year == tmpYear] <- var(tmp)
        df_CoefVar_NHSS[,tmpYname][df_CoefVar_NHSS$year == tmpYear] <- sd(tmp) / mean(tmp)
    }
}

# 4.2 print info
cat("\nThe county-level coef of variance in every year of NHSS dataset are:\n")
print(df_CoefVar_NHSS)
cat("\nThe county-level variances in every year of NHSS dataset are:\n")
print(df_Variance_NHSS)

# -------------------------------------
## SECTION 5: COLLECT INEQUALITY INDICES & OUTPUT TO CSV
# NOTE: in this section, we firstly bind all kinds of inequality indices to a single dframe,
#       then, output the large dframe to a csv file

# 5.1 table joining
df_InequalIdx_NHSS <- data.frame( Index = rep(c("Gini","Theil-I","Theil-II","Variance","Coef of Variance"), each = 3) )
df_InequalIdx_NHSS <- cbind(
    df_InequalIdx_NHSS, rbind( df_Gini_NHSS, df_Theil1_NHSS, df_Theil2_NHSS, df_Variance_NHSS, df_CoefVar_NHSS )
)
# 5.2 garbage recycling
rm( df_Gini_NHSS, df_Theil1_NHSS, df_Theil2_NHSS, df_Variance_NHSS, df_CoefVar_NHSS )
rm( tmpYear, tmpYname, tmpdf )
# 5.3 output to csv
write.csv( df_InequalIdx_NHSS, file = paste(sep="",envNHSS$Output,"InequalityIdx_NHSS.csv") )

# -------------------------------------
## INTERVAL: MATCH CITY GIS DATA WITH NHSS DATASET
# NOTE: in this section, we merge the China city/County GIS data with df_NHSS;
#       the data have been preloaded by mapplots.r, and named as MapGIScity;
#       in the next section, the dataset will be used to plot
# NOTE: we do not plot PCs -> health outcomes, therefore, just use df_NHSS!
tmpdf <- plyr::join( df_NHSS, MapGIScity[,c("ID","long","lat")], by = "ID" )  # we use GB (country standard) codes of city/county to match by
tmpdf <- na.omit(tmpdf)

cat("\nPlotting ...\n")
# -------------------------------------
## SECTION 6: COLORED MAPS OF HEALTH OUTCOMES
# NOTE: in this section, we plot colored maps for every health outcomes,
#       where every health outcome (in three years) are seen as a pool data,
#       to see the (geographic) distribution (descriptively, intuitively) of the outcomes;
#       using ggplot2, and tools in mapplots.r
#       finally, we output these maps/figures to PDF figures
# NOTE: please refer to mapplots.r for more information & technical details
# NOTE: in final plots, the size of circle indicates the values of independents, where the color indicates the values of health outcomes
#       the plots work to display the relationship between X & Y
# --------
# 6.1 colored maps
# 1998,2003,2008 data are now seen as a pool dataset
for(tmpYname in envNHSS$Ynames){
    # 6.1.1 temp slice: pooled health outcome & province tags
    tmp <- tmpdf[,c(tmpYname,envNHSS$flagProv, "ID", "long", "lat" )]
        # # 6.1.2 averaging by procinve tags (weighted by sample size/population of each county)
        # tmp <- sqldf::sqldf(paste(sep="",
        #       "SELECT DISTINCT SUM(",envNHSS$flagSampleSize," * ",tmpYname,") / SUM(",envNHSS$flagSampleSize,") AS ",tmpYname,
        #       ", ",envNHSS$flagProv," FROM tmp GROUP BY ",envNHSS$flagProv
        # ))
        # 6.1.3 using mapplots.r to get an instance of current map figure
        # eval(parse(text=paste(sep="",
        #     "tmpfig <- func_MapProv( tmp$",tmpYname,", tmp$",envNHSS$flagProv,", vecName = \"",tmpYname,"\" )"
        # )))
    # income -> health outcome
    tmpfig1 <- func_MapCityXY( tmpdf$income, tmpdf[,tmpYname], tmpdf$lat, tmpdf$long,
                               Xname = "income", Yname = tmpYname , FontSize = 5,
                               ColorScale = c("skyblue","red"), CircleScale = c(2,15) )
    # edu -> health outcome
    tmpfig2 <- func_MapCityXY( tmpdf$edu, tmpdf[,tmpYname], tmpdf$lat, tmpdf$long,
                               Xname = "edu", Yname = tmpYname , FontSize = 5,
                               ColorScale = c("skyblue","red"), CircleScale = c(2,15) )
    # output figures
    eval(parse(text=paste(sep="",
       "func_SaveMap2PDF( tmpfig1, \"",envNHSS$Output,"Map_income_2_",tmpYname,".pdf\" ) "
    )))
    eval(parse(text=paste(sep="",
       "func_SaveMap2PDF( tmpfig2, \"",envNHSS$Output,"Map_edu_2_",tmpYname,".pdf\" ) "
    )))
}

# 6.2 print info
cat("\n We have created colored maps for health outcomes and output them to the assigned output directory as PDF figures\n")

# --------------------------------------
# GARBAGE COLLECTION
rm( tmp, tmpfig1, tmpfig2, tmpYname, tmpdf )

#
# # -------------------------------------
# ## SECTION 7: COLORED MAPS OF CORE INDEPENDENTS (INCOME & EDU)
# # NOTE: please refer to mapplots.r for more information & technical details
#
# # 7.1 colored maps
# # 1998,2003,2008 data are now seen as a pool dataset
# for(tmpXname in envNHSS$Xcore){
#     # 7.1.1 temp slice: pooled health outcome & province tags
#     tmp <- df_NHSS[,c(tmpXname,envNHSS$flagProv,envNHSS$flagSampleSize)]
#     # 7.1.2 averaging by procinve tags
#     tmp <- sqldf::sqldf(paste(sep="",
#                               "SELECT DISTINCT SUM(",envNHSS$flagSampleSize," * ",tmpXname,") / SUM(",envNHSS$flagSampleSize,") AS ",tmpXname,
#                               ", ",envNHSS$flagProv," FROM tmp GROUP BY ",envNHSS$flagProv
#     ))
#     # 7.1.3 using mapplots.r to get an instance of current map figure
#     eval(parse(text=paste(sep="",
#                           "tmpfig <- func_MapProv( tmp$",tmpXname,", tmp$",envNHSS$flagProv,", vecName = \"",tmpXname,"\" )"
#     )))
#     # 7.1.4 output figures
#     eval(parse(text=paste(sep="",
#                           "func_SaveMap2PDF( tmpfig, \"",envNHSS$Output,"Map_",tmpXname,".pdf\" ) "
#     )))
# }
#
# # 7.2 print info
# cat("\nWe have created colored maps for core independents (income & edu) and output them to the assigned output directory as PDF figures\n")


#' @title calcShipCategory
#'
#' @description
#' Assigns EPA regulatory ship category from engine displacement
#'
#' @param mainEngineBore Diameter of the main engine cylinder (mm) (vector of
#' numericals)
#' @param mainEngineStroke Main engine stroke length, the distance travelled by
#' the piston in each cycle (mm) (vector of numericals)
#'
#'@details
#'For more information, see Section 3.3.2.1 of the Port Emissions Inventory
#'Guidance.
#'
#' @return \code{shipCategory} (vector of ints). Valid values are: \itemize{
#' \item 1
#' \item 2
#' \item 3
#' }
#'
#' @references
#'\href{https://nepis.epa.gov/Exe/ZyPURL.cgi?Dockey=P1005ZGH.TXT}{EPA. 2009.
#'"Regulatory impact analysis: Control of emissions air pollution from category
#'3 marine diesel engines." Ann Arbor, MI: Office of Transportation and Air
#'Quality. US Environmental Protection Agency.}
#'
#'\href{https://nepis.epa.gov/Exe/ZyPDF.cgi/?Dockey=P10024CN.PDF}{EPA. 2008.
#'"Regulatory impact analysis: Control of emissions of air pollution from
#'locomotive engines and marine compression ignition engines less than 30 liters
#'per cylinder." Ann Arbor, MI: Office of Transportation and Air
#'Quality. US Environmental Protection Agency.}
#'
#' \href{https://nepis.epa.gov/Exe/ZyPDF.cgi?Dockey=P10102U0.pdf}{EPA. 2020.
#' "Port Emissions Inventory Guidance: Methodologies for Estimating
#' Port-Related and Goods Movement Mobile Source Emissions." Ann Arbor, MI:
#' Office of Transportation and Air Quality. US Environmental Protection Agency.}
#'
#' @examples
#' calcShipCategory(400, 240)
#'
#' @export

calcShipCategory<- function(mainEngineBore, mainEngineStroke){

# Calculate displacement in L/cyl
  main.engine.displacement <- ((pi / 4) * (mainEngineBore/10) ^ 2 *
                                    (mainEngineStroke / 10)) / 1000

# Categorize Ships (C1/C2/C3)
shipCategory <- ifelse(main.engine.displacement>= 30, 3,
                                  ifelse(main.engine.displacement >= 7, 2,
                                         ifelse(main.engine.displacement > 0 , 1,
                                                NA)))
return(shipCategory)
}


require(expm)
require(survival)

dyn.load("lk_rec.so")
dyn.load("lk_rec_surv.so")
dyn.load("lk_comp_joint.so")
dyn.load("lk_comp_joint_surv.so")
dyn.load("dlk_comp_joint.so")
dyn.load("dlk_comp_joint_surv.so")

est_incomp_joint_HMcontinuous = function(Y,TT,X,W,tv,dev,k,type=c("norm","bino"),btv=NULL,reltol=1e-6,inits=NULL,
                                         verbose=FALSE,LC=FALSE,surv=TRUE,maxit=5000,stderr=FALSE){

# Estimate the joint model with continuous HM chain described in
# "A shared-parameter continuous-time hidden Markov and survival model for longitudinal data
#  with informative drop-out"
# by F.Bartolucci and A.Farcomeni (2018)
#
# INCOMPLETE DATA CASE
#
# INPUT:
#
# Y       = matrix of responses of dimensino n x Tmax (Tmax = maximum number of individual observations)
# TT      = corresponding matrix of observation times
# X       = array of covariates for longitudinal process of dimension n x TT x nX (nX = number of covariates X)
# W       = array of covariates (including itercept) for survival process of dimension n x nW (nW = number of covariates V)
# tv      = vector of surivival times of lenght n
# dev     = vector of indicators of censoring of lenght n
# k       = number of latent states
# type    = type of response variables (norm or bino)
# btv     = list of equally spaced points defining time windows for the likelihood approximation
# reltol  = relative tollerance for checking log-likelihood convergence
# inits   = list of initial values of the parameters
# verbose = to disply partial output
# LC      = to fit the latent class (instead of latent Markov) model
# surv    = if survival time is provided
#
# OUTPUT:
#
# lk      = maximum log-likelihood at convergence
# Q       = estimated infinitesimal transition matrix
# piv     = estimated vector of initial probabilities
# bev     = estiamted vector of regression coefficients of the longitudinal responses
# nu      = estiamted parater in the baseline hazard function
# xiv     = estimated vector of support points for the latent distribution
# phi     = estimated association coeffiecint for the survivial process
# psiv    = estimated vector of
# si2     = estimated variance parameter
# pP      = time-specific posterior probabilities of each state
# Pi      = discrete time transition matrix
# it      = final number of iterations
# tim     = computing time
# bic     = BIC
# Ji      = observed information matrix (optional)
# sev     = overall vector of standard errors (optional)
# selav   = standard errors for vector lambda (optional)
# seLa    = standard errors for matrix Lambda (optional)
# sebev   = standard errors for bev (optional)
# seeta   = standard error for eta (optional)
# sexiv   = standard errors for xiv (optional)
# sephi   = standard error for phi (optional)
# sepsiv  = standard errors for psi (optional)
# sega    = standard error for gamma (optional)

# Preliminaries
	type = match.arg(type)
	censv = 1-dev
	t0 = proc.time()[3]
	npX = dim(X)[3]          # number of covariates on longitudinal process
	if(surv) npW = ncol(W)   # number of covariates on survival process
	n = dim(Y)[1]            # sample size
	jv = rowSums(!is.na(Y))  # number of longitudinal observations
	mjv = max(jv)            # maximum number of longitudinal observations
	if(!surv) tv = apply(TT,1,max,na.rm=TRUE)

# starting values
	Q = matrix(0,k,k)
	if(!is.null(inits)){
		if(LC) Q = matrix(0,k,k) else Q = inits$Q
		piv = inits$piv
		bev = inits$bev
		if(surv){
			nu = inits$nu
			phi = inits$phi
			psiv = inits$psiv
		}else{
			nu = phi = psiv = NULL
		}
		xiv = inits$xiv
		if(type=="bino") si2 = NULL else si2 = inits$si2
	}
	if(is.null(inits)){
		if(!LC){
			Q = matrix(1/(k-1),k,k)
			diag(Q) = -1
		}
		if(type=="norm"){
			jnk = lm(Y[,1]~.,data=data.frame(X[,1,]))
			xiv = jnk$coef[1]+seq(-k,k,length=k)
			bev = jnk$coef[-1]
			res = jnk$residuals
			si2 = var(res)
		}
		if(type=="bino"){
			yv = as.vector(Y)
			XX = matrix(X,n*mjv,npX)
			jnk = glm(yv~XX,family=binomial)
			pred = predict(jnk)
			ppred0 = exp(-pred)/(1+exp(-pred))
			ppred1 = exp(pred)/(1+exp(pred))
			res = yv[!is.na(yv)]*(-pred-log(1-ppred1)/ppred1)+(1-yv[!is.na(yv)])*(-pred+log(1-ppred0)/ppred0)
			si2 = NULL
		}
		out = kmeans(res,centers=k,iter.max=100)
		bev = jnk$coef[-1]
		ind = order(out$centers)
		xiv = jnk$coef[1]+out$centers[ind]
		piv = out$size[ind]/sum(out$size)
		if(surv){
			su = survreg(Surv(tv,ifelse(censv,0,1))~.,data=data.frame(W[,-1]))
			nu = 1/su$scale
			phi = 0
			jnk = su$coef
			jnk[-1] = jnk[-1]*nu
			jnk[1] = exp(jnk[1])
			psiv = -jnk
		}else{
			nu = phi = psiv = NULL
		}
	}

# compute transition matrix
	if(is.null(btv)){
		a = min(c(0.1,tv))
		btv = seq(0,max(tv),a)
	}else{
		a = btv[2]-btv[1]
	}
	M = length(btv)
	if(LC){
		Pi = diag(k)
	}else{
		Pi = diag(k)+a*Q
		Tmp = a*Q
		it = 1
		while(any(abs(Tmp)>0) & it<1000){
			it = it+1
			Tmp = (a/it)*(Tmp%*%Q)
			Pi = Pi+Tmp
		}
	}

# expand matrix of responses and array of covariates
	YM = matrix(0,n,M)
	XM = array(0,c(n,M,npX))
	PM = matrix(0,n,M)
	for(i in 1:n) for(j in 1:jv[i]){
		m = sum(TT[i,j]>=btv)
		YM[i,m] = Y[i,j]
		XM[i,m,] = X[i,j,]
		PM[i,m] = 1
	}

# compute log-likelihood
	if(is.null(si2)) si2f = as.double(0) else si2f = si2
	if(surv){
		out = .Fortran("lk_rec_surv",k=as.integer(k),piv=piv,Pi=Pi,npX=as.integer(npX),bev=bev,
		                nu=nu,xiv=xiv,phi=phi,npW=as.integer(npW),psiv=psiv,si2=si2f,
		                typ=type,n=as.integer(n),XM=XM,W=W,YM=YM,M=as.integer(M),PM=as.integer(PM),
		                btv=btv,dev=as.integer(1-censv),tv=as.double(tv),lk=0,Pp1=array(999,c(k,M,n)),Pp2=matrix(0,k,k))
	}else{
		out = .Fortran("lk_rec",k=as.integer(k),piv=piv,Pi=Pi,npX=as.integer(npX),bev=bev,
		                xiv=xiv,si2=si2f,typ=type,n=as.integer(n),XM=XM,YM=YM,M=as.integer(M),PM=as.integer(PM),
		                btv=btv,tv=as.double(tv),lk=0,Pp1=array(999,c(k,M,n)),Pp2=matrix(0,k,k),ver=as.integer(rep(0,n)))
	}
	lk = out$lk
	it = 0; lko = lk
	if(verbose) print(c(0,lk,phi))
	while(((lk-lko)/abs(lko)>reltol | it==0) & it<maxit){
		it = it+1; lko = lk
# M-step
# update piv and Pi
		piv = rowSums(out$Pp1[,1,])/n
		if(!LC) Pi = (1/rowSums(out$Pp2))*out$Pp2
# update other parameters
		if(surv) par = c(bev,log(nu),xiv,phi,psiv) else par = c(bev,xiv,phi,psiv)
		if(it==1){
			if(surv){
	 			par1 = optim(par,lk_comp_joint_HM2fort,method="Nelder-Mead",piv=piv,Pi=Pi,si2=si2,type=type,
	      		             dev=1-censv,tv=tv,XM=XM,W=W,YM=YM,PM=PM,btv=btv,Pp1=out$Pp1,
		     	             control=list(fnscale=-1))
		   	}else{
	 			par1 = optim(par,lk_comp_joint_HM2fort,method="Nelder-Mead",piv=piv,Pi=Pi,si2=si2,type=type,
	      		             tv=tv,XM=XM,YM=YM,PM=PM,btv=btv,Pp1=out$Pp1,surv=FALSE,
	      		             control=list(fnscale=-1))
		   	}
			par = par1$par
		}else{
			if(surv){
				par2 = optim(par,lk_comp_joint_HM2fort,dlk_comp_joint_HM2fort,method="BFGS",piv=piv,Pi=Pi,si2=si2,
	     		             type=type,dev=1-censv,tv=tv,XM=XM,W=W,YM=YM,PM=PM,btv=btv,Pp1=out$Pp1,
		     	             control=list(fnscale=-1,reltol=10^-12))
			}else{
				par2 = optim(par,lk_comp_joint_HM2fort,dlk_comp_joint_HM2fort,method="BFGS",piv=piv,Pi=Pi,si2=si2,
	     		             type=type,tv=tv,XM=XM,YM=YM,PM=PM,btv=btv,Pp1=out$Pp1,surv=FALSE,
		     	             control=list(fnscale=-1,reltol=10^-12))
			}
			par = par2$par
		}
		bev = par[1:npX]
		if(surv){
			nu = exp(par[npX+1]); xiv = par[(npX+2):(k+npX+1)]
			phi = par[k+npX+2]; psiv = par[(k+npX+3):length(par)]
		}else{
			xiv = par[(npX+1):(k+npX)]
		}
		if(type=="norm"){
			num = 0
			for(i in 1:n) for(m in 1:M) if(PM[i,m]==1){
				muv = xiv+as.vector(XM[i,m,]%*%bev)
				num = num+(YM[i,m]-muv)^2%*%out$Pp1[,m,i]
			}
			si2 = as.vector(num/sum(jv))
			if(is.na(si2)){
				print(si2)
				browser()
			}
		}
# compute log-likelihood
		if(is.null(si2)) si2f = as.double(0) else si2f = si2
		if(surv){
			out = .Fortran("lk_rec_surv",k=as.integer(k),piv=piv,Pi=Pi,npX=as.integer(npX),bev=bev,
		                           nu=nu,xiv=xiv,phi=phi,npW=as.integer(npW),psiv=psiv,si2=si2f,
		                           typ=type,n=as.integer(n),XM=XM,W=W,YM=YM,M=as.integer(M),PM=as.integer(PM),
		                           btv=btv,dev=as.integer(1-censv),tv=as.double(tv),lk=0,Pp1=array(999,c(k,M,n)),Pp2=matrix(0,k,k))
		}else{
			out = .Fortran("lk_rec",k=as.integer(k),piv=piv,Pi=Pi,npX=as.integer(npX),bev=bev,
		                           xiv=xiv,si2=si2f,typ=type,n=as.integer(n),XM=XM,YM=YM,M=as.integer(M),PM=as.integer(PM),
		                           btv=btv,tv=as.double(tv),lk=0,Pp1=array(999,c(k,M,n)),Pp2=matrix(0,k,k))
		}
		lk = out$lk
		if(verbose & it%%100==0) print(c(it,lk,lk-lko,phi,proc.time()[3]-t0))
	}
	if(verbose) print(c(it,lk,lk-lko,phi,proc.time()[3]-t0))

# compute standard errors
	if(stderr){
		lav = log(piv[-1]/piv[1])
		par_tot = lav
		if(!LC) for(u in 1:k) par_tot = c(par_tot,log(Pi[u,-u]/Pi[u,u]))
		if(surv) par_tot = c(par_tot,bev,eta=log(nu),xiv,phi,psiv) else par_tot = c(par_tot,bev,xiv,phi,psiv)
		if(type=="norm") par_tot = c(par_tot,ga=log(si2))
		out = lk_obs(par_tot,YM,XM,W,tv,dev,k,type,btv,M,PM,LC,surv,npX,npW,n)
		dlkn = rep(0,length(par_tot))
		Ji = NULL
		for(j in 1:length(par_tot)){
			par_tot1 = par_tot; par_tot1[j] = par_tot1[j]+10^-6
			out1 = lk_obs(par_tot1,YM,XM,W,tv,dev,k,type,btv,M,PM,LC,surv,npX,npW,n)
			dlkn[j] = (out1$lk-out$lk)*10^6
			Ji = cbind(Ji,(out1$dlk-out$dlk)*10^6)
		}
		# print(c(lk,out$lk))
		# print(round(cbind(dlkn,out$dlk,dlkn-out$dlk),6))
		Ji = -(Ji+t(Ji))/2            # information matrix
		sev = sqrt(abs(diag(solve(Ji))))   # vector of standard errors
		tmp = sev
		selav = tmp[1:(k-1)]; tmp = tmp[-(1:(k-1))]
		if(!LC){
			seLa = NULL
			for(u in 1:k){
				seLa = rbind(seLa,tmp[1:(k-1)])
				tmp = tmp[-(1:(k-1))]
			}
		}
		sebev = tmp[1:npX]; tmp = tmp[-(1:npX)]
		if(surv){
			seeta = tmp[1]; tmp = tmp[-1]
		}
		sexiv = tmp[1:k]; tmp = tmp[-(1:k)]
		sephi = tmp[1]; tmp = tmp[-1]
		sepsiv = tmp[1:npW]; tmp = tmp[-(1:npW)]
		if(type=="norm"){
			sega = tmp
		}else{
			sega = NULL
		}
	}

# final output
	if(!LC) Q = logm(Pi %^% (1/a))
	np = length(par)
	if(!LC) np = np+k^2-1
	if(type=="norm") np = np+1
	bic = -2*lk+log(n)*np
	out = list(lk=lk,Q=Q,piv=piv,bev=bev,nu=nu,xiv=xiv,phi=phi,psiv=psiv,si2=si2,pP=out$Pp1,Pi=Pi,it=it,tim=proc.time()[3]-t0,bic=bic)
	if(stderr){
		out$Ji = Ji; out$sev = sev; out$selav = selav; out$seLa = seLa; out$sebev = sebev; out$seeta = seeta
		out$sexiv = sexiv; out$sephi = sephi; out$sepsiv = sepsiv; out$sega = sega
	}
	return(out)

}
# ---------------------------------------------------------------------------------------------------
lk_comp_joint_HM2fort <- function(par,piv,Pi,si2,type,XM,W,YM,PM,btv,Pp1,dev,tv,surv=TRUE){

# to compute the second and third components of the expect complete log-likelihood

# preliminaries
	k = length(piv)
	npX = dim(XM)[3]
	if(surv) npW = ncol(W)
	bev = par[1:npX]
	if(surv){
		lnu = par[npX+1]; xiv = par[(npX+2):(k+npX+1)]; phi = par[k+npX+2]
		nu = exp(lnu); psiv = par[(k+npX+3):length(par)]
	}else{
		xiv = par[(npX+1):(k+npX)]
	}
	n = nrow(YM)
	M = length(btv)
# call fortran function
	if(is.null(si2)) si2f = as.double(0) else si2f = si2
	if(surv){
		out = .Fortran("lk_comp_joint_surv",bev=bev,lnu=lnu,xiv=xiv,phi=phi,psiv=psiv,k=as.integer(k),piv=piv,Pi=Pi,si2=si2f,
		                   typ=type,n=as.integer(n),M=as.integer(M),npX=as.integer(npX),XM=XM,npW=as.integer(npW),
	    	               W=W,YM=YM,PM=as.integer(PM),btv=btv,Pp1=Pp1,dev=as.integer(dev),tv=as.double(tv),lk=0)
	}else{
		out = .Fortran("lk_comp_joint",bev=bev,xiv=xiv,k=as.integer(k),piv=piv,Pi=Pi,si2=si2f,
		                   typ=type,n=as.integer(n),M=as.integer(M),npX=as.integer(npX),XM=XM,
	    	               YM=YM,PM=as.integer(PM),btv=btv,Pp1=Pp1,tv=as.double(tv),lk=0)
	}
	lk = out$lk
# return output
	return(lk)

}
# ---------------------------------------------------------------------------------------------------
dlk_comp_joint_HM2fort <- function(par,piv,Pi,si2,type,XM,W,YM,PM,jv,btv,Pp1,dev,tv,surv=TRUE){

# to compute the derivative of lk_comp_joint_HM
# preliminaries
	k = length(piv)
	npX = dim(XM)[3]
	if(surv) npW = ncol(W)
	bev = par[1:npX]
	if(surv){
		lnu = par[npX+1]; xiv = par[(npX+2):(k+npX+1)]; phi = par[k+npX+2]
		nu = exp(lnu); psiv = par[(k+npX+3):length(par)]
	}else{
		xiv = par[(npX+1):(k+npX)]
	}
	n = nrow(YM)
	M = length(btv)
# call Fortran
# call fortran function
	if(is.null(si2)) si2f = as.double(0) else si2f = as.double(si2)
	if(surv){
		out = .Fortran("dlk_comp_joint_surv",bev=bev,lnu=lnu,xiv=xiv,phi=phi,psiv=psiv,k=as.integer(k),piv=piv,Pi=Pi,si2=si2f,
		                   typ=type,n=as.integer(n),M=as.integer(M),npX=as.integer(npX),XM=XM,npW=as.integer(npW),
	    	               W=W,YM=YM,PM=as.integer(PM),btv=btv,Pp1=Pp1,dev=as.integer(dev),tv=as.double(tv),
	    	               dlk=rep(0,npX+1+k+1+npW))
	}else{
		out = .Fortran("dlk_comp_joint",bev=bev,xiv=xiv,k=as.integer(k),piv=piv,Pi=Pi,si2=si2f,
		                   typ=type,n=as.integer(n),M=as.integer(M),npX=as.integer(npX),XM=XM,
	    	               YM=YM,PM=as.integer(PM),btv=btv,Pp1=Pp1,tv=as.double(tv),dlk=rep(0,npX+k))
	}
	dlk = out$dlk
# return output
	return(dlk)

}
# ---------------------------------------------------------------------------------------------------
lk_comp_joint_HM2 <- function(par,piv,Pi,si2,type,XM,W,YM,PM,btv,Pp1,dev,tv,surv=TRUE){

# to compute the second and third components of the expect complete log-likelihood

# preliminaries
	k = length(piv)
	npX = dim(XM)[3]
	bev = par[1:npX]
	if(surv){
		lnu = par[npX+1]; xiv = par[(npX+2):(k+npX+1)]; phi = par[k+npX+2]
		nu = exp(lnu); psiv = par[(k+npX+3):length(par)]
	}else{
		xiv = par[(npX+1):(k+npX)]
	}
	n = nrow(YM)
	M = length(btv)
	if(surv){
		btvnu = btv^nu
		dbtvnu = diff(btvnu)
	}
	if(type=="norm") lsi2pi = log(si2)+log(2*pi)
# compute functions
	l2 = 0
	if(surv){
		l3 = 0
		tmp00 = xiv*phi
	}
	if(surv) Wp = W%*%psiv
	for(i in 1:n){
		if(surv){
			tmp0 = tmp00+Wp[i]
			etmp0 = exp(tmp0)
		}
		mi = sum(btv<=tv[i])
		for(m in 1:mi){
			if(PM[i,m]==1){
				muv = xiv+as.vector(XM[i,m,]%*%bev)
				if(type=="norm") l2 = l2-1/2*(lsi2pi+(YM[i,m]-muv)^2/si2)%*%Pp1[,m,i]
				if(type=="bino"){
					muv = exp(muv)/(1+exp(muv))
					l2 = l2+(YM[i,m]*log(muv)+(1-YM[i,m])*log(1-muv))%*%Pp1[,m,i]
				}
			}
		}
		if(surv){
			if(mi>1) for(m in 2:mi) l3 = l3-dbtvnu[m-1]*(etmp0%*%Pp1[,m-1,i])
			l3 = l3+(dev[i]*(lnu+(nu-1)*log(tv[i])+tmp0)-etmp0*(tv[i]^nu-btvnu[mi]))%*%Pp1[,mi,i]
		}
	}
	if(surv) lk = as.vector(l2+l3) else lk = l2
	return(lk)

}
# ---------------------------------------------------------------------------------------------------
dlk_comp_joint_HM2 <- function(par,piv,Pi,si2,type,XM,W,YM,PM,jv,btv,Pp1,dev,tv,surv=TRUE){

# to compute the derivative of lk_comp_joint_HM
# preliminaries
	k = length(piv)
	npX = dim(XM)[3]
	if(surv) npW = ncol(W)
	bev = par[1:npX]
	if(surv){
		lnu = par[npX+1]; xiv = par[(npX+2):(k+npX+1)]; phi = par[k+npX+2]
		nu = exp(lnu); psiv = par[(k+npX+3):length(par)]
	}else{
		xiv = par[(npX+1):(k+npX)]
	}
	n = nrow(YM)
	M = length(btv)
	if(surv){
		btvnu = btv^nu
		dbtvnu = diff(btvnu)
	}
# compute functions
	dl2bev = rep(0,2)
	if(surv) dl3nu = 0
	dl2xiv = dl3xiv = rep(0,k)
	if(surv){
		dl3phi = 0
		dl3psiv = rep(0,npW)
	}
	for(i in 1:n){
		if(surv) tmp0 = exp(xiv*phi+as.vector(W[i,]%*%psiv))
		mi = sum(btv<=tv[i])
		for(m in which(PM[i,1:mi]==1)){
			muv = xiv+as.vector(XM[i,m,]%*%bev)
			if(type=="norm"){
				dl2bev = dl2bev+XM[i,m,]*as.vector(((YM[i,m]-muv)/si2)%*%Pp1[,m,i])
				dl2xiv = dl2xiv+(YM[i,m]-muv)/si2*Pp1[,m,i]
			}
			if(type=="bino"){
				muv = exp(muv)/(1+exp(muv))
				tmp1 = (YM[i,m]-muv)*Pp1[,m,i]
				dl2bev = dl2bev+XM[i,m,]*sum(tmp1)
				dl2xiv = dl2xiv+tmp1
			}
		}
		if(surv){
			if(mi>1) for(m in 2:mi){
				if(btv[m-1]>0) tmp = log(btv[m-1]) else tmp = 0
				dl3nu = dl3nu-(tmp0*(log(btv[m])*btvnu[m]-tmp*btvnu[m-1]))%*%Pp1[,m-1,i]
				dl3xiv = dl3xiv-(tmp0*dbtvnu[m-1])*phi*Pp1[,m-1,i]
				dl3phi = dl3phi-(tmp0*dbtvnu[m-1]*xiv)%*%Pp1[,m-1,i]
				dl3psiv = dl3psiv-W[i,]*as.vector((tmp0*dbtvnu[m-1])%*%Pp1[,m-1,i])
			}
			if(btv[mi]>0) tmp = log(btv[mi]) else tmp = 0
			dl3nu = dl3nu+(dev[i]*(1/nu+log(tv[i]))-tmp0*(log(tv[i])*tv[i]^nu-tmp*btvnu[mi]))%*%Pp1[,mi,i]
			dl3xiv = dl3xiv+(dev[i]*phi-tmp0*(tv[i]^nu-btvnu[mi])*phi)*Pp1[,mi,i]
			dl3phi = dl3phi+(dev[i]*xiv-tmp0*(tv[i]^nu-btvnu[mi])*xiv)%*%Pp1[,mi,i]
			dl3psiv = dl3psiv+W[i,]*as.vector((dev[i]-tmp0*(tv[i]^nu-btvnu[mi]))%*%Pp1[,mi,i])
		}
	}
# output
	if(surv) dlk = c(dl2bev,dl3nu*nu,dl2xiv+dl3xiv,dl3phi,dl3psiv) else dlk = c(dl2bev,dl2xiv)
	return(dlk)

}
# ---------------------------------------------------------------------------------------------------
lk_rec_joint_HM2 <- function(piv,Pi,bev,nu,xiv,phi,psiv,si2,type,XM,W,YM,btv,dev,tv,surv=TRUE){

# preliminaries
	n = nrow(YM)
	k = length(xiv)
	M = length(btv)
	if(type=="norm") si = sqrt(si2)
	if(surv){
		btvnu = btv^nu
		dbtvnu = diff(btvnu)
	}
# recursion for each individual
	lk = 0
	Pp1 = array(NA,c(k,M,n))
	Pp2 = matrix(0,k,k)
	for(i in 1:n){
		mi = sum(btv<=tv[i])
		if(surv) tmp0 = exp(xiv*phi+as.vector(W[i,]%*%psiv))
# FORWARD RECURSION
		Frec = matrix(0,k,mi)
# first time
		muv = xiv+as.vector(XM[i,1,]%*%bev)
		if(type=="norm") frec = piv*dnorm(YM[i,1],muv,si)
		if(type=="bino"){
			muv = exp(muv)/(1+exp(muv))
			frec = piv*(YM[i,1]*muv+(1-YM[i,1])*(1-muv))
		}
		Frec[,1] = frec
# following time occasions
		if(mi>1) for(m in 2:mi){
			if(surv){
				tmp = exp(-tmp0*dbtvnu[m-1])
				frec = as.vector(t(Pi)%*%(frec*tmp))
			}else{
				frec = as.vector(t(Pi)%*%frec)
			}
			if(!is.na(YM[i,m])){
				muv = xiv+as.vector(XM[i,m,]%*%bev)
				if(type=="norm") frec = frec*dnorm(YM[i,m],muv,si)
				if(type=="bino"){
					muv = exp(muv)/(1+exp(muv))
					frec = frec*(YM[i,m]*muv+(1-YM[i,m])*(1-muv))
				}
			}
			Frec[,m] = frec
		}
		if(surv){
			tmp = exp(-tmp0*(tv[i]^nu-btvnu[mi]))
			if(dev[i]==1) tmp = tmp*nu*tv[i]^(nu-1)*tmp0
			frec = as.vector(frec*tmp)
		}
		fi = sum(frec)
# accumulate log-likelihood
		lk = lk+log(fi)

# BACKWARD RECURSION
		Grec = matrix(0,k,mi)
# last time occasion
		if(surv){
			grec = exp(-tmp0*(tv[i]^nu-btvnu[mi]))
			if(dev[i]==1) grec = grec*nu*tv[i]^(nu-1)*tmp0
		}else{
			grec = rep(1,k)
		}
		Grec[,mi] = grec
# previous time occasions
		if(mi>1) for(m in (mi-1):1){
			if(surv) tmp = exp(-tmp0*dbtvnu[m])
			if(!is.na(YM[i,m+1])){
				muv = xiv+as.vector(XM[i,m+1,]%*%bev)
				if(type=="norm") grec = grec*dnorm(YM[i,m+1],muv,si)
				if(type=="bino"){
					muv = exp(muv)/(1+exp(muv))
					grec = grec*(YM[i,m+1]*muv+(1-YM[i,m+1])*(1-muv))
				}
			}
			if(surv) grec = tmp*as.vector(Pi%*%grec) else grec = as.vector(Pi%*%grec)
			Grec[,m] = grec
		}
# Posterior probabilities for the first time occasion
		Pp1[,1:mi,i] = Frec*Grec/fi
# Joint Posterior probabilities for two occasions
		if(mi>1) for(m in 1:(mi-1)){
			if(surv) tmp = exp(-tmp0*dbtvnu[m])
			if(!is.na(YM[i,m+1])){
				muv = xiv+as.vector(XM[i,m+1,]%*%bev)
				if(type=="norm") tmp1 = dnorm(YM[i,m+1],muv,si)
				if(type=="bino"){
					muv = exp(muv)/(1+exp(muv))
					tmp1 = YM[i,m+1]*muv+(1-YM[i,m+1])*(1-muv)
				}
			}else{
				tmp1 = 1
			}
			if(surv){
				Tmp = ((Frec[,m]*tmp)%o%(Grec[,m+1]*tmp1))*Pi/fi
			}else{
				Tmp = (Frec[,m]%o%(Grec[,m+1]*tmp1))*Pi/fi
			}
			Pp2 = Pp2+Tmp
		}
	}
# output
	out = list(lk=lk,Pp1=Pp1,Pp2=Pp2)

}
# ---------------------------------------------------------------------------------------------------
lk_obs <- function(par_tot,YM,XM,W,tv,dev,k,type,btv,M,PM,LC,surv,npX,npW,n)
{

# reconstruct parameters
    G = diag(k)[,-1]
    if(k==2) {G=matrix(G)}
	piv = exp(G%*%par_tot[1:(k-1)]); piv = as.vector(piv/sum(piv))
	par_tot = par_tot[-(1:(k-1))]
	if(!LC){
		Pi = NULL
		for(u in 1:k){
                    G = diag(k)[,-u]
                    if(k==2) {G=matrix(G)}
			tmp = exp(G%*%par_tot[1:(k-1)]); tmp = as.vector(tmp/sum(tmp))
			Pi = rbind(Pi,tmp)
			par_tot = par_tot[-(1:(k-1))]
		}
	}
	bev = par_tot[1:npX]; par_tot = par_tot[-(1:npX)]
	if(surv){
		eta = par_tot[1]; nu = exp(eta); par_tot = par_tot[-1]
	}
	xiv = par_tot[1:k]; par_tot = par_tot[-(1:k)]
	phi = par_tot[1]; par_tot = par_tot[-1]
	psiv = par_tot[1:npW]; par_tot = par_tot[-(1:npW)]
	if(type=="norm"){
		ga = par_tot; si2 = exp(ga)
	}else{
		si2 = NULL
	}

# compute log-likelihood
	if(is.null(si2)) si2f = as.double(0) else si2f = si2
	if(surv){
		out = .Fortran("lk_rec_surv",k=as.integer(k),piv=piv,Pi=Pi,npX=as.integer(npX),bev=bev,
		                nu=nu,xiv=xiv,phi=phi,npW=as.integer(npW),psiv=psiv,si2=si2f,
		                typ=type,n=as.integer(n),XM=XM,W=W,YM=YM,M=as.integer(M),PM=as.integer(PM),
		                btv=btv,dev=as.integer(dev),tv=as.double(tv),lk=0,Pp1=array(999,c(k,M,n)),Pp2=matrix(0,k,k))
	}else{
		out = .Fortran("lk_rec",k=as.integer(k),piv=piv,Pi=Pi,npX=as.integer(npX),bev=bev,
		                xiv=xiv,si2=si2f,typ=type,n=as.integer(n),XM=XM,YM=YM,M=as.integer(M),PM=as.integer(PM),
		                btv=btv,tv=as.double(tv),lk=0,Pp1=array(999,c(k,M,n)),Pp2=matrix(0,k,k),ver=as.integer(rep(0,n)))
	}
	lk = out$lk
# compute posterior probabilities
	Pp1 = out$Pp1; Pp2 = out$Pp2
	Pp1[Pp1==999] = NA
# preliminaries
	if(surv){
		btvnu = btv^nu
		dbtvnu = diff(btvnu)
	}
# compute functions
    G = diag(k)[,-1]
    if(k==2) {G=matrix(G)}
	bnv = rowSums(Pp1[,1,])
	dl1lav = t(G)%*%(bnv-n*piv)
	dl1La = NULL
	for(u in 1:k){
            G = diag(k)[,-u]
            if(k==2) {G=matrix(G)}
		dl1La = c(dl1La,t(G)%*%(Pp2[u,]-sum(Pp2[u,])*Pi[u,]))
	}
	dl2bev = rep(0,2)
	if(surv) dl3nu = 0
	if(type=="norm") dl2ga = 0
	dl2xiv = dl3xiv = rep(0,k)
	if(surv){
		dl3phi = 0
		dl3psiv = rep(0,npW)
	}
	for(i in 1:n){
		if(surv) tmp0 = exp(xiv*phi+as.vector(W[i,]%*%psiv))
		mi = sum(btv<=tv[i])
		for(m in which(PM[i,1:mi]==1)){
			muv = xiv+as.vector(XM[i,m,]%*%bev)
			if(type=="norm"){
				dl2bev = dl2bev+XM[i,m,]*as.vector(((YM[i,m]-muv)/si2)%*%Pp1[,m,i])
				dl2xiv = dl2xiv+(YM[i,m]-muv)/si2*Pp1[,m,i]
				dl2ga = dl2ga-(1-(YM[i,m]-muv)^2/si2)%*%Pp1[,m,i]/2
			}
			if(type=="bino"){
				muv = exp(muv)/(1+exp(muv))
				tmp1 = (YM[i,m]-muv)*Pp1[,m,i]
				dl2bev = dl2bev+XM[i,m,]*sum(tmp1)
				dl2xiv = dl2xiv+tmp1
			}
		}
		if(surv){
			if(mi>1) for(m in 2:mi){
				if(btv[m-1]>0) tmp = log(btv[m-1]) else tmp = 0
				dl3nu = dl3nu-(tmp0*(log(btv[m])*btvnu[m]-tmp*btvnu[m-1]))%*%Pp1[,m-1,i]
				dl3xiv = dl3xiv-(tmp0*dbtvnu[m-1])*phi*Pp1[,m-1,i]
				dl3phi = dl3phi-(tmp0*dbtvnu[m-1]*xiv)%*%Pp1[,m-1,i]
				dl3psiv = dl3psiv-W[i,]*as.vector((tmp0*dbtvnu[m-1])%*%Pp1[,m-1,i])
			}
			if(btv[mi]>0) tmp = log(btv[mi]) else tmp = 0
			dl3nu = dl3nu+(dev[i]*(1/nu+log(tv[i]))-tmp0*(log(tv[i])*tv[i]^nu-tmp*btvnu[mi]))%*%Pp1[,mi,i]
			dl3xiv = dl3xiv+(dev[i]*phi-tmp0*(tv[i]^nu-btvnu[mi])*phi)*Pp1[,mi,i]
			dl3phi = dl3phi+(dev[i]*xiv-tmp0*(tv[i]^nu-btvnu[mi])*xiv)%*%Pp1[,mi,i]
			dl3psiv = dl3psiv+W[i,]*as.vector((dev[i]-tmp0*(tv[i]^nu-btvnu[mi]))%*%Pp1[,mi,i])
		}
	}
	if(surv) dlk = c(dl1lav,dl1La,dl2bev,dl3nu*nu,dl2xiv+dl3xiv,dl3phi,dl3psiv) else dlk = c(dl2bev,dl2xiv)
	if(type=="norm") dlk = c(dlk,dl2ga)
# output
	out = list(lk=lk,dlk=dlk)

}


#' @title calcKristCr
#'
#' @description Calculate the residual resistance coefficient (\code{Cr})
#' (dimensionless) using the Kristensen method with Harvald regressions.
#'
#'@param shipType Ship type (vector of strings, see \code{\link{calcShipType}})
#'@param M Fineness/slenderness coefficient (vector of numericals,
#'dimensionless) (see \code{\link{calcShipM}})
#'@param froudeNum Froude number (vector of numericals, dimensionless) (see
#'\code{\link{calcFroudeNum}})
#'@param actualDraft Actual draft (vector of numericals, m)
#'@param breadth Moulded breadth (vector of numericals, m)
#'@param Cp Prismatic coefficient (vector of numericals, dimensionless) (see
#'\code{\link{calcCp}})
#'@param tankerBulkCarrierGCargoShipTypes Ship types specified in input
#'\code{shipTypes} to be modeled as tankers, bulk carriers and general cargo
#'vessels (vector of strings)
#'
#' @details
#'Note: Uses Guldhammer Residual Resistance Calculation and adjusts for bulbous bow, and
#'breadth/draft ratio.
#'
#' This method this requires ship types to be grouped. Use the
#' \code{tankerBulkCarrierGCargoShipTypes} grouping
#' parameters to provide these ship type groupings. Any ship types not include
#' this grouping will be considered as miscellaneous vessels.
#'
#' Actual draft is typically obtained from sources such as AIS messages or ship
#' records.
#'
#' @return \code{Cr} (vector of numericals, dimensionless)
#'
#' @references
#'Kristensen, H. O. and Lutzen, M. 2013. "Prediction of Resistance and Propulsion
#'Power of Ships."
#'
#'\href{https://gitlab.gbar.dtu.dk/oceanwave3d/Ship-Desmo}{Kristensen, H. O.
#'"Ship-Desmo-Tool." https://gitlab.gbar.dtu.dk/oceanwave3d/Ship-Desmo}
#'
#'@seealso \itemize{
#'\item \code{\link{calcShipM}}
#'\item \code{\link{calcFroudeNum}}
#'\item \code{\link{calcCp}}
#'\item \code{\link{calcShipType}}
#'}
#'
#' @family Kristensen Calculations
#' @family Resistance Calculations
#'
#' @examples
#' calcKristCr("bulk.carrier", 5.2, 0.1199268, 12.48, 32.25, 0.81)
#'
#' @export

calcKristCr <- function(shipType, M, froudeNum, actualDraft, breadth, Cp,
                        tankerBulkCarrierGCargoShipTypes=c("tanker","general.cargo","chemical.tanker","liquified.gas.tanker","oil.tanker","other.tanker","bulk.carrier")
){

E<- (
  (1.35-(0.23*M)+0.012*(M^2))+
    (1.5*(froudeNum^1.8))+
    ( 0.0011*(M^9.1)*(froudeNum^(2*M-3.7)))
)*
  (0.98+(2.5/((M-2)^4)))+
  ((M-5)^4)*((froudeNum-0.1)^4)

G<- ((7-0.09*((M)^2))*(((5*Cp)-2.5)^2))/
  ((600*((froudeNum-0.315)^2)+1)^1.5)

H<-exp(
  80*(froudeNum-(0.04+(0.59*Cp))-0.015*(M-5))
)

  #Calculate Cr
  #Source: Kristensen(2016) using Guldhammer(1978) method
  Cr<-((
        (1.35-(0.23*M)+0.012*(M^2))+
        (1.5*(froudeNum^1.8))+
        ( 0.0011*(M^9.1)*(froudeNum^(2*M-3.7)))
        )*
        (0.98+(2.5/((M-2)^4)))+
        ((M-5)^4)*((froudeNum-0.1)^4)+
        (
          ((7-0.09*((M)^2))*(((5*Cp)-2.5)^2))/
          ((600*((froudeNum-0.315)^2)+1)^1.5)
        )+
        exp(
            80*(froudeNum-(0.04+(0.59*Cp))-0.015*(M-5))
            )+
        ( 180*(froudeNum^3.7)*exp((20*Cp)-16))+
        #correction for vessels where B/T!=2.5
        (.16*((breadth/actualDraft)-2.5))
      )/1000

  #Apply Cr Correction for bulbous bow shape according to ship type.
  #Bulbous bows have become more common since these regressions were
  #made, they cut down on a lot of resistance
  Cr<- ifelse( shipType %in%tankerBulkCarrierGCargoShipTypes,

               Cr+ pmax(-0.4, -0.1-1.6*froudeNum)/1000,

               Cr+ pmax(-50,250*froudeNum-90)* (Cr/100)

                )#end ifelse

  return(Cr)

}


# lhl.r
# Church two step breakpoint algorithm, from Church et al, 1994.
# Author: Emmanuel Alcalá
# Date: 5 may 2019
#
# ------------------------------------------------------------------------------
# This is an improved version of lhl_old. It is just marginally faster, but
# cleaner and more intuitive. It sets first all possible combinations of s1
# and s2 (start and stop), then it computes the differences between the overall
# rate r and the state rates r1, r2 and r3, weighted by their respective
# durations t1, t2 and t3. Then it computes the index (i.e., an integer of the
# position in the vector) at which the differences were maximized, with which
# all other metrics can be solved.
# inputs:
#        - r_times: every time at which there was a response
#        - trial_duration: the duration of peal trial
# ------------------------------------------------------------------------------

lhl = function(r_times, trial_duration) {
  # Number of responses
  nr = length(r_times)
  # overall response rate
  r = nr/trial_duration
  # Vector of putative starts (all time of responses except the last)
  s1 = r_times[1:(nr - 1)]
  # Vector of putative stops (except the first)
  s2 = r_times[2:nr]

  # Get all combinations of s1 and s2 with s1 less than s2.
  #   For example, if s1 can be 1, 2, 3, and s2 = 2, 3, 4,
  #   the combinations are (1, 2), (1, 3), (1, 4), (2, 3), (2, 4), (3, 4)
  grid = expand.grid.jc(s1,s2) # matrix of nx2; n = all possible combinations
  grid = grid[grid[ ,1] < grid[ ,2], ] # select n with s1 < s2
  # Vector of 0 to store results
  A = numeric(nrow(grid))
  # Computations of differences between the r and r(i), with i = 1, 2, 3
  # which are the rates at the low, high and low states.
  # Store the results of the differences in A[j]
  for (j in 1:nrow(grid)) { # loop trough all combinations of s1 & s2
    t1 = grid[j, 1]
    t2 = grid[j, 2] - t1
    t3 = trial_duration - grid[j, 2]
    # Previous versions looked if some t equals 0 and assigned 0 to rates
    # to avoid division by 0, but isn' necessary: there are 0 responses
    # with duration 0!
    r1 = sum(r_times <= t1)/t1
    r2 = sum(r_times > t1 & r_times <= grid[j, 2])/t2
    r3 = sum(r_times > grid[j, 2])/t3

    A[j] = sum(t1*(r - r1), t2*(r2 - r), t3*(r - r3), na.rm = T)

  }
  # Get the index which maximizes A
  argmaxA = which.max(A)
  # Vector of 2 at which A were maximum. column 1 is s1, 2 is s2
  s1s2 = grid[argmaxA, ]
  # High rate state duration: s2 - s1
  t2 = s1s2[2] - s1s2[1]
  # Low rate 1
  r1 = sum(r_times <= s1s2[1])/s1s2[1]
  # Low rate 2
  r3 = sum(r_times > s1s2[2])/(trial_duration - s1s2[2])
  # High rate
  hrate = sum(r_times > s1s2[1] & r_times <= s1s2[2])/t2
  # Middle time
  mid = sum(s1s2)/2
  # Return results
  data.frame(start = s1s2[1],
             stop = s1s2[2],
             spread = t2,
             hrate = hrate,
             mid = mid,
             r1 = r1,
             r3 = r3)
}


#!/usr/bin/Rscript

# R (http://www.r-project.org/) script
# from https://en.wikipedia.org/wiki/R_(programming_language)#Basic_syntax

install.packages("caTools")  # install external package
library(caTools)             # external package providing write.gif function
jet.colors <- colorRampPalette(c("red", "blue", "#007FFF", "cyan", "#7FFF7F",
                                 "yellow", "#FF7F00", "red", "#7F0000"))
dx <- 1500                    # define width
dy <- 1400                    # define height
C  <- complex(real = rep(seq(-2.2, 1.0, length.out = dx), each = dy),
              imag = rep(seq(-1.2, 1.2, length.out = dy), dx))
C <- matrix(C, dy, dx)       # reshape as square matrix of complex numbers
Z <- 0                       # initialize Z to zero
X <- array(0, c(dy, dx, 20)) # initialize output 3D array
for (k in 1:20) {            # loop with 20 iterations
  Z <- Z^2 + C               # the central difference equation
  X[, , k] <- exp(-abs(Z))   # capture results
}
write.gif(X, "out/mandelbrot.gif", col = jet.colors, delay = 100)


rdirichlet = function(n, a) {
  l = length(a)
  x = matrix(rgamma(n * l, a), ncol = l, byrow = T)
  s = x %*% rep(1, l)
  x / as.vector(s)
}