#########################################################################
# Simulation trial

trial <- function(tau=1.2) {
  args = commandArgs(trailingOnly=TRUE)
  config_id <- args[1] %>% as.numeric
  trial_id <- Sys.getenv('SLURM_ARRAY_TASK_ID') %>% as.numeric
  obs <- .obs(config_id)
  y <- obs$y; X <- obs$X; Z <- obs$Z; D <- obs$D; u <- obs$u
  n <- nrow(X); px <- ncol(X); pz <- ncol(Z)
  Sigma_z <- obs$Sigma_z
  Alpha0 <- obs$Alpha0
  sigma0_u <- obs$sigma0_u; sigma0_v <- obs$sigma0_v
  beta0 <- obs$beta0

  # first stage estimation
  lambda_j.Alpha0hat <- .lambda_j.Alpha0hat(X, Z, sigma0_v, tune_type = "CV")
  lambda_j <- lambda_j.Alpha0hat$lambda_j; Alpha0hat <- lambda_j.Alpha0hat$Alpha0hat

  # second-stage estimation
  D.hat <- Z %*% Alpha0hat
  lambda.beta_Lasso <- .lambda.beta_Lasso(y, D.hat, sigma0_u, no_pen_ids=c())
  lambda <- lambda.beta_Lasso$lambda; beta_Lasso_D.hat <- lambda.beta_Lasso$beta_Lasso

  # relaxed inverse estimation
  Sigma_d.hat <- .Sigmahat(D.hat)
  Id <- diag(1, px) # Identity matrix

  # Calculate JM estimator
  res_JM <- .Theta.hat_JM(Sigma_d.hat, n)
  Theta.hat_JM <- res_JM$Theta.hat; mus_JM <- res_JM$mus
  mu_stars_JM <- map_dbl(1:px,
    ~ { (Sigma_d.hat %*% Theta.hat_JM[.,] - Id[.,]) %>% abs %>% max})
  objs_JM <- map_dbl(1:px,
    ~ { Theta.hat_JM[.,] %>% abs %>% sum })

  # Calculate CLIME estimator
  mus_CLIME <- find_mus(Sigma_d.hat) * tau
  Theta.hat_CLIME <- .Theta.hat_CLIME(Sigma_d.hat, mus_CLIME)
  mu_stars_CLIME <- map_dbl(1:px,
    ~ { (Sigma_d.hat %*% Theta.hat_CLIME[.,] - Id[.,]) %>% abs %>% max})
  objs_CLIME <- map_dbl(1:px,
    ~ { Theta.hat_CLIME[.,] %>% abs %>% sum })

  # de-biased second-stage lasso estimation
  beta_debiased_CLIME <- .beta_debiased(y, X, D.hat, beta_Lasso_D.hat, Theta.hat_CLIME)
  beta_debiased_JM <- .beta_debiased(y, X, D.hat, beta_Lasso_D.hat, Theta.hat_JM)

  # Predict the second-stage noise
  u.hat <- y - X %*% beta_Lasso_D.hat
  sd_u.hat <- u.hat^2 %>% mean %>% sqrt

  # To record/compute estimate of Theta_jj (and Theta_jj)
  Theta <- solve(t(Alpha0) %*% Sigma_z %*% Alpha0)

  # Calculate remainder terms
  rem_1_CLIME <- (Theta.hat_CLIME-Theta)%*%t(D)%*%u/sqrt(n)
  rem_1_JM <- (Theta.hat_JM-Theta)%*%t(D)%*%u/sqrt(n)
  rem_2_CLIME <- Theta.hat_CLIME%*%t(D.hat-D)%*%u/sqrt(n)
  rem_2_JM <- Theta.hat_JM%*%t(D.hat-D)%*%u/sqrt(n)
  rem_3_CLIME <- Theta.hat_CLIME%*%t(D.hat)%*%(X-D.hat)%*%(beta_Lasso_D.hat-beta0)/sqrt(n)
  rem_3_JM <- Theta.hat_JM%*%t(D.hat)%*%(X-D.hat)%*%(beta_Lasso_D.hat-beta0)/sqrt(n)
  rem_4_CLIME <-sqrt(n)*(Theta.hat_CLIME%*%Sigma_d.hat-Id)%*%(beta_Lasso_D.hat-beta0)
  rem_4_JM <-sqrt(n)*(Theta.hat_JM%*%Sigma_d.hat-Id)%*%(beta_Lasso_D.hat-beta0)

  # Calculate main term
  w <- Theta%*%t(D)%*%u/sqrt(n)

  # estimator data
  df_est <- data.frame(
    config_id = rep(config_id, 3*px),
    trial_id = rep(trial_id, 3*px),
    estimator = c(rep("Debiased_CLIME", px), rep("Debiased_JM", px), rep("Lasso", px)),
    j = rep(1:px, 3),
    estimate_j = c(beta_debiased_CLIME, beta_debiased_JM, beta_Lasso_D.hat),
    beta0_j = rep(beta0, 3),
    SE1 = c(.SE1(Sigma_d.hat, Theta.hat_CLIME, sd_u.hat),
            .SE1(Sigma_d.hat, Theta.hat_JM, sd_u.hat),
            rep(NA, px)),
    SE2 = c(.SE2(D.hat, Theta.hat_CLIME, u.hat),
            .SE2(D.hat, Theta.hat_JM, u.hat),
            rep(NA, px)),
    SE3 = c(.SE3(Theta.hat_CLIME, sd_u.hat),
            .SE3(Theta.hat_JM, sd_u.hat),
            rep(NA, px)),
    Theta_jj = c(diag(Theta), diag(Theta), rep(NA, px)),
    Theta.hat_jj = c(diag(Theta.hat_CLIME), diag(Theta.hat_JM), rep(NA, px)),
    mu = c(mus_CLIME, mus_JM, rep(NA, px)),
    mu_stars = c(mu_stars_CLIME, mu_stars_JM, rep(NA, px)),
    objs = c(objs_CLIME, objs_JM, rep(NA, px)),
    w = c(w, w, rep(NA, px)),
    rem_1 = c(rem_1_CLIME, rem_1_JM, rep(NA, px)),
    rem_2 = c(rem_2_CLIME, rem_2_JM, rep(NA, px)),
    rem_3 = c(rem_3_CLIME, rem_3_JM, rep(NA, px)),
    rem_4 = c(rem_4_CLIME, rem_4_JM, rep(NA, px)),
    lambda_j = rep(lambda_j, 3)
  )
  # df_est <- data.frame(
  #   config_id = rep(config_id, 2*px),
  #   trial_id = rep(trial_id, 2*px),
  #   estimator = c(rep("Debiased_CLIME", px), rep("Lasso", px)),
  #   j = rep(1:px, 2),
  #   estimate_j = c(beta_debiased_CLIME, beta_Lasso_D.hat),
  #   beta0_j = rep(beta0, 2),
  #   SE1 = c(.SE1(Sigma_d.hat, Theta.hat_CLIME, sd_u.hat),
  #           rep(NA, px)),
  #   SE2 = c(.SE2(D.hat, Theta.hat_CLIME, u.hat),
  #           rep(NA, px)),
  #   SE3 = c(.SE3(Theta.hat_CLIME, sd_u.hat),
  #           rep(NA, px)),
  #   Theta_jj = c(diag(Theta), rep(NA, px)),
  #   Theta.hat_jj = c(diag(Theta.hat_CLIME), rep(NA, px)),
  #   mu_stars = c(mu_stars_CLIME, rep(NA, px)),
  #   objs = c(objs_CLIME, rep(NA, px))
  #   # vhat = rep(vhat, 2),
  #   lambda_j = rep(lambda_j, 2)
  # )

  # statistics
  mu_star_CLIME <- (Theta.hat_CLIME %*% Sigma_d.hat - diag(1, ncol(Sigma_d.hat))) %>% abs %>% max
  mu_star_JM <- (Theta.hat_JM %*% Sigma_d.hat - diag(1, ncol(Sigma_d.hat))) %>% abs %>% max
  mse_Debiased_CLIME <- (y - X %*% beta_debiased_CLIME)^2 %>% mean
  mse_Debiased_JM <- (y - X %*% beta_debiased_JM)^2 %>% mean
  mse_Lasso_D.hat <- (y - X %*% beta_Lasso_D.hat)^2 %>% mean

  # estimation statistics data
  df_stats <- data.frame(
    config_id = rep(config_id, 3),
    trial_id = rep(trial_id, 3),
    estimator = c("Debiased_CLIME", "Debiased_JM", "Lasso"),
    mse = c(mse_Debiased_CLIME, mse_Debiased_JM, mse_Lasso_D.hat),
    sd_u = rep(sd_u.hat, 3),
    # mu = c(mu, NA),
    tau = c(tau, NA, NA),
    mu_star = c(mu_star_CLIME, mu_star_JM, NA),
    # trial_cvg = c(trial_cvg, NA),
    lambda = rep(lambda, 3)
  )
  # df_stats <- data.frame(
  #   config_id = rep(config_id, 2),
  #   trial_id = rep(trial_id, 32),
  #   estimator = c("Debiased_CLIME", "Lasso"),
  #   mse = c(mse_Debiased_CLIME, mse_Lasso_D.hat),
  #   sd_u = rep(sd_u.hat, 2),
  #   # mu = c(mu, NA),
  #   tau = c(tau, NA),
  #   mu_star = c(mu_star_CLIME, NA),
  #   # trial_cvg = c(trial_cvg, NA),
  #   lambda = rep(lambda, 2)
  # )

  write.csv(df_est, paste("res/est_", config_id, "_", trial_id, ".csv", sep=""))
  write.csv(df_stats, paste("res/stats_", config_id, "_", trial_id, ".csv", sep=""))
}


# qvalue() function extracted from the package "qvalue"
# by Alan Dabney and John D. Storey, with assistance from Gregory R. Warnes
#
# DESCRIPTION file follows:
#
# Package: qvalue
# Title: Q-value estimation for false discovery rate control
# Version: 1.32.0
# Author: Alan Dabney <adabney@u.washington.edu> and John D. Storey
# <jstorey@u.washington.edu>, with assistance from Gregory R.
# Warnes <gregory_r_warnes@groton.pfizer.com>
#
# Description: This package takes a list of p-values resulting from the
# simultaneous testing of many hypotheses and estimates their
# q-values. The q-value of a test measures the proportion of
# false positives incurred (called the false discovery rate)
# when that particular test is called significant. Various plots
# are automatically generated, allowing one to make sensible
# significance cut-offs. Several mathematical results have
# recently been shown on the conservative accuracy of the
# estimated q-values from this software. The software can be
# applied to problems in genomics, brain imaging, astrophysics,
# and data mining.
# Maintainer: John D. Storey <jstorey@u.washington.edu>
# 	Imports: graphics, stats
# License: LGPL
# biocViews: MultipleComparisons
# Packaged: 2012-10-02 02:49:23 UTC; biocbuild
# InstallableEverywhere: yes

qvalue <- function(p=NULL, lambda=seq(0,0.90,0.05), pi0.method="smoother", fdr.level=NULL, robust=FALSE,
									 gui=FALSE, smooth.df = 3, smooth.log.pi0 = FALSE) {

	#This is just some pre-processing
	if(is.null(p))  ## change by Alan
	{qvalue.gui(); return("Launching point-and-click...")}
	if(gui & !interactive())  ## change by Alan
		gui = FALSE

	if(min(p)<0 || max(p)>1) {
		if(gui) ## change by Alan:  check for GUI
			eval(expression(postMsg(paste("ERROR: p-values not in valid range.", "\n"))), parent.frame())
		else
			print("ERROR: p-values not in valid range.")
		return(0)
	}
	if(length(lambda)>1 && length(lambda)<4) {
		if(gui)
			eval(expression(postMsg(paste("ERROR: If length of lambda greater than 1, you need at least 4 values.",
																		"\n"))), parent.frame())
		else
			print("ERROR: If length of lambda greater than 1, you need at least 4 values.")
		return(0)
	}
	if(length(lambda)>1 && (min(lambda) < 0 || max(lambda) >= 1)) { ## change by Alan:  check for valid range for lambda
		if(gui)
			eval(expression(postMsg(paste("ERROR: Lambda must be within [0, 1).", "\n"))), parent.frame())
		else
			print("ERROR: Lambda must be within [0, 1).")
		return(0)
	}
	m <- length(p)
	#These next few functions are the various ways to estimate pi0
	if(length(lambda)==1) {
		if(lambda<0 || lambda>=1) { ## change by Alan:  check for valid range for lambda
			if(gui)
				eval(expression(postMsg(paste("ERROR: Lambda must be within [0, 1).", "\n"))), parent.frame())
			else
				print("ERROR: Lambda must be within [0, 1).")
			return(0)
		}

		pi0 <- mean(p >= lambda)/(1-lambda)
		pi0 <- min(pi0,1)
	}
	else {
		pi0 <- rep(0,length(lambda))
		for(i in 1:length(lambda)) {
			pi0[i] <- mean(p >= lambda[i])/(1-lambda[i])
		}

		if(pi0.method=="smoother") {
			if(smooth.log.pi0)
				pi0 <- log(pi0)

			spi0 <- smooth.spline(lambda,pi0,df=smooth.df)
			pi0 <- predict(spi0,x=max(lambda))$y

			if(smooth.log.pi0)
				pi0 <- exp(pi0)
			pi0 <- min(pi0,1)
		}
		else if(pi0.method=="bootstrap") {
			minpi0 <- min(pi0)
			mse <- rep(0,length(lambda))
			pi0.boot <- rep(0,length(lambda))
			for(i in 1:100) {
				p.boot <- sample(p,size=m,replace=TRUE)
				for(i in 1:length(lambda)) {
					pi0.boot[i] <- mean(p.boot>lambda[i])/(1-lambda[i])
				}
				mse <- mse + (pi0.boot-minpi0)^2
			}
			pi0 <- min(pi0[mse==min(mse)])
			pi0 <- min(pi0,1)
		}
		else {  ## change by Alan: check for valid choice of 'pi0.method' (only necessary on command line)
			print("ERROR: 'pi0.method' must be one of 'smoother' or 'bootstrap'.")
			return(0)
		}
	}
	if(pi0 <= 0) {
		if(gui)
			eval(expression(postMsg(
				paste("ERROR: The estimated pi0 <= 0. Check that you have valid p-values or use another lambda method.",
							"\n"))), parent.frame())
		else
			print("ERROR: The estimated pi0 <= 0. Check that you have valid p-values or use another lambda method.")
		return(0)
	}
	if(!is.null(fdr.level) && (fdr.level<=0 || fdr.level>1)) {  ## change by Alan:  check for valid fdr.level
		if(gui)
			eval(expression(postMsg(paste("ERROR: 'fdr.level' must be within (0, 1].", "\n"))), parent.frame())
		else
			print("ERROR: 'fdr.level' must be within (0, 1].")
		return(0)
	}
	#The estimated q-values calculated here
	u <- order(p)

	# change by Alan
	# ranking function which returns number of observations less than or equal
	qvalue.rank <- function(x) {
		idx <- sort.list(x)

		fc <- factor(x)
		nl <- length(levels(fc))
		bin <- as.integer(fc)
		tbl <- tabulate(bin)
		cs <- cumsum(tbl)

		tbl <- rep(cs, tbl)
		tbl[idx] <- tbl

		return(tbl)
	}

	v <- qvalue.rank(p)

	qvalue <- pi0*m*p/v
	if(robust) {
		qvalue <- pi0*m*p/(v*(1-(1-p)^m))
	}
	qvalue[u[m]] <- min(qvalue[u[m]],1)
	for(i in (m-1):1) {
		qvalue[u[i]] <- min(qvalue[u[i]],qvalue[u[i+1]],1)
	}
	#The results are returned
	if(!is.null(fdr.level)) {
		retval <- list(call=match.call(), pi0=pi0, qvalues=qvalue, pvalues=p, fdr.level=fdr.level, ## change by Alan
									 significant=(qvalue <= fdr.level), lambda=lambda)
	}
	else {
		retval <- list(call=match.call(), pi0=pi0, qvalues=qvalue, pvalues=p, lambda=lambda)
	}
	class(retval) <- "qvalue"
	return(retval)
}


