# If you use this code or parts of it, please cite the following paper
# Bruhin A., E. Fehr, D. Schunk (2018): "The Many Faces of Human Sociality: Uncovering the Distribution and Stability of Social Preferences",
# Journal of the European Economic Association, forthcoming.
# ===========================================================================================================================================

rm(list=ls()) # Clear the memory
library(compiler)
set.seed(51)

# Data-Preparation-Function
f.dataprep <- function(dat, persid, label_choice_x, label_indicators_x, label_indicators_y, label_self_x, label_other_x, label_self_y, label_other_y, nafiller=1) {
	tid <- table(dat[,persid])    #Table of individuals
	pid <- as.numeric(names(tid)) #Vector of individual identifiers
	j   <- length(tid)            #Number of individuals
	n   <- max(tid)               #Maximum number of choices per individual
	lk  <- length(label_indicators_x)        #Number of "left" regressors
	rk  <- length(label_indicators_y)        #Number of "right" regressors
	{if (lk != rk) {stop("Length of label_indicators_x and label_indicators_y differ!")}}
	k   <- lk                     #Number of regressors

	indicators_x  <- matrix(1, n*j, k)  #"left" regressor matrix
	indicators_y <- matrix(1, n*j, k)  #"right" regressor matrix

	nlist      <- list(c(), pid)
	self_x     <- matrix(0, n, j, dimnames=nlist) # Payoff matrix, player A, "X"
	other_x    <- matrix(0, n, j, dimnames=nlist) # Payoff matrix, player B, "X"
	self_y     <- matrix(0, n, j, dimnames=nlist) # Payoff matrix, player A, "Y"
	other_y    <- matrix(0, n, j, dimnames=nlist) # Payoff matrix, player B, "Y"
	choice_x <- matrix(0, n, j, dimnames=nlist) # Matrix indicating the choice (1=left, 0=right)
	sel        <- matrix(0, n, j, dimnames=nlist) # Matrix indicating whether an observation is present (panel may be unbalanced!)
	rm(nlist)

	#Fill the data and selection matrices individual by individual (could handle an unbalanced panel)
	for (i in 1:j) {
		sel[,i]        <- c(rep(1, tid[i]), rep(0, n-tid[i]))

		idx            <- dat[,persid] == pid[i]
		self_x[,i]     <- c(dat[idx, label_self_x],   rep(nafiller, n-tid[i]))
		other_x[,i]    <- c(dat[idx, label_other_x],  rep(nafiller, n-tid[i]))
		self_y[,i]     <- c(dat[idx, label_self_y],   rep(nafiller, n-tid[i]))
		other_y[,i]    <- c(dat[idx, label_other_y],  rep(nafiller, n-tid[i]))
		choice_x[,i] <- c(dat[idx, label_choice_x], rep(nafiller, n-tid[i]))

		{if (k > 1) {
			m1 <- as.matrix(dat[idx, label_indicators_x])
			m2 <- matrix(0, n-nrow(m1), ncol(m1))
			m  <- rbind(m1, m2)
			indicators_x[(n*(i-1)+1):(i*n),] <- m
			rm (m1,m2)

			m1 <- as.matrix(dat[idx, label_indicators_y])
			m2 <- matrix(0, n-nrow(m1), ncol(m1))
			m  <- rbind(m1, m2)
			indicators_y[(n*(i-1)+1):(i*n),] <- m
			rm(m1,m2)
		}}
	}
	list(tid=tid, pid=pid, n=n, j=j, k=k, indicators_x=indicators_x, indicators_y=indicators_y, self_x=self_x, other_x=other_x, self_y=self_y, other_y=other_y, choice_x=choice_x, sel=sel)
}
f.dataprep <- cmpfun(f.dataprep)

# Component Density Function
f.compdens <- function(theta.v, choice_x, self_x, other_x, self_y, other_y, indicators_x, indicators_y, sel, nc ,k, n, j) {
	theta.m <- matrix(theta.v, k+1, nc)

	linindex_x <- indicators_x  %*% theta.m[1:k,] # Linear index of X-choice, a (n*j x nc)-matrix
	linindex_y <- indicators_y %*% theta.m[1:k,]  # Linear index of Y-choice, a (n*j x nc)-matrix

	ret     <- matrix(0, j, nc)

	# Loop over all nc preference types
	for (c in 1:nc) {
		li_x   <- matrix(linindex_x[,c], n, j)
		li_y   <- matrix(linindex_y[,c], n, j)
		choicesens <- theta.m[k+1,c]
		cs_li_x   <- choicesens * ((1-li_x)*self_x + li_x*other_x)  #Type-specific deterministic utility of chosing "X" times choice sensitivity
		cs_li_y   <- choicesens * ((1-li_y)*self_y + li_y*other_y)  #Type-specific deterministic utility of chosing "Y" times choice sensitivity

		probs <- ( (exp(cs_li_x)/(exp(cs_li_x)+exp(cs_li_y)))^choice_x * (exp(cs_li_y)/(exp(cs_li_x)+exp(cs_li_y)))^(1-choice_x) )^sel

		for (i in 1:j) {
			ret[i,c] <- prod(probs[,i])
		}
	}
	ret
}
f.compdens <- cmpfun(f.compdens)

# Log Likelihood of the Mixture Model and its Gradient
f.totloglik <- function(psi, choice_x, self_x, other_x, self_y, other_y, indicators_x, indicators_y, sel, nc, k, n, j) {
	lpsi    <- length(psi)
	theta.v <- psi[1:(lpsi-nc+1)]
	mix     <- psi[(lpsi-nc+2):lpsi]
	mix     <- c(mix, 1-sum(mix))
	m       <- matrix(mix, j, nc, byrow=T)

	sum(log(rowSums( m*f.compdens(theta.v, choice_x, self_x, other_x, self_y, other_y, indicators_x, indicators_y, sel, nc, k, n, j))))
}
f.totloglik <- cmpfun(f.totloglik)

f.totloglik.gradient <- function(psi, choice_x, self_x, other_x, self_y, other_y, indicators_x, indicators_y, sel, nc, k, n, j) {
	lpsi    <- length(psi)
	theta.v <- psi[1:(lpsi-nc+1)]
	theta.m <- matrix(theta.v, k+1, nc)
	mix     <- psi[(lpsi-nc+2):lpsi]
	mix     <- c(mix, 1-sum(mix))
	m       <- matrix(mix, j, nc, byrow=T)

	nj           <- n*j
	choice_x.k   <- matrix(as.vector(choice_x), nj, k)
	sel.k        <- matrix(as.vector(sel)     , nj, k)
	self_x.k     <- matrix(as.vector(self_x)  , nj, k)
	other_x.k    <- matrix(as.vector(other_x) , nj, k)
	self_y.k     <- matrix(as.vector(self_y)  , nj, k)
	other_y.k    <- matrix(as.vector(other_y) , nj, k)

	cdens.m <- f.compdens(theta.v, choice_x, self_x, other_x, self_y, other_y, indicators_x, indicators_y, sel, nc ,k, n, j)
	denom.outerderiv   <- 1/matrix(rowSums(m*cdens.m), n, j, byrow=T)
	denom.outerderiv.k <- matrix(as.vector(denom.outerderiv), nj, k)

	ret.theta <- rep(0, nc*(k+1))
	ret.mix   <- rep(0, nc)

	for (c in 1:nc) {
		choicesens <- theta.m[k+1, c]
		beta  <- theta.m[1:k, c]

		li_x   <- matrix(as.vector(indicators_x %*% beta), n, j) # linear index X-choice
		li_y   <- matrix(as.vector(indicators_y %*% beta), n, j) # linear index Y-choice

		ul    <- (1-li_x)*self_x + li_x*other_x # Deterministic utility from X-choice
		ur    <- (1-li_y)*self_y + li_y*other_y # Deterministic utility from Y-choice
		ul.k  <- matrix(as.vector(ul), nj, k)
		ur.k  <- matrix(as.vector(ur), nj, k)

		innerderiv   <- mix[c]*matrix(cdens.m[,c],n,j,byrow=T)/((exp(choicesens*ul)/(exp(choicesens*ul)+exp(choicesens*ur)))^choice_x * (exp(choicesens*ur)/(exp(choicesens*ul)+exp(choicesens*ur)))^(1-choice_x))
		innerderiv.k <- matrix(as.vector(innerderiv), nj, k)

		choicesens.grad <- sel * denom.outerderiv * innerderiv * ((1/(exp(choicesens*ul)+exp(choicesens*ur))^2) * ( (exp(choicesens*ul)/(exp(choicesens*ul)+exp(choicesens*ur)))^choice_x * exp(choicesens*(ul + 2*ur))*(-ur + ul) * (exp(choicesens*ur)/(exp(choicesens*ul)+exp(choicesens*ur)))^(-choice_x) * (choice_x * exp(-choicesens*ul) - exp(-choicesens*ur) + choice_x * exp(-choicesens*ur)) ))
		choicesens.grad <- sum(choicesens.grad)

		beta.grad <- sel.k * denom.outerderiv.k * innerderiv.k * ((-1/(exp(choicesens*ul.k)+exp(choicesens*ur.k))^2) * ( (exp(choicesens*ul.k)/(exp(choicesens*ul.k)+exp(choicesens*ur.k)))^choice_x.k * choicesens*exp(choicesens*(ul.k + 2*ur.k))*(indicators_x*(self_x.k-other_x.k) + indicators_y*(other_y.k-self_y.k)) * (exp(choicesens*ur.k)/(exp(choicesens*ul.k)+exp(choicesens*ur.k)))^(-choice_x.k) * (choice_x.k * exp(-choicesens*ul.k) - exp(-choicesens*ur.k) + choice_x.k * exp(-choicesens*ur.k)) ))
		beta.grad <- colSums(beta.grad)

		mix.grad  <- sum(cdens.m[,c] * denom.outerderiv[1,])

		ret.theta[((c-1)*(k+1)+1):(c*(k+1))] <- c(beta.grad, choicesens.grad)
		ret.mix[c]                           <- mix.grad
	}
	ret.mix <- ret.mix - ret.mix[nc]
	ret     <- c(ret.theta, ret.mix[1:(nc-1)])

	ret
}
f.totloglik.gradient <- cmpfun(f.totloglik.gradient)

# Log Likelihood Part in the M-Step and its Gradient
f.mloglik <- function(theta.v, choice_x, self_x, other_x, self_y, other_y, indicators_x, indicators_y, sel, nc, k, n, j, tau) {
	sum(tau * log(f.compdens(theta.v, choice_x, self_x, other_x, self_y, other_y, indicators_x, indicators_y, sel, nc, k, n, j)))
}
f.mloglik <- cmpfun(f.mloglik)

f.mloglik.pregrad <- function(theta.c, weight, choice_x, self_x, other_x, self_y, other_y, indicators_x, indicators_y, sel, nc, k, n, j) {
	choicesens <- theta.c[k+1]
	beta       <- theta.c[1:k]

	weights <- matrix(weight, n, j, byrow=T) #The individual weights (used for the taus)

	li_x   <- matrix(as.vector(indicators_x %*% beta), n, j) # linear index X-choice
	li_y   <- matrix(as.vector(indicators_y %*% beta), n, j) # linear index Y-choice

	ul    <- (1-li_x)*self_x + li_x*other_x # Deterministic utility from X-choice
	ur    <- (1-li_y)*self_y + li_y*other_y # Deterministic utility from Y-choice

	denom <- 1/((exp(choicesens*ul)/(exp(choicesens*ul)+exp(choicesens*ur)))^choice_x * (exp(choicesens*ur)/(exp(choicesens*ul)+exp(choicesens*ur)))^(1-choice_x))

	#The gradient of choicesens (scalar)
	choicesens.grad <- sel * weights * denom * ((1/(exp(choicesens*ul)+exp(choicesens*ur))^2) * ( (exp(choicesens*ul)/(exp(choicesens*ul)+exp(choicesens*ur)))^choice_x * exp(choicesens*(ul + 2*ur))*(-ur + ul) * (exp(choicesens*ur)/(exp(choicesens*ul)+exp(choicesens*ur)))^(-choice_x) * (choice_x * exp(-choicesens*ul) - exp(-choicesens*ur) + choice_x * exp(-choicesens*ur)) ))
	choicesens.grad <- sum(choicesens.grad)

	#The gradient (k x 1)-vector of beta
	nj           <- n*j
	choice_x.k <- matrix(as.vector(choice_x), nj, k)
	denom.k      <- matrix(as.vector(denom)     , nj, k)
	weights.k    <- matrix(as.vector(weights)   , nj, k)
	ul.k         <- matrix(as.vector(ul)        , nj, k)
	ur.k         <- matrix(as.vector(ur)        , nj, k)
	sel.k        <- matrix(as.vector(sel)       , nj, k)
	self_x.k   <- matrix(as.vector(self_x)  , nj, k)
	other_x.k   <- matrix(as.vector(other_x)  , nj, k)
	self_y.k  <- matrix(as.vector(self_y) , nj, k)
	other_y.k  <- matrix(as.vector(other_y) , nj, k)

	beta.grad <- sel.k * weights.k * denom.k * ((-1/(exp(choicesens*ul.k)+exp(choicesens*ur.k))^2) * ( (exp(choicesens*ul.k)/(exp(choicesens*ul.k)+exp(choicesens*ur.k)))^choice_x.k * choicesens*exp(choicesens*(ul.k + 2*ur.k))*(indicators_x*(self_x.k-other_x.k) + indicators_y*(other_y.k-self_y.k)) * (exp(choicesens*ur.k)/(exp(choicesens*ul.k)+exp(choicesens*ur.k)))^(-choice_x.k) * (choice_x.k * exp(-choicesens*ul.k) - exp(-choicesens*ur.k) + choice_x.k * exp(-choicesens*ur.k)) ))
	beta.grad <- colSums(beta.grad)

	c(beta.grad, choicesens.grad)
}
f.mloglik.pregrad <- cmpfun(f.mloglik.pregrad)

f.mloglik.grad <- function(theta.v, choice_x, self_x, other_x, self_y, other_y, indicators_x, indicators_y, sel, nc, k, n, j, tau) {
	ret <- rep(0, (k+1)*nc)

	for (c in 1:nc) {
		ret[((k+1)*(c-1)+1):((k+1)*c)] <- f.mloglik.pregrad(theta.v[((c-1)*(k+1)+1):(c*(k+1))], tau[,c], choice_x, self_x, other_x, self_y, other_y, indicators_x, indicators_y, sel, nc, k, n, j)
	}

	ret
}
f.mloglik.grad <- cmpfun(f.mloglik.grad)

