library(igraph)
library(sna)
library(mgcv)
library(assortnet)

### functions
## SET UP FUNCTIONS
# x <- difference
sig <- function(x,sociality,assortment) { 
	if (assortment < 0) {
		return ( 1/(1 + exp(-20*(abs(x)-(1-sociality)))) )
	} else {
		return ( 1/(1 + exp(-20*(abs(x)-sociality))) )
	}
}

si <- function(x, assortment, sociality) {
	return(-assortment*(sig(x,sociality,assortment)-0.5)+0.5)
}

# N individuals, m moving distance (% distance), r range of attraction, thresh threshold edge value
generate_network <- function(N, m, r, thresh, binary=FALSE) {
	x <- runif(N)
	y <- runif(N)

	distances <- sqrt(outer(x,x,FUN="-")^2 + outer(y,y,FUN="-")^2)
	diag(distances) <- NA

	x <- x-m*(x-x[apply(distances,1,which.min)])
	y <- y-m*(y-y[apply(distances,1,which.min)])

	distances <- sqrt(outer(x,x,FUN="-")^2 + outer(y,y,FUN="-")^2)

	space.network <- exp(-(distances^2) / r)
	diag(space.network) <- 0
	space.network[space.network < thresh] <- 0
	
	if (binary)
		space.network[space.network > 0] <- 1
	
	return(space.network)
}


# network properties
  m <- 0.15  # moving distance (% distance) has little influence on clustering
  r <- 0.025  # range of attraction has lots of influence on clustering
  thresh <- 0.05

#### run simulations for homophily
  psi <- 4
  assortment.values <- seq(0, 1, 0.02) 
  r <- (0.3^2)
  default.sociality <- 0.3  # 0.3 for homophily works well, 0.2 for heterophily works well
  include.assortment <- TRUE  # whether to weight the social network by preferential affiliation
  N <- 50 
  n.replicates <- 50
  beta = 0.2
  

# to store results results
results.inds <- list()
results.ids <- list()
results.nets <- list()
network.density <- list()
network.assortment <- list()
network.clustering <- list()
network.modularity <- list()

for (l in 1:length(assortment.values)) {
	assortment <- assortment.values[l]
	## Initiate lists
	results.inds[[length(results.inds)+1]] <- list()
	results.ids[[length(results.ids)+1]] <- list()
	results.nets[[length(results.nets)+1]] <- list()
	network.density[[length(network.density)+1]] <- vector()
	network.assortment[[length(network.assortment)+1]] <- vector() 
	network.clustering[[length(network.clustering)+1]] <- vector()
	network.modularity[[length(network.modularity)+1]] <- vector()
	
	for (i in 1:n.replicates) {

		## SET UP POPULATION
	
		inds <- data.frame(id=1:N)
		inds$Ai <- runif(n=N,min=-1,max=1) #individual direct breeding values
		inds$Ei <- rnorm(n=N,mean=0,sd=0.25) #individual environmental effects
		inds$Soc <- default.sociality
		inds$APi <- inds$Ai + inds$Ei # phenotype before interactions.
	
		## BUILD MATRIX
		ids <- data.frame(expand.grid(ID1=inds$id,ID2=inds$id))
		ids <- ids[-which(ids[,1]==ids[,2]),]
		ids$APj <- inds$APi[ids$ID2]
		ids$APi <- inds$APi[ids$ID1]
		ids$Aj <- inds$Ai[ids$ID2]
		ids$Ai <- inds$Ai[ids$ID1]
	
		# generate the baseline network
		base.network <- generate_network(N, m, r, thresh)
		weights <- si(outer((inds$APi),(inds$APi),"-"),assortment,outer(inds$Soc,inds$Soc,"+"))
		if (include.assortment) {
			updated.network <- base.network * weights
		} else {
			updated.network <- base.network
		}
		ids$edge.w <- updated.network[cbind(ids$ID1,ids$ID2)]

		ids$PjSi <- ids$APj * ids$edge.w
		ids$AjSi <- ids$Aj * ids$edge.w
		
		# Update parameters of interest
		inds$PjSi <- tapply(ids$AjSi,ids$ID1,FUN=mean) #individual social environment.
		inds$AjSi <- tapply(ids$AjSi,ids$ID1,FUN=mean) #individual breeding value of social environment.
		inds$Aj <- tapply(ids$Aj,ids$ID1,FUN=mean) #individual average interactant phenotype
		inds$Si <- tapply(ids$edge.w,ids$ID1,FUN=mean) #individual average edge strength
		inds$IGE <- psi*inds$PjSi #indirect effects.
		inds$SPi <- inds$Ai + inds$Ei + inds$IGE #individual phenotype after interaction.
                inds$tot.Ai <- inds$Ai * (1 + inds$Si*psi) #individual total breeding value.
		inds$Degree <- rowSums(updated.network)
		ids$SPi <- inds$SPi[ids$ID1]
		ids$SPj <- inds$SPi[ids$ID2]
		
		## SAVE NETWORK
		net <- matrix(NA, nrow=N, ncol=N)
		net[cbind(ids$ID1,ids$ID2)] <- ids$edge.w
		diag(net) <- 0
		assortment.out <- assortment.continuous(updated.network, inds$APi, weighted=TRUE, SE=FALSE)$r
		network.assortment[[length(network.assortment)]] <- c(network.assortment[[length(network.assortment)]],assortment.out)
		network.density[[length(network.density)]] <- c(network.density[[length(network.density)]],gden(updated.network,mode="graph"))
		## STORE RESULTS
		results.inds[[length(results.inds)]][[i]] <- inds
		results.ids[[length(results.ids)]][[i]] <- ids
		results.nets[[length(results.nets)]][[i]] <- net
	} }


