########################### BayesKin ###########################
## Bayesian estimates of rigid body kinematics
## Andy Pohl
## UofC - Faculty of Kinesiology

## Written by: Andy Pohl
## UofC - Faculty of Kinesiology
## June-Dec 2020
## Revision 1: May 2020
################################################################

##################### SupplementE.r  ###################
## Performs additional analysis examining the robustness of Bayesian models to
## increases in measurement noise
################################################################

## Preliminaries
rm(list = ls())

WORKING_DIR = "/media/data/ActiveProjects/P01_BayesianKinematics/PublishedCode/BayesKin/src/SupplementE" # Replace with appropriate location of working directory on local PC
SAVE_MDL = TRUE # assume save
INIT_TYPE = 'true_vals' # either 'true_vals', 'ls_est', 'random
NITS = 1000 # number of iterations.

setwd(WORKING_DIR)

# External libraries
library('dplyr')
library('latex2exp')

source('./library.r') # Retrieve function library

set.seed(1)  # Set seed for reproducibility

########################### Section A  #########################
## Specify simulation parameters
################################################################

## 1) Specify parameters
n.links = 1                 # Number of links.
seg.length = c(0.45)        # Length of segment specified in m. NB: Pataky pg 2.
r_true = c(0.07, 0.03)      # Origin location in m.
theta_true = c(-55)         # Rotation angle in deg. NB: specified to avoid 0/360 issue
# and mirroring solutions.
sigma_true = c(0.1, 1.5, 5, 10, 20,40)*(1/1000) # Measurement Noise in m
links = list(gen_link(seg.length = seg.length[1],
                      plate.center = 0.7*seg.length[1])) # Link geometry

########################### Section B  #########################
## Run simulations
    # run NITS iterations per sigma level
################################################################
results = matrix(NA, nrow = length(sigma_true)*NITS*5, ncol = 7)
row_idx =1
for(i in 1:length(sigma_true)){
    true_vals = list(r_true = r_true, theta_true = theta_true, sigma_true = sigma_true[i])

    # Generate posture
    posture = gen_posture(links, r = true_vals$r_true, theta = true_vals$theta_true)
    seeds = 1:NITS # Generate seed values for each iteration

    for(j in 1:NITS){
        seed = seeds[j]
        set.seed(seed)
        print(sprintf("--------------------Iteration %.0f of %.0f  -  seed = %.0f---------------------\n",j, NITS, seed))
        # gen obs
        y = gen_obs(posture, true_vals$sigma_true)
        #plot_system(posture = posture, y=y)
        print("Compute LS solution.")
        LS_result = LS_soln(y, links,
                            inits = c(true_vals$r_true, true_vals$theta_true, true_vals$sigma_true),
                            init_type = 'true_vals')
        results[row_idx,] = c(true_vals$sigma_true, 1, seed, LS_result$r.hat[1], LS_result$r.hat[2], LS_result$theta.hat[1],LS_result$sigma.hat)
        row_idx = row_idx+1

        Bayes1_result = Bayes_soln(y=y,
                                   links = links,
                                   mdlfile = './BayesModel_vague1.jags', # Vague 1
                                   init_type = INIT_TYPE,
                                   true_vals = c(true_vals$r_true,true_vals$theta_true, true_vals$sigma_true),
                                   ls_est = c(LS_result$r.hat, LS_result$theta.hat, LS_result$sigma.hat))
        Bayes1_result = summary(Bayes1_result$fit)
        results[row_idx,] = c(true_vals$sigma_true, 2, seed, Bayes1_result$statistics['r[1]',1], Bayes1_result$statistics['r[2]',1], Bayes1_result$statistics['theta',1],  Bayes1_result$statistics['sigma',1])
        row_idx = row_idx +1

        Bayes2_result = Bayes_soln(y=y,
                                   links = links,
                                   mdlfile = './BayesModel_vague2.jags', # Vague 2
                                   init_type = INIT_TYPE,
                                   true_vals = c(true_vals$r_true,true_vals$theta_true, true_vals$sigma_true),
                                   ls_est = c(LS_result$r.hat, LS_result$theta.hat, LS_result$sigma.hat))
        Bayes2_result = summary(Bayes2_result$fit)
        results[row_idx,] = c(true_vals$sigma_true, 3, seed, Bayes2_result$statistics['r[1]',1], Bayes2_result$statistics['r[2]',1], Bayes2_result$statistics['theta',1],  Bayes2_result$statistics['sigma',1])
        row_idx = row_idx +1

        Bayes4_result = Bayes_soln(y=y,
                                   links = links,
                                   mdlfile = './BayesModel_p3.jags', # Weakly Informative
                                   init_type = INIT_TYPE,
                                   true_vals = c(true_vals$r_true,true_vals$theta_true, true_vals$sigma_true),
                                   ls_est = c(LS_result$r.hat, LS_result$theta.hat, LS_result$sigma.hat))
        Bayes4_result = summary(Bayes4_result$fit)
        results[row_idx,] = c(true_vals$sigma_true, 4, seed, Bayes4_result$statistics['r[1]',1], Bayes4_result$statistics['r[2]',1], Bayes4_result$statistics['theta[1]',1],  Bayes4_result$statistics['sigma',1])
        row_idx = row_idx +1

        Bayes5_result = Bayes_soln(y=y,
                                   links = links,
                                   mdlfile = './BayesModel_p2.jags', # Informative
                                   init_type = INIT_TYPE,
                                   true_vals = c(true_vals$r_true,true_vals$theta_true, true_vals$sigma_true),
                                   ls_est = c(LS_result$r.hat, LS_result$theta.hat, LS_result$sigma.hat))

        Bayes5_result = summary(Bayes5_result$fit)
        results[row_idx,] = c(true_vals$sigma_true, 5, seed, Bayes5_result$statistics['r[1]',1], Bayes5_result$statistics['r[2]',1], Bayes5_result$statistics['theta',1],Bayes5_result$statistics['sigma',1])
        row_idx = row_idx +1

    }
}

if(SAVE_MDL){
    saveRDS(results, "./DeltaVarResults.rds")
}
results = readRDS('./DeltaVarResults.rds')

########################### Section C  #########################
## Process results and generate figures
################################################################
models = c('LS', 'Bayes Vague 1', "Bayes Vague 2",
           'Bayes Weakly Informative', 'Bayes Informative')

abs_error = matrix(NA, nrow = length(sigma_true)*length(models)*4, ncol = 6)
abs_error = data.frame(abs_error)
names(abs_error) = c('sigma', 'Model', 'Parameter', 'lwr', 'med', 'upr')

perf = matrix(NA, nrow = 3*length(sigma_true)*length(models), ncol = 4)
perf = data.frame(perf)
names(perf) = c('sigma', 'Model', 'Parameter', 'Perf')
row_idx = 1
row_idx2 = 1
for(i in 1:length(sigma_true)){
    sigma_result = results[results[,1]==sigma_true[i],]
    for(j in 1:length(models)){ # Summarize abs error of each model
        mdl_result = sigma_result[sigma_result[,2] == j,]
        abs_error[row_idx,1] = unique(mdl_result[,1])
        abs_error[row_idx,2] = unique(mdl_result[,2])
        abs_error[row_idx,3] = 'rx'
        abs_error[row_idx, 4:6] = quantile(abs(mdl_result[,4] - r_true[1]), probs = c(0.025, 0.5, 0.975))*1000
        row_idx = row_idx+1

        abs_error[row_idx,1] = unique(mdl_result[,1])
        abs_error[row_idx,2] = unique(mdl_result[,2])
        abs_error[row_idx,3] = 'ry'
        abs_error[row_idx, 4:6] = quantile(abs(mdl_result[,5] - r_true[2]), probs = c(0.025, 0.5, 0.975))*1000
        row_idx = row_idx+1

        abs_error[row_idx,1] = unique(mdl_result[,1])
        abs_error[row_idx,2] = unique(mdl_result[,2])
        abs_error[row_idx,3] = 'theta1'
        abs_error[row_idx, 4:6] = quantile(abs(mdl_result[,6] - theta_true[1]), probs = c(0.025, 0.5, 0.975))
        row_idx = row_idx+1

        abs_error[row_idx,1] = unique(mdl_result[,1])
        abs_error[row_idx,2] = unique(mdl_result[,2])
        abs_error[row_idx,3] = 'sigma'
        abs_error[row_idx, 4:6] = quantile(abs(mdl_result[,7] - sigma_true[i]), probs = c(0.025, 0.5, 0.975))*1000
        row_idx = row_idx +1

        if(j >=2){ # Compute performance of Bayes vs LS.
            perf[row_idx2,1] = unique(mdl_result[,1])
            perf[row_idx2,2] = unique(mdl_result[,2])
            perf[row_idx2,3] = 'rx'
            perf[row_idx2,4] = mean(abs(mdl_result[,4] - r_true[1]) < abs(sigma_result[sigma_result[,2] == 1,4] - r_true[1]))
            row_idx2 = row_idx2 +1

            perf[row_idx2,1] = unique(mdl_result[,1])
            perf[row_idx2,2] = unique(mdl_result[,2])
            perf[row_idx2,3] = 'ry'
            perf[row_idx2,4] = mean(abs(mdl_result[,5] - r_true[2]) < abs(sigma_result[sigma_result[,2] == 1,5] - r_true[2]))
            row_idx2 = row_idx2 +1

            perf[row_idx2,1] = unique(mdl_result[,1])
            perf[row_idx2,2] = unique(mdl_result[,2])
            perf[row_idx2,3] = 'theta1'
            perf[row_idx2,4] = mean(abs(mdl_result[,6] - theta_true[1]) < abs(sigma_result[sigma_result[,2] == 1,6] - theta_true[1]))
            row_idx2 = row_idx2 +1

            perf[row_idx2,1] = unique(mdl_result[,1])
            perf[row_idx2,2] = unique(mdl_result[,2])
            perf[row_idx2,3] = 'sigma'
            perf[row_idx2,4] = mean(abs(mdl_result[,7] - sigma_true[i]) < abs(sigma_result[sigma_result[,2] == 1,7] - sigma_true[i]))
            row_idx2 = row_idx2 +1
        }
    }
}
abs_error$sigma = abs_error$sigma*1000
perf$sigma = perf$sigma*1000

########################### Section D  #########################
## Generate figures and tables
################################################################

# Color scheme
colors = list(grey = rgb(97/255, 97/255,97/255, 0.2),
              orange = rgb(241/255, 136/255,5/255, 1),
              blue = rgb(18/255, 78/255,120/255,  0.2),
              light_blue = rgb(93/255, 173/255,226/255,  0.2),
              red = rgb(162/255, 0/255,33/255, 0.2),
              green = rgb(76/255, 159/255, 112/255, 0.2),
              purple = rgb(217/255, 187/255, 249/255, 0.2),
              pink = rgb(217/255, 3/255, 104/255, 0.2))

colors2 = list(grey = rgb(97/255, 97/255,97/255, 0.2),
               orange = rgb(241/255, 136/255,5/255, 1),
               blue = rgb(18/255, 78/255,120/255, 1),
               light_blue = rgb(93/255, 173/255,226/255,  1),
               red = rgb(162/255, 0/255,33/255, 1),
               green = rgb(76/255, 159/255, 112/255, 1),
               purple = rgb(217/255, 187/255, 249/255, 1),
               pink = rgb(217/255, 3/255, 104/255, 1))
point_cols = list(colors$pink, colors$blue,colors$light_blue, colors$purple,colors$green)
point_cols2 = list(colors2$pink, colors2$blue,colors2$light_blue, colors2$purple,colors2$green)

## -----------------------------------------------------------------------------
# Figure 1 'Line' plot of absolute error accross strain.
## -----------------------------------------------------------------------------

ylabs = c(TeX("Absolute Error $\\hat{r}_x \\, (mm)$"), TeX("Absolute Error $\\hat{r}_y \\, (mm)$"), TeX("Absolute Error $\\hat{\\theta}_1 \\, (deg)$"), TeX("Absolute Error $\\hat{\\sigma} \\, (mm)$"))
titles = c(TeX("$r_x$"), TeX("$r_y$"), TeX('$\\theta_1'), TeX("$\\sigma$"))
titles = c("(a)","(b)",'(c)',"(d)")

parameters = c('rx', 'ry', 'theta1', 'sigma')
ylims = list(c(0, 250), c(0,250), c(0, 60), c(0, 40))
# j = 1:4 parameters
# i = 1:5 models

par(mfrow=c(1,4))
for(j in 1:4){
    for(i in 1:5){
        parm_abs_error = abs_error %>% filter(Parameter ==parameters[j])

        mdl_abs_error = parm_abs_error %>% filter(Model ==i)
        jitter = (2-(5-i)) * 0.05*mdl_abs_error[,1]

        plot_col2 = point_cols2[[i]]
        plot_col1 = point_cols[[i]]
        if(i ==1){
            plot(mdl_abs_error$sigma + jitter, mdl_abs_error$med,type = 'b', pch = 16, col =plot_col2,
                 xlab = TeX("Noise Level: $\\sigma \\, (mm)$"), ylab = ylabs[j], bty='n',
                 xlim = c(-0.02, 50), ylim = ylims[[j]],
                 xaxt ='n')
            segments(x0=mdl_abs_error$sigma + jitter, y0 =  mdl_abs_error$lwr, x1 = mdl_abs_error$sigma + jitter, y1 = mdl_abs_error$upr, lty = 3, col = plot_col2)
            axis(side = 1, at = (sigma_true*1000), as.character(sigma_true*1000), las=2)
        }else{
            lines(mdl_abs_error$sigma + jitter, mdl_abs_error$med,type = 'b', pch = 16, col =plot_col2)
            segments(x0=mdl_abs_error$sigma + jitter, y0 =  mdl_abs_error$lwr, x1 = mdl_abs_error$sigma + jitter, y1 = mdl_abs_error$upr, lty = 3, col = plot_col2)

        }
        title(titles[j], font.main=1)

    }
}
legend(0, 40, legend = models[1:5],
       col = c(point_cols2[[1]],point_cols2[[2]], point_cols2[[3]], point_cols2[[4]], point_cols2[[5]]),
       lty = 1, bty='n')

## -----------------------------------------------------------------------------
# Figure 2 'Line' plot of proportion of sims where Bayes had < error than LS
## -----------------------------------------------------------------------------

titles = c(TeX("(a) \t $r_x$"),TeX("(b) \t $r_y$"),TeX('(c) \t $\\theta_1$'),TeX("(d) \t $\\sigma$"))
par(mfrow=c(1,4))
for(j in 1:4){
    for(i in 2:5){
        parm_perf = perf %>% filter(Parameter ==parameters[j])
        mdl_parm_perf = parm_perf %>% filter(Model ==i)
        jitter = 0

        plot_col2 = point_cols2[[i]]
        if(i ==2){
            plot(mdl_parm_perf$sigma + jitter, mdl_parm_perf$Perf,type = 'b', pch = 16, col =plot_col2,
                 xlab = TeX("Noise Level: $\\sigma \\, (mm)$"), ylab = "Proportion of Simulations", bty='n',
                 xlim = c(-0.02, 50), ylim = c(0.4,1),
                 xaxt ='n')
            axis(side = 1, at = (sigma_true*1000), as.character(sigma_true*1000), las=2)
        }else{
            lines(mdl_parm_perf$sigma + jitter, mdl_parm_perf$Perf,type = 'b', pch = 16, col =plot_col2)

        }
        title(titles[j], main.font=1)

    }
}
legend(0, 0.5, legend = models[2:5],
       col = c(point_cols2[[2]], point_cols2[[3]], point_cols2[[4]], point_cols2[[5]]),
       lty = 1, bty='n')