# E-Step Function
f.estep <- function(theta.v, choice_x, self_x, other_x, self_y, other_y, indicators_x, indicators_y, sel, nc, k, n, j, mix) {
	m <- matrix(mix, j, nc, byrow=T)

	numerator   <- m * f.compdens(theta.v, choice_x, self_x, other_x, self_y, other_y, indicators_x, indicators_y, sel, nc, k, n, j)
	denominator <- matrix(rowSums(numerator), j, nc)

	ret <- matrix(numerator/denominator, j, nc, dimnames=list(colnames(sel), paste("comp",1:nc,sep="")))
	ret
}
f.estep <- cmpfun(f.estep)

# S-Step-Function
f.sstep <- function(em.tau,j,nc){
	sem.tau <- matrix(0,j,nc)
	for (i in 1:j){
		sem.tau[i,] <- t(rmultinom(1,1,em.tau[i,]))
	}
	sem.tau
}
f.sstep <- cmpfun(f.sstep)

# M-Step Function
f.mstep <- function(theta.v0, choice_x, self_x, other_x, self_y, other_y, indicators_x, indicators_y, sel, nc, k, n, j, tau) {
	ret <- optim(theta.v0, f.mloglik, f.mloglik.grad, method="BFGS", control=list(maxit=10000, trace=0, fnscale=-1), hessian=F, choice_x=choice_x, self_x=self_x, other_x=other_x, self_y=self_y, other_y=other_y, indicators_x=indicators_x, indicators_y=indicators_y, sel=sel, nc=nc, k=k, n=n, j=j, tau=tau)

	mix <- colMeans(tau)

	list(mix=mix, theta.v=ret$par)
}
f.mstep <- cmpfun(f.mstep)

# Annealing-Part-Function for the SAEM-Algorithm
f.annealing <- function(em.mix, em.vm, sem.mix, sem.vm, iter, b=20){
	{if (iter <= b)
		{q <- cos((iter/b)*acos(b/100))}
	 else
		{q <- (b/100)*sqrt(b/iter)}
	}

	mix <- q*sem.mix+(1-q)*em.mix
	vm  <- q*sem.vm+(1-q)*em.vm

	list(mix=mix, vm=vm)
}
f.annealing <- cmpfun(f.annealing)

# Sandwitch part of the clustered standard errors
f.sandwich <- function(psi, choice_x, self_x, other_x, self_y, other_y, indicators_x, indicators_y, sel, nc, k, n, j) {
	lpsi    <- length(psi);
	theta.v <- psi[1:(lpsi-nc+1)];
	theta.m <- matrix(theta.v, k+1, nc);
	mix     <- psi[(lpsi-nc+2):lpsi];
	mix     <- c(mix, 1-sum(mix));
	m       <- matrix(mix, j, nc, byrow=T);

	# Individual gradient
	nj           <- n*j;
	choice_x.k <- matrix(as.vector(choice_x), nj, k);
	sel.k        <- matrix(as.vector(sel)       , nj, k);
	self_x.k   <- matrix(as.vector(self_x)  , nj, k);
	other_x.k   <- matrix(as.vector(other_x)  , nj, k);
	self_y.k  <- matrix(as.vector(self_y) , nj, k);
	other_y.k  <- matrix(as.vector(other_y) , nj, k);

	cdens.m <- f.compdens(theta.v, choice_x, self_x, other_x, self_y, other_y, indicators_x, indicators_y, sel, nc ,k, n, j);
	denom.outerderiv   <- 1/matrix(rowSums(m*cdens.m), n, j, byrow=T);
	denom.outerderiv.k <- matrix(as.vector(denom.outerderiv), nj, k);

	ret.theta <- matrix(0, nj, nc*(k+1));
	ret.mix   <- matrix(0, nj, nc);

	for (c in 1:nc) {
		gamma <- theta.m[k+1, c];
		beta  <- theta.m[1:k, c];

		lli   <- matrix(as.vector(indicators_x  %*% beta), n, j); #Left linear index
		rli   <- matrix(as.vector(indicators_y %*% beta), n, j); #Right linear index

		ul    <- (1-lli)*self_x  + lli*other_x;  #Utility from "left" choice
		ur    <- (1-rli)*self_y + rli*other_y; #Utility from "right" choice
		ul.k  <- matrix(as.vector(ul), nj, k);
		ur.k  <- matrix(as.vector(ur), nj, k);

		innerderiv   <- mix[c]*matrix(cdens.m[,c],n,j,byrow=T)/((exp(gamma*ul)/(exp(gamma*ul)+exp(gamma*ur)))^choice_x * (exp(gamma*ur)/(exp(gamma*ul)+exp(gamma*ur)))^(1-choice_x));
		innerderiv.k <- matrix(as.vector(innerderiv), nj, k);

		gamma.grad <- sel * denom.outerderiv * innerderiv * ((1/(exp(gamma*ul)+exp(gamma*ur))^2) * ( (exp(gamma*ul)/(exp(gamma*ul)+exp(gamma*ur)))^choice_x * exp(gamma*(ul + 2*ur))*(-ur + ul) * (exp(gamma*ur)/(exp(gamma*ul)+exp(gamma*ur)))^(-choice_x) * (choice_x * exp(-gamma*ul) - exp(-gamma*ur) + choice_x * exp(-gamma*ur)) ));

		beta.grad <- sel.k * denom.outerderiv.k * innerderiv.k * ((-1/(exp(gamma*ul.k)+exp(gamma*ur.k))^2) * ( (exp(gamma*ul.k)/(exp(gamma*ul.k)+exp(gamma*ur.k)))^choice_x.k * gamma*exp(gamma*(ul.k + 2*ur.k))*(indicators_x*(self_x.k-other_x.k) + indicators_y*(other_y.k-self_y.k)) * (exp(gamma*ur.k)/(exp(gamma*ul.k)+exp(gamma*ur.k)))^(-choice_x.k) * (choice_x.k * exp(-gamma*ul.k) - exp(-gamma*ur.k) + choice_x.k * exp(-gamma*ur.k)) ));

		mix.grad  <- cdens.m[,c] * denom.outerderiv[1,]

		ret.theta[,((c-1)*(k+1)+1):(c*(k+1))] <- cbind(beta.grad, as.vector(gamma.grad));
		ret.mix[,c]                           <- as.vector(matrix(mix.grad/n,n,j,byrow=T));
	};

	ret.mix <- ret.mix - matrix(ret.mix[,nc], nj, nc)
	ret     <- cbind(ret.theta, ret.mix[,1:(nc-1)]);

	#Sandwich part
	sandwich <- matrix(0, lpsi, lpsi)
	for (i in 1:j) {
		gradi <- colSums(ret[((i-1)*n+1):(i*n),])
		sandwich <- sandwich + (gradi %o% gradi)
	}

	qc <- (nj-1) / (nj-lpsi) * j/(j-1) #df adjustment

	list(qc=qc, sandwich=sandwich)
}
f.sandwich <- cmpfun(f.sandwich)

# Function to compute clustered standard errors, relies on f.sandwich
f.clusterse <- function(estobj, dat, persid, label_choice_x, label_indicators_x, label_indicators_y, self_x, other_x, self_y, other_y, nc, sortrow=1) {
	prep         <- f.dataprep(dat, persid, label_choice_x, label_indicators_x, label_indicators_y, self_x, other_x, self_y, other_y, nafiller=1);
	n            <- prep$n;
	j            <- prep$j;
	k            <- prep$k;
	indicators_x <- prep$indicators_x;
	indicators_y <- prep$indicators_y;
	self_x       <- prep$self_x;
	other_x      <- prep$other_x;
	self_y       <- prep$self_y;
	other_y      <- prep$other_y;
	choice_x     <- prep$choice_x;
	sel          <- prep$sel;
	rm(prep);

	ll1 <- f.totloglik(estobj$ret$par, choice_x, self_x, other_x, self_y, other_y, indicators_x, indicators_y, sel, nc, k, n, j)
	if (abs(ll1 - estobj$ret$value)>0.001) stop(paste("Saved loglik", ll1, "and computed loglik", estobj$value, "differ"))

	sandwich <- f.sandwich(estobj$ret$par, choice_x, self_x, other_x, self_y, other_y, indicators_x, indicators_y, sel, nc, k, n, j)

	lpsi <- length(estobj$ret$par)
	vcm <- solve(-estobj$ret$hessian)
	vcm <- sandwich$qc * (vcm %*% sandwich$sandwich %*% vcm)

	se <- sqrt(diag(vcm))
	lastmix.se <- sqrt( rep(1,nc-1) %*% vcm[(lpsi-(nc-2)):lpsi, (lpsi-(nc-2)):lpsi] %*% rep(1,nc-1) )
	psi.se <- c(se,lastmix.se)
	psi    <- estobj$ret$par

	#Extract the point estimates
	theta.v    <- psi[1:(lpsi-nc+1)];
	mix        <- psi[(lpsi-nc+2):lpsi];
	mix        <- c(mix, 1-sum(mix));

	mix.m      <- matrix(mix,       1, nc, dimnames=list(c("mixture"), paste("comp",1:nc,sep="")));
	theta.m    <- matrix(theta.v, k+1, nc, dimnames=list(c(paste(label_indicators_x,label_indicators_y),"choicesens"), paste("comp",1:nc,sep="")));

	#Extract the standard errors
	theta.v.se <- psi.se[1:(lpsi-nc+1)];
	mix.se     <- psi.se[(lpsi-nc+2):(lpsi+1)];

	mix.m.se   <- matrix(mix.se,       1, nc, dimnames=list(c("mixture"), paste("comp",1:nc,sep="")));
	theta.m.se <- matrix(theta.v.se, k+1, nc, dimnames=list(c(paste(label_indicators_x,label_indicators_y),"choicesens"), paste("comp",1:nc,sep="")));

	pe.out <- rbind(mix.m,    theta.m);
	se.out <- rbind(mix.m.se, theta.m.se);

	#Sort the groups according to their size to avoid problems of label switching
	sortidx  <- order(pe.out[sortrow,]);
	pe.out   <- pe.out[, sortidx];
	se.out   <- se.out[, sortidx];

	p.out  <- 2*(1-pnorm(abs(pe.out)/se.out));

	cat("POINT ESTIMATES\n")
	print(round(pe.out,3));
	cat("\n------------------------------------------------\n");
	cat("ROBUST STANDARD ERRORS\n")
	print(round(se.out,3));
	cat("\n------------------------------------------------\n");
	cat("ROBUST P-VALUES\n")
	print(round(p.out,3));

	list(se=se.out, p=p.out, vcm=vcm)
}
f.clusterse <- cmpfun(f.clusterse)

# Main Estimation Function
f.estim <- function(dat, v0path, persid, nc, label_choice_x, label_indicators_x, label_indicators_y, label_self_x, label_other_x, label_self_y, label_other_y, saem=F, loglik1=NA, maxemiter=50, diffcrit=5e-3, sortrow=1) {

	# Prepares the data and stores it
	prep       <- f.dataprep(dat, persid, label_choice_x, label_indicators_x, label_indicators_y, label_self_x, label_other_x, label_self_y, label_other_y, nafiller=1)
	n          <- prep$n
	j          <- prep$j
	k          <- prep$k
	indicators_x      <- prep$indicators_x
	indicators_y     <- prep$indicators_y
	self_x   <- prep$self_x
	other_x   <- prep$other_x
	self_y  <- prep$self_y
	other_y  <- prep$other_y
	choice_x <- prep$choice_x
	sel        <- prep$sel
	rm(prep)

	# Prepares the start values
	v0      <- read.table(v0path, sep=",", header=F)[,1]
	theta.v <- runif(nc*(1+k),0.8,1.2) * rep(v0,nc) # Random variation of the start values from the pooled model
	aa      <- runif(nc, 0.5, 2)
	mix     <- aa/sum(aa) # Initial mixture
	rm(aa)

	# Initial EM control flow parameters
	iter  <- 1
	diff  <- 1
	llold <- 0
	llnew <- 1

	# EM/SAEM algorithm (see Dempster et al. (1977) and book by McLachlan (2000))
	cat("Starting EM type maximization...\n")
	while (iter < maxemiter+1 & abs(diff) > diffcrit) {
		cat("\n  EM-Type Iteration:", iter,"of max.",maxemiter,"\n")
		cat("\n     working on E-step.\n")
		tau.em <- f.estep(theta.v, choice_x, self_x, other_x, self_y, other_y, indicators_x, indicators_y, sel, nc, k, n, j, mix)
		cat("     working on M-step (EM algorithm).\n")
		mstep <- f.mstep(theta.v, choice_x, self_x, other_x, self_y, other_y, indicators_x, indicators_y, sel, nc, k, n, j, tau.em)
		theta.v.em <- mstep$theta.v
		mix.em     <- mstep$mix

		if (saem==T) {
			cat("\n     working on S-step.\n")
			tau.sem <- f.sstep(tau.em, j, nc)

			cat("     working on M-step (SEM algorithm).\n")
			mstep <- f.mstep(theta.v, choice_x, self_x, other_x, self_y, other_y, indicators_x, indicators_y, sel, nc, k, n, j, tau.sem)
			theta.v.sem <- mstep$theta.v
			mix.sem     <- mstep$mix

			ann <- f.annealing(mix.em, theta.v.em, mix.sem, theta.v.sem, iter, b=20)
			mix     <- ann$mix
			theta.v <- ann$v
		 } else {
		 	mix     <- mix.em
		 	theta.v <- theta.v.em
		 }

		# Update the control flow parameters
		iter  <- iter+1
		llold <- llnew
		psi   <- c(theta.v, mix[1:(length(mix)-1)])
		llnew <- f.totloglik(psi, choice_x, self_x, other_x, self_y, other_y, indicators_x, indicators_y, sel, nc, k, n, j)
		diff  <- llnew-llold
		cat("\n  Achieved Log Likelihood:", round(llnew,4),"\n\n  ----------\n")
	}
	cat("\n...done. Parameters after EM-type estimation:\n")
	mix.m   <- matrix(mix,       1, nc, dimnames=list(c("mixture"), paste("comp",1:nc,sep="")))
	theta.m <- matrix(theta.v, k+1, nc, dimnames=list(c(paste(label_indicators_x,label_indicators_y),"choicesens"), paste("comp",1:nc,sep="")))
	out     <- rbind(mix.m, theta.m)

	#Direct ML estimation
	ret  <- optim(psi,f.totloglik, f.totloglik.gradient, method="BFGS", control=list(maxit=100000,trace=1,REPORT=20, fnscale=-1), hessian=T, choice_x=choice_x, self_x=self_x, other_x=other_x, self_y=self_y, other_y=other_y, indicators_x=indicators_x, indicators_y=indicators_y, sel=sel, nc=nc, k=k, n=n, j=j)

	psi     <- ret$par #Extract the MLE
	lpsi    <- length(psi)
	psi.se <- sqrt(diag(solve(-ret$hessian))) #Standard errors

	#Extract the point estimates
	theta.v    <- psi[1:(lpsi-nc+1)]
	mix        <- psi[(lpsi-nc+2):lpsi]
	mix        <- c(mix, 1-sum(mix))

	mix.m      <- matrix(mix,       1, nc, dimnames=list(c("mixture"), paste("comp",1:nc,sep="")))
	theta.m    <- matrix(theta.v, k+1, nc, dimnames=list(c(paste(label_indicators_x,label_indicators_y),"choicesens"), paste("comp",1:nc,sep="")))

	#Extract the standard errors
	theta.v.se <- psi.se[1:(lpsi-nc+1)]
	mix.se   <- psi.se[(lpsi-nc+2):lpsi]
	mix.se   <- c(mix.se,NA)

	mix.m.se   <- matrix(mix.se,       1, nc, dimnames=list(c("mixture"), paste("comp",1:nc,sep="")))
	theta.m.se <- matrix(theta.v.se, k+1, nc, dimnames=list(c(paste(label_indicators_x,label_indicators_y),"choicesens"), paste("comp",1:nc,sep="")))

	tau     <- f.estep(theta.v, choice_x, self_x, other_x, self_y, other_y, indicators_x, indicators_y, sel, nc, k, n, j, mix)

	aic <- -2*ret$value + 2*lpsi
	bic <- -2*ret$value + log(sum(sel))*lpsi

	tauv   <- as.vector(tau)
	nec    <- sum(na.omit(-tauv*log(tauv)))/(ret$value-loglik1)

	i.out  <- matrix(c(ret$value, lpsi, sum(sel), sortrow, aic, bic, nec), , 1, dimnames=list(c("Loglik:", "# of parameters:", "# of observations:", "Sortrow (to avoid label switching):", "AIC:", "BIC:", "NEC:"), c()))
	pe.out <- rbind(mix.m,    theta.m)
	se.out <- rbind(mix.m.se, theta.m.se)

	#Sort the groups according to their size to avoid problems of label switching
	sortidx  <- order(pe.out[sortrow,])
	pe.out   <- pe.out[, sortidx]
	se.out   <- se.out[, sortidx]
	tau      <- tau[, sortidx]

	p.out  <- 2*(1-pnorm(abs(pe.out)/se.out))

	cat("\n------------------------------------------------\n")
	cat("GENERAL INFORMATION\n")
	print(round(i.out,3))
	cat("\n------------------------------------------------\n")
	cat("POINT ESTIMATES (NOT CLUSTERED)\n")
	print(round(pe.out,3))
	cat("\n------------------------------------------------\n")
	cat("STANDARD ERRORS (NOT CLUSTERED)\n")
	print(round(se.out,3))
	cat("\n------------------------------------------------\n")
	cat("P-VALUES\n")
	print(round(p.out,3))
	cat("\n------------------------------------------------\n")

	list(pe=pe.out, se=se.out, p=p.out, info=i.out, ret=ret, tau=tau, success=T)
}
f.estim <- cmpfun(f.estim)

