# writeModel2.R can be used to automate many of the MSSMs, particularly those
# that involve 1 < m < n signals

# These next 4 lines don't really change anything, but they need to be here - sometimes
# BUGS is looking for the collection operator when it won't be there

# To summarize, there are N sites being monitored, and there may be M subpopulations (1 <= M <= N).
# The vector 'groups' assigns sites to subpopulations, and the remaining vectors allow for parameters
# to be shared among subpopulations.  For example, if 2 subpopulations have the same process error vars,
# but different growth rates and observation errors, U.groups = c(1,2), Q.groups = c(1,1), R.groups = c(1,2)
groups <- c(groups,0)
Q.groups <- c(Q.groups,0)
R.groups <- c(R.groups,0)
U.groups <- c(U.groups,0)


###################################
# THIS SECTION WRITES THE BUGS CODE
###################################
Str1 <- ""
Str2 <- ""
Str3 <- ""
for(j in 1:numClones) {
   Str1 <- paste(Str1,"groupStates",j,"[1,i] <- 0;\n",sep="")
   #Str1 <- paste(Str1,"groupStates",j,"[1,i] ~ dunif(4,10);\n",sep="")
   #Str2 <- paste(Str2,"mu",j,"[i,j] <- (groupStates",j,"[i-1,j] + U.temp[U.groups[j]]);\n",sep="")
   Str2 <- paste(Str2,"mu",j,"[i,j] <- (inprod(B[j,],groupStates",j,"[i-1,]) + U.temp[U.groups[j]]);\n",sep="")
   Str3 <- paste(Str3,"groupStates",j,"[i,1:M] ~ dmnorm(mu",j,"[i,1:M], invQ[,]);\n",sep="")
}