## -----------------------------------------------------------------------------
# Figure 3 Scatter Plot of LS estimator vs Bayes estimator
## -----------------------------------------------------------------------------
axis_lims = list(c(-0.2*1000, 0.5*1000), c(-0.5*1000, 0.2*1000), c(-150,100), c(0*1000, 0.12*1000))
true_vals_vec = c(r_true*1000, theta_true, sigma_true[6]*1000)
results_40 = results[results[,1]==40/1000,] # This can be changed to explore other noise levels
results_40[,c(4,5,7)] = results_40[,c(4,5,7)]*1000
ylabs = c("Vague 1", "Vague 2", "Weakly Informative", "Informative")
titles = c(TeX("$r_x \\, (mm)$"), TeX("$r_y \\, (mm)$"), TeX('$\\theta_1 \\, (deg)'), TeX("$\\sigma \\, (mm)$"))

par(mfrow=c(4,4), mai = c(0.6, 0.6, 0.0, 0))
for(i in 2:5){
    for(j in 1:4){
        ls_result = results_40[results_40[,2]==1, 3+j]
        mdl_result = results_40[results_40[,2]==i, 3+j]
        b_ls = sqrt((mdl_result-true_vals_vec[j])^2) - sqrt((ls_result - true_vals_vec[j])^2)

        b_ls = ifelse(b_ls<0, 1,2)
        cols = list(rgb(52/255, 152/255, 219/255, 0.2),
                    colors$pink)
        mycol = rep("", length(b_ls))
        for(ii in 1:length(b_ls)){
            mycol[ii] = cols[[b_ls[ii]]]
        }
        if(j ==1 & i ==5){
            plot(ls_result, mdl_result, pch = 16, col = mycol,
                 #asp = 1,
                 xlim = axis_lims[[j]], ylim = axis_lims[[j]],
                 xlab = "LS est",
                 ylab = ylabs[i-1],
                 bty='n')
        }else if (j >1 & i ==5 ){
            plot(ls_result, mdl_result, pch = 16, col = mycol,
                 #asp = 1,
                 xlim = axis_lims[[j]], ylim = axis_lims[[j]],
                 xlab = "LS est",
                 ylab = NA,
                 bty='n')
        }else if(j==1 & i < 5){
            plot(ls_result, mdl_result, pch = 16, col = mycol,
                 #asp = 1,
                 xlim = axis_lims[[j]], ylim = axis_lims[[j]],
                 xlab = NA,
                 ylab = ylabs[i-1],
                 bty='n')
        }else{
            plot(ls_result, mdl_result, pch = 16, col = mycol,
                 #asp = 1,
                 xlim = axis_lims[[j]], ylim = axis_lims[[j]],
                 xlab = NA, ylab = NA,
                 bty='n')
        }
        if(i ==2){
            title(titles[j], line =-1)
        }

        abline(0,1, col = colors$orange)
        points(true_vals_vec[j],true_vals_vec[j], pch=18, cex = 2,col = colors2$orange)

    }
}

########################### Section E  #########################
## Additional Figures not referenced in supplement
################################################################

# Scatter plot of error distribution for each model/noise level
ylabs = list(r1=TeX("$r_1 \\,  (mm)$"), r2=TeX("$r_2  \\, (mm)$"), theta1=TeX("$\\theta_1 \\, (^o)$"),
             sigma=TeX("$\\sigma \\, (mm)$"))
ylims = list(c(-0.1,0.6), c(-0.6, 0.4), c(-180, 60), c(-0.03, 0.12))

par(mfrow = c(4, 6),mai = c(0, 0.6, 0.0, 0), oma=c(0,0,0,0))
for(j in 1:4){
    for(i in 1:6){
        # i= 1 # noise level (1:6)
        # j = 1 # parameter (1:4)
        true_val_tmp = c(r_true, theta_true, sigma_true[i])

        sigma_results = results[(results[,1]==sigma_true[i]), ]
        plot_col = rep("", nrow(sigma_results))
        for(ii in 1:nrow(sigma_results)){
            plot_col[ii] = point_cols[[sigma_results[ii,2]]]
        }

        if(i != 1){
            plot(sigma_results[,2]+rnorm(nrow(sigma_results), 0, 0.1), sigma_results[,3+j],
                 pch = 16, col = plot_col,
                 ylab = NA, yaxt = 'n',
                 ylim =ylims[[j]],
                 xlab = NA, bty='n', xaxt = 'n')
        }else{
            plot(sigma_results[,2]+rnorm(nrow(sigma_results), 0, 0.1), sigma_results[,3+j],
                 pch = 16, col = plot_col,
                 ylab = ylabs[[j]],
                 ylim =ylims[[j]],
                 xlab = NA, bty='n', xaxt = 'n')
        }
        abline(h = true_val_tmp[j], col = colors$orange)

        if(j ==1){
            title(TeX(sprintf("$\\sigma$ = %.1fmm$", c(sigma_true[i]*1000))), line = -2)
        }
    }
}


compute.model.average <- function(B,P,params)
{
	region_enrichment <- c(params[3],params[2],params[4],params[1],params[5],params[2],params[6],params[1],0)
	combined.adj <- (abs(B)+params[8]) * (abs(P)+params[7])
	Q.test.I <- combined.adj; Q.test.I[which(B <= 0)] <- 0; Q.test.I[which(P <= 0)] <- 0;
	Q.test.II <- combined.adj; Q.test.II[which(B >= 0)] <- 0; Q.test.II[which(P <= 0)] <- 0;
	Q.test.III <- combined.adj; Q.test.III[which(B >= 0)] <- 0; Q.test.III[which(P >= 0)] <- 0;
	Q.test.IV <- combined.adj; Q.test.IV[which(B <= 0)] <- 0; Q.test.IV[which(P >= 0)] <- 0;
	Q.test.PDE <- combined.adj; Q.test.PDE[which(B != 0)] <- 0; Q.test.PDE[which(P <= 0)] <- 0;
	Q.test.NDE <- combined.adj; Q.test.NDE[which(B != 0)] <- 0; Q.test.NDE[which(P >= 0)] <- 0;
	Q.test.PB <- combined.adj; Q.test.PB[which(B <= 0)] <- 0; Q.test.PB[which(P != 0)] <- 0;
	Q.test.NB <- combined.adj; Q.test.NB[which(B >= 0)] <- 0; Q.test.NB[which(P != 0)] <- 0;
	Q.test <- region_enrichment[1] * Q.test.I + region_enrichment[2] * Q.test.PDE + region_enrichment[3] * Q.test.II + region_enrichment[4] * Q.test.PB + region_enrichment[5] * Q.test.III + region_enrichment[6] * Q.test.NDE + region_enrichment[7] * Q.test.IV + region_enrichment[8] * Q.test.NB
	Q.test
}

compute.model.average.new <- function(B,D,params)
{
	# params
	# (I,II,III,IV,B,D,Cb,Cd)
	Bs <- abs(B) + params[7]
	Ds <- abs(D) + params[8]
	M <- matrix(0,nrow=dim(B)[1],ncol=dim(B)[2])
	rindx <- intersect(which(B>0),which(D>0))								# I
	M[rindx] <- abs(Bs[rindx]) * abs(Ds[rindx]) * params[1]
	rindx <- intersect(which(B<0),which(D>0))								# II
	M[rindx] <- abs(Bs[rindx]) * abs(Ds[rindx]) * params[2]
	rindx <- intersect(which(B<0),which(D<0))								# III
	M[rindx] <- abs(Bs[rindx]) * abs(Ds[rindx]) * params[3]
	rindx <- intersect(which(B>0),which(D<0))								# IV
	M[rindx] <- abs(Bs[rindx]) * abs(Ds[rindx]) * params[4]
	rindx <- intersect(which(B!=0),which(D==0))							# B
	M[rindx] <- abs(Bs[rindx]) * abs(Ds[rindx]) * params[5]
	rindx <- intersect(which(B==0),which(D!=0))							# D
	M[rindx] <- abs(Bs[rindx]) * abs(Ds[rindx]) * params[6]
	M
}

plot.model.average.ranking <- function(B,P,params,n=10000,fname)
{
  region_enrichment <- c(params[3],params[2],params[4],params[1],params[5],params[2],params[6],params[1],0)
  combined.adj <- (abs(B)+params[8]) * (abs(P)+params[7])
  Q.test.I <- combined.adj; Q.test.I[which(B <= 0)] <- 0; Q.test.I[which(P <= 0)] <- 0;
  Q.test.II <- combined.adj; Q.test.II[which(B >= 0)] <- 0; Q.test.II[which(P <= 0)] <- 0;
  Q.test.III <- combined.adj; Q.test.III[which(B >= 0)] <- 0; Q.test.III[which(P >= 0)] <- 0;
  Q.test.IV <- combined.adj; Q.test.IV[which(B <= 0)] <- 0; Q.test.IV[which(P >= 0)] <- 0;
  Q.test.PDE <- combined.adj; Q.test.PDE[which(B != 0)] <- 0; Q.test.PDE[which(P <= 0)] <- 0;
  Q.test.NDE <- combined.adj; Q.test.NDE[which(B != 0)] <- 0; Q.test.NDE[which(P >= 0)] <- 0;
  Q.test.PB <- combined.adj; Q.test.PB[which(B <= 0)] <- 0; Q.test.PB[which(P != 0)] <- 0;
  Q.test.NB <- combined.adj; Q.test.NB[which(B >= 0)] <- 0; Q.test.NB[which(P != 0)] <- 0;
  Q.test <- region_enrichment[1] * Q.test.I + region_enrichment[2] * Q.test.PDE + region_enrichment[3] * Q.test.II + region_enrichment[4] * Q.test.PB + region_enrichment[5] * Q.test.III + region_enrichment[6] * Q.test.NDE + region_enrichment[7] * Q.test.IV + region_enrichment[8] * Q.test.NB

  regions <- matrix(0,ncol=dim(B)[2],nrow=dim(B)[1])
  regions[intersect(which(B > 0),which(P > 0))] <- 1	#IV
  regions[intersect(which(B < 0),which(P > 0))] <- 2	#III
  regions[intersect(which(B < 0),which(P < 0))] <- 3	#II
  regions[intersect(which(B > 0),which(P < 0))] <- 4	#I
  regions[intersect(which(B == 0),which(P != 0))] <- 5	#P
  regions[intersect(which(B != 0),which(P == 0))] <- 6	#B

  r.freq <- matrix(0,ncol=n,nrow=6)
  r.nzi <- which(regions!=0)
  q.nz <- Q.test[r.nzi]
  r.nz <- regions[r.nzi]
  r.rnk <- r.nz[sort.list(q.nz,decreasing=TRUE)]
  for (i in 2:n)
  {
    r.freq[,i] <- r.freq[,i-1]
    r.freq[r.rnk[i],i] <- r.freq[r.rnk[i],i] + 1
  }
  r.frac <- matrix(0,ncol=n,nrow=6)
  r.frac[1,] <- r.freq[1,]
  for (i in 2:6)
  {
    r.frac[i,] <- r.freq[i,] + r.frac[i-1,]
  }
  r.frac <- t(t(r.frac) / (seq(1,n)-1))
	jpeg(filename=fname,width=1800,height=1800,quality=100,res=300)
  plot(seq(1,n)-1,rep(1,times=n),ylim=c(0,1),type='n',xlab="Interaction rank",ylab="Cumulative representation")
  xx <- c(seq(1,n),seq(n,1))-1

  cols <- c(rgb(0.7,0,0.7),rgb(0.8,0.3,0.3),rgb(1,0.6,0.2),rgb(0.3,0.8,0.3),rgb(0.6,1,1),rgb(0.2,0.6,1))

  yy <- c(rep(0,times=n),rev(r.frac[1,]))
  polygon(xx,yy,col=cols[1])
  for (i in 1:5)
  {
    yy <- c(r.frac[i,],rev(r.frac[i+1,]))
    polygon(xx,yy,col=cols[i+1])
  }
	dev.off()
}


#'@title calcKristCaa
#'
#'@description Calculate air resistance (\code{Caa}) (dimensionless) using the
#'Kristensen method.
#'
#'@param shipType Ship type (vector of strings, see \code{\link{calcShipType}}).
#'Must align with \code{tankerBulkCarrierGCargoShipTypes} and
#' \code{containerShipTypes} groupings
#'@param dwt Ship maximum deadweight tonnage (vector of numericals, tonnage)
#'@param tankerBulkCarrierGCargoShipTypes Ship types specified in input
#'\code{shipTypes} to be modeled as tankers, bulk carriers and general cargo
#'vessels (vector of strings)
#'@param containerShipTypes Ship types specified in input \code{shipTypes} to be
#'modeled as container ships (vector of strings)
#'
#' @details
#' Models the effect of realistic hull roughness on resistance, which is not
#' captured in the frictional and residual resistance coefficients from tank
#' towing operations.
#'
#' This method this requires ship types to be grouped. Use the
#' \code{tankerBulkCarrierGCargoShipTypes}, \code{containerShipTypes} grouping
#' parameters to provide these ship type groupings. Any ship types not included
#' in these groupings will be considered as miscellaneous vessels.
#'
#' NOTE: within the container ship section of this calculation, estimations are
#' made of the ships TEU:\itemize{
#' \item Feeder: TEU = ((dwt/15.19)^(1/0.9814))
#' \item Panamax: TEU = ((dwt/28.81)^(1/0.902))
#' \item PostPanamax: TEU = ((dwt/37)^(1/0.875))
#' }
#' These come from: H.O. Kristensen (2016), "Revision of statistical analysis and determination of
#' regression formulas for main dimensions of container ships based on data from Clarkson".
#' Kristensen's SHIP DESMO model also uses another dwt/teu relation for post-panamax container ships
#' with breadth > 49m. This equation has no inverse to map dwt to TEU and thus all post panamaxs are
#' assumed to have the same dwt/teu relation for all breadths, described above.
#'
#' @return \code{Caa} (vector of numericals, dimensionless)
#'
#' @references
#'Kristensen, H. O. and Lutzen, M. 2013. "Prediction of Resistance and Propulsion
#'Power of Ships."
#'
#'Kristensen, H. O. 2016. "Revision of statistical analysis and determination of
#'regression formulas for main dimensions of container ships based on data from
#'Clarkson."
#'
#'\href{https://gitlab.gbar.dtu.dk/oceanwave3d/Ship-Desmo}{Kristensen, H. O.
#'"Ship-Desmo-Tool." https://gitlab.gbar.dtu.dk/oceanwave3d/Ship-Desmo}
#'
#'@family Kristensen Calculations
#'@family Resistance Calculations
#'
#' @examples
#' calcKristCaa(c("bulk.carrier","container.ship","other.tanker"),c(70000,191144,20000))
#' calcKristCaa(c("bulk.carrier","container.ship","other.tanker"),c(70000,191144,20000),
#'              tankerBulkCarrierGCargoShipTypes=c("bulk.carrier","other.tanker"))
#'
#' @export