# Load choice data from both experiments
dat1     <- read.table("choices_exp1.csv",  sep=",", header=T)
dat2     <- read.table("choices_exp2.csv",  sep=",", header=T)

# Remove erratic subjects (see 2nd paragraph of section 4; to replicate estimate individual parameters)
dropped  <- read.table("dropped_subjects_section4paragraph2.csv", sep=",", header=F)[,1]
dat1     <- dat1[(dat1$sid %in% dropped) == F, ]
dat2     <- dat2[(dat2$sid %in% dropped) == F, ]

# Estimating Finite Model with nc=3 Types using start values from aggregate estimation
cat("\n\n\n================== ESTIMATING SESSION 1 ================================\n")
pe1 <- f.estim(dat1, "svExp1_1cl.csv", "sid", nc=3, "choice_x", c("s_x","r_x","q","v"), c("s_y","r_y","q","v"), "self_x", "other_x", "self_y", "other_y", saem=T, loglik1=-5472.314, maxemiter=30, diffcrit=1e-3,sortrow=1)

cat("\n\n\n================== ESTIMATING SESSION 2 ================================\n")
pe2 <- f.estim(dat2, "svExp2_1cl.csv", "sid", nc=3, "choice_x", c("s_x","r_x","q","v"), c("s_y","r_y","q","v"), "self_x", "other_x", "self_y", "other_y", saem=T, loglik1=-4540.739, maxemiter=30, diffcrit=1e-3,sortrow=1)

cat("\n\nTo extract the individual ex-post probabilities of type membership type pe1$tau or pe2$tau.\n")

# Computing Robust Standard Errors
cat("\n\n\n================== ROBUST STANDARD ERRORS: SESSION 1 ===================\n")
rse1 <- f.clusterse(pe1, dat1, "sid", "choice_x", c("s_x","r_x","q","v"), c("s_y","r_y","q","v"), "self_x", "other_x", "self_y", "other_y", 3, sortrow=1)
cat("\n\n\n================== ROBUST STANDARD ERRORS: SESSION 2 ===================\n")
rse2 <- f.clusterse(pe2, dat2, "sid", "choice_x", c("s_x","r_x","q","v"), c("s_y","r_y","q","v"), "self_x", "other_x", "self_y", "other_y", 3, sortrow=1)

# Save the R workspace
save.image("finmix_est_3cl.RData")


suppressPackageStartupMessages(library(dplyr))
suppressPackageStartupMessages(library(ggplot2))

# Set themes
file_extensions <- c(".png", ".pdf")

plot_theme <- theme(
    title = element_text(size = 9),
    axis.title = element_text(size = 9),
    legend.text = element_text(size = 7),
    legend.title = element_text(size = 9),
    legend.key.size = unit(0.5, "cm"),
    strip.text = element_text(size = 10),
    strip.background = element_rect(colour="black", fill="#fdfff4")
)

# Load L2 distances per MOA
cp_file <- file.path("data", "MOA_LSA_zscores.tsv")
cp_df <- readr::read_tsv(cp_file)

cp_df$model <- factor(
    cp_df$model,
    levels = c("Vanilla", "Beta", "MMD", "PCA", "Complete")
)

head(cp_df)

# Compare performance means across architectures
cp_mean_df <- cp_df %>%
    dplyr::group_by(model) %>%
    dplyr::mutate(
        model_mean = round(mean(zscore), 2)
    ) %>%
    dplyr::select(model, model_mean) %>%
    dplyr::distinct()

cp_mean_df

plot_gg <- (
    ggplot(cp_df, aes(x=zscore))
    + geom_histogram(aes(y = ..density..), bins=20)
    + geom_density()
    + geom_text(data=cp_mean_df, aes(x = -2, y = 0.35, label = paste0("Mean: ", model_mean)))
    + geom_vline(data=cp_mean_df, aes(xintercept = model_mean), color = "red", linetype="dashed")
    + facet_wrap("model", nrow = 5)
    + geom_vline(xintercept = 0, linetype = "dashed", color = "blue")
    + theme_bw()
    + plot_theme
    + ylab("Density")
    + xlab("Z-score of L2 distances between real and simulations\nfor predicted polypharmacology state compared to 10 random permutations")

)
plot_gg

# Save figure
output_file_base <- file.path("output", "sup_fig_lsa_zscores")
for (file_extension in file_extensions) {
    output_file <- paste0(output_file_base, file_extension)
    ggplot2::ggsave(output_file, plot_gg, dpi = 500, width = 6, height = 8)
}


#Jenny Smith

#July 11, 2017

#purpose: plot the highest genes fold-changes

barplot <- function(limma.res, decile=NULL) {

  dec <- quantile(limma.res$logFC, probs = seq(0,1,length= 11), type=5)
  limma.res$gene <- rownames(limma.res)

  if (is.null(decile)){
    limma.res <- limma.res
  }else{
    decile.90 <- dec[10]
    decile.10 <- dec[2]
    limma.res <- limma.res[which(limma.res$logFC >= decile.90 | limma.res$logFC <= decile.10), ]
  }
  limma.res$cut <- ifelse(limma.res$logFC > 0, "up", "dn")

  ggplot(limma.res, aes(x=reorder(limma.res$gene, limma.res$logFC), y=logFC, fill=cut)) +
    geom_bar(stat = "identity") + theme_JS + labs(title="Differentially Regulated Genes: Largest FC",x="") +
    scale_fill_manual(values = c(up="red", dn="darkgreen"))
}


# variable name for 16S data is 'd'
# basically same script as for yeast

if(!exists('d')){
  load(paste(my.path,'data/tiyani_pup_vs_ys.Rdata', sep=''))
}

if(!exists('test.ty.eff')){
	conds.ty <- c(rep("Pup", 161), rep("YS", 212))

	t <- aldex.clr(d, conds=conds.ty)
	t.e <- aldex.effect(t)

	# reference TP names
	ty.ol <- rownames(t.e)[which(t.e$overlap < 0.05)]
	ty.diff <- rownames(t.e)[which(abs(t.e$diff.btw) > 1)]
	ty.eff <- rownames(t.e)[which(abs(t.e$effect) > 1)]

	ty.diff.ol <- intersect(ty.ol, ty.diff)
	ty.diff.eff <- intersect(ty.eff, ty.diff)
	ty.ol.eff <- intersect(ty.eff, ty.ol)

	# an example FP dataset
	n.samples = 5
	test.col <- c(sample(1:161, n.samples), sample(162:373, n.samples))
	test.conds <- c(rep("Pup", n.samples), rep("YS", n.samples))
	test.data <- d[,test.col]

	t.s <- aldex.clr(test.data, test.conds)
	t.s.e <- aldex.effect(t.s)

	test.ty.ol <- rownames(t.s.e)[which(t.s.e$overlap < 0.1)]
	test.ty.diff <- rownames(t.s.e)[which(abs(t.s.e$diff.btw) > 1)]
	test.ty.eff <- rownames(t.s.e)[which(abs(t.s.e$effect) > 1)]
} # end test for test.ty.eff


require("shiny")
require("reshape")
require("ggplot2")
require("plotly")
require("ggfortify")

PCA.tabPanel<-function(profile){
	#tabPanel("PCA",verbatimTextOutput("PCA"))
  tabPanel("PCA",div(id = "eggHead.container"
          , div(class = "plotlybars-wrapper"
            , div( class="plotlybars"
              , div(class="plotlybars-bar b1")
              , div(class="plotlybars-bar b2")
              , div(class="plotlybars-bar b3")
              , div(class="plotlybars-bar b4")
              , div(class="plotlybars-bar b5")
              , div(class="plotlybars-bar b6")
              , div(class="plotlybars-bar b7")
            )
            , div(class="plotlybars-text"
              , p("loading")
            )
          )
          ,plotlyOutput("PCA")),
			h5("variance"),plotOutput("PCA_variance_plot"))
}

PCA.sidebarPanel<-function(profile,input,session){

  conditionalPanel(condition = "input.tabs1 == 'PCA'",
				selectInput("pca_xc","Select x component",colnames(PCA.pcadata()$x)[1:min(5,ncol(PCA.pcadata()$x))],selected="PC1"),
				selectInput("pca_yc","Select y component",colnames(PCA.pcadata()$x)[1:min(5,ncol(PCA.pcadata()$x))],selected="PC2")
                #selectInput("PCA.type","Test",levels(PCA.data()$variable))

	)

}

#PCA.prepareData<-function(input){
  PCA.mergeData<-reactive({
			data<-read.table(fileLocate("matrix.tsv"),header=T,sep="\t",check.names = F)
			data<-merge(sampleInfo,data,by="SampleID")
			data
	})
  PCA.sampleData<-reactive({
	sdata=share.mergeData()[,c(1:ncol(sampleInfo))]
  })
  PCA.pcadata <- reactive({
			data=share.mergeData()[,-c(1:ncol(sampleInfo))]
			data=data[,which(apply(data, 2, var, na.rm=TRUE)!=0)] #remove constant
			for(i in 1:ncol(data)){
				data[is.na(data[,i]), i] <- mean(data[,i], na.rm = TRUE) #set NA as average
			}
			prcomp(data,scale=T)
     })
	PCA.variance <- reactive({
		vars <- apply(PCA.pcadata()$x, 2, var)
		props <- vars / sum(vars)

	})
	PCA.drawdata <- reactive({
		if(file.exists(fileLocate("pca.tsv"))){
			data<-read.delim(fileLocate("pca.tsv"),check.names=F)
			data<-merge(sampleInfo,data,by="SampleID")
		}
		else{
			data<-PCA.pcadata()$x[,1:min(5,ncol(PCA.pcadata()$x))]
			data<-cbind(PCA.sampleData(),data)
			rdata=data
			rdata$SampleID=sampleRIDs[rdata$SampleID]
			write.table(rdata[,-c(2:ncol(sampleInfo))],fileLocate("pca.tsv"),sep="\t",row.names=F)
			data<-cbind(PCA.sampleData(),data[,c(input$pca_xc,input$pca_yc)])

		}
		data$SampleID=sampleRIDs[data$SampleID]
		data
	})
#}

#PCA.preparePlot<-function(input,output){

  output$PCA <- renderPlotly({
   # p<-autoplot(PCA.pcadata(), data=sampleInfo,colour=input$group_by,samid="SampleID",frame = TRUE, frame.type = 'norm')
	p<-ggplot(PCA.drawdata(),aes_string(x=input$pca_xc,y=input$pca_yc,color=input$group_by))+
		geom_point(aes(samid=SampleID))+stat_ellipse()
	ggplotly(p)
  })
	output$PCA_variance_plot <- renderPlot({
      plot(PCA.pcadata(),type="l")
    })
#}


# Usage eDNAindex(x, sample, taxa, nReads) all elements without quotes