## --------------------
## --------------------
## --------------------
## This file contains functions mainly for cross validation of Adaptive Penalized Tensor Decomposition.
## Most function are only suitable for three dimension tensor. For higher-order tensor, codes should be modified adaptively.
## --------------------
## --------------------
## --------------------

## --------------------
## Initialization for cross validation.
## --------------------
prepare_adaptive_adj <- function(tnsr, num_components, Niter = 1500, tol = 1e-05, year = 5, trace = FALSE, ku, kv, kw, splits = 10) {
    original <- tnsr
    Test <- tnsr[, , (dim(tnsr)[3] - year + 1):dim(tnsr)[3]]
    V1 <- list()
    S1 <- list()
    In1 <- list()
    mean1 <- list()
    sd1 <- list()
    tnsr <- tnsr[, , 1:(dim(tnsr)[3] - year)]
    for (m in 1:splits) {
        message("Training on years: ", paste(1:(dim(tnsr)[3] - splits - year + m), collapse = "+"))
        message("Test on years: ", paste((dim(tnsr)[3] - splits - year + m + 1):((dim(tnsr)[3] - splits - year + m + 1) + year - 1), collapse = "+"))

        temp_s <- tnsr@data[, , 1:(dim(tnsr)[3] - splits - year + m)]
        S1[[m]] <- as.tensor((temp_s - array(mean(temp_s), dim = dim(temp_s))) / sd(temp_s))
        mean1[[m]] <- c(mean(temp_s))
        sd1[[m]] <- c(sd(temp_s))

        V1[[m]] <- as.tensor(tnsr@data[, , (dim(tnsr)[3] - splits - year + m + 1):((dim(tnsr)[3] - splits - year + m + 1) + year - 1)])

        make_init <- gen_startvals(num_components, S1[[m]])

        message("Fitting unpenalized tensor...")
        In1[[m]] <- multiple_tf_simple(
            tnsr = S1[[m]], d_hat = NULL, num_components = num_components,
            u_init = make_init$u, v_init = make_init$v, w_init = make_init$w,
            Niter = Niter, tol = tol,
            c1 = rep(0, num_components), c2 = rep(0, num_components), c3 = rep(0, num_components),
            ku = ku, kv = kv, kw = kw, trace = trace
        )
        message("Fitting completed.")
    }

    ## generate D
    Dut <- list()
    Dvt <- list()
    Dwt <- list()
    for (m in 1:splits) {
        Dut[[m]] <- create_D(In1[[m]]$u, ord = (ku + 1))
        Dvt[[m]] <- create_D(In1[[m]]$v, ord = (kv + 1))
        Dwt[[m]] <- create_D(In1[[m]]$w, ord = (kw + 1))
    }

    return(list(
        Dut = Dut, Dvt = Dvt, Dwt = Dwt, In1 = In1, S1 = S1, V1 = V1, year = year,
        num_components = num_components, mean1 = mean1, sd1 = sd1, Test = Test, original = original
    ))
}

## --------------------
## Function for cross validation.
## --------------------
multiple_cv_adaptive_adj <- function(init, grid, ku = NULL, kv = NULL, kw = NULL, tol = 5e-05, Niter_small = 1500, splits = 10) {

    # trian data
    # input tenor should be training +validation tensor
    # third dimension should be year dimension

    tstart <- Sys.time()
    S1 <- init$S1
    V1 <- init$V1
    Dut <- init$Dut
    Dvt <- init$Dvt
    Dwt <- init$Dwt
    num_components <- init$num_components
    year <- init$year
    In1 <- init$In1
    mean1 <- init$mean1
    sd1 <- init$sd1

    sfInit(parallel = TRUE, cpus = 6)
    # load
    sfLibrary(rTensor)
    sfLibrary(genlasso)
    sfLibrary(forecast)
    sfLibrary(mgcv) # gam
    sfLibrary(Hmisc)
    sfExport(
        "multiple_D_adaptive", "PTD_D", "In1", "tol", "product", "S1", "cal_pmse_gam_adj",
        "cal_pmse_arima_adj", "cal_pmse_linearextra_adj", "num_components", "V1", "year", "Niter_small", "Dut", "Dvt", "Dwt", "mean1", "sd1"
    )

    best <- 0
    meanPMSE_gam <- meanPMSE_linextrap <- meanPMSE_arima <- NULL # record each combination's mean of the prediction RMSEs (averge of 10 )
    allPMSE_gam <- allPMSE_linextrap <- allPMSE_arima <- NULL

    for (i in 1:nrow(grid)) {
        message("Tuning parameters: ", paste(grid[i, ], collapse = "/"))
        sfExport("i", "grid")

        parallel1 <- function(m, num_components, In1, grid, Niter_small, tol, Dut, Dvt, Dwt) {
            temp <- multiple_D_adaptive(
                tnsr = S1[[m]], num_components = num_components, d_hat = In1[[m]]$d,
                u_init = In1[[m]]$u, v_init = In1[[m]]$v, w_init = In1[[m]]$w,
                c1 = rep(grid[i, 1], num_components), c2 = rep(grid[i, 2], num_components), c3 = rep(grid[i, 3], num_components),
                Niter = Niter_small, tol = tol, Du = Dut[[m]], Dv = Dvt[[m]], Dw = Dwt[[m]], trace = FALSE
            )
        }

        tempresult <- sfLapply(1:splits, parallel1, num_components, In1, grid, Niter_small, tol, Dut, Dvt, Dwt)
        sfExport("tempresult")

        temprmse_gam <- as.numeric(sfLapply(1:splits, cal_pmse_gam_adj, year = year))
        temprmse_linextrap <- as.numeric(sfLapply(1:splits, cal_pmse_linearextra_adj, year = year))
        temprmse_arima <- as.numeric(sfLapply(1:splits, cal_pmse_arima_adj, year = year))
        allPMSE_gam <- rbind(allPMSE_gam, temprmse_gam)
        allPMSE_arima <- rbind(allPMSE_arima, temprmse_arima)
        allPMSE_linextrap <- rbind(allPMSE_linextrap, temprmse_linextrap)

        # bias and variance
        meanPMSE_gam <- c(meanPMSE_gam, mean(temprmse_gam))
        meanPMSE_linextrap <- c(meanPMSE_linextrap, mean(temprmse_linextrap))
        meanPMSE_arima <- c(meanPMSE_arima, mean(temprmse_arima))

        # print(meanPMSE_gam)
        # print(meanPMSE_arima)
        # print(meanPMSE_linextrap)
    }

    ## stop parallel
    sfStop()
    tend <- Sys.time()
    message("Fitting completed!")

    colnames(grid) <- c("U", "V", "W")
    performance <- data.frame(grid,
        meanPMSE_gam = meanPMSE_gam, meanPMSE_linextrap = meanPMSE_linextrap, meanPMSE_arima = meanPMSE_arima,
        mean_diffunpenPMSE_gam = meanPMSE_gam - meanPMSE_gam[1], mean_diffunpenPMSE_linextrap = meanPMSE_linextrap - meanPMSE_linextrap[1],
        mean_diffunpenPMSE_arima = meanPMSE_arima - meanPMSE_arima[1]
    )
    performance$sd_diffunpenPMSE_gam <- apply(t(t(allPMSE_gam) - allPMSE_gam[1, ]), 1, sd)
    performance$sd_diffunpenPMSE_arima <- apply(t(t(allPMSE_arima) - allPMSE_arima[1, ]), 1, sd)
    performance$sd_diffunpenPMSE_linextrap <- apply(t(t(allPMSE_linextrap) - allPMSE_linextrap[1, ]), 1, sd)

    out <- list(
        best_parameters_gam = grid[which.min(performance$meanPMSE_gam), ],
        best_parameters_linextrap = grid[which.min(performance$meanPMSE_linextrap), ],
        best_parameters_arima = grid[which.min(performance$meanPMSE_arima), ],
        performance = performance, init = init, time = tend - tstart
    )
    return(out)
}


seasonalize <- function(arr, maxes = TRUE){
  # this function finds the max or min by climatological season (DJF, MAM, JJA, SON), assuming a 3-d array with month as the second dimension

  if(maxes) FUN <- max else FUN <- min

  nSeas <- 4
  seasonIndices <- list(DJF = c(1, 2, 12), MAM = 3:5, JJA = 6:8, SON = 9:11)
  dec <- 12
  nYr <- dim(arr)[1]

  arr[2:nYr, dec, ] <- arr[1:(nYr - 1), dec, ]
  arr[1, dec, ] <- NA

  tmp <- array(NA, c(nYr, nSeas, dim(arr)[3]))
  for(j in 1:nSeas){
    tmp[ , j, ] <- apply(arr[, seasonIndices[[j]], , drop = FALSE], c(1, 3), FUN, na.rm = FALSE)
  }
  return(tmp)
}

is.wholenumber <- function(x, tol = .Machine$double.eps^0.5)  abs(x - round(x)) < tol

validateIndices <- function(model, nVar){
  if(min(model) < 0 | max(model) > nVar)
    return(FALSE)
  if(sum(is.wholenumber(model)) < length(model))
    return(FALSE)
  if(length(unique(model)) != length(model))
    return(FALSE)
  return(TRUE)
}

normalize <- function(vec, shift = NULL, lower = NULL, upper = NULL){
  if(is.null(shift))
    shift <- mean(vec)
  if(is.null(lower))
    lower <- min(vec)
  if(is.null(upper))
    upper <- max(vec)
  (vec - shift) / (upper - lower)
}

returnValue <- function(fit, rvInterval = 20, rvCovariates = NULL){

  if(!(class(fit) == 'gev.fit' | class(fit) == 'pp.fit'))
    stop('Function only handles gev.fit and pp.fit objects')

#  if(!identical(fit$link[[1]], identity) || !identical(shlink, identity) ||  == 'gev.fit' | class(fit) == 'pp.fit'))
  #  stop('Function only handles gev.fit and pp.fit objects')

  links = substring(fit$link, 3, nchar(fit$link) - 1)
  links = unlist(strsplit(links, ", "))

  if(links[1] != "identity" || links[3] != "identity" || !(links[2] %in% c("identity", "exp")))
    stop("Function only computes return values for identity link functions for location and shape and for identity or exponential link functions for scale")

  locationIndices <- 1:(length(fit$model[[1]]) + 1)
  pos = 1 + max(locationIndices)
  scaleIndices <- pos:(pos+length(fit$model[[2]]))
  pos = 1 + max(scaleIndices)
  shapeIndices <- pos:(pos+length(fit$model[[3]]))

  location <- eval(as.name(links[1]))(sum(c(1, rvCovariates[fit$model[[1]]]) * fit$mle[locationIndices]))
  scale <- eval(as.name(links[2]))(sum(c(1, rvCovariates[fit$model[[2]]]) * fit$mle[scaleIndices]))
  shape <- eval(as.name(links[3]))(sum(c(1, rvCovariates[fit$model[[3]]]) * fit$mle[shapeIndices]))

  yp <- -log(1 - 1/rvInterval)
  ypshape <- yp^(-shape)
  zp <- location - (scale/shape) * (1 - ypshape)

  if(links[2] == "identity"){
    scaleTerm = 1
  } else{  # exponential link
    scaleTerm = scale
  }

  grad <- c(c(1, rvCovariates[fit$model[[1]]]),
    -(scaleTerm/shape) * (1 - ypshape) * c(1, rvCovariates[fit$model[[2]]]),
    ((scale/shape^2) * (1 - ypshape) - (scale/shape) * ypshape * log(yp)) * c(1, rvCovariates[fit$model[[3]]]))

  se <- sqrt(t(grad) %*% fit$cov %*% grad)

  return(c(zp, se))
}