calcKristCaa<-function(shipType,dwt,
                       tankerBulkCarrierGCargoShipTypes=c("tanker","general.cargo","chemical.tanker","liquified.gas.tanker","oil.tanker","other.tanker","bulk.carrier"),
                       containerShipTypes=c("container.ship")
){
  Caa<- ifelse(shipType%in%tankerBulkCarrierGCargoShipTypes,
               ifelse(#case 1:"Small","Handysize","Handymax",
                 dwt <= 55000,0.07,
                 ifelse(#case 2:"Panamax","Aframax","Suezmax"
                   dwt <= 200000,0.05,
                   #case 3: "VLCC","VLBC"
                   0.04
                 )#end case 3
               )# end case 2
               # end case 1
  #end non container ship group
  ,ifelse(shipType %in% containerShipTypes,#start container ship group
  ifelse( #case 1 Feeder
    dwt <= 35000, pmax(0.28*((dwt/15.19)^(1/0.9814))^-0.126,0.09),
    ifelse(#case 2 Panamax
      dwt <= 60000, pmax(0.28*((dwt/28.81)^(1/0.902))^-0.126,0.09),
      pmax(0.28*((dwt/37)^(1/0.875))^-0.126,0.09) #PostPanamax
    )# end case 2
  ),
  NA# end case 1
  #end container ship group
                            )
  )

  return(Caa/1000)
}


# This is a minimal version of the "sandwich" package suitable for use in this course
# It is based on the "sandwich" package: https://cran.r-project.org/web/packages/sandwich/index.html
# Modified 06-03-2021 and this file is licensed under the GPL-2 / GPL-3 license

sandwich <- function(x, bread. = bread, meat. = meat, ...)
{
  if(is.list(x) && !is.null(x$na.action)) class(x$na.action) <- "omit"
  if(is.function(bread.)) bread. <- bread.(x)
  if(is.function(meat.)) meat. <- meat.(x, ...)
  n <- NROW(estfun(x))
  return(1/n * (bread. %*% meat. %*% bread.))
}

meat <- function(x, adjust = FALSE, ...)
{
  if(is.list(x) && !is.null(x$na.action)) class(x$na.action) <- "omit"
  psi <- estfun(x, ...)
  k <- NCOL(psi)
  n <- NROW(psi)
  rval <- crossprod(as.matrix(psi))/n
  if(adjust) rval <- n/(n-k) * rval
  rownames(rval) <- colnames(rval) <- colnames(psi)
  return(rval)
}

vcovOPG <- function(x, adjust = FALSE, ...)
{
  if(is.list(x) && !is.null(x$na.action)) class(x$na.action) <- "omit"
  psi <- estfun(x, ...)
  k <- NCOL(psi)
  n <- NROW(psi)
  rval <- chol2inv(qr.R(qr(psi)))
  if(adjust) rval <- n/(n-k) * rval
  rownames(rval) <- colnames(rval) <- colnames(psi)
  return(rval)
}

vcovHC <- function(x, ...) {
  UseMethod("vcovHC")
}

vcovHC.default <- function(x,
  type = c("HC3", "const", "HC", "HC0", "HC1", "HC2", "HC4", "HC4m", "HC5"),
  omega = NULL, sandwich = TRUE, ...)
{
  type <- match.arg(type)
  rval <- meatHC(x, type = type, omega = omega)
  if(sandwich) rval <- sandwich(x, meat. = rval, ...)
  return(rval)
}

meatHC <- function(x,
  type = c("HC3", "const", "HC", "HC0", "HC1", "HC2", "HC4", "HC4m", "HC5"),
  omega = NULL, ...)
{
  ## ensure that NAs are omitted
  if(is.list(x) && !is.null(x$na.action)) class(x$na.action) <- "omit"

  ## extract X
  X <- model.matrix(x)
  if(any(alias <- is.na(coef(x)))) X <- X[, !alias, drop = FALSE]
  attr(X, "assign") <- NULL
  n <- NROW(X)

  ## get hat values and residual degrees of freedom
  diaghat <- try(hatvalues(x), silent = TRUE)
  df <- n - NCOL(X)

  ## the following might work, but "intercept" is also claimed for "coxph"
  ## res <- if(attr(terms(x), "intercept") > 0) estfun(x)[,1] else rowMeans(estfun(x)/X, na.rm = TRUE)
  ## hence better rely on
  ef <- estfun(x, ...)
  res <- rowMeans(ef/X, na.rm = TRUE)

  ## handle rows with just zeros
  all0 <- apply(abs(ef) < .Machine$double.eps, 1L, all)
  res[all0] <- 0
  ## in case of lm/glm and type = "const" re-obtain the working residuals
  if(any(all0) && substr(type, 1L, 1L) == "c") {
    if(inherits(x, "glm")) {
      res <- as.vector(residuals(x, "working")) * weights(x, "working")
      if(!(substr(x$family$family, 1L, 17L) %in% c("poisson", "binomial", "Negative Binomial"))) {
        res <- res * sum(weights(x, "working"), na.rm = TRUE) / sum(res^2, na.rm = TRUE)
      }
    } else if(inherits(x, "lm")) {
      res <- as.vector(residuals(x))
      if(!is.null(weights(x))) res <- res * weights(x)
    }
  }

  ## if omega not specified, set up using type
  if(is.null(omega)) {
    type <- match.arg(type)
    if(type == "HC") type <- "HC0"
    switch(type,
      "const" = { omega <- function(residuals, diaghat, df) rep(1, length(residuals)) * sum(residuals^2)/df },
      "HC0"   = { omega <- function(residuals, diaghat, df) residuals^2 },
      "HC1"   = { omega <- function(residuals, diaghat, df) residuals^2 * length(residuals)/df },
      "HC2"   = { omega <- function(residuals, diaghat, df) residuals^2 / (1 - diaghat) },
      "HC3"   = { omega <- function(residuals, diaghat, df) residuals^2 / (1 - diaghat)^2 },
      "HC4"   = { omega <- function(residuals, diaghat, df) {
        n <- length(residuals)
	p <- as.integer(round(sum(diaghat),  digits = 0))
	delta <- pmin(4, n * diaghat/p)
        residuals^2 / (1 - diaghat)^delta
      }},
      "HC4m"  = { omega <- function(residuals, diaghat, df) {
        gamma <- c(1.0, 1.5) ## as recommended by Cribari-Neto & Da Silva
        n <- length(residuals)
	p <- as.integer(round(sum(diaghat),  digits = 0))
	delta <- pmin(gamma[1], n * diaghat/p) + pmin(gamma[2], n * diaghat/p)
        residuals^2 / (1 - diaghat)^delta
      }},
      "HC5"   = { omega <- function(residuals, diaghat, df) {
        k <- 0.7 ## as recommended by Cribari-Neto et al.
        n <- length(residuals)
	p <- as.integer(round(sum(diaghat),  digits = 0))
	delta <- pmin(n * diaghat/p, pmax(4, n * k * max(diaghat)/p))
        residuals^2 / sqrt((1 - diaghat)^delta)
      }}
    )
  }

  ## process omega
  if(is.function(omega)) omega <- omega(res, diaghat, df)
  rval <- sqrt(omega) * X
  rval <- crossprod(rval)/n

  return(rval)
}

vcovHC.mlm <- function(x,
  type = c("HC3", "const", "HC", "HC0", "HC1", "HC2", "HC4", "HC4m", "HC5"),
  omega = NULL, sandwich = TRUE, ...)
{
  ## compute meat "by hand" because meatHC() can not deal with
  ## residual "matrices"

  ## ensure that NAs are omitted
  if(is.list(x) && !is.null(x$na.action)) class(x$na.action) <- "omit"

  ## regressors, df, hat values
  X <- model.matrix(x)
  attr(X, "assign") <- NULL
  if(any(alias <- apply(coef(x), 1, function(z) all(is.na(z))))) X <- X[, !alias, drop = FALSE]
  n <- NROW(X)
  diaghat <- hatvalues(x)
  df <- n - NCOL(X)

  ## residuals
  res <- residuals(x)

  ## if omega not specified, set up using type
  if(is.null(omega)) {
    type <- match.arg(type)
    if(type == "HC") type <- "HC0"
    switch(type,
      "const" = {
        warning("not implemented for 'mlm' objects, using vcov() instead")
	return(vcov(x))
       },
      "HC0"   = { omega <- function(residuals, diaghat, df) residuals^2 },
      "HC1"   = { omega <- function(residuals, diaghat, df) residuals^2 * length(residuals)/df },
      "HC2"   = { omega <- function(residuals, diaghat, df) residuals^2 / (1 - diaghat) },
      "HC3"   = { omega <- function(residuals, diaghat, df) residuals^2 / (1 - diaghat)^2 },
      "HC4"   = { omega <- function(residuals, diaghat, df) {
        n <- length(residuals)
	p <- as.integer(round(sum(diaghat),  digits = 0))
	delta <- pmin(4, n * diaghat/p)
        residuals^2 / (1 - diaghat)^delta
      }},
      "HC4m"  = { omega <- function(residuals, diaghat, df) {
        gamma <- c(1.0, 1.5)
        n <- length(residuals)
	p <- as.integer(round(sum(diaghat),  digits = 0))
	delta <- pmin(gamma[1], n * diaghat/p) + pmin(gamma[2], n * diaghat/p)
        residuals^2 / (1 - diaghat)^delta
      }},
      "HC5"   = { omega <- function(residuals, diaghat, df) {
        k <- 0.7
        n <- length(residuals)
	p <- as.integer(round(sum(diaghat),  digits = 0))
	delta <- pmin(n * diaghat/p, pmax(4, n * k * max(diaghat)/p))
        residuals^2 / sqrt((1 - diaghat)^delta)
      }}
    )
  }

  ## process omega
  omega <- apply(res, 2, function(r) omega(r, diaghat, df))

  ## compute crossproduct X' Omega X
  rval <- lapply(1:ncol(omega), function(i) as.vector(sign(res[,i]) * sqrt(omega[,i])) * X)
  rval <- do.call("cbind", rval)
  rval <- crossprod(rval)/n
  colnames(rval) <- rownames(rval) <- colnames(vcov(x))

  ## if necessary compute full sandwich
  if(sandwich) rval <- sandwich(x, meat. = rval, ...)
  return(rval)
}

                 estfun <- function(x, ...)
{
  UseMethod("estfun")
}

estfun.lm <- function(x, ...)
{
  xmat <- model.matrix(x)
  xmat <- naresid(x$na.action, xmat)
  if(any(alias <- is.na(coef(x)))) xmat <- xmat[, !alias, drop = FALSE]
  wts <- weights(x)
  if(is.null(wts)) wts <- 1
  res <- residuals(x)
  rval <- as.vector(res) * wts * xmat
  attr(rval, "assign") <- NULL
  attr(rval, "contrasts") <- NULL
  if(is.zoo(res)) rval <- zoo(rval, index(res), attr(res, "frequency"))
  if(is.ts(res)) rval <- ts(rval, start = start(res), frequency = frequency(res))
  return(rval)
}

estfun.mlm <- function(x, ...)
{
  xmat <- model.matrix(x)
  xmat <- naresid(x$na.action, xmat)
  wts <- weights(x)
  if(is.null(wts)) wts <- 1
  res <- residuals(x)
  cf <- coef(x)
  rval <- lapply(1:NCOL(res), function(i) {
    rv <- as.vector(res[,i]) * wts * xmat
    colnames(rv) <- paste(colnames(cf)[i], colnames(rv), sep = ":")
    rv
  })
  rval <- do.call("cbind", rval)
  attr(rval, "assign") <- NULL
  attr(rval, "contrasts") <- NULL
  if(any(alias <- is.na(as.vector(cf)))) rval <- rval[, !alias, drop = FALSE]
  if(is.zoo(res)) rval <- zoo(rval, index(res), attr(res, "frequency"))
  if(is.ts(res)) rval <- ts(rval, start = start(res), frequency = frequency(res))
  return(rval)
}

estfun.glm <- function(x, ...)
{
  xmat <- model.matrix(x)
  xmat <- naresid(x$na.action, xmat)
  if(any(alias <- is.na(coef(x)))) xmat <- xmat[, !alias, drop = FALSE]
  wres <- as.vector(residuals(x, "working")) * weights(x, "working")
  dispersion <- if(substr(x$family$family, 1, 17) %in% c("poisson", "binomial", "Negative Binomial")) 1
    else sum(wres^2, na.rm = TRUE)/sum(weights(x, "working"), na.rm = TRUE)
  rval <- wres * xmat / dispersion
  attr(rval, "assign") <- NULL
  attr(rval, "contrasts") <- NULL
  res <- residuals(x, type = "pearson")
  if(is.ts(res)) rval <- ts(rval, start = start(res), frequency = frequency(res))
  if(is.zoo(res)) rval <- zoo(rval, index(res), attr(res, "frequency"))
  return(rval)
}

estfun.rlm <- function(x, ...)
{
  xmat <- model.matrix(x)
  xmat <- naresid(x$na.action, xmat)
  wts <- weights(x)
  if(is.null(wts)) wts <- 1
  res <- residuals(x)
  psi <- function(z) x$psi(z) * z
  rval <- as.vector(psi(res/x$s)) * wts * xmat
  attr(rval, "assign") <- NULL
  attr(rval, "contrasts") <- NULL
  if(is.ts(res)) rval <- ts(rval, start = start(res), frequency = frequency(res))
  if(is.zoo(res)) rval <- zoo(rval, index(res), attr(res, "frequency"))
  return(rval)
}

estfun.polr <- function(x, ...)
{
  ## link processing
  mueta <- x$method
  if(mueta == "logistic") mueta <- "logit"
  mueta <- make.link(mueta)$mu.eta

  ## observations
  xmat <- model.matrix(x)[, -1L, drop = FALSE]
  n <- nrow(xmat)
  k <- ncol(xmat)
  m <- length(x$zeta)
  mf <- model.frame(x)
  y <- as.numeric(model.response(mf))
  w <- model.weights(mf)
  if(is.null(w)) w <- rep(1, n)

  ## estimates
  prob <- x$fitted.values[cbind(1:n, y)]
  xb <- if(k >= 1L) as.vector(xmat %*% x$coefficients) else rep(0, n)
  zeta <- x$zeta
  lp <- cbind(0, mueta(matrix(zeta, nrow = n, ncol = m, byrow = TRUE) - xb), 0)

  ## estimating functions
  rval <- matrix(0, nrow = n, ncol = k + m + 2L)
  if(k >= 1L) rval[, 1L:k] <- (-xmat * as.vector(lp[cbind(1:n, y + 1L)] - lp[cbind(1:n, y)]))
  rval[cbind(1:n, k + y)] <- -as.vector(lp[cbind(1:n, y)])
  rval[cbind(1:n, k + y + 1L)] <- as.vector(lp[cbind(1:n, y + 1L)])
  rval <- rval[, -c(k + 1L, k + m + 2L), drop = FALSE]
  rval <- w/prob * rval

  ## dimnames and return
  dimnames(rval) <- list(rownames(xmat), c(colnames(xmat), names(x$zeta)))
  return(rval)
}

