library(deSolve)
library(ggplot2)
library(rstan)

# ALL files should be in the same working directory. The user should set this directory under the 'path' varialbe
options(max.print=999999)
path <- "" ###### USER INPUT HERE #####
setwd(path)
load("results.RData")
stan_fit <- extract(fit)

# !!!!!!!!!  In exposure concentration, the last value will not be taken into account by the model, so it better be zero !!!!!!!!!!!!!!!!!!!
##########################################
# Function for creating parameter vector #
##########################################

create.params <- function(input, stan_fit, stochastic = TRUE){
  with( as.list(input),{

    # List with names of all possible compartments
    comps <- list("Lungs" = "Lungs" ,"Liver"="Liver", "Spleen"="Spleen", "Kidneys"="Kidneys", "Heart"="Heart",  "Brain"="Brain",
                  "Uterus"="Uterus", "Skeleton"="Skeleton", "Skin"="Skin", "Soft"="Soft") # List with names of all possible compartments

    ### Density of tissues/organs
    d_tissue <- 1 #g/ml
    d_skeleton <- 1.92 #g/ml
    d_adipose <- 0.940 #g/ml

    Q_total <- (1.54*weight^0.75)*60 # Total Cardiac Output (ml/h)

    Total_Blood <- 0.06*weight+0.77 # Total blood volume (ml)

    #Arterial blood volume
    Vart <- 0.15*Total_Blood #(ml)

    #Veins blood volume
    Vven <-0.64*Total_Blood #(ml)

    #Tissue weight fraction
    Tissue_fractions <- c(0.5, 3.66, 0.2, 0.73, 0.33, 0.57, 0.011, 10, 19.03, NA)/100 # % of BW. Na values refers to the volume of the rest organs(RoB)

    #Regional blood flow fraction
    Regional_flow_fractions <- c(100, 17.4, 1.22, 14.1, 4.9, 2, 1.1, 12.2, 5.8, NA)/100 # % of total cardiac output

    #Capillary volume fractions (fractions of tissue volume)
    Capillary_fractions <- c(0.36, 0.21, 0.22, 0.16, 0.26, 0.03, 0.04, 0.04, 0.02, 0.04) # fraction of tissue volume
    # Where NA, it is the same value as Rest of Body due to luck od data for uterus and adipose

    #Macrophage content (g per gram of tissue)
    Macrophage_fraction <- c(0.04, 0.1, 0.3, 0.02, 0.02, 0.04, 0.04, 0.04, 0.02, 0.02)
    Macrophage_fraction_blood <- 0.01

    CLE_ua <- 0
    CLE_muc <-  0.01181567
    k_ua_br <- 0
    CLE_hep <-  1.52e-03
    CLE_muc_cor <- 1
    k_ab0 <- 1.45
    k_de <- 6.7e-19
    k_ab0_spl <- 0.5

    W_tis <- rep(0,length(comps))
    V_tis <- rep(0,length(comps))
    V_cap <- rep(0,length(comps))
    W_macro <- rep(0,length(comps))
    Q <- rep(0,length(comps))

    for (i in 1:(length(comps)-1)) {
      control <- comps[i]

      Tissue_fractions[i] <- ifelse(is.na(control), NA, Tissue_fractions[i])
      Regional_flow_fractions[i] <- ifelse(is.na(control), NA, Regional_flow_fractions[i])
      Capillary_fractions[i] <- ifelse(is.na(control), NA, Capillary_fractions[i])
      Macrophage_fraction[i] <- ifelse(is.na(control), NA, Macrophage_fraction[i])

      ### Calculation of tissue weights
      W_tis[i] <- weight*Tissue_fractions[i]
      W_macro[i] <- W_tis[i]*Macrophage_fraction[i]

      ###Calculation of tissue volumes

      if (i==9){
        V_tis[i] <- W_tis[i]/d_skeleton
      } else{
        V_tis[i] <- W_tis[i]/d_tissue
      }

      ###Calculation of capillary volumes
      V_cap[i] <- V_tis[i]*Capillary_fractions[i]

      ###Calculation of regional blood flows
      Q[i] <- Q_total*Regional_flow_fractions[i]
    }

    ### Calculations for "Rest of Body" compartment
    W_tis[10] <- weight - sum(W_tis[1:(length(W_tis)-1)], na.rm = TRUE)
    V_tis[10] <- W_tis[10]/d_adipose     #(considering that the density of the rest tissues is 1 g/ml)
    Q[10] <- Q_total - sum(Q[2:(length(Q)-1)],na.rm = TRUE)
    V_cap[10] <- V_tis[10]*Capillary_fractions[10]
    # V_cap[9] <- Total_Blood - Vven - Vart - sum(V_cap[1:(length(V_cap)-1)], na.rm = TRUE) #this is problematic because it produces negative number
    W_macro[10] <- W_tis[10]*Macrophage_fraction[10]
    #Capillary_fractions[1] <- V_cap[1]/V_tis[1]

    Macro_blood <- Macrophage_fraction_blood*Total_Blood
    Wm_al <-  W_macro[1]

    if(stochastic == TRUE){
            x_fast <- exp(rnorm(1,mean(stan_fit$theta_tr[,1]), sd(stan_fit$theta_tr[,1])))
            CLE_ur<-  exp(rnorm(1,mean(stan_fit$theta_tr[,2]), sd(stan_fit$theta_tr[,2])))
            uptake <- exp(rnorm(1,mean(stan_fit$theta_tr[,3]), sd(stan_fit$theta_tr[,3])))
            k_alpc_tb <- exp(rnorm(1,mean(stan_fit$theta_tr[,4]), sd(stan_fit$theta_tr[,4])))
            k_lu_al <-  exp(rnorm(1,mean(stan_fit$theta_tr[,5]), sd(stan_fit$theta_tr[,5])))
            k_al_lu <- exp(rnorm(1,mean(stan_fit$theta_tr[,6]), sd(stan_fit$theta_tr[,6])))
            uptake_al <- exp(rnorm(1,mean(stan_fit$theta_tr[,7]), sd(stan_fit$theta_tr[,7])))
            uptake_skel <- exp(rnorm(1,mean(stan_fit$theta_tr[,8]), sd(stan_fit$theta_tr[,8])))
            P_lu <- exp(rnorm(1,mean(stan_fit$theta_tr[,9]), sd(stan_fit$theta_tr[,9])))
            P_ki <- exp(rnorm(1,mean(stan_fit$theta_tr[,10]), sd(stan_fit$theta_tr[,10])))
            P_br <- exp(rnorm(1,mean(stan_fit$theta_tr[,11]), sd(stan_fit$theta_tr[,11])))
            P_ut <- exp(rnorm(1,mean(stan_fit$theta_tr[,12]), sd(stan_fit$theta_tr[,12])))
            P_skin <- exp(rnorm(1,mean(stan_fit$theta_tr[,13]), sd(stan_fit$theta_tr[,13])))
            P_soft_skel_ht <- exp(rnorm(1,mean(stan_fit$theta_tr[,14]), sd(stan_fit$theta_tr[,14])))
            P_li_spl <- exp(rnorm(1,mean(stan_fit$theta_tr[,15]), sd(stan_fit$theta_tr[,15])))

    }else{
      x_fast <- exp(mean(stan_fit$theta_tr[,1]))
      CLE_ur<-  exp(mean(stan_fit$theta_tr[,2]))
      uptake <- exp(mean(stan_fit$theta_tr[,3]))
      k_alpc_tb <- exp(mean(stan_fit$theta_tr[,4]))
      k_lu_al <-  exp(mean(stan_fit$theta_tr[,5]))
      k_al_lu <- exp(mean(stan_fit$theta_tr[,6]))
      uptake_al <- exp(mean(stan_fit$theta_tr[,7]))
      uptake_skel <- exp(mean(stan_fit$theta_tr[,8]))
      P_lu <- exp(mean(stan_fit$theta_tr[,9]))
      P_ki <- exp(mean(stan_fit$theta_tr[,10]))
      P_br <- exp(mean(stan_fit$theta_tr[,11]))
      P_ut <- exp(mean(stan_fit$theta_tr[,12]))
      P_skin <- exp(mean(stan_fit$theta_tr[,13]))
      P_soft_skel_ht <- exp(mean(stan_fit$theta_tr[,14]))
      P_li_spl <- exp(mean(stan_fit$theta_tr[,15]))

    }

    e1 <- truncnorm::rtruncnorm(1, a=0, b=Inf, mean =  mean(stan_fit$sigma[,1]) , sd =  sd(stan_fit$sigma[,1]))
    e2 <- truncnorm::rtruncnorm(1, a=0, b=Inf, mean =  mean(stan_fit$sigma[,2]) , sd =  sd(stan_fit$sigma[,2]))
    e3 <- truncnorm::rtruncnorm(1, a=0, b=Inf, mean =  mean(stan_fit$sigma[,3]) , sd =  sd(stan_fit$sigma[,3]))

    uptake_spl <- uptake
    P_li <- P_li_spl
    P_spl <- P_li_spl
    P_ht <- P_soft_skel_ht
    P_soft <- P_soft_skel_ht
    P_skel <- P_soft_skel_ht

    x_lu <- x_fast
    x_li <- x_fast
    x_spl <- x_fast
    x_ki <- x_fast
    x_ht <- x_fast
    x_br <- x_fast
    x_ut <- x_fast
    x_skel <- x_fast
    x_skin <- x_fast
    x_soft <- x_fast

    return(list("Vlu_tis"= V_tis[1], "Vli_tis"= V_tis[2], "Vspl_tis"= V_tis[3], "Vki_tis"= V_tis[4],"Vht_tis"= V_tis[5],
                "Vbr_tis"= V_tis[6], "Vut_tis"= V_tis[7], "Vskel_tis"= V_tis[8], "Vskin_tis"= V_tis[9], "Vsoft_tis"= V_tis[10],
                "Vven"=Vven, "Vart"=Vart, "V_blood"=Total_Blood,

                "Vlu_cap"= V_cap[1], "Vli_cap"= V_cap[2], "Vspl_cap"= V_cap[3], "Vki_cap"= V_cap[4], "Vht_cap"= V_cap[5],
                "Vbr_cap"= V_cap[6], "Vut_cap"= V_cap[7], "Vskel_cap"= V_cap[8], "Vskin_cap"= V_cap[9], "Vsoft_cap"= V_cap[10],

                "Q_lu"= Q[1], "Q_li"= Q[2], "Q_spl"= Q[3], "Q_ki"= Q[4], "Q_ht"= Q[5], "Q_br"= Q[6], "Q_ut"= Q[7], "Q_skel"= Q[8],
                "Q_skin"= Q[9], "Q_soft"= Q[10],  "Q_total"=Q_total,

                "Wm_blood" = Macro_blood, "Wm_lu"= W_macro[1], "Wm_li"= W_macro[2], "Wm_spl"= W_macro[3], "Wm_ki"= W_macro[4],
                "Wm_ht"= W_macro[5], "Wm_br"= W_macro[6], "Wm_ut"= W_macro[7], "Wm_skel"= W_macro[8], "Wm_skin"= W_macro[9],
                "Wm_soft"= W_macro[10],

                "Inhaled.vol.rate" = Inhaled.vol.rate, "dep.ua" = dep.ua, "dep.tb" = dep.tb, "dep.al" = dep.al,
                "exposure.concentration" = exposure.concentration, "exposure.time" = exposure.time,

                "CLE_ua" = CLE_ua, "CLE_muc" = CLE_muc,
                "x_fast" = x_fast ,   "CLE_ur" = CLE_ur, "CLE_hep" = CLE_hep,
                "CLE_muc_cor" = CLE_muc_cor, "k_ab0" = k_ab0, "k_de" = k_de,
                "k_ab0_spl" = k_ab0_spl, "uptake" = uptake,
                "Wm_al" = Wm_al, "k_ua_br" = k_ua_br, "k_alpc_tb" = k_alpc_tb, "k_lu_al" = k_lu_al,
                "k_al_lu" = k_al_lu,   "uptake_al" =uptake_al,
                "uptake_skel" = uptake_skel, "uptake_spl" = uptake_spl,

                "P_lu"= P_lu, "x_lu" = x_lu, "P_li" = P_li, "x_li" = x_li, "P_spl" =P_spl,
                "x_spl" = x_spl, "P_ki" = P_ki, "x_ki" = x_ki, "P_ht" = P_ht, "x_ht" = x_ht,
                "P_br" = P_br, "x_br" = x_br, "P_ut" = P_ut, x_ut = x_ut,
                "P_skel" = P_skel, "x_skel" = x_skel, "P_skin" = P_skin, "x_skin" = x_skin,
                "P_soft" = P_soft, "x_soft" = x_soft,

                "error1" = e1, "error2" = e2, "error3" = e3
    ))

  })
}

### store them once

#################################################
# Function for creating initial values for ODEs #
#################################################

create.inits <- function(parameters){
  with( as.list(parameters),{
    Mua<- 0; Mtb <- 0;Mal <- 0; Mal_pc <- 0; Mlu_cap <- 0;
    Mlu_tis <- 0; Mlu_pc <- 0; Mli_cap <- 0; Mli_tis <- 0;
    Mli_pc <- 0; Mspl_cap <- 0; Mspl_tis <- 0; Mspl_pc <- 0;
    Mki_cap <- 0; Mki_tis <- 0; Mki_pc <- 0; Mht_cap <- 0;
    Mht_tis <- 0; Mht_pc <- 0; Mbr_cap <- 0; Mbr_tis <- 0;
    Mbr_pc <- 0; Mut_cap <- 0; Mut_tis <- 0; Mut_pc <- 0;
    Mskel_cap <- 0; Mskel_tis <- 0; Mskel_pc <- 0; Mskin_cap <- 0;
    Mskin_tis <- 0; Mskin_pc <- 0; Msoft_cap <- 0; Msoft_tis <- 0;
    Msoft_pc <- 0; Art_blood <- 0; Ven_blood <- 0; Mblood_pc <- 0;
    Mfeces <- 0; Murine <- 0;

    return(c("Mua" = Mua, "Mtb" = Mtb, "Mal" = Mal,"Mal_pc" = Mal_pc,
             "Mlu_cap" = Mlu_cap, "Mlu_tis" = Mlu_tis, "Mlu_pc" = Mlu_pc, "Mli_cap" = Mli_cap,
             "Mli_tis" = Mli_tis, "Mli_pc" = Mli_pc, "Mspl_cap" = Mspl_cap, "Mspl_tis" = Mspl_tis,
             "Mspl_pc" = Mspl_pc, "Mki_cap" = Mki_cap, "Mki_tis" = Mki_tis, "Mki_pc" = Mki_pc,
             "Mht_cap" = Mht_cap, "Mht_tis" = Mht_tis, "Mht_pc" = Mht_pc, "Mbr_cap" = Mbr_cap,
             "Mbr_tis" = Mbr_tis, "Mbr_pc" = Mbr_pc,  "Mut_cap" = Mut_cap, "Mut_tis" = Mut_tis,
             "Mut_pc" = Mut_pc, "Mskel_cap" = Mskel_cap, "Mskel_tis" = Mskel_tis, "Mskel_pc" = Mskel_pc,
             "Mskin_cap" = Mskin_cap, "Mskin_tis" = Mskin_tis, "Mskin_pc" = Mskin_pc, "Msoft_cap" = Msoft_cap,
             "Msoft_tis" = Msoft_tis, "Msoft_pc" = Msoft_pc, "Art_blood" = Art_blood, "Ven_blood" = Ven_blood,
             "Mblood_pc" = Mblood_pc, "Mfeces" = Mfeces, "Murine" =Murine))
  })
}