########################
# change in mean phenotype ~ homophily
########################
jpeg(file = 'figure.5.jpeg', height = 1000, width = 1000, pointsize = 6, quality = 100, res = 600)
par(mar = c(4, 6, 2, 0.5))
ass <- vector()
mean.ige <- vector()
for (i in 1:length(results.inds)) {
        for (j in 1:length(results.inds[[i]])) {
                ass <- c(ass, network.assortment[[i]][[j]])
                mean.ige <- c(mean.ige, mean(results.inds[[i]][[j]]$SPi-results.inds[[i]][[j]]$APi)) } }
plot(ass, mean.ige, xlim=c(0,1), ylim=c(-0.2,0.2), xlab="", ylab="", axes = F
		,pch=19, cex=0.2, col="#66666666")
mtext(side = 1, 'Network homophily', line = 2.5, cex = 1)
mtext(side = 2, expression(paste('Change in phenotypic mean (', bar(z), ')')), line = 2.8, cex = 1)
box()
axis(side = 1, at = seq(0, 1, 0.2), las = 1, cex = 0.8)
axis(side = 2, at = seq(-0.2, 0.2, 0.1), las = 1)
points(seq(min(ass, na.rm = T), max(ass, na.rm = T), 0.001), predict(gam(mean.ige ~ s(ass, bs = 'cs')), newdata = list(ass = seq(min(ass, na.rm = T), max(ass, na.rm = T), 0.001))), type = 'l', col = 'red')
text(0, 0.16*0.85, pos = 4, c(expression(paste(psi['g'], ' =   ', 4))), cex = 1.2)
text(0, 0.18*0.95, pos = 4, c(expression(paste(N['group'], ' =  ', 50))), cex = 1.2)
dev.off()

##########################
#### multipanel
##########################
jpeg(file = 'figure.6.jpeg', height = 2000, width = 2000, pointsize = 6, quality = 100, res = 600)
par(mfrow = c(2,2))
par(mar = c(4, 6, 2, 0.5))
ass <- vector()
var.ige <- vector()
for (i in 1:length(results.inds)) {
        for (j in 1:length(results.inds[[i]])) {
                ass <- c(ass, network.assortment[[i]][[j]])
                var.ige <- c(var.ige, var(results.inds[[i]][[j]]$IGE)/var(results.inds[[i]][[j]]$APi)) } }
plot(ass, var.ige, xlim=c(0,1), ylim=c(0,0.12), axes = F, xlab="", ylab=""
		,pch=19, cex=0.2, col="#66666666")  
