setwd("/Users/eike/Documents/R/BLUP_Problem")
rm(list=ls())
graphics.off()
set.seed(0)
library(MASS)
library(lattice)
library(lme4,keep.source=F);

newsim = 0;
if(newsim){
####################################
#now we'll have a balanced design
L = 10; #number of levels of the predictor
N = 1;  #number of measurements
v = 30; #number of participants
o = L*N*v #number of objects

#Parameters
beta    = c(0,0); #fixed effect parameters
sigma   = 1; 	  #std dev of residual

#data frame of balanced design
a = data.frame(matrix(c(sort(rep(1:v,L*N)),rep(rep((1:L)-mean(1:L),N),v),rep(0,o)),o,3))
names(a)<-c("vp","x","y")
a$vp<-factor(a$vp)

#simulate
LR  = 100000						#number of different true correlations to simulate between -.9 and .9
RHO = seq(-.9,.9,1.8/(LR-1))		#true correlations to simulate 
sim = 1 							#number of simulations per correlation
B   = matrix(rep(0,2*sim*v*LR),nrow=sim*v*LR,ncol=2) #ranef: parameters minus BLUPs
F   = matrix(rep(0,2*sim*LR),nrow=sim*LR,ncol=2)     #fixef: parameters minus estimates
P   = matrix(rep(0,6*sim*LR),nrow=sim*LR,ncol=6)     #Psi: estimated std's and corr and std's and corr of BLUPs
S   = rep(0,sim*LR)								     #Sigma: estimated std of epsilon
K   = S											     #Rho: true correlations of models

k   = 0
r	= rep(0,o)	#random component of the model
e	= r			#residual of the model

#fixed effect
X<-cbind(rep(1,o),a$x)
y<-X%*%beta 

for (rho in RHO){
	Psi  = matrix(c(1,rho,rho,1),2,2);      #std and corr of random effects
	 	
	for (i in 1:sim){

		#random effect
		b       = mvrnorm(n=v,c(0,0),Psi);#,empirical=TRUE
		for (j in 1:v){
    			id<-a[,1]==j
    			#random noise
				e[id] = mvrnorm(n=length(id),0,sigma*sigma)#,empirical=TRUE
				r[id] = rep(b[j,1],sum(id)) + b[j,2]*a$x[id]
		}

		#data complete:
		a$y<- y + r + e
	 	
	 	#save one example data set
	 	if (i==1 & rho==RHO[round(length(RHO)/3)]) {A = a}
	 			 	
	 	#lmer	
	 	a.r<-lmer(y ~ x + (x | vp),data=a)
	 	
	    B[(1:v)+(i-1)*v+k*v,1:2]	= as.matrix(ranef(a.r)$vp-b)
	 	F[i+k,1:2]             		= as.matrix(fixef(a.r))-beta
	 	s         		       		= as.matrix(VarCorr(a.r)$vp)
	 	P[i+k,1]			   		= sqrt(s[1,1])
	 	P[i+k,2]			   		= sqrt(s[2,2])
	 	P[i+k,3]			   		= s[1,2]/(P[i+k,1]*P[i+k,2])
	 	P[i+k,4:5]			   		= sd(coef(a.r)$vp[c(1,2)])
		P[i+k,6]			   		= cor(coef(a.r)$vp[c(1,2)])[2];
		S[i+k]				   		= attr(VarCorr(a.r),"sc")
		K[i+k]						= rho
	}
	k    = k + sim	
}

save(B,F,P,S,K,file="simul1.Rdata")
write(K,file="K1",ncolumns = 1)
write(t(P),file="P1",ncolumns = 6)
} else load("simul1.Rdata")

K1 = K;
P1 = P;

