## model Minnesota data
mod.mn <- function(dat.mn, varout = NULL, runmod = F) {
    
    df1 <- dat.mn[[1]]
    mndat.stream <- dat.mn[[2]]
    
    ## select hucs present in lakes and stream datasets
    hucall <- c(unique(df1$huc12), unique(mndat.stream$huc12))
    hucsel <- hucall[duplicated(hucall)]

    ## select hucs from lake data
    incvec <- rep(F, times = nrow(df1))
    for (i in hucsel) incvec <- incvec| i==df1$huc12
    df1 <- df1[incvec,]

    ## center chl
    df1$chl <- log(df1$resultmeasurevalue)
    chlmn <- mean(df1$chl)
    df1$chl.sc <- df1$chl - chlmn
    print(summary(df1$chl.sc))
    print(sd(df1$chl.sc))

    ## set up factors
    df1$monitoringlocationidentifier <- factor(df1$monitoringlocationidentifier)
    df1$huc12 <- factor(df1$huc12)
    df1$sitenum <- as.numeric(df1$monitoringlocationidentifier)
    df1$hucnum <- as.numeric(df1$huc12)

    lookup <- unique.data.frame(df1[, c("sitenum", "hucnum")])
    lookup <- lookup[order(lookup$sitenum),]

    ## select tp data
    incvec <- rep(F, times = nrow(mndat.stream))
    for (i in hucsel) incvec <- incvec | mndat.stream$huc12 == i
    mndat.stream <- mndat.stream[incvec,]

    ## center TP data
    tpmn <- mean(mndat.stream$tp.random)
    mndat.stream$tp.sc <- mndat.stream$tp.random - tpmn
    mndat.stream$huc12 <- factor(mndat.stream$huc12)
    mndat.stream$hucnum <- as.numeric(mndat.stream$huc12)
    print(max(mndat.stream$hucnum))

    print(chlmn)
    print(tpmn)

    datstan <- list(n = nrow(df1), nsite = max(df1$sitenum),
                    nhuc = max(df1$hucnum), chl = df1$chl.sc,
                    hucnum = lookup$hucnum,
                    sitenum = df1$sitenum,
                    n2 = nrow(mndat.stream),
                    hucnum2 = mndat.stream$hucnum,
                    tp = mndat.stream$tp.sc)
    print(str(datstan))


    modstan <- '
        data {
            int n;              // number of chl samples
            int nsite;          // number of sites with chl
            int nhuc;           // number of hucs with chl
            vector[n] chl;      // lake chl measurements
            int hucnum[nsite];  // hucs associated with each chl site
            int sitenum[n];     // sites associated with each chl sample
            int n2;             // number of tp measurements
            int hucnum2[n2];    // hucs associated with each tp measurement
            vector[n2] tp;      // stream tp measurements
        }
        parameters {
            // parameters for logistic relationship 
            real b1;
            real b2;
            real<lower = 0> k;
            real p0;
            // parameters to define different components of variability
            vector[nsite] etasite;
            vector[nhuc] etahuc;
            real<lower = 0> sig[3];
            real mu;
            vector[nhuc] etahuc2;
            real<lower = 0> sigtp[2];
        }
        transformed parameters {
            vector[nsite] chlsite;  // estimate of mean chl at site
            vector[nhuc] chlhuc;    // estimate of mean chl for huc
            vector[nhuc] tphuc;     // estimate of mean TP for huc
            tphuc = mu + sigtp[1]*etahuc2;  // model for tp for each huc
            // logistic model linking chl[huc] to tp[huc]
            for (i in 1:nhuc)
               chlhuc[i] = b1 + (b2-b1)/(1+exp(-k*(tphuc[i] - p0))) +
                           etahuc[i]*sig[1];
            chlsite = chlhuc[hucnum] + etasite*sig[2];
        }
        model {
            mu ~ normal(0,3);
            etasite ~ std_normal();
            etahuc ~ std_normal();
            etahuc2 ~ std_normal();
            sig ~ cauchy(0,3);
            sigtp ~ cauchy(0,3);
            // weakly informative priors
           b1 ~ normal(-1.5,2);
           b2 ~ normal(1.5,2);
           k ~ normal(1,1);
           p0 ~ normal(0,2);

            tp ~ normal(tphuc[hucnum2], sigtp[2]);
            chl ~ normal(chlsite[sitenum], sig[3]);
        }
    '

    if (runmod) {
        require(rstan)
        rstan_options(auto_write = TRUE)
        nchains <- 4
        options(mc.cores = nchains)
        fit <- stan(model_code = modstan, data = datstan, iter = 1400,
                    warmup = 600, thin = 1, chains = nchains)
        return(fit)
    }
    else {
        ## make the TP-Chl plot
        tphuc <- apply(varout$tphuc, 2, mean)
        chlhuc <- apply(varout$chlhuc, 2, mean)

        b1 <- mean(varout$b1)
        b2 <- mean(varout$b2)
        k <- mean(varout$k)
        p0 <- mean(varout$p0)
        sig <- apply(varout$sig, 2, mean)
        sigtp <- apply(varout$sigtp, 2, mean)
        
        tiff(width = 4, height =3, pointsize = 8, units = "in",
             res = 600, type = "cairo", file = "sr.tif")
        par(mar = c(4,4,1,1), mgp = c(2.3,1,0), mgp = c(2.3,1,0))
        plot(tphuc + tpmn, chlhuc + chlmn, axes = F,
             pch = 21, col = "grey", bg ="white",
             xlab = expression(Stream~TP[HUC]~(mu*g/L)),
             ylab = expression(Lake~Chl[HUC]~(mu*g/L)))
        logtick.exp(0.001, 10, c(1,2), c(T,T))
        xnew <- seq(min(tphuc), max(tphuc), length = 40)

        siglakehuc <- sqrt(sig[1]^2 + sig[2]^2)
        mn <- b1 + (b2-b1)/(1+ exp(-k*(xnew - p0))) + chlmn
        up <- mn + 0.67*siglakehuc[1]
        dn <- mn - 0.67*siglakehuc[1]
        lines(xnew + tpmn, up, lty = "dashed")
        lines(xnew + tpmn, dn, lty = "dashed")
        lines(xnew + tpmn, mn)

        dev.off()
    }
}

fitout <- mod.mn(dat.mn, runmod =T)
#varout.mn <- extract(fitout, pars  = c("tphuc", "chlhuc", "b1", "b2", "k",
#                           "p0", "sig", "sigtp"))        
#mod.mn(dat.mn,varout = varout.mn, runmod = F)