mtext(side = 1, 'Network homophily', line = 2.5, cex = 1)
mtext(side = 2, expression(paste('Variance in indirect effects (', sigma['g']^2, ')')), line = 2.8, cex = 1)
box()
axis(side = 1, at = seq(0, 1, 0.2), las = 1, cex = 0.8)
axis(side = 2, at = seq(0, 0.12, 0.03), las = 1)
points(seq(min(ass, na.rm = T), max(ass, na.rm = T), 0.001), predict(gam(var.ige ~ s(ass, bs = 'cs')), newdata = list(ass = seq(min(ass, na.rm = T), max(ass, na.rm = T), 0.001))), type = 'l', col = 'red', lwd = 2)
text(0, 0.115*0.85, pos = 4, expression(paste(psi['g'], ' =   ', 4)), cex = 1.2)
text(0, 0.115*0.95, pos = 4, expression(paste(N['group'], ' = ', 50)), cex = 1.2)
mtext('a)', side = 1, adj = -0.5, line = -15.4, cex = 2)

par(mar = c(4, 6, 2, 0.5))
ass <- vector()
cor <- vector()
for (i in 1:length(results.inds)) {
        for (j in 1:length(results.inds[[i]])) {
                ass <- c(ass, network.assortment[[i]][[j]])
                cor <- c(cor,cor(results.inds[[i]][[j]]$Ai, results.inds[[i]][[j]]$AjSi))}}
plot(ass, cor, xlim=c(0,1), ylim=c(-1,1), xlab="", ylab="", axes = F
		,pch=19, cex=0.2, col="#66666666") 
mtext(side = 1, 'Network homophily', line = 2.5, cex = 1)
mtext(side = 2, 'Direct - indirect genetic effects', line = 3.7, cex = 1)
mtext(side = 2, expression(paste('correlation (r'[paste('a,', bar(sa),'\'')], ')')), line = 2.3, cex = 1)
box()
axis(side = 1, at = seq(0, 1, 0.2), las = 1, cex = 0.8)
axis(side = 2, at = seq(-1, 1, 0.5), las = 1)
points(seq(min(ass, na.rm = T), max(ass, na.rm = T), 0.001), predict(gam(cor ~ s(ass, bs = 'cs')), newdata = list(ass=seq(min(ass, na.rm = T), max(ass, na.rm = T), 0.001))), type = 'l', col = 'red', lwd = 2)
mtext('b)', side = 1, adj = -0.5, line = -15.4, cex = 2)

# phenotypic variance ~ homophily
par(mar = c(4, 6, 2, 0.5))
ass <- vector()
var.pi <- vector()
for (i in 1:length(results.inds)) {
        for (j in 1:length(results.inds[[i]])) {
                ass <- c(ass, network.assortment[[i]][[j]])
                var.pi <- c(var.pi,
			(var(results.inds[[i]][[j]]$SPi)-var(results.inds[[i]][[j]]$APi))/var(results.inds[[i]][[j]]$APi)  )}}
plot(ass, var.pi, xlim=c(0,1), ylim=c(-0.06,0.75), axes = F, xlab="", ylab="" 
		,pch=19, cex=0.2, col="#66666666")
mtext(side = 1, 'Network homophily', line = 2.5, cex = 1)
#mtext(side = 2, expression(paste('Change in phenotypic variance (', sigma[bar('z')]^2, ')')), line = 2.8, cex = 1)
mtext(side = 2, 'Relative change', line = 3.7, cex = 1)
mtext(side = 2, expression(paste('in phenotypic variance (', sigma[bar('z')]^2, ')')), line = 2.3, cex = 1)
box()
axis(side = 1, at = seq(0, 1, 0.2), las = 0.8)
axis(side = 2, at = seq(0, 0.75, 0.25), las = 1)
points(seq(min(ass, na.rm = T), max(ass, na.rm = T), 0.001), predict(gam(var.pi ~ s(ass, bs = 'cs')), newdata = list(ass=seq(min(ass, na.rm = T), max(ass, na.rm = T), 0.001))), type = 'l', col = 'red', lwd = 2)
mtext('c)', side = 1, adj = -0.5, line = -15.4, cex = 2)