# FIRST CASE FOR Q (UNCONSTRAINED COVARIANCES BETWEEN M GROUPS)
modelQ = paste("
# no constraints
	invQ[1:M , 1:M]  ~ dwish(priorQ[1:M,1:M], M)	
	Q[1:M , 1:M] <- inverse(invQ[1:M,1:M])
	
	# these are the individual state vectors
      for(i in 1:M) {\n",
	Str1
	,"}
	for(i in 2:numY) {
	   # cycle through years
	   for(j in 1:M) {\n",
		 Str2
,"
	   }
",
	   Str3
	   ,"
	}
",sep="")


Str1 <- ""
Str2 <- ""
for(j in 1:numClones) {
   Str1 <- paste(Str1,"groupStates",j,"[1,i] <- 0;\n",sep="");
   Str2 = paste(Str2,"mu",j,"[i,j] <- (inprod(B[j,],groupStates",j,"[i-1,]) + U.temp[U.groups[j]]);\n",sep="")
   Str2 = paste(Str2,"groupStates",j,"[i,j] ~ dnorm(mu",j,"[i,j],invQ[j]);\n",sep="")
}

# SECOND CASE FOR Q, INDEPEDENT VARIANCES
modelQ[2] = paste("
# independent and potentially different
	# Q is cov matrix for process error	
      #Qprcsn <- 1; 
	for(i in 1:Q.G) {
           sigma.Q[i] ~ dunif(0.0001,2);
           sigma2.Q[i] <- sigma.Q[i]*sigma.Q[i];
      }
	for(j in 1:M) {
     	Q[j] <- sigma2.Q[Q.groups[j]];
		invQ[j] <- 1/Q[j];
     }

	# first year
    	for(i in 1:M) {\n",
		Str1,
"	}
	# remaining years
	for(i in 2:numY) {
	   # cycle through years
	   for(j in 1:M) { \n",
           Str2,
"	   }
	}
",sep="")


Str1 <- ""
Str2 <- ""
Str3 <- ""
for(j in 1:numClones) {
   Str1 <- paste(Str1,"groupStates",j,"[1,i] <- 0;\n",sep="")
   Str2 <- paste(Str2,"mu",j,"[i,j] <- (inprod(B[j,],groupStates",j,"[i-1,]) + U.temp[U.groups[j]]);\n",sep="");
   Str3 <- paste(Str3,"groupStates",j,"[i,1:M] ~ dmnorm(mu",j,"[i,1:M], invQ[,]);\n",sep="");
}

# THIRD CASE FOR Q, EQUAL VARIANCE & COVARIANCE
modelQ[3] = paste("
# not-independent and different
	# Q is cov matrix for process error	
      #Qprcsn <- 1; 
 	# 1st will be variance, 2nd will be covariance
	for(i in 1:2) {
           sigma.Q[i] ~ dunif(0.0001,2);
           sigma2.Q[i] <- sigma.Q[i]*sigma.Q[i];
	}
	# first row
   	Q[1,1] <- sigma2.Q[1];
	for(i in 2:M) { Q[1,i] <- sigma2.Q[2];}
	# deal with remaining rows
        for(i in 2:M) {
	   # for each column (j)
 	   for(j in 1:(i-1)) {	Q[i,j] <- sigma2.Q[2]; }
           # diagonal
	   Q[i,i] <- sigma2.Q[1];
           # to right of diagonal
	   for(j in (i+1):M) {  Q[i,j] <- sigma2.Q[2]; }
	}
	invQ[1:M,1:M] <- inverse(Q[1:M,1:M]);
	
	# these are the individual state vectors
\n
      for(i in 1:M) {",

	Str1,
"}
	for(i in 2:numY) {
	   # cycle through years
	   for(j in 1:M) {\n",
Str2,"
	   }
",
Str3,
"
	}
",sep="")

Ystr <- "";
Ystr1 <- "";
for(j in 1:numClones) {
   p<-paste("DATA",j,"[i, 1:N] ~ dmnorm(states",j,"[i,1:N], invR[1:N,1:N])","\n",sep="");
   Ystr <- paste(Ystr,p,sep="");
   p<-paste("states",j,"[i,j] <- groupStates",j,"[i,groups[j]] + A[j];","\n",sep="");
   Ystr1 <- paste(Ystr1,p,sep="");
}
modelR = paste("
# no constraints
	invR[1:N , 1:N]  ~ dwish(priorR[1:N,1:N], N)	
	R[1:N , 1:N] <- inverse(invR[1:N,1:N])
	# the same observation error matrix
	for (i in 1:numY) { 
	   for(j in 1:N) {\n",
		Ystr1,
"}\n",
	   Ystr,
"
	}
",sep="")


Ystr <- "";
Ystr1<-"";
for(j in 1:numClones) {
   p<-paste("DATA",j,"[i, j] ~ dnorm(states",j,"[i,j], invR[j])","\n",sep="");
   Ystr <- paste(Ystr,p,sep="");
   p<-paste("states",j,"[i,j] <- groupStates",j,"[i,groups[j]] + A[j];","\n",sep="");
   Ystr1 <- paste(Ystr1,p,sep="");
}
modelR[2] = paste("
# independent and different
	# R is cov matrix for observation error	
	#Rprcsn <- 1; 
	for(i in 1:R.G) {
            sigma.R[i] ~ dunif(0.0001,2);
		sigma2.R[i] <- sigma.R[i]*sigma.R[i];
	}
	for(j in 1:N) {
     	R[j] <- sigma2.R[R.groups[j]];
		invR[j] <- 1/R[j];
		for(i in 1:numY) {\n",
		Ystr1,
	      Ystr,	
"		}
     }
",sep="")


Ystr <- "";
Ystr1<-"";
for(j in 1:numClones) {
   p<-paste("DATA",j,"[i, 1:N] ~ dmnorm(states",j,"[i,1:N], invR[1:N,1:N])","\n",sep="");
   Ystr <- paste(Ystr,p,sep="");
   p<-paste("states",j,"[i,j] <- groupStates",j,"[i,groups[j]] + A[j];","\n",sep="");
   Ystr1 <- paste(Ystr1,p,sep="");
}
modelR[3] = paste("
# not-independent and different
	# R is cov matrix for observation error	
      #Rprcsn <- 1; 
 	# 1st will be variance, 2nd will be covariance
	for(i in 1:2) {
            sigma.R[i] ~ dunif(0.0001,2);
            sigma2.R[i] <- sigma.R[i]*sigma.R[i];
	}
	# first row
   	R[1,1] <- sigma2.R[1];
	for(i in 2:N) { R[1,i] <- sigma2.R[2];}
	# deal with remaining rows
        for(i in 2:N) {
	   # for each column (j)
 	   for(j in 1:(i-1)) {	R[i,j] <- sigma2.R[2]; }
        # diagonal
	   R[i,i] <- sigma2.R[1];
        # to right of diagonal
	   for(j in (i+1):N) {  R[i,j] <- sigma2.R[2]; }
	}
	invR[1:N,1:N] <- inverse(R[1:N,1:N]);
	# the same observation error matrix
	for (i in 1:numY) { 
	   for(j in 1:N) {\n",
		Ystr1,
"}\n",
	   Ystr,
"
	}	
",sep="")

# Create a string for the B matrix
if(f1 == 1) {
    #Bstr <- paste("        for(i in 1:M) {B[i,i] <- 1}\n",sep="")       # no interactions
     Bstr <- paste("        for(i in 1:M) { \n","        B[i,i] <- 1\n","        for(j in 1:(i-1)) {B[i,j] <- 0}\n","        for(j in (i+1):M) {B[i,j] <- 0}\n","        }\n",sep="")       # no interactions
}
if(f1 == 2) {
    #Bstr <- paste("        B0 ~ dunif(0,1)\n","for(i in 1:M) {B[i,i] <- B0}\n",sep="")       # identity
    Bstr <- paste("        B0 ~ dunif(0,1)\n","        for(i in 1:M) { \n","        B[i,i] <- B0\n","        for(j in 1:(i-1)) {B[i,j] <- 0}\n","        for(j in (i+1):M) {B[i,j] <- 0}\n","        }\n",sep="")
}
if(f1 == 3) {
    #Bstr <- paste("        for(i in 1:M) {B[i,i] ~ dunif(0,1)}\n",sep="")       # identity
    Bstr <- paste("        for(i in 1:M) { \n","        B[i,i] ~ dunif(0,1)\n","        for(j in 1:(i-1)) {B[i,j] <- 0}\n","        for(j in (i+1):M) {B[i,j] <- 0}\n","        }\n",sep="")
}
if(f1 == 4) {
    Bstr <- paste("        for(i in 1:M) {\n","        for(j in 1:M) {\n","        B[i,j] ~ dunif(0,1)\n","        }\n","}\n",sep="")       # unconstrained
}

Str1 <- ""
#for(j in 1:numClones) {
   Str1 <- paste(Str1,"           A[i] ~ dunif(-10,10);\n",sep="")
#}
modelB = paste("        #Put bounds on interaction matrix\n",Bstr,"
	# this is the multivariate mean vector describing growth          
	for(i in 1:U.G) { 
	   U.temp[i] ~ dunif(-1, 1);
	}
	# U.G above is the actual estimated parameters; these need to be assigned to groups
	for(i in 1:M) {
	   U.temp2[i] <- U.temp[U.groups[i]]   
	}
	# the group growth rates need to be assigned to sites to return back
	for(i in 1:N) {
	   U[i] <- U.temp2[groups[i]];
\n",
	   Str1,
"
      }
",sep="");

###################################
#  THIS WRITES THE MODEL TO A TEXT FILE, model.txt
###################################
modelString <- paste("model {",modelB,modelQ[f2],modelR[f3],"}",sep="\n");
write(modelString,file="model.txt",append=FALSE)

priorQ <- matrix(0,M,M); 
priorR <- matrix(0,N,N);
for(i in 1:M) {
   priorQ[i,i]<-1;
}
for(i in 1:N) {
   priorR[i,i]<-1;
}

###################################
# WRITE DATA TO A TEXT FILE AS A LIST...THIS IS UGLY, BUT IT'S THE
# ONLY OPTION - BUGS DOESN'T LIKE MIXING 3D ARRAYS / MVN 
###################################
for(i in 1:numClones) {
   assign(paste("DATA",i,sep=""),DATA);
}

str <- "DATA1";
if(numClones > 1) {
  for(i in 2:numClones) {
    p <- paste("DATA",i,sep="");
    str <- c(str,p)
  }
}

##################################
# EACH DATA FILE IS UNIQUE
##################################
if(f2==1 && f3 == 1) { 
   str2 <- c("numY","N","priorQ","priorR","M","U.groups","U.G","groups")
   data <- as.list(c(str,str2));
}
if(f2==2 && f3 == 1) {
   str2 <- c("numY","N","priorR","M","Q.groups","Q.G","U.groups","U.G","groups")
   data <- as.list(c(str,str2));
}
if(f2==3 && f3 == 1) {
   str2 <- c("numY","N","priorR","M","U.groups","U.groups","U.G","groups")
   data <- as.list(c(str,str2));
}
if(f2==1 && f3 == 2) { 
   str2 <- c("numY","N","priorQ","R.G","R.groups","M","U.groups","U.G","groups")
   data <- as.list(c(str,str2));   
}
if(f2==2 && f3 == 2) {
   str2 <- c("numY","N","R.G","R.groups","M","Q.groups","Q.G","U.groups","U.G","groups")
   data <- as.list(c(str,str2));
}
if(f2==3 && f3 == 2) {
   str2 <- c("numY","N","R.G","R.groups","M","U.groups","U.G","groups")
   data <- as.list(c(str,str2));
}
if(f2==1 && f3 == 3) { 
   str2 <- c("numY","N","priorQ","M","U.groups","U.G","groups")
   data <- as.list(c(str,str2));
}
if(f2==2 && f3 == 3) {
   str2 <- c("numY","N","M","Q.groups","Q.G","U.groups","U.G","groups")
   data <- as.list(c(str,str2));
}
if(f2==3 && f3 == 3) {
   str2 <- c("numY","N","M","U.groups","U.G","groups")
   data <- as.list(c(str,str2));
}

# These next 4 lines don't really change anything, but they need to be here - the 
# first 4 lines of this file added an element to each vector, now it is deleted
#groups <- groups[1:(length(groups)-1)]
#Q.groups <- Q.groups[1:(length(Q.groups)-1)]
#R.groups <- R.groups[1:(length(R.groups)-1)]
#U.groups <- U.groups[1:(length(U.groups)-1)]