#################################################
# Function for creating events #
#################################################
create.events<- function(parameters){
  with( as.list(parameters),{

    exposure.time <- seq(from=0, to=exposure.time , by=exposure.time /100) #inhalation time in hours (When was this applies)
    exposure.concentration <- rep(exposure.concentration,length(exposure.time))  # in mg/m^3

    lexposure <- length(exposure.concentration)
    ltimes <- length(exposure.time)

    add.tb <- rep(0,lexposure-1)
    add.al <- rep(0,lexposure-1)
    cur.time <- exposure.time[1]

    for (i in 1:(ltimes-1)){
      # Calculate the interval of exposure
      interval <- exposure.time[i+1] - cur.time
      add.tb[i] = interval * exposure.concentration[i] * (Inhaled.vol.rate/1000) * dep.tb * 1000 # deposited mass in microgram in trancheobronchial region
      add.al[i] = interval * exposure.concentration[i] * (Inhaled.vol.rate/1000) * dep.al * 1000 # deposited mass in microgram in alveolar region
      cur.time <- exposure.time[i+1]
    }

    if (lexposure == ltimes){
      events <- list(data = rbind(data.frame(var = "Mtb",  time = exposure.time[2:ltimes],
                                             value = add.tb, method = c("add")),
                                  data.frame(var = "Mal",  time = exposure.time[2:ltimes],
                                             value = add.al, method = c("add"))

      ))
    }else{
      stop("The user must provide t")
    }

    return(events)
  })
}

###################
# Custom function #
###################

custom.func <- function(){
  return()
}

#################
# ODEs system #
#################

ode.func <- function(time, Initial.values, Parameters, custom.func){
  with( as.list(c(Initial.values, Parameters)),{

    # concentrations in tissues
    #C_lu <- Lu_tissue/W_lu
    Clu_tis  <-  Mlu_tis/Vlu_tis
    Clu_cap <-  Mlu_cap/Vlu_cap
    Cli_tis  <-  Mli_tis/Vli_tis
    Cli_cap <-  Mli_cap/Vli_cap
    Cspl_tis  <-  Mspl_tis/Vspl_tis
    Cspl_cap  <-  Mspl_cap/Vspl_cap
    Cki_tis  <-  Mki_tis/Vki_tis
    Cki_cap  <-  Mki_cap/Vki_cap
    Cht_tis  <-  Mht_tis/Vht_tis
    Cht_cap  <-  Mht_cap/Vht_cap
    Cbr_tis  <-  Mbr_tis/Vbr_tis
    Cbr_cap  <-  Mbr_cap/Vbr_cap
    Cut_tis  <-  Mut_tis/Vut_tis
    Cut_cap  <-  Mut_cap/Vut_cap
    Cskel_tis  <-  Mskel_tis/Vskel_tis
    Cskel_cap  <-  Mskel_cap/Vskel_cap
    Cskin_tis  <-  Mskin_tis/Vskin_tis
    Cskin_cap  <-  Mskin_cap/Vskin_cap
    Csoft_tis  <-  Msoft_tis/Vsoft_tis
    Csoft_cap  <-  Msoft_cap/Vsoft_cap

    Cart <- Art_blood/(0.19*V_blood)
    Cven <- Ven_blood/(0.81*V_blood)

    # Uptake rates by phagocytizing cells
    kluab <- k_ab0*(1-(Mlu_pc/(Wm_lu*uptake)))
    kliab <- k_ab0*(1-(Mli_pc/(Wm_li*uptake)))
    ksplab <- k_ab0_spl*(1-(Mspl_pc/(Wm_spl*uptake_spl)))
    kkiab <- k_ab0*(1-(Mki_pc/(Wm_ki*uptake)))
    khtab <- k_ab0*(1-(Mht_pc/(Wm_ht*uptake)))
    kbrab <- k_ab0*(1-(Mbr_pc/(Wm_br*uptake)))
    kutab <- k_ab0*(1-(Mut_pc/(Wm_ut*uptake)))
    kskelab <- k_ab0*(1-(Mskel_pc/(Wm_skel*uptake_skel)))
    kskinab <- k_ab0*(1-(Mskin_pc/(Wm_skin*uptake)))
    ksoftab <- k_ab0*(1-(Msoft_pc/(Wm_soft*uptake)))
    kbloodab <- k_ab0*(1-(Mblood_pc/(Wm_blood*uptake)))
    kalab <- k_ab0*(1-(Mal_pc/(Wm_al*uptake_al)))

    #Upper airways
    dMua = -(k_ua_br*Mua) - (CLE_ua * Mua)

    #Trancheobronchial region
    dMtb = (k_alpc_tb * Mal_pc) - (CLE_muc * CLE_muc_cor* Mtb)

    #Alveolar region
    dMal = (k_lu_al * Mlu_tis) - (k_al_lu * Mal) - (Mal*kalab - Mal_pc*k_de)
    dMal_pc = (Mal*kalab - Mal_pc*k_de) - (k_alpc_tb * Mal_pc)

    #Lungs
    dMlu_cap  =  (Cven*Q_lu)- (Clu_cap*Q_lu) + (x_lu*Q_lu)*Clu_tis/P_lu - (x_lu*Q_lu)*Clu_cap ;
    dMlu_tis  =  - (x_lu*Q_lu)*Clu_tis/P_lu + (x_lu*Q_lu)*Clu_cap - (Vlu_tis*Clu_tis*kluab - Mlu_pc*k_de) -
      (k_lu_al * Mlu_tis) + (k_al_lu * Mal) ; #Lung interstitium
    dMlu_pc  = (Vlu_tis*Clu_tis*kluab - Mlu_pc*k_de);

    #Liver
    dMli_cap  = (Cart*Q_li) + (Cspl_cap*Q_spl) - (Cli_cap*(Q_li+Q_spl)) + (x_li*Q_li)*Cli_tis/P_li - (x_li*Q_li)*Cli_cap;  #capillary
    dMli_tis = - (x_li*Q_li)*Cli_tis/P_li + (x_li*Q_li)*Cli_cap -(Vli_tis*Cli_tis*kliab - Mli_pc*k_de) - (CLE_hep*Mli_tis) ; #tissue
    dMli_pc  = (Vli_tis*Cli_tis*kliab - Mli_pc*k_de) ; #Phagocytized

    #Spleen
    dMspl_cap  = (Cart*Q_spl) - (Cspl_cap*Q_spl) +  (x_spl*Q_spl)*Cspl_tis/P_spl - (x_spl*Q_spl)*Cspl_cap ;  #capillary
    dMspl_tis =  - (x_spl*Q_spl)*Cspl_tis/P_spl + (x_spl*Q_spl)*Cspl_cap - (Vspl_tis*Cspl_tis*ksplab - Mspl_pc*k_de)  ; #tissue
    dMspl_pc  = (Vspl_tis*Cspl_tis*ksplab - Mspl_pc*k_de) ; #seq

    #Kidneys
    dMki_cap  = (Cart*Q_ki) - (Cki_cap*Q_ki) +  (x_ki*Q_ki)*Cki_tis/P_ki - (x_ki*Q_ki)*Cki_cap- (CLE_ur*Mki_cap) ;  #capillary
    dMki_tis =  - (x_ki*Q_ki)*Cki_tis/P_ki + (x_ki*Q_ki)*Cki_cap - (Vki_tis*Cki_tis*kkiab - Mki_pc*k_de)  ; #tissue
    dMki_pc  = (Vki_tis*Cki_tis*kkiab - Mki_pc*k_de) ; #seq

    #Heart
    dMht_cap  = (Cart*Q_ht) - (Cht_cap*Q_ht) +  (x_ht*Q_ht)*Cht_tis/P_ht - (x_ht*Q_ht)*Cht_cap ;  #capillary
    dMht_tis =  - (x_ht*Q_ht)*Cht_tis/P_ht + (x_ht*Q_ht)*Cht_cap - (Vht_tis*Cht_tis*khtab - Mht_pc*k_de)  ; #tissue
    dMht_pc  = (Vht_tis*Cht_tis*khtab - Mht_pc*k_de) ; #seq

    #Brain
    dMbr_cap  = (Cart*Q_br) - (Cbr_cap*Q_br) +  (x_br*Q_br)*Cbr_tis/P_br - (x_br*Q_br)*Cbr_cap ;  #capillary
    dMbr_tis =  - (x_br*Q_br)*Cbr_tis/P_br + (x_br*Q_br)*Cbr_cap - (Vbr_tis*Cbr_tis*kbrab - Mbr_pc*k_de) + (k_ua_br*Mua)  ; #tissue
    dMbr_pc  =  (Vbr_tis*Cbr_tis*kbrab - Mbr_pc*k_de) ; #seq

    #Uterus
    dMut_cap  = (Cart*Q_ut) - (Cut_cap*Q_ut) +  (x_ut*Q_ut)*Cut_tis/P_ut - (x_ut*Q_ut)*Cut_cap ;  #capillary
    dMut_tis =  - (x_ut*Q_ut)*Cut_tis/P_ut + (x_ut*Q_ut)*Cut_cap - (Vut_tis*Cut_tis*kutab - Mut_pc*k_de)  ; #tissue
    dMut_pc  = (Vut_tis*Cut_tis*kutab - Mut_pc*k_de) ; #seq

    #Skeleton
    dMskel_cap  = (Cart*Q_skel) - (Cskel_cap*Q_skel) +  (x_skel*Q_skel)*Cskel_tis/P_skel - (x_skel*Q_skel)*Cskel_cap ;  #capillary
    dMskel_tis =  - (x_skel*Q_skel)*Cskel_tis/P_skel + (x_skel*Q_skel)*Cskel_cap - (Vskel_tis*Cskel_tis*kskelab - Mskel_pc*k_de)  ; #tissue
    dMskel_pc  = (Vskel_tis*Cskel_tis*kskelab - Mskel_pc*k_de) ; #seq

    #Skin
    dMskin_cap  = (Cart*Q_skin) - (Cskin_cap*Q_skin) +  (x_skin*Q_skin)*Cskin_tis/P_skin - (x_skin*Q_skin)*Cskin_cap ;  #capillary
    dMskin_tis =  - (x_skin*Q_skin)*Cskin_tis/P_skin + (x_skin*Q_skin)*Cskin_cap - (Vskin_tis*Cskin_tis*kskinab - Mskin_pc*k_de)  ; #tissue
    dMskin_pc  = (Vskin_tis*Cskin_tis*kskinab - Mskin_pc*k_de) ; #seq

    #Soft
    dMsoft_cap  = (Cart*Q_soft) - (Csoft_cap*Q_soft) +  (x_soft*Q_soft)*Csoft_tis/P_soft - (x_soft*Q_soft)*Csoft_cap ;  #capillary
    dMsoft_tis =  - (x_soft*Q_soft)*Csoft_tis/P_soft + (x_soft*Q_soft)*Csoft_cap - (Vsoft_tis*Csoft_tis*ksoftab - Msoft_pc*k_de)  ; #tissue
    dMsoft_pc  = (Vsoft_tis*Csoft_tis*ksoftab - Msoft_pc*k_de) ; #seq

    #Blood
    dArt_blood = (Clu_cap*Q_lu) - Cart* (Q_li + Q_spl+ Q_ki + Q_ht + Q_br + Q_ut + Q_skel + Q_skin + Q_soft) -
      (0.19*V_blood*Cart*kbloodab - 0.19* Mblood_pc*k_de);
    dVen_blood = (Cli_cap*(Q_li+Q_spl)) +  (Cki_cap*Q_ki) + (Cht_cap*Q_ht) + (Cbr_cap*Q_br) + (Cut_cap*Q_ut) + (Cskel_cap*Q_skel) +
      (Cskin_cap*Q_skin) + (Csoft_cap*Q_soft) - (Cven*Q_lu)  -  (0.81*V_blood*Cven*kbloodab - 0.81* Mblood_pc*k_de)
    dMblood_pc <- ((0.19*V_blood*Cart + 0.81*V_blood*Cven)*kbloodab - Mblood_pc*k_de)

    #Feces
    dMfeces = (CLE_hep*Mli_tis) + (CLE_ua * Mua) + (CLE_muc* CLE_muc_cor * Mtb)  ;

    #Urine
    dMurine = (CLE_ur*Mki_cap) ;

    Total_lungs <- Mal + Mal_pc + Mlu_cap + Mlu_tis + Mlu_pc + Mtb

    list(c(dMua = dMua, dMtb = dMtb, dMal = dMal,dMal_pc = dMal_pc,
           dMlu_cap = dMlu_cap, dMlu_tis = dMlu_tis, dMlu_pc = dMlu_pc, dMli_cap = dMli_cap,
           dMli_tis = dMli_tis, dMli_pc = dMli_pc, dMspl_cap = dMspl_cap, dMspl_tis = dMspl_tis,
           dMspl_pc = dMspl_pc, dMki_cap = dMki_cap, dMki_tis = dMki_tis, dMki_pc = dMki_pc,
           dMht_cap = dMht_cap, dMht_tis = dMht_tis, dMht_pc = dMht_pc, dMbr_cap = dMbr_cap,
           dMbr_tis = dMbr_tis, dMbr_pc = dMbr_pc,  dMut_cap = dMut_cap, dMut_tis = dMut_tis,
           dMut_pc = dMut_pc, dMskel_cap = dMskel_cap, dMskel_tis = dMskel_tis, dMskel_pc = dMskel_pc,
           dMskin_cap = dMskin_cap, dMskin_tis = dMskin_tis, dMskin_pc = dMskin_pc, dMsoft_cap = dMsoft_cap,
           dMsoft_tis = dMsoft_tis, dMsoft_pc = dMsoft_pc, dArt_blood = dArt_blood, dVen_blood = dVen_blood,
           dMblood_pc = dMblood_pc, dMfeces = dMfeces, dMurine = dMurine), Total_lungs = Total_lungs)
  })
}

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

# User must provide a string declaring the path where the vpc_data file is stored
path <- ""
input.data <- openxlsx::read.xlsx(paste(path,"biodist_data.xlsx", sep = ""),
                                  sheet=3,colNames = TRUE, rowNames = TRUE)

# the following numbers derive from normalisation of the MPPD numbers 0.12 and 0.52
tb.depo = 0.1875
al.depo = 0.8125
ua.depo = 0
# use average input data to represent all rats
input <-list("exposure.concentration" = rowMeans(input.data[3,]), "exposure.time" = 2,
              "weight" = 277, "dep.ua" = 0, "dep.tb" = tb.depo*rowMeans(input.data[2,]),
              "dep.al" =  al.depo*rowMeans(input.data[2,]),
              "Inhaled.vol.rate" = rowMeans(input.data[10,]))