# difference in return values for single time series, different covariate values
returnValueDiff <- function(fit, rvInterval = 20, rvCovariates = NULL){
  if(!(class(fit) == 'gev.fit' || class(fit) == 'pp.fit'))
    stop('Function only handles gev.fit and pp.fit objects')

  links = substring(fit$link, 3, nchar(fit$link) - 1)
  links = unlist(strsplit(links, ", "))

  if(links[1] != "identity" || links[3] != "identity" || !(links[2] %in% c("identity", "exp")))
    stop("Function only computes return value differences for identity link functions for location and shape and for identity or exponential link functions for scale")

  locationIndices <- 1:(length(fit$model[[1]]) + 1)
  pos = 1 + max(locationIndices)
  scaleIndices <- pos:(pos+length(fit$model[[2]]))
  pos = 1 + max(scaleIndices)
  shapeIndices <- pos:(pos+length(fit$model[[3]]))

  location <- eval(as.name(links[1]))(cbind(rep(1, 2), rvCovariates[ , fit$model[[1]]]) %*% fit$mle[locationIndices])
  scale <- eval(as.name(links[2]))(cbind(rep(1, 2), rvCovariates[ , fit$model[[2]]]) %*% fit$mle[scaleIndices])
  shape <- eval(as.name(links[3]))(cbind(rep(1, 2), rvCovariates[ , fit$model[[3]]]) %*% fit$mle[shapeIndices])

  yp <- -log(1 - 1/rvInterval)
  ypshape <- yp^(-shape)
  zpDiff <- diff(location - (scale/shape) * (1 - ypshape))

  if(links[2] == "identity"){
    scaleTerm = rep(1, 2)
  } else{  # exponential link
    scaleTerm = rep(scale, 2)
  }

  grad <- c(0, rvCovariates[2, fit$model[[1]]] - rvCovariates[1, fit$model[[1]]],
    -(scaleTerm[2]/shape[2]) * (1 - ypshape[2]) * c(1, rvCovariates[2, fit$model[[2]]]) +
    (scaleTerm[1]/shape[1]) * (1 - ypshape[1]) * c(1, rvCovariates[1, fit$model[[2]]]),
    ((scale[2]/shape[2]^2) * (1 - ypshape[2]) - (scale[2]/shape[2]) * ypshape[2] * log(yp) ) * c(1, rvCovariates[2, fit$model[[3]]]) -
    ((scale[1]/shape[1]^2) * (1 - ypshape[1]) - (scale[1]/shape[1]) * ypshape[1] * log(yp) ) * c(1, rvCovariates[1, fit$model[[3]]]))

  se <- sqrt(t(grad) %*% fit$cov %*% grad)

  return(c(zpDiff, se))
}

removeRuns <- function(vals, days, maxes = TRUE){
  n <- length(vals)
  if(maxes) FUN <- max else FUN <- min
  if(n > 1){
    checkBack <- c(TRUE, days[2:n] - days[1:(n-1)] != 1)
    checkFwd <- c(days[1:(n-1)] - days[2:n] == -1, FALSE)
    seqStarts <- which(checkBack & checkFwd)
    for(start in seqStarts){
      pos = start
      while(pos < n && days[pos + 1] - days[start] == pos + 1 - start) pos <- pos + 1
      tmp <- vals[start:pos]
      jitter = rnorm(length(tmp), 0, 1e-12)  # amounts to randomly choosing a max when values are equal
      tmp[tmp + jitter < FUN(tmp + jitter)] <- NA
      vals[start:pos] <- tmp
    }
  }
  return(vals)
}

# alternative: function for removing multiple exceedances in blockLen-day windows; surprisingly this is much slower
withinBlockScreenAlt <- function(vals, days, blockLen = 10, maxes = TRUE){
  n <- length(vals)
  if(n > 1){
    jitter = rnorm(n, 0, 1e-12) # randomly choose which value in block to retain when there are ties
    if(maxes) FUN <- max else FUN <- min
    block <- data.frame(block = cut(days, seq(0, max(days)+blockLen, by = blockLen)))
    tmp <- tapply(vals + jitter, block, FUN)
    blockExtrema <- data.frame(block = names(tmp), extrema = tmp)
    comparators = merge(block, blockExtrema, all.x = TRUE, all.y = FALSE)
    vals[vals + jitter < comparators$extrema] = NA
  }
  return(vals)
}

withinBlockScreen <-  function(vals, days, blockLen = 10, maxes = TRUE){
  if(!identical(days, sort(days)))
    stop('data must be sorted by days in ascending order')
  n <- length(vals)
  if(n > 1){
    i <- 1
    while(i < length(vals)){
      start <- i
      while( (days[i+1]-1) %/% blockLen == (days[i]-1) %/% blockLen &&  i < length(vals)){
        i <- i + 1
      }
      if(i > start){
        jitter = rnorm(i - start + 1, 0, 1e-12)
        tmp = vals[start:i]
        tmp[tmp + jitter < max(tmp + jitter)] = NA
        vals[start:i] = tmp
      }
      i = i + 1
    }
  }
  return(vals)
}

lik.ratio.test = function(fit1, fit2){
  if(!(class(fit1) %in% c("gev.fit", "pp.fit")) ||
     !(class(fit2) %in% c("gev.fit", "pp.fit")))
    stop("Fit objects must be of class 'gev.fit' or 'pp.fit'")
  dev = 2 * abs(fit1$nllh - fit2$nllh)
  df = abs(length(unlist(fit1$model)) - length(unlist(fit2$model)))
  return(list(df = df, negLogLik = c(fit1$nllh, fit2$nllh), chisq = dev, p = 1 - pchisq(dev, df)))
}

#linkManip = function(mulink, siglink, shlink)
# return(deparse(substitute(c(mulink, siglink, shlink))))


## Functions for ecological model

date2doy <- function(yr, mo, dy){
  yr <- as.character(yr)
  mo <- as.character(mo)
  dy <- as.character(dy)

  dt <- paste(c(yr,'-',mo,'-',dy), collapse='')
  doy <- strftime(dt, format='%j')
  return(doy)
}

define_global_variables <- function(){

  # Site constants
  PAR2SWIR <<- 0.43 # PAR to shortwave radiation ratio
  STD_MERAD <<- -75.0 # Time zone longitude
  LATI <<- 35.9736 # Site latitude (CHANGE TO MMS)
  LONGI <<- -79.1004 # Site longitude (CHANGE TO MMS)
  hs <<- 20.0 # Screen height
  h2 <<- 16.2 # Default canopy height
  PA <<- 101325.0 # Air pressure in Pa
  CP <<- 1010.0 # Specific heat of air (J kg-1 K-1)
  CA <<- 410.0 # Atmospheric CO2 concentration (ppm)

  # Physical constants
  SO <<- 1367.0 # Solar constant (W m-2)

  # Leaf constants
  alpha_par <<- 0.80
  alpha_nir <<- 0.20
  rho_s <<- 0.10
  x <<- 1.0 # Leaf angle distribution parameter
  omega <<- 1.0 # Clumping index
  m <<- 8.0 # Ball-Berry model slope
  g1_blayer <<- 0.005 # Leaf boundary layer conductance in m/s
  g0 <<- 0.010 # Parameter in Leuning's gs function in mol m-2 s-1

  # Soil constants
  Ksat <<- 0.6 # Saturated conductance (m day-1)
  wilting_p <<- 0.08 # Wilting poit soil water content (unitless)
  porosity <<- 0.54 # Porosity at depth zero (unitless)
  p1 <<- 0.478 # Air entry pressure in meters of water
  p2 <<- 0.186 # Pore size index (unitless)
  p_0 <<- 4000.0 # Porosity decay in m-1
  m_z <<- 0.24 # Saturated conductivity decay parameter (in m-1)
  soil_depth <<- 0.325 # Soil depth in rooting zone

}

canopy_aerodynamic_resistance <- function(u_h, h, h_o){
  h_u <- 2.0
  cn <- 0.002 # wind attenuation coefficient

  # Compute zero plane displacement
  d_o <- 0.7 * h_o

  # Compute the roughness length
  zo_o <- 0.1 * h_o
  u_o <- compute_toc_wind(u_h, h, h_o)
  u_m <- (u_o * exp(cn * (0.5 * h_o)/h_o - 1))
  u_m <- sapply(u_m, function(x) max(x,0))

  if(h<d_o) h <- d_o + 0.0001

  ra <- (log((h_o-d_o)/zo_o)/0.41)^2 / u_m
  ra[u_m==0] <- 9999999

  return(ra)
}

canopy_bwb_stomatal_conductance <- function(psn, Tair, a1, vpda){
  svp <- 610.7 * exp(17.38 * Tair / (239.0+Tair))
  rh <- 1.0 - vpda/svp
  gc <- g0 + a1 * psn * rh/CA

  # convert mol m-2 s-1 to m s-1
  gc <- gc * (22.4 / 1000.0)

  # convert gc (Co2) to gs (H2O)
  gs <- 1.56 * gc

  return(data.frame(gc=gc, gs=gs))
}

compute_bwb_farq_psn <- function(psnp, apar_wm2, Tair, vpda, fstheta, leaftype){
  svp <- 610.7 * exp(17.38 * Tair / (239.0 + Tair))
  rh <- 1.0 - vpda/svp
  apar <- apar_wm2*4.55
  Ca <- CA

  # calculate atmospheric O2 in Pa, assuming 21% O2 by volume
  O2 <- psnp$O2
  Ko <- psnp$Ko

  # Convert Pa to ppm
  Kc <- psnp$Kc * 1e6/PA
  gamma <- psnp$gamma * 1e6/PA

  ppe <- psnp$ppe
  if(leaftype==1){
    Vmax <- psnp$Vmax_sunlit
    Rd <- psnp$Rd_sunlit
  } else {
    Vmax <- psnp$Vmax_shaded
    Rd <- psnp$Rd_shaded
  }

  Jmax <- 2.1 * Vmax

  g1 <- 0.01 + m * rh/Ca
  K <- Kc * (1.0 + O2/Ko)

  # Calculate rubisco-limited assimilation rate
  aa <- Ca*g1 + g1*K - 1.0
  bb <- Ca*g0 + g0*K - Vmax*Ca*g1 + Vmax + Vmax*g1*gamma + Rd*Ca*g1 - Rd + Rd*g1*K
  cc = -Vmax*Ca*g0 + Vmax*gamma*g0 + Rd*Ca*g0 + Rd*K*g0

  det <- bb*bb - 4.0*aa*cc
  Av <- (-bb + sqrt(det)) / (2.0*aa)
  Av[det<0] <- -1.0

  # Calculate RuBP regeneration limited assimilation rate
  aa <- 0.7
  bb <- -Jmax - (apar/ppe)
  cc <- Jmax * apar/ppe
  J <- (-bb - sqrt(bb*bb - 4*aa*cc)) / (2.0/aa)

  aa <- 4.5*Ca*g1 - 4.5 + 10.5*gamma*g1
  bb <- 4.5*Ca*g0 + 10.5*gamma*g0 - J*g1*Ca + J + gamma*g1*J + 4.5*Rd*Ca*g1-4.5*Rd+10.5*gamma*g1*Rd
  cc <- -J*Ca*g0+4.5*Rd*Ca*g0+10.5*gamma*g0*Rd

  det <- bb*bb - 4.0*aa*cc
  Aj <- (-bb+sqrt(det))/(2.0*aa)
  Aj[det<0] <- -1.0

  # Calculate net assimilation rate
  A = pmin(Aj,Av)
  A[Av == -1.0 & Aj == -1.0] <- 0.0
  A[Av == -1.0 & Aj != -1.0] <- Aj[Av == -1.0 & Aj != -1.0]
  A[Aj == -1.0 & Av != -1.0] <- Av[Aj == -1.0 & Av != -1.0]

  return(A)
}

compute_canopy_intercepted_shortwave_rad <- function(par_D, par_d, nir_D, nir_d, lai, Kb, Kd){
  Q_sc_par <- par_D*(exp(-Kb*sqrt(alpha_par)*lai)-exp(-Kb*lai))/2.0
  Q_d_par <- par_d*(1.0-exp(-Kd*sqrt(alpha_par)*lai))/(Kd*sqrt(alpha_par)*lai)
  Q_par_soil <- par_D*exp(-Kb*sqrt(alpha_par)*lai)+par_d*exp(-Kd*sqrt(alpha_par)*lai)
  Q_nir_soil <- nir_D*exp(-Kb*sqrt(alpha_nir)*lai)+nir_d*exp(-Kd*sqrt(alpha_nir)*lai)

  sunlit_apar <- alpha_par*(par_D*Kb+ Q_sc_par + Q_d_par)
  shaded_apar <- alpha_par*(Q_sc_par + Q_d_par)

  Q_sc_nir <- nir_D*(exp(-Kb*sqrt(alpha_nir)*lai)-exp(-Kb*lai))/2.0
  Q_d_nir <- par_d*(1.0-exp(-Kd*sqrt(alpha_par)*lai))/(Kd*sqrt(alpha_par)*lai)

  sunlit_anir <- alpha_nir*(nir_D*Kb+Q_sc_nir+Q_d_nir)
  shaded_anir <- alpha_nir*(Q_sc_nir+Q_d_nir)

  rad <- data.frame(sunlit_apar=sunlit_apar,
              shaded_apar=shaded_apar,
              sunlit_anir=sunlit_anir,
              shaded_anir=shaded_anir)
  return(rad)
}

