//   rwm.java
//   Copyright 2006 by Jeffrey Rosenthal
//   All rights reserved.
//   jeff@math.toronto.edu
//   http://probability.ca/jeff/

import java.util.*;
import java.lang.Math;
import java.applet.*;
import java.awt.*;
import java.awt.image.*;

public class rwm extends Applet implements Runnable {

  int i, j, k;
  int ix, iproposal;
  int igamma = 1;
  double SUM, mean, tmpdouble;
  int iteration;
  boolean REDRAW = true;
  boolean FORCEDREDRAW = true;
  boolean SHOWLINES = true;
  int speedcontrol = 2;
  int MINNUMSTATES = 3;
  int MAXNUMSTATES = 10;
  double idens[] = new double[MAXNUMSTATES+1];
  int denscount[] = new int[MAXNUMSTATES+1];
  int numstates = 6;
  double smalldensval = 0.1;
  double densmean;
  double diff;
  int MINBLUERECT = 2;
  int BARWIDTH = 5;
  double tmpsum;

  int xzero = 100;
  int xspacing = 420/numstates;
  int yzero = 300;
  int yspacing = 40;
  int xaxislength = 550;
  int yaxislength = 250;
  int axiswidth = 3;
  int xmarkspacing = 1;
  double ymarkspacing = 0.1;
  int markrad = 10;
  int circlerad = 12;
  int stateheight = yzero + 60;
  int propheight = stateheight + 25;
  int segment, itmp;
  int textx = 600;
  int texty = 30;
  int textinc = 28;
  boolean accept, JUMPONE, RECOUNT;
  boolean keyjustpressed = false;
  boolean RESTARTING = true;
  int ADAPTLEVEL = 1;

  Image holdImage;
  Graphics holdGraphics;
  // Font strongfont = new Font ("TimesRoman", Font.BOLD, 18);

  public void init() {

    holdImage = createImage(size().width, size().height);
    holdGraphics = holdImage.getGraphics();
    // repaint();

  }

  public void run() {

    while (true) {

// System.out.println("Here");

      if (RESTARTING) {
          setdens();
          zerocount();
	  if (ADAPTLEVEL > 1)
            igamma = 2;
          ix = numstates - 1;
          // segment = 1;
          JUMPONE = RECOUNT = false;
          REDRAW = true;
          RESTARTING = false;
      } else if (RECOUNT) {
          zerocount();
	  RECOUNT = false;
      }

// System.out.println("Here");

      // Start fresh iteration.
      iteration++;
      SUM = SUM + ix;
      denscount[ix] = denscount[ix] + 1;

// System.out.println("Here");

      // Compute the proposal and accept/reject.
      itmp =  2*igamma + 1;
      iproposal = ix;
// System.out.println("igamma = " + igamma + " ; itmp = " + itmp + " ; ix = " + ix);
      while (iproposal == ix) {
        iproposal = ix + (int) (Math.random() * itmp) - igamma;
// System.out.println("Here");
      }

// System.out.println("Here");
      if ( (iproposal>=1) && (iproposal<=numstates) &&
             (idens[ix] * Math.random() < idens[iproposal]) ) {
        accept = true;
      } else {
        accept = false;
      }
// System.out.println("ix=" + ix + ";  igamma=" + igamma + ";  iproposal=" + iproposal + "; accept=" + accept);

// System.out.println("Here");

      // Display the various segments of the iteration.
      for (segment=1; segment<=5; segment++)
        updatenpause();

// System.out.println("Here");

      // Update state.
      if (accept)
        ix = iproposal;

// System.out.println("Here");

    }

  }

  public void updatenpause() {

    if ( (!RESTARTING) && (!RECOUNT)
                && ((speedcontrol<5) || (segment==1) || keyjustpressed)
    		&& ((segment<5) || accept) ) {

      itmp = 10*(speedcontrol-5)*(speedcontrol-5);
      if ( (speedcontrol<=5) || keyjustpressed ||
				(iteration == (iteration/itmp)*itmp) )
        repaint();
      keyjustpressed = false;

      if (speedcontrol<=1) {
        JUMPONE = false;
        while ( (speedcontrol<=1) && (JUMPONE==false) 
                  && (!RESTARTING) && (!RECOUNT) ) {
	  dosleep(50);
	}
      } else if (speedcontrol==2) {
	  dosleep(1000);
	  // for (j=0; j<2; j++) {
// System.out.println("Pause #" + j);
	    // if (!keyjustpressed) {
	    // if (true) {
	      // dosleep(500);
	    // }
	  // }
      } else if (speedcontrol==3) {
	  dosleep(250);
      } else if (speedcontrol==4) {
	  dosleep(50);
      } else if (speedcontrol==5) {
	  dosleep(40);
      } else if (speedcontrol>=6) {
	  dosleep(20);
      }

      REDRAW = FORCEDREDRAW;
      JUMPONE = false;

    }

  }