estfun.clm <- function(x, ...)
{
  if(x$threshold != "flexible") stop("only flexible thresholds implemented at the moment")

  ## link processing
  mueta <- make.link(x$link)$mu.eta

  ## observations
  xmat <- model.matrix(x)
  if(length(xmat) > 1L) stop("estimating functions for scale regression not implemented yet")
  xmat <- xmat$X[, -1L, drop = FALSE]
  n <- nrow(xmat)
  k <- ncol(xmat)
  m <- length(x$alpha)
  mf <- model.frame(x)
  y <- as.numeric(model.response(mf))
  w <- model.weights(mf)
  if(is.null(w)) w <- rep(1, n)

  ## estimates
  prob <- x$fitted.values
  xb <- if(k >= 1L) as.vector(xmat %*% x$beta) else rep(0, n)
  zeta <- x$alpha
  lp <- cbind(0, mueta(matrix(zeta, nrow = n, ncol = m, byrow = TRUE) - xb), 0)

  ## estimating functions
  rval <- matrix(0, nrow = n, ncol = k + m + 2L)
  if(k >= 1L) rval[, 1L:k] <- (-xmat * as.vector(lp[cbind(1:n, y + 1L)] - lp[cbind(1:n, y)]))
  rval[cbind(1:n, k + y)] <- -as.vector(lp[cbind(1:n, y)])
  rval[cbind(1:n, k + y + 1L)] <- as.vector(lp[cbind(1:n, y + 1L)])
  rval <- rval[, -c(k + 1L, k + m + 2L), drop = FALSE]
  rval <- w/prob * rval

  ## dimnames, re-order and return
  dimnames(rval) <- list(rownames(xmat), c(colnames(xmat), names(x$alpha)))
  ix <- if(k >= 1L) c((k + 1L):(k + m), 1L:k) else 1L:m
  return(rval[, ix, drop = FALSE])
}

estfun.coxph <- function(x, ...)
{
  wts <- x$weights
  if(is.null(wts)) wts <- 1
  wts * residuals(x, type = "score", ...)
}

estfun.survreg <- function(x, ...)
{
  mf <- model.frame(x)
  xmat <- model.matrix(terms(x), mf)
  wts <- model.weights(mf)
  if(is.null(wts)) wts <- 1
  res <- residuals(x, type = "matrix")
  rval <- as.vector(res[,"dg"]) * wts * xmat
  if(NROW(x$var) > length(coef(x))) {
    rval <- cbind(rval, res[,"ds"])
    colnames(rval)[NCOL(rval)] <- "Log(scale)"
  }
  attr(rval, "assign") <- NULL

  return(rval)
}

estfun.nls <- function(x, ...)
{
  rval <- as.vector(x$m$resid()) * x$m$gradient()
  colnames(rval) <- names(coef(x))
  rval
}

estfun.hurdle <- function(x, ...) {
  ## extract data
  Y <- if(is.null(x$y)) model.response(model.frame(x)) else x$y
  X <- model.matrix(x, model = "count")
  Z <- model.matrix(x, model = "zero")
  beta <- coef(x, model = "count")
  gamma <- coef(x, model = "zero")
  fulltheta <- x$theta

  offset <- x$offset
  if(is.list(offset)) {
    offsetx <- offset$count
    offsetz <- offset$zero
  } else {
    offsetx <- offset
    offsetz <- NULL
  }
  if(is.null(offsetx)) offsetx <- 0
  if(is.null(offsetz)) offsetz <- 0
  if(x$dist$zero == "binomial") linkobj <- make.link(x$link)
  wts <- weights(x)
  if(is.null(wts)) wts <- 1
  Y1 <- Y > 0

  ## count component: working residuals
  eta <- as.vector(X %*% beta + offsetx)
  mu <- exp(eta)
  theta <- fulltheta["count"]

  wres_count <- as.numeric(Y > 0) * switch(x$dist$count,
    "poisson" = {
      (Y - mu) - exp(ppois(0, lambda = mu, log.p = TRUE) -
        ppois(0, lambda = mu, lower.tail = FALSE, log.p = TRUE) + eta)
    },
    "geometric" = {
      (Y - mu * (Y + 1)/(mu + 1)) - exp(pnbinom(0, mu = mu, size = 1, log.p = TRUE) -
        pnbinom(0, mu = mu, size = 1, lower.tail = FALSE, log.p = TRUE) - log(mu + 1) + eta)
    },
    "negbin" = {
      (Y - mu * (Y + theta)/(mu + theta)) - exp(pnbinom(0, mu = mu, size = theta, log.p = TRUE) -
        pnbinom(0, mu = mu, size = theta, lower.tail = FALSE, log.p = TRUE) +
	log(theta) - log(mu + theta) + eta)
    })

  ## zero component: working residuals
  eta <- as.vector(Z %*% gamma + offsetz)
  mu <- if(x$dist$zero == "binomial") linkobj$linkinv(eta) else exp(eta)
  theta <- fulltheta["zero"]

  wres_zero <- switch(x$dist$zero,
    "poisson" = {
      ifelse(Y1, exp(ppois(0, lambda = mu, log.p = TRUE) -
        ppois(0, lambda = mu, lower.tail = FALSE, log.p = TRUE) + eta), -mu)
    },
    "geometric" = {
      ifelse(Y1, exp(pnbinom(0, mu = mu, size = 1, log.p = TRUE) -
        pnbinom(0, mu = mu, size = 1, lower.tail = FALSE, log.p = TRUE) - log(mu + 1) + eta), -mu/(mu + 1))
    },
    "negbin" = {
      ifelse(Y1, exp(pnbinom(0, mu = mu, size = theta, log.p = TRUE) -
        pnbinom(0, mu = mu, size = theta, lower.tail = FALSE, log.p = TRUE) +
        log(theta) - log(mu + theta) + eta), -mu * theta/(mu + theta))
    },
    "binomial" = {
      ifelse(Y1, 1/mu, -1/(1-mu)) * linkobj$mu.eta(eta)
    })

  ## compute gradient from data
  rval <- cbind(wres_count * wts * X, wres_zero * wts * Z)
  colnames(rval) <- names(coef(x))
  rownames(rval) <- rownames(X)
  return(rval)
}

estfun.zeroinfl <- function(x, ...) {
  ## extract data
  Y <- if(is.null(x$y)) model.response(model.frame(x)) else x$y
  X <- model.matrix(x, model = "count")
  Z <- model.matrix(x, model = "zero")
  beta <- coef(x, model = "count")
  gamma <- coef(x, model = "zero")
  theta <- x$theta

  offset <- x$offset
  if(is.list(offset)) {
    offsetx <- offset$count
    offsetz <- offset$zero
  } else {
    offsetx <- offset
    offsetz <- NULL
  }
  if(is.null(offsetx)) offsetx <- 0
  if(is.null(offsetz)) offsetz <- 0
  linkobj <- make.link(x$link)
  wts <- weights(x)
  if(is.null(wts)) wts <- 1
  Y1 <- Y > 0

  eta <- as.vector(X %*% beta + offsetx)
  mu <- exp(eta)
  etaz <- as.vector(Z %*% gamma + offsetz)
  muz <- linkobj$linkinv(etaz)

  ## density for y = 0
  clogdens0 <- switch(x$dist,
    "poisson" = -mu,
    "geometric" = dnbinom(0, size = 1, mu = mu, log = TRUE),
    "negbin" = dnbinom(0, size = theta, mu = mu, log = TRUE))
  dens0 <- muz * (1 - as.numeric(Y1)) + exp(log(1 - muz) + clogdens0)

  ## working residuals
  wres_count <- switch(x$dist,
    "poisson" = ifelse(Y1, Y - mu, -exp(-log(dens0) + log(1 - muz) + clogdens0 + log(mu))),
    "geometric" = ifelse(Y1, Y - mu * (Y + 1)/(mu + 1), -exp(-log(dens0) +
      log(1 - muz) + clogdens0 - log(mu + 1) + log(mu))),
    "negbin" = ifelse(Y1, Y - mu * (Y + theta)/(mu + theta), -exp(-log(dens0) +
      log(1 - muz) + clogdens0 + log(theta) - log(mu + theta) + log(mu))))
  wres_zero <- ifelse(Y1, -1/(1-muz) * linkobj$mu.eta(etaz),
    (linkobj$mu.eta(etaz) - exp(clogdens0) * linkobj$mu.eta(etaz))/dens0)

  ## compute gradient from data
  rval <- cbind(wres_count * wts * X, wres_zero * wts * Z)
  colnames(rval) <- names(coef(x))
  rownames(rval) <- rownames(X)
  return(rval)
}

estfun.mlogit <- function(x, ...)
{
  x$gradient
}

                 bread <- function(x, ...)
{
  UseMethod("bread")
}

bread.default <- function(x, ...) {
  nobs0(x) * vcov0(x, ...)
}

bread.lm <- function(x, ...)
{
  if(!is.null(x$na.action)) class(x$na.action) <- "omit"
  sx <- summary.lm(x)
  return(sx$cov.unscaled * as.vector(sum(sx$df[1L:2L])))
}

bread.mlm <- function(x, ...)
{
  if(!is.null(x$na.action)) class(x$na.action) <- "omit"
  cf <- coef(x)
  rval <- summary.lm(x)
  rval <- kronecker(
    structure(diag(ncol(cf)), .Dimnames = rep.int(list(colnames(cf)), 2L)),
    structure(rval$cov.unscaled,  .Dimnames = rep.int(list(rownames(cf)), 2L)) * as.vector(sum(rval$df[1L:2L])),
    make.dimnames = TRUE
  )
  return(rval)
}

bread.glm <- function(x, ...)
{
  if(!is.null(x$na.action)) class(x$na.action) <- "omit"
  sx <- summary(x)
  wres <- as.vector(residuals(x, "working")) * weights(x, "working")
  dispersion <- if(substr(x$family$family, 1L, 17L) %in% c("poisson", "binomial", "Negative Binomial")) 1
    else sum(wres^2)/sum(weights(x, "working"))
  return(sx$cov.unscaled * as.vector(sum(sx$df[1L:2L])) * dispersion)
}

bread.nls <- function(x, ...)
{
  if(!is.null(x$na.action)) class(x$na.action) <- "omit"
  sx <- summary(x)
  return(sx$cov.unscaled * as.vector(sum(sx$df[1L:2L])))
}

bread.polr <- function(x, ...)
{
  vcov(x) * x$n
}

bread.clm <- function(x, ...)
{
  vcov(x) * x$n
}

bread.survreg <- function(x, ...)
  length(x$linear.predictors) * x$var

bread.gam <- function(x, ...)
{
  if(!is.null(x$na.action)) class(x$na.action) <- "omit"
  sx <- summary(x)
  sx$cov.unscaled * sx$n
}

bread.coxph <- function(x, ...)
{
  rval <- x$var * x$n
  dimnames(rval) <- list(names(coef(x)), names(coef(x)))
  return(rval)
}

bread.hurdle <- function(x, ...)
{
  x$vcov * x$n
}

bread.zeroinfl <- function(x, ...)
{
  x$vcov * x$n
}

bread.mlogit <- function(x, ...)
{
  if(!is.null(x$na.action)) class(x$na.action) <- "omit"
  vcov(x) * length(residuals(x))
}

bread.rlm <- function(x, ...)
{
  if(!is.null(x$na.action)) class(x$na.action) <- "omit"
  xmat <- model.matrix(x)
  xmat <- naresid(x$na.action, xmat)
  wts <- weights(x)
  if(is.null(wts)) wts <- 1
  res <- residuals(x)
  psi_deriv <- function(z) x$psi(z, deriv = 1)
  rval <- sqrt(abs(as.vector(psi_deriv(res/x$s)/x$s))) * wts * xmat
  rval <- chol2inv(qr.R(qr(rval))) * nrow(xmat)
  return(rval)
}


#' knots2ms
#'
#' Conversion speed from knots in meter per second.
#'
#' @param numeric  knots Speed in knots.
#' @return
#'
#'
#' @author  Istituto di Biometeorologia Firenze Italy  Alfonso Crisci \email{a.crisci@@ibimet.cnr.it}
#' @keywords  knots2mph
#'
#' @export
#'
#'
#'
#'
knots2ms<-function(knots)
{
  return(knots * 0.514444);
}


#' Deterministic model function
#'
#' Deterministic daily age structured model
#'
#' @param params A list with the following parameters included:
#'
#' ini.prev = percentage of individuals infected at the start (scalar value between 0 and 1),
#'
#' n.age.cats = number of age categories to monitor. Currently makes the most sense to be by decade for 10 categories (must be scalar value greater than 3).
#'
#' n.i.cats = number of I subcategories
#'
#' n.e.cats = number of E subcategories
#'
#' e.move = rate of movement in the exposed categories (scalar values between 0 and 1).
#'
#' i.move = rate of movement in the infectious categories (scalar between 0 and 1)
#'
#' beta = transmission coefficient
#'
#' contacts = symmetrical matrix of n.age.cats*n.age.cats of the contact rate between age categories. 1 = average, 1.1 = 10% increase, .9 = 10% decrease.
#'
#' n0 = initial population size (scalar value greater than 0)
#'
#' n.days = number of days to run the model (scalar value greater than 2),
#'
#' severity is a vector of length n.age.cats of the proportion of I that goes to the hospital
#'
#' d.no is the proportion of severe cases that die per day w/o hospital care (currently scalar)
#'
#' d.yes is the proportion that die with hospital care
#'
#' beds is the # of beds or ICUs that are available
#'
#' @return A list with 3 outputs:
#'
#' 1. List of counts of the # of individuals in S, E, I, R categories over time with the following columns: age category, day of simulation, sub category (for E and I), n = number of individuals
#'
#' 2. List of outcomes-- Dt = deaths, Ht = hospitalizations needed
#'
#'  Lists for Dt and Ht with the following columns: age category, day of the simulation, n = # of individuals
#'
#' 3. R0 = basic disease reproductive number
#'
#' @importFrom popbio stable.stage
#' @importFrom dplyr rename mutate
#' @importFrom reshape2 melt
#' @importFrom stats rgamma
#' @importFrom magrittr %>%
#' @examples
#' params <- list(ini.prev = 0.001, n.age.cats = 10,  n.e.cats = 6, n.i.cats = 5,
#'               e.move = 0.99, i.move = 0.99, beta = .35, theta = 1,
#'               n0 = 10000, n.days = 360, d.no = 0.05, d.yes = 0.01, beds = 20,
#'               contacts = matrix(1, nrow = 10, ncol = 10),
#'               severity = c(0, 0, 0.001, 0.002, 0.004, 0.010, 0.020, 0.044,
#'               0.093, 0.200))
#'
#' out <- det_model(params)
#'
#' @export