# genotypic variance ~ homophily
par(mar = c(4, 6, 2, 0.5))
ass <- vector()
var.gi <- vector()
for (i in 1:length(results.inds)) {
        for (j in 1:length(results.inds[[i]])) {
                ass <- c(ass, network.assortment[[i]][[j]])
		var.gi <- c(var.gi,
			(var(results.inds[[i]][[j]]$tot.Ai)-var(results.inds[[i]][[j]]$Ai))/var(results.inds[[i]][[j]]$Ai)  )}}
plot(ass, var.gi, xlim=c(0,1), ylim=c(0, 1.6), axes = F, xlab="", ylab="" 
		,pch=19, cex=0.2, col="#66666666")
mtext(side = 1, 'Network homophily', line = 2.5, cex = 1)
#mtext(side = 2, expression(paste('Change in genetic variance (', sigma['A']^2, ')')), line = 2.8, cex = 1)
mtext(side = 2, 'Relative change', line = 3.7, cex = 1)
mtext(side = 2, expression(paste('in genetic variance (', sigma['A']^2, ')')), line = 2.3, cex = 1)
box()
axis(side = 1, at = seq(0, 1, 0.2), las = 0.8)
axis(side = 2, at = seq(0, 1.5, 0.3), las = 1)
points(seq(min(ass, na.rm = T), max(ass, na.rm = T), 0.001), predict(gam(var.gi ~ s(ass, bs = 'cs')), newdata = list(ass=seq(min(ass, na.rm = T), max(ass, na.rm = T), 0.001))), type = 'l', col = 'red', lwd = 2)
mtext('d)', side = 1, adj = -0.5, line = -15.4, cex = 2)
dev.off()

########################
#response to selection by homophily
########################
jpeg(file = 'figure.7.jpeg', height = 1000, width = 1000, pointsize = 6, quality = 100, res = 600)
par(mar = c(4, 6, 2, 0.5))
ass <- vector()
response <- vector()
for (i in 1:length(results.inds)) {
        for (j in 1:length(results.inds[[i]])) {
                ass <- c(ass, network.assortment[[i]][[j]])
		response <- c(response, beta * cov(results.inds[[i]][[j]]$tot.Ai, results.inds[[i]][[j]]$SPi)) }}
plot(response ~ ass, xlim=c(0,1), ylim=c(0,0.17), axes = F, xlab="", ylab=""
                ,pch=19, cex=0.2, col="#66666666")
mtext(side = 1, 'Network homophily', line = 2.5, cex = 1)
mtext(side = 2, expression(paste('Response to selection (', Delta,'z', ')', sep = '')), line = 2.8, cex = 1)
axis(side = 1, at = seq(0, 1, 0.2), las = 1, cex = 0.8)
axis(side = 2, at = seq(0, 0.16, 0.04), las = 1, cex = 0.8)
box()
text(1, 0.005, pos = 2, expression(paste(psi['g'], ' =   ', 4)), cex = 1.2)
text(1, 0.02, pos = 2, expression(paste(N['group'], ' = ', 50)), cex = 1.2)
text(1, 0.038, pos = 2, expression(paste(beta, ' = ', 0.2)), cex = 1.2)

points(seq(min(ass, na.rm = T),max(ass, na.rm = T),0.001), predict(gam(response ~ s(ass, bs = 'cs')), newdata = list(ass = seq(min(ass, na.rm = T),max(ass, na.rm = T),0.001))), type = 'l', col = 'red', lwd = 2)
text(1, 0.35*0.8, pos = 2, expression(paste(psi['g'], ' = ', 0.4)), cex = 1.2)
text(1, 0.35*0.9, pos = 2, expression(paste(N['group'], ' = ', 50)), cex = 1.2)
dev.off()


##### compute standardised psi
psi.standardised <- array(dim = length(results.inds), data = NA)
for (i in 2:length(results.inds)){
psi.standardised[i] <-
coef(lm( scale(unlist(lapply(results.inds[[i]], '[', 'SPi'))) ~
scale(unlist(lapply(results.inds[[i]], '[', 'PjSi')))
))[2] }