  public void paint(Graphics g) {

if (REDRAW) {

    // setFont(strongfont);

    // Background rectangle.
    g.setColor(Color.pink);
    g.fillRect(0,0,size().width,size().height);

    // Draw axes.
    g.setColor(Color.red);
    g.fillRect(xzero-axiswidth/2,yzero-yaxislength,axiswidth,yaxislength);
    g.fillRect(xzero,yzero-axiswidth/2,xaxislength,axiswidth);
    g.drawString("state", xzero+xaxislength-30, yzero+40);

    // Draw the axis arrows.
    g.drawLine(xzero+xaxislength+2,yzero,xzero+xaxislength-20,yzero-10);
    g.drawLine(xzero+xaxislength+2,yzero,xzero+xaxislength-20,yzero+10);
    g.drawLine(xzero,yzero-yaxislength-2,xzero-10,yzero-yaxislength+20);
    g.drawLine(xzero,yzero-yaxislength-2,xzero+10,yzero-yaxislength+20);

    // Mark some x-axis marks.
    for(k=0; k<=(xaxislength-30)/xspacing; k++) {
      tmpdouble = Math.floor(100*k*xmarkspacing)/100;
      // itmp = (int) Math.floor(xzero - k*xmarkspacing);
      itmp = xzero + k*xspacing;
      // g.drawString("" + tmpdouble, itmp-10, yzero+markrad+15);
      g.drawString("" + k*xmarkspacing, itmp-3, yzero+markrad+15);
      g.drawLine(itmp, yzero-markrad, itmp, yzero+markrad);
    }

    // Mark some y-axis marks.
    for(k=0; k<=(yaxislength-30)/yspacing; k++) {
      tmpdouble = Math.floor(100*k*ymarkspacing)/100;
      itmp = yzero - k*yspacing;
      g.drawString("" + tmpdouble, xzero - markrad-30, itmp+5);
      g.drawLine(xzero - markrad, itmp, xzero + markrad, itmp);
    }

    // Draw graph of density.
    g.setColor(Color.blue);
    for (j=1; j<=numstates; j++) {
      itmp = (int)Math.floor(idens[j]*yspacing/ymarkspacing);
      // fillcircle(g, xzero + j*xspacing, yzero-itmp, 5);
      if (itmp<MINBLUERECT)
        itmp = MINBLUERECT;
      g.fillRect(xzero+j*xspacing+2,yzero-itmp,BARWIDTH, itmp);
    }
    g.drawString("target", xzero-80, yzero-yaxislength-10);
    // g.drawString("  density", xzero-80, yzero-yaxislength+0);
    if (SHOWLINES) {
      // And line for mean of density.
      // itmp = (int)Math.floor(xzero+densmean*xspacing/xmarkspacing);
      // g.drawLine(itmp, yzero-yaxislength, itmp, yzero-yaxislength+30);
    }

} else {    // end of if(REDRAW)
  g.setColor(Color.pink);
  g.fillRect(xzero, yzero-50, 200, 40);
}

    // Top text matter.
    g.setColor(Color.black);
    mean = SUM / ((double)iteration);
    g.drawString("Mean state: " + Math.floor(1000*mean)/1000, textx, texty);
    g.drawString("Iteration: " + iteration, textx, texty+textinc);
    g.drawString("State: " + ix, textx, texty+2*textinc);
    g.drawString("Gamma: " + igamma, textx, texty+3*textinc);
    g.drawString("t[2]: " + idens[2], textx, texty+4*textinc);
    // Compute total variation distance between two distributions.
    diff = 0.0;
    for (j=1; j<=numstates; j++)
      diff = diff + Math.abs( idens[j] - ((double)denscount[j]) / iteration );
    g.drawString("TV dist: " + Math.floor(500000*diff)/1000000,
					textx, texty+5*textinc);
    g.drawString("Animation: " + speedcontrol, textx, texty+6*textinc);
    if (SHOWLINES) {
      for (j=1; j<=numstates; j++) {
        itmp = (int)Math.floor(((double)denscount[j]) / iteration 
      						* yspacing/ymarkspacing);
        g.fillRect(xzero+j*xspacing-BARWIDTH-2,yzero-itmp,BARWIDTH, itmp);
      }
    }

    // Draw current state.
    g.setColor(Color.black);
    if ( (segment==5) && (accept==true) )
      fillcircle(g, xzero + iproposal*xspacing, stateheight, circlerad);
    else
      fillcircle(g, xzero + ix*xspacing, stateheight, circlerad);

    if (segment == 2) {

      // Draw proposal distribution dots.
      g.setColor(Color.white);
      for (j=ix-igamma; j<=ix+igamma; j++) {
	if (j != ix)
          fillcircle(g, xzero + j*xspacing, propheight, circlerad);
      }

    } else if (segment == 3) {

        // Draw actual proposal dot.
        g.setColor(Color.yellow);
        fillcircle(g, xzero + iproposal*xspacing, propheight, circlerad);

    } else if (segment >= 4) {

          // Indicate accept/reject.
	  if (accept)
            g.setColor(Color.green);
	  else
            g.setColor(Color.red);
          fillcircle(g, xzero + iproposal*xspacing, propheight, circlerad);

    }

  }

