mod.nat <- function(dat.nat, varout = NULL, varout.mn = NULL, runmod = F) {

    ## set up data for a bayes run
    bayeset <- function(envdata, df.chl, varname = c("ptl", "chl", "huc8")) {

        df.chl <- na.omit(df.chl)
        envdata <- na.omit(envdata)
        ## find hucs with both TP and chl data
        huc.all <- c(unique(envdata[, varname[3]]),
                     unique(df.chl[, varname[3]]))
        huc.p <- huc.all[duplicated(huc.all)]
        cat("Number of hucs:", length(huc.p), "\n")
        ## select samples collected from the selected hucs
        incvec <- rep(F, times = nrow(envdata))
        for (i in huc.p) incvec <- incvec | envdata[, varname[3]] == i
        envdata <- envdata[incvec,]
        incvec <- rep(F, times = nrow(df.chl))
        for (i in huc.p) incvec <- incvec|df.chl[, varname[3]]== i
        df.chl <- df.chl[incvec,]
        ## make huc number a factor for running in stan
        envdata[, varname[3]] <- factor(envdata[, varname[3]])
        df.chl[, varname[3]] <- factor(df.chl[, varname[3]])
        envdata$hucnum <- as.numeric(envdata[, varname[3]])
        df.chl$hucnum <- as.numeric(df.chl[, varname[3]])

        ## center all the data prior to running stan
        mn.p <- mean(log(envdata[, varname[1]]))
        envdata$x.cent <- log(envdata[, varname[1]]) - mn.p
        mn.chl <- mean(log(df.chl[, varname[2]]))
        df.chl$y.cent <- log(df.chl[, varname[2]]) - mn.chl
        return(list(envdata, df.chl, mn.p, mn.chl))
    }

    ## stan code for fitting hierachical bayesian model
    modstan <- '
      data {
         int np;                 // number of TP measurements
         vector[np] p;           // TP measurements
         int hucnum1[np];        // huc number associated with each TP
         int nc;                 // number of Chl measurements
         vector[nc] chl;         // Chl measurements
         int hucnum2[nc];        // huc number associated with each Chl
         int nhuc;               // number of hucs
     }
     parameters {
         real mu_p;
         real<lower = 0> sigp[2];
         vector[nhuc] etap;
         real b1;
         real b2;
         real<lower = 0> k;
         real p0;
         vector[nhuc] chuc;
         real<lower = 0> sigchl[2];
    }
     transformed parameters {
         vector[nhuc] phuc;        // estimate of TP[huc]
         vector[nhuc] chlmn;       // estimate of Chl[huc]
         phuc = mu_p + etap*sigp[1];   // model linking TP[huc] to grand mean

         // logistic model linking TP[huc] to Chl[huc]
         for (i in 1:nhuc)
             chlmn[i] = b1 + (b2-b1)/(1+exp(-k*(phuc[i] - p0)));  
     }
     model {
         b1 ~ normal(-1.5,2);    // weakly informative priors
         b2 ~ normal(1.5,2);
         k ~ normal(1, 1);
         p0 ~ normal(0,2);
         mu_p ~ normal(0,3);
         sigp ~ cauchy(0,3);
         etap ~ std_normal();
         sigchl ~ cauchy(0,3);

         // sampling statement linking TP observation to TP[huc]
         p ~ normal(phuc[hucnum1], sigp[2]);      
         // sampling statement linking Chl[huc] to grand mean
         chuc ~ student_t(4,chlmn, sigchl[1]);
         // sampling statement linking Chl observation to Chl[huc]
         chl ~ normal(chuc[hucnum2], sigchl[2]);
     }
    '
    
    bout <- bayeset(envdata = dat.nat[[1]],
                    df.chl = dat.nat[[2]],
                    varname = c("ptl", "chl", "huc8"))

    print(str(bout))

    datstan <- list(np = nrow(bout[[1]]), p = bout[[1]]$x.cent,
                    hucnum1 = bout[[1]]$hucnum,
                    nc = nrow(bout[[2]]), chl = bout[[2]]$y.cent,
                    hucnum2 = bout[[2]]$hucnum,
                    nhuc = max(bout[[1]]$hucnum))
    print(str(datstan))

    if (runmod) {
        require(rstan)
        rstan_options(auto_write = TRUE)
        nchains <- 4
        options(mc.cores = nchains)

        fit <- stan(model_code = modstan,
                    data = datstan, iter = 1200, chains = nchains,
                    warmup =600, thin= 1)

        return(fit)
    }
    else {
        ## make plots
        chuc <- apply(varout$chuc, 2, mean)
        phuc <- apply(varout$phuc, 2, mean)
        dftemp <- data.frame(hucnum = 1:length(phuc), phuc = as.vector(phuc),
                             chuc = as.vector(chuc))

        grey.t <- adjustcolor("grey", alpha.f = 0.5)
        tiff(width = 4, height =3, pointsize = 8, units = "in",
             res = 600, type = "cairo", file = "sr.nat.tif")
        par(mar = c(4,4,1,1), mgp = c(2.3,1,0))
        xnew <- seq(min(dftemp$phuc), max(dftemp$phuc), length = 40)
        predout <- matrix(NA, ncol = 3, nrow = length(xnew))
        predout2 <- matrix(NA, ncol = 3, nrow = length(xnew))

        ## mean values for chl and tp from the MN data set
        chl.mn.mn <- 2.309
        tp.mn.mn <- 4.498
        ## define same set of TP values for scaling of MN model
        xnew2 <- xnew + bout[[3]] - tp.mn.mn

        for (i in 1:length(xnew)) {
            y <- varout$b1 + (varout$b2-varout$b1)/
                (1+exp(-varout$k*(xnew[i] - varout$p0)))
            predout[i,] <- quantile(y, prob = c(0.025, 0.5, 0.975))
            y2 <- varout.mn$b1 + (varout.mn$b2 - varout.mn$b1)/
                (1 + exp(-varout.mn$k*(xnew2[i] - varout.mn$p0)))
            predout2[i,] <- quantile(y2, prob = c(0.025, 0.5, 0.975))

        }

        sigchl <- apply(varout$sigchl, 2, mean)

        up <- predout[,1]
        dn <- predout[,3]

        plot(dftemp$phuc + bout[[3]], dftemp$chuc + bout[[4]], axes = F,
             xlab = expression(Stream~TP[HUC]~(mu*g/L)),
             ylab = expression(Lake~Chl[HUC]~(mu*g/L)),
             pch = 21, col = "grey", bg = "white")
        logtick.exp(0.001, 10, c(1,2), c(F,F))
        #lines(xnew + bout[[3]], predout[,2] + bout[[4]])
        polygon(c(xnew, rev(xnew)) + bout[[3]],
                c(up, rev(dn)) + bout[[4]], col = grey.t,
                border = NA)

        addMN <- TRUE
        if (addMN) {
            lines(xnew2 + tp.mn.mn, predout2[,1] + chl.mn.mn)
            lines(xnew2 + tp.mn.mn, predout2[,3] + chl.mn.mn)
        }
        dev.off()

    }

    return()
}

fitout <- mod.nat(dat.nat, runmod = T)
varout.nat <- extract(fitout, pars = c("chuc", "phuc", "b1", "b2", "k", "p0",
                                       "sigchl"))
mod.nat(dat.nat, varout = varout.nat, varout.mn = varout.mn, runmod = F)