compute_canopy_net_rad <- function(rad_short, rad_long, toc_rad, sunlit_lai, shaded_lai, Kb, Kd){

  lai <- sunlit_lai + shaded_lai

  Rnet_sunlit <- rad_short$sunlit_apar + rad_short$sunlit_anir +
    rad_long$floor * (1.0 - exp(-Kd*lai)) / (Kd*lai) +
    rad_long$air * (1.0 - exp(-Kd*lai)) / (Kd*lai) +
    -2.0 * rad_long$canopy * (1.0 - exp(-Kd*lai)) / (Kd*lai)

  Rnet_shaded <- rad_short$shaded_apar + rad_short$shaded_anir +
    rad_long$floor * (1.0 - exp(-Kd*lai)) / (Kd*lai) +
    rad_long$air * (1.0 - exp(-Kd*lai)) / (Kd*lai) +
    -2.0 * rad_long$canopy * (1.0 - exp(-Kd*lai)) / (Kd*lai)

  Rnet_floor <- (1.0 - rho_s) * (toc_rad$par_D * exp(-Kb*sqrt(alpha_par)*lai)+
                                  toc_rad$par_d * exp(-Kd*sqrt(alpha_par)*lai)+
                                  toc_rad$nir_D * exp(-Kb*sqrt(alpha_nir)*lai)+
                                  toc_rad$nir_d * exp(-Kd*sqrt(alpha_nir)*lai))+
                                  rad_long$air * exp(-Kd*lai)+
                                  rad_long$canopy * (1.0-exp(-Kd*lai))-
                                  rad_long$floor

  Rnet_canopy <- sunlit_lai*Rnet_sunlit + shaded_lai*Rnet_shaded
  Rnet_stand <- Rnet_canopy + Rnet_floor

  out <- data.frame(Rnet_sunlit = Rnet_sunlit,
                    Rnet_shaded = Rnet_shaded,
                    Rnet_floor = Rnet_floor,
                    Rnet_canopy = Rnet_canopy,
                    Rnet_stand = Rnet_stand)
  return(out)

}

compute_longwave_rad <- function(Tair, Tsoil, vpd, cloud){
  # Compute emitted longwave radiation from the canopy, floor, and air

  N = cloud
  e_s = 610.7 * exp(17.38 * Tair / ( 239.0 + Tair))
  e_a = (e_s-vpd)/100.0

  epsilon_air = 1.24 * (e_a/(Tair+273.13))^(1.0/7.0)
  epsilon_air =(1.0-0.84*N)*epsilon_air + 0.84*N

  Ln_air = epsilon_air*(5.6704e-8) * (Tair+273.13)^4.0
  Ln_canopy = 1.0*(5.6704e-8) * (Tair+273.13)^4.0
  Ln_floor = 0.97*(5.6704e-8) * (Tsoil+273.13)^4.0

  emitted_rad <- data.frame(air=Ln_air, canopy=Ln_canopy, floor=Ln_floor, air_emissivity=epsilon_air)
  return(emitted_rad)

}

compute_potential_exfiltration <- function(Sr, depth_s, K_sat, Ksat_decay, psi_air_entry, pore_size_index, pore_decay, pore0){
  # Estimate mean porosity
  porosity_average <- pore0 * pore_decay * (1.0-exp(-1.0*depth_s/pore_decay))

  # Estimate mean saturated conductivity
  if (Ksat_decay > 0){
    Ksat_average <- Ksat_decay* Ksat *(1.0-exp(-1.0*depth_s/Ksat_decay))
  } else Ksat_average <- Ksat

  S <- max(0,min(Sr,1))

  # Plug everything into the equation for max infiltration
  potential_exfiltration <- S ^ ((1.0 / (2.0*pore_size_index))+2.0) *
                            sqrt((8.0 * porosity_average *
                            Ksat_average * psi_air_entry) /
                            (3.0 * (1.0 + 3.0 * pore_size_index) *
                            (1.0 + 4.0 * pore_size_index)))

  potential_exfiltration = min(0.001, potential_exfiltration)
  pe_wm2 = potential_exfiltration * 1e6 * 597.3 * 4.18 / (24.0*3600.0)
  return(pe_wm2)
}

compute_psn_parameters <- function(t, Kb, L_sunlit, L_shaded){
  # Local static variables
  Kc25 <- 404.0  # (ubar) MM const carboxylase, 25 deg C
  q10Kc <- 2.1 # (DIM) Q_10 for kc
  Ko25 <- 248.0 # (mbar) MM const oxygenase, 25 deg C
  q10Ko <- 1.2 # (DIM) Q_10 for ko
  act25 <- 3.6 # (umol/mgRubisco/min) Rubisco activity
  q10act <- 2.4 # (DIM) Q_10 for Rubisco activity
  Kn <- 0.52 # extinction coefficient for Nitrogen distribution
  cica <- 0.67 # Ci to Ca ratio
  alpha <- 0.10 # quantum yield efficiency

  L <- L_sunlit + L_shaded

  # Calculate atmospheric O2 in Pa, assumes 21% O2 by volume
  O2 <- 0.21 * PA

  # Correct kinetic constants for temperature, and do unit conversions
  Ko <- Ko25 * q10Ko^((t-25.0)/10.0)
  Ko <- Ko * 100.0 # mbar --> Pa

  Kc <- Kc25 * (1.8*q10Kc)^((t-25.0)/10.0) / q10Kc
  Kc[t>15.0] <- Kc25 * q10Kc^((t[t>15.0]-25.0)/10.0)
  act <- act25 * (1.8*q10act)^((t-25.0)/10.0) / q10act
  act[t>15.0] <- act25 * q10act^((t[t>15.0]-25.0)/10.0)

  Kc <- Kc * 0.10 # ubar --> Pa
  act <- act * 1e6/60.0 # umol/mg/min --> umol/kg/s

  # Calculate gamma (Pa), assumes Vomax/Vcmax = 0.21
  gamma <- 0.5 * 0.21 * Kc * O2 / Ko

  # Calculate Vmax from leaf nitrogen data and Rubisco activity
  Vmax25 <- 59.0 # umol/m2leaf/s at the top of the canopy (Lai et al, 2002, PCE, 25:1095-1119)

  Vmax25_canopy <- L * Vmax25 * (1.0 - exp(-L*Kn)) / (Kn*L)
  Vmax25_sunlit <- L * Vmax25 * (1.0 - exp(-Kn-Kb*L)) / (Kn+Kb*L)
  Vmax25_shaded <- (Vmax25_canopy - Vmax25_sunlit) / L_shaded
  Vmax25_sunlit <- Vmax25_sunlit / L_sunlit

  Vmax_sunlit <- Vmax25_sunlit * exp(0.051*(t-25.0))/(1.0+exp(0.205*(t-41.0)))
  Vmax_shaded <- Vmax25_shaded * exp(0.051*(t-25.0))/(1.0+exp(0.205*(t-41.0)))

  Rd_sunlit <- 0.015 * Vmax_sunlit
  Rd_shaded <- 0.015 * Vmax_shaded

  ppe <- 1.0 / (alpha*(4.0*cica*CA+2.0*gamma)/(cica*CA-gamma))

  psn_para <- data.frame(gamma=gamma,
                   Vmax_sunlit=Vmax_sunlit,
                   Vmax_shaded=Vmax_shaded,
                   Ko=Ko,
                   Kc=Kc,
                   O2=O2,
                   Rd_sunlit=Rd_sunlit,
                   Rd_shaded=Rd_shaded,
                   ppe=ppe)

  return(psn_para)
}

compute_soil_surface_resistance <- function(theta){
  gs <- 1.0 / (-83000.0 * theta + 16100.0)
  gs[theta > 0.185] <- 0.001429

  rs <- 1.0 / gs
  return(rs)
}

compute_sun_angles <- function(lat, lon, yr, mo, dt, hr, mi){
  solar_time <- hr + mi/60.0 + (lon-STD_MERAD)/15.0
  hangle <- (12.0-solar_time)*15.0*pi/180.0

  # Compute Julian day
  jday <- as.numeric(mapply(date2doy, yr=yr, mo=mo, dy=dt))

  # Sun declination angle
  declangle <- 23.45 * sin(2.0*pi*(284.0+jday)/365.0) * pi/180.0

  # sin(sun_elevation)
  sinh <- sin(lat*pi/180.0)*sin(declangle)+cos(lat*pi/180.0)*cos(declangle)*cos(hangle)

  zenith <- acos(sinh)
  azimuth <- asin(cos(declangle)*sin(hangle)/cos(pi/2.0-zenith))

  out <- data.frame(jday = jday, zenith = zenith, sun_decl = declangle, hourangle = hangle)
  return(out)
}

compute_toc_down_rad <- function(par, sun_zenith, jday){
  tau_t <- (par/4.55)*(1.0/PAR2SWIR)/(SO*(1.0+0.033*cos(2.0*pi*(jday-10.0)/365.0)))

  # Constants from Weiss & Norman 1985, Eq. 12
  A <- 0.9
  B <- 0.7
  C <- 0.88
  D <- 0.68

  # Cosine of the solar zenith angle
  cos_theta <- cos(sun_zenith)
  cos_theta[cos_theta<0] <- 0

  # Potential visible (V) and NIR (N) radiation at top of atmosphere
  S_V <- 0.473*1367.0*(1.0+0.033*cos(2.0*pi*(jday-10)/365))
  S_N <- 0.527*1367.0*(1.0+0.033*cos(2.0*pi*(jday-10)/365))

  # Potential direct visible on a horizontal surface
  R_DV <- S_V*exp(-0.185/cos_theta)*cos_theta

  # Potential diffuse visible on a horizontal surface
  R_dV <- 0.4*(S_V-R_DV)*cos_theta

  # Total potential visible radiation on a horizontal surface
  R_V <- R_DV + R_dV

  # Absorbed NIR by atmospheric water vapor
  cos_theta_noZero <- cos_theta
  cos_theta_noZero[cos_theta==0] <- 0.0000000000000001 # avoid division by zero
  R_aN <- S_N * 10^(-1.195+0.4459*log10(1.0/cos_theta_noZero)-0.0345*(log10(1.0/cos_theta_noZero)^2.0))

  # Potential direct NIR on a horizontal surface
  R_DN <- (S_N*exp(-0.06/cos_theta)-R_aN)*cos_theta
  R_DN[R_DN<0] <- 0 # Can't have negative direct radiation

  # Potetial diffuse NIR on a horizontal surface
  R_dN <- 0.6*(S_N-R_DN-R_aN)*cos_theta

  # Total potential NIR on a horizontal surface
  R_N <- R_DN + R_dN

  # Ratio of actual total radiation to potential total radiation on the ground
  RATIO <- (par/4.55)*(1.0/PAR2SWIR) / (R_V+R_N)

  # Actual fraction of direct visible radiation
  RATIO[RATIO>A] <- A
  f_DV <- (R_DV/(R_V*0.928)) * (1.0 - ((A-RATIO)/B) ^ (2.0/3.0))
  f_DV[R_V==0] <- 0

  # Actual fraction of direct NIR radiation
  RATIO[RATIO>C] <- C
  f_DN <- (R_DN/R_N) * (1.0 - ((C-RATIO)/D) ^ (2.0/3.0))
  f_DN[R_N==0] <- 0

  rad <- data.frame(par_D = f_DV * par/4.55,
                    par_d = par/4.55 - f_DV * par/4.55,
                    nir_D = f_DN * (par/4.55) * ((1-PAR2SWIR)/PAR2SWIR),
                    nir_d = (par/4.55) * ((1-PAR2SWIR)/PAR2SWIR) - f_DN * (par/4.55) * ((1-PAR2SWIR)/PAR2SWIR),
                    cloud = max(1.0-tau_t/0.7, 0))

  return(rad)

}

compute_toc_wind <- function(u_h, h, z){
  # Returns wind speed at height z in stratum
  #   from RHESSys, which was from BGC 4.11.

  # Compute winds
  d_o <- 0.63 * z
  z_o <- 0.1 * z

  # Wind speed at toc log decrease from screen height to toc
  u_z <- u_h * log((z-d_o)/z_o) / log((h-d_o)/z_o)
  return(u_z)
}

penman_monteith <- function(Tair, Rnet, gs, ga, vpda){
  # Calculate forest canopy ET using penman-monteith equation

    # dt:     delta t, a small temperature increase
    # rho:    air density (kg/m2)
    # lhvap:  latent heat of water (J/kg)
    # s:      Slope of saturated vapor pressure - temp
    # es:     Saturated vapor pressure at t2 in Pa
    # ET:     Evapotranspiration in W/m2
    # gamma:  Psychrometric constant
    # z0:     Surface roughness
    # d0:     Zero plane displacement

  # Density of air (rho) as a f'n of air temp
  rho <- 1.292 - (0.00428*Tair)

  # Latent heat of vaporization as a f'n of air temp
  lhvap = 2.5023e6 - 2430.54*Tair

  # Saturation vapor pressures at Tair (Pa)
  es <- 610.7 * exp(17.38 * Tair / ( 239.0 + Tair))

  # Slope of es-T curve at Tair (Pa/deg C) (from Campbell & Norman 1998)
  s <- 17.38*239.0*es / (239.0+Tair)^2

  # Calculate gamma
  gamma <- CP * PA / ( 0.622*lhvap )

  # Evaporation in W/m2
  ET <- ((s*Rnet) + (rho*CP*vpda*ga)) / (gamma*(1.0 + ga/gs) +s)
}


