# Square mixing model example, BX Semmens April 09, 2013
# Steps... 
# 1) specify population means and variances for mixtures and source
# 2) define sample sizes, draw random deviates to simulate data
# 3) specify fully Bayesian model for JAGS (Ward et al.2010), write file
# 4) run model, plot posteriors

library(runjags)
library(R2WinBUGS)
library(R2jags)
library(gtools)
library(gdata)
library(ggplot2)

set.seed(850) #make any given result repeatable...

# Specify means and variances for sources and mixtures ###############################

# establish source means (4 sources (rows), delC in 1st column, delN in 2nd column5
src_means <- matrix(c(-20,5,-20,15,-10,15,-10,5), nrow = 4, ncol=2, byrow=TRUE)

# establish source variances (4 sources (rows), delC in 1st column, delN in 2nd column)
# NOTE: sd val mimics the sd values reported for mangrove litter by Kon et al. 2008 (Fish Sci.,75:715–722)
src_sds <- matrix(c(1.12,1.12,1.12,1.12,1.12,1.12,1.12,1.12), nrow = 4, ncol=2, byrow=TRUE)

#establsh mixture means (delC in 1st column, delN in 2nd column)
mix_means <- c(-15,10)

#establsh mixture variances (delC in 1st column, delN in 2nd column)
mix_sds <- c(1.12,1.12)

# Simulate data #####################################################################
N <- 5 #sample size
num.src <- 4 #number of sources
num.iso <- 2 # number of tracers
S <- array(0, dim=c(num.src,num.iso,N))
X <- array(0, dim=c(num.iso,N))

for(i in 1:N) {
	for(src in 1:num.src) {
   		for(iso in 1:num.iso) {
   			S[src,iso,i] <- rnorm(1, mean = src_means[src,iso], sd = src_sds[src,iso] ) #source data
   			X[iso,i] <- rnorm(1, mean = mix_means[iso], sd = mix_sds[iso] )		    #mixture data
		}
   	}
}

#PLOT THE DATA #######################################################################
plot(S[1,1,],S[1,2,],ylim=c(1,19),xlim=c(-23,-7),pch =0) 	#first source
points(S[2,1,],S[2,2,],pch=2) 								#second source
points(S[3,1,],S[3,2,],pch=5) 								#third source
points(S[4,1,],S[4,2,],pch=6) 								#fourth source
points(X[1,],X[2,],pch = 16)							  	# mixture


# SPECIFY MODEL AND WRITE TO FILE ####################################################
cat ("model {
   # SPECIFY PRIORS FOR DIET PROPORTIONS AND SOURCE VALUES
   p[1:num.src] ~ ddirch(alpha[]);  # these are diet proportions
   for(i in 1:num.src){
   	for (j in 1:num.iso){
   	   	u[i,j]~dunif(-50,50) # mean of source isotope vals
   	   	u.prec[i,j]~dgamma(.01,.01) # precision of source iso values
   	   	u.sigma2[i,j] <- 1 / u.prec[i,j] # variance of source iso vals
   	}
   }
   
   # NOW CALCULATE MIXTURE EXPECTATION
   for(i in 1:num.src) {
      p2[i] <- p[i]*p[i]; # these are weights (diet proportions) for variances
   }

   # for each isotope, calculate the predicted mixtures
   
   for(iso in 1:num.iso) {

         mix.mu[iso] <- inprod(u[,iso],p[]);
         mix.var[iso] <- inprod(u.sigma2[,iso],p2[]);
         mix.totalVar[iso] <- mix.var[iso] ;
         mix.prcsn[iso] <- 1/(mix.totalVar[iso]);
   }

   # NOW CALC THE LIKELIHOOD OF THE DATA
   for(i in 1:N) {
      for(iso in 1:num.iso) {
         X[iso,i] ~ dnorm(mix.mu[iso], mix.prcsn[iso]); #likelihood of mixture data
         for (src in 1:num.src){
         	S[src,iso,i] ~ dnorm(u[src,iso],u.prec[src,iso]) #likelihood of the sourc data
         }
      }
   }

}",file="full_Bayes_model.txt")

