# drawchain() -- a function to draw a Markov chain diagram

# by Jeffrey S. Rosenthal (www.probability.ca).

# usage: drawchain(P), where P is an N by N matrix of transition probabilites.


library(diagram)
library(plotrix)

thepoint = function(rot, rad) {
  theangle = 2*pi*rot
  thex = 0.5+rad*sin(theangle)
  they = 0.5+rad*cos(theangle)
  return( c(thex,they) )
}
circleat = function(pp, rad) {
  draw.circle(pp[1], pp[2], rad)
}
textat = function(pp, thetext, ...) {
  text(pp[1], pp[2], thetext, ...)
}

drawtwochain = function(trans) {
  N=sqrt(length(trans))
  P = t( matrix(trans, nrow=N) )
  print(P)
  plot.new()
  p1 = c(0.4,0.5)
  p2 = c(0.6,0.5)
  if (P[1,2]!="0") {
    curvedarrow(p1,p2,curve=-0.5, segment=c(0.15,0.8), arr.pos=0.8)
    p3 = (p1+p2)/2 + c(0,0.135)
    textat(p3, P[1,2])
  }
  if (P[2,1]!="0") {
    curvedarrow(p2,p1,curve=-0.5, segment=c(0.15,0.8), arr.pos=0.8)
    p3 = (p1+p2)/2 + c(0,-0.135)
    textat(p3, P[2,1])
  }
  if (P[1,1]!="0") {
    uu = p1 + c(0,0.01)
    ll = p1 - c(0,0.01)
    curvedarrow(uu,ll,curve=8, segment=c(0.07,0.85), arr.pos=0.88)
    p3 = p1 + c(-0.1,0.04)
    textat(p3, P[1,1])
  }
  if (P[2,2]!="0") {
    uu = p2 + c(0,0.01)
    ll = p2 - c(0,0.01)
    curvedarrow(uu,ll,curve=-8, segment=c(0.07,0.85), arr.pos=0.88)
    p3 = p2 + c(0.1,0.04)
    textat(p3, P[2,2])
  }
  circleat(p1, 0.03)
  textat(p1, "1", font=2)
  circleat(p2, 0.03)
  textat(p2, "2", font=2)
}  # End of "drawchain()" definition.


drawchain = function(trans) {

N=sqrt(length(trans))

if (N==2) {
  drawtwochain(trans)
} else {

P = t( matrix(trans, nrow=N) )
print(P)
plot.new()
mainrad = 0.15
circlerad = 0.03
for (i in 1:N) {
  # Draw the states.
  mainp = thepoint(i/N, mainrad)
  circleat(mainp, circlerad)
  textat(mainp, i, font=2)
  if (N>=4) ind=1 else ind=0
  # Draw transitions from i to i+1
  p1 = thepoint((i+0.1)/N, mainrad+0.00)
  p2 = thepoint((i+0.9)/N, mainrad+0.00)
  p3 = thepoint((i+0.5)/N, mainrad+0.05+0.01*ind)
  if (i==N) thetext = P[i,1] else thetext = P[i,i+1]
  if (thetext!="0") {
    textat(p3, thetext)
    curvedarrow(p1,p2,curve=-0.3, segment=c(0.1,0.8), arr.pos=0.85)
  }
  # Draw transitions from i to i-1
  p1 = thepoint((i-0.05)/N, mainrad-0.00)
  p2 = thepoint((i-0.99)/N, mainrad-0.00)
  p3 = thepoint((i-0.5)/N, mainrad-0.06-0.03*ind)
  if (i==1) thetext = P[i,N] else thetext = P[i,i-1]
  if (thetext!="0") {
    textat(p3, thetext)
    curvedarrow(p1,p2,curve=-0.1, segment=c(0.25,0.63), arr.pos=0.68)
  }
  # Draw transitions from i to i
  p1 = thepoint((i-0.05)/N, mainrad+0.03)
  p2 = thepoint((i+0.05)/N, mainrad+0.03)
  p3 = thepoint(i/N, mainrad+0.18-0.04*ind)
  thetext = P[i,i]
  if (thetext!="0") {
    textat(p3, thetext)
    curvedarrow(p1,p2,curve=-3, segment=c(0.03,0.87), arr.pos=0.92)
  }
}

}

}  # End of "drawchain()" definition.


# EXAMPLES:
drawchain( c("2/3","1/3","3/4","1/4") )
drawchain( c("0","1/2","1/2","1/3","1/3","1/3","1/4","1/4","1/2") )
drawchain( c("1/2","1/3","1/6","1/3","0","2/3","0","1/4","3/4") )