#' convert_P.r
#' For phosphorus, but general two step linear converter.
#' Lets say you want to convert a value, A, to a different, value, C.
#' You have data relating A to B, and B to C.
#' This will take those datasets, and a vector of A values you want to convert.
#' It will return a list, where "mean" is the mean predicted value of c, given a, and "sd" is the predictions standard deviation.
#'
#' @param to_convert vector of values to convert from a->c
#' @param d1         dataframe that relates a to b, a in the first column, b in the second.
#' @param d2         dataframe that relates b to c, b in the first column, c in the second.
#' @param n.sim      number of simulations to generate mean and sd from. 1000 is default and fine.
#'
#' @return           returns a list where mean is the mean predicted value of c, given a, and sd is the predictions standard deviation.
#' @export
#'
#' @examples
#' #phosphorus data converter.
#' a <- runif(100, 0, 1)
#' b <- a*1.2 + rnorm(length(a), sd = 0.1)
#' c <- b*0.9 + rnorm(length(b), sd = 0.1)
#' d1 <- data.frame(a,b)
#' d2 <- data.frame(b,c)
#' to_convert <- runif(25, 0, 1)
#' convert_P(to_convert,d1,d2)
convert_P <- function(to_convert,d1,d2,n.sim=1000, log10_flip = F){
  #setup predictors for JAGS.----
  y1 <- d1[,2]
  x1 <- d1[,1]
  x1 <- cbind(rep(1,length(x1)),x1)
  y2 <- d2[,2]
  x2 <- d2[,1]
  x2 <- cbind(rep(1,length(x2)),x2)

  #JAGS data object.
  jd <- list(N = length(y1), N.pred = ncol(x1),
             y1=y1, y2=y2, x1=x1, x2=x2)

  #Define JAGS model.----
  jags.model = "
  model{
  #covariate parameter priors.
  for(k in 1:N.pred){
  m1[k] ~ dnorm(0,0.001)
  m2[k] ~ dnorm(0,0.001)
  }
  #data model precision priors.
  tau1 <- pow(sigma1, -2)
  sigma1 ~ dunif(0, 100)
  tau2 <- pow(sigma2, -2)
  sigma2 ~ dunif(0, 100)

  #regression models.
  for(i in 1:N){
  mu1[i] <- inprod(x1[i,], m1)
  mu2[i] <- inprod(x2[i,], m2)
   y1[i] ~ dnorm(mu1[i], tau1)
   y2[i] ~ dnorm(mu2[i], tau2)
  }
  } #end JAGS model loop.
  "

  #Fit JAGS model.----
  cat('Fitting JAGS model...\n')
  suppressWarnings(
    jags.out <- run.jags(jags.model,
                         data=jd,
                         n.chains=3,
                         adapt = 500,
                         burnin = 1000,
                         sample = 1000,
                         monitor=c('m1','m2'),
                         method = 'rjags')
  )
  jout <- summary(jags.out)

  #grab mcmc of model parameters.
  mcmc <- do.call(rbind, jags.out$mcmc)

  #convert to_convert data from a to c, propagating uncertainty!----
  convert.out <- list()
  for(i in 1:n.sim){
    par <- mcmc[sample(nrow(mcmc),1),]
    b <- par[2]*to_convert + par[1]
    c <- par[4]*b + par[3]
    convert.out[[i]] <- c
  }

  #wrap up output to return.----
  convert.out <- do.call(cbind, convert.out)
  #log10 variable?
  suppressWarnings(
    if(log10_flip == T){
      convert.out <- log10(convert.out)
    }
  )
  mu <- rowMeans(convert.out, na.rm = T)
  mu.sd <- apply(convert.out, 1, sd, na.rm = T)
  output <- list(mu, mu.sd)
  names(output) <- c('mean','sd')

  #return output, end function.----
  return(output)
} #end function.


################################################################################
# Part of the R/EpiILM package
#
# AUTHORS:
#         Vineetha Warriyar. K. V. <vineethawarriyar.kod@ucalgary.ca>,
#         Waleed Almutiry <wkmtierie@qu.edu.sa>, and
#         Rob Deardon <robert.deardon@ucalgary.ca>
#
#
# Algorithm based on:
#                Deardon R, Brooks, S. P., Grenfell, B. T., Keeling, M. J., Tildesley,
#                M. J., Savill, N. J., Shaw, D. J.,  Woolhouse, M. E. (2010).
#                Inference for individual level models of infectious diseases in large
#                populations. Statistica Sinica, 20, 239-261.
#
#
# Free software under the terms of the GNU General Public License, version 2,
# a copy of which is available at http://www.r-project.org/Licenses/.
################################################################################

epilike <- function(object, tmin = NULL, tmax, sus.par, trans.par = NULL,
                    beta = NULL, spark = NULL, Sformula = NULL, Tformula = NULL) {

if (!is(object, "epidata")) {
   stop("The object must be in a class of \"epidata\"", call. = FALSE)
} else {

  # Error checks for input arguments
  if (any(is.null(object$type) | !(object$type %in% c("SI", "SIR"))) == TRUE) {
    stop("epilike: Specify type as \"SI\" or \"SIR\" ", call. = FALSE)
  }

  ns <- length(sus.par)
  if (!is.null(trans.par)){
    nt <- length(trans.par)
    flag.trans <- 1
    phi <- trans.par
  } else if (is.null(trans.par)) {
    nt <- 1
    flag.trans <- 0
    phi <- 1
    trans.par <- 1
  }

  if (!is.vector(object$inftime)) {
    stop('epilike: inftime is not a vector')
  }
  n <- length(object$inftime)

  if (all(is.null(object$contact) &  is.null(object$XYcoordinates)) == TRUE) {
      stop('epilike: Specify contact network or x, y coordinates')
  }

  if (all(is.null(object$remtime) & object$type == "SIR") == TRUE) {
    stop(' epilike: Specify removal times')
  }

  if(object$type == "SIR") {
      infperiod <- object$remtime - object$inftime
  }

  if (is.null(tmin)) {
    tmin <- 1
  }
  if (is.null(spark)) {
    spark <- 0
  }

  if (is.null(object$contact)) {
    ni <- length(beta)
    if (is.null(beta)) {
      stop("epilike: A scalar value of the spatial parameter beta must be specified", call. = FALSE)
    }
    if (ni != 1) {
      stop("epilike: The input of beta has more than one value while the considered distance-based ILM needs only one spatial parameter, beta.", call. = FALSE)
    }
  } else {
    if (length(object$contact)/(n * n) == 1) {
      ni <- 1
      if (is.null(beta)) {
        beta <- 1
      } else {
        stop("epilike: As the model has only one contact network, The model does not need a network parameter beta, beta must be assigned to its default values = NULL", call. = FALSE)
      }
    } else if (length(object$contact)/(n * n) > 1) {
      ni <- length(beta)
      if (is.null(beta)) {
        stop("epilike: A vector values of the network parameters beta must be specified", call. = FALSE)
      }
      if (length(object$contact)/(n * n) != ni) {
        stop('epilike: Dimension of beta and the number of contact networks are not matching')
      }
    }
    network <- array(object$contact, c(n, n, ni))
  }

  # formula for susceptibility function
  	if (!is.null(Sformula)) {
  		covmat.sus <- model.matrix(Sformula)

  		if (all((ncol(covmat.sus) == length(all.vars(Sformula))) & (ns != length(all.vars(Sformula)))) == TRUE) {
  			stop("epilike: Check Sformula (no intercept term) and the dimension of susceptibility parameters", call. = FALSE)
  		} else if (all((ncol(covmat.sus) > length(all.vars(Sformula))) & (ns != ncol(covmat.sus))) == TRUE) {
  			stop("epilike: Check Sformula (intercept term) and the dimension of susceptibility parameters", call. = FALSE)
  		}
  	} else {
  		if (ns == 1) {
  			covmat.sus <- matrix(1.0, nrow = n, ncol = ns)
  		}
  	}

    # formula for transmissibility function
    if (flag.trans == 1) {
      if (is.null(Tformula)) {
        stop("epilike: Tformula is missing. It has to be specified with no intercept term and number of columns equal to the length of trans.par", call. = FALSE)
      } else if (!is.null(Tformula)) {
    		covmat.trans <- model.matrix(Tformula)

    		if (all((ncol(covmat.trans) == length(all.vars(Tformula))) & (nt != length(all.vars(Tformula)))) == TRUE) {
    			stop("epilike: Check Tformula. It has to be with no intercept term and number of columns equal to the length of trans.par", call. = FALSE)
    		}
      }
    } else if (flag.trans == 0) {
    	covmat.trans <- matrix(1.0, nrow = n, ncol = nt)
    }

# Calling fortran subroutines for Purely Spatial models : SI and SIR

  if (all((object$type == "SI") & is.null(object$contact)) == TRUE) {

    tmp1 <- .Fortran("like",
                   x = as.vector(object$XYcoordinates[,1], mode = "double"),
                   y = as.vector(object$XYcoordinates[,2], mode = "double"),
                   tau = as.vector(object$inftime, mode = "integer"),
                   n = as.integer(n),
                   tmin = as.integer(tmin),
                   tmax = as.integer(tmax),
                   ns = as.integer(ns),
                   nt = as.integer(nt),
                   ni = as.integer(ni),
                   alpha = as.vector(sus.par, mode = "double"),
                   phi = as.vector(trans.par, mode = "double"),
                   beta = as.vector(beta, mode = "double"),
                   spark = as.double(spark),
                   covmatsus = matrix(as.double(covmat.sus), ncol = ncol(covmat.sus), nrow = n),
                   covmattrans = matrix(as.double(covmat.trans), ncol = ncol(covmat.trans), nrow = n),
                   val = as.double(0))

  } else if (all((object$type == "SIR") & is.null(object$contact)) == TRUE) {

    tmp1 <- .Fortran("likesir",
                    x = as.vector(object$XYcoordinates[,1], mode = "double"),
                    y = as.vector(object$XYcoordinates[,2], mode = "double"),
                    tau = as.vector(object$inftime, mode = "integer"),
                    lambda = as.vector(infperiod, mode = "integer"),
                    n = as.integer(n),
                    tmin = as.integer(tmin),
                    tmax = as.integer(tmax),
                    ns = as.integer(ns),
                    nt = as.integer(nt),
                    ni = as.integer(ni),
                    alpha = as.vector(sus.par, mode = "double"),
                    phi = as.vector(trans.par, mode = "double"),
                    beta = as.vector(beta, mode = "double"),
                    spark = as.double(spark),
                    covmatsus = matrix(as.double(covmat.sus), ncol = ncol(covmat.sus), nrow = n),
                    covmattrans = matrix(as.double(covmat.trans), ncol = ncol(covmat.trans), nrow = n),
                    val = as.double(0))

  } else if (all((object$type == "SI") & !is.null(object$contact)) == TRUE) {
    tmp1 <- .Fortran("likecon",
                    tau = as.vector(object$inftime, mode = "integer"),
                    n = as.integer(n),
                    ns = as.integer(ns),
                    nt = as.integer(nt),
                    ni = as.integer(ni),
                    tmin = as.integer(tmin),
                    tmax = as.integer(tmax),
                    alpha = as.vector(sus.par, mode = "double"),
                    phi = as.vector(trans.par, mode = "double"),
                    beta = as.vector(beta, mode = "double"),
                    spark = as.double(spark),
                    covmatsus = matrix(as.double(covmat.sus), ncol = ncol(covmat.sus), nrow = n),
                    covmattrans = matrix(as.double(covmat.trans), ncol = ncol(covmat.trans), nrow = n),
                    network = as.vector(network),
                    val = as.double(0))

  } else if (all((object$type == "SIR") & !is.null(object$contact)) == TRUE) {

    # Calling fortran subroutines for Contact network models: SI and SIR

    if (length(object$contact)/(n * n) != ni) {
      stop('epilike:  Dimension of beta  and the number of contact networks are not matching')
    }
    network <- array(object$contact, c(n, n, ni))

    tmp1 <- .Fortran("likeconsir",
                      tau = as.vector(object$inftime, mode = "integer"),
                      lambda = as.vector(infperiod, mode = "integer"),
                      n = as.integer(n),
                      ns = as.integer(ns),
                      nt = as.integer(nt),
                      ni = as.integer(ni),
                      tmin = as.integer(tmin),
                      tmax = as.integer(tmax),
                      alpha = as.vector(sus.par, mode = "double"),
                      phi = as.vector(trans.par, mode = "double"),
                      beta = as.vector(beta, mode = "double"),
                      spark = as.double(spark),
                      covmatsus = matrix(as.double(covmat.sus), ncol = ncol(covmat.sus), nrow = n),
                      covmattrans = matrix(as.double(covmat.trans), ncol = ncol(covmat.trans), nrow = n),
                      network = as.vector(network),
                      val = as.double(0))
  }
  return(tmp1$val)
}# End of function
}


#' Multivariate metropolis (i.e., metropolis within gibbs) sampler
#' @param target Target function (returning log unnormalised posterior density);
#' 	this function should take the parameter as its first argument and a data list as its second
#' @param initial Initial parameter vector
#' @param data Data to pass to the target
#' @param iter Number of iterations
#' @param scale Scale vector for the proposal distributions, if NA, the algorithm will adaptively
#' 		choose scales
#' @param n_adapt Number of iterations for adaptation, if needed
#' @param ... Additional named parameters to pass to m_adapt
#'
#' @return A list, with three components: 'chain' is the markov chain, 'scale'
#' 		is the scale parameter used, and 'accept' is the acceptance rate
metropolis = function(target, initial, data, iter = 1000, scale = NA, ...) {

	# check that the target distribution is defined at the initial value
	if(!is.finite(target(initial, data)))
		stop("Target distribution undefined at the initial value, choose better inits")

	# adaptation, if needed
	if(any(is.na(scale)))
		scale = m_adapt(target, initial, data, ...)

	# do the sampling
	result = m_sample(target, initial, data, scale, iter)

	# save the scale as well, because this allows resuming the chain
	result$scale = scale
	return(result)
}

#' Worker function to actually do the sampling
#' @param target Target function (returning log unnormalised posterior density);
#' 	this function should take the parameter as its first argument and a data list as its second
#' @param state Initial state vector (parameter values)
#' @param data Data to pass to the target
#' @param scale Scale for the proposal distributions
#' @param iter Number of iterations
m_sample = function(target, state, data, scale, iter) {
	# number of parameters
	k = length(state)

	# set up the markov chain
	chain = matrix(NA, nrow=iter, ncol=k)
	colnames(chain) = names(state)

	# keep track of how many times we accept the proposals
	# the acceptance rate is an important diagnostic
	accept = rep(0, k)

	for(i in 1:(iter)) {
		# for multivariate problems, it's helpful to keep track of the previous state as well
		prev_state = state

		# for each iteration, we need to randomize the order of parameters
		jj = sample(k)

		# loop over parameters one at a time
		for(j in jj) {
			# propose a value and accept/reject, updating the current state with the new value
			state = m_propose(state, scale, target, data, j)
			if(state[j] != prev_state[j]) {
				# if the value changed, that was an acceptance, increment accept variable
				accept[j] = accept[j] + 1
			}
		}
		# after we have proposed for each parameter, update the chain
		chain[i, ] = state
	}
	return(list(chain = chain, accept = accept/iter))
}

