#
# testem.S
#
# demo program for complicated models
#
# Splus test code of the EM algorithm for mixture models
# 	copyright by Shotaro Akaho <s.akaho@aist.go.jp>
#                    AIST Neuroscience Research Instutute, Japan 
# Notice:
#  1. Only noncommercial use is allowed
#  2. Attach this header in the case of redistribution
#  3. If you find any bugs, it is strongly recommended to inform them to me.

h1 <- gaussdiag(c(3, 3), c(1, 0.5))
h2 <- gaussdiag(c(-3, 3), c(1, 0.5))
h3 <- gaussdiag(c(0, -3), c(0.5, 1))
hh1 <- mixture(list(h1, h2, h3))

test.typeI <- function(scale = c(1, 1), shift = c(0, 0),
  n = 1000, k = 10, model = hh1, xlim = c(-5,5), ylim = c(-5,5), id="file") {
  model.t <- typeI(model, scale = scale, shift = shift)
  rsample <- rprob(model.t, n)
  model.i <- typeI(model, scale = c(1, 1), shift = c(0, 0))
  test.general(model.t, model.i, rsample, k, xlim, ylim, id)
}

test.typeIm <- function(scale = list(c(1, 1), c(1, 1)),
  shift = list(c(0, 0), c(5, 0)), n = 1000, k = 10, model = hh1, 
  xlim=c(-10,10), ylim=c(-10,10), id="file") {
  model.i <- model.t <- list()
  for(i in 1:length(scale)) {
    model.t[[i]] <- typeI(model, scale=scale[[i]], shift=shift[[i]])
    model.i[[i]] <- typeI(model, scale=c(1, 1), shift=c(5 * (2 * i - 3), 0))
  }
  model.t <- mixture(model.t)
  rsample <- rprob(model.t, n)
  model.i <- mixture(model.i)
  test.general(model.t, model.i, rsample, k, xlim, ylim, id)
}

g1 <- gaussconst(c(0, 0), sd=.8)
g2 <- gaussconst(c(3, 3), sd=1)
g3 <- gaussconst(c(-3,3), sd=1)
g4 <- gaussconst(c(0,-3), sd=1)
gg1 <- mixture(list(g1, g2, g3, g4))

test.typeII <- function(sclrot = c(1, 0), shift = c(0, 0), 
  n = 1000, k = 10, model = gg1, xlim=c(-5,5), ylim=c(-5,5), id="file") {
  model.t <- typeII(model, sclrot=sclrot, shift=shift)
  rsample <- rprob(model.t, n)
  model.i <- typeII(model, sclrot=c(1, 0), shift=c(0, 0))
  test.general(model.t, model.i, rsample, k, xlim, ylim, id)
}

test.typeIIm <- function(sclrot = list(c(1, 0), c(1, 0)),
  shift = list(c(0, 0), c(3, 0)), n = 1000, k = 10, model = gg1, 
  xlim=c(-10,10), ylim=c(-10,10), id="file") {
  model.i <- model.t <- list()
  for(i in 1:length(sclrot)) {
    model.t[[i]] <- typeII(model, sclrot=sclrot[[i]], shift=shift[[i]])
    model.i[[i]] <- typeII(model, sclrot=c(1, 0), shift=c(5 * (2 * i - 3), 0))
  }
  model.t <- mixture(model.t)
  model.i <- mixture(model.i)
  rsample <- rprob(model.t, n)
  test.general(model.t, model.i, rsample, k, xlim, ylim, id)
}