# Produce one solution with mean values
params_determ <- create.params(input,stan_fit, stochastic = FALSE)
inits <- create.inits(params_determ)
events <- create.events(params_determ)

# Here we create a first solution to get the dimensions because the total instances will be
# equal to the dimension of sample_time plus the extra events which occured at times
# not provided in the sample time vector
sample_time <- c(1e-08, 1e-04, 1e-03, 1e-02, 1e-01, 0.5, 1, 2, 6, 12, 26, 48, 72, 144, 256, 300,
                 400, 500, 600 , 700)
solution <-  ode(times = sample_time,  func = ode.func, y = inits, parms = params_determ,
                 custom.func = custom.func, method="lsodes",  events = events)
solution <- as.data.frame(solution[solution[,1]  %in% sample_time,1:40])
deterministic.df <- matrix(rep(NA, 17*dim(solution)[1]), ncol = 17)
deterministic.df[,1] <- solution[,1]
deterministic.df[,2] <- solution[,3]
deterministic.df[,3] <- solution[,4]
deterministic.df[,4] <-  solution[,5]
deterministic.df[,5] <- (solution[,7] + solution[,8])
deterministic.df[,6] <- (solution[,10] + solution[,11])
deterministic.df[,7] <- (solution[,13] + solution[,14])
deterministic.df[,8] <- (solution[,16] + solution[,17])
deterministic.df[,9] <- (solution[,19] + solution[,20])
deterministic.df[,10] <- (solution[,22] + solution[,23])
deterministic.df[,11] <- (solution[,25] + solution[,26])
deterministic.df[,12] <- (solution[,28] + solution[,29])
deterministic.df[,13] <- (solution[,31] + solution[,32])
deterministic.df[,14] <- (solution[,34] + solution[,35])
deterministic.df[,15] <- (solution[,36] + solution[,37]+ solution[,38])
deterministic.df[,16] <- solution[,39]
deterministic.df[,17] <- solution[,40]

colnames(deterministic.df) <- c( "Time", "Trachea", "BALF", "BALC", "Lavaged lungs", "Liver", "Spleen",
                         "Kidneys", "Heart", "Brain", "Uterus", "Skeleton", "Skin", "Soft tissues",
                         "Blood", "Feces", "Urine")

Nsim = 1000 # number of simulations
Nrat = 20 # number of virtual rats
ltime <- dim(solution)[1] # number of solution instances
data_time <- solution[,"time"]
Ncomp <- dim(solution)[2] # Number of compartments to plot
pred_con_tra <-array(rep(NA,Nrat*ltime*Nsim),dim=c(Nrat,ltime,Nsim))
pred_con_balf <-array(rep(NA,Nrat*ltime*Nsim),dim=c(Nrat,ltime,Nsim))
pred_con_balc <-array(rep(NA,Nrat*ltime*Nsim),dim=c(Nrat,ltime,Nsim))
pred_con_lav <-array(rep(NA,Nrat*ltime*Nsim),dim=c(Nrat,ltime,Nsim))
pred_con_li <-array(rep(NA,Nrat*ltime*Nsim),dim=c(Nrat,ltime,Nsim))
pred_con_spl <-array(rep(NA,Nrat*ltime*Nsim),dim=c(Nrat,ltime,Nsim))
pred_con_ki <-array(rep(NA,Nrat*ltime*Nsim),dim=c(Nrat,ltime,Nsim))
pred_con_ht <-array(rep(NA,Nrat*ltime*Nsim),dim=c(Nrat,ltime,Nsim))
pred_con_br <-array(rep(NA,Nrat*ltime*Nsim),dim=c(Nrat,ltime,Nsim))
pred_con_ut <-array(rep(NA,Nrat*ltime*Nsim),dim=c(Nrat,ltime,Nsim))
pred_con_skel <-array(rep(NA,Nrat*ltime*Nsim),dim=c(Nrat,ltime,Nsim))
pred_con_skin <-array(rep(NA,Nrat*ltime*Nsim),dim=c(Nrat,ltime,Nsim))
pred_con_soft <-array(rep(NA,Nrat*ltime*Nsim),dim=c(Nrat,ltime,Nsim))
pred_con_blood <-array(rep(NA,Nrat*ltime*Nsim),dim=c(Nrat,ltime,Nsim))
pred_con_feces <-array(rep(NA,Nrat*ltime*Nsim),dim=c(Nrat,ltime,Nsim))
pred_con_urine <-array(rep(NA,Nrat*ltime*Nsim),dim=c(Nrat,ltime,Nsim))

failed = list()

for (sim in 1:Nsim){
  for (rat in 1:Nrat){
    print(paste("rat   ",rat, "Nsim    ",sim))

        pred <- matrix(rep(NA, Ncomp*ltime), nrow = ltime)
        # parameter to inform about failed attempts
        failed_attempt <- 0
        # if params result in odes not solving, resample and resolve
        while( (dim(pred)[1] != ltime) | sum(is.na(pred)) ) {

        #params are selected stochastically because create.params samples from stan estimates
        params <- create.params(input,stan_fit)
        events <- create.events(params)
        sol  <- ode(times = sample_time,  func = ode.func, y = inits, parms = params,
                    custom.func = custom.func, method="lsodes",  events = events,
                    maxsteps = 50000)
        pred <- sol[sol[,1] %in% sample_time,]
        pred
        }

        for(t in 1:ltime){
        ###Total amount of NPs in each organ
        # Amount in trachea
        pred_con_tra[rat,t,sim] <- truncnorm::rtruncnorm(1, a=0, b=Inf, mean = pred[t,3] , sd = params$error1)
        # Amount in Balf
        pred_con_balf[rat,t,sim] <- truncnorm::rtruncnorm(1, a=0, b=Inf, mean = pred[t,4] , sd = params$error1/2)

        # Amount in Balc
        pred_con_balc[rat,t,sim] <- truncnorm::rtruncnorm(1, a=0, b=Inf, mean = pred[t,5] , sd = params$error1/2)

        # Amount in lavaged lungs
        pred_con_lav[rat,t,sim]<-   truncnorm::rtruncnorm(1, a=0, b=Inf, mean = (pred[t,7] + pred[t,8])  , sd = params$error1)

        # Amount in liver
        pred_con_li[rat,t,sim] <- truncnorm::rtruncnorm(1, a=0, b=Inf, mean =  (pred[t,10] + pred[t,11]) , sd = params$error2)

        # Amount in spleen
        pred_con_spl[rat,t,sim]<-  truncnorm::rtruncnorm(1, a=0, b=Inf, mean = (pred[t,13] + pred[t,14]) , sd = params$error3)

        # Amount in kidneys
        pred_con_ki[rat,t,sim] <-  truncnorm::rtruncnorm(1, a=0, b=Inf, mean = (pred[t,16] + pred[t,17]) , sd = params$error2)

        # Amount in heart
        pred_con_ht[rat,t,sim] <- truncnorm::rtruncnorm(1, a=0, b=Inf, mean = (pred[t,19] + pred[t,20]) , sd = params$error3)

        # Amount in brain
        pred_con_br[rat,t,sim] <- truncnorm::rtruncnorm(1, a=0, b=Inf, mean = (pred[t,22] + pred[t,23]) , sd = params$error3)

        # Amount in uterus
        pred_con_ut[rat,t,sim] <- truncnorm::rtruncnorm(1, a=0, b=Inf, mean = (pred[t,25] + pred[t,26]) , sd = params$error3)

        # Amount in skeleton
        pred_con_skel[rat,t,sim] <- truncnorm::rtruncnorm(1, a=0, b=Inf, mean = (pred[t,28] + pred[t,29])  , sd = params$error2)

        # Amount in skin
        pred_con_skin[rat,t,sim] <- truncnorm::rtruncnorm(1, a=0, b=Inf, mean = (pred[t,31] + pred[t,32])  , sd = params$error2)

        # Amount in soft tissues
        pred_con_soft[rat,t,sim] <-  truncnorm::rtruncnorm(1, a=0, b=Inf, mean = (pred[t,34] + pred[t,35])  , sd = params$error2)

        # Amount in blood
        pred_con_blood[rat,t,sim] <- truncnorm::rtruncnorm(1, a=0, b=Inf, mean = (pred[t,36] + pred[t,37]+ pred[t,38]), sd = params$error2)

        # Amount in  feces
        pred_con_feces[rat,t,sim] <- truncnorm::rtruncnorm(1, a=0, b=Inf, mean = pred[t,39], sd = params$error1)

        # Amount in urine
        pred_con_urine[rat,t,sim] <- truncnorm::rtruncnorm(1, a=0, b=Inf, mean = pred[t,40], sd = params$error1)
        }
       }
}

bag_of_data <- list( pred_con_tra, pred_con_balf, pred_con_balc, pred_con_lav, pred_con_li,
                     pred_con_spl, pred_con_ki, pred_con_ht, pred_con_br, pred_con_ut, pred_con_skel,
                     pred_con_skin, pred_con_soft, pred_con_blood, pred_con_feces, pred_con_urine)
names(bag_of_data) <- c( "Trachea", "BALF", "BALC", "Lavaged lungs", "Liver", "Spleen",
                         "Kidneys", "Heart", "Brain", "Uterus", "Skeleton", "Skin", "Soft tissues",
                         "Blood", "Feces", "Urine")
comp_names <- names(bag_of_data)

# File where data are stored
path <- ""
### load data
biodist_data <- openxlsx::read.xlsx(paste(path,"biodist_data.xlsx", sep = ""),
                                    sheet=1,colNames = TRUE,rowNames = TRUE)
sd_data <- openxlsx::read.xlsx(paste(path,"biodist_data.xlsx", sep = ""),
                               sheet=2,colNames = TRUE,rowNames = TRUE)

feces <- openxlsx::read.xlsx(paste(path,"feces.xlsx", sep = ""),
                             sheet=2,colNames = TRUE)[,2]*input.data[1,5]
urine<- openxlsx::read.xlsx(paste(path,"urine.xlsx", sep = "")
                            , sheet=2,colNames = TRUE)[,2]*input.data[7,5]

excreta_time <- c(3.5, 7, 10.5, 14, 17.5, 21, 24.5, 27)*24
biodist_time <- c(2, 6, 26, (7*24+2), (28*24+2))

excreta.df <- as.data.frame(cbind(excreta_time, feces, urine))
colnames(excreta.df) <- c("Time", "Feces", "Urine")

data.df <- data.frame(cbind(biodist_time, biodist_data))
colnames(data.df) <- c("Time", colnames(biodist_data))

sd.data.df <- data.frame(cbind(biodist_time, sd_data))
colnames(data.df) <- c("Time", colnames(sd_data))

# Set the path where the plots will be stored
setwd(path)

counter <-1
for (dat in bag_of_data) {
    # Get the compartment name to be plotted
     comp_name <- comp_names[counter]
     # Define the save name
     save_name <- paste0(comp_name, "_vpc.png", sep="")
     pred_con <- dat

     ##########################
     # Calculation of 95 % CI of the 50th quantile
     con<-matrix(rep(NA,ltime*Nsim), nrow = ltime)
     ymin<- rep(NA,ltime)
     ymax<- rep(NA,ltime)

     for (i in 1:ltime){
       for (j in 1:Nsim){
         con[i,j]<-quantile(pred_con[,i,j],probs=0.5)
       }
     }

     for (i in 1:ltime){
        con[i,]<-sort(con[i,])
        ymin[i]<-quantile(con[i,],probs=0.025)
        ymax[i]<-quantile(con[i,],probs=0.975)
     }

    q50<-data.frame(time=data_time,lo=ymin,hi=ymax)
    ############################
    # Calculation of 95 % CI of the 5th quantile

    con<-matrix(rep(NA,ltime*Nsim), nrow = ltime)
    ymin<- rep(NA,ltime)
    ymax<- rep(NA,ltime)

    for (i in 1:ltime){
      for (j in 1:Nsim){
        con[i,j]<-quantile(pred_con[,i,j],probs=0.05)
      }
    }

    for (i in 1:ltime){
      con[i,]<-sort(con[i,])
      ymin[i]<-quantile(con[i,],probs=0.025)
      ymax[i]<-quantile(con[i,],probs=0.975)
    }

    q05<-data.frame(time=data_time,lo=ymin,hi=ymax)
    #################################
    # Calculation of 95 % CI of the 95th quantile
    con<-matrix(rep(NA,ltime*Nsim), nrow = ltime)
    ymin<- rep(NA,ltime)
    ymax<- rep(NA,ltime)

    for (i in 1:ltime){
      for (j in 1:Nsim){
        con[i,j]<-quantile(pred_con[,i,j],probs=0.95)
      }
    }

    for (i in 1:ltime){
      con[i,]<-sort(con[i,])
      ymin[i]<-quantile(con[i,],probs=0.025)
      ymax[i]<-quantile(con[i,],probs=0.975)
    }

    q95<-data.frame(time=data_time,lo=ymin,hi=ymax)
  ############################
    # Mean values of predictions from deterministic.df model
    use_comp <- colnames(deterministic.df)[counter+1]
    mean_pred = data.frame(deterministic.df[,c("Time", use_comp)])
    colnames(mean_pred) = c("time", "value")

  if(!(comp_name %in% c("Feces", "Urine"))){
    use_comp <- colnames(data.df)[counter+1]
    df1 <- as.data.frame(cbind(data.df[,c("Time",use_comp)], sd.data.df[,use_comp]))
    colnames(df1) <- c("Time", "Mean", "Sd")
    my_plot <- ggplot() +
          geom_line(data=mean_pred, aes(x=time, y=value, linetype = " Prediction mean"),
                    size=1.2)+
          geom_ribbon(data=q50,aes(x=time, ymin = lo, ymax = hi, fill = "50th percentile"),
                      inherit.aes=FALSE, alpha = 0.2) +
          geom_ribbon(data=q05,aes(x=time, ymin = lo, ymax = hi, fill = "5th percentile"),
                      inherit.aes=FALSE, alpha = 0.2) +
          geom_ribbon(data=q95,aes(x=time, ymin = lo, ymax = hi, fill="95th percentile"),
                      inherit.aes=FALSE,alpha = 0.2) +
         geom_point(data = df1, aes(x = Time, y=Mean),size=5)+
         geom_errorbar(data = df1, aes(x = Time, ymin=ifelse((Mean-Sd)>0,Mean-Sd,0), ymax=Mean+Sd),size=1)+
         labs(title = rlang::expr(!!comp_name), y = "TiO2 (ug)", x = "Time (in hours)") +
         scale_fill_manual("Prediction Ribbons", values = c( "50th percentile" = 1,
                                                  "5th percentile" = 2,"95th percentile" = 3)) +
        scale_linetype_manual("Line", values = c(" Prediction mean" = 1))+
        scale_colour_manual("Point", values = c("Biodistribtion data" = 1))+
        theme(plot.title =element_text(hjust = 0.5, size=30, face="bold"),
              axis.title.y =element_text(hjust = 0.5, size=20, face="bold"),
              axis.text.y=element_text(size=18),
              axis.title.x =element_text(hjust = 0.5, size=20, face="bold"),
              axis.text.x=element_text(size=18),
              legend.title=element_text(hjust = 0.01, size=20),
              legend.text=element_text(size=18))
        png(rlang::expr(!!save_name), width = 15, height = 10, units = 'in', res = 500)
        print(my_plot)
  } else{
      observed <- excreta.df[,c("Time",colnames(excreta.df)[counter-13])]
      colnames(observed) <- c("Time", "mean")
      my_plot <- ggplot(observed, aes(x=Time, y=mean, colour="Biodistribtion data"))+
        geom_point(shape=19, size=4) +
        #geom_bar(aes(ymin=ifelse((mean-sd)>0,mean-sd,0), ymax=mean+sd), width=.1) +
        geom_line(data=mean_pred, aes(x=time, y=value, linetype = " Prediction mean"),
                  size=1.2)+
        geom_ribbon(data=q50,aes(x=time, ymin = lo, ymax = hi, fill = "50th percentile"),
                    inherit.aes=FALSE, alpha = 0.2) +
        geom_ribbon(data=q05,aes(x=time, ymin = lo, ymax = hi, fill = "5th percentile"),
                    inherit.aes=FALSE, alpha = 0.2) +
        geom_ribbon(data=q95,aes(x=time, ymin = lo, ymax = hi, fill="95th percentile"),
                    inherit.aes=FALSE,alpha = 0.2) +
        labs(title = rlang::expr(!!comp_name), y = "TiO2 (ug)", x = "Time (in hours)") +
        scale_fill_manual("Prediction Ribbons", values = c( "50th percentile" = 1,
                                                "5th percentile" = 2,"95th percentile" = 3)) +
        scale_linetype_manual("Line", values = c(" Prediction mean" = 1))+
        scale_colour_manual("Point", values = c("Biodistribtion data" = 1))+
        theme(plot.title =element_text(hjust = 0.5, size=30, face="bold"),
              axis.title.y =element_text(hjust = 0.5, size=20, face="bold"),
              axis.text.y=element_text(size=18),
              axis.title.x =element_text(hjust = 0.5, size=20, face="bold"),
              axis.text.x=element_text(size=18),
              legend.title=element_text(hjust = 0.01, size=20),
              legend.text=element_text(size=18))
      png(rlang::expr(!!save_name), width = 15, height = 10, units = 'in', res = 500)
      print(my_plot)

  }
  dev.off()
  counter <- counter +1
}