eDNAindex<- function(df, Sample_column, OTU_column, Counts_column, Biological.replicate, ... ){ #

  require(tidyverse)
  require(rlang)

  Sample_column <- rlang::enquo(Sample_column)
  OTU_column    <- rlang::enquo(OTU_column)
  Counts_column <- rlang::enquo(Counts_column)
  Biological.replicate <- rlang::enquo(Biological.replicate)
  #vars.to.add   <- rlang::enquos(...)

  df %>%
    ungroup() %>%
    select (!!Sample_column, !!OTU_column, ...) %>%
    group_by(!!Sample_column, !!OTU_column) %>%
    summarise_all(first) -> matching.df

  # if(ncol(matching.df) > 2){
  #   print("Adding extra vars")
  #
  #   matching.df %>%
  #     select(-!!Sample_column, -!!OTU_column) %>%
  #     sapply(., typeof) -> types.of.vars
  #
  #  if(str_detect(types.of.vars, "double")){
  #
  #   matching.df %>%
  #     summarise_if(is.numeric,mean) -> num.vars
  #  }
  #
  #   if(str_detect(types.of.vars, "character")){
  #
  #     matching.df %>%
  #       summarise_if(is.character,first) -> char.vars
  #   }
  #
  #
  #
  #   left_join(num.vars, char.vars) -> matching.df
  # }

  if (quo_name(Biological.replicate) %in% colnames(df)){

    print ("Averaging ratios between Biological replicates")

    df %>%

      group_by(!!Sample_column, !!OTU_column, !!Biological.replicate) %>%

      summarise ( sumreads = sum(!!Counts_column)) %>%  # This sums technical replicates

      group_by(!!Sample_column,!!Biological.replicate) %>%

      mutate (Tot = sum(sumreads),
              Row.prop = sumreads / Tot)  %>%                      # This creates the proportion on each biological replicate

      group_by(!!Sample_column) %>%

      mutate (nreps = length(unique(!!Biological.replicate))) %>%

      group_by(!!Sample_column, !!OTU_column) %>%

      summarise (mean.prop = sum (Row.prop) / max(nreps))   %>%

      group_by (!!OTU_column) %>%

      mutate (Colmax = max (mean.prop),
              Normalized.reads = mean.prop / Colmax) %>%

      dplyr::select( -Colmax, -mean.prop) -> output
    #return(output)
  }else{

  print("Calculating eDNAindex directly")

  # IF THERE ARE TECHNICAL REPLICATES, WE NEED TO SUM THOSE VALUES FIRST
  #return(!!Technical_replicate)

  df %>%

    group_by(!!Sample_column, !!OTU_column) %>%

    summarise (sumreads = sum(!!Counts_column)) %>% # In case the sample column is of a higher group than the nrows

    group_by(!!Sample_column) %>%

    mutate (Tot = sum(sumreads),
            Row.prop = sumreads / Tot) %>%

    group_by (!!OTU_column) %>%

    mutate (Colmax = max (Row.prop),
            Normalized.reads = Row.prop / Colmax) %>%
    dplyr::select(-Tot, -Row.prop, -Colmax, -sumreads) -> output #note, specifying dplyr::select to avoid conflict w MASS package
  }

  left_join(output, matching.df) -> output

  return (output)
}


rm(list=ls())

file.eigenvect   <- commandArgs()[4]
file.annotation  <- commandArgs()[5]

eigenvect <- read.table(file=file.eigenvect, skip=1, col.names=c("FID", "IID", "PC1", "PC2", "PC3", "PC4", "PC5", "PC6", "PC7", "PC8", "PC9", "PC10", "batch"))
annotation       <- read.table(file=file.annotation, header=T, sep="\t")

#####matched <- match(as.character(eigenvect$indivID), as.character(annotation$individualID))

# ------------------------------ #
# -- draw histograms of PC1-6 -- #
# ------------------------------ #

library('sm')
library('plotrix')
#color_code <- c(1:(length(levels(annotation$diagnosis))-1))
color_code <- c(1:(length(levels(annotation$diagnosis))))
#legend_code=c("Control", "CD", "UC", "AS", "PS", "PSC")
legend_code= sort(unique(annotation$diagnosis))

for (i in seq(3,12)) {
    png(file=paste(file.eigenvect, ".histPC", i-2, ".png", sep=""), width=900, height=720)
    par(mfrow=c(2,2))

    a = eigenvect[annotation$diagnosis == legend_code[1], i]
    b = eigenvect[annotation$diagnosis == legend_code[2], i]
    #a = eigenvect[annotation$diagnosis == "Control", i]
    #b = eigenvect[annotation$diagnosis == "CD", i]
    #c = eigenvect[annotation$diagnosis == "UC", i]
    #d = eigenvect[annotation$diagnosis == "AS", i]
    #e = eigenvect[annotation$diagnosis == "PS", i]
    #f = eigenvect[annotation$diagnosis == "PSC", i]
    a = a[!is.na(a)]
    b = b[!is.na(b)]
    #c = c[!is.na(c)]
    #d = d[!is.na(d)]
    #e = e[!is.na(e)]
    #f = f[!is.na(f)]

    print(c(min(a),max(a)))

    #multhist(list(a,b,c,d,e,f), col=color_code, xlab=paste("PC", i-2, sep=""), ylab="Counts", freq=T)
    multhist(list(a,b), col=color_code, xlab=paste("PC", i-2, sep=""), ylab="Counts", freq=T)
    legend("topright", legend=legend_code, fill=color_code)

    #multhist(list(a,b,c,d,e,f), col=color_code, xlab=paste("PC", i-2, sep=""), ylab="Density", freq=F)
    multhist(list(a,b), col=color_code, xlab=paste("PC", i-2, sep=""), ylab="Density", freq=F)
    legend("topright", legend=legend_code, fill=color_code)

    #disease.f <- factor(annotation$diagnosis, levels=c("Control", "CD", "UC", "AS", "PS", "PSC"), labels=c("Control", "CD", "UC", "AS", "PS", "PSC"))
    disease.f <- factor(annotation$diagnosis, levels=legend_code, labels=legend_code)
    colfill <- c(1:(1+length(levels(disease.f))))
    #sm.density.compare(eigenvect[,i], annotation$diagnosis, xlab=paste("PC", i-2, sep=""))
    sm.density.compare(eigenvect[,i], annotation$diagnosis, xlab=paste("PC", i-2, sep=""), col=color_code)
    legend("topright", levels(disease.f), fill=colfill)

    dev.off()
}


[comment This code is written in REBOL 3.]
[comment This code is written in REBOL 3.]
append engScript_a engScript_b
append engScript_a engScript_b
(4 == 5)
engScript_b: to-hash [[ f o o   1   b a r   2 ]]
engScript_a: "haha"
engScript_b: to-hash [[ f o o   1   b a r   2 ]]
engScript_add: func [engScript_a [integer!], engScript_b [integer!]] [
    return (engScript_a + engScript_b)
]
engScript_multiply: func [engScript_a [integer!], engScript_b [integer!]] [
    return (engScript_a * engScript_b)
]
engScript_subtract: func [engScript_a [integer!], engScript_b [integer!]] [
    return (engScript_a - engScript_b)
]
engScript_squared: func [engScript_a [integer!]] [
    return (engScript_a * engScript_a)
]


# purpose: merge and vis the Diallel pheno
# Jinliang Yang
# 2.1.2012

setwd("/Users/yangjl/Documents/Heterosis_GWAS/pheno2011");
f1 <- read.csv("pheno_diallel_master_BLUE.csv");
parents <- read.csv("pheno_namparents_BLUE.csv")
heterosis <- read.csv("pheno_diallel_master_BLUE_heterosis.csv")
sca <- read.csv("pheno_diallel_master_BLUP_SCA.csv")
gca <- read.csv("pheno_diallel_master_BLUP_GCA.csv")
f1$type <- "F1"
parents$type <- "P"

pheno <- rbind(f1, parents)
### Normalize the pheno ###
for (i in 2:8){
  big <- max(pheno[, i])
  small <- min(pheno[,i])
  pheno[, i] <- (pheno[,i]-small)/(big-small);
}

####################Trait per se ##########################################
par(mfrow=c(2,4))
cols <- c("red", "bisque4","blue3","cadetblue4","chartreuse4","chocolate4", "darkslategray")
for (i in 2:8){
  plot(density(subset(pheno, type=="F1")[,i]), col=cols[i-1], xlab="Normalized Scale", lwd=3, xlim=c(-0.3, 1.3), ylim=c(0, 8), main=names(pheno)[i])
  lines(density(subset(pheno, type=="P")[,i]), col=cols[i-1], lty=2, lwd=3)
  #abline(v=subset(pheno, Genotype=="B73")[,i], col=cols[i-2], lwd=3)
}

### Plot the normalized density plot ###
plot(density(subset(pheno, type=="F1")$KRN), main="All traits", xlab="Normalized Scale", col="red", lwd=2, xlim=c(-0.3, 1.3), ylim=c(0, 8))
lines(density(subset(pheno, type=="P")$KRN), col="red", lty=2, lwd=2)

for (i in 3:8){
  lines(density(subset(pheno, type=="F1")[,i]), col=cols[i-1], lwd=2)
  lines(density(subset(pheno, type=="P")[,i]), col=cols[i-1], lty=2, lwd=2)
  #abline(v=subset(pheno, Genotype=="B73")[,i], col=cols[i-2], lwd=3)
}
#abline(v=subset(pheno, Genotype=="B73")$KRN, col="red", lwd=3)
###################### HPH & MPH ########################################
### Normalizing the heterosis ###
for (i in 6:26){
  big <- max(heterosis[, i])
  small <- min(heterosis[,i])
  heterosis[, i] <- (heterosis[,i]-0)/(big-small);
}

par(mfrow=c(2,4))
cols <- c("red", "bisque4","blue3","cadetblue4","chartreuse4","chocolate4", "darkslategray")
hph <- heterosis[, c(13,15,17,19,21,23,25)]
mph <- heterosis[, c(14,16,18,20,22,24,26)]
nm <- c("KRN HPH & MPH", "CL HPH & MPH", "CW HPH & MPH", "CD HPH & MPH", "AKW HPH & MPH", "TKW HPH & MPH", "KC HPH & MPH")
for (i in 1:7){
  plot(density(mph[,i]), col=cols[i], lty=2, lwd=3, main=nm[i])
  lines(density(hph[,i]), col=cols[i], lwd=3)
  #abline(v=subset(pheno, Genotype=="B73")[,i], col=cols[i-2], lwd=3)
}

source("~/Documents/Codes/reusableRcode.r")
pairs(heterosis[, 6:26], text.panel = diag, upper.panel=panel.smooth,
      lower.panel=panel.cor, gap=0, main="Correlation Plots of the HPH")

####################### GCA ##############################################################################
### Normalizing the heterosis ###
for (i in 2:8){
  big <- max(gca[, i])
  small <- min(gca[,i])
  gca[, i] <- (gca[,i]-small)/(big-small);
}

gca$mean <- apply(gca[,2:8], 1, mean)
gca <- merge(namid, gca, by.x="pop", by.y="Genotype")
gca <- gca[order(gca$mean, decreasing=TRUE),]

#######
#cols <- c("greenyellow", "honeydew3","indianred1", "khaki2", "lightblue3","mediumorchid3", "mediumorchid1")
cols <- c("lightcoral", "bisque4","royalblue3","plum","chartreuse4","chocolate4", "yellow3")
par(mfrow=c(1,1), las=1,col.axis="royalblue")
hh <- t(gca[,3:9]);
mybarcol <- "gray20"
mp <- barplot(hh, beside = TRUE,
              col = cols,
              names.arg = gca$parent,
              legend = rownames(hh), ylim= c(0,1),
              main = "GCA of NAM parents", font.main = 4,
              col.sub = mybarcol,
              cex.names = 0.8)

#segments(mp, hh, mp, hh + 2*sqrt(1000*hh/100), col = mybarcol, lwd = 1.5)
#stopifnot(dim(mp) == dim(hh))# corresponding matrices
mtext(side = 1, at = colMeans(mp), line = 2,
      text = round(colMeans(hh),2), col = "red")

####################### SCA ##############################################################################
### Normalizing the heterosis ###
for (i in 2:8){
  big <- max(sca[, i])
  small <- min(sca[,i])
  sca[, i] <- (sca[,i]-0)/(big-small);
}

#######
par(mfrow=c(2,4))
#cols <- c("greenyellow", "honeydew3","indianred1", "khaki2", "lightblue3","mediumorchid3", "mediumorchid1")
cols <- c("lightcoral", "bisque4","royalblue3","plum","chartreuse4","chocolate4", "yellow3")
for (i in 2:8){
  plot(density(sca[,i]), col=cols[i-1], lty=2, lwd=3, main=names(sca)[i], xlim=c(-1, 1),ylim=c(0,4), xlab="normalized scale")
  lines(density(sca[,i]), col=cols[i-1], lwd=3)
  #abline(v=subset(pheno, Genotype=="B73")[,i], col=cols[i-2], lwd=3)
}

################################################################


###
### Copyright (C) 2017 Rafael Laboissière
###
### This program is free software: you can redistribute it and/or modify it
### under the terms of the GNU General Public License as published by the
### Free Software Foundation, either version 3 of the License, or (at your
### option) any later version.
###
### This program is distributed in the hope that it will be useful, but
### WITHOUT ANY WARRANTY; without even the implied warranty of
### MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
### General Public License for more details.
###
### You should have received a copy of the GNU General Public License along
### with this program.  If not, see <http://www.gnu.org/licenses/>.

### * Boxplot the results

### ** All conditions, collapsed objects

### *** Threshold
pdf (file = "obj-stab-threshold-all.pdf")
par (mar = c (11, 4, 1, 1))
boxplot (threshold ~ chair * stimulus * background, data = obj.stab.psycho,
         las = 2, ylab = "threshold (degree)")
dev.off ()

### *** Slope
pdf (file = "obj-stab-slope-all.pdf")
par (mar = c (11, 4, 1, 1))
### **** Inclination at threshold is 25*a
boxplot (abs (25 * slope) ~ chair * stimulus * background, log = "y",
         data = obj.stab.psycho, las = 2, ylab = "slope (%/degree)")
dev.off ()

### ** Object condition
df <- droplevels (subset (obj.stab.psycho, stimulus == "object"))

### *** Threshold
pdf (file = "obj-stab-threshold-object.pdf")
par (mar = c (11, 4, 1, 1))
boxplot (threshold ~ object * chair * background, data = df,
         las = 2, ylab = "threshold (degree)")
dev.off ()

### *** Slope
pdf (file = "obj-stab-slope-object.pdf")
par (mar = c (11, 4, 1, 1))
### **** Inclination at threshold is 25*a
boxplot (abs (25 * slope) ~ object * chair * background, log = "y",
         data = df, las = 2, ylab = "slope (%/degree)")
dev.off ()


#' chol
#'
#' Cholesky factorization for a float vector/matrix.
#'
#' @param x
#' A float vector/matrix.
#'
#' @return
#' A float vector/matrix.
#'
#' @examples
#' library(float)
#'
#' s = flrunif(10, 3)
#' cp = crossprod(s)
#' chol(cp)
#'
#' @useDynLib float R_chol_spm
#' @name chol
#' @rdname chol
NULL

chol_float32 = function(x)
{
  ret = .Call(R_chol_spm, DATA(x))
  float32(ret)
}

#' @rdname chol
#' @export
setMethod("chol", signature(x="float32"), chol_float32)