if(newsim){
##############################
#simul of unbalanced design
#now we'll have an unbalanced design
N<-N*2  #maximum number of measurements
o = L*N*v

#data frame of balanced design
rm(a)
a = data.frame(matrix(c(sort(rep(1:v,L*N)),rep(rep((1:L)-mean(1:L),N),v),rep(0,o)),o,3))
names(a)<-c("vp","x","y")
a$vp<-factor(a$vp)

#now kick p% of the rows
p  = 50
id = sort(unique(round(runif(5*o,min=1,max=o)))[1:(round((1-p/100)*o))])
a  = a[id,]
o  = nrow(a)

#simulate
k   = 0
r	= rep(0,o)	#random component of the model
e	= r			#residual of the model

#fixed effect
X<-cbind(rep(1,o),a$x)
y<-X%*%beta 

for (rho in RHO){
	Psi  = matrix(c(1,rho,rho,1),2,2);
	 	
	for (i in 1:sim){
		
		#random effect
		b       = mvrnorm(n=v,c(0,0),Psi);#,empirical=TRUE
		for (j in 1:v){
    			id<-a[,1]==j
    			#random noise
				e[id] = mvrnorm(n=length(id),0,sigma*sigma)#,empirical=TRUE
				r[id] = rep(b[j,1],sum(id)) + b[j,2]*a$x[id]    
		}

		#data complete:
		a$y<- y + r + e
	 	
	 	#save one example data set
	 	if (i==1 & rho==RHO[round(length(RHO)/3)]) {A = a}
	 	
	 	#lmer	
	 	a.r<-lmer(y ~ x + (x | vp),data=a)
	 	
	    B[(1:v)+(i-1)*v+k*v,1:2]	= as.matrix(ranef(a.r)$vp-b)
	 	F[i+k,1:2]             		= as.matrix(fixef(a.r))-beta
	 	s         		       		= as.matrix(VarCorr(a.r)$vp)
	 	P[i+k,1]			   		= sqrt(s[1,1])
	 	P[i+k,2]			   		= sqrt(s[2,2])
	 	P[i+k,3]			   		= s[1,2]/(P[i+k,1]*P[i+k,2])
	 	P[i+k,4:5]			   		= sd(coef(a.r)$vp[c(1,2)])
		P[i+k,6]			   		= cor(coef(a.r)$vp[c(1,2)])[2];
		S[i+k]				   		= attr(VarCorr(a.r),"sc")
		K[i+k]						= rho
	}
	k    = k + sim	
}

save(B,F,P,S,K,file="simul2.Rdata")
write(K,file="K2",ncolumns = 1)
write(t(P),file="P2",ncolumns = 6)
} else load("simul2.Rdata")


K2 = K;
P2 = P;
w = 5000;

n1  = length(K1)
n2  = length(K2)

par(mfrow=c(1,3))	
id1 = order(K1)
id2 = order(K2)
x1  = K1[id1];
y1  = P1[id1,3]-x1;
x2  = K2[id2];
y2  = P2[id2,3]-x2;

st1 = seq(w,n1,by=5);
xf1 = filter(x1,rep(1/w,w),side=1)
yf1 = filter(y1,rep(1/w,w),side=1)

st2 = seq(w,n2,by=5);
xf2 = filter(x2,rep(1/w,w),side=1)
yf2 = filter(y2,rep(1/w,w),side=1)

plot(0,0,col="white",main="estimated correlations",xlab="true correlation",ylab="estimate - true",ylim=c(-0.05,0.05));
lines(xf1[st1],yf1[st1],col="black",lwd=2)
lines(xf2[st2],yf2[st2],col="red",lwd=2)

y1  = P1[id1,6]-x1
y2  = P2[id2,6]-x2
plot(0,0,col="white",main="posterior correlations (from BLUPs)",xlab="true correlation",ylab="BLUPestimate - true",ylim=c(-0.05,0.05)) #plot corr's of BLUPs (y) against actual (x)
yf1 = filter(y1,rep(1/w,w),side=1)
yf2 = filter(y2,rep(1/w,w),side=1)
lines(xf1[st1],yf1[st1],col="black",lwd=2)
lines(xf2[st2],yf2[st2],col="red",lwd=2)


y1  = P1[id1,6]-P1[id1,3]
y2  = P2[id2,6]-P2[id2,3]
plot(0,0,col="white",main="estimated correlations",xlab="true correlation",ylab="BLUPestimate - estimate",ylim=c(-0.05,0.05)) #plot difference of both (y) against actual (x)
yf1 = filter(y1,rep(1/w,w),side=1)
yf2 = filter(y2,rep(1/w,w),side=1)
lines(xf1[st1],yf1[st1],col="black",lwd=2)
lines(xf2[st2],yf2[st2],col="red",lwd=2)