runTime <- function(combs, maxDoseT, missp, Rscale, doseType = "random", Ne = 1.5*10^5){

    mutinf <- return_mutinf()
    druginf <- return_druginf()

    hours <- list()
    hours[[combs[1]]] <- 0
    hours[[combs[2]]] <- 0
    hours[[combs[3]]] <- 0

    #There are three dose missing types:
    #These patterns are NOT explored in this MS, but I will include them anyway

    #1) Totally random, with dose taking probability p
    if(doseType == "random"){
        missedDoses <- runif(maxDoseT, 0, 1) < missp

	#Each year, they have a missp probability of a week long treatment interruption (i.e., 336 time periods)
	missedYears <- runif(10, 0, 1) < missp
	#Each year is 48 * 7 * 4 * 12 = 16128 time periods
	for(i in 1:10){
	      if(i == TRUE){
	      	   #choose a time within the year
		   missedDoseStart <- 16128 * (i - 1) + sample(1:16128, 1)
		   missedDoses[missedDoseStart:(missedDoseStart + 336 * 2)] <- 0
	      }
	}

        for(i in 2:maxDoseT){
            if(missedDoses[i] == TRUE){
                for(D in combs){
                    hours[[D]] <- append(hours[[D]], tail(hours[[D]], n = 1) + 0.5)
                }
            }else{
                for(D in combs){
                    di <- druginf[[which(names(druginf) == D)]]
                    dosing <- di[which(names(di) == "dosing")]
                    hour <- tail(hours[[D]], n = 1)
                    if(i %% (24/dosing) == 0){
                        hour = 0
                    }else{
                        hour = hour + 0.5
                    }
                    hours[[D]] <- append(hours[[D]], hour)
                }
            }
        }

    }

    #Weekends off, weekdays on (Five on, Two off)
    if(doseType == "FOTO"){

        #Weekday index?
        #240 half hour chunks in 5 days
        #96 half hour chunks in a weekend
        missedDoses <- rep(c(rep(FALSE, 240), rep(TRUE, 96)),
                           ceiling(maxDoseT/(240+96)))[1:maxDoseT]

        for(i in 2:maxDoseT){
            if(missedDoses[i] == TRUE){
                for(D in combs){
                    hours[[D]] <- append(hours[[D]], tail(hours[[D]], n = 1) + 0.5)
                }
            }else{
                for(D in combs){
                    di <- druginf[[which(names(druginf) == D)]]
                    dosing <- di[which(names(di) == "dosing")]
                    hour <- tail(hours[[D]], n = 1)
                    if(i %% (24/dosing) == 0){
                        hour = 0
                    }else{
                        hour = hour + 0.5
                    }
                    hours[[D]] <- append(hours[[D]], hour)
                }
            }
        }

    }

    #The dosing approach that maximizes the time at which resistant strains have higher R than
    # sensitive strains
    if(doseType == "resistantAdvantage"){

          dosingPeriods <- foreach(D = combs, .combine = "rbind")%do%{

            di <- druginf[[which(names(druginf) == D)]]
            mi <- mutinf[[which(names(druginf) == D)]]

            IC_50 <- di[which(names(di) == "IC_50")]
            hl <- di[which(names(di) == "hl")]
            m <- di[which(names(di) == "slope")]
            C_max <- di[which(names(di) == "c_max")]
            dosing <- di[which(names(di) == "dosing")]

            s <- as.numeric(mi[which(names(mi) == "s")])
            mu <- as.numeric(mi[which(names(mi) == "mu")])
            rho <- as.numeric(mi[which(names(mi) == "rho")])
            sigma <- as.numeric((mi[which(names(mi) == "sigma")]))

            mut <- 1
            checkPeriod <- 1000

            #clunky
            relHours <- bind_rows(bind_cols(t =  1:checkPeriod, nhours = seq(0.5, checkPeriod/2, by  = 0.5),
                                mut = rep(1, checkPeriod)),
                      bind_cols(t =  1:checkPeriod, nhours = seq(0.5, checkPeriod/2, by  = 0.5),
                                mut = rep(0, checkPeriod)))   %>%
                      mutate(drug = D) %>%
                      mutate( R = (1 - mut * s)/(1 + (decay_instant(nhours, hl, C_max )/(IC_50 * (1 - mut*(1 - rho))))^(m * (1 + mut * (sigma)))))

            worstT <- (relHours %>% spread(mut, R) %>%
                 filter(`1` <= `0`) %>% slice(n = 1))$t

            missedDoses <- rep(c(rep(TRUE, worstT - 1), rep(FALSE, 1)),
                           ceiling(maxDoseT/worstT))[1:maxDoseT]

            for(i in 2:maxDoseT){

                if(missedDoses[i] == TRUE){
                    hours[[D]] <- append(hours[[D]], tail(hours[[D]], n = 1) + 0.5)

                }else{
                    hour = 0
                    hours[[D]] <- append(hours[[D]], hour)

                }
            }
        }
    }

    Rinf <- foreach(D = combs, .combine = "rbind")%do%{

         di <- druginf[[which(names(druginf) == D)]]
         mi <- mutinf[[which(names(druginf) == D)]]

         IC_50 <- di[which(names(di) == "IC_50")]
         hl <- di[which(names(di) == "hl")]
         m <- di[which(names(di) == "slope")]
         C_max <- di[which(names(di) == "c_max")]
         dosing <- di[which(names(di) == "dosing")]

         s <- as.numeric(mi[which(names(mi) == "s")])
         mu <- as.numeric(mi[which(names(mi) == "mu")])
         rho <- as.numeric(mi[which(names(mi) == "rho")])
         sigma <- as.numeric((mi[which(names(mi) == "sigma")]))

        mut <- 1
                                        #clunky
         bind_rows(bind_cols(t =  1:maxDoseT, nhours = hours[[D]], mut = rep(1, maxDoseT)),
              bind_cols(t =  1:maxDoseT, nhours = hours[[D]], mut = rep(0, maxDoseT)))   %>%
              mutate(drug = D) %>%
              mutate( R = (1 - mut * s)/(1 + (decay_instant(nhours, hl, C_max )/(IC_50 * (1 - mut*(1 - rho))))^(m * (1 + mut * (sigma)))))

    }

    toTrack <- tbl_df(expand.grid(A = c(1, 0), B = c(1, 0), C = c(1, 0)))
    names(toTrack) = combs
    druginds <- toTrack %>% mutate(ind = 1:n()) %>% gather(drug, mut, -ind)

   #Now, we go through and compute R at each of these timepoints for each of the resistance profiles
    allRs <- left_join(druginds, Rinf, by = c("drug", "mut")) %>%
        group_by(ind, t) %>% summarize(R = Rscale * prod(R), .groups = 'drop')

    arsplit <- allRs %>% group_split(t)

    N <- Ne
    A_i <- 3000 #3000
    lambda <- Ne * 1 * 10/(10 - 1)
    dy <- 1
    dt <- .02
    mu <- .00001

    mutinf <- foreach(D = combs, .combine = 'rbind')%do%{

        mi <- mutinf[[which(names(mutinf) == D)]]
        s <- mi[which(names(mi) == "s")]
        mu <- mi[which(names(mi) == "mu")]
        rho <- mi[which(names(mi) == "rho")]
        sigma <- mi[which(names(mi) == "sigma")]

        return(c(D, mu, s, sigma, rho))
    }

    colnames(mutinf) <- c("drug", "mu", "s", "sigma", "rho")
    mutinf <- tbl_df(mutinf)

    init.n <- left_join(druginds, mutinf, by = "drug") %>%
        mutate(mu_over_s = as.numeric(mu)/as.numeric(s) ) %>%
            group_by(ind) %>% mutate(mu_over_s = ifelse(mut == 1, mu_over_s, 1)) %>%
                summarize(lam = prod(mu_over_s), .groups = "drop")  %>%
                    mutate(i.N = round(lam*N))

    tmp <- init.n %>% filter(ind < 8)
    init.n <- init.n %>% mutate(i.N = ifelse(ind == 8, N - sum(tmp$i.N), i.N)) %>%
        mutate(lam = ifelse(ind == 8, 1 - sum(tmp$lam), lam))

                                        #new cells
    ys <- matrix(0, nrow = 8, ncol = maxDoseT + 1)
    ys[,1] <- init.n$i.N

    toMerge <- toTrack %>% mutate(ind = 1:n())
    names(toMerge) <- c("m1", "m2", "m3", "ind")

    ref <- bind_rows(
        bind_cols(ind = 0, name = "m1"),
        bind_cols(ind = 0, name = "m2"),
        bind_cols(ind = 0, name = "m3"))

    start_time <- Sys.time()
    i <- 1

    res <- init.n$lam
    while(i < maxDoseT & ! stopCondition(i, ys[,i]) ){

        emergingCells <- rpois(8, A_i * dt * res)

        rs <- arsplit[[i]]$R

        yvals <- ys[,i]

        denom <- lambda + sum(( dy * rs * yvals))
        rates <- (yvals*dy*lambda*rs*dt)/denom

        newys <- rpois(8, rates)

        mutFrom.s1 <- rbinom(8, newys, as.numeric(mutinf$mu[1]))
        mutFrom.s2 <- rbinom(8, newys, as.numeric(mutinf$mu[2]))
        mutFrom.s3 <- rbinom(8, newys, as.numeric(mutinf$mu[3]))

        mutFrom <- rep(0, 8)
        mutTo <- rep(0, 8)

        for(ind in 1:8){

            if(newys[ind] > 0){
                m1 <- sample(1:newys[ind], mutFrom.s1[ind], replace = FALSE)
                m2 <- sample(1:newys[ind], mutFrom.s2[ind], replace = FALSE)
                m3 <- sample(1:newys[ind], mutFrom.s3[ind], replace = FALSE)

                base <- toTrack[ind,]

                mutsTo <- bind_rows(bind_rows(
                    bind_cols(ind = m1, name = rep("m1", length(m1))),
                    bind_cols(ind = m2, name = rep("m2", length(m2))),
                    bind_cols(ind = m3, name = rep("m3", length(m3)))) %>%
                        mutate(i = 1), ref %>% mutate(i =0)) %>%
                            spread(name, i) %>% group_by(m1, m2, m3) %>%
                                summarize(n = n(), .groups = "drop")  %>%
                                    filter(!(m1 == 0 & m2 == 0 & m3 == 0)) %>%
                                        mutate(m1 = ifelse(is.na(m1), 0, m1),
                                               m2 = ifelse(is.na(m2), 0, m2),
                                               m3 = ifelse(is.na(m3), 0, m3))

                #So, we actually want something different here - we want to mutate from
                #based on the index we are considering
                #Here, ind = 5, or 110. Our mutation is the same (second position),
                #so we should go to 100, instead of 000 -> 010

                myInd <- ind
                flips <- toMerge %>% filter(ind == myInd)

                mutsTo <- mutsTo %>% mutate(m1 =  abs(flips$m1 - m1),
                                            m2 =  abs(flips$m2 - m2),
                                            m3 =  abs(flips$m3 - m3))

                mutFrom[ind] <- sum(mutsTo$n)
                mutTo <- mutTo + (left_join(toMerge, mutsTo, by = c("m1", "m2", "m3")) %>%
                                      mutate(n = ifelse(is.na(n), 0, n)))$n

            }

        }

        ys[,i + 1] <- rbinom(8, emergingCells + newys + yvals - mutFrom + mutTo, exp(-dy * dt))

        i <- i + 1

    }

    end_time <- Sys.time()
    print(end_time - start_time)

    conc_by_drug <- Rinf

    R_by_geno <- left_join(druginds, Rinf, by = c("drug", "mut")) %>%
        group_by(ind, t) %>% summarize(R = Rscale * prod(R), .groups = 'drop')

    return(list(f = ys, drugs = conc_by_drug, R = R_by_geno))
}

stopCondition <- function(t, ysi){

    #every 6 months
    if(t %% (2*24*7*4*3) == 0  ){

        if(sum(ysi[1:8]) >= 5000){
            return(TRUE)
        }
        return(FALSE)
    }
    return(FALSE)

}

##############################################
################ Code to run #################
##############################################

#Relates to parallelization - feel free to comment
ncores <- as.numeric(Sys.getenv('SLURM_CPUS_ON_NODE'))
registerDoParallel(ncores)

toTrack <- tbl_df(expand.grid(A = c(1, 0), B = c(1, 0), C = c(1, 0)))
druginds <- toTrack %>% mutate(ind = 1:n()) %>% gather(drug, mut, -ind)
indinf <- toTrack %>% mutate(comb = paste0(A, B, C)) %>%
    mutate(ind = 1:8)  %>% select(comb, ind)

#2 periods/hour, 24 hours/day, 7 days/week, 4 weeks/month, 12 months/year * 10 years
maxDoseT <- 2 * 24 * 7 * 4 * 12 * 10
ind <- 1

dirp <- "../dat/time_model/"
print(paste0("making directory: ", dirp))

system(paste0("mkdir ", dirp))