"get.residuals" <- function(asreml.obj, units="ignore")
{
  asr4 <- isASRemlVersionLoaded(4, notloaded.fault = TRUE)

  options <- c("addtoresiduals", "ignore")
  unit.opt <- options[check.arg.values(units, options)]
  if (length(asreml.obj$aom) == 0)
  {
    res <- as.vector(asreml.obj$residuals)
    attr(res, which="restype") <- "Residuals"
    stdres <- FALSE
  } else
  {
    res <- as.vector(asreml.obj$aom$R[,2])
    attr(res, which="restype") <- "Standardized conditional residuals"
    stdres <- TRUE
  }
  if (unit.opt == "addtoresiduals")
  {
    #Check whether units is in random model
#    if (asr4)
      termno <- findterm("units", rownames(summary(asreml.obj)$varcomp))
#    else
#      termno <- findterm("units", rownames(summary(asreml.obj)$varcomp))
    if (termno > 0)
    {
      if (asr4)
      {
        if (stdres)
          uBLUP <- asreml.obj$aom$G[,2]
        else
          uBLUP <- asreml.obj$coefficients$random
      } else
      {
        if (!stdres)
          ucoeff <- 1
        else
        {
#          ucoeff <- asreml::summary.asreml(asreml.obj)$varcomp[termno, "component"]
          ucoeff <- summary(asreml.obj)$varcomp[termno, "component"]
          ucoeff <- sqrt(ucoeff/(ucoeff^2 -1))
        }
        uBLUP <- ucoeff*asreml.obj$coefficients$random
      }
      uBLUP <- uBLUP[grep("units_", substr(names(uBLUP), 1, 6), fixed=TRUE)]
      names(uBLUP) <- NULL
      res <- res + uBLUP
    }
  }
  res[is.nan(res)] <- NA
  return(res)
}

"variofaces.asreml" <- function(asreml.obj, means=NULL, V = NULL, nsim=100, seed = NULL,
                                extra.matrix = NULL, ignore.terms = NULL,
                                fixed.spline.terms = NULL,
                                bound.exclusions = c("F","B","S","C"),
                                tolerance = 1E-10, units = "ignore",
                                update = TRUE, trace = FALSE,
                                graphics.device=NULL, ncores = detectCores(), ...)
#function to do the face variogram plots, including envelopes, described by
#Stefanova et al (2010)
#asreml.obj is an asreml object from a call to asreml in which the data argument
#   must have been set.
#means is a vector of predictions for fixed terms - should include spline terms
#   if null, fitted values from fitted model are used.
#V is the fitted variance matrix i.e. having the appropriate pattern and values
#   given the model fitted and the estimates of the parameters obtained
#nsim is the number of data sets to be simulated in obtaining the envelopes
#   Note that only the results from data sets that converge in the fitting are
#        included in computing the envelope
#tolerance is the value such that eigenvalues less than it are consdered to be zero
#... parameters to supply to asreml calls within variofaces.asreml
{
  asr4 <- isASRemlVersionLoaded(4, notloaded.fault = TRUE)
  #Check that have a valid object of class asreml
  validasr <- validAsreml(asreml.obj)
  if (is.character(validasr))
    stop(validasr)

  options <- c("addtoresiduals", "ignore")
  unit.opt <- options[check.arg.values(units, options)]
  #Check V
  if (is.null(V))
  {
    if (!asr4)
      stop("It appears that asreml4 is not loaded and so V must be supplied to variofaces.asreml.")
    V <- estimateV(asreml.obj, extra.matrix = extra.matrix, ignore.terms = ignore.terms,
                   fixed.spline.terms = fixed.spline.terms,
                   bound.exclusions = bound.exclusions)
    n <- nrow(V)
  }
  else #check supplied V
  {
    if (!isSymmetric(V))
      stop("Variance matrix must be symmetric")
    n <- length(asreml.obj$residuals)
    if (!is.null(means) & length(means) != n)
      stop("The lengths of means  and the response variable are not the same")
    if (!all(dim(V) == c(n, n)))
      stop("V is not a square matrix whose order equals the length of the response variable")
  }

  #use eigenvalue decomposition to establish transformation matrix
  eigdecomp <- eigen(V, symmetric = TRUE)
  eigenval <- eigdecomp$values
  if (!all(eigenval >= -tolerance * abs(max(eigenval))))
     stop("Variance matrix is not nonnegative definite")
  R <- eigdecomp$vectors %*% diag(sqrt(pmax(eigenval, 0)), n)

  #check and get info in supplied call
  call <- asreml.obj$call
  if (!("data" %in% names(call)))
    stop("variofaces.asreml assumes that data has been set in call to asreml")
  env.dat <- eval(call$data)

  #set up call
  object.sim <- asreml.obj
  call <- object.sim$call
  elno <- grep("data", names(call))
  languageEl(call, which=elno) <- quote(env.dat)
  languageEl(call$fixed, which=2) <- quote(y.sim)
  call$trace <- trace

 #Deal with the R.param and G.param arguments
 if (update)
 {
   #If update, set R.param and G.param
   languageEl(call, which = "R.param") <- asreml.obj$R.param
   languageEl(call, which = "G.param") <- asreml.obj$G.param
 } else
 {
   #If R.param and G.param already set, make them NULL
   if (!is.null(languageEl(call, which = "R.param")))
     languageEl(call, which = "R.param") <- NULL
   if (!is.null(languageEl(call, which = "G.param")))
     languageEl(call, which = "G.param") <- NULL
 }

 #deal with args coming via ...
 #- this will overwite previously set values, except data and models are protectd
 tempcall <- list(...)
 if (length(tempcall))
 {
   for (z in names(tempcall))
     if (z == "data")
       env.dat <- data
     else
       if (z %in% c("fixed", "random", "residual", "sparse"))
         stop("attempt to change model to be fitted")
     else
       languageEl(call, which = z) <- tempcall[[z]]
 }

 #investigate residual term to see if it is a two-factor term or a
 #three-factor term with one factor indexing sections
 if (asr4)
   rterm <- languageEl(call, which = "residual")
 else
   rterm <- languageEl(call, which = "rcov")
 #Check for multiple terms
 rterm.form <- as.formula(rterm)
 res.specials <- c("dsum")
 common.specials <-  c("id", "diag", "us",
                       "ar1", "ar2", "ar3", "sar","sar2",
                       "ma1", "ma2", "arma", "exp", "gau",
                       "cor", "corb", "corg")
 other.specials <- c("at", "sph", "chol", "ante", "sfa", "facv")
 rterm.obj <- as.terms.object(rterm, asreml.obj,
                              specials=c(res.specials, common.specials,
                                         paste(common.specials[4:length(common.specials)],"v",sep=""),
                                         paste(common.specials[4:length(common.specials)],"h",sep=""),
                                         other.specials))
 if (length(labels(rterm.obj)) != 1)
 {
   if (asr4)
     stop("In analysing ",asreml.obj$formulae$fixed[[2]],
          ", the residual model must involve a single term")
   else
     stop("In analysing ",asreml.obj$fixed.formula[[2]],
          ", the residual model must involve a single term")
 }
 if (asr4)
   kspecial <- attr(rterm.obj, which = "specials")$dsum
 else
 {
   if (length(labels(rterm.obj)) != 1)
     stop("In analysing ",object.sim$fixed.formula[[2]],
          ", the rcov model must involve a single term")
   kspecial <- attr(rterm.obj, which = "specials")$at
 }

 if (is.null(kspecial))
 {
   grid.facs <- rownames(attr(rterm.obj, which = "factors"))
   grid.facs <- unlist(lapply(grid.facs, rmFunction))
   sections <- 1
   fac.sec <- NULL
 } else
 {
   if (asr4)
   {
     rterm <- labels(rterm.obj)
     rterm <- substr(rterm, start=6, stop=nchar(rterm))
     rterm <- substr(rterm, start=1, stop =(nchar(rterm)-1))
     trim <- function (x) gsub("^\\s+|\\s+$", "", x)
     if (grepl(",", rterm, fixed=TRUE))
     {
       rterm <- trim(strsplit(rterm, ",", fixed=TRUE)[[1]])
     }
     kmodel <- grep("~", rterm, fixed=TRUE)
     if (length(kmodel) != 1)
       stop("residual model appears to lack a formula")
     rterm <- rterm[kmodel]
     rterm <- strsplit(rterm, "~", fixed=TRUE)[[1]][2]
     var <- trim(strsplit(rterm, "|", fixed=TRUE)[[1]])
     var <- trim(var)
     if (length(var) != 2)
       stop("residual model appears not to include `|' operator")
     fac.sec <- var[2]
     sections <- levels(env.dat[[fac.sec]])
     rterm <- as.formula(paste("~",var[1],sep=""))
     rterm.obj <- as.terms.object(rterm, asreml.obj,
                                  specials=c("dsum", "at", "ar1", "ar2", "ar3", "sar","sar2",
                                             "ma1", "ma2", "arma", "exp", "gau",
                                             "cor", "corb", "corg", "diag", "us",
                                             "chol", "ante", "sfa", "facv", "fa", "rr"))
     if (length(labels(rterm.obj)) != 1)
       stop("In analysing ",asreml.obj$formulae$fixed[[2]],
            ", the residual model must involve a single term")
     grid.facs <- rownames(attr(rterm.obj, which = "factors"))
     grid.facs <- unlist(lapply(grid.facs, rmFunction))
   } else
   {
     grid.facs <- rownames(attr(rterm.obj, which = "factors"))
     grid.facs <- unlist(lapply(grid.facs, rmFunction))
     if (length(kspecial) != 1)
       stop("Can only have a single factor defining sections of the data (using at)")
     fac.sec <- grid.facs[kspecial[1]]
     sections <- levels(env.dat[[fac.sec]])
     grid.facs <- grid.facs[-kspecial[1]]
   }
 }
 if (length(grid.facs) != 2)
   stop("Must have two dimensions in the residual model")

 #Set up for simulation
 #There will be two lists, one for each face
 #The list for each face  has number of sections components
 #Each component will be an nsim x nfaci matrix, where i is the face
 fac1 <- grid.facs[1]
 fac2 <- grid.facs[2]
 if (any(table(env.dat[[fac1]])==0) | any(table(env.dat[[fac2]])==0))
   stop("Some levels in the the supplied factors are not observed")
 n1 <- length(levels(env.dat[[fac1]]))
 n2 <- length(levels(env.dat[[fac2]]))
 ns1 <- nsim*n1
 ns2 <- nsim*n2
 env.var <- vector(mode = "list", length = 2)
 names(env.var) <- grid.facs
 env.var[[1]] <- vector(mode = "list", length = length(sections))
 names(env.var[[1]]) <- sections
 env.var[[2]] <- vector(mode = "list", length = length(sections))
 names(env.var[[2]]) <- sections

 #generate nsim data sets and save variogram face data
  conv <- FALSE
  res.dat <- data.frame(env.dat[[fac1]],env.dat[[fac2]])
  names(res.dat)[1:2] <- grid.facs
  res.dat[grid.facs] <- lapply(res.dat[grid.facs], as.numfac)
  if (!is.null(fac.sec))
  {
    res.dat <- data.frame(env.dat[[fac.sec]], res.dat)
    names(res.dat)[1] <- fac.sec
  }
  #Set up expectation
  if (is.null(means))
  {
    if (asr4)
      mu <- fitted.values(asreml.obj)
    else
      mu <- fitted(asreml.obj) #asreml::fitted.asreml(asreml.obj)
  }
  else
    mu <- means

  #Get observed residuals
  res.dat$res <- get.residuals(asreml.obj, units = units)
  restype <- attr(res.dat$res, which = "restype")
  r <- res.dat$res

  cl <- makeCluster(ncores)
  registerDoParallel(cl)
  #process seed argument
  if (!is.null(seed))
  {
    RNGkind("L'Ecuyer-CMRG")
    set.seed(seed)
    ## start ncores workers
    s <- .Random.seed
    for (i in 1:ncores) {
      s <- nextRNGStream(s)
      # send s to worker i as .Random.seed
    }
  }

  if (asr4)
  {
    env.var <- foreach (i = 1:nsim, .packages = c("asreml","asremlPlus"))  %dopar%
                        {
                          asreml::asreml.options(fail = "soft")
                          while (!conv)
                          {
                            env.dat <- within(env.dat,
                                              {y.sim <- as.vector(mu + R %*% rnorm(n))})
                            sim.asreml <- eval(call)
                            conv <- sim.asreml$converge
                          }
                          conv=FALSE
                          #Get residuals
                          res.dat <- within(res.dat,
                                            {res <- get.residuals(sim.asreml, units = units)})
                          if (restype != attr(res.dat$res, which = "restype"))
                            warning("Observed and simulated residuals are not of the same type")

                          #Get variogram faces for this simulated data set
                          for (k in sections)
                          {
                            if (is.null(fac.sec))
                              sect.dat <- res.dat
                            else
                              sect.dat <- res.dat[res.dat[[fac.sec]]==k, -1]
                            #      sect.dat <- as.matrix(sect.dat)
                            sim.var <- asreml::asr_varioGram(sect.dat)
                            names(sim.var)[1:2] <- grid.facs
                            env.var[[1]][[k]] <- sim.var[sim.var[[fac2]]==0,]$gamma
                            env.var[[2]][[k]] <- sim.var[sim.var[[fac1]]==0,]$gamma
                          }
                          env.var
                        }
  } else
  {
    env.var <- foreach (i = 1:nsim, .packages = c("asreml","asremlPlus"))  %dopar%
                        {
                          loadASRemlVersion(3, lib.loc = "D:\\Analyses\\R oldpkg")
                          while (!conv)
                          {
                            env.dat <- within(env.dat,
                                              {y.sim <- as.vector(mu + R %*% rnorm(n))})
                            sim.asreml <- eval(call)
                            conv <- sim.asreml$converge
                          }
                          conv=FALSE
                          #Get residuals
                          res.dat$res <- get.residuals(sim.asreml, units = units)
                          if (restype != attr(res.dat$res, which = "restype"))
                            warning("Observed and simulated residuals are not of the same type")
                          #Get variogram faces for this simulated data set
                          for (k in sections)
                          {
                            if (is.null(fac.sec))
                              sect.dat <- res.dat
                            else
                              sect.dat <- res.dat[res.dat[[fac.sec]]==k, -1]
                            #      sect.dat <- as.matrix(sect.dat)
                            sim.var <- asreml::asreml.variogram(sect.dat)
                            names(sim.var)[1:2] <- grid.facs
                            env.var[[1]][[k]] <- sim.var[sim.var[[fac2]]==0,]$gamma
                            env.var[[2]][[k]] <- sim.var[sim.var[[fac1]]==0,]$gamma
                          }
                          env.var
                        }
  }
  stopCluster(cl)
  if (abs(mean(r -res.dat$res, na.rm = TRUE)) > 1e-6)
    stop("res.dat$res has changed")

  #Do face plots for fac1 and fac2 in each section
  face.1 <- data.frame(matrix(nrow = 0, ncol=5))
  colnames(face.1) <- c("factor","observed","X2.5.","X50.","X97.5.")
  face.2 <- data.frame(matrix(nrow = 0, ncol=5))
  colnames(face.2) <- c("factor","observed","X2.5.","X50.","X97.5.")

  #Get the observed variogram and form the elements of the variofaces plot
  for (k in sections)
  {
    if (is.null(fac.sec))
      sect.dat <- res.dat
    else
      sect.dat <- res.dat[res.dat[[fac.sec]]==k, -1]
    #    sect.dat <- as.matrix(sect.dat)
    if (asr4)
      object.var <- asreml::asr_varioGram(sect.dat)
    else
      object.var <- asreml::asreml.variogram(sect.dat)
    names(object.var)[1:2] <- grid.facs

    #Form variogram for current section
    sect.var <- vector(mode = "list", length = 2)
    names(sect.var) <- grid.facs
    sect.var[[1]] <- as.matrix(as.data.frame(lapply(env.var,
                                                    function(x, k) {x[[1]][[k]]},
                                                    k = k)))
    sect.var[[2]] <- as.matrix(as.data.frame(lapply(env.var,
                                                    function(x, k) {x[[2]][[k]]},
                                                    k = k)))

    #Form data.frame for current section
    face.1 <- rbind(face.1,
                    data.frame(factor = object.var[object.var[[fac2]]==0,][[fac1]],
                               observed = object.var[object.var[[fac2]]==0,]$gamma,
                               t(apply(sect.var[[1]], 1, quantile,
                                       probs=c(0.025, 0.50, 0.975),
                                       na.rm = TRUE)),
                               stringsAsFactors = FALSE))
    face.2 <- rbind(face.2,
                    data.frame(factor = object.var[object.var[[fac1]]==0,][[fac2]],
                               observed = object.var[object.var[[fac1]]==0,]$gamma,
                               t(apply(sect.var[[2]], 1, quantile,
                                       probs=c(0.025, 0.50, 0.975),
                                       na.rm = TRUE)),
                               stringsAsFactors = FALSE))
  }
  names(face.1)[1] <- fac1
  names(face.2)[1] <- fac2

  if (!is.null(fac.sec))
  {
    face.1[[fac.sec]] <- factor(rep(1:length(sections), each=(n1)), labels=sections)
    face.2[[fac.sec]] <- factor(rep(1:length(sections), each=(n2)), labels=sections)
  }

  #Do plots
  if (!is.null(graphics.device) )
    do.call(graphics.device, list(record = FALSE))

  p <- ggplot(data=face.1) +
    theme_bw() +
    geom_line(aes_string(x= fac1, y = "observed"), size=1) +
    geom_point(aes_string(x= fac1, y = "observed"), size=3) +
    geom_line(aes_string(x=fac1, y = "X2.5."), colour="red") +
    geom_line(aes_string(x=fac1, y = "X97.5."), colour="red") +
    geom_line(aes_string(x=fac1, y = "X50."), colour="blue") +
    labs(title = paste("Variogram face of",restype,"for",fac1,sep=" "),
         x = paste(fac1,"differences", sep=" "),y=NULL)
  if (!is.null(fac.sec))
    p <- p + facet_wrap(as.formula(paste("~",fac.sec)), ncol=2, scales="free")
  print(p)

  p <- ggplot(data=face.2) +
    theme_bw() +
    geom_line(aes_string(x= fac2, y = "observed"), size=1) +
    geom_point(aes_string(x= fac2, y = "observed"), size=3) +
    geom_line(aes_string(x=fac2, y = "X2.5."), colour="red") +
    geom_line(aes_string(x=fac2, y = "X97.5."), colour="red") +
    geom_line(aes_string(x=fac2, y = "X50."), colour="blue") +
    labs(title = paste("Variogram face of",restype,"for",fac2,sep=" "),
         x = paste(fac2,"differences", sep=" "), y=NULL)
  if (!is.null(fac.sec))
    p <- p + facet_wrap(as.formula(paste("~",fac.sec)), ncol=2, scales="free")
  print(p)

  #return data frames containing the variogram values on which the plots are based
  invisible(list(face1 = face.1, face2 = face.2))
}