det_model <- function(params) {
  #### warnings ####
  # write the list objects to the local environment
  for (v in 1:length(params)) assign(names(params)[v], params[[v]])
  if(exists("ini.prev")==FALSE){
    message("initial prevalence is missing, using default value")
    ini.ad.m.prev <- 0.01
  }
  if(exists("n.age.cats")==FALSE){
    message("# of age categories is missing, using default value")
    n.age.cats <- 10
  }

  if(exists("e.move")==FALSE){
    message("e.move is missing, using default value")
    e.move <- 0.4
  }

  if(exists("i.move")==FALSE){
    message("i.move is missing, using default value")
    i.move <- 0.4
  }

  if(exists("beta")==FALSE){
    message("transmission beta is missing, using default value")
    beta <- 0.08
  }

  if(exists("theta")==FALSE){
    message("theta is missing, using default value")
    theta <- 1
  }

  if(exists("n0")==FALSE){
    message("initial population size n0 is missing, using default value")
    n0 <- 10000
  }

  if(exists("n.days")==FALSE){
    message("n.days is missing, using default value")
    n.days <- 360
  }

  ##### check parameter values #####
  if(ini.prev < 0) warning("ini.prev must >=0")
  if(ini.prev > 1) warning("ini.prev must be <= 1")

  if(n.age.cats < 4) warning("n.age.cats must be 4 or more")
  if(e.move < 0) warning("e.move must be between 0 and 1")
  if(e.move > 1) warning("e.move must be between 0 and 1")
  if(i.move < 0) warning("i.move must be between 0 and 1")
  if(i.move > 1) warning("i.move must be between 0 and 1")
  if(beta < 0) warning("beta cannot be negative")
  if(n0 <= 0) warning("n0 must be positive")
  if(n.days <= 0) warning("n.days must be positive")
  if(n.e.cats <= 2) warning("n.e.cats must be >= 2")
  if(n.i.cats <= 2) warning("n.i.cats must be >= 2")

  #### INITIAL CONDITIONS ####

  #NOTE!!! currently making up these survival and repro values
  #NOTE!!! n.age.cats not fully flexible, best if n.age.cats = 10 (1 for each decade)
  prop.aging <- 1/n.age.cats # proportion that age out of a category (approx.)
  mort <- rep(0.001, n.age.cats - 4)
  old.mort <- rep(0.1, 4) # last 4 categories

  remain <- rep(1, n.age.cats) - prop.aging - c(mort, old.mort)
  repro <- 0.08

  # Create the matrix to start the population at stable age dist
  # Probably should just fix n.age.cats and feed in the current US age dist.
  M <- matrix(rep(0, n.age.cats * n.age.cats), nrow = n.age.cats)
  diag(M) <- remain
  M[row(M) == (col(M) + 1)] <- prop.aging # off diagonal
  M[n.age.cats, n.age.cats] <- 1-old.mort[1] # old category mortality

  # insert the fecundity vector for prebirth census
  # only used for the stable age dist
  M[1, 2:4] <- repro * (1-mort[1]) * 0.5

  #lambda(M)
  #plot(stable.stage(M))

  # pre-allocate the output matrices
  tmp <- matrix(0, nrow = n.age.cats, ncol = n.days)
  St <- tmp  # susceptible
  Et <- array(rep(tmp), dim = c(n.age.cats, n.days, n.e.cats)) # exposed
  It <- array(rep(tmp), dim = c(n.age.cats, n.days, n.i.cats))  # infectious
  Rt <- tmp # recovered

  # for tracking purposes
  Dt <- tmp # disease deaths
  Ht <- tmp  # hospitalizations

  # Intializing with the stable age distribution.
  St[, 1] <- popbio::stable.stage(M) * n0 * (1 - ini.prev)

  # equally allocating prevalence across ages,
  # currently putting them in the 1st I category
  It[, 1, 1] <- popbio::stable.stage(M) * n0 * ini.prev

  # calculate R0 (not sure about this)
  R0 <- (1-exp (-(beta * n0) / n0^theta)) *
    mean(rgamma(1000, n.i.cats, i.move))
  # assumes no mortality during infectious period and no effect of
  # hospitalization or containment
  # doesn't account for contact matrix. Probably need an NGM approach

  #### MODEL ####
  for (t in 2:(n.days)) {
    # bring forward previous month, then modify it
    St[, t]   <- St[, t - 1]
    Et[, t, ] <- Et[, t - 1, ]
    It[, t, ] <- It[, t - 1, ]
    Rt[, t]   <- Rt[, t - 1]

    # Move E individuals forward in their subcategories
    e.movers <- Et[, t, ] * e.move

    Et[, t, 1] <- Et[, t, 1] - e.movers[, 1]
    Et[, t, 2:n.e.cats] <- Et[, t, 2:n.e.cats] - e.movers[, 2:n.e.cats] +
      e.movers[, 1:(n.e.cats-1)]

    # Exposed become infectious
    It[, t, 1] <- It[, t, 1] + e.movers[ , n.e.cats]

    # Move I individuals forward in their subcategories
    i.movers <- It[, t, ] * i.move

    It[, t, 1] <- It[, t, 1] - i.movers[, 1]
    It[, t, 2:n.i.cats] <- It[, t, 2:n.i.cats] - i.movers[, 2:n.i.cats] +
                            i.movers[, 1:(n.i.cats-1)]

    # infectious individs exit from the last I cat into recovered
    Rt[ , t] <- Rt[ , t] + i.movers[ , n.i.cats]

    ##### Transmission #####
    I.current <- rowSums(It[ ,t, ]) # current # of I's by age
    Nall <- sum(St[, t]) + sum(Et[, t, ]) + sum(It[, t, ]) + sum(Rt[, t]) #total

    # probably a faster way to do this
    cases <- St[, t] * contacts %*%
                       (1 - exp( -(beta * (as.matrix(I.current)/ Nall ^ theta))))

    St[, t] <- St[, t] - cases
    Et[, t, 1] <- Et[, t, 1] + cases

    ##### Hospitalizations needed #####
    Ht[ ,t] <- severity * I.current

    ##### Deaths #####
    # hospital beds are allocated proportion to case load
    # maybe modify in the future for different options
    hos.yes <- min(c(beds * Ht[,t]/sum(Ht[,t]), Ht[,t])) # hospitalized
    hos.no <- Ht[,t] - hos.yes # not able to be hospitalized

    # deaths depend on whether they are hospitalized
    Dt[ ,t] <- hos.no * d.no + hos.yes * d.yes

    # remove those deaths from the I category
    # currently not skewed by time in the infectious class
    It[ , t, ] <- It[ , t, ] - (It[ , t, ] / I.current) * Dt[ ,t]
  }

  ##### organize the output #####
  St <- as.data.frame(t(St))
  names(St) <- paste("age", seq(1,n.age.cats), sep = ".")
  St$day <- 1:n.days

  St.long <- melt(St, id.vars = c("day")) %>%
    rename(age = variable, n = value)
  St.long$age <- as.numeric(St.long$age)
  St.long <- St.long[,c("age", "day", "n")]

  Rt <- as.data.frame(t(Rt))
  names(Rt) <- paste("age", seq(1,n.age.cats), sep = ".")
  Rt$day <- 1:n.days

  Rt.long <- melt(Rt, id.vars = c("day")) %>%
    rename(age = variable, n = value)
  Rt.long$age <- as.numeric(Rt.long$age)
  Rt.long <- Rt.long[,c("age", "day", "n")]

  Dt <- as.data.frame(t(Dt))
  names(Dt) <- paste("age", seq(1,n.age.cats), sep = ".")
  Dt$day <- 1:n.days

  Dt.long <- melt(Dt, id.vars = c("day")) %>%
    rename(age = variable, n = value)
  Dt.long$age <- as.numeric(Dt.long$age)
  Dt.long <- Dt.long[,c("age", "day", "n")]

  Ht <- as.data.frame(t(Ht))
  names(Ht) <- paste("age", seq(1,n.age.cats), sep = ".")
  Ht$day <- 1:n.days

  Ht.long <- melt(Ht, id.vars = c("day")) %>%
    rename(age = variable, n = value)
  Ht.long$age <- as.numeric(Ht.long$age)
  Ht.long <- Ht.long[,c("age", "day", "n")]

    # convert array to matrix format
  Et.long <- melt(Et) %>%
    rename(age = Var1, day = Var2, subcat = Var3, n = value)
  It.long <- melt(It) %>%
    rename(age = Var1, day = Var2, subcat = Var3, n = value)

  counts <- list(St = St.long, Et = Et.long, It = It.long, Rt = Rt.long)
  outcomes <- list(Dt = Dt.long, Ht = Ht.long)
  output <- list(counts = counts, outcomes = outcomes, R0 = R0)
}


## Created on Sunday, August 29, 2021 at 11:17pm EDT by Weekend Editor on WeekendEditorMachine.
## Copyright (c) 2021, SomeWeekendReading.blog.  All rights reserved.  As if you care.

library("gsDesign")                                    # For ciBinomial()
library("plyr")                                        # For ddply()

toolsDir <- "../../tools"                              # Various tools available from author
source(file.path(toolsDir, "graphics-tools.r"))        # Various graphics hacks
source(file.path(toolsDir, "pipeline-tools.r"))        # Various graphics hacks

##
## Take summary data on Israeli population (vax, unvax, part vax) stratified by age.
## Compute efficacy and 95% confidence interval by posterior Beta function + Monte Carlo method.
##

doit <- function(## Inputs
                 dataDir  = ".",
                 dataFile = "2021-08-29-covid-simpson-Israeli_data_August_15_2021-vax-summary.tsv",

                 ## Outputs
                 resultsDir = ".",
                 txFile     = "2021-08-29-covid-simpson-ve-confidence-intervals-by-age.txt",
                 plotFile   = "2021-08-29-covid-simpson-ve-confidence-intervals-by-age.png") {

  loadData <- function(dataDir, dataFile) {            # Load the Israeli summary data of Morris
    f <- file.path(dataDir, dataFile)                  # Full pathname to data
    cat(sprintf("* Loading data from %s", f))          # Announce what we're doing
    vaxData <- read.table(f, header = TRUE, sep = "\t")# Read the data
    cat(sprintf("\n  - Found this dataframe:\n"))      # Capture what we got to the transcript
    print(vaxData)                                     # It's short enough that this is doable
    cat("\n")                                          # Final newline
    vaxData                                            # Return the dataframe
  }                                                    #

  addEfficacyAndCI <- function(vaxData) {              #

#    efficacyCI <- function(Ntrt, NtrtInf, Ncnt, NcntInf,
#                                  n = 1000000, probs = c(0.025, 0.500, 0.975)) {
#      ## *** This gets negative CI's when case counts are 0
#      ## Median as best guess, 2.5% and 97.5% as confidence limits
#      alphaNum   <- NtrtInf + 1                        # Numerator Beta distribution, alpha
#      betaNum    <- Ntrt - NtrtInf + 1                 #  and beta parameters (treatment)
#      alphaDenom <- NcntInf + 1                        # Denominator Beta distribution, alpha
#      betaDenom  <- Ncnt - NcntInf + 1                 #  and beta parameters (control)
#      effs       <- 100.0 * (1 - rbeta(n, alphaNum, betaNum) / rbeta(n, alphaDenom, betaDenom))
#      round(quantile(effs, probs = probs), digits = 1) # Get just the quantiles requested
#    }                                                  #

    efficacyCI <- function(Ntrt, NtrtInf, Ncnt, NcntInf, probs = c(0.025, 0.500, 0.975)) {
      ## Frequentist method: risk rates are scaled binomial, so we use the binomial
      ## confidence interval function in the gsDesign package.
      ci <- 100.0 * (1 - rev(ciBinomial(NtrtInf, NcntInf, Ntrt, Ncnt, scale = "RR")))
      c("2.5%"  = ci[[1]],                             #
        "ve"    = 100.0 * (1 - (NtrtInf / Ntrt) / (NcntInf / Ncnt)),
        "97.5%" = ci[[2]])                             #
    }                                                  #

    cat(sprintf("* Adding vaccine efficacies and 95%% confidence limits"))

    vaxDataEffs <- merge(vaxData,                      #
                         ddply(vaxData, "Age.Cohort", function(acdf) {
                           stopifnot(nrow(acdf) == 1)  # Defensive programming: unique age cohorts?
                           ves  <- efficacyCI(acdf[1, "Vax.Popln"],
                                              acdf[1, "Vax.Severe"],
                                              acdf[1, "Unvax.Popln"],
                                              acdf[1, "Unvax.Severe"])
                           pves <- efficacyCI(acdf[1, "Part.Vax.Popln"],
                                              acdf[1, "Part.Vax.Severe"],
                                              acdf[1, "Unvax.Popln"],
                                              acdf[1, "Unvax.Severe"])
                           data.frame(## Efficacy for the fully vaccinated
                                      VE      = ves[["ve"]],
                                      VE.LCL  = ves[["2.5%"]],
                                      VE.UCL  = ves[["97.5%"]],
                                      ## Efficacy for the partly vaccinated
                                      PVE     = pves[["ve"]],
                                      PVE.LCL = pves[["2.5%"]],
                                      PVE.UCL = pves[["97.5%"]])
                         }),                           #
                         by = "Age.Cohort")            #

    cat("\n")                                          # Final newline
    vaxDataEffs                                        # Return augmented dataframe
  }                                                    #

  plotVEs <- function(vaxDataEffs, resultsDir, plotFile) {
    f <- file.path(resultsDir, plotFile)               #
    withPNG(f, 400, 400, FALSE, function() {           #
      withPars(function() {                            #
        cols <- c("blue", "green")
        matplot(x = 1 : nrow(vaxDataEffs),
                y = subset(vaxDataEffs, select = c(VE)), #, PVE)),
                type = "b", pch = 21, col = "black", bg = cols,
                xaxt = "n", xlab = NA, ylab = "Vaccine Efficacy (%)", ylim = c(0, 100),
                main = "Vaccine efficacy by age cohort")
        sapply(1 : nrow(vaxDataEffs), function(rn) {
          drawErrorBar(x0 = rn,
                       y0 = vaxDataEffs[[rn, "VE.LCL"]],
                       x1 = rn,
                       y1 = vaxDataEffs[[rn, "VE.UCL"]],
                       col = "blue")
#          drawErrorBar(x0 = rn,
#                       y0 = vaxDataEffs[[rn, "PVE.LCL"]],
#                       x1 = rn,
#                       y1 = vaxDataEffs[[rn, "PVE.UCL"]],
#                       col = "green")
        })
#        legend("bottomright", inset = 0.05, bg = "antiquewhite",
#               pch = 21, pt.bg = cols, legend = c("Fully vaccinated", "Partially vaccinated"))
        axis(side = 1, at = 1 : nrow(vaxDataEffs), labels = as.character(vaxDataEffs$"Age.Cohort"))
        title(xlab = "Age Cohort", line = 3.5)         # Horizontal axis title
      }, pty = "m",                                    # Maximal plotting area
         bg  = "white",                                # White background
         mar = c(5, 3, 2, 1),                          # Pull in on margins a bit
         mgp = c(1.7, 0.5, 0),                         # Axis title, label, tick
         las = 2,                                      # Axis labels perp to axis
         ps  = 16)                                     # Larger type size for file capture
    })                                                 #
    cat(sprintf("* Plotted to %s\n", f))               # Capture to transcript
    TRUE                                               # Flag that it's done
  }                                                    #

  withTranscript(dataDir, resultsDir, txFile, "Israeli COVID-19 VE Stratified by Age", function() {

    heraldPhase("Loading Data")                        #
    maybeAssign("vaxData", function() { loadData(dataDir, dataFile) })

    heraldPhase("Computing efficacies and confidence interval")
    maybeAssign("vaxDataEffs", function() { addEfficacyAndCI(vaxData) })

    cat("\n"); print(vaxDataEffs) # ***

    heraldPhase("Making plot")                         #
    maybeAssign("plotDone", function() { plotVEs(vaxDataEffs, resultsDir, plotFile) })

    invisible(NA)                                      # Return nothing of interest
  })                                                   # Done with transcript capture
}                                                      # Done overall