test.general <- function(model.t, model.i, rsample, k, xlim, ylim, id) {
  like <- rep(0, k + 1)
  like.max <- like.em(model.t, rsample)
  plot.em2(model.i, rsample, xlim, ylim, id)
  cat("like.max : ", like.max, "\n")
  for(i in 1:k) {
    like[i] <- like.em(model.i, rsample)
    cat("Step", i, ":", like[i], "\n")
    model.i <- estimate(model.i, rsample)
  }
  like[k + 1] <- like.em(model.i, rsample)
  plot.em2(model.i, rsample, xlim, ylim, id, add=T)
  postscript(paste(id, "-2.eps", sep=""))
  par(mar=c(5,4,1,1)+.1)
  plot(c(0:k,0), c(like,like.max), type="n", xlab="step", ylab="likelihood")
  lines(0:k, like)
  par(lty=2)
  lines(c(0,k), rep(like.max, 2))
  par(lty=1)
  like
}

plot.circ <- function(mu, sd, xlim, ylim, add=F) {
  if(add) {
    n <- 32
    th <- (0:n) / n * 2 * 3.141592
    x <- mu[1] + sd[1] * cos(th)
    y <- mu[2] + sd[2] * sin(th)
    lines(x, y)
  } else {
    n <- 32
    for(i in 0:(n/2)) {
      th1 <- (2 * i) / n * 2 * 3.141592
      th2 <- (2 * i + 1) / n * 2 * 3.141592
      x <- mu[1] + sd[1] * cos(c(th1, th2))
      y <- mu[2] + sd[2] * sin(c(th1, th2))
      lines(x, y)
    }
  }
}

plot.em2 <- function(model, dat, xlim=NULL, ylim=NULL, id = id, add=F) {
  x <- dat[, 1]
  y <- dat[, 2]
  if(is.null(xlim)) xlim <- c(min(x), max(x))
  if(is.null(ylim)) ylim <- c(min(y), max(y))
  if(xlim[1] > min(x)) xlim[1] <- min(x)
  if(xlim[2] < max(x)) xlim[2] <- max(x)
  if(ylim[1] > min(y)) ylim[1] <- min(y)
  if(ylim[2] < max(y)) ylim[2] <- max(y)
  if(!add) {
    postscript(paste(id, "-1.eps", sep=""))
    par(mar=c(5,4,1,1)+.1)
    plot(xlim, ylim, type="n", xlab="x", ylab="y")
    points(x, y, pch=".")
  }
  plot.em3(model, xlim, ylim, add)
}

plot.em3 <- function(model, xlim, ylim, add) {
  k <- length(model$models)
  if(class(model) == "typeI") {
    for(i in 1:k) {
      m <- model$models[[i]]
      scale <- model$scale
      shift <- model$shift
      plot.circ((m$mu - shift) / scale, 2 * m$sd / sqrt(scale), xlim, ylim, add)  
    }
  } else if(class(model) == "typeII") {
    for(i in 1:k) {
      m <- model$models[[i]]
      sclrot <- model$sclrot
      shift <- model$shift
      hh <- sclrot[1]^2 + sclrot[2]^2
      mu1 <- (sclrot[1] * (m$mu[1] - shift[1]) - 
        sclrot[2] * (m$mu[2] - shift[2])) / hh
      mu2 <- (sclrot[2] * (m$mu[1] - shift[1]) + 
        sclrot[1] * (m$mu[2] - shift[2])) / hh
      sd <- 2 * m$sd / sqrt(hh)
      plot.circ(c(mu1,mu2),c(sd,sd), xlim, ylim, add)
    }
  } else if(class(model) == "mixture") {
    for(i in 1:k) {
      plot.em3(model$models[[i]], xlim, ylim, add)
    }
  } else {
    cat(class(model)) 
    stop("unsupported")
  }
}

set.seed(0)
ps.options(width=4,height=4)
test.typeI(c(1.5,1.2),c(4,1),id="fig-1")
test.typeIm(list(c(1.5,1.2),c(0.7,1.5)),list(c(4,1),c(-6,-2)),id="fig-2")
test.typeII(c(1.5,1),c(4,1),id="fig-3")
test.typeIIm(list(c(1.5,1),c(0.7,-0.5)),list(c(4,1),c(-6,-2)),id="fig-4")