"simulate.asreml" <- function(object, nsim=100, seed = NULL, means=NULL, V,
                              tolerance = 1E-10, update = TRUE, trace = FALSE,
                              which = "data", units = "ignore", ncores = detectCores(),
                              ...)
  #function to obtain simulated data corresponding to a fitted asreml model
  #object is an asreml object from a call to asreml in which the data argument
  #   must have been set.
  #means is a vector of predictions for fixed terms - should include spline terms
  #   if null, fitted values from fitted model are used.
  #V is the fitted variance matrix i.e. having the appropriate pattern and values
  #   given the model fitted and the estimates of the parameters obtained
  #nsim is the number of data sets to be simulated in obtaining the envelopes
  #   Note that only the results from data sets that converge in the fitting are
  #        included in computing the envelope
  #tolerance is the value such that eigenvalues less than it are consdered to be zero
#... parameters to supply to plot functions called within variofaces.asreml
{
  asr4 <- isASRemlVersionLoaded(4, notloaded.fault = TRUE)
  #Check that have a valid object of class asreml
  validasr <- validAsreml(object)
  if (is.character(validasr))
    stop(validasr)

  options <- c("data", "fitted", "residuals", "all")
  opt <- options[unlist(lapply(which, check.arg.values, options=options))]
  if ("all" %in% opt)
    opt <-  c("data", "fitted", "residuals")
  options <- c("addtoresiduals", "ignore")
  unit.opt <- options[check.arg.values(units, options)]
  n <- length(object$residuals)
  if (!is.null(means) & length(means) != n)
    stop("The lengths of means  and the response variable are not the same")
  #Check V
  if (!isSymmetric(V))
    stop("Variance matrix must be symmetric")
  if (!all(dim(V) == c(n, n)))
    stop("V is not a square matrix whose order equals the length of the response variable")
  #use eigenvalue decomposition to establish transformation matrix
  eigdecomp <- eigen(V, symmetric = TRUE)
  eigenval <- eigdecomp$values
  if (!all(eigenval >= -tolerance * abs(max(eigenval))))
    stop("Variance matrix is not nonnegative definite")
  R <- eigdecomp$vectors %*% diag(sqrt(pmax(eigenval, 0)), n)

  #check and get info in supplied call
  call <- object$call
  if (!("data" %in% names(call)))
    stop("simulate.asreml assumes that data has been set in call to asreml")
  env.dat <- eval(call$data)

  #set up call
  object.sim <- object
  call <- object.sim$call
  elno <- grep("data", names(call))
  languageEl(call, which=elno) <- quote(env.dat)
  languageEl(call$fixed, which=2) <- quote(y.sim)
  call$trace <- trace

  #Deal with the R.param and G.param arguments
  if (update)
  {
    #If update, set R.param and G.param
    languageEl(call, which = "R.param") <- object$R.param
    languageEl(call, which = "G.param") <- object$G.param
  } else
  {
    #If R.param and G.param already set, make them NULL
    if (!is.null(languageEl(call, which = "R.param")))
      languageEl(call, which = "R.param") <- NULL
    if (!is.null(languageEl(call, which = "G.param")))
      languageEl(call, which = "G.param") <- NULL
  }

  #deal with args coming via ...
  #- this will overwite previously set values, except data and models are protectd
  tempcall <- list(...)
  if (length(tempcall))
  {
    for (z in names(tempcall))
      if (z == "data")
        env.dat <- data
      else
      {
        if (z %in% c("fixed", "random", "residual", "rcov", "sparse"))
          stop("attempt to change model to be fitted")
        else
          languageEl(call, which = z) <- tempcall[[z]]
      }
  }

  #Set up for simulation and get observed residuals and fitted values
  #There will be a data.fame for each quantity saved, as well as the supplied
  #data.frame to which the requested quantites have been added
  if ("fitted"%in% opt)
  {
    if (asr4)
      env.dat$fitted <- fitted.values(object)
    else
      env.dat$fitted <- fitted(object)
#      env.dat$fitted <- asreml::fitted.asreml(object)
  }
  if ("residuals"%in% opt)
  {
    #Get observed residuals
    env.dat$residuals <- get.residuals(object, units = units)
    restype <- attr(env.dat$residuals, which = "restype")
  }

  #Set up expectation
  if (is.null(means))
  {
    if (asr4)
      mu <- fitted.values(object)
    else
      mu <- fitted(object)
#      mu <- asreml::fitted.asreml(object)
  }
  else
    mu <- means

  #Check type of residuals in simulation
  env.dat$y.sim <- as.vector(mu + R %*% rnorm(n))
  sim.asreml <- eval(call)
  if ("residuals" %in% opt)
    if (((length(sim.asreml$aom) == 0) & (restype == "Standardized conditional residuals")) |
          ((length(sim.asreml$aom) != 0) & (restype == "Residuals")))
      warning("Observed and simulated residuals are not of the same type\n",
              "- check setting of aom")

  #generate nsim data sets and save variogram face data
  conv <- FALSE
  cl <- makeCluster(ncores)
  registerDoParallel(cl)
  #process seed argument
  if (!is.null(seed))
  {
    RNGkind("L'Ecuyer-CMRG")
    set.seed(seed)
    ## start ncores workers
    s <- .Random.seed
    for (i in 1:ncores) {
      s <- nextRNGStream(s)
      # send s to worker i as .Random.seed
    }
  }
  if (asr4)
  {
    sim <- foreach (i = 1:nsim, .combine=rbind,
                    .packages = c("asreml","asremlPlus"))  %dopar%
                    {
                      asreml::asreml.options(fail = "soft")
                      while (!conv)
                      {
                        env.dat <- within(env.dat,
                                          {y.sim <- as.vector(mu + R %*% rnorm(n))})
                        sim.asreml <- eval(call)
                        conv <- sim.asreml$converge
                      }
                      conv=FALSE
                      sim <- vector(mode = "list", length = length(opt))
                      names(sim) <- opt
                      if ("data"%in% opt)
                        sim[["data"]] <- env.dat$y.sim
                      #Get residuals
                      if ("residuals"%in% opt)
                        sim[["residuals"]] <- get.residuals(sim.asreml, units=units)
                      if ("fitted"%in% opt)
                        sim[["fitted"]] <- fitted.values(sim.asreml)
                      sim
                    }
  } else
  {
    sim <- foreach (i = 1:nsim, .combine=rbind,
                    .packages = c("asreml","asremlPlus"))  %dopar%
                    {
                      loadASRemlVersion(3, lib.loc = "D:\\Analyses\\R oldpkg")
                      while (!conv)
                      {
                        env.dat <- within(env.dat,
                                          {y.sim <- as.vector(mu + R %*% rnorm(n))})
                        sim.asreml <- eval(call)
                        conv <- sim.asreml$converge
                      }
                      conv=FALSE
                      sim <- vector(mode = "list", length = length(opt))
                      names(sim) <- opt
                      if ("data"%in% opt)
                        sim[["data"]] <- env.dat$y.sim
                      #Get residuals
                      if ("residuals"%in% opt)
                        sim[["residuals"]] <- get.residuals(sim.asreml, units=units)
                      if ("fitted"%in% opt)
                        sim[["fitted"]] <- fitted(sim.asreml)
#                        sim[["fitted"]] <- asreml::fitted.asreml(sim.asreml)
                     sim
                    }
  }
  stopCluster(cl)

  out <- vector("list", length = 0)
  if ("residuals" %in% opt | "fitted" %in% opt)
  {
    env.dat <- env.dat[,-match("y.sim",names(env.dat))]
    out[["observed"]] <- env.dat
  }
  if ("data" %in% opt)
  {
    out[["data"]] <- as.data.frame(lapply(1:nsim,
                                          function(k, sim){sim[k, "data"]},
                                          sim = sim))
    names(out[["data"]]) <- paste("data",1:nsim,sep=".")
  }

  if ("fitted" %in% opt)
  {
    out[["fitted"]] <- as.data.frame(lapply(1:nsim,
                                            function(k, sim){sim[k, "fitted"]},
                                            sim = sim))
    names(out[["fitted"]]) <- paste("fitted",1:nsim,sep=".")
  }
  if ("residuals" %in% opt)
  {
    out[["residuals"]] <- as.data.frame(lapply(1:nsim,
                                               function(k, sim){sim[k, "residuals"]},
                                               sim = sim))
    names(out[["residuals"]]) <- paste("residuals",1:nsim,sep=".")
  }
  #return data frames containing the variogram values on which the plots are based
  invisible(out)
}