#' Runs the adaptation phase of the metropolis sampler
#' @param target Target function
#' @param state Initial vector of k parameter values
#' @param data Data to pass to the target
#' @param chunk_size Number of iterations in each chunk before adapting
#' @param max_chunk Maximum number of chunks before stopping adaptation
#' 		or divide if the acceptance rate is too small; should be a vector of length k
#' @param accept_range Target range of acceptance rates
#' @param scale Optional initial guess for the scale vector, default 1
#' @param scale_step Optional, how large the adaptation steps are, default 1.1
#' @return Scale parameter after adaptation
m_adapt = function(target, state, data, chunk_size = 100, max_chunk = 100,
			accept_range = c(0.2, 0.4), scale, scale_step) {
	k = length(state)
	accept_rate = rep(0, k)
	if(missing(scale))
		scale = rep(1, k)
	if(missing(scale_step)) {
		scale_step = rep(1.1, k)
	} else if(any(scale_step <= 1)) {
		stop("Scale step must be greater than one")
	}
	nchunk = 0
	while(nchunk < max_chunk & (any(accept_rate < min(accept_range)) |
							any(accept_rate > max(accept_range)))) {
		chunk = m_sample(target, state, data, scale, iter = chunk_size)
		accept_rate = chunk$accept
		state = chunk$chain[nrow(chunk$chain),]

		for(j in 1:k) {
			if(accept_rate[j] < min(accept_range)) {
				scale[j] = scale[j] / scale_step[j]
			} else if(accept_rate[j] > min(accept_range)) {
				scale[j] = scale[j] * scale_step[j]
			}
		}
		nchunk = nchunk + 1
	}
	cat("Finished adaptation after", nchunk, "chunks of", chunk_size,
		"iterations; acceptance rate:\n", "   ", accept_rate, "\n")
	return(scale)
}

#' Propose and select candidate values for the metropolis-hastings algorithm
#' @param state Current state of the chain
#' @param scale Scale of the proposal distribution
#' @param target Target distribution density function
#' @param data Data for the target distribution
#' @param j Index of the state/scale vectors to propose, others will be held constant
#' @return New value in the chain
m_propose <- function(state, scale, target, data, j) {
	# choose a candidate value from the proposal distribution (rnorm)
	candidate = state
	candidate[j] = rnorm(1, state[j], scale[j])

	# compute the log posterior density at the current state and the candidate value
	ld_cand = target(candidate, data)
	ld_current = target(state, data)

	# compute the acceptance probability r, making sure the density is finite
	# r is the ratio of densities: (d_cand / d_current)
	# because r should be the exponential of the difference of log densities
	if(!is.finite(ld_cand)) {
		r = 0
	} else {
		r = exp(ld_cand - ld_current)
	}

	# now choose a value such that the chance of choosing d_cand is equal to r
	# the easiest way is to draw a random uniform number U on (0,1) and compare it to r
	# pr(U < r) = r (acceptance), pr(U >= r) = 1 - r (rejection)
	# could also use sample() to do this
	U = runif(1)
	if(U < r) {
		val = candidate
	} else {
		val = state
	}
	return(val)
}

#' Computes a multivariate hpdi interval
#'
#' Assumes a unimodal posterior distribution
#' @param samples A matrix of mcmc samples
#' @param posterior A function returning the log unnormalized posterior density
#' @param data The data used to compute the samples
#' @param density The (minimum) probability density contained in the interval
hpdi = function(samples, posterior, data, density = 0.9) {

	# compute the indices of all possible intervals
	n = nrow(samples)
	n_include = ceiling(density * n)
	lower = 1:(n - n_include)
	upper = lower + n_include

	# sort the samples by their posterior density
	dens = apply(samples, 1, posterior, data = data)
	ind = order(dens)
	dens = dens[ind]
	samples = samples[ind, ]

	# compute the width for each variable, then the area in the parameter dimension
	area = mapply(function(l, u, samp)
		prod(apply(samp, 2, function(x) x[u[i]] - x[l[i]])),
		l = lower, u = upper, moreArgs = list(samp = samples))

	i = which.min(area)
	cbind(samples[lower[i], ], samples[upper[i], ])
}


## --------------------
## --------------------
## --------------------
## This file contains some basic functions for implementing Adaptive Penalized Tensor Decomposition.
## Most function are only suitable for three dimension tensor. For higher-order tensor, codes should be modified adaptively.
## --------------------
## --------------------
## --------------------

## --------------------
## inital values
## k @int
## --------------------
inits <- function(k) {
    return(t(as.matrix(rnorm(k, 0, 1))))
}

## --------------------
## tensor product
## --------------------
product <- function(X, ind, u, v, w) {
    u <- t(as.matrix(u))
    v <- t(as.matrix(v))
    w <- t(as.matrix(w))

    if (length(ind) == 2) {
        lizt <- list("mat2" = u, "mat3" = v)
        u <- ttl(X, lizt, ms = ind)
        return(as.vector(u@data))
    }

    lizt <- list("mat" = u, "mat2" = v, "mat3" = w)
    u <- ttl(X, lizt, ms = ind)
    return(as.vector(u@data))
}

## --------------------
## generate inital value for tensor decompostion if initial value is null.
## --------------------
gen_startvals <- function(num_factors, tnsr) {
    u <- matrix(inits(dim(tnsr)[1] * num_factors), nrow = num_factors)
    v <- matrix(inits(dim(tnsr)[2] * num_factors), nrow = num_factors)
    w <- matrix(inits(dim(tnsr)[3] * num_factors), nrow = num_factors)
    for (i in 1:num_factors) {
        u[i, ] <- u[i, ] / norm(as.matrix(u[i, ]), "F")
    }
    for (i in 1:num_factors) {
        v[i, ] <- v[i, ] / norm(as.matrix(v[i, ]), "F")
    }
    for (i in 1:num_factors) {
        w[i, ] <- w[i, ] / norm(as.matrix(w[i, ]), "F")
    }

    return(list(u = u, v = v, w = w))
}

## --------------------
## generate mortality tensor.
## --------------------
generate_mortality_tensor <- function(x, log = TRUE, std = TRUE) {
    x <- as.data.table(x)
    len <- ncol(x)
    circu_m <- data.matrix(x[Cause == "Circulatory system"][, (len - 18):len])
    c_m <- data.matrix(x[Cause == "Cancer"][, (len - 18):len])
    r_m <- data.matrix(x[Cause == "Respiratory system"][, (len - 18):len])
    e_m <- data.matrix(x[Cause == "External causes"][, (len - 18):len])
    i_m <- data.matrix(x[Cause == "Infectious and parasitic diseases"][, (len - 18):len])
    o_m <- data.matrix(x[Cause == "Others"][, (len - 18):len])
    ## define the tensor and apply log transformation
    if (log) {
        m_tensor <- array(NA, dim = c(6, 19, (nrow(x) / 6)))
        m_tensor[1, , ] <- t(log(circu_m))
        m_tensor[2, , ] <- t(log(c_m))
        m_tensor[3, , ] <- t(log(r_m))
        m_tensor[4, , ] <- t(log(e_m))
        m_tensor[5, , ] <- t(log(i_m))
        m_tensor[6, , ] <- t(log(o_m))
    } else {
        m_tensor <- array(NA, dim = c(6, 19, (nrow(x) / 6)))
        m_tensor[1, , ] <- t(circu_m)
        m_tensor[2, , ] <- t(c_m)
        m_tensor[3, , ] <- t(r_m)
        m_tensor[4, , ] <- t(e_m)
        m_tensor[5, , ] <- t(i_m)
        m_tensor[6, , ] <- t(o_m)
    }
    ###
    if (std) {
        tensor1 <- (m_tensor - array(mean(m_tensor), dim = dim(m_tensor))) / sd(m_tensor)
        return(as.tensor(tensor1))
    } else {
        tensor1 <- m_tensor
        return(as.tensor(tensor1))
    }
}


holes_liu_dlh <- function(space, data, sidesmall, reorder='yes') {
    ## Finds MHR (maximal empty hyper-rectangle) in a defined space
    ## Inputs should potentially be transposed and scaled if it is desired that:
    ## (1) an extent at a low value is equal to an extent at a high value
    ##     (e.g., pressure = 200, 400, 800, 1600 should be input as
    ##      log(pressure) if holes at 200 and 1600 should be considered
    ##      the same distance from 400 and 800, respectively)
    ## (2) an extent in 1 parameter is considered equal to the
    ##     same extent in another parameter (e.g., scale from 0 to 1)
    ##
    ## Input:  data      = dataframe containing datapoints
    ##                     format: 1st column = name of datapoint
    ##                             nth column = parameter values
    ##         space     = 2 row dataframe defining region of interest
    ##                     format: same as for data except name of "datapoint"
    ##                             in is SL in 1st row for the lower bound
    ##                             and   SU in 2nd row for the upper bound
    ##         sidesmall = MHR is too small to consider if
    ##                     volume < sidesmall^(# of parameters)
    ##         reorder   = reorder list of MHRs from largest to smallest
    ## Output: MHR         = sorted list of candidate MHRs from largest to smallest
    ##         MHRdropped  = sorted list of deleted because of being too small or unbounded
    ##         MHRreplaced = list of MHRs deleted because a datapoint exists inside it so not empty
    ##         volumesmall = sidesmall^(# of parameters)

    ## ## debug
    ## browser()

    ## ## debug
    ## space <- spaces
    ## data  <- dfs
    ## sidesmall <- 7
    ## reorder <- 'yes'

    ## rename data to df for use within function
    df <- data

    ## number of parameters
    nparam <- ncol(space) - 1
    end    <- nparam * 2

    ## Define first MHR h as the entire space and insert into the set of MHRs t
    ## MHR to be defined using lower and upper bounds
    ## hL <- data.frame(t(rep('SL', nparam)), stringsAsFactors=FALSE)
    ## names(hL) <- names(space[2:(nparam+1)])
    hL <- space[1, 2:(nparam+1)]
    names(hL) <- paste(names(hL), "_l", sep="")
    ## hU <- data.frame(t(rep('SU', nparam)), stringsAsFactors=FALSE)
    ## names(hU) <-  names(space[2:(nparam+1)])
    hU <- space[2, 2:(nparam+1)]
    names(hU) <- paste(names(hU), "_u", sep="")
    h <- cbind(hL,hU)
    ## add a list of additional bounding points on the MHR
    ## that may be needed to keep from inappropriately throwing out an MHR
    ## h$bounds <- list(bounds = as.character(h[1,]))
    ## h$bounds <- list(bounds = 'NA')
    t <- h

    ## define size for an MHR that is too small to consider
    volumespace <- size(hlocation(h, df, space))
    volumesmall <- sidesmall^nparam
    if (volumesmall > volumespace/4 ) {
        volumesmall <- volumespace/4
        cat('\n**** WARNING: sidesmall too big so reduced to 1/4 of space of interest ****\n')
    }

    ## initialize variables and dataframes
    hcount <- 1
    treplaced <- delete(t, t[1,], end)
    tdropped   <- treplaced

    ## For each point in database, search whether it is in t
    cat('number of datapoints', nrow(df),'\n')
    for (irow in 1:nrow(df)) {
        ##browser()
        ## if (irow >= 6) browser()
        pname <- as.character(df$name[irow])
        cat('starting datapoint', pname,'; number of MHRs =',nrow(t),'\n')
        ## print(pname)
        ## identify list of MHRs (each is an h) in the set of MHRs t that contain x
        ## for loop will modify t so need a temporary dataframe tprev with the previous values of t
        ## that the for loop can work through
        tprev <- t
        for ( hi in 1:nrow(tprev) ) {
            ## cat('irow=',irow,'; hi=',hi,'; jdim=',jdim,'\n')
            ## if (irow==12 & hi==130 & jdim==2) browser()
            h <- tprev[hi,]
            testloc <- ContainmentSearch(h, pname, df, space)$location
            if (testloc == 'on boundary') {
                ## insert x into set of bounding points in that surface
                ## t[hi,]$bounds[[1]] <- append(h$bounds[[1]], pname)
                ## ignore because point on boundary does not cause MHR to subdivide

            } else if (testloc == 'inside') {

                ## delete MHR h from collection of MHRs t
                t     <- delete(t, h, end)

                ## add MHR h to list of deleted MHRs for later tracing if needed
                treplaced <- insert(treplaced, h)

                ## break MHR h into 2 separate smaller MHRs along each parameter
                ## add the two new MHRs, ha and hb, to the set of MHRs t
                ## for each dimension
                for (jdim in 1:nparam) {
                    ## print(pname)
                    ## print(jdim)
                    ## start with the MHR of interest
                    ha <- h
                    hb <- h
                    ## print(h)

                    ## split h in each dimension, jdim
                    ## first assign location to the split location
                    datapoint <- df[df$name == pname,]
                    ha[jdim+nparam] <- datapoint[1 + jdim]
                    hb[jdim]        <- datapoint[1 + jdim]

                    ## identify name of the parent MHR
                    parent <- stri_split_fixed(str = row.names(h), pattern = "_", n = 2)[[1]][1]

                    ## add new MHR ha if big enough and not completely contained in an existing MHR
                    haloc  <- hlocation(ha, df, space)
                    sizeha <- size(haloc)
                    ## if (pname == 'point3' & jdim == 2) browser()
                    ## name the MHR
                    hcount <- hcount+1
                    row.names(ha) <- paste(hcount, '_jdim', jdim, "L", "_parent_", parent, sep="")
                    ## add MHR to t if big enough and drop if too small
                    if ( sizeha >= volumesmall ) {
                        if ( nrow(t) == 0 ) {
                            ## t is empty so add ha
                            t <- rbind(t, ha)
                        } else if (  inexisting(t, haloc) == 'no' ) {
                            ## add ha because it is not fully contained inside another MHR
                            t <- rbind(t, ha)
                        }
                    } else {
                        ## drop the MHR
                        tdropped <- rbind(tdropped, ha)
                    }

                    ## add new MHR ha if big enough and not completely contained in an existing MHR
                    hbloc  <- hlocation(hb, df, space)
                    sizehb <- size(hbloc)
                    hcount <- hcount+1
                    row.names(hb) <- paste(hcount, '_jdim', jdim, "U", "_parent_", parent, sep="")
                    if ( sizehb >= volumesmall ) {
                        if ( nrow(t) == 0 ) {
                            t <- rbind(t, hb)
                        } else if (  inexisting(t, hbloc) == 'no' ) {
                            t <- rbind(t, hb)
                        }
                    } else {
                        tdropped <- rbind(tdropped, hb)
                    }

                }
            }
            ## else testloc == 'outside' so look at the next MHR in t
        }
    }

    ## identify location of each side of MHRs
    ## strip off the boundary column for later lapply function
    tloc <- t[1:end]
    for (i in 1:nrow(tloc)) {
        ## tloc <- apply(t, 1, function(h) hlocation(h, df, space))
        h     <- tloc[i,]
        tloc[i,] <- as.numeric( hlocation(h, df, space) )
    }
    tloc
    ## tloc was still all character for some reason so need to convert to numeric
    tloc[] <- lapply(tloc, as.numeric)

    ## add column of MHR sizes
    tloc$size <- size(tloc)
    t$size    <- tloc$size

    ## reorder by size
    if (reorder == 'yes') {
        tloc     <- tloc    [order(tloc$size,     decreasing=TRUE),]
        t        <- t       [order(t$size,        decreasing=TRUE),]
    }

    ## do the same for the dataframe of MHRs that were too small
    if (nrow(tdropped) > 0) {
        tlocdropped <- tdropped
        for (i in 1:nrow(tlocdropped)) {
            h     <- tlocdropped[i,]
            tlocdropped[i,] <- as.numeric( hlocation(h, df, space) )
        }
        tlocdropped[] <- lapply(tlocdropped, as.numeric)   # tlocdropped was all character
        tdropped$size <- size(tlocdropped)
        if (reorder == 'yes') tdropped <- tdropped[order(tdropped$size, decreasing=TRUE),]
    }

    return(list(MHRloc=tloc, MHR=t, MHRdropped=tdropped,
                MHRreplaced=treplaced, volumesmall=volumesmall))

}