#This refers to the random nature at which doses are missed
treatmentType <- "random"
combs <- c("3TC", "D4T", "NFV")

#These nested loops draw from our cluster set up
#As a test to see if things are working properly, I'd recommend lowering
#these numbers
foreach(ind = 1:20)%dopar%{
    foreach(ind = 1:75 )%do%{

        R00 <- 10

        Ne <- floor(10^rnorm(1, 5.2, .5))
        while(Ne < 10^4 | Ne > 10^6){
            Ne <- floor(10^rnorm(1, 5.2, .5))
        }
        rseed <- sample(1:100000, 1)

        #Adherence
        pv <- runif(1, 0, 1)

        stringToPrint <- paste(paste0(combs, collapse = "-"),
                               treatmentType,
                               pv,
                               R00,
                               ind, Ne, rseed, sep  = "_")

        dat <- runTime(combs = combs, maxDoseT = maxDoseT, missp = pv,
                      Rscale = R00, doseType = treatmentType, Ne = Ne)

        relInds <- c(1, which(1:ncol(dat$f) %% (2 * 24 * 7 * 4 ) == 0))

        write.table(dat$f[,relInds], paste0(dirp, stringToPrint, "_f.csv"),
                    quote = FALSE, sep = ",",
                    row.names = FALSE, col.names = FALSE)

    }
}


#Jenny Smith

#May 8, 2017

#Purpose: Create a PCA, and MDS ordination plots given expression data and genes of interest.

plotPCoA <- function(expnData,phenovector, title="",colorCodes=NULL, geneList=NULL, size=3){
  #expnData is normalized counts, typically log2 scale.

  #factor is the name of the factor column
  suppressPackageStartupMessages(library(vegan))
  suppressPackageStartupMessages(library(ggplot2))
  suppressPackageStartupMessages(library(dplyr))

  #Ensure correct order of patients in both datasets
  expnData <- expnData[ ,intersect(names(phenovector), colnames(expnData))]
  phenovector <- phenovector[intersect(names(phenovector),colnames(expnData))]

  if (! is.null(geneList)){
    expnData <- t(expnData[geneList, ])
  }else{
    expnData <- t(expnData) #note: must remove all zero count genes or will  fail on an error
  }

  PCoA <- capscale(expnData ~ 1, distance = "bray", add=TRUE)
  scores <- data.frame(scores(PCoA, display="sites"),
                       group=phenovector) %>%
            dplyr::arrange(desc(group))

  p <- ggplot(scores, aes(x=MDS1, MDS2)) +
    geom_point(aes(color=scores[,"group"]), size=size, alpha=0.75) +
    theme_numX +
    labs(title=title)

  if(!is.null(colorCodes)){
    p <- p +
      scale_color_manual(values=colorCodes)
  }
  #If wanted to add ellipses and convex hulls use the code below

  # k <- 4
  # clust <- kmeans(scores[,1:2], k)
  # groups <- as.factor(clust$cluster)
  #
  # scores <- scores %>%
  #   mutate(group=groups,
  #          Status=clinData[rownames(expnData),factor])
  # head(scores)
  # mds.plot <- ggscatter(scores, x="MDS1", y="MDS2",
  #                       size=3.5,
  #                       color="Status",
  #                       group="group",
  #                       palette = "Set1",
  #                       ellipse =FALSE,
  #                       ellipse.type = "norm",
  #                       repel = TRUE)
  # #https://stats.stackexchange.com/questions/22805/how-to-draw-neat-polygons-around-scatterplot-regions-in-ggplot2/22855
  # find_hull <- function(df) df[chull(df$MDS1, df$MDS2), ]
  # hulls <- ddply(scores,"group",find_hull )
  #
  # #add hulls or ellipses
  # mds.plot <- mds.plot +
  #   # stat_ellipse(data = scores,
  #   # mapping = aes(group=group, fill=group), geom = "polygon", alpha=0.15, color="black", level=0.99)
  #   geom_polygon(data=hulls, aes(fill=group), alpha=0.15) +
  #   scale_fill_manual(values = rainbow_hcl(4)) +
  #   theme(text = element_text(size=20))

  # aov <- aov(scores$MDS1 ~ df[,factor]) #NOTE: This is only valid for balanced experimental designs! Equal # of obs in each factor level.

  list <- list(PCoA,scores, p)
  names(list) <- c("PCoA","scores","plot")

  return(list)
}

#from https://github.com/mikelove/DESeq2/blob/master/R/plots.R
#Want to return the whole scores matrix so can examine 3d pca plots.
plotPCA.DESeq.mod <- function(object, intgroup="condition", ntop=500, returnData=FALSE, PC3=FALSE)
{
  library(matrixStats)
  # calculate the variance for each gene
  rv <- rowVars(assay(object))

  # select the ntop genes by variance
  select <- order(rv, decreasing=TRUE)[seq_len(min(ntop, length(rv)))]

  # perform a PCA on the data in assay(x) for the selected genes
  pca <- prcomp(t(assay(object)[select,]))

  # the contribution to the total variance for each component
  percentVar <- pca$sdev^2 / sum( pca$sdev^2 )

  if (!all(intgroup %in% names(colData(object)))) {
    stop("the argument 'intgroup' should specify columns of colData(dds)")
  }

  intgroup.df <- as.data.frame(colData(object)[, intgroup, drop=FALSE])

  # add the intgroup factors together to create a new grouping factor
  group <- if (length(intgroup) > 1) {
    factor(apply( intgroup.df, 1, paste, collapse=":"))
  } else {
    colData(object)[[intgroup]]
  }

  # assembly the data for the plot - first 10 PCs
  # d <- data.frame(PC1=pca$x[,1], PC2=pca$x[,2], group=group, intgroup.df, name=colnames(object))
  n <- min(10, ncol(as.data.frame(pca$x)))
  d <- data.frame(as.data.frame(pca$x)[,1:n], group=group, intgroup.df, name=colnames(object))

  if (returnData) {
    attr(d, "percentVar") <- percentVar[1:10]
    rot <- pca$rotation[,1:10] #for first 10 PCs
    dat <- list("scores"=d,"rotation"=rot)
    return(dat)
  }

  if(PC3){
    ggplot(data=d, aes_string(x="PC1", y="PC3", color="group")) +
      geom_point(size=3, alpha=0.75) +
      xlab(paste0("PC1: ",round(percentVar[1] * 100),"% variance")) +
      ylab(paste0("PC3: ",round(percentVar[3] * 100),"% variance"))

  }else{
    ggplot(data=d, aes_string(x="PC1", y="PC2", color="group")) +
      geom_point(size=3, alpha=0.75) +
      xlab(paste0("PC1: ",round(percentVar[1] * 100),"% variance")) +
      ylab(paste0("PC2: ",round(percentVar[2] * 100),"% variance"))
    # coord_fixed()
  }
}

#Updated on 6/9/17 to use variance stabilized transformed data as input (not center scaled log2, like in princomp)
PCA <- function(expnData,phenovector,title="",round=TRUE,colorCodes=NULL,
                ntop=500,PC3=FALSE, GOI=NULL){

  suppressPackageStartupMessages(library(DESeq2))
  library(ggplot2)
  #expnData is the raw counts (not normalized) has patient IDs as colnames and genes as rownames.

  # countData <- expnData[,match(names(phenovector), colnames(expnData))]
  samples <- intersect(names(phenovector), colnames(expnData))
  countData <- expnData[,samples]
  phenovector <- phenovector[samples]

  countData <- round(countData, digits = 0)
  colData <- as.data.frame(phenovector)

  if(length(unique(phenovector)) < 2){
    dds <- DESeqDataSetFromMatrix(countData = countData,
                                  colData = colData,
                                  design = ~ 1)
  }else{
    dds <- DESeqDataSetFromMatrix(countData = countData,
                                  colData = colData,
                                  design = ~ phenovector)
  }

  #create deseq2 dataset and perform variance stabilized transformation for sample to sample comparisons
  dds <- dds[ rowSums(counts(dds)) > 10, ]
  varianceStab <- vst(dds, blind = TRUE)

  #if given a list of genes of interest
  if (! is.null(GOI)){
    GOI <- intersect(GOI, rownames(assay(varianceStab)))
  }else{
    GOI <- 1:nrow(varianceStab)
  }

  # Create a PCA plot
  plot.1 <- plotPCA.DESeq.mod(varianceStab[GOI,], intgroup = "phenovector", ntop = ntop, PC3=FALSE) +
    theme_numX +
    labs(title=title)

  if(PC3){
    plot.2 <- plotPCA.DESeq.mod(varianceStab[GOI,], intgroup = "phenovector", ntop = ntop, PC3=TRUE) +
      theme_numX +
      labs(title=title)
  }

  #change points to custom colors if provided
  if (!is.null(colorCodes)){
    plot.1 <- plot.1 +
      scale_color_manual(values=colorCodes)

    if(exists("plot.2")){
      plot.2 <- plot.2 +
        scale_color_manual(values=colorCodes)
    }

  }

  #PCA data frame with the wieghts/loadings and eigen vectors
  pca.dat <- plotPCA.DESeq.mod(varianceStab[GOI,], intgroup = "phenovector", ntop = ntop,
                               returnData=TRUE)

  #Final Results object
  res <- list(dds, varianceStab,pca.dat$scores,pca.dat$rotation, plot.1)
  names(res) <- c("dds", "vst","pca_data","pca_loadings", "pca_plot")

  if(is.character(GOI)){
    res[["GOI"]] <- GOI
  }

  if(PC3){
    res[["pca_plot2"]] <- plot.2
  }

  return(res)
}

pca_custom <- function(expnData,CDE,fillCol, colorCol, colorCode=NULL, PC3=FALSE,
                       single.col.outline=FALSE, toHighlight=NULL, ellipse=FALSE){
  library(tibble)
  library(dplyr)
  #expn data is log2, normalized counts
  #CDE has patients as rownames.
  #fillCol == character string of column name for fill colors
  #colorCol == character string of column name for border colors
  #colorCode is an option named vector of colors. If specifying color and fill manually, create a list with  length 2, with names c("fill","color")
  #single.col.outline is T/F for wether to susbet to dataframe for the border colors.
  #toHighlight is a character vector for if single.col.outline == TRUE. Gives the factor to highlight with borders from the colorColumn.
  #ellipse is T/F for an ellipse based on fillCol column
  expnData <- expnData[,intersect(rownames(CDE),colnames(expnData))]

  # print(dim(expnData))
  pca <- prcomp(t(expnData), scale=TRUE)
  summ <- summary(pca)

  scores <- as.data.frame(pca$x) %>%
    rownames_to_column("USI") %>%
    inner_join(., dplyr::select(CDE,USI=matches("USI$"), everything()), by="USI") %>%
    dplyr::select(USI,fillCol, colorCol, everything())

  #Plot function for  PC1 and either PC2 or anyother
  pca.plot_function <- function(scores,PC){

    idx <- as.numeric(gsub("[A-Za-z]{2}","", PC))

    pca.plot  <- ggplot(scores, aes_string(x="PC1", y=PC)) +
      labs(x=paste("PC1: ", round(summ$importance[2,1], digits=3)*100, "% variance"),
           y=paste(paste0(PC,": "), round(summ$importance[2,idx], digits=3)*100, "% variance")) +
      theme_numX  +
      theme(legend.text = element_text(size=14),
            legend.title = element_text(size=16))

    if(single.col.outline){
      pca.plot <- pca.plot +
        geom_point(size=5,stroke=0.2, alpha=1,shape=21,color="white",
                   aes_string(fill=fillCol)) +
        geom_point(data=subset(scores, scores[,colorCol] == toHighlight),
                   aes_string(x="PC1", y=PC, color=colorCol),
                   size=4, stroke=1, alpha=1,shape=21)

    }else{
      pca.plot <- pca.plot +
        geom_point(size=5, stroke=2, alpha=0.85,shape=21,
                   aes_string(fill=fillCol, color=colorCol))
    }

    if(!is.null(colorCode)){

        if(is.list(colorCode)){
          pca.plot <- pca.plot +
            scale_fill_manual(values=colorCode[["fill"]]) +
            scale_color_manual(values=colorCode[["color"]])
        }else{
          pca.plot <- pca.plot +
            scale_fill_manual(values=colorCode)
        }
    }

    if(ellipse){
        pca.plot <- pca.plot +
          stat_ellipse(data=scores, type="norm",
                       aes_string(x="PC1", y=PC, fill=fillCol),
                       geom="polygon", alpha=0.1)
    }

    return(pca.plot)
  }

  #PC1 and PC2 plot
  pc1.pc2 <- pca.plot_function(scores=scores,PC="PC2")

  #PC1 and PC3 plot
  if(PC3){
    pc1.pc3 <- pca.plot_function(scores=scores, PC="PC3")

  }

  res <- list("pca"=pca,"scores"=scores,"plot.1"=pc1.pc2)

  if(PC3){
    res[["plot.2"]] <- pc1.pc3
  }

  return(res)

}


# Fonction R à traduire en Python
get_pca <- function(res.pca, element = c("var", "ind")){
  elmt <- match.arg(element)
  if(elmt =="var") get_pca_var(res.pca)
  else if(elmt == "ind") get_pca_ind(res.pca)
}

get_pca_ind<-function(res.pca, ...){

  # FactoMineR package
  if(inherits(res.pca, c('PCA'))) ind <- res.pca$ind

  # ade4 package
  else if(inherits(res.pca, 'pca') & inherits(res.pca, 'dudi')){
    ind.coord <- res.pca$li
    # get the original data
    data <- res.pca$tab
    data <- t(apply(data, 1, function(x){x*res.pca$norm} ))
    data <- t(apply(data, 1, function(x){x+res.pca$cent}))
    ind <- .get_pca_ind_results(ind.coord, data, res.pca$eig,
                                res.pca$cent, res.pca$norm)
  }

  # stats package
  else if(inherits(res.pca, 'princomp')){
    ind.coord <- res.pca$scores
    data <- .prcomp_reconst(res.pca)
    ind <- .get_pca_ind_results(ind.coord, data, res.pca$sdev^2,
                                res.pca$center, res.pca$scale)

  }
  else if(inherits(res.pca, 'prcomp')){
    ind.coord <- res.pca$x
    data <- .prcomp_reconst(res.pca)
    ind <- .get_pca_ind_results(ind.coord, data, res.pca$sdev^2,
                                res.pca$center, res.pca$scale)
  }
  # ExPosition package
  else if (inherits(res.pca, "expoOutput") & inherits(res.pca$ExPosition.Data,'epPCA')){
    res <- res.pca$ExPosition.Data
    ind <- list(coord = res$fi,  cos2 = res$ri, contrib = res$ci*100)
  }
  else stop("An object of class : ", class(res.pca),
            " can't be handled by the function get_pca_ind()")

  class(ind)<-c("factoextra", "pca_ind")

  ind
}