"plotVariofaces.data.frame" <- function(data, residuals, restype="Residuals", ...)
  #function to do the face variogram plots, including envelopes, described by
  #Stefanova et al (2010)
  #data is a data.frame with 3 or 4 columns. If there are 4, the first is taken to index
  #     sections of the data, the second and the third to index the x and y dimensions
  #     underlying the observations and the fourth the observed residuals. If there are 3,
  #     then the same as for 4 columns, except that the factor index the sections has
  #     been omitted.
  #residuals is a data frame with the same initial columns as data, the observed residuals
  #     having been ommitted. Following the initial columns are nsim sets of residuals
  #     derived from simulated data.
  #restype is a character string describing the type of residuals supplied.
{
  asr4 <- isASRemlVersionLoaded(4, notloaded.fault = TRUE)

  if (!is.data.frame(data) | !is.data.frame(residuals))
    stop("Both data and residuals should be data frames")
  n <- nrow(data)
  nvars <- ncol(data)
  if (nvars < 3 & nvars > 4)
    stop("The number of columns in data is not 3 or 4")
  nsim <- ncol(residuals) - nvars + 1
  if (nsim < 1)
    stop("Could not find columns with simulated residuals")
  #Check compatibility
  if (n != nrow(residuals))
    stop("The number of rows in data and residuals are not equal")

  #Check that initial columns in data and residuals match
  if (!setequal(names(data)[1:(nvars-1)], names(residuals)[1:(nvars-1)]))
    stop("The names of the initial ",nvars-1," columns do not match")
  if (!all(unlist(lapply(data[1:(nvars-1)], is.factor))))
    stop("Some of the initial ",nvars-1," columns of data are not factors")
  if (!all(unlist(lapply(residuals[1:(nvars-1)], is.factor))))
    stop("Some of the initial ",nvars-1," columns of residuals are not factors")
  for (i in 1:(nvars-1))
    if (!all(data[i] == residuals[i]))
      stop("The values in column ",i," of data and residuals do not match")

  #Identify columns in data
  if (nvars == 3)
  {
    fac.sec <- NULL
    sections <- 1
    grid.facs <- colnames(data)[1:2]
    resid <- colnames(data)[3]
  }
  else
  {
    fac.sec <- colnames(data)[1]
    sections <- levels(data[[fac.sec]])
    grid.facs <- colnames(data)[2:3]
    resid <- colnames(data)[4]
  }

  #Form variogram faces for each set of residuals
  #There will be two lists, one for each face
  #The list for each face  has number of sections components
  #Each component will be an nsim x nfaci matrix, where i is the face
  fac1 <- grid.facs[1]
  fac2 <- grid.facs[2]
  if (any(table(data[[fac1]])==0) | any(table(data[[fac2]])==0))
    stop("Some levels in the the supplied factors are not observed")
  n1 <- length(levels(data[[fac1]]))
  n2 <- length(levels(data[[fac2]]))
  ns1 <- nsim*n1
  ns2 <- nsim*n2
  variog <- vector(mode = "list", length = 2)
  variog[[1]] <- vector(mode = "list", length = length(sections))
  names(variog[[1]]) <- sections
  variog[[1]] <- lapply(variog[[1]][sections],
                        function(variog, nrow, ncol)
                          variog <- matrix(rep(0, nrow*ncol), nrow=nrow, ncol=ncol),
                        nrow=nsim, ncol=n1)
  variog[[2]] <- vector(mode = "list", length = length(sections))
  names(variog[[2]]) <- sections
  variog[[2]] <- lapply(variog[[2]][sections],
                        function(variog, nrow, ncol)
                          variog <- matrix(rep(0, nrow*ncol), nrow=nrow, ncol=ncol),
                        nrow=nsim, ncol=n2)
  data[grid.facs] <- lapply(data[grid.facs], as.numfac)
  residuals[grid.facs] <- lapply(residuals[grid.facs], as.numfac)

  for (k in sections)
  {
    sect.dat <- data.frame(data[fac1], data[fac2])
    if (!is.null(fac.sec))
      sect.dat <- sect.dat[data[[fac.sec]] == k]
    names(sect.dat)[1:2] <- grid.facs
    for (i in 1:nsim)
    {
      if (is.null(fac.sec))
        sect.dat$res <- as.vector(residuals[[(nvars+i-1)]])
      else
        sect.dat$res <- residuals[data[[fac.sec]] == k, (nvars+i-1)]
      if (asr4)
        sim.var <- asreml::asr_varioGram(sect.dat)
      else
        sim.var <- asreml::asreml.variogram(sect.dat)

      names(sim.var)[1:2] <- grid.facs
      variog[[1]][[k]][i,] <- sim.var[sim.var[[fac2]]==0,]$gamma
      variog[[2]][[k]][i,] <- sim.var[sim.var[[fac1]]==0,]$gamma
    }
  }

  #Do face plots for fac1 and fac2 in each section
  face.1 <- data.frame(matrix(nrow = 0, ncol=5))
  colnames(face.1) <- c("factor","observed","X2.5.","X50.","X97.5.")
  face.2 <- data.frame(matrix(nrow = 0, ncol=5))
  colnames(face.2) <- c("factor","observed","X2.5.","X50.","X97.5.")

  #Get the observed variogram and form the elements of the variofaces plot
  for (k in sections)
  {
    if (is.null(fac.sec))
      sect.dat <- data
    else
      sect.dat <- data[data[[fac.sec]]==k, -1]
    if (asr4)
      object.var <- asreml::asr_varioGram(sect.dat)
    else
      object.var <- asreml::asreml.variogram(sect.dat)
    names(object.var)[1:2] <- grid.facs

    #Form data.frame for current section
    face.1 <- rbind(face.1,
                    data.frame(factor = object.var[object.var[[fac2]]==0,][[fac1]],
                               observed = object.var[object.var[[fac2]]==0,]$gamma,
                               t(apply(variog[[1]][[k]], 2, quantile,
                                       probs=c(0.025, 0.50, 0.975),
                                       na.rm = TRUE)),
                               stringsAsFactors = FALSE))
    face.2 <- rbind(face.2,
                    data.frame(factor = object.var[object.var[[fac1]]==0,][[fac2]],
                               observed = object.var[object.var[[fac1]]==0,]$gamma,
                               t(apply(variog[[2]][[k]], 2, quantile,
                                       probs=c(0.025, 0.50, 0.975),
                                       na.rm = TRUE)),
                               stringsAsFactors = FALSE))
  }
  names(face.1)[1] <- fac1
  names(face.2)[1] <- fac2

  if (!is.null(fac.sec))
  {
    face.1[[fac.sec]] <- factor(rep(1:length(sections), each=(n1)), labels=sections)
    face.2[[fac.sec]] <- factor(rep(1:length(sections), each=(n2)), labels=sections)
  }

  #Do plots
  p <- ggplot(data=face.1) +
    theme_bw() +
    geom_line(aes_string(x= fac1, y = "observed"), size=1) +
    geom_point(aes_string(x= fac1, y = "observed"), size=3) +
    geom_line(aes_string(x=fac1, y = "X2.5."), colour="red") +
    geom_line(aes_string(x=fac1, y = "X97.5."), colour="red") +
    geom_line(aes_string(x=fac1, y = "X50."), colour="blue") +
    labs(title = paste("Variogram face of",restype,"for",fac1,sep=" "),
         x = paste(fac1,"differences", sep=" "),y=NULL)
  if (!is.null(fac.sec))
    p <- p + facet_wrap(as.formula(paste("~",fac.sec)), ncol=2, scales="free")
  print(p)

  p <- ggplot(data=face.2) +
    theme_bw() +
    geom_line(aes_string(x= fac2, y = "observed"), size=1) +
    geom_point(aes_string(x= fac2, y = "observed"), size=3) +
    geom_line(aes_string(x=fac2, y = "X2.5."), colour="red") +
    geom_line(aes_string(x=fac2, y = "X97.5."), colour="red") +
    geom_line(aes_string(x=fac2, y = "X50."), colour="blue") +
    labs(title = paste("Variogram face of",restype,"for",fac2,sep=" "),
         x = paste(fac2,"differences", sep=" "), y=NULL)
  if (!is.null(fac.sec))
    p <- p + facet_wrap(as.formula(paste("~",fac.sec)), ncol=2, scales="free")
  print(p)

  #return data frames containing the variogram values on which the plots are based
  invisible(list(face1 = face.1, face2 = face.2))
}


#################################################################################
#
# Mark Cembrowski, Janelia Research Campus, Apr 7 2015
#
# This looks for the number of genes enriched in the various CA region sof the
# hippocampus.
#
#################################################################################

caRegionEnrichVsFold <- function(fpkmThres=10,countThres=20,doFpkm=T){
	if(doFpkm){
		theThres <- fpkmThres
	}else{
		theThres <- countThres
	}

	theFolds <- c(3:10)
	df <- as.data.frame(matrix(nrow=theFolds,ncol=3))

	theMask <- setdiff(colnames(fpkmPoolMat),c('ca3_d','ca2','ca1_d'))

	for (ii in 1:length(theFolds)){
		numCa3Enrich <- length(enrichedGenes('ca3_d',fpkmThres=theThres,
			foldThres=theFolds[ii],avgPass=T,mask=theMask,doFpkm=doFpkm))
		numCa2Enrich <- length(enrichedGenes('ca2',fpkmThres=theThres,
			foldThres=theFolds[ii],avgPass=T,mask=theMask,doFpkm=doFpkm))
		numCa1Enrich <- length(enrichedGenes('ca1_d',fpkmThres=theThres,
			foldThres=theFolds[ii],avgPass=T,mask=theMask,doFpkm=doFpkm))
	if(theFolds[ii]>9.5){
print(idToSym(enrichedGenes('ca1_d',fpkmThres=theThres,
                        foldThres=theFolds[ii],avgPass=T,mask=theMask,doFpkm=doFpkm)))
}
		df[ii,] <- c(numCa3Enrich,numCa2Enrich,numCa1Enrich)
	}

	df <- cbind(theFolds,df)
	colnames(df) <- c('theFolds','ca3','ca2','ca1')

	# Plot.
	gg <- ggplot(df,aes(x=theFolds))
	gg <- gg + geom_line(aes(y=ca2),colour='black')
	gg <- gg + geom_line(aes(y=ca3),colour='red')
	gg <- gg + geom_line(aes(y=ca1),colour='blue')
	gg <- gg + theme_bw()
	gg <- gg + expand_limits(y=0)
	gg <- gg + xlab('Fold difference') + ylab('Number of genes')
	gg <- gg + ggtitle('Blue:CA1,Black:CA2,red:CA3')
	print(gg)

	invisible(df)
}


## calculate variance explained in total by all three cell fractions

source("rmdConfig.mrc")
setwd(dataDir)
load(normData)

## remove rs probes
celltypenormbeta<-celltypenormbeta[grep("rs", rownames(celltypenormbeta), invert = TRUE),]

## exclude IRF8+ve as not present in sufficient sample numbers
pheno<-pheno[pheno$Cell.type %in% c("Double -ve","NeuN +ve","Sox10 +ve", "Total"),]
celltypenormbeta<-celltypenormbeta[,pheno$Basename]
pheno$Cell.type<-factor(pheno$Cell.type)

## reformat as a list with matrix for each cell type, with samples in the same order
allSamples<-unique(pheno$Indidivual.ID)

betas.list<-list("Total" = celltypenormbeta[,which(pheno$Cell.type == "Total")[match(allSamples, pheno$Indidivual.ID[which(pheno$Cell.type == "Total")])]],
"NeuN" = celltypenormbeta[,which(pheno$Cell.type == "NeuN +ve")[match(allSamples, pheno$Indidivual.ID[which(pheno$Cell.type == "Total")])]],
"Sox10" = celltypenormbeta[,which(pheno$Cell.type == "Sox10 +ve")[match(allSamples, pheno$Indidivual.ID[which(pheno$Cell.type == "Total")])]],
"DNeg" = celltypenormbeta[,which(pheno$Cell.type == "Double -ve")[match(allSamples, pheno$Indidivual.ID[which(pheno$Cell.type == "Total")])]])

## take total as outcome then all other cell types as variables

out<-matrix(data = NA, ncol = 2, nrow = nrow(betas.list[["Total"]]))
rownames(out)<-rownames(betas.list[["Total"]])
colnames(out)<-c("MultipleR2", "AdjustedR2")

for(i in 1:nrow(out)){
	model<-lm(betas.list[["Total"]][i,] ~ betas.list[["NeuN"]][i,] + betas.list[["Sox10"]][i,] + betas.list[["DNeg"]][i,])
	out[i,2]<-summary(model)$adj.r.squared
	out[i,1]<-summary(model)$r.squared
}

save(out, file="Analysis/CellType/TotalVarExplainAllCT.rdata")


#install.packages("IRkernel")
#install.packages('dplyr')
#library(dplyr)
#install.packages("htmlwidgets")
#install.packages("ggplot2")
#install.packages('pheatmap')
#install.packages("devtools")
library(devtools)
#install_github("hrbrmstr/waffle")
library(ggplot2)
#library(waffle)
library(htmlwidgets)
library(pheatmap)
library(grid)
library(gridExtra)

heat_map <- function(){
  clean_df <- read.csv('out.summary.csv',header = TRUE, stringsAsFactors = FALSE)
  clean_df[clean_df == 'yes'] <- 1
  clean_df[clean_df == 'no'] <- 0
  rownames(clean_df) <- clean_df$name
  plot_df <- data.matrix(clean_df[,2:57])
  hmcols<- colorRampPalette(c("red","blue"))(2)

  p =pheatmap(plot_df, cluster_row=T,cluster_cols =F,color=hmcols,
     fontsize=10, fontsize_col=5,scale="none",legend_breaks=c(0,0.25,0.5,0.75,1),
        legend_labels=c('','Not match with cluster','','Match with cluster',''),
         show_colnames = T,show_rownames=F)
  p
}

scatter_plot <- function(){
  d=data.frame(Sensitivity=c(1,0.6803,0.494,0.975,0.947,0.724,0.988,0.95),
   Specificity=c(0.831,0.508,0.975,0.248,0.924,0.859,0.953,0.61),
   Method=c('Ariba','Ariba','Ariba','Ariba','Mykrobe','Mykrobe','Mykrobe','Mykrobe'),
   Drug=c('Rifampicin','Ethambutol','Isoniazid','Pyrazinamide','Rifampicin','Ethambutol','Isoniazid','Pyrazinamide'),
   info=c('Rifampicin_Ariba','Ethambutol_Ariba','Isoniazid_Ariba','Pyrazinamide_Ariba','Rifampicin_Mykrobe','Ethambutol_Mykrobe','Isoniazid_Mykrobe','Pyrazinamide_Mykrobe'))
  ggplot()  + xlim(0, 1) + ylim(0, 1) +
   geom_point(data=d, mapping=aes(x=Specificity, y=Sensitivity, color=Drug, shape=Method),size=4) +
   ggtitle("Prediction Tools Comparison") + scale_color_brewer(palette="Set1")+
    theme(plot.title = element_text(hjust = 0.5))+
   geom_text(aes(label = d$info),size = 2,x=d$Specificity-0.088, y=d$Sensitivity)+ theme_minimal() +
   theme(legend.position = c(0.2, 0.4))
}

grid_plot <- function(){
  #A-RIFA
  p1 <- waffle(c(Match = 600, Not_Match = 56,Not_Available=0), rows = 16, keep = FALSE,size = .2,colors = c("blue","red",'grey'),
  legend_pos = c(16,16))
  #A-ETH
  p2 <- waffle(c(Match = 413, Not_Match = 243,Not_Available=0), rows = 16, keep = FALSE,size = .2,colors = c("blue","red",'grey'),
  legend_pos = c(16,16))
#A-ISO
  p3 <- waffle(c(Match = 519, Not_Match = 137,Not_Available=0), rows = 16, keep = FALSE,size = .2,colors = c("blue","red",'grey'),
  legend_pos = c(16,16))
#A-PYR
  p4 <- waffle(c(Match = 539, Not_Match = 113, Not_Available=4), rows = 16, keep = FALSE,size = .2,colors = c("blue","red",'grey'),
  legend_pos = c(16,16))
#M-RIFA
  p5 <- waffle(c(Match = 613, Not_Match = 43,Not_Available=0), rows = 16, keep = FALSE,size = .2,colors = c("blue","red",'grey'),
  legend_pos = c(16,16))
#M-ETH
  p6 <- waffle(c(Match = 502, Not_Match = 154,Not_Available=0), rows = 16, keep = FALSE,size = .2,colors = c("blue","red",'grey'),
  legend_pos = c(16,16))
#M-ISO
  p7 <- waffle(c(Match = 634, Not_Match = 22,Not_Available=0), rows = 16, keep = FALSE,size = .2,colors = c("blue","red",'grey'),
  legend_pos = c(16,16))
#M-PYR
  p8 <- waffle(c(Match = 578, Not_Match = 74,Not_Available=4), rows = 16, keep = FALSE,size = .2,colors = c("blue","red",'grey'),
  legend_pos = c(16,16))

  col1 <- arrangeGrob(p2, p3, p4,p1, top="Ariba",ncol=1)
  col2 <- arrangeGrob(p6,p7,p8,p5, top='Mykrobe',ncol=1)
  col3 <- textGrob(x = 0.5, y = 0.48, paste("Ethambutol\n\n\n\n\n", "Isoniazid\n\n\n\n",
        "\nPyrazinamide\n\n\n\n\n", "Rifampicin"),gp=gpar(fontsize=9,font=8))
  legd <- legendGrob(c("Match with phenotype", "Not match with phenotype", "Not Available"), pch=22,
   nrow=1, byrow=TRUE, gp=gpar(col = c("blue","red",'grey'), fill = c("blue","red",'grey')))
  grid.arrange(col1,col3,col2, legd, widths=c(1,0.2,1), heights=c(4,0.5),top = "Accuracy Comparison",
             layout_matrix=rbind(c(1,2,3), c(4,4,4)))
}

