/* Global variables and constants */
// constants to change
final int NUMBER_OF_MASSES = 20;       // must be >= 2
final float GRAVITY_CONSTANT = 10.0;  // in m/s^2
final float ROD_LENGTH = 1.0;         // in meters
final float RATIO_MASS_DIAMETER_TO_ROD_LENGTH = 1.0/10.0;
final float TIMESTEP = 1.0/30.0;     // in seconds, equivalent to h in RK4 description

// these constants calculated in setup()
// not for user manipulation
float ROD_LENGTH_PIXELS;
float MASS_DIAMETER_PIXELS;
float CENTER_PIXEL_X;  // origin is at the ceiling
float CENTER_PIXEL_Y;
float RATIO_G_TO_L;
PFont displayFont;

// current state array:
// has 2 arrays, each array has a subarray with NUMBER_OF_MASSES elements
// element 0 contains theta, element 1 contains thetadot
// i.e. current_state[1][k] = thetadot_k
float[][] currentState = new float[2][NUMBER_OF_MASSES];

void setup()
{
  frameRate(30);
  size(400, 400);
  smooth();
  displayFont = loadFont("LucidaConsole-12.vlw");  // load font for text
  
  // calculate constants
  // must fit NUMBER_OF_MASSES rods in the window and have space left over
  ROD_LENGTH_PIXELS = 17.0/20.0 * ((float) height / NUMBER_OF_MASSES);
  MASS_DIAMETER_PIXELS = ROD_LENGTH_PIXELS * RATIO_MASS_DIAMETER_TO_ROD_LENGTH;
  CENTER_PIXEL_X = width/2;
  CENTER_PIXEL_Y = 1.0/20.0*height;
  RATIO_G_TO_L = GRAVITY_CONSTANT / ROD_LENGTH;
  
  setInitialConditions();  
}

// sets up initial conditions for the differential equations
// angles must be set up in radians
void setInitialConditions()
{
  for(int i = 0; i < currentState.length; i++)
    for(int j = 0; j < currentState[0].length; j++)
      currentState[i][j] = 0.0;
      
  // last mass given initial velocity
  currentState[1][NUMBER_OF_MASSES-1] = -5;
}

void drawCurrentState()
{
  // draw ceiling
  strokeWeight(1);
  stroke(0, 70);
  fill(0, 50);
  rect(1.0/20.0*width, 0, 9.0/10.0*width, CENTER_PIXEL_Y);
  
  // draw pendulum
  float originX = CENTER_PIXEL_X;  // initial value (origin at ceiling)
  float originY = CENTER_PIXEL_Y;  // initial value
  float massX, massY;              // first iteration: position of first mass
  for(int i = 0; i < NUMBER_OF_MASSES; i++)
  {
    // calculate
    massX = originX + ROD_LENGTH_PIXELS * sin(currentState[0][i]);
    massY = originY + ROD_LENGTH_PIXELS * cos(currentState[0][i]);
    
    // draw line
    strokeWeight(2);
    stroke(0, 100);
    line(originX, originY, massX, massY);
    
    // draw mass
    strokeWeight(0);
    fill(0, 170);
    ellipse(massX, massY, MASS_DIAMETER_PIXELS, MASS_DIAMETER_PIXELS);
    
    // update origin
    originX = massX;
    originY = massY;
  }
}

// input to function is a state array rk4_stateselect_input, equivalent to x in RK4 description
// output is state array rk4_stateselect_output, equivaleng to f(x) in RK4 description
// describes N differential equations for an N-mass pendulum
float[][] rk4_stateselect_input = new float[2][NUMBER_OF_MASSES];
float[][] rk4_stateselect_output = new float[2][NUMBER_OF_MASSES];
void rk4_stateselect_function()    // equivalent to f() in RK4 description
{
  // first half of x vector
  for(int i = 0; i < NUMBER_OF_MASSES; i++)
    rk4_stateselect_output[0][i] = rk4_stateselect_input[1][i];
  
  // second half of x vector, relation a bit more complicated
  rk4_stateselect_output[1][0] = RATIO_G_TO_L * (-sin(rk4_stateselect_input[0][0]) + sin(rk4_stateselect_input[0][1]));
  for(int i = 1; i < NUMBER_OF_MASSES - 1; i++)
    rk4_stateselect_output[1][i] = RATIO_G_TO_L * (sin(rk4_stateselect_input[0][i-1]) - 2.0 * sin(rk4_stateselect_input[0][i]) + sin(rk4_stateselect_input[0][i+1]));
  rk4_stateselect_output[1][NUMBER_OF_MASSES - 1] = RATIO_G_TO_L * (sin(rk4_stateselect_input[0][NUMBER_OF_MASSES - 2]) - 2.0 * sin(rk4_stateselect_input[0][NUMBER_OF_MASSES - 1]));
}

