##################################
#      R code for question 7     #
##################################

# We will be using the code from the computer practicals to compute
# the invariant distribution 
invariant.distribution <- function(K) {
  mu <- eigen(t(K))$vectors[,1]     # Find the eigenvector corresponding
                                    # to the largest eigenvalue (i.e. 1) 
  mu <- mu / sum(mu)                # Normalise distribution
  as.numeric(mu)                    # Return mu (with complex part (=0) removed)
}

# Define cross-tabulation matrix
C <- rbind( c(50, 19, 26, 8, 7, 11, 6, 2),
            c(16, 40, 34, 18, 11, 20, 8, 3),
            c(12, 35, 65, 66, 35, 88, 23, 21),
            c(11, 20, 58, 110, 40, 183, 64, 32),
            c(2, 8, 12, 23, 25, 46, 28, 12),
            c(12, 28, 102, 162, 90, 554, 230, 177),
            c(0, 6, 19, 40, 21, 158, 143, 71),
            c(0, 3, 14, 32, 15, 126, 91, 106))

# Absolute frequencies of occupations of fathers is row-wise sum
fathers <- apply(C, 1, sum)
# Absolute frequencies of occupations of sons is column-wise sum
sons <- apply(C, 2, sum)

# Compute the relative frequencies
freq.fathers <- fathers / sum(fathers)
freq.sons <- sons / sum(sons)
print(freq.fathers)
print(freq.sons)

# To turn C into a transition kernel we need to divide each row by its sum
K = sweep(C, 1, fathers, "/")

# Compute the invariant distribution
mu = invariant.distribution(K)
print(mu)

# Visualise the result using a barplot
colors <- c("darkblue","red","yellow")
barplot(rbind(freq.fathers,freq.sons,mu), beside=TRUE, col=colors,
        names.arg=paste("Category",1:8))
legend("topleft", fill=colors, legend=c("Actual distribution of fathers",
                                     "Actual distribution of sons",
                                     "Invariant distribution"))

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

##################################
#      R code for question 8     #
##################################

# Simulates a path of length n+1 from the Wright-Fisher model
# Initial distribution is given by the array initial (length 2)
simulate.wright.fisher <- function(initial, n) {
  result <- matrix(nrow=n+1, ncol=2)             # Create matrix to store result
  result[1,] <- initial                          # Set first line to initial dist'n
  two.N <- sum(initial)                          # Compute 2*N
  for (t in 1:n) {
    result[t+1,1] <- rbinom(1, size=two.N, prob=result[t,1]/two.N)
                                                # Draw X[t+1]|X[t]=x[t]
    result[t+1,2] <- two.N - result[t+1,1]      # Compute 2*N - X[t+1] 
  }
  result
}

wf.path <- simulate.wright.fisher(c(300,100), 1000)
wf.path

matplot(wf.path, type="o", col=3:4, pch=15:16, lty=1,
        xlab="Time", ylab="Count")
legend("topleft", col=3:4, lty=1, pch=15:16, 
       legend=c("Allele 1","Allele 2"))