# NOW SEND THE MODEL TO JAGS ###########################################################
# These are the parameters that need to be set by the user
  mcmc.burn <- as.integer(10000) # burn in
  mcmc.chainLength <- as.integer(20000)  # post-burn
  mcmc.thin = 10; # thinning the chain
  mcmc.chains = 3 # needs to be at least 2 for DIC
  alpha=rep(1,num.src) # these are used to specify uninformative prior for Dirichlet
  
        
  jags.data = list("N", "num.src", "num.iso", "S", "X","alpha")  
  jags.params = c("p","u","u.prec")

  model.loc = "full_Bayes_model.txt"
  model.run = jags(jags.data, parameters.to.save= jags.params, model.file=model.loc, n.chains = mcmc.chains, 
  			n.burnin = mcmc.burn, n.thin = mcmc.thin, n.iter = mcmc.chainLength, DIC = FALSE) 
  
  attach.jags(model.run)

# NOW PLOT RESULTS ######################################################################
 
 # Draftman's Plot 
library(IDPmisc) 
ipairs(p)
  
# Polygon plot set up using Fry's methods
SmeanX=c(0,0,0,0)
SmeanY=c(0,0,0,0)
SdangerX=c(0,0,0,0)
SdangerY=c(0,0,0,0)
centroidX = 0
centroidY = 0
danger.pcnt=0.50

#calculate the src polygon, danger zone mixture means and the centroid ####################
for(src in 1:num.src) {
   			SmeanX[src] = mean(S[src,1,1:N]) #mean of source data in x axis
   			SmeanY[src] = mean(S[src,2,1:N]) #mean of source data in y axis
   			Mx = X[1,1:N] #mixture mean
   			My = X[2,1:N] #mixture mean
   			centroidX = centroidX + SmeanX[src] * 1/num.src 
   			centroidY = centroidY + SmeanY[src] * 1/num.src 
		}

boundz<- (1-1/num.src) * danger.pcnt + (1/num.src)
# now calculate the danger polygon
for(src in 1:num.src) {
splitter <-rep((1-boundz)/(num.src-1),num.src)	
splitter[src]<- boundz
SdangerX[src] =	sum(SmeanX * splitter)
SdangerY[src] =	sum(SmeanY * splitter)
}

# Plot the polygons

plot(S[1,1,],S[1,2,],ylim=c(1,19),xlim=c(-23,-7),pch =0) 	#first source
points(S[2,1,],S[2,2,],pch=2) 								#second source
points(S[3,1,],S[3,2,],pch=5) 								#third source
points(S[4,1,],S[4,2,],pch=6) 								#fourth source
points(X[1,],X[2,],pch = 16)								# mixture		
polygon(SmeanX,SmeanY, col=NA, lty=1, lwd=1, border="black")
polygon(SdangerX,SdangerY, col=NA, lty=2, lwd=1, border="black")


# Make plot of % resolved for illustrative purposes ########################################

library(plotrix)
p.range =(apply(p,2,max)-apply(p,2,min)) /2
p.center = (apply(p,2,max)+apply(p,2,min))/2

 plotCI(p.center,1:num.src,p.range,pt.bg=par("bg"),pch=NA,err="x",gap=0.06,xlim=c(0,1))
text(p.center,1:num.src,paste(round(100*p.center*2, 1), "%", sep=""))

# Finally, make histograms ###############
srclist <- c("Source 1","Source 2","Source 3","Source 4")
Dataz <- data.frame(
    src.name=rep(srclist,each=(dim(p)[1])),
    value=c(p[,1],p[,2],p[,3],p[,4])
    )

ggplot(Dataz,aes(value,fill=src.name)) + geom_density(alpha=0.2)