// Runge-Kutta 4th order routine
// updates the state into the currentState array
// uses rk4_a, rk4_b, rk4_c, rk4_d arrays as temp variables
float[][] rk4_a = new float[2][NUMBER_OF_MASSES];
float[][] rk4_b = new float[2][NUMBER_OF_MASSES];
float[][] rk4_c = new float[2][NUMBER_OF_MASSES];
float[][] rk4_d = new float[2][NUMBER_OF_MASSES];
void rk4_stateupdate()
{
  // set up input for a
  for(int i = 0; i < 2; i++)
    for(int j = 0; j < NUMBER_OF_MASSES; j++)
      rk4_stateselect_input[i][j] = currentState[i][j];
  // find value of f at that point
  rk4_stateselect_function();
  // return to a
  for(int i = 0; i < 2; i++)
    for(int j = 0; j < NUMBER_OF_MASSES; j++)
      rk4_a[i][j] = rk4_stateselect_output[i][j];
  
  // set up input for b
  for(int i = 0; i < 2; i++)
    for(int j = 0; j < NUMBER_OF_MASSES; j++)
      rk4_stateselect_input[i][j] = currentState[i][j] + TIMESTEP/2.0 * rk4_a[i][j];
  // find value of f at that point
  rk4_stateselect_function();
  // return to b
  for(int i = 0; i < 2; i++)
    for(int j = 0; j < NUMBER_OF_MASSES; j++)
      rk4_b[i][j] = rk4_stateselect_output[i][j];
  
  // set up input for c
  for(int i = 0; i < 2; i++)
    for(int j = 0; j < NUMBER_OF_MASSES; j++)
      rk4_stateselect_input[i][j] = currentState[i][j] + TIMESTEP/2.0 * rk4_b[i][j];
  // find value of f at that point
  rk4_stateselect_function();
  // return to c
  for(int i = 0; i < 2; i++)
    for(int j = 0; j < NUMBER_OF_MASSES; j++)
      rk4_c[i][j] = rk4_stateselect_output[i][j];
  
  // set up input for d
  for(int i = 0; i < 2; i++)
    for(int j = 0; j < NUMBER_OF_MASSES; j++)
      rk4_stateselect_input[i][j] = currentState[i][j] + TIMESTEP * rk4_c[i][j];
  // find value of f at that point
  rk4_stateselect_function();
  // return to d
  for(int i = 0; i < 2; i++)
    for(int j = 0; j < NUMBER_OF_MASSES; j++)
      rk4_d[i][j] = rk4_stateselect_output[i][j];
  
  // update state
  for(int i = 0; i < 2; i++)
    for(int j = 0; j < NUMBER_OF_MASSES; j++)
      currentState[i][j] += TIMESTEP/6.0 * (rk4_a[i][j] + 2.0 * rk4_b[i][j] + 2.0 * rk4_c[i][j] + rk4_d[i][j]);
}

void draw() {
  background(250);
  drawCurrentState();
  rk4_stateupdate();
  
  // draw framerate and text
  fill(0, 150);
  textFont(displayFont);
  text("N-mass pendulum (chain) simulation", 5, height - 20);
  text("Right-click to restart", 5, height - 10);
  text("Framerate: " + floor(frameRate), width - 100, height - 10);
}

// event handlers
void mouseReleased()
{
  if(mouseButton == RIGHT)
    setInitialConditions();
}