###################################################################
#
# R source code providing generalized bilinear mixed effects modeling
# for social network data. For more information on the model fitting
# procedure, see
# "Bilinear Mixed-Effects Models for Dyadic Data" (Hoff 2003)
# http://www.stat.washington.edu/www/research/reports/2003/tr430.pdf
#
# This research was supported by ONR grant N00014-02-1-1011
#
# Copyright (C) 2003 Peter D. Hoff
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place - Suite 330,
# Boston, MA  02111-1307, USA.
#
##################################################################

###Last modification : 2/22/04

###Preconditions of the main function "gbme"

##objects with no defaults
#Y : a square matrix (n*n)

##objects with defaults
#Xd : an n*n*rd array representing dyad-specific predictors
#Xs : an n*rs matrix of sender-specific predictors
#Xr : an n*rr matrix of receiver-specific predictors

# fam : "gaussian" "binomial", or "poisson"
# if fam="binomial" then an (n*n) matrix N of trials is needed
#
# pi.Sab = (s2a0,sab0,s2b0,v0) : parameters of inverse wishart dist
# pi.s2u = vector of length 2: parameters of an inv-gamma distribution
# pi.s2v = vector of length 2: parameters of an inv-gamma distribution
# pi.s2z = k*2 matrix : parameters of k inv-gamma distributions
#
# pim.bd  : vector of length rd, the prior mean for beta.d
# piS.bd  : a rd*rd pos def matrix, the prior covariance for beta.d
#
# pim.b0sr : vector of length 1+rs+rr, prior mean for (beta0,beta.s,beta.r)
# piS.b0sr  : a (1+rs+rr)*(1+rs+rr) pos def matrix,
#             the prior variance for (beta0,beta.s,beta.r)

# k     : dimension of latent space
# directed : T or F. Is the data directed or undirected? If set to F,
#                    estimation proceeds assuming Y[i,j]==Y[j,i],
#                    and any supplied value of Xr is not used
#                    also, only s2a0 and v0 in pi.Sab need be supplied,
#                    and pi.s2v is not used. rr is zero, so pim.b0sr is
#                    length 1+rs

# starting values
# beta.d   #a vector of length rd
# beta.u   #a vector of length 1+rs+rr
# s        #a vector of length n
# r        #a vector of length n
# z        #an n*k matrix

###the main  function
gbme<-function(
 #data
 Y,
 Xd=array(dim=c(n,n,0)),
 Xs=matrix(nrow=dim(Y)[1],ncol=0),
 Xr=matrix(nrow=dim(Y)[1],ncol=0),

 #model specification
 fam="gaussian",
 N=NULL,
 k=0,
 directed=T,

#compute starting values and emp Bayes for priors?
 estart=T,

#priors  - provide if estart=F
 pi.Sab=NULL,
 pi.s2u=NULL,
 pi.s2v=NULL,
 pi.s2z=NULL,
 pim.bd=NULL,
 piS.bd=NULL,
 pim.b0sr=NULL,
 piS.b0sr=NULL,

 #starting values - provide if estart=F
 beta.d=NULL,
 beta.u=NULL,
 s=NULL,
 r=NULL,
 z=NULL,

 #random seed,
 seed=0,

 #details of output
 NS=100000,         #number of scans of mcmc to run
 odens=100,         #output density
 owrite=T,          #write output to a file?
 ofilename="OUT",   #name of output file
 oround=3,          #rounding of output
 zwrite=(k>0),      #write z to  file
 zfilename="Z",     #outfile for z
 sdigz=3,           #rounding for z
 awrite=F,bwrite=F,
 afilename="A",
 bfilename="B"
	       )
{

###set seed
set.seed(seed)

###dimensions of everything
n<<-dim(Y)[1]
rd<<-dim(Xd)[3]
rs<<-dim(Xs)[2]
rr<<-dim(Xr)[2]
diag(Y)<<-rep(0,n)
###

###column names for output file
if(directed==T) {
cnames<-c("k","scan","ll",paste("bd",seq(1,rd,length=rd),sep="")[0:rd],"b0",
                          paste("bs",seq(1,rs,length=rs),sep="")[0:rs],
                          paste("br",seq(1,rr,length=rr),sep="")[0:rr],
  "s2a","sab","s2b","s2e","rho", paste("s2z",seq(1,k,length=k),sep="")[0:k] )
                }

if(directed==F) {
cnames<-c("k","scan","ll",paste("bd",seq(1,rd,length=rd),sep="")[0:rd],"b0",
                          paste("bs",seq(1,rs,length=rs),sep="")[0:rs],
                "s2a","s2e",paste("s2z",seq(1,k,length=k),sep="")[0:k] )
                }

###design matrix for unit-specific predictors
if(directed==T) {
X.u<<-rbind( cbind( rep(.5,n), Xs, matrix(0,nrow=n,ncol=rr)),
            cbind( rep(.5,n),  matrix(0,nrow=n,ncol=rs), Xr ) )
                 }

if(directed==F) { X.u<<-cbind( rep(.5,n), Xs )  }
###

###construct an upper triangular matrix (useful for later computations)
tmp<-matrix(1,n,n)
tmp[lower.tri(tmp,diag=T)]<-NA
UT<<-tmp
###

###get starting values and empirical bayes priors using a glm-type approach
if(estart==T) {tmp<-gbme.glmstart(Y,Xd,Xs,Xr,N=N,fam,k,directed)
               priors<-tmp$priors ; startv<-tmp$startv
              }

if(directed==T) {
 pi.Sab<-priors$pi.Sab; pi.s2u<-priors$pi.s2u; pi.s2v<-priors$pi.s2v
 pi.s2z<-priors$pi.s2z ; pim.bd<-priors$pim.bd ; piS.bd<-priors$piS.bd
 pim.b0sr<-priors$pim.b0sr ; piS.b0sr<-priors$piS.b0sr
 beta.d<-startv$beta.d ; beta.u<-startv$beta.u
 s<-startv$s ; r<-startv$r ; z<-startv$z
               }
###

if(directed == F) {  rho<-1  ;sv<-0 ;  Xr<-Xs ; rr<-rs
                  pi.s2u<-priors$pi.s2u ; pi.s2z<-priors$pi.s2z
                  pim.bd<-priors$pim.bd ; piS.bd<-priors$piS.bd
                  beta.d<-startv$beta.d ; beta.u<-startv$beta.u ; z<-startv$z
                  s<-startv$s ;
                  pi.s2a<-priors$pi.s2a ;
                  pim.b0s<-priors$pim.b0s ;  piS.b0s<-priors$piS.b0s
                  Xr<-NULL ; rr<-0 ; r<-s*0   }

###Matrices for ANOVA  on Xd, s, r
tmp<-TuTv(n)   #unit level effects design matrix for u=yij+yji, v=yij-yji
Tu<-tmp$Tu
Tv<-tmp$Tv

if(directed==F) { Tu<-2*Tu[,1:n] }

tmp<-XuXv(Xd)   #regression design matrix for u=yij+yji, v=yij-yji
Xu<-tmp$Xu
Xv<-tmp$Xv

XTu<<-cbind(Xu,Tu)
XTv<<-cbind(Xv,Tv)

tXTuXTu<<-t(XTu)%*%XTu
tXTvXTv<<-t(XTv)%*%XTv
###

###redefining hyperparams
if(directed==T){
Sab0<-matrix( c(pi.Sab[1],pi.Sab[2],pi.Sab[2],pi.Sab[3]),nrow=2,ncol=2)
v0<-pi.Sab[4]   }
###

###initializing error matrix
E<-matrix(0,nrow=n,ncol=n)
###

###if using the linear link, then theta=Y
if(fam=="gaussian"){ theta<-Y }
###

##### The Markov chain Monte Carlo algorithm
for(ns in 1:NS){

###update value of theta
if(fam!="gaussian"){theta<-theta.betaX.d.srE.z(beta.d,Xd,s,s*(1-directed)+r*directed,E,z)  }

###impute any missing values if gaussian
if(fam=="gaussian" & any(is.na(Y))) {
  rho<-(pi.s2u[2]-pi.s2v[2])/(pi.s2u[2]+pi.s2v[2])
  se<-(pi.s2u[2]+pi.s2v[2])/4
  mu<-theta.betaX.d.srE.z(beta.d,Xd,s,s*(1-directed)+r*directed,E*0,z)
  imiss<-(1:n)[ is.na(Y%*%rep(1,n))  ]
  for(i in imiss){
    for(j in (1:n)[is.na(Y[i,])]) {
    if( is.na(Y[i,j]) & is.na(Y[j,i]) )       {
      smp<-rmvnorm(c(mu[i,j],mu[j,i]),matrix( se*c(1,rho,rho,1),nrow=2,ncol=2))
      theta[i,j]<-smp[1] ; theta[j,i]<-smp[2] }
else{theta[i,j]<-rmvnorm( mu[i,j]+rho*(theta[j,i]-mu[j,i]), sqrt(se*(1-rho^2)))}
                                   }
                  }
                                    }

###Update regression part
tmp<-uv.E(theta-z%*%t(z)) #the regression part
u<-tmp$u                  #u=yij+yji,  i<j
v<-tmp$v                  #v=yij-yji,  i<j

#First get variance - covariance terms
if(directed==T) {
su<-rse.beta.d.gibbs(pi.s2u[1],pi.s2u[2],u,XTu,s,r,beta.d)  #for rho, se
sv<-rse.beta.d.gibbs(pi.s2v[1],pi.s2v[2],v,XTv,s,r,beta.d)
rho<-(su-sv)/(su+sv)
se<-(su+sv)/4
                 }
if(directed==F) {
su<-rse.beta.d.gibbs(pi.s2u[1],pi.s2u[2],u,XTu,s,NULL,beta.d) #for se
se<-su/4       }

sr.hat<-X.u%*%beta.u                                        #Sab
a<-s-sr.hat[1:n]
b<-r-sr.hat[n+(1:n)]
if( directed==T ) {Sab<-rSab.gibbs(a,b,Sab0,v0)}
if( directed==F ) { s2a<-1/rgamma( 1, pi.s2a[1]+n/2 , pi.s2a[2]+sum(a^2)/2 ) }

#dyad specific regression coef and unit level effects
mu<-c(pim.bd, X.u%*%beta.u)    #"prior" mean for (beta.d,s,r)
if(directed==T){
beta.d.sr<-rbeta.d.sr.gibbs(u,v,su,sv,piS.bd,mu,Sab)
beta.d<-beta.d.sr$beta.d ; s<-beta.d.sr$s ; r<-beta.d.sr$r
#regression coef for unit level effects
beta.u<-rbeta.sr.gibbs(s,r,X.u,pim.b0sr,piS.b0sr,Sab)
                }
if(directed==F) {
beta.d.s<-rbeta.d.s.gibbs(u,su,piS.bd,mu,s2a)
beta.d<-beta.d.s$beta.d ; s<-beta.d.s$s
#regression coef for unit level effects
beta.u<-rbeta.s.gibbs(s,X.u,pim.b0s,piS.b0s,s2a)
                 }
###

###bilinear effects
if(k>0){

#update variance
#tmp<-eigen(z%*%t(z))
#z<-tmp$vec[,1:k]%*%sqrt(diag(tmp$val[1:k],nrow=k)) #make cols of z orthogonal
sz<-1/rgamma(k,pi.s2z[,1]+n/2,pi.s2z[,2]+diag( t(z)%*%z)/2)

#Gibbs for zs, using regression
res<-theta-theta.betaX.d.srE.z(beta.d,Xd,s,s*(1-directed)+r*directed,0*E,0*z) #theta-(beta*x+a+b)
s.ares<-se*(1+rho)/2
for(i in sample(1:n))      {
ares<-(res[i,-i]+res[-i,i])/2
Zi<-z[-i,]
Szi<-chol2inv(chol(diag(1/(sz),nrow=k) +
     t(Zi)%*%Zi/s.ares)) #cond variance of z[i,]
muzi<-Szi%*%( t(Zi)%*%ares/s.ares )                      #cond mean of z[i,]
z[i,]<-t(rmvnorm(muzi,Szi)) }
         }
###
if(k==0){sz<-NULL }

#update E
if( fam!="gaussian"){
E<-theta-theta.betaX.d.srE.z(beta.d,Xd,s,s*(1-directed)+r*directed,E*0,z)
      #current E|beta.d,Xd,s,r,z

UU<-VV<-Y*0

u<-rnorm(n*(n-1)/2,0,sqrt(su))
v<-rnorm(n*(n-1)/2,0,sqrt(sv))

UU[!is.na(UT)]<-u
tmp<-t(UU)
tmp[!is.na(UT)]<-u
UU<-t(tmp)

VV[!is.na(UT)]<-v
tmp<-t(VV)
tmp[!is.na(UT)]<--v
VV<-t(tmp)

Ep<-(UU+VV)/2
theta.p<-theta-E+Ep

diag(theta.p)<-diag(theta)<-NA

##cases
if(fam=="poisson") { mup<-exp(theta.p) ;  mu<-exp(theta)
                     M.lhr<-dpois(Y,mup,log=T) - dpois(Y,mu,log=T)   }

if(fam=="binomial"){ mup<-exp(theta.p)/(1+exp(theta.p))
                          mu<-exp(theta)/(1+exp(theta))
         M.lhr<-dbinom(Y,N,mup,log=T) - dbinom(Y,N,mu,log=T)   }

   M.lhr[ is.na(M.lhr) ]<-0
   M.lhr<-(M.lhr+t(M.lhr))/( 2^(1-directed) )
   RU<-matrix(log(runif(n*n)),nrow=n,ncol=n)
   RU[is.na(UT)]<-0  ; RU<-RU+t(RU)
   E[M.lhr>RU]<-Ep[M.lhr>RU]
                 }

####output
if(ns%%odens==0){

#compute loglikelihoods of interest
if(fam=="poisson") { lpy.th<-sum( dpois(Y,exp(theta),log=T),na.rm=T) }
if(fam=="binomial"){ lpy.th<-sum( dbinom(Y,N,exp(theta)/(1+exp(theta)),
                                    log=T),na.rm=T) }

if(fam=="gaussian"){
E<-theta-theta.betaX.d.srE.z(beta.d,Xd,s,s*(1-directed)+r*directed,E*0,z)
tmp<-uv.E(theta-z%*%t(z)) #the regression part
u<-tmp$u                  #u=yij+yji,  i<j
v<-tmp$v                  #v=yij-yji,  i<j
if(directed==T) {
lpy.th<-sum(dnorm(u,0,sqrt(su),log=T) + dnorm(v,0,sqrt(sv),log=T)) }
if(directed==F) { lpy.th<-2*sum(dnorm(u,0,sqrt(su),log=T))  }
                    }
if(directed==F) {lpy.th<-lpy.th/2
                 out<-round(c(k,ns,lpy.th,beta.d,beta.u,s2a,se,sz[0:k]),oround) }
if(directed==T) {
out<-round(c(k,ns,lpy.th,beta.d,beta.u,Sab[1,1],Sab[1,2],Sab[2,2],
             se,rho,sz[0:k]),oround) }

if(owrite==T) {
    if(ns==odens) { write.table(t(out),file=ofilename,quote=F,
                     row.names=F,col.names=cnames) }
    if(ns>odens)  { write.table(t(out),file=ofilename,append=T,quote=F,
                    row.names=F,col.names=F) }
              }
if(owrite==F) {cat(out,"\n") }
if(zwrite==T) { write.table(signif(z,sdigz),file=zfilename,
                      append=T*(ns>odens),quote=F,row.names=F,col.names=F) }

if(awrite==T){write.table(round(t(a),oround),file=afilename,append=T*(ns>odens),quote=F,
                    row.names=F,col.names=F)  }
if(bwrite==T & directed==T){write.table(round(t(b),oround),
                    file=bfilename,append=T*(ns>odens),quote=F,
                    row.names=F,col.names=F) }

                 }

}}

