K <- rbind(c(0.9, 0.1),                   # Define exaple transition kernel
           c(0.3, 0.7)) 
lam <- c(0.2,0.8)                         # Set (arbitrary) initial distribution

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

# Task 1

# Compute m-step transition kernel
K.m.step <- function(K, m=1) {
  K.m <- diag(nrow(K))                   # Create identity matrix of correct size
  if (m>0)
    for (i in 1:m) 
      K.m <- K.m %*% K                   # Repeatedly multiply it by K
  K.m
}
      
K.m.step(K,2)                            # Print 2-step transition kernel

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

# Task 2

# Compute distribution of states at time m
distribution.at.m <- function(K, lambda, m=1) {
  t(lambda)%*%K.m.step(K,m)             # Compute lambda'K(m)
}

distribution.at.m(K, lam, 2)            # Print distribution of states at t=2

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

# Task 3

# 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
  mu
}

mu <- invariant.distribution(K)         # Compute invariant distribution
mu

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

# Compute the invariant distribution (alternative implementation)
invariant.distribution.alternative <- function(K) {
  n <- nrow(K)                          # Get number of rows
  A <- rbind(t(K)-diag(n),1)            # Compute K'-I and add row of ones
  b <- rep(0,n+1)                       # Set b = [0 ... 0 1]
  b[n+1] <- 1                  
  mu <- qr.solve(A,b)                   # Solve this system of equations
  mu  
}

mu <- invariant.distribution.alternative(K)         
mu                                      # Compute invariant distribution

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

# Task 4

n <- 20
distn <- matrix(nrow=n+1,ncol=2)       # Create empty matrix to store result
for (i in 1:nrow(distn))
  distn[i,] <- distribution.at.m(K, lam, i-1)    
                                       # Compute distribution of states
                                       # for t=0..n
matplot(0:n,distn, type="b", pch=15:16, col=1:2, lty=1, ylim=0:1)
abline(h=mu,col=1:2,lty=3)

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

# Task 5

# Draw path of length n from the Markov chain
sample.chain <- function(K, lambda, n) {
  sample.path <- numeric(n)             # Set up vector to hold results
  sample.path[1] <- sample(nrow(K), 1, prob=lambda)
                                        # Draw X[0] from the initial distribution
  for (t in 2:n)                        # Draw remaining X[t] from conditional
                                        # distribution if X[t+1] | X[t]=x[t]
    sample.path[t] <- sample(nrow(K), 1, prob=K[sample.path[t-1],])
  sample.path
}

mySample <- sample.chain(K, lam, 1000)

plot(mySample, type="b")

# Task 6
table(mySample)