## --------------------
## --------------------
## --------------------
## These codes implement the MLR for short-term forecast(5 year). For mid-term, long-term or other dataset, these should be modified adaptively.
## --------------------
## --------------------
## --------------------
rm(list = ls())
library(nnet)
library(tidyverse)

library(reshape2)
library(data.table)
library(rTensor)
set.seed(082020)

## --------------------
## Data attributes
## --------------------
popm <- fread("population_male.csv")
usam <- fread("mortality rate for male with comparability ratio.csv")
colnames(popm)[4:22] <- c(
    "infant", "1-4",
    "5-9", "10-14",
    "15-19", "20-24",
    "25-29", "30-34",
    "35-39", "40-44",
    "45-49", "50-54",
    "55-59", "60-64",
    "65-69", "70-74",
    "75-79", "80-84",
    "85+"
)
colnames(usam)[4:22] <- c(
    "infant", "1-4",
    "5-9", "10-14",
    "15-19", "20-24",
    "25-29", "30-34",
    "35-39", "40-44",
    "45-49", "50-54",
    "55-59", "60-64",
    "65-69", "70-74",
    "75-79", "80-84",
    "85+"
)

###### process population data
popn <- reshape2::melt(popm[, c(2, 4:22)], id.vars = c("Year"))
colnames(popn) <- c("Year", "agegroups", "pop")
popn <- setorder(popn, "Year")
popn$Year <- popn$Year - 1949

###### process mortality rate
cir <- usam[Cause == "Circulatory system"][, c(2, 4:22)]
cir1 <- melt(cir, id.vars = "Year")
setorder(cir1, "Year", "variable", "value")
colnames(cir1) <- c("Year", "agegroups", "Circulatory system")
cir1$Year <- cir1$Year - 1949

can <- usam[Cause == "Cancer"][, c(2, 4:22)]
can1 <- melt(can, id.vars = "Year")
setorder(can1, "Year", "variable", "value")
colnames(can1) <- c("Year", "agegroups", "Cancer")
can1$Year <- can1$Year - 1949

rs <- usam[Cause == "Respiratory system"][, c(2, 4:22)]
rs1 <- melt(rs, id.vars = "Year")
setorder(rs1, "Year", "variable", "value")
colnames(rs1) <- c("Year", "agegroups", "Respiratory system")
rs1$Year <- rs1$Year - 1949

iapd <- usam[Cause == "Infectious and parasitic diseases"][, c(2, 4:22)]
iapd1 <- melt(iapd, id.vars = "Year")
setorder(iapd1, "Year", "variable", "value")
colnames(iapd1) <- c("Year", "agegroups", "Infectious and parasitic diseases")
iapd1$Year <- iapd1$Year - 1949

ec <- usam[Cause == "External causes"][, c(2, 4:22)]
ec1 <- melt(ec, id.vars = "Year")
setorder(ec1, "Year", "variable", "value")
colnames(ec1) <- c("Year", "agegroups", "External causes")
ec1$Year <- ec1$Year - 1949

oth <- usam[Cause == "Others"][, c(2, 4:22)]
oth1 <- melt(oth, id.vars = "Year")
setorder(oth1, "Year", "variable", "value")
colnames(oth1) <- c("Year", "agegroups", "Others")
oth1$Year <- oth1$Year - 1949

probs <- cbind(cir1, can1[, 3], rs1[, 3], iapd1[, 3], ec1[, 3], oth1[, 3])
sur1 <- oth1
colnames(sur1) <- c("Year", "agegroups", "Survival")
sur1$Survival <- 1 - rowSums(probs[, 3:8])
probs <- cbind(sur1, probs[, 3:8])
colnames(probs)[5:9] <- c("Cancer", "Respiratory system", "Infectious and parasitic diseases", "External causes", "Others")

##### counts
resp <- left_join(probs, popn, by = c("Year", "agegroups"))
head(resp)

resp2 <- (resp[, 3:9] * resp$pop) %>%
    round() %>%
    as.matrix()

(rowSums(resp2) - resp$pop) %>%
    summary() ## Sanity check. All of these should be very close to zero
rownames(probs) <- rownames(resp) <- rownames(resp2) <- rownames(popn) <- NULL

## --------------------
## Fit multinomial logistic regression and compare results
## --------------------
years <- (1:58) %>%
    scale(center = FALSE, scale = FALSE)
agegroups <- c(
    "infant", "1-4",
    "5-9", "10-14",
    "15-19", "20-24",
    "25-29", "30-34",
    "35-39", "40-44",
    "45-49", "50-54",
    "55-59", "60-64",
    "65-69", "70-74",
    "75-79", "80-84",
    "85+"
)
dat <- data.frame(years = rep(years, each = length(unique(resp$agegroups))), agegroups = rep(agegroups, length(unique(resp$Year))))
head(dat)

make_offset <- matrix(resp$pop, nrow = nrow(probs), ncol = ncol(resp2))
# simplest
fit_mlr1s <- multinom(resp2[1:1007, ] ~ years + agegroups + offset(make_offset[1:1007, ]), data = dat[1:1007, ], maxit = 1000)
# single
fit_mlr1 <- multinom(resp2[1:1007, ] ~ years + agegroups + years:agegroups + offset(make_offset[1:1007, ]), data = dat[1:1007, ], maxit = 1000)
# quadratic
fit_mlr2 <- multinom(resp2[1:1007, ] ~ years + agegroups + I(years):agegroups + I(years^2) + I(years^2):agegroups + offset(make_offset[1:1007, ]), data = dat[1:1007, ], maxit = 1000)
# cubic
fit_mlr3 <- multinom(resp2[1:1007, ] ~ years + agegroups + I(years^2) + I(years^2):agegroups + I(years):agegroups + I(years^3) + I(years^3):agegroups + offset(make_offset[1:1007, ]), data = dat[1:1007, ], maxit = 5000)

summary(fit_mlr1s)
summary(fit_mlr1)
summary(fit_mlr2)
summary(fit_mlr3)

coefs1s <- coef(fit_mlr1s)
coefs1 <- coef(fit_mlr1)
coefs2 <- coef(fit_mlr2)
coefs3 <- coef(fit_mlr3)

## --------------------
## --------------------
## fit
## --------------------
## --------------------
fit_dat <- data.frame(years = rep(c(1:53), each = length(unique(resp$agegroups))), agegroups = rep(agegroups, 53))
X_fit1s <- model.matrix(~ years + agegroups, data = fit_dat)
X_fit1 <- model.matrix(~ years + agegroups + years:agegroups, data = fit_dat)
X_fit2 <- model.matrix(~ years + agegroups + I(years):agegroups + I(years^2) + I(years^2):agegroups, data = fit_dat)
X_fit3 <- model.matrix(~ years + agegroups + I(years^2) + I(years^2):agegroups + I(years):agegroups + I(years^3) + I(years^3):agegroups, data = fit_dat)

## --------------------
## simplest
## --------------------
fit_probs1s <- tcrossprod(X_fit1s, coefs1s) %>%
    exp()
fit_probs1s <- cbind(1 / (1 + rowSums(fit_probs1s)), fit_probs1s / (1 + rowSums(fit_probs1s)))
colnames(fit_probs1s)[1] <- "Survival"
fit_probs1s <- cbind(fit_dat, fit_probs1s[, 2:7])
colnames(fit_probs1s)[1] <- "Year"

df_wide1s <- as.data.table(fit_probs1s) %>% data.table::dcast(Year ~ agegroups, value.var = colnames(as.data.table(fit_probs1s))[3:8])

fit1s <- df_wide1s %>%
    gather(variable, value, -Year) %>%
    mutate(
        Cause = sub("_.*", "", variable),
        variable = sub(".*_", "", variable)
    ) %>%
    spread(variable, value)

fit1s <- fit1s[, colnames(usam[, c(2:22)])]

## --------------------
## single
## --------------------
fit_probs1 <- tcrossprod(X_fit1, coefs1) %>%
    exp()
fit_probs1 <- cbind(1 / (1 + rowSums(fit_probs1)), fit_probs1 / (1 + rowSums(fit_probs1)))
colnames(fit_probs1)[1] <- "Survival"
fit_probs1 <- cbind(fit_dat, fit_probs1[, 2:7])
colnames(fit_probs1)[1] <- "Year"

df_wide1 <- as.data.table(fit_probs1) %>% data.table::dcast(Year ~ agegroups, value.var = colnames(as.data.table(fit_probs1))[3:8])

fit1 <- df_wide1 %>%
    gather(variable, value, -Year) %>%
    mutate(
        Cause = sub("_.*", "", variable),
        variable = sub(".*_", "", variable)
    ) %>%
    spread(variable, value)

fit1 <- fit1[, colnames(usam[, c(2:22)])]
##### difference for single
fnorm(generate_mortality_tensor(fit1) - generate_mortality_tensor(usam[1:(53 * 6), ]))

## --------------------
## quadratic
## --------------------
fit_probs2 <- tcrossprod(X_fit2, coefs2) %>%
    exp()
fit_probs2 <- cbind(1 / (1 + rowSums(fit_probs2)), fit_probs2 / (1 + rowSums(fit_probs2)))
colnames(fit_probs2)[1] <- "Survival"
fit_probs2 <- cbind(fit_dat, fit_probs2[, 2:7])
colnames(fit_probs2)[1] <- "Year"

df_wide2 <- as.data.table(fit_probs2) %>% data.table::dcast(Year ~ agegroups, value.var = c("Circulatory system", "Cancer", "Respiratory system", "Infectious and parasitic diseases", "External causes", "Others"))

fit2 <- df_wide2 %>%
    gather(variable, value, -Year) %>%
    mutate(
        Cause = sub("_.*", "", variable),
        variable = sub(".*_", "", variable)
    ) %>%
    spread(variable, value)

fit2 <- fit2[, colnames(usam[, c(2:22)])]
##### difference for quadratic
fnorm(generate_mortality_tensor(fit2) - generate_mortality_tensor(usam[1:(53 * 6), ]))

## --------------------
## cubic
## --------------------
fit_probs3 <- tcrossprod(X_fit3, coefs3) %>%
    exp()
fit_probs3 <- cbind(1 / (1 + rowSums(fit_probs3)), fit_probs3 / (1 + rowSums(fit_probs3)))
colnames(fit_probs3)[1] <- "Survival"
fit_probs3 <- cbind(fit_dat, fit_probs3[, 2:7])
colnames(fit_probs3)[1] <- "Year"

df_wide3 <- as.data.table(fit_probs3) %>% data.table::dcast(Year ~ agegroups, value.var = c("Circulatory system", "Cancer", "Respiratory system", "Infectious and parasitic diseases", "External causes", "Others"))

fit3 <- df_wide3 %>%
    gather(variable, value, -Year) %>%
    mutate(
        Cause = sub("_.*", "", variable),
        variable = sub(".*_", "", variable)
    ) %>%
    spread(variable, value)

fit3 <- fit3[, colnames(usam[, c(2:22)])]
##### difference for quadratic
fnorm(generate_mortality_tensor(fit3) - generate_mortality_tensor(usam[1:(53 * 6), ]))

## --------------------
## --------------------
## predict
## --------------------
## --------------------
pre_dat <- data.frame(years = rep(c(54:58), each = length(unique(resp$agegroups))), agegroups = rep(agegroups, 5))
X_pred1s <- model.matrix(~ years + agegroups, data = pre_dat)
X_pred1 <- model.matrix(~ years + agegroups + years:agegroups, data = pre_dat)
X_pred2 <- model.matrix(~ years + agegroups + I(years):agegroups + I(years^2) + I(years^2):agegroups, data = pre_dat)
X_pred3 <- model.matrix(~ years + agegroups + I(years^2) + I(years^2):agegroups + I(years):agegroups + I(years^3) + I(years^3):agegroups, data = pre_dat)

## --------------------
## simplest
## --------------------
pred_probs1s <- tcrossprod(X_pred1s, coefs1s) %>%
    exp()
pred_probs1s <- cbind(1 / (1 + rowSums(pred_probs1s)), pred_probs1s / (1 + rowSums(pred_probs1s)))
colnames(pred_probs1s)[1] <- "Survival"
pred_probs1s <- cbind(pre_dat, pred_probs1s[, 2:7])
colnames(pred_probs1s)[1] <- "Year"