get_pca_var<-function(res.pca){
  # FactoMineR package
  if(inherits(res.pca, c('PCA'))) var <- res.pca$var
  # ade4 package
  else if(inherits(res.pca, 'pca') & inherits(res.pca, 'dudi')){
    var <- .get_pca_var_results(res.pca$co)
  }
  # stats package
  else if(inherits(res.pca, 'princomp')){
    # Correlation of variables with the principal component
    var_cor_func <- function(var.loadings, comp.sdev){var.loadings*comp.sdev}
    var.cor <- t(apply(res.pca$loadings, 1, var_cor_func, res.pca$sdev))
    var <- .get_pca_var_results(var.cor)
  }
  else if(inherits(res.pca, 'prcomp')){
    # Correlation of variables with the principal component
    var_cor_func <- function(var.loadings, comp.sdev){var.loadings*comp.sdev}
    var.cor <- t(apply(res.pca$rotation, 1, var_cor_func, res.pca$sdev))
    var <- .get_pca_var_results(var.cor)
  }
  # ExPosition package
  else if (inherits(res.pca, "expoOutput") & inherits(res.pca$ExPosition.Data,'epPCA')){
    res <- res.pca$ExPosition.Data
    data_matrix <- res$X
    factor_scores <- res$fi
    var.coord <- var.cor <- stats::cor(res$X, res$fi) # cor(t(data_matrix), factor_scores)
    var.coord <- replace(var.coord, is.na(var.coord), 0)
    var <- list(coord = var.coord, cor = var.coord, cos2 = res$rj, contrib = res$cj*100)
  }
  else stop("An object of class : ", class(res.pca),
            " can't be handled by the function get_pca_var()")
  class(var)<-c("factoextra", "pca_var")
  var
}

# compute all the results for individuals : coord, cor, cos2, contrib
# ind.coord : coordinates of variables on the principal component
# pca.center, pca.scale : numeric vectors corresponding to the pca
# center and scale respectively
# data : the orignal data used during the pca analysis
# eigenvalues : principal component eigenvalues
.get_pca_ind_results <- function(ind.coord, data, eigenvalues, pca.center, pca.scale ){

  eigenvalues <- eigenvalues[1:ncol(ind.coord)]

  if(pca.center[1] == FALSE) pca.center <- rep(0, ncol(data))
  if(pca.scale[1] == FALSE) pca.scale <- rep(1, ncol(data))

  # Compute the square of the distance between an individual and the
  # center of gravity
  getdistance <- function(ind_row, center, scale){
    return(sum(((ind_row-center)/scale)^2))
  }
  d2 <- apply(data, 1,getdistance, pca.center, pca.scale)

  # Compute the cos2
  cos2 <- function(ind.coord, d2){return(ind.coord^2/d2)}
  ind.cos2 <- apply(ind.coord, 2, cos2, d2)

  # Individual contributions
  contrib <- function(ind.coord, eigenvalues, n.ind){
    100*(1/n.ind)*(ind.coord^2/eigenvalues)
  }
  ind.contrib <- t(apply(ind.coord, 1, contrib,  eigenvalues, nrow(ind.coord)))

  colnames(ind.coord) <- colnames(ind.cos2) <-
    colnames(ind.contrib) <- paste0("Dim.", 1:ncol(ind.coord))

  rnames <- rownames(ind.coord)
  if(is.null(rnames)) rnames <- as.character(1:nrow(ind.coord))
  rownames(ind.coord) <- rownames(ind.cos2) <- rownames(ind.contrib) <- rnames

  # Individuals coord, cos2 and contrib
  ind = list(coord = ind.coord,  cos2 = ind.cos2, contrib = ind.contrib)
  ind
}

# compute all the results for variables : coord, cor, cos2, contrib
# var.coord : coordinates of variables on the principal component
.get_pca_var_results <- function(var.coord){

  var.cor <- var.coord # correlation
  var.cos2 <- var.cor^2 # variable qualities

  # variable contributions (in percent)
  # var.cos2*100/total Cos2 of the component
  comp.cos2 <- apply(var.cos2, 2, sum)
  contrib <- function(var.cos2, comp.cos2){var.cos2*100/comp.cos2}
  var.contrib <- t(apply(var.cos2,1, contrib, comp.cos2))

  colnames(var.coord) <- colnames(var.cor) <- colnames(var.cos2) <-
    colnames(var.contrib) <- paste0("Dim.", 1:ncol(var.coord))

  # Variable coord, cor, cos2 and contrib
  list(coord = var.coord, cor = var.cor, cos2 = var.cos2, contrib = var.contrib)
}


#=========================================================
#
# Judas detection probability - discrete survival analysis
#
# Dave Ramsey (22/11/2021)
#=========================================================

library(tidyverse)
library(lubridate)
library(sf)
library(purrr)
library(cmdstanr)
library(posterior)
source("R/judas functions.r")

#==================================================
# Santiago Judas goats
#==================================================

goat_dyads<- read_csv("data/Santiago_judas_dyads.csv")
goat_dyads<- goat_dyads %>% mutate(id1 = factor(as.character(id1)), id2 = factor(as.character(id2)))
goat_ids<- read_csv("data/Santiago_judas_ID.csv")
goat_ids<- goat_ids %>% mutate(ID = factor(as.character(ID)))

id1<- as.vector(unclass(goat_dyads$id1))
id2<- as.vector(unclass(goat_dyads$id2))

dyad<- as.vector(unclass(factor(goat_dyads$Dyad)))
ndyads<- max(dyad)
N<- nrow(goat_dyads)
M<- nrow(goat_ids)
X<- model.matrix(~ distance + Sex1 + Trans1, data=goat_dyads)[,-1]
K<- ncol(X)
sex<- model.matrix(~Sex + Trans, goat_ids)[,-1]

#---- Calculate home range scale (sigma) from Judas locations ----

goat_locs<- read_csv("data/santiago_judas_locs.csv")
goat_locs<- goat_locs %>% mutate(ID = factor(as.character(ID)))
goat_locs<- goat_locs %>% group_by(ID) %>% nest() %>% arrange(ID)

goat_chr<- goat_locs %>% mutate(CNorm=map(data, calc_hr))
goat_chr<- goat_chr %>% unnest(CNorm) %>% select(-data)

maxD<- goat_chr$sigma * 2.45

#---- Discrete survival analysis -----

set.seed(123)

mod<- cmdstan_model("stan/discrete_survival.stan")

data<- list(N=N,M=M,K=K,X=X,id1=id1,id2=id2,
            time=as.integer(goat_dyads$time),
            ev=as.integer(goat_dyads$event),sex=sex,maxD=maxD)

## MCMC settings (for illustration - increase for serious inference)
ni <- 500
nt <- 1
nb <- 500
nc <- 2

inits <- lapply(1:nc, function(i)
  list(b1= -3, beta=c(-0.5,1,1,1), sigma_id1=runif(1), sigma_id2=runif(1), q=runif(N, 0, 0.1)))

out<- mod$sample(
  data=data,
  init = inits,
  iter_warmup = nb,
  iter_sampling = ni,
  chains=nc,
  parallel_chains = nc,
  refresh = 50)

out$summary(c("b1","beta","sigma_id1","sigma_id2"))

out$summary("pave")

out$save_object(file="out/goat_surv.rds")

#==================================================
# Santa Cruz Island Judas pigs
#==================================================

pig_dyads<- read_csv("data/Santacruz_judas_dyads.csv")
pig_dyads<- pig_dyads %>% mutate(id1 = factor(as.character(id1)), id2 = factor(as.character(id2)))
pig_ids<- pig_dyads %>% select(ID=id1,Sex=Sex1) %>% distinct() %>% arrange(ID)

id1<- as.vector(unclass(pig_dyads$id1))
id2<- as.vector(unclass(pig_dyads$id2))

dyad<- as.vector(unclass(factor(pig_dyads$Dyad)))
ndyads<- max(dyad)
N<- nrow(pig_dyads)
M<- nrow(pig_ids)
X<- model.matrix(~ distance + Sex1, data=pig_dyads)[,-1]
K<- ncol(X)
sex<- as.matrix(ifelse(pig_ids$Sex=="B", 0, 1))

#--- Calculate home range scale (sigma) from Judas locations
pig_locs<- read_csv("data/santacruz_judas_locs.csv")
# rescale coordinates to km
pig_locs<- pig_locs %>% mutate(ID = factor(as.character(ID)),East=East/1000, North=North/1000)
pig_locs<- pig_locs %>% group_by(ID) %>% nest() %>% arrange(ID)

pig_chr<- pig_locs %>% mutate(CNorm=map(data, calc_hr))
pig_chr<- pig_chr %>% unnest(CNorm) %>% select(-data)

maxD<- pig_chr$sigma * 2.45

#------------------------------------------------------
# Discrete survival analysis
#-----------------------------------------------------------

mod<- cmdstan_model("stan/discrete_survival.stan")

data<- list(N=N,M=M,K=K,X=X,id1=id1,id2=id2,
            time=as.integer(pig_dyads$time),
            ev=as.integer(pig_dyads$event),sex=sex,maxD=maxD)

## MCMC settings
ni <- 500
nt <- 1
nb <- 500
nc <- 2

inits <- lapply(1:nc, function(i)
  list(b1= -3, beta=c(-0.5,1), sigma_id1=runif(1), sigma_id2=runif(1), q=runif(N, 0, 0.1)))

out<- mod$sample(
  data=data,
  init = inits,
  iter_warmup = nb,
  iter_sampling = ni,
  chains=nc,
  parallel_chains = nc,
  refresh = 50)

out$summary(c("b1","beta","sigma_id1","sigma_id2"))

out$summary("pave")

out$save_object(file="out/pig_surv.rds")

#---------------------------------------------------
# Plots
#---------------------------------------------------

# Predicted probability of detection at activity centres (i.e. distance = 0)

out<- readRDS("out/goat_surv.rds")

fit<- as_draws_rvars(out$draws(c("beta0","beta")))

beta0<- draws_of(fit$beta0)
beta<- draws_of(fit$beta)
beta1<- beta[,1]
beta2<- beta[,-1]

n<- nrow(goat_ids)
XX<- model.matrix(~Sex + Trans, goat_ids)[,-1]

plam<- matrix(NA,nrow=dim(beta)[1],ncol=n)
for(i in 1:n){
  plam[,i]<- 1-exp(-exp(beta0[,i] + beta2 %*% XX[i,]))
}

plam<- as.data.frame(plam)
names(plam)<- goat_ids$ID
plam$samp<- 1:nrow(plam)

plam<- plam %>% pivot_longer(-samp, names_to = "ID", values_to = "p")
plam<- left_join(plam, goat_ids)
plam<- plam %>% mutate(Sex = factor(Sex, levels=c("Female","Super","Male")),
                       Trans = factor(Trans, labels = c("Not Translocated","Translocated")))

win.graph(10,5)
plam %>% ggplot(aes(ID, p, color=interaction(Sex,Trans))) +
  geom_boxplot(fill="grey90",outlier.color = NA) +
  labs(x="Judas ID", y="Detection probability at centre of UD") +
  theme_bw() +
  theme(axis.text.x = element_blank(),
        legend.position = "bottom",
        legend.title = element_blank())

#--------------------------------------------------------
# Predicted detection curves with distance from activity centre
#--------------------------------------------------------
fit<- as_draws_rvars(out$draws(c("beta0","beta")))

beta0<- E(fit$beta0)
beta<- E(fit$beta)
beta1<- beta[1]
beta2<- beta[-1]

n<- nrow(goat_ids)
distance<- seq(0,20,0.1)
XX<- model.matrix(~Sex + Trans, goat_ids)[,-1]

plam<- matrix(NA,nrow=length(distance), ncol=n)

for(i in 1:n){
  plam[,i]<- 1-exp(-exp(beta0[i] + beta1*distance + as.vector(XX[i,] %*% beta2)))
}
plam<- as.data.frame(plam)
names(plam)<- goat_ids$ID
plam<- plam %>% mutate(Distance=distance)

plam<- plam %>% pivot_longer(-Distance, names_to = "ID", values_to = "p")
plam<- left_join(plam, goat_ids)
plam<- plam %>% mutate(Sex = factor(Sex, levels=c("Female","Super","Male")),
                       Trans = factor(Trans, labels = c("Not Translocated","Translocated")))

win.graph(8,7)
plam %>% ggplot(aes(Distance, p, group=ID)) +
  geom_line(color="grey50") +
  facet_grid(rows=vars(Sex), cols=vars(Trans)) +
  labs(x="Distance (km)",y="Probability of association") +
  theme_bw() +
  theme(legend.position = "none",
        strip.text = element_text(face="bold"))

 #--------------------------------------------------------------------
 # Unconditional detection probability (Eqn. 2)
 #--------------------------------------------------------------------

fit<- summarise_draws(out$draws("pave"), mean, ~quantile(.x, c(0.05,0.95)))

plam<- goat_ids %>% mutate(pave= fit$mean, lcl = fit$`5%`, ucl=fit$`95%`)
plam<- plam %>% mutate(Sex = factor(Sex, levels=c("Female","Super","Male")),
                       Trans = factor(Trans, labels = c("Not Translocated","Translocated")))

win.graph(10,5)
 plam<- plam %>% arrange(pave)
 plam %>%  mutate(ID = fct_reorder(ID, pave), Sex) %>%
   ggplot(aes(x=Trans, y=pave, color=Sex)) +
   geom_pointrange(aes(ymin=lcl,ymax=ucl), position=position_dodge2(width=0.5)) +
   facet_wrap(~Sex) +
   labs(x="",y="Probability of association",color="Status") +
   theme_bw() +
   theme(legend.position = "none",
         strip.text = element_text(face="bold"),
         axis.title = element_text(face="bold", size=12),
         axis.text.x = element_text(size=11))


# Snow crab --- Areal unit modelling of habitat  -- no reliance upon stmv fields

# -------------------------------------------------
# Part 1 -- construct basic parameter list defining the main characteristics of the study
# require(aegis)

  year.assessment = 2021
  require(bio.snowcrab)   # loadfunctions("bio.snowcrab")

  # choose one:
  if (0) {
    # default: best DIC 31457.19, WAIC 32405.81
    # P(y) = Pois(y|y>0) ; p = exp(theta) / ( 1 + exp(theta) )
    family="poisson"
    carstm_model_label = "tesselation"

    # estimates suggest 30% prob of 0 overdispersion: DIC: 36196.11, WAIC 36850.04
    # P(y) = p 1_{y=0} + (1-p) Pois(y|y>0) ; pois is of positive valued only
    family="zeroinflatedpoisson0"
    carstm_model_label = "tesselation_zip0"

    # unstable ... will not complete due to singularities ... might be overparamterized
    # P(y) = p 1_{y=0} + (1-p) Pois(y) ; pois is of all values
    family="zeroinflatedpoisson1"
    carstm_model_label = "tesselation_zip1"

  }

  p = snowcrab_parameters(
    project_class="carstm",
    yrs=2000:year.assessment,
    areal_units_type="tesselation",
#    areal_units_constraint_ntarget = 20,
#    areal_units_constraint_nmin = 5,
    family=family,
    carstm_model_label = carstm_model_label,
    selection = list(type = "number")
  )