sorted_grid <- function() {
    t=read.csv('match_plot_sort_mykrobe.txt',sep='\t')
 colors <- colorRampPalette(c("blue", "grey", "red"))(3)
 ggplot(t[1:5248,], aes(column, row, fill = match_not)) +
  geom_tile(colour = "white") +
  facet_grid(drug~method)  +theme_minimal()+
  labs(x="", y="") + ggtitle("Side by Side Comparison")+
  theme(plot.title = element_text(hjust = 0.5))+
  theme(legend.position="bottom") +
  guides(fill=guide_legend(title=""))+
  scale_fill_manual(values=colors)+  theme(strip.text.x = element_text(size = c(10)))+
  theme(strip.text.y = element_text(size = c(7)))
}


setwd("prslev/result/")

source('d://tools/qqunif.r')

################# liability threshold

a = read.table('probit_a/h2p_0.1_nsnp_all_snps_vr_0.01_freqr_0.001_k_0.01_Nsample_10000_pi0_0',header=T)
b = read.table('probit_b/h2p_0.1_nsnp_all_snps_vr_0.01_freqr_0.001_k_0.01_Nsample_10000_pi0_0',header=T)
c = read.table('probit_c/h2p_0.1_nsnp_all_snps_vr_0.01_freqr_0.001_k_0.01_Nsample_10000_pi0_0',header=T)

qqunif(a$case_inverse_wilcox_p, BH=F, CI=F, BF=F, xlim=c(0,3), ylim=c(0,7), cex.axis=2.0, cex.lab=2.0, cex=1.8)
par(new=T)
qqunif(b$case_inverse_wilcox_p, BH=F, CI=F, BF=F, xlim=c(0,3), ylim=c(0,7), col="red", cex.axis=2.0, cex.lab=2.0, cex=1.8)
par(new=T)
qqunif(c$case_inverse_wilcox_p, BH=F, CI=F, BF=F, xlim=c(0,3), ylim=c(0,7), col="blue", cex.axis=2.0, cex.lab=2.0, cex=1.8)

legend('topleft', c("Negative selection","LDAK","Polygenic"),
             col=c('red','blue','black'),lty=c(1,1,1), cex=1.8)

################# logit

a = read.table('probit_a/h2p_0.2_nsnp_all_snps_vr_0.01_freqr_0.001_k_0.01_Nsample_10000_pi0_0',header=T)
b = read.table('probit_b/h2p_0.2_nsnp_all_snps_vr_0.01_freqr_0.001_k_0.01_Nsample_10000_pi0_0',header=T)
c = read.table('probit_c/h2p_0.2_nsnp_all_snps_vr_0.01_freqr_0.001_k_0.01_Nsample_10000_pi0_0',header=T)

qqunif(a$case_inverse_wilcox_p, BH=F, CI=F, BF=F, xlim=c(0,3), ylim=c(0,7), cex.axis=2.0, cex.lab=2.0)
par(new=T)
qqunif(b$case_inverse_wilcox_p, BH=F, CI=F, BF=F, xlim=c(0,3), ylim=c(0,7), col="red", cex.axis=2.0, cex.lab=2.0)
par(new=T)
qqunif(c$case_inverse_wilcox_p, BH=F, CI=F, BF=F, xlim=c(0,3), ylim=c(0,7), col="blue", cex.axis=2.0, cex.lab=2.0)

legend('topleft', c("Negative selection","LDAK","Polygenic"),
             col=c('red','blue','black'),lty=c(1,1,1), cex=1.8)

################# logit vs probit, negative selection

m = read.table('probit_b/h2p_0.1_nsnp_all_snps_vr_0.01_freqr_0.001_k_0.01_Nsample_10000_pi0_0',header=T)
n = read.table('logit_b/h2p_0.1_nsnp_all_snps_vr_0.01_freqr_0.001_k_0.01_Nsample_10000_pi0_0',header=T)

qqunif(m$case_inverse_wilcox_p, BH=F, CI=F, BF=F, xlim=c(0,3), ylim=c(0,7), col="red", cex.axis=2.0, cex.lab=2.0)
par(new=T)
qqunif(n$case_inverse_wilcox_p, BH=F, CI=F, BF=F, xlim=c(0,3), ylim=c(0,7), col="blue", cex.axis=2.0, cex.lab=2.0)

legend('topleft', c("Liability threshold","Logit"),
             col=c('red','blue'),lty=c(1,1,1), cex=1.8)


#Area of interest
wdpaid <- '-8.02_51.76_-2.63_55.79'
#wdpaid <- 'minlon_minlat_maxlon_maxlat'
wdpaidsplit <- unlist(strsplit(wdpaid, "[_]"))
xmin <- as.numeric(wdpaidsplit[1])
ymin <- as.numeric(wdpaidsplit[2])
xmax <- as.numeric(wdpaidsplit[3])
ymax <- as.numeric(wdpaidsplit[4])

#Library
library(dplyr)
library(ggplot2)
library(rgdal)
library(rgeos)
library(RGeostats)
library(viridis)
#Libraries for mask
library(rworldmap)
library(rworldxtra)

# Etendre le tableau et rajouter lignes de densit? = 0
ttsel <- Sp %>% ungroup() %>% dplyr::select(group, Year, meanLon, meanLat, Abun) %>% distinct()
codif <- ttsel %>% dplyr::select( Year, meanLon, meanLat) %>% distinct()
allstrat <- ttsel %>% tidyr::expand(group,Year)
pipo <- dplyr::left_join(allstrat, codif) %>% dplyr::left_join(ttsel) %>% dplyr::mutate(Abun = ifelse(is.finite(Abun), Abun, 0))
Sp0 <- dplyr::left_join(codif, pipo) %>% na.omit()

#Creation du polygon de krigeage

# 1st Polygon
liste <- with(Sp0, chull(meanLon, meanLat))
hull <- Sp0[liste, c("meanLon", "meanLat")]
Poly <- Polygon(hull)

# Create SpatialPolygons objects
SpPoly<- SpatialPolygons(list(Polygons(list(Poly), "SpPoly")))
buff <- raster::buffer(SpPoly, 0.1)

# Cut object along coast
coast <- rworldmap::getMap(resolution = "high")
coast <- raster::crop(coast, raster::extent(xmin-0.1,xmax+0.1,ymin-0.1,ymax+0.1))

res <- gDifference(buff, coast)
PolyCut <- fortify(res)

# Put polygon in good format for later use
toto <- PolyCut[PolyCut$piece==1,]
db.poly <- polygon.create(toto[,c(1,2)])

data.db <- db.create(Sp0)
data.db <- db.locate(data.db, c("meanLon","meanLat"),"x")
data.db <- db.locate(data.db, "Abun","z")

# Boucle vario/krigeage sur tout les clusters (avec logarithme)

Longitude <- numeric()
Latitude <- numeric()
Prediction.std <- numeric()
Variance.std <- numeric()
Year <- numeric()
Cluster <- character()
Prediction <- numeric()
Variance <- numeric()

# Grille de krigeage
grid.db <- db.grid.init(data.db,nodes=c(100,100))
poly <- polygon.create(x=PolyCut[PolyCut$group=="1.1",1],y=PolyCut[PolyCut$group=="1.1",2]) #RGeostats polygon
grid.db <- db.polygon(grid.db, poly)

# Moving or unique neighbouhood?
nei1 <- neigh.create(ndim=2,type=0) #unique
nei2 <- neigh.create(ndim=2,type=2,nmini=2,nmaxi=8,radius=1) # moving

for (j in unique(pull(Sp0[,"group"]))){

  data.db <- db.create(Sp0[Sp0$group==j,]) #Choix du cluster
  data.db <- db.locate(data.db,names=c("meanLon","meanLat"),"x")
  data.db <- db.locate(data.db,names=c("Abun"),"z")

  # Standartisation par an
  data.db.std <- data.db
  data.db.std <- db.locate(data.db.std,"Year","code")

  # Vario moyen
  vg.data.std <- vario.calc(data.db.std, lag=0.05, nlag=floor(max(dist(data.db@items[,c(3,4)]))/2)/0.05,opt.code=1,tolcode=0)
  #plot(vg.data.std,npairdw=T,inches=0.1,las=1,add=T,col=2,lwd=2)
  vario1 <- vg.data.std
  vg.mod <- model.auto(vario=vario1, struct=c(1:5), npairdw=TRUE, title= paste0("Communauté", " ", j), size=45,  col="red", inches=.05)

  data.db.std <- db.locate(data.db.std,"Year",NA)

  for (i in sort(unique(data.db[,"Year"]))){

    sel <- data.db.std[,"Year"]== i

    if (all(sel==F))
    {Longitude <- c(Longitude, kres@items$x1[kres@items$Polygon==TRUE])
    Latitude <- c(Latitude, kres@items$x2[kres@items$Polygon==TRUE])
    Prediction <- c(Prediction,rep(NA, length(kres@items$x1[kres@items$Polygon==TRUE])))
    Variance <- c(Variance, rep(NA, length(kres@items$x1[kres@items$Polygon==TRUE])))
    Year <- c(Year, rep(i, length(kres@items$x1[kres@items$Polygon==TRUE])))
    Cluster <- c(Cluster, rep(j, length(kres@items$x1[kres@items$Polygon==TRUE])))}else{

      # Ordinary kriging, unique neighbourhood
      sel <- data.db.std[,"Year"]== i
      db.kriege <- db.create(data.db.std[sel])
      db.kriege <- db.locate(db.kriege,names=c("meanLon","meanLat"),"x")
      db.kriege <- db.locate(db.kriege,names=c("Abun"),"z")
      kres <- kriging(db.kriege, grid.db, model = vg.mod, neigh = nei1, uc=c("1"), mean=NA)

      Longitude <- c(Longitude, kres@items$x1[kres@items$Polygon==TRUE])
      Latitude <- c(Latitude, kres@items$x2[kres@items$Polygon==TRUE])
      Prediction <- c(Prediction, kres@items$Kriging.Abun.estim[kres@items$Polygon==TRUE])
      Variance <- c(Variance, kres@items$Kriging.Abun.stdev[kres@items$Polygon==TRUE])
      Year <- c(Year, rep(i, length(kres@items$x1[kres@items$Polygon==TRUE])))
      Cluster <- c(Cluster, rep(j, length(kres@items$x1[kres@items$Polygon==TRUE])))
    }

  }

}
Kriege.logdens<- data.frame(Longitude=Longitude, Latitude=Latitude, Prediction=Prediction, Variance=Variance, Year=Year, Cluster=factor(Cluster))
#save(Kriege.logdens, file="data/krigeage log.RData")

names(Kriege.logdens)[6]<- "Community"
Kriege.logdens$Community<- as.numeric(Kriege.logdens$Community)

Longitude<- numeric()
Latitude<- numeric()
Clust<- numeric()
Community<- numeric()

for (j in unique(data.frame(Kriege.logdens)[,"Community"])){
  Tab1<- Kriege.logdens[Kriege.logdens$Community==j,] %>% dplyr::select(-Variance, -Community)
  Tab2<- pivot_wider(Tab1, names_from = Year, values_from = Prediction)
  metaTab<- Tab2 %>% dplyr::select(Longitude, Latitude)
  Tab2<- Tab2 %>% dplyr::select(-c(Longitude, Latitude))

  #Classification

  distance<- dist(Tab2)
  tree<- fastcluster::hclust(distance, method="ward.D2")

  # Dendrogrammes
  zones<- cutree(tree, 5)

  toto<- cbind(metaTab, Clust=factor(zones))
  toto<- left_join(toto, Kriege.logdens[Kriege.logdens$Community==j,], by=c("Longitude", "Latitude"))
  toto<- toto %>% dplyr::select(Longitude, Latitude, Clust, Community)

  Longitude<- c(Longitude, toto$Longitude)
  Latitude<- c(Latitude, toto$Latitude)
  Clust<- c(Clust, toto$Clust)
  Community<- c(Community, toto$Community)

  tata <- left_join(toto, cbind(metaTab, Tab2))
  tata <- pivot_longer(tata, cols=c(5:dim(tata)[2]), names_to="Year", values_to = "Prediction")
  tata <- tata %>% dplyr::group_by(Year, Clust) %>% dplyr::summarise(Prediction=mean(Prediction))

  ggtata<-  ggplot(tata)+
    geom_point(aes(x=Year, y=Prediction, col=Clust))+
    geom_line(aes(x=Year, y=Prediction, col=Clust, group=Clust))+
    theme_minimal()+
    facet_wrap(.~Clust)

  #print(ggtata)

  tete<- tata
  names(tete)[2]<- "Zones"
  tete<- tata %>% ungroup() %>% dplyr::group_by(Clust) %>% dplyr::summarise(Prediction= mean(Prediction))

  titi<- toto %>% left_join(tete, by="Clust")

}
Tabfaunefin<- data.frame(Longitude=Longitude, Latitude=Latitude, Clust=as.factor(Clust), Community=Community)

for(i in 1:(length(keep)+2)){
  Tabfaunefin$Community[Tabfaunefin$Community==i] <- keep[i]
}

ggplot(Tabfaunefin)+
  geom_raster(aes(x=Longitude,y=Latitude,fill=Clust))+facet_wrap(~Community)+scale_fill_viridis_d()+theme_minimal()+
  geom_polygon(data=PolyCut[PolyCut$group==1.1,],aes(x=long, y=lat, group=group), col="black", fill=NA)+
  ggtitle("Classificiation of biological groups in selected area")