df_wide1ps <- as.data.table(pred_probs1s) %>% data.table::dcast(Year ~ agegroups, value.var = c("Circulatory system", "Cancer", "Respiratory system", "Infectious and parasitic diseases", "External causes", "Others"))

pred1s <- df_wide1ps %>%
    gather(variable, value, -Year) %>%
    mutate(
        Cause = sub("_.*", "", variable),
        variable = sub(".*_", "", variable)
    ) %>%
    spread(variable, value)

pred1s <- pred1s[, colnames(usam[, c(2:22)])]
##### difference
fnorm(generate_mortality_tensor(pred1s) - generate_mortality_tensor(usam[(53 * 6 + 1):(58 * 6), ]))

## --------------------
## single
## --------------------
pred_probs1 <- tcrossprod(X_pred1, coefs1) %>%
    exp()
pred_probs1 <- cbind(1 / (1 + rowSums(pred_probs1)), pred_probs1 / (1 + rowSums(pred_probs1)))
colnames(pred_probs1)[1] <- "Survival"
pred_probs1 <- cbind(pre_dat, pred_probs1[, 2:7])
colnames(pred_probs1)[1] <- "Year"

df_wide1p <- as.data.table(pred_probs1) %>% data.table::dcast(Year ~ agegroups, value.var = c("Circulatory system", "Cancer", "Respiratory system", "Infectious and parasitic diseases", "External causes", "Others"))

pred1 <- df_wide1p %>%
    gather(variable, value, -Year) %>%
    mutate(
        Cause = sub("_.*", "", variable),
        variable = sub(".*_", "", variable)
    ) %>%
    spread(variable, value)

pred1 <- pred1[, colnames(usam[, c(2:22)])]
##### difference
fnorm(generate_mortality_tensor(pred1) - generate_mortality_tensor(usam[(53 * 6 + 1):(58 * 6), ]))

## --------------------
## quadratic
## --------------------
pred_probs2 <- tcrossprod(X_pred2, coefs2) %>%
    exp()
pred_probs2 <- cbind(1 / (1 + rowSums(pred_probs2)), pred_probs2 / (1 + rowSums(pred_probs2)))
colnames(pred_probs2)[1] <- "Survival"
pred_probs2 <- cbind(pre_dat, pred_probs2[, 2:7])
colnames(pred_probs2)[1] <- "Year"

df_wide2p <- as.data.table(pred_probs2) %>% data.table::dcast(Year ~ agegroups, value.var = c("Circulatory system", "Cancer", "Respiratory system", "Infectious and parasitic diseases", "External causes", "Others"))

pred2 <- df_wide2p %>%
    gather(variable, value, -Year) %>%
    mutate(
        Cause = sub("_.*", "", variable),
        variable = sub(".*_", "", variable)
    ) %>%
    spread(variable, value)

pred2 <- pred2[, colnames(usam[, c(2:22)])]
##### difference
fnorm(generate_mortality_tensor(pred2) - generate_mortality_tensor(usam[(53 * 6 + 1):(58 * 6), ]))

## --------------------
## cubic
## --------------------

pred_probs3 <- tcrossprod(X_pred3, coefs3) %>%
    exp()
pred_probs3 <- cbind(1 / (1 + rowSums(pred_probs3)), pred_probs3 / (1 + rowSums(pred_probs3)))
colnames(pred_probs3)[1] <- "Survival"
pred_probs3 <- cbind(pre_dat, pred_probs3[, 2:7])
colnames(pred_probs3)[1] <- "Year"

df_wide3p <- as.data.table(pred_probs3) %>% data.table::dcast(Year ~ agegroups, value.var = c("Circulatory system", "Cancer", "Respiratory system", "Infectious and parasitic diseases", "External causes", "Others"))

pred3 <- df_wide3p %>%
    gather(variable, value, -Year) %>%
    mutate(
        Cause = sub("_.*", "", variable),
        variable = sub(".*_", "", variable)
    ) %>%
    spread(variable, value)

pred3 <- pred3[, colnames(usam[, c(2:22)])]
##### difference
fnorm(generate_mortality_tensor(pred3) - generate_mortality_tensor(usam[(53 * 6 + 1):(58 * 6), ]))


# This is a primal-dual log barrier form of the
# interior point LP solver of Lustig, Marsten and Shanno ORSA J Opt 1992.
# It is a projected Newton primal-dual logarithmic barrier method which uses
# the predictor-corrector approach of Mehrotra for the mu steps.
# For the sake of brevity we will call it a Frisch-Newton algorithm.
# The primary difference between this code and the previous version fna.r is
# that we assume a feasible starting point so p,d,b feasibility gaps = 0.
# Problem:
# 	min c'x s.t. Ax=b, 0<=x<=u
#
# The linear system we are trying to solve at each interation is:
# 		 Adx = 0
# 	     dx + ds = 0
# 	A'dy -dw +dz = 0
# 	   Xdz + Zdx = me - XZe - DXDZe
# 	   Sdw + Wds = me - SWe - DSDWe
# But the algorithm proceeds in two steps the first of which is to solve:
# 		 Adx = 0
# 	     dx + ds = 0
# 	A'dy -dw +dz = 0
# 	   Xdz + Zdx =  - XZe
# 	   Sdw + Wds =  - SWe
# and then to make some refinement of mu and modify the implied Newton step.
# Denote dx,dy,dw,ds,dz as the steps for the respective variables x,y,w,s,z, and
# the corresponding upper case letters are the diagonal matrices respectively.
#
# To illustrate the use of the function we include a calling routine to
# compute solutions to the linear quantile regression estimation problem.
# See the associated S function rqfn for details on the calling sequence.
# On input:
# 	a is the p by n matrix X'
# 	y is the n-vector of responses
# 	u is the n-vector of upper bounds
#       d is an n-vector of ones
#       wn is an n-vector of ones in the first n elements
# 	beta is a scaling constant, conventionally .99995
# 	eps is a convergence tolerance, conventionally 1d-07
# On output:
# 	a,y are unaltered
# 	wp contains the solution  coefficient vector in the first p elements
# 	wn contains the residual vector in the first n elements
#
#
subroutine rqfn(n,p,a,y,rhs,d,u,beta,eps,wn,wp,aa,nit,info)
integer n,p,info,nit(3)
double precision a(p,n),y(n),rhs(p),d(n),u(n),wn(n,10),wp(p,p+3),aa(p,p)
double precision one,beta,eps
parameter( one = 1.0d0)
call fna(n,p,a,y,rhs,d,u,beta,eps,wn(1,1),wn(1,2),
	wp(1,1),wn(1,3),wn(1,4),wn(1,5), wn(1,6),
	wp(1,2),wn(1,7),wn(1,8),wn(1,9),wn(1,10),wp(1,3), wp(1,4),aa,nit,info)
return
end
subroutine fna(n,p,a,c,b,d,u,beta,eps,x,s,y,z,w,
		dx,ds,dy,dz,dw,dsdw,dxdz,rhs,ada,aa,nit,info)

integer n,p,pp,i,info,nit(3)
double precision a(p,n),c(n),b(p)
double precision zero,one,mone,big,ddot,dmax1,dmin1,dasum
double precision deltap,deltad,beta,eps,cx,by,uw,uz,mu,mua,acomp,rdg,g
double precision x(n),u(n),s(n),y(p),z(n),w(n),d(n),rhs(p),ada(p,p)
double precision aa(p,p),dx(n),ds(n),dy(p),dz(n),dw(n),dxdz(n),dsdw(n)

parameter( zero  = 0.0d0)
parameter( half  = 0.5d0)
parameter( one   = 1.0d0)
parameter( mone   = -1.0d0)
parameter( big  = 1.0d+20)

# Initialization:  We try to follow the notation of LMS
# On input we require:
#
# 	c = n-vector of marginal costs (-y in the rq problem)
# 	a = p by n matrix of linear constraints (x' in rq)
# 	b = p-vector of rhs ((1-tau)x'e in rq)
#       u = upper bound vector ( e in rq)
# 	beta = barrier parameter, LMS recommend .99995
# 	eps = convergence tolerance, LMS recommend 10d-8
#
# 	the integer vector nit returns iteration counts
# 	the integer info contains an error code from the Cholesky in stepy
# 		info = 0 is fine
# 		info < 0 invalid argument to dposv
# 		info > 0 singular matrix
nit(1)=0
nit(2)=0
nit(3)=n
pp=p*p
# Start at the OLS estimate for the parameters
call dgemv('N',p,n,one,a,p,c,1,zero,y,1)
call stepy(n,p,a,d,y,aa,info)
if(info != 0) return
# Save sqrt of aa' for future use for confidence band
do i=1,p{
	do j=1,p
		ada(i,j)=zero
        ada(i,i)=one
	}
call dtrtrs('U','T','N',p,p,aa,p,ada,p,info)
call dcopy(pp,ada,1,aa,1)
# put current residual vector in s (temporarily)
call dcopy(n,c,1,s,1)
call dgemv('T',p,n,mone,a,p,y,1,one,s,1)
# Initialize remaining variables
# N.B. x must be initialized on input: for rq as (one-tau) in call coordinates
do i=1,n{
	d(i)=one
	if(dabs(s(i)) < eps){
		z(i) = dmax1( s(i),zero) + eps
		w(i) = dmax1(-s(i),zero) + eps
		}
	else {
		z(i) = dmax1( s(i),zero)
		w(i) = dmax1(-s(i),zero)
		}
	s(i)=u(i)-x(i)
	}
cx = ddot(n,c,1,x,1)
by = ddot(p,b,1,y,1)
uw = dasum(n,w,1)
uz = dasum(n,z,1)
# rdg =  (cx - by + uw)/(one + dabs( by - uw))
# rdg =  (cx - by + uw)/(one + uz + uw)
rdg =  (cx - by + uw)
while(rdg > eps) {
	nit(1)=nit(1)+1
	do i =1,n{
		d(i) = one/(z(i)/x(i) + w(i)/s(i))
		ds(i)=z(i)-w(i)
		dx(i)=d(i)*ds(i)
		}
	call dgemv('N',p,n,one,a,p,dx,1,zero,dy,1)#rhs
	call dcopy(p,dy,1,rhs,1)#save rhs
	call stepy(n,p,a,d,dy,ada,info)
	if(info != 0) return
	call dgemv('T',p,n,one,a,p,dy,1,mone,ds,1)
	deltap=big
	deltad=big
	do i=1,n{
		dx(i)=d(i)*ds(i)
		ds(i)=-dx(i)
		dz(i)=-z(i)*(dx(i)/x(i) + one)
		dw(i)=w(i)*(dx(i)/s(i) - one)
		dxdz(i)=dx(i)*dz(i)
		dsdw(i)=ds(i)*dw(i)
		if(dx(i)<0)deltap=dmin1(deltap,-x(i)/dx(i))
		if(ds(i)<0)deltap=dmin1(deltap,-s(i)/ds(i))
		if(dz(i)<0)deltad=dmin1(deltad,-z(i)/dz(i))
		if(dw(i)<0)deltad=dmin1(deltad,-w(i)/dw(i))
		}
	deltap=dmin1(beta*deltap,one)
	deltad=dmin1(beta*deltad,one)
	if(deltap*deltad<one){
		nit(2)=nit(2)+1
		acomp=ddot(n,x,1,z,1)+ddot(n,s,1,w,1)
		g=acomp+deltap*ddot(n,dx,1,z,1)+
			deltad*ddot(n,dz,1,x,1)+
			deltap*deltad*ddot(n,dz,1,dx,1)+
			deltap*ddot(n,ds,1,w,1)+
			deltad*ddot(n,dw,1,s,1)+
			deltap*deltad*ddot(n,ds,1,dw,1)
		mu=acomp/dfloat(2*n)
		mua=g/dfloat(2*n)
		mu=mu*(mua/mu)**3
		#if(acomp>1) mu=(g/dfloat(n))*(g/acomp)**2
		#else mu=acomp/(dfloat(n)**2)
		do i=1,n{
			dz(i)=d(i)*(mu*(1/s(i)-1/x(i))+
				dx(i)*dz(i)/x(i)-ds(i)*dw(i)/s(i))
			}
		call dswap(p,rhs,1,dy,1)
		call dgemv('N',p,n,one,a,p,dz,1,one,dy,1)#new rhs
		call dpotrs('U',p,1,ada,p,dy,p,info)
		call daxpy(p,mone,dy,1,rhs,1)#rhs=ddy
		call dgemv('T',p,n,one,a,p,rhs,1,zero,dw,1)#dw=A'ddy
		deltap=big
		deltad=big
		do i=1,n{
			dx(i)=dx(i)-dz(i)-d(i)*dw(i)
			ds(i)=-dx(i)
			dz(i)=mu/x(i) - z(i)*dx(i)/x(i) - z(i) - dxdz(i)/x(i)
			dw(i)=mu/s(i) - w(i)*ds(i)/s(i) - w(i) - dsdw(i)/s(i)
			if(dx(i)<0)deltap=dmin1(deltap,-x(i)/dx(i))
			else deltap=dmin1(deltap,-s(i)/ds(i))
			if(dz(i)<0)deltad=dmin1(deltad,-z(i)/dz(i))
			if(dw(i)<0)deltad=dmin1(deltad,-w(i)/dw(i))
			}
		deltap=dmin1(beta*deltap,one)
		deltad=dmin1(beta*deltad,one)
		}
	call daxpy(n,deltap,dx,1,x,1)
	call daxpy(n,deltap,ds,1,s,1)
	call daxpy(p,deltad,dy,1,y,1)
	call daxpy(n,deltad,dz,1,z,1)
	call daxpy(n,deltad,dw,1,w,1)
	cx=ddot(n,c,1,x,1)
	by=ddot(p,b,1,y,1)
	uw = dasum(n,w,1)
	uz = dasum(n,z,1)
	#rdg=(cx-by+uw)/(one+dabs(by-uw))
	#rdg=(cx-by+uw)/(one+uz+uw)
	rdg=(cx-by+uw)
	}
# return residuals in the vector x
call daxpy(n,mone,w,1,z,1)
call dswap(n,z,1,x,1)
return
end