# ------------------------------------------------
# Part 2 -- spatiotemporal statistical model

  if ( spatiotemporal_model ) {

      if (0) {
        # polygon structure:: create if not yet made
        # adjust based upon RAM requirements and ncores
          # create if not yet made
        for (au in c("cfanorth", "cfasouth", "cfa4x", "cfaall" )) plot(polygon_managementareas( species="snowcrab", au))
        xydata = snowcrab.db( p=p, DS="areal_units_input", redo=TRUE )

        sppoly = areal_units( p=p, hull_alpha=15, redo=TRUE, verbose=TRUE )  # create constrained polygons with neighbourhood as an attribute
        plot( sppoly[, "npts"]  )

        MS = NULL

        # p$carstm_model_label = "tesselation_overdispersed"   # default is the name of areal_units_type
        # p$family  = "zeroinflatedpoisson0" #  "binomial",  # "nbinomial", "betabinomial", "zeroinflatedbinomial0" , "zeroinflatednbinomial0"
        # p$carstm_model_inla_control_familiy = NULL

      }

      sppoly = areal_units( p=p )  # to reload

      # -------------------------------------------------
      M = snowcrab.db( p=p, DS="carstm_inputs", redo=TRUE )  # will redo if not found
      M = NULL; gc()

      fit = carstm_model(
        p=p,
        data='snowcrab.db( p=p, DS="carstm_inputs" )',
        posterior_simulations_to_retain="predictions" ,
        scale_offsets = TRUE,  # required to stabilize  as offsets are so small, required for : inla.mode = "experimental"
        # redo_fit = FALSE,  # only to redo sims and extractions
        # toget="predictions",  # this updates a specific subset of calc
        control.inla = list( strategy='adaptive' ),  # strategy='laplace', "adaptive" int.strategy="eb"
        num.threads="4:2"
      )

      if (0) {
        # control.compute=list(smtp="default", dic=TRUE, waic=TRUE, cpo=FALSE, config=TRUE, return.marginals.predictor=TRUE)
        # control.fixed=list(prec=1,prec.intercept=1)
        # 151 configs and long optim .. 19 hrs
        # fit = carstm_model( p=p, DS="carstm_modelled_fit")

        # extract results
        # very large files .. slow
          fit = carstm_model( p=p, DS="carstm_modelled_fit" )  # extract currently saved model fit
          fit$summary$dic$dic
          fit$summary$dic$p.eff

          plot(fit)
          plot(fit, plot.prior=TRUE, plot.hyperparameters=TRUE, plot.fixed.effects=FALSE )

      }

      res = carstm_model( p=p, DS="carstm_modelled_summary"  ) # to load currently saved results

      map_centre = c( (p$lon0+p$lon1)/2 - 0.5, (p$lat0+p$lat1)/2 -0.8 )
      map_zoom = 5

      plot_crs = p$aegis_proj4string_planar_km

      require(tmap)

      additional_features =
        tm_shape( aegis.polygons::area_lines.db( DS="cfa.regions", returntype="sf", project_to=plot_crs ), projection=plot_crs ) +
          tm_lines( col="slategray", alpha=0.75, lwd=2)   +
        tm_shape( aegis.bathymetry::isobath_db(  depths=c( seq(0, 400, by=50), 1000), project_to=plot_crs  ), projection=plot_crs ) +
          tm_lines( col="slategray", alpha=0.5, lwd=0.5) +
        tm_shape( aegis.coastline::coastline_db( DS="eastcoast_gadm", project_to=plot_crs ), projection=plot_crs ) +
          tm_polygons( col="lightgray", alpha=0.5 , border.alpha =0.5)

      (additional_features)

      vn=c( "random", "space", "combined" )
      vn=c( "random", "spacetime", "combined" )
      vn="predictions"  # numerical density (km^-2)

      tmatch="2015"

      # densities
      carstm_map(  res=res, vn=vn, tmatch=tmatch,
          palette="-RdYlBu",
          plot_elements=c(  "compass", "scale_bar", "legend" ),
          additional_features=additional_features,
          tmap_zoom= c(map_centre, map_zoom),
          title =paste( vn, paste0(tmatch, collapse="-"), "no/km^2"  )
      )

      # map all :
      outputdir = file.path( p$modeldir, p$carstm_model_label, "predicted.numerical.densitites" )
      if ( !file.exists(outputdir)) dir.create( outputdir, recursive=TRUE, showWarnings=FALSE )

      vn="predictions"
      toplot = carstm_results_unpack( res, vn )
      brks = pretty(  quantile(toplot[,,"mean"], probs=c(0,0.975), na.rm=TRUE )  )

      for (y in res$time ){
          tmatch = as.character(y)
          fn_root = paste("Predicted_numerical_abundance", paste0(tmatch, collapse="-"), sep="_")
          fn = file.path( outputdir, paste(fn_root, "png", sep=".") )

            o = carstm_map(  res=res, vn=vn, tmatch=tmatch,
              breaks =brks,
              palette="-RdYlBu",
              plot_elements=c(    "compass", "scale_bar", "legend" ),
              additional_features=additional_features,
              title=paste("Predicted numerical density (no./km^2) ", paste0(tmatch, collapse="-") ),
              map_mode="plot",
              scale=0.75,
              outformat="tmap",
              outfilename=fn
            )

      }

      fit = meanweights_by_arealunit_modelled( p=p, redo=TRUE, returntype="carstm_modelled_fit" )  ## used in carstm_output_compute

         if (0) {
          fit = meanweights_by_arealunit_modelled( p=p, returntype="carstm_modelled_fit" )

          res = meanweights_by_arealunit_modelled( p=p, returntype="carstm_modelled_summary" )  ## used in carstm_output_compute

          map_centre = c( (p$lon0+p$lon1)/2 - 0.5, (p$lat0+p$lat1)/2 -0.8 )
          map_zoom = 5

          plot_crs = p$aegis_proj4string_planar_km

          additional_features =
            tm_shape( aegis.polygons::area_lines.db( DS="cfa.regions", returntype="sf", project_to=plot_crs ), projection=plot_crs ) +
              tm_lines( col="slategray", alpha=0.75, lwd=2)   +
            tm_shape( aegis.bathymetry::isobath_db( depths=c( seq(50, 400, by=50), 500), project_to=plot_crs  ), projection=plot_crs ) +
              tm_lines( col="slategray", alpha=0.5, lwd=0.5) +
            tm_shape( aegis.coastline::coastline_db( DS="eastcoast_gadm", project_to=plot_crs ), projection=plot_crs ) +
              tm_polygons( col="lightgray", alpha=0.5 , border.alpha =0.5)

          (additional_features)

          vn="predictions"
          tmatch = "2020"

          outputdir = file.path( p$modeldir, p$carstm_model_label, "predicted.mean.weight" )
          if ( !file.exists(outputdir)) dir.create( outputdir, recursive=TRUE, showWarnings=FALSE )

          fn_root = paste("Predicted_mean_size", paste0(tmatch, collapse="-"), sep="_")

          fn = file.path( outputdir, paste(fn_root, "png", sep=".") )

          carstm_map(  res=res, vn=vn, tmatch=tmatch,
              palette="-RdYlBu",
              plot_elements=c(  "compass", "scale_bar", "legend" ),
              additional_features=additional_features,
              tmap_zoom= c(map_centre, map_zoom),
              title =paste("Predicted  mean weight of individual (kg)", paste0(tmatch, collapse="-") )
          )

        }

      snowcrab.db(p=p, DS="carstm_output_compute" )

      RES = snowcrab.db(p=p, DS="carstm_output_timeseries" )

      bio = snowcrab.db(p=p, DS="carstm_output_spacetime_biomass" )
      num = snowcrab.db(p=p, DS="carstm_output_spacetime_number" )

      outputdir = file.path( p$modeldir, p$carstm_model_label, "aggregated_biomass_timeseries" )

      if ( !file.exists(outputdir)) dir.create( outputdir, recursive=TRUE, showWarnings=FALSE )

      ( fn = file.path( outputdir, "cfa_all.png") )
      png( filename=fn, width=3072, height=2304, pointsize=12, res=300 )
        plot( cfaall ~ yrs, data=RES, lty="solid", lwd=4, pch=20, col="slateblue", type="b", ylab="Biomass index (kt)", xlab="")
        lines( cfaall_lb ~ yrs, data=RES, lty="dotted", lwd=2, col="slategray" )
        lines( cfaall_ub ~ yrs, data=RES, lty="dotted", lwd=2, col="slategray" )
      dev.off()

      ( fn = file.path( outputdir, "cfa_south.png") )
      png( filename=fn, width=3072, height=2304, pointsize=12, res=300 )
        plot( cfasouth ~ yrs, data=RES, lty="solid", lwd=4, pch=20, col="slateblue", type="b", ylab="Biomass index (kt)", xlab="")
        lines( cfasouth_lb ~ yrs, data=RES, lty="dotted", lwd=2, col="slategray" )
        lines( cfasouth_ub ~ yrs, data=RES, lty="dotted", lwd=2, col="slategray" )
      dev.off()

      ( fn = file.path( outputdir, "cfa_north.png") )
      png( filename=fn, width=3072, height=2304, pointsize=12, res=300 )
        plot( cfanorth ~ yrs, data=RES, lty="solid", lwd=4, pch=20, col="slateblue", type="b", ylab="Biomass index (kt)", xlab="")
        lines( cfanorth_lb ~ yrs, data=RES, lty="dotted", lwd=2, col="slategray" )
        lines( cfanorth_ub ~ yrs, data=RES, lty="dotted", lwd=2, col="slategray" )
      dev.off()

      ( fn = file.path( outputdir, "cfa_4x.png") )
      png( filename=fn, width=3072, height=2304, pointsize=12, res=300 )
        plot( cfa4x ~ yrs, data=RES, lty="solid", lwd=4, pch=20, col="slateblue", type="b", ylab="Biomass index (kt)", xlab="")
        lines( cfa4x_lb ~ yrs, data=RES, lty="dotted", lwd=2, col="slategray" )
        lines( cfa4x_ub ~ yrs, data=RES, lty="dotted", lwd=2, col="slategray" )
      dev.off()

      # map it ..mean density

      sppoly = areal_units( p=p )  # to reload

      vn = paste("biomass", "predicted", sep=".")

      outputdir = file.path( p$modeldir, p$carstm_model_label, "predicted.biomass.densitites" )

      if ( !file.exists(outputdir)) dir.create( outputdir, recursive=TRUE, showWarnings=FALSE )

      B = apply( bio, c(1,2), mean )

      brks = pretty( quantile( B[], probs=c(0,0.975) )* 10^6 )

      for (i in 1:length(p$yrs) ){
        y = as.character( p$yrs[i] )
        sppoly[,vn] = B[,y]* 10^6
        fn = file.path( outputdir , paste( "biomass", y, "png", sep=".") )

          carstm_map(  sppoly=sppoly, vn=vn,
            breaks=brks,
            additional_features=additional_features,
            title=paste("Predicted biomass density", y ),
            outfilename=fn
          )
      }

      plot( fit, plot.prior=TRUE, plot.hyperparameters=TRUE, plot.fixed.effects=FALSE )
      plot( fit$marginals.hyperpar$"Phi for space_time", type="l")  # posterior distribution of phi nonspatial dominates
      plot( fit$marginals.hyperpar$"Precision for space_time", type="l")
      plot( fit$marginals.hyperpar$"Precision for setno", type="l")

    (res$summary)

  }

  ##########

  if (fishery_model) {

    # you need a stan installation on your system as well (outside of R), and the R-interface "cmdstanr":
    # install.packages("cmdstanr", repos = c("https://mc-stan.org/r-packages/", getOption("repos")))

    require(cmdstanr)

    p$fishery_model = fishery_model( DS = "logistic_parameters", p=p, tag=p$areal_units_type )
    p$fishery_model$stancode = stan_initialize( stan_code=fishery_model( p=p, DS="stan_surplus_production" ) )
    str( p$fishery_model)
    p$fishery_model$stancode$compile()
    to_look = c("K", "r", "q", "qc", "logtheta")
#    to_look = c("K", "r", "q", "qc" )

      if (0) {
        # testing other samplers and optimizsers ... faster , good for debugging

        # (penalized) maximum likelihood estimate (MLE)
        fit_mle =  p$fishery_model$stancode$optimize(data =p$fishery_model$standata, seed = 123)
        fit_mle$summary( to_look )
        u = stan_extract( as_draws_df(fit_mle$draws() ) )

        mcmc_hist(fit$draws("K")) + vline_at(fit_mle$mle(), size = 1.5)

        # Variational Bayes
        fit_vb = p$fishery_model$stancode$variational( data =p$fishery_model$standata, seed = 123, output_samples = 4000)
        fit_vb$summary(to_look)
        fit_vb$cmdstan_diagnose()
        fit_vb$cmdstan_summary()

        u = stan_extract( as_draws_df(fit_vb$draws() ) )

        bayesplot_grid(
          mcmc_hist(fit$draws("K"), binwidth = 0.025),
          mcmc_hist(fit_vb$draws("K"), binwidth = 0.025),
          titles = c("Posterior distribution from MCMC", "Approximate posterior from VB")
        )

        color_scheme_set("gray")
        mcmc_dens(fit$draws("K"), facet_args = list(nrow = 3, labeller = ggplot2::label_parsed ) ) + facet_text(size = 14 )
        # mcmc_hist( fit$draws("K"))

        # obtain mcmc samples from vb solution
        res_vb = fishery_model(
          DS="logistic_model",
          p=p,
          tag=p$areal_units_type,
          fit = fit_vb
        )

        names(res_vb$mcmc)

      }

      fit = p$fishery_model$stancode$sample(
        data=p$fishery_model$standata,
        iter_warmup = 15000,
        iter_sampling = 10000,
        seed = 123,
        chains = 3,
        parallel_chains = 3,  # The maximum number of MCMC chains to run in parallel.
        max_treedepth = 16,
        adapt_delta = 0.99,
        refresh = 1000
      )

      if (0) {
        fit = fishery_model( p=p,   DS="fit", tag=p$areal_units_type )  # to load samples (results)
        fit$summary(c("K", "r", "q", "qc"))
        print( fit, max_rows=30 )
        fit$cmdstan_diagnose()
        fit$cmdstan_summary()

      }

      fit$summary(to_look)

      # save fit and get draws
      res = fishery_model( p=p, DS="logistic_model", tag=p$areal_units_type, fit=fit )       # from here down are params for cmdstanr::sample()

      # frequency density of key parameters
      fishery_model( DS="plot", vname="K", res=res )
      fishery_model( DS="plot", vname="r", res=res )
      fishery_model( DS="plot", vname="q", res=res, xrange=c(0.5, 2.5))
      fishery_model( DS="plot", vname="FMSY", res=res  )
      # fishery_model( DS="plot", vname="bosd", res=res  )
      # fishery_model( DS="plot", vname="bpsd", res=res  )

      # timeseries
      fishery_model( DS="plot", type="timeseries", vname="biomass", res=res  )
      fishery_model( DS="plot", type="timeseries", vname="fishingmortality", res=res)

      # Summary table of mean values for inclusion in document
      biomass.summary.table(x)

      # Harvest control rules
      fishery_model( DS="plot", type="hcr", vname="default", res=res  )
      fishery_model( DS="plot", type="hcr", vname="simple", res=res  )

      # diagnostics
      # fishery_model( DS="plot", type="diagnostic.errors", res=res )
      # fishery_model( DS="plot", type="diagnostic.phase", res=res  )

      NN = res$p$fishery_model$standata$N

      # bosd
      plot.new()
      layout( matrix(c(1,2,3), 3, 1 ))
      par(mar = c(4.4, 4.4, 0.65, 0.75))
      for (i in 1:3) plot(density(res$mcmc$bosd[,i] ), main="")
      ( qs = apply(  res$mcmc$bosd[,], 2, quantile, probs=c(0.025, 0.5, 0.975) ) )

      # bpsd
      plot.new()
      layout( matrix(c(1,2,3), 3, 1 ))
      par(mar = c(4.4, 4.4, 0.65, 0.75))
      for (i in 1:3) plot(density(res$mcmc$bpsd[,i] ), main="")
      ( qs = apply(  res$mcmc$bpsd[,], 2, quantile, probs=c(0.025, 0.5, 0.975) ) )

      # rem_sd
      plot.new()
      layout( matrix(c(1,2,3), 3, 1 ))
      par(mar = c(4.4, 4.4, 0.65, 0.75))
      for (i in 1:3) plot(density(res$mcmc$rem_sd[,i] ), main="")
      ( qs = apply(  res$mcmc$rem_sd[,], 2, quantile, probs=c(0.025, 0.5, 0.975) ) )

    # qc
      plot.new()
      layout( matrix(c(1,2,3), 3, 1 ))
      par(mar = c(4.4, 4.4, 0.65, 0.75))
      for (i in 1:3) plot(density(res$mcmc$qc[,i] ), main="")
      ( qs = apply(  res$mcmc$qc[,], 2, quantile, probs=c(0.025, 0.5, 0.975) ) )

     # b0
      plot.new()
      layout( matrix(c(1,2,3), 3, 1 ))
      par(mar = c(4.4, 4.4, 0.65, 0.75))
      for (i in 1:3) plot(density(res$mcmc$b0[,i] ), main="")
      ( qs = apply(  res$mcmc$b0[,], 2, quantile, probs=c(0.025, 0.5, 0.975) ) )

      # K
      plot.new()
      layout( matrix(c(1,2,3), 3, 1 ))
      par(mar = c(4.4, 4.4, 0.65, 0.75))
      for (i in 1:3) plot(density(res$mcmc$K[,i] ), main="")
      ( qs = apply(  res$mcmc$K[,], 2, quantile, probs=c(0.025, 0.5, 0.975) ) )

      # R
      plot.new()
      layout( matrix(c(1,2,3), 3, 1 ))
      par(mar = c(4.4, 4.4, 0.65, 0.75))
      for (i in 1:3) plot(density(res$mcmc$r[,i] ), main="")
      ( qs = apply(  res$mcmc$r[,], 2, quantile, probs=c(0.025, 0.5, 0.975) ) )

      # q
      plot.new()
      layout( matrix(c(1,2,3), 3, 1 ))
      par(mar = c(4.4, 4.4, 0.65, 0.75))
      for (i in 1:3) plot(density(res$mcmc$q[,i] ), main="")
      ( qs = apply(  res$mcmc$q[,], 2, quantile, probs=c(0.025, 0.5, 0.975) ) )

      # FMSY
      plot.new()
      layout( matrix(c(1,2,3), 3, 1 ))
      par(mar = c(4.4, 4.4, 0.65, 0.75))
      for (i in 1:3) plot(density(res$mcmc$FMSY[,i] ), main="")
      ( qs = apply(  res$mcmc$FMSY[,], 2, quantile, probs=c(0.025, 0.5, 0.975) ) )

      # densities of biomass estimates for the year.assessment
      plot.new()
      layout( matrix(c(1,2,3), 3, 1 ))
      par(mar = c(4.4, 4.4, 0.65, 0.75))
      for (i in 1:3) plot(density(res$mcmc$B[,NN,i] ), main="")
      ( qs = apply(  res$mcmc$B[,NN,], 2, quantile, probs=c(0.025, 0.5, 0.975) ) )

      # densities of biomass estimates for the previous year
      plot.new()
      layout( matrix(c(1,2,3), 3, 1 ))
      par(mar = c(4.4, 4.4, 0.65, 0.75))
      for (i in 1:3) plot(density( res$mcmc$B[,NN-1,i] ), main="")
      ( qs = apply(  res$mcmc$B[,NN-1,], 2, quantile, probs=c(0.025, 0.5, 0.975) ) )

      # densities of F in assessment year
      plot.new()
      layout( matrix(c(1,2,3), 3, 1 ))
      par(mar = c(4.4, 4.4, 0.65, 0.75))
      for (i in 1:3) plot(density(  res$mcmc$F[,NN,i] ), xlim=c(0.01, 0.6), main="")
      ( qs = apply(  res$mcmc$F[,NN,], 2, quantile, probs=c(0.025, 0.5, 0.975) ) )
      ( qs = apply(  res$mcmc$F[,NN,], 2, mean ) )

      # densities of F in previous year
      plot.new()
      layout( matrix(c(1,2,3), 3, 1 ))
      par(mar = c(4.4, 4.4, 0.65, 0.75))
      for (i in 1:3) plot(density(  res$mcmc$F[,NN-1,i] ), xlim=c(0.01, 0.6), main="")
      ( qs = apply(  res$mcmc$F[,NN-1,], 2, quantile, probs=c(0.025, 0.5, 0.975) ) )
      ( qs = apply(  res$mcmc$F[,NN-1,], 2, mean ) )

      # F for table ---
      summary( res$mcmc$F, median)
  }