#####################End of main function : below are helper functions

####
TuTv<-function(n){
Xu<-Xv<-NULL
for(i in 1:(n-1)){
tmp<-tmp<-NULL
if( i >1 ){ for(j in 1:(i-1)){ tmp<-cbind(tmp,rep(0,n-i)) } }
tmp<-cbind(tmp,rep(1,n-i))
tmpu<-cbind(tmp,diag(1,n-i)) ; tmpv<-cbind(tmp,-diag(1,n-i))
Xu<-rbind(Xu,tmpu) ; Xv<-rbind(Xv,tmpv)
         }

list(Tu=cbind(Xu,Xu),Tv=cbind(Xv,-Xv))
                   }
####

XuXv<-function(X){
Xu<-Xv<-NULL
if(dim(X)[3]>0){
for(r in 1:dim(X)[3]){
xu<-xv<-NULL
for(i in 1:(n-1)){
for(j in (i+1):n){ xu<-c(xu,X[i,j,r]+X[j,i,r])
                   xv<-c(xv,X[i,j,r]-X[j,i,r]) }}
Xu<-cbind(Xu,xu)
Xv<-cbind(Xv,xv)  }
                }
list(Xu=Xu,Xv=Xv)}

###
uv.E<-function(E){
u<- c(  t( (  E + t(E) )  *UT ) )
u<-u[!is.na(u)]

v<-c(  t( (  E - t(E) )  *UT ) )
v<-v[!is.na(v)]
list(u=u,v=v)}
####

####
theta.betaX.d.srE.z<-function(beta.d,X.d,s,r,E,z){
m<-dim(X.d)[3]
mu<-matrix(0,nrow=length(s),ncol=length(s))
if(m>0){for(l in 1:m){ mu<-mu+beta.d[l]*X.d[,,l] }}
tmp<-mu+re(s,r,E,z)
diag(tmp)<-0
tmp}
####

####
re<-function(a,b,E,z){
n<-length(a)
matrix(a,nrow=n,ncol=n,byrow=F)+matrix(b,nrow=n,ncol=n,byrow=T)+E+z%*%t(z) }
####

####
rbeta.d.sr.gibbs<-function(u,v,su,sv,piS.bd,mu,Sab){
del<-Sab[1,1]*Sab[2,2]-Sab[1,2]^2
iSab<-rbind(cbind( diag(rep(1,n))*Sab[2,2]/del ,-diag(rep(1,n))*Sab[1,2]/del),
          cbind( -diag(rep(1,n))*Sab[1,2]/del,diag(rep(1,n))*Sab[1,1]/del) )
rd<-dim(as.matrix(piS.bd))[1]

if(dim(piS.bd)[1]>0){
cov.beta.sr<-matrix(0,nrow=rd,ncol=2*n)
iS<-rbind(cbind(solve(piS.bd),cov.beta.sr),
   cbind(t(cov.beta.sr),iSab)) }
else{iS<-iSab}

Sig<-chol2inv( chol(iS + tXTuXTu/su + tXTvXTv/sv ) )
    #this may have a closed form expression

#theta<-Sig%*%( (t(XTu)%*%u)/su + (t(XTv)%*%v)/sv + iS%*%mu)
theta<-Sig%*%( t(  (u%*%XTu)/su + (v%*%XTv)/sv ) + iS%*%mu)

beta.sr<-rmvnorm(theta,Sig)
list(beta.d=beta.sr[(rd>0):rd],s=beta.sr[rd+1:n],r=beta.sr[rd+n+1:n]) }
####

####
rse.beta.d.gibbs<-function(g0,g1,x,XTx,s,r,beta.d){
n<-length(s)
1/rgamma(1, g0+choose(n,2)/2,g1+.5*sum( (x-XTx%*%c(beta.d,s,r))^2 ) ) }
####

####
rbeta.sr.gibbs<-function(s,r,X.u,pim.b0sr,piS.b0sr,Sab) {
del<-Sab[1,1]*Sab[2,2]-Sab[1,2]^2
iSab<-rbind(cbind( diag(rep(1,n))*Sab[2,2]/del ,-diag(rep(1,n))*Sab[1,2]/del),
          cbind( -diag(rep(1,n))*Sab[1,2]/del,diag(rep(1,n))*Sab[1,1]/del) )

S<-solve( solve(piS.b0sr) + t(X.u)%*%iSab%*%X.u )
mu<-S%*%(  solve(piS.b0sr)%*%pim.b0sr+ t(X.u)%*%iSab%*%c(s,r))
rmvnorm( mu,S)
                                               }
####

####
rSab.gibbs<-function(a,b,S0,v0){
n<-length(a)
ab<-cbind(a,b)
Sn<-S0+ (t(ab)%*%ab)
solve(rwish(solve(Sn),v0+n) )
                               }

####
rmvnorm<-function(mu,Sig2){
R<-t(chol(Sig2))
R%*%(rnorm(length(mu),0,1)) +mu }

####
rwish<-function(S0,nu){
S<-S0*0
for(i in 1:nu){ z<-rmvnorm(rep(0,dim(as.matrix(S0))[1]), S0)
                S<-S+z%*%t(z)  }
		S }

###  Procrustes transformation: rotation and reflection
proc.rr<-function(Y,X){
k<-dim(X)[2]
A<-t(Y)%*%(  X%*%t(X)  )%*%Y
eA<-eigen(A,symmetric=T)
Ahalf<-eA$vec[,1:k]%*%diag(sqrt(eA$val[1:k]),nrow=k)%*%t(eA$vec[,1:k])

t(t(X)%*%Y%*%solve(Ahalf)%*%t(Y)) }
###

gbme.glmstart<-function(Y,Xd,Xs,Xr,N=NULL,fam,k,directed){
#generate starting values and emp bayes priors from
#approximate mles

if(fam!="binomial"){
y<-x<-NULL
for (i in 1:n) {
for (j in (1:n)[-i]) {
y<-c(y,Y[i,j])
x<-rbind(x,c(i,j,Xd[i,j,])) }}
                    }

if(fam=="binomial"){
y<-x<-NULL
for (i in 1:n) {
for (j in (1:n)[-i]) {
y<-rbind(y,c(Y[i,j], N[i,j]-Y[i,j] ))
x<-rbind(x,c(i,j,Xd[i,j,])) }}
                    }

if(rd>0) {
 fit1<-glm(y~-1+as.factor(x[,1])+as.factor(x[,2])+x[,-(1:2)],family=fam) }
if(rd==0){ fit1<-glm(y~-1+as.factor(x[,1])+as.factor(x[,2]),family=fam) }

s<-fit1$coef[1:n]
r<-c(0,fit1$coef[n+1:(n-1)])
if(rd>0){
beta.d<-fit1$coef[(2*n):length(fit1$coef)]               #beta.d
        }

la<-lm( s~-1+cbind(rep(.5,n),Xs)); a<-la$res
lb<-lm( r~-1+cbind(rep(.5,n),Xr)); b<-lb$res

b0<-(la$coef[1]+lb$coef[1])/2

s1<-s+( b0-la$coef[1])/2                 #s
r1<-r+( b0-lb$coef[1])/2                 #r

la<-lm( s1~-1+cbind(rep(.5,n),Xs)); a<-la$res
lb<-lm( r1~-1+cbind(rep(.5,n),Xr)); b<-lb$res

beta.u<-c(lb$coef[1], la$coef[-1],lb$coef[-1])
Sab<-var(cbind(a,b))                    #Sab

Res<-matrix(0,nrow=n,ncol=n)
l<-1
for( m in 1:dim(x)[1] ) {
if( !is.na(y[m]) ) { Res[x[m,1],x[m,2]]<-fit1$res[l] ; l<-l+1 }
                         }
######
Res[Res>quantile(Res,.95,na.rm=T) ] <- quantile(Res,.95,na.rm=T)
diag(Res)<-rep(0,n)

if(k>0){
for(i in 1:50){
tmp<-eigen( .5*(Res+t(Res) ))
z<-tmp$vec[,1:k] %*%diag(sqrt(tmp$val[1:k]),nrow=k)
diag(Res)<-diag(z%*%t(z))
                 }
sz<-var(c(z))                         #z,sz
E<-Res-z%*%t(z)                           #E
          }
if(k==0){E<-Res ; z<-matrix(0,nrow=n,ncol=1) ;sz<-.1}

u<-E+t(E)
u<-c(u[!is.na(UT)])                      #su
su<-var(u)                               #sv
v<-E-t(E)
v<-c(v[!is.na(UT)])
sv<-var(v)

rho<-(su-sv)/(su+sv)
se<-(su+sv)/4

###construct priors centered on these empirical values
pi.Sab<-c(Sab[1,1],Sab[2,1],Sab[2,2],4)
pi.s2u<-c(2,su)
pi.s2v<-c(2,sv)
pi.s2z<-cbind( rep(2,k), diag( t(z)%*%z)/n )

if(rd>0){
pim.bd<-c(beta.d)
piS.bd<-as.matrix((n*(n-1))^2*summary(fit1)$cov.unscaled[(2*n):length(fit1$coef),
                                          (2*n):length(fit1$coef)] )
         }

if(rd==0){ beta.d<-pim.bd<-NULL ;  piS.bd<-matrix(nrow=0,ncol=0) }

if(directed==T) {
pim.b0sr<-beta.u
piS.b0sr<-rbind(cbind( summary(la)$cov[-1,-1], matrix(0,nrow=rs,ncol=rr)),
# cbind(matrix(0,nrow=rs,ncol=rr),summary(lb)$cov[-1,-1]))
# error found 7-7-6
  cbind(matrix(0,nrow=rr,ncol=rs),summary(lb)$cov[-1,-1]))

piS.b0sr<-rbind(c(summary(la)$cov[1,-1],summary(lb)$cov[1,-1]),piS.b0sr)
piS.b0sr<-cbind(c(summary(la)$cov[1,1:(rs+1)],
                  summary(lb)$cov[1,-1]),piS.b0sr)
piS.b0sr<-as.matrix(piS.b0sr*n*n)

pim.b0s<-NULL ; piS.b0s<-NULL  ;pi.s2a<-NULL
                }
if(directed==F) { pim.b0sr<-NULL ; piS.b0sr<-NULL ; pi.Sab<-NULL
                  s<-(s+r)/2  ; r<-s*0
                  tmp2<-lm(s~-1+cbind( rep(.5,n),Xs))
                  beta.u<-tmp2$coef ; pi.s2a<-c(2,var(tmp2$res)) ;
                  pim.b0s<-tmp2$coef ; piS.b0s<-summary(tmp2)$cov*n*n }

###
list( priors=list(pi.Sab=pi.Sab,pi.s2u=pi.s2u,pi.s2v=pi.s2v,pi.s2z=pi.s2z,
                  pim.bd=pim.bd,piS.bd=piS.bd,
                  pim.b0sr=pim.b0sr,piS.b0sr=piS.b0sr,
                  pim.b0s=pim.b0s, piS.b0s=piS.b0s , pi.s2a=pi.s2a) ,
      startv=list(beta.d=beta.d,beta.u=beta.u,s=s,r=r,z=z)
     )
}

####
rbeta.d.s.gibbs<-function(u,su,piS.bd,mu,s2a){
iSa<-diag(rep(1/s2a,n))

if(dim(piS.bd)[1]>0){
cov.beta.s<-matrix(0,nrow=rd,ncol=n)
iS<-rbind(cbind(solve(piS.bd),cov.beta.s),
   cbind(t(cov.beta.s),iSa)) }
else{iS<-iSa}

Sig<-chol2inv( chol(iS + tXTuXTu/su ) )
    #this may have a closed form expression

theta<-Sig%*%( t(  (u%*%XTu)/su) + iS%*%mu)

beta.s<-rmvnorm(theta,Sig)
list(beta.d=beta.s[(rd>0):rd],s=beta.s[rd+1:n]) }
####

####
rbeta.s.gibbs<-function(s,X.u,pim.b0s,piS.b0s,s2a) {
iSa<-diag(rep(1/s2a,n))

S<-solve( solve(piS.b0s) + t(X.u)%*%iSa%*%X.u )
mu<-S%*%(  solve(piS.b0s)%*%pim.b0s+ t(X.u)%*%iSa%*%s)
rmvnorm( mu,S)
                                               }