  public void setdens() {

    idens[1] = idens[3] = idens[5] = 0.17;
    idens[2] = smalldensval;
    idens[4] = idens[6] = 0.25;
    if (numstates > 6) {
      for (j=7; j<=numstates; j++)
        idens[j] = 0.1;
    }

    tmpsum = 0.0;
    for (j=1; j<=numstates; j++)
      tmpsum = tmpsum + idens[j];
    for (j=1; j<=numstates; j++)
      idens[j] = idens[j] / tmpsum;

    densmean = 0.0;
    for (j=1; j<=numstates; j++)
      densmean = densmean + j * idens[j];
    if (numstates == 4) {
      yspacing = 30;
      ymarkspacing = 0.1;
    } else if (numstates == 3) {
      yspacing = 40;
      ymarkspacing = 0.2;
    } else {
      yspacing = 40;
      ymarkspacing = 0.1;
    }

  }

  public void zerocount() {

      SUM = mean = 0.0;
      iteration = 0;
      for (j=1; j<=numstates; j++)
        denscount[j] = 0;

  }

  public boolean action(Event e, Object o) {

    return true;
  }

  public void update(Graphics g) {
    paint(holdGraphics);
    g.drawImage(holdImage, 0, 0, this);
  }

  public void dosleep(long nummilisecs) {
      try {
        Thread.currentThread().sleep(nummilisecs);
      } catch (InterruptedException e) {
      }
  }

  public void fillcircle(Graphics g, int xx, int yy, int rr) {
      g.fillOval(xx-rr, yy-rr, 2*rr, 2*rr);
  }

  public boolean keyDown(Event evt, int key) {
    char keystroke = (char) key;

    if ( (keystroke >= '0') && (keystroke <= '9') ) {
      speedcontrol = keystroke - '0';
      if (keystroke > '0')
        JUMPONE = true;
    }

    if (keystroke == 'r') {  // restart
      RESTARTING = true;
      // init();
    }

    if (keystroke == 'z') {  // zero the counts
      // RECOUNT = true;
      zerocount();
    }

    if (keystroke == 's') {  // toggle whether to show the lines
      SHOWLINES = !SHOWLINES;
    }

    if ( (keystroke == '+') && (numstates<MAXNUMSTATES) ) { 
        // Increment the number of states.
        numstates++;
        xspacing = 420/numstates;
	smalldensval = Math.min( smalldensval, 1.0 / numstates );
        RESTARTING = true;
    }

    if ( (keystroke == '-') && (numstates>MINNUMSTATES) ) { 
        // Decrement the number of states.
        numstates--;
        xspacing = 420/numstates;
        RESTARTING = true;
    }

    if ( keystroke == '>' ) {  // Increase dens[2].
	smalldensval = Math.min( 10*smalldensval, 1.0 / numstates );
        xspacing = 420/numstates;
        RESTARTING = true;
    }

    if ( keystroke == '<' ) {  // Decrease dens[2].
	if (smalldensval > 0.1)
	  smalldensval = 0.1;
	else
	  smalldensval = smalldensval / 10.0;
        xspacing = 420/numstates;
        RESTARTING = true;
    }

    if (keystroke == 'o') {  // Always set igamma = 1.
      ADAPTLEVEL = 1;
      igamma = 1;
    }

    if (keystroke == 't') {  // Always set igamma = 2.
      ADAPTLEVEL = 1;
      igamma = 2;
    }

    if (keystroke == 'p') {  // Increment igamma.
      igamma++;
    }

    if (keystroke == 'm') {  // Decrement igamma.
      if (igamma >= 2)
        igamma--;
    }

    if (keystroke == 'F') {  // Set igamma to 50.
      igamma = 50;
    }

    if (keystroke == 'A') {  // Set ix to 1.
      ix = iproposal = 1;
    }

    if (keystroke == 'B') {  // Set ix to numstates.
      ix = iproposal = numstates;
    }

    // if (keystroke == 'l') {  // redraw screen
      // REDRAW = true;
    // }

    // if (keystroke == 'L') {  // toggle forced redraw
      // FORCEDREDRAW = !FORCEDREDRAW;
    // }

    keyjustpressed = true;
    repaint();
    return true;
  }

  Thread t;

  public void start() {
    t = new Thread(this);
    t.start();
  }

  public void stop() {
    t.stop();
    t = null;
  }

}