# end


Ghuber <- function (u, k = 30, deriv = 0)
{
    if (!deriv)
    {
        return(pmin(1, k/abs(u)))
    } else {
        return(abs(u) <= k)
    }
}

#' Univariate LDSC
#'
#' Imported here to help estimate sample overlap between studies
#'
#' @param Z summary Z-statistics for M variants
#' @param r2 average reference LD scores for M variants
#' @param N GWAS sample size for each variant (could be different across variants)
#' @param W variant weight
#'
#' @keywords internal
#' @return model fit
ldsc_h2_internal <- function(Z, r2, N, W=NULL)
{
    if (is.null(W))
    {
        W <- rep(1, length(Z))
    }
    tau <- (mean(Z^2) - 1) / mean(N * r2)
    Wv <- 1 / (1 + tau * N * r2)^2
    id <- which(Z^2 > 30)
    if (length(id) > 0)
    {
        Wv[id] <- sqrt(Wv[id])
    }
    mod <- MASS::rlm(I(Z^2) ~ I(N * r2), weight = W * Wv,
        psi = Ghuber, k = 30)
    return(summary(mod))
}

#' Bivariate LDSC
#'
#' Imported here to help estimate sample overlap between studies
#'
#' @param Zs Mx2 matrix of summary Z-statistics for M variants from two GWAS
#' @param r2 average reference LD scores for M variants
#' @param N1 sample size for the 1st GWAS
#' @param N2 sample size for the 2nd GWAS
#' @param Nc overlapped sample size between the two GWAS
#' @param W variant weight
#' @param h1 hsq for trait 1
#' @param h2 hsq for trait 2
#'
#' @return List of models
#' @references
#' Bulik-Sullivan,B.K. et al. (2015) An atlas of genetic correlations across human diseases and traits. Nat. Genet. 47, 1236–1241.
#'
#' Guo,B. and Wu,B. (2018) Principal component based adaptive association test of multiple traits using GWAS summary statistics. bioRxiv 269597; doi: 10.1101/269597
#'
#' Gua,B. and Wu,B. (2019) Integrate multiple traits to detect novel trait-gene association using GWAS summary data with an adaptive test approach. Bioinformatics. 2019 Jul 1;35(13):2251-2257. doi: 10.1093/bioinformatics/bty961.
#'
#' https://github.com/baolinwu/MTAR
#' @keywords internal
ldsc_rg_internal <- function(Zs, r2, h1, h2, N1, N2, Nc=0, W=NULL)
{
    if(is.null(W))
    {
        W = rep(1,length(r2))
    }

    Y <- Zs[,1] * Zs[,2]

    X <- (sqrt(N1) * sqrt(N2) + sqrt(Nc/N1 * N2)) * r2
    N1r2 <- N1 * r2
    N2r2 <- N2 * r2
    r0 <- 0

    ## 1st round
    if(any(Nc > 0))
    {
        rcf <- as.vector(MASS::rlm(Y ~ X, psi = Ghuber)$coef)
        r0 <- rcf[1]
        gv <- rcf[-1]
    } else {
        gv <- as.vector(MASS::rlm(Y ~ X-1, psi = Ghuber)$coef)
    }

    ## 2nd round
    Wv <- 1 / ((h1 * N1r2 + 1) * (h2 * N2r2 + 1) + (X * gv + r0)^2)
    id <- which(abs(Zs[,1] * Zs[,2]) > 30)
    if(length(id) > 0)
    {
        Wv[id] <- sqrt(Wv[id])
    }

    if(any(Nc > 0))
    {
        rcf <- MASS::rlm(Y ~ X, weight = W * Wv, psi = Ghuber, k = 30)
    } else {
        rcf <- MASS::rlm(Y ~ X - 1, weight = W * Wv, psi = Ghuber, k = 30)
    }
    return(summary(rcf))
}

#' Univariate LDSC
#'
#' Imported here to help estimate sample overlap between studies
#'
#' @param id ID to analyse
#' @param ancestry ancestry of traits 1 and 2 (AFR, AMR, EAS, EUR, SAS) or 'infer' (default) in which case it will try to guess based on allele frequencies
#' @param snpinfo Output from ieugwasr::afl2_list("hapmap3"), or NULL for it to be done automatically
#' @param splitsize How many SNPs to extract at one time. Default=20000
#'
#' @export
#' @return model fit
#' @references
#' Bulik-Sullivan,B.K. et al. (2015) An atlas of genetic correlations across human diseases and traits. Nat. Genet. 47, 1236–1241.
#'
#' Guo,B. and Wu,B. (2018) Principal component based adaptive association test of multiple traits using GWAS summary statistics. bioRxiv 269597; doi: 10.1101/269597
#'
#' Gua,B. and Wu,B. (2019) Integrate multiple traits to detect novel trait-gene association using GWAS summary data with an adaptive test approach. Bioinformatics. 2019 Jul 1;35(13):2251-2257. doi: 10.1093/bioinformatics/bty961.
#'
#' https://github.com/baolinwu/MTAR
ldsc_h2 <- function(id, ancestry="infer", snpinfo = NULL, splitsize=20000)
{
    if(is.null(snpinfo))
    {
        snpinfo <- ieugwasr::afl2_list("hapmap3")
    }

    snpinfo <- snpinfo %>%
        dplyr::filter(complete.cases(.))

    d <- extract_split(snpinfo$rsid, id, splitsize) %>%
        ieugwasr::fill_n() %>%
        dplyr::mutate(z = beta / se) %>%
        dplyr::select(rsid, z = z, n = n, eaf) %>%
        dplyr::filter(complete.cases(.))

    stopifnot(nrow(d) > 0)

    if(ancestry == "infer")
    {
        ancestry <- ieugwasr::infer_ancestry(d, snpinfo)$pop[1]
    }

    d <- snpinfo %>%
        dplyr::select(rsid, l2=paste0("L2.", ancestry)) %>%
        dplyr::inner_join(., d, by="rsid") %>%
        dplyr::filter(complete.cases(.))

    return(ldsc_h2_internal(d$z, d$l2, d$n))
}

#' Bivariate LDSC
#'
#' Imported here to help estimate sample overlap between studies
#'
#' @param id1 ID 1 to analyse
#' @param id2 ID 2 to analyse
#' @param ancestry ancestry of traits 1 and 2 (AFR, AMR, EAS, EUR, SAS) or 'infer' (default) in which case it will try to guess based on allele frequencies
#' @param snpinfo Output from ieugwasr::afl2_list("hapmap3"), or NULL for it to be done automatically
#' @param splitsize How many SNPs to extract at one time. Default=20000
#'
#' @export
#' @return model fit
ldsc_rg <- function(id1, id2, ancestry="infer", snpinfo = NULL, splitsize=20000)
{
    if(is.null(snpinfo))
    {
        snpinfo <- ieugwasr::afl2_list("hapmap3")
    }

    x <- extract_split(snpinfo$rsid, c(id1, id2), splitsize)
    d1 <- subset(x, id == id1) %>%
        ieugwasr::fill_n() %>%
        dplyr::mutate(z = beta / se) %>%
        dplyr::select(rsid, z1 = z, n1 = n, eaf) %>%
        dplyr::filter(complete.cases(.))

    stopifnot(nrow(d1) > 0)

    d2 <- subset(x, id == id2) %>%
        ieugwasr::fill_n() %>%
        dplyr::mutate(z = beta / se) %>%
        dplyr::select(rsid, z2 = z, n2 = n, eaf) %>%
        dplyr::filter(complete.cases(.))

    stopifnot(nrow(d2) > 0)

    if(ancestry == "infer")
    {
        ancestry1 <- ieugwasr::infer_ancestry(d1, snpinfo)
        ancestry2 <- ieugwasr::infer_ancestry(d2, snpinfo)
        if(ancestry1$pop[1] != ancestry2$pop[1])
        {
            stop("d1 ancestry is ", ancestry1$pop[1], " and d2 ancestry is ", ancestry2$pop[1])
        }
        ancestry <- ancestry1$pop[1]
    }

    d1 <- snpinfo %>%
        dplyr::select(rsid, l2=paste0("L2.", ancestry)) %>%
        dplyr::inner_join(., d1, by="rsid")

    d2 <- snpinfo %>%
        dplyr::select(rsid, l2=paste0("L2.", ancestry)) %>%
        dplyr::inner_join(., d2, by="rsid")

    h1 <- ldsc_h2_internal(d1$z1, d1$l2, d1$n1, d1$l2)
    h2 <- ldsc_h2_internal(d2$z2, d2$l2, d2$n2, d1$l2)

    dat <- dplyr::inner_join(d1, d2, by="rsid") %>%
        dplyr::mutate(
            l2 = l2.x,
            n1 = as.numeric(n1),
            n2 = as.numeric(n2),
            rhs = l2 * sqrt(n1 * n2)
        )

    gcov <- dat %>%
        {
            ldsc_rg_internal(
                Zs = cbind(.$z1, .$z2),
                r2 = .$l2,
                h1 = h1$coefficients[2,1] * nrow(d1),
                h2 = h2$coefficients[2,1] * nrow(d2),
                N1 = .$n1,
                N2 = .$n2,
                W = .$l2
            )
        }
    return(list(
        gcov = gcov,
        h1=h1,
        h2=h2,
        rg = (gcov$coefficients[1,1] * nrow(dat)) / sqrt(h1$coefficients[2,1] * nrow(d1) * h2$coefficients[2,1] * nrow(d2))
    ))
}

extract_split <- function(snplist, id, splitsize=20000)
{
    nsplit <- round(length(snplist)/splitsize)
    split(snplist, 1:nsplit) %>%
        pbapply::pblapply(., function(x)
        {
            ieugwasr::associations(x, id, proxies=FALSE)
        }) %>% dplyr::bind_rows()
}