model_snowwet <- function (Swet_t1 = 0.0,
         precip = 0.0,
         Snowaccu = 0.0,
         Mrf = 0.0,
         M = 0.0,
         Sdry = 0.0){
    #'- Name: SnowWet -Version: 1.0, -Time step: 1
    #'- Description:
    #'            * Title: water in liquid state in the snow cover calculation
    #'            * Author: STICS
    #'            * Reference: doi:http://dx.doi.org/10.1016/j.agrformet.2014.05.002
    #'            * Institution: INRA
    #'            * Abstract: water in liquid state in the snow cover
    #'- inputs:
    #'            * name: Swet_t1
    #'                          ** description : water in liquid state in the snow cover in previous day
    #'                          ** inputtype : variable
    #'                          ** variablecategory : state
    #'                          ** datatype : DOUBLE
    #'                          ** default : 0.0
    #'                          ** min : 0.0
    #'                          ** max : 100.0
    #'                          ** unit : mmW
    #'                          ** uri :
    #'            * name: precip
    #'                          ** description : current precipitation
    #'                          ** inputtype : variable
    #'                          ** variablecategory : auxiliary
    #'                          ** datatype : DOUBLE
    #'                          ** default : 0.0
    #'                          ** min : 0.0
    #'                          ** max : 5000.0
    #'                          ** unit : mmW
    #'                          ** uri :
    #'            * name: Snowaccu
    #'                          ** description :  snowfall accumulation
    #'                          ** inputtype : variable
    #'                          ** variablecategory : rate
    #'                          ** datatype : DOUBLE
    #'                          ** default : 0.0
    #'                          ** min :
    #'                          ** max :
    #'                          ** unit : mmW/d
    #'                          ** uri :
    #'            * name: Mrf
    #'                          ** description : liquid water in the snow cover in the process of refreezing
    #'                          ** inputtype : variable
    #'                          ** variablecategory : rate
    #'                          ** datatype : DOUBLE
    #'                          ** default : 0.0
    #'                          ** min :
    #'                          ** max :
    #'                          ** unit : mmW/d
    #'                          ** uri :
    #'            * name: M
    #'                          ** description : snow in the process of melting
    #'                          ** inputtype : variable
    #'                          ** variablecategory : rate
    #'                          ** datatype : DOUBLE
    #'                          ** default : 0.0
    #'                          ** min :
    #'                          ** max :
    #'                          ** unit : mmW/d
    #'                          ** uri :
    #'            * name: Sdry
    #'                          ** description : water in solid state in the snow cover
    #'                          ** inputtype : variable
    #'                          ** variablecategory : state
    #'                          ** datatype : DOUBLE
    #'                          ** default : 0.0
    #'                          ** min : 0.0
    #'                          ** max : 500.0
    #'                          ** unit : mmW
    #'                          ** uri :
    #'- outputs:
    #'            * name: Swet
    #'                          ** description : water in liquid state in the snow cover
    #'                          ** variablecategory : state
    #'                          ** datatype : DOUBLE
    #'                          ** min : 0.0
    #'                          ** max : 500.0
    #'                          ** unit : mmW
    #'                          ** uri :
    Swet <- 0.0
    if (Mrf <= Swet_t1)
    {
        tmp_swet <- Swet_t1 + precip - Snowaccu + M - Mrf
        frac_sdry <- 0.1 * Sdry
        if (tmp_swet < frac_sdry)
        {
            Swet <- tmp_swet
        }
        else
        {
            Swet <- frac_sdry
        }
    }
    return (list('Swet' = Swet))
}


#' @title interpolate.xy.robust
#' @description Interpolate some xy data used by netmensuration and mode fifnding methods.
#' @family abysmally documented
#' @author  Jae Choi, \email{jae.choi@dfo-mpo.gc.ca}
#' @export
interpolate.xy.robust = function( xy, method, target.r2=0.9, mv.win=10, trim=0.05,
  probs=c(0.025, 0.975), loess.spans=seq( 0.25, 0.01, by=-0.01 ), inla.model="rw2",
  smoothing.kernel=kernel( "modified.daniell", c(2,1)), nmax=3, inla.h=0.1, inla.diagonal=0.01, inla.ngroups=400 ) {

  # simple interpolation methods
  # target.r2 == target prediction R^2

  if (FALSE) {
    x=seq(-5,5,by=0.1)
    y=sin(x) + runif(length(x))^2
    xy = data.frame( x=x, y=y )
    target.r2=0.9
    nmax=5  # max number of times to try to reduce data set
    probs=c(0.025, 0.975)
    loess.spans=seq( 0.2, 0.01, by=-0.01 )
    inla.model="rw2"
    smoothing.kernel=kernel( "modified.daniell", c(2,1))
  }

  if ( is.vector(xy) ) {
    nz = length(xy)
    z = data.frame( x=1:nz, y=xy )
  } else {
    if (ncol(xy) == 2 ) {
      nz = nrow(xy)
      z = data.frame( xy)
      names(z) = c("x", "y" )
    } else {
      stop( "Error: y or c(x,y) expected")
    }
  }

  z$p= NA
  z$noise = FALSE
  z$loc = 1:nz  # row index used in seq lin method

  if (method == "moving.window" ) {
    z$p = z$y
    for (nw in mv.win:1) {
      win = -nw:nw
      z$p = NA
      for( i in 1:nz ) {
        iw = i+win
        iw = iw[ which( iw >= 1 & iw <=nz ) ]
        z$p[i] = mean( z$y[iw], na.rm=TRUE )
      }
      z$diff = abs( z$p - z$y )
      qnts = quantile( z$diff, probs=(1-trim), na.rm=TRUE )  # remove upper quantile
      ii = which( z$diff > qnts )
      if (length(ii) > 0) z$p[ ii ] = NA
      z$p[ii] = approx( x=z$x, y=z$p, xout=z$x[ii], method="linear", rule=2 )$y
      rsq = cor( z$p, z$y, use="pairwise.complete.obs" )^2
      if (!is.na(rsq)) {
        if (rsq > target.r2 ) break()
      }
    }
  }

  if (method =="sequential.linear") {
    # check adjacent values and reject if differences are greater than a given range of quantiles (95%)

      m = z # locs are the location indices for z$y ... to bring back into "z"
      m = m[ which(is.finite( m$y)) , ]
      lg = 1 # lag
      nw = nrow( m )  # staring number of data points
      nw_target =  nw * (1-trim)  # proportion of data to retain (not trim)

      while ( nw > nw_target ) {
        yd0 = c( diff( m$y, lag=lg ) , rep( 0, lg )) # diffs left-right # padding with 0's  (no change)
        dm = quantile( yd0, probs=probs, na.rm=TRUE )
        n = which( yd0 < dm[1] | yd0 > dm[2] )
        if (length(n) == 0 ) break()
        if ( length(n) > 0 ) m = m[-n, ]  # remove na's and then assess again
        nw = nrow( m )
      }

      # m$locs contains indices of "good" data .. setdiff to get the "bad" .. set to NA and then
      #   infill missing with interpolated values
      ii = setdiff( 1:nz, unique(m$loc) )
      if (length(ii) > 0)  {
        z$y[ii] = NA
        afun = approxfun( x=z$x, y=z$y,  method="linear", rule=2  )
        z$y[ii] = afun( z$x[ii] )
      }
    z$p = z$y
  }

  if (method =="local.variance") {
    # local variance estimates used to flag strange areas .. where local variance is too high .. remove the data

    z$p = z$y
    for (nw in mv.win:1) {
      z$Zvar = NA
      win = c( -nw : nw )
      for (i in 1:nz ) {
        iw = i+win
        iw = iw[ which(iw >=1 & iw<= nz)]
        if (length(iw)> 3)  z$Zvar[i] = var( z$p[iw], na.rm=TRUE )
      }
      qnts = quantile( z$Zvar, probs=(1-trim), na.rm=TRUE )  # remove upper quantile
      ii = which( z$Zvar > qnts )
      if (length(ii) > 0) z$p[ ii ] = NA
      afun = approxfun( x=z$x, y=z$p,  method="linear", rule=2  )
      z$p[ii] = afun( z$x[ii] )
      rsq = cor( z$p, z$y, use="pairwise.complete.obs" )^2
      if (!is.na(rsq)){
        if (rsq > target.r2 ) break()
      }
    }
  }

  if (method == "loess" ) {
    var0 = var( z$y, na.rm=TRUE)
    for( sp in loess.spans ) {
      # ID a good span where loess solution ~ real data

      loess.mod = try( loess( y ~ x, data=z, span=sp, control=loess.control(surface="direct"),na.action = na.exclude ), silent=TRUE )
      if ( "try-error" %in% class(loess.mod) ) next()
      z$p = predict( loess.mod, newdata=z )
      rsq = cor( z$p, z$y, use="pairwise.complete.obs" )^2
      if (!is.na(rsq)) {
        if (rsq > target.r2 ) {
         # if ( var( z$p, na.rm=TRUE) < var0 )
		break()
        }
      }
    }
  }

  if (method == "inla") {

    require(INLA)
    ymean = mean( z$y, na.rm=TRUE )
    z$y = z$y - ymean
    nn = abs( diff( z$x) )
    dd = median( nn[nn>0], na.rm=TRUE )
    z$xiid = z$x = jitter( z$x, amount=dd / 20 ) # add noise as inla seems unhappy with duplicates in x?

    rsq = 0
    nw = length( which(is.finite( z$y)))
    nw0 = nw + 1
    count = 0

    ndat = nrow( z )
    #if ( ndat < inla.ngroups ) inla.ngroups = ndat

    # m = get("inla.models", INLA:::inla.get.inlaEnv())
    # m$latent$rw2$min.diff = NULL
    # assign("inla.models", m, INLA:::inla.get.inlaEnv())

    # inla's default RW2 resolution is too coarse, causing incomplete solution:
    # alter a bit as this is potentially a smoothed series
    #menv = get("inla.models", INLA:::inla.get.inlaEnv())
    #menv$latent$rw2$min.diff =1e-4
    # assign("inla.models", menv, INLA:::inla.get.inlaEnv() )
    h = 0.02
    while ( nw != nw0 ) {
      count = count + 1
      if (count > nmax ) break() # this is CPU expensive ... try only a few times

      FM = formula(  y ~ f( inla.group(xiid, n=inla.ngroups), model="iid")
                       + f( inla.group(x, n=inla.ngroups), model=inla.model )
                       )
      v = NULL
      v = try( inla( FM, data=z,
            control.predictor=list (compute=TRUE ),
            control.compute=list(config=TRUE),
            control.results=list(return.marginals.random=TRUE, return.marginals.predictor=TRUE ),
            control.inla = list( h=h ), # increase in case values are too close to zero
            # control.mode = list( restart=TRUE, result=v ), # restart from previous estimates
            quantiles=NULL,  # no quants to save storage requirements
            verbose=FALSE )
        , silent=TRUE )

      if ( "try-error" %in% class(v) ) next()

      # make sure Eigenvalues of Hessian are appropriate (>0)
      if ( v$mode$mode.status > 0 ) {
        h = h + 0.02
        next()
      }

      z$p = v$summary.fitted.values$mean
      rsq = cor( z$p, z$y, use="pairwise.complete.obs" )^2

      if (!is.na(rsq)) if (rsq > target.r2) break()

      # drop a few extreme data points and redo
      iid = v$summary.random$xiid[["mean"]]
      qiid = quantile( iid, probs=probs, na.rm=TRUE )
      i = which( iid > qiid[2] | iid < qiid[1] )
      if (length( i) > 0 ) z$y[i] = z$p[i]
      nw0 = nw
      nw = length( which(is.finite( z$y))) # to check for convergence
    }
    # return to original scale
    z$y = z$y + ymean
    z$p = z$p + ymean
  }

  if (method=="kernel.density") {
    # default method is kernel( "modified.daniell", c(2,1)),
    # a 2X smoothing kernel: 2 adjacent values and then 1 adjacent on the smoothed series .. i.e. a high-pass filter
    nd = length(smoothing.kernel$coef) - 1 # data points will be trimmed from tails so need to recenter:

    # test in case x increments are not uniform
    vv = unique( diff( z$x ) )
    xr = min(vv[vv>0])
    uu = unique( round( vv, abs(log10(xr))))
    if (length(uu) > 1 ) {
      # interpolate missing/skipped data before smoothing
      ifunc = approxfun( x= z$x, y=z$y, method="linear", rule=2 )
      r0 = range( z$x, na.rm=TRUE )
      fx = seq( r0[1], r0[2], by= min(uu) )
      fy = ifunc( fx)
    } else if (length(uu) ==1 ) {
      fx = z$x
      fy = z$y
    }
    fn = length( fx)
    si = (nd+1):(fn-nd)
    fp = rep( NA, fn)
    fp[si] = kernapply( as.vector(z$y), smoothing.kernel )

    pfunc = approxfun( x=fx, y=fp, method="linear", rule=2)
    z$p = pfunc( z$x)
    # plot( p ~ x, z, type="l" )

    # constrain range of predicted data to the input data range
    TR =  quantile( z$y, probs=probs, na.rm=TRUE  )
    toolow = which( z$p < TR[1] )
    if ( length(toolow) > 0 )  z$p[toolow] = TR[1]
    toohigh = which( z$p > TR[2] )
    if ( length(toohigh) > 0 ) z$p[toohigh] = TR[2]
    # lines( y ~ x, z, col="green" )
  }

  if (method=="tukey" ) {
    z$p =smooth( z$y )
  }

  if (method=="simple.linear") {
    pfunc = approxfun( x=z$x, y=z$y, method="linear", rule=2)
    z$p = pfunc( z$x)
  }

  if (method=="gam") {
    require(mgcv)
    gmod = gam( y ~ s(x) , data=z)
    z$p = predict( gmod, z, type="response" )
  }

  if (method=="smooth.spline") {
    # similar to inla .. duplicated x is problematic for smooth.spline ... add small noise
    nn = abs( diff( z$x) )
    dd = median( nn[nn>0], na.rm=TRUE )
    z$x = jitter( z$x, amount=dd / 20) # add noise as inla seems unhappy with duplicates in x?
    rsq = 0
    nw = length( which(is.finite( z$y)))
    nw0 = nw + 1
    count = 0
    while ( nw != nw0 ) {
      count = count + 1
      if ( count > nmax ) break()  # in case of endless loop
      uu = smooth.spline( x=z$x, y=z$y, keep.data=FALSE, control.spar=list(tol=dd / 20) )
      if ( length(uu$x) != nrow(z)  ) {
        afun = approxfun( x=uu$x, y=uu$y,  method="linear", rule=2  )
        vv = afun( z$x )
        z$p = vv$y
      } else {
        z$p = uu$y
      }
      rsq = cor( z$p, z$y, use="pairwise.complete.obs" )^2
      if (!is.na(rsq)){
        if (rsq > target.r2) break()
      }
      # drop a few extreme data points and redo
      iid = z$y - z$p
      qiid = quantile( iid, probs=probs, na.rm=TRUE )
      i = which( iid > qiid[2] | iid < qiid[1] )
      if (length( i) > 0 ) {
        z$y[i] = z$p[i]
      }
      nw0 = nw
      nw = length( which(is.finite( z$y))) # to check for convergence
    }
  }

  return(z$p)

}