/**
 * Hunter Adams (vha3@cornell.edu)
 * 
 * HARDWARE CONNECTIONS
 *  - GPIO 16 ---> VGA Hsync
 *  - GPIO 17 ---> VGA Vsync
 *  - GPIO 18 ---> 330 ohm resistor ---> VGA Green lo-bit |__ both wired to 150 ohm to ground 
 *  - GPIO 19 ---> 220 ohm resistor ---> VGA Green hi_bit |   and to VGA Green
 *  - GPIO 20 ---> 330 ohm resistor ---> VGA Blue
 *  - GPIO 21 ---> 330 ohm resistor ---> VGA Red
 *  - RP2040 GND ---> VGA GND
 *
 * RESOURCES USED
 *  - PIO state machines 0, 1, and 2 on PIO instance 0
 *  - DMA channels 0, 1, 2, and 3
 *  - 153.6 kBytes of RAM (for pixel color data)
 *
 * Protothreads v1.1.1
 * Threads:
 * core 0:
 * Graphics demo
 * blink LED25 
 * core 1:
 * Toggle gpio 4 
 * Serial i/o 
 */
// ==========================================
// === VGA graphics library
// ==========================================
#include "vga16_graphics.h"
#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include "pico/stdlib.h"
#include "hardware/pio.h"
#include "hardware/dma.h"
// // Our assembled programs:
// // Each gets the name <pio_filename.pio.h>
// #include "hsync.pio.h"
// #include "vsync.pio.h"
// #include "rgb.pio.h"

// ==========================================
// === protothreads globals
// ==========================================
#include "hardware/sync.h"
#include "hardware/timer.h"
#include "pico/multicore.h"
#include "string.h"
// protothreads header
#include "pt_cornell_rp2040_v1_1_1.h"

// ==================================================
// === Lattice Boltzmann code from Daniel V. Schroeder
// ==================================================
/*
A lattice-Boltzmann fluid simulation in JavaScript, using HTML5 canvas for graphics	
	Copyright 2013, Daniel V. Schroeder
  >>> Modifed for C on  Pi Pico by Bruce Land 2022 <<<
	Permission is hereby granted, free of charge, to any person obtaining a copy of 
	this software and associated data and documentation (the "Software"), to deal in 
	the Software without restriction, including without limitation the rights to 
	use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies 
	of the Software, and to permit persons to whom the Software is furnished to do 
	so, subject to the following conditions:

	The above copyright notice and this permission notice shall be included in all 
	copies or substantial portions of the Software.

	THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, 
	INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A 
	PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR 
	ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR 
	OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR 
	OTHER DEALINGS IN THE SOFTWARE.

	Except as contained in this notice, the name of the author shall not be used in 
	advertising or otherwise to promote the sale, use or other dealings in this 
	Software without prior written authorization.
  */
 // Global floatiables:	

	int pxPerSquare = 2 ;
													// width of plotted grid site in pixels
	#define xdim  70			// grid dimensions for simulation
	#define ydim 30 
	
	char tracerCheck = true ;
	char running = false;						// will be true when running
	int stepCount = 0;
	int startTime = 0;
	float four9ths = 4.0 / 9.0;					// abbreviations
	float one9th = 1.0 / 9.0;
	float one36th = 1.0 / 36.0;
	int barrierCount = 0;
	int barrierxSum = 0;
	int barrierySum = 0;
	float barrierFx = 0.0;						// total force on all barrier sites
	float barrierFy = 0.0;
	int time = 0;								// time (in simulation step units) since data collection started
	char showingPeriod = false;
	
	// Create the arrays of fluid particle densities, etc. (using 1D arrays for speed):
	// To index into these arrays, use x + y*xdim, traversing rows first and then columns.
	float n0 [xdim*ydim];			// microscopic densities along each lattice direction
	float nN [xdim*ydim];
	float nS [xdim*ydim];
	float nE [xdim*ydim];
	float nW [xdim*ydim];
	float nNE [xdim*ydim];
	float nSE [xdim*ydim];
	float nNW [xdim*ydim];
	float nSW [xdim*ydim];
	float rho [xdim*ydim];			// macroscopic density
	float ux  [xdim*ydim];			// macroscopic x-velocity
	float uy  [xdim*ydim];          // macroscopic y-velocity
	//float curl  [xdim*ydim];
	char barrier  [xdim*ydim];		// boolean array of barrier locations

  //#define nColors 16;

  // Initialize tracers (but don't place them yet):
	//#define nTracers  144
	//int tracerX [nTracers];
	//int tracerY [nTracers];

  #define fluid_speed 0.10
  #define fluid_viscosity 0.005
  #define stepsPerFrame 1

// =========================================
// LB init function
// =========================================
void init_LB(void){
	// Initialize to a steady rightward flow with no barriers:
	for (int y=0; y<ydim; y++) {
		for (int x=0; x<xdim; x++) {
			barrier[x+y*xdim] = false;
		}
	}
	// Create a simple linear "wall" barrier (intentionally a little offset from center):
	#define barrierSize  8
	for (int y=(ydim/2)-barrierSize; y<=(ydim/2)+barrierSize; y++) {
		int x =(ydim/3);
		barrier[x+y*xdim] = true;
	}

  	//for (int t=0; t<nTracers; t++) {
	//	tracerX[t] = 0; tracerY[t] = 0;
	//}
} // end init_LB

// Set all densities in a cell to their equilibrium values for a given velocity and density:
// (If density is omitted, it's left unchanged.)
void setEquil(int x, int y, float newux, float newuy, float newrho) {
	int i = x + y*xdim;
	//if (typeof newrho == 'undefined') {
		//newrho = rho[i];
	//}
	float ux3 = 3 * newux;
	float uy3 = 3 * newuy;
	float ux2 = newux * newux;
	float uy2 = newuy * newuy;
	float uxuy2 = 2 * newux * newuy;
	float u2 = ux2 + uy2;
	float u215 = 1.5 * u2;
	n0[i]  = four9ths * newrho * (1                              - u215);
	nE[i]  =   one9th * newrho * (1 + ux3       + 4.5*ux2        - u215);
	nW[i]  =   one9th * newrho * (1 - ux3       + 4.5*ux2        - u215);
	nN[i]  =   one9th * newrho * (1 + uy3       + 4.5*uy2        - u215);
	nS[i]  =   one9th * newrho * (1 - uy3       + 4.5*uy2        - u215);
	nNE[i] =  one36th * newrho * (1 + ux3 + uy3 + 4.5*(u2+uxuy2) - u215);
	nSE[i] =  one36th * newrho * (1 + ux3 - uy3 + 4.5*(u2-uxuy2) - u215);
	nNW[i] =  one36th * newrho * (1 - ux3 + uy3 + 4.5*(u2-uxuy2) - u215);
	nSW[i] =  one36th * newrho * (1 - ux3 - uy3 + 4.5*(u2+uxuy2) - u215);
	rho[i] = newrho;
	ux[i] = newux;
	uy[i] = newuy;
}

// Function to initialize or re-initialize the fluid, based on speed slider setting:
void initFluid() {
	// Amazingly, if I nest the y loop inside the x loop, Firefox slows down by a factor of 20
	float u0 = fluid_speed;
	for (int y=0; y<ydim; y++) {
		for (int x=0; x<xdim; x++) {
			setEquil(x, y, u0, 0, 1);
			//curl[x+y*xdim] = 0.0;
		}
	}
//paintCanvas();
} // end initFluid

// Set the fluid floatiables at the boundaries, according to the current slider value:
void setBoundaries() {
	float u0 = fluid_speed;
	for (int x=0; x<xdim; x++) {
		setEquil(x, 0, u0, 0, 1);
		setEquil(x, ydim-1, u0, 0, 1);
	}
	for (int y=1; y<ydim-1; y++) {
		setEquil(0, y, u0, 0, 1);
		setEquil(xdim-1, y, u0, 0, 1);
	}
} // end setBoundaries()

// Collide particles within each cell (here's the physics!):
void collide(void) {
	float viscosity = fluid_viscosity;	// kinematic viscosity coefficient in natural units
	float omega = 1 / (3*viscosity + 0.5);		// reciprocal of relaxation time
	for (int y=1; y<ydim-1; y++) {
		for (int x=1; x<xdim-1; x++) {
			int i = x + y*xdim;		// array index for this lattice site
			float thisrho = n0[i] + nN[i] + nS[i] + nE[i] + nW[i] + nNW[i] + nNE[i] + nSW[i] + nSE[i];
			rho[i] = thisrho;
			float thisux = (nE[i] + nNE[i] + nSE[i] - nW[i] - nNW[i] - nSW[i]) / thisrho;
			ux[i] = thisux;
			float thisuy = (nN[i] + nNE[i] + nNW[i] - nS[i] - nSE[i] - nSW[i]) / thisrho;
			uy[i] = thisuy ;
			float one9thrho = one9th * thisrho;		// pre-compute a bunch of stuff for optimization
			float one36thrho = one36th * thisrho;
			float ux3 = 3 * thisux;
			float uy3 = 3 * thisuy;
			float ux2 = thisux * thisux;
			float uy2 = thisuy * thisuy;
			float uxuy2 = 2 * thisux * thisuy;
			float u2 = ux2 + uy2;
			float u215 = 1.5 * u2;
			//n0[i]  += omega * (four9ths*thisrho * (1                        - u215) - n0[i]);
			nE[i]  += omega * (   one9thrho * (1 + ux3       + 4.5*ux2        - u215) - nE[i]);
			nW[i]  += omega * (   one9thrho * (1 - ux3       + 4.5*ux2        - u215) - nW[i]);
			nN[i]  += omega * (   one9thrho * (1 + uy3       + 4.5*uy2        - u215) - nN[i]);
			nS[i]  += omega * (   one9thrho * (1 - uy3       + 4.5*uy2        - u215) - nS[i]);
			nNE[i] += omega * (  one36thrho * (1 + ux3 + uy3 + 4.5*(u2+uxuy2) - u215) - nNE[i]);
			nSE[i] += omega * (  one36thrho * (1 + ux3 - uy3 + 4.5*(u2-uxuy2) - u215) - nSE[i]);
			nNW[i] += omega * (  one36thrho * (1 - ux3 + uy3 + 4.5*(u2-uxuy2) - u215) - nNW[i]);
			nSW[i] += omega * (  one36thrho * (1 - ux3 - uy3 + 4.5*(u2+uxuy2) - u215) - nSW[i]);
			n0[i]   = thisrho - (nE[i]+nW[i]+nN[i]+nS[i]+nNE[i]+nSE[i]+nNW[i]+nSW[i]);
		}
	}
	for (int y=1; y<ydim-2; y++) {
		nW[xdim-1+y*xdim] = nW[xdim-2+y*xdim];		// at right end, copy left-flowing densities from next row to the left
		nNW[xdim-1+y*xdim] = nNW[xdim-2+y*xdim];
		nSW[xdim-1+y*xdim] = nSW[xdim-2+y*xdim];
	}
} // end collide

// Move particles along their directions of motion:
void stream(void) {
	barrierCount = 0; barrierxSum = 0; barrierySum = 0;
	barrierFx = 0.0; barrierFy = 0.0;
	for (int y=ydim-2; y>0; y--) {			// first start in NW corner...
		for (int x=1; x<xdim-1; x++) {
			nN[x+y*xdim] = nN[x+(y-1)*xdim];			// move the north-moving particles
			nNW[x+y*xdim] = nNW[x+1+(y-1)*xdim];		// and the northwest-moving particles
		}
	}
	for (int y=ydim-2; y>0; y--) {			// now start in NE corner...
		for (int x=xdim-2; x>0; x--) {
			nE[x+y*xdim] = nE[x-1+y*xdim];			// move the east-moving particles
			nNE[x+y*xdim] = nNE[x-1+(y-1)*xdim];		// and the northeast-moving particles
		}
	}
	for (int y=1; y<ydim-1; y++) {			// now start in SE corner...
		for (int x=xdim-2; x>0; x--) {
			nS[x+y*xdim] = nS[x+(y+1)*xdim];			// move the south-moving particles
			nSE[x+y*xdim] = nSE[x-1+(y+1)*xdim];		// and the southeast-moving particles
		}
	}
	for (int y=1; y<ydim-1; y++) {				// now start in the SW corner...
		for (int x=1; x<xdim-1; x++) {
			nW[x+y*xdim] = nW[x+1+y*xdim];			// move the west-moving particles
			nSW[x+y*xdim] = nSW[x+1+(y+1)*xdim];		// and the southwest-moving particles
		}
	}
	for (int y=1; y<ydim-1; y++) {				// Now handle bounce-back from barriers
		for (int x=1; x<xdim-1; x++) {
			if (barrier[x+y*xdim]) {
				int index = x + y*xdim;
				nE[x+1+y*xdim] = nW[index];
				nW[x-1+y*xdim] = nE[index];
				nN[x+(y+1)*xdim] = nS[index];
				nS[x+(y-1)*xdim] = nN[index];
				nNE[x+1+(y+1)*xdim] = nSW[index];
				nNW[x-1+(y+1)*xdim] = nSE[index];
				nSE[x+1+(y-1)*xdim] = nNW[index];
				nSW[x-1+(y-1)*xdim] = nNE[index];
				// Keep track of stuff needed to plot force vector:
				barrierCount++;
				barrierxSum += x;
				barrierySum += y;
				barrierFx += nE[index] + nNE[index] + nSE[index] - nW[index] - nNW[index] - nSW[index];
				barrierFy += nN[index] + nNE[index] + nNW[index] - nS[index] - nSE[index] - nSW[index];
			}
		}
	}
} // end stream

/*
// Move the tracer particles:
void moveTracers(void) {
	for (int t=0; t<nTracers; t++) {
		//float roundedX = Math.round();
		//float roundedY = Math.round();
		int index = tracerX[t] + tracerY[t]*xdim;
		tracerX[t] += ux[index];
		tracerY[t] += uy[index];
		if (tracerX[t] > xdim-1) {
			tracerX[t] = 0;
			//tracerY[t] = Math.random() * ydim;
		}
	}
}
*/

// Simulate function executes a bunch of steps and then schedules another call to itself:
void simulate(void) {
	//
	setBoundaries();	
	// Execute a bunch of time steps:
	for (int step=0; step<stepsPerFrame; step++) {
		collide();
		stream();
		//moveTracers();
		//lastBarrierFy = barrierFy;
	}
	//paintCanvas();

	if (running) {
		stepCount += stepsPerFrame;
	}
	char stable = true;
	for (int x=0; x<xdim; x++) {
		int index = x + (ydim/2)*xdim;	// look at middle row only
		if (rho[index] <= 0) stable = false;
	}
	//if (!stable) {
		//window.alert("The simulation has become unstable due to excessive fluid speeds.");
	//	startStop();
	//	initFluid();
	//}	
}

	/*
	// Initialize the tracer particles:
	function initTracers() {
		if (tracerCheck.checked) {
			float nRows = Math.ceil(Math.sqrt(nTracers));
			float dx = xdim / nRows;
			float dy = ydim / nRows;
			float nextX = dx / 2;
			float nextY = dy / 2;
			for (float t=0; t<nTracers; t++) {
				tracerX[t] = nextX;
				tracerY[t] = nextY;
				nextX += dx;
				if (nextX > xdim) {
					nextX = dx / 2;
					nextY += dy;
				}
			}
		}
		paintCanvas();
	}
*/

/*
	// Compute the curl (actually times 2) of the macroscopic velocity field, for plotting:
	function computeCurl() {
		for (float y=1; y<ydim-1; y++) {			// interior sites only; leave edges set to zero
			for (float x=1; x<xdim-1; x++) {
				curl[x+y*xdim] = uy[x+1+y*xdim] - uy[x-1+y*xdim] - ux[x+(y+1)*xdim] + ux[x+(y-1)*xdim];
			}
		}
	}

	*/


	// Add a barrier at a given grid coordinate location:
	void addBarrier(int x, int y) {
		if ((x > 1) && (x < xdim-2) && (y > 1) && (y < ydim-2)) {
			barrier[x+y*xdim] = true;
		}
	}

	// Remove a barrier at a given grid coordinate location:
	void removeBarrier(int x, int y) {
		if (barrier[x+y*xdim]) {
			barrier[x+y*xdim] = false;
			//paintCanvas();
		}
	}

	// Clear all barriers:
	void clearBarriers(void) {
		for (int x=0; x<xdim; x++) {
			for (int y=0; y<ydim; y++) {
				barrier[x+y*xdim] = false;
			}
		}
		//paintCanvas();
	}

// ==================================================
// === Lattice Boltzmann demo -- RUNNING on core 0
// ==================================================


static PT_THREAD (protothread_graphics(struct pt *pt)) {
    PT_BEGIN(pt);
    // the protothreads interval timer
    PT_INTERVAL_INIT() ;

    // Draw some filled rectangles
    fillRect(64, 0, 176, 50, BLUE); // blue box
    fillRect(250, 0, 176, 50, ORANGE); // red box
    fillRect(435, 0, 176, 50, GREEN); // green box

    // Write some text
    setTextColor(WHITE) ;
    setCursor(65, 0) ;
    setTextSize(1) ;
    writeString("Raspberry Pi Pico") ;
    setCursor(65, 10) ;
    writeString("Graphics primitives demo") ;
    setCursor(65, 20) ;
    writeString("Hunter Adams") ;
    setCursor(65, 30) ;
    writeString("vha3@cornell.edu") ;
    setCursor(445, 10) ;
    setTextColor(BLACK) ;
    setTextSize(1) ;
    writeString("Protothreads rp2040 v1.1.1") ;
    setCursor(445, 20) ;
    writeString("Mod for 16 colors (brl4)") ;
    setTextColor(WHITE) ;

    char video_buffer[32];

    // the Lattice-boltzmaann setup
    init_LB();
    // initialize to steady rightward flow
    initFluid();		
    // put in a vertical barrier
	// about 1/3 of width
	for (int y=10; y<20; y++) {
		addBarrier(10, y) ;
	}
    // 

    while(true) {
      // A brief nap
      PT_YIELD_INTERVAL(30000) ;
      // compute one diffusion iteration 
      int elapsed_time = PT_GET_TIME_usec();

      // main simulation loop 
	  simulate() ;

	  for (int x=0; x<xdim; x++) {
			for (int y=0; y<ydim; y++) {
				short speed = 50*sqrt(ux[x+y*xdim]*ux[x+y*xdim]+uy[x+y*xdim]*uy[x+y*xdim]);
				drawPixel(2*x+100, 2*y+100, speed) ;
				drawPixel(2*x+101, 2*y+100, speed) ;
				drawPixel(2*x+100, 2*y+101, speed) ;
				drawPixel(2*x+101, 2*y+101, speed) ;
			}
		}

      elapsed_time = PT_GET_TIME_usec()-elapsed_time ;
      sprintf(video_buffer, "Frame time= %d mSec \r\n", elapsed_time/1000);
      setCursor(280, 20) ;
      setTextColor2(BLACK, ORANGE) ;
      setTextSize(1) ;
      writeString(video_buffer) ;
   }

   PT_END(pt);
} // graphics thread

// ==================================================
// === toggle25 thread on core 0
// ==================================================
// the on-board LED blinks
static PT_THREAD (protothread_toggle25(struct pt *pt))
{
    PT_BEGIN(pt);
    static bool LED_state = false ;
    
     // set up LED p25 to blink
     gpio_init(25) ;	
     gpio_set_dir(25, GPIO_OUT) ;
     gpio_put(25, true);
     // data structure for interval timer
     PT_INTERVAL_INIT() ;

      while(1) {
        // yield time 0.1 second
        //PT_YIELD_usec(100000) ;
        PT_YIELD_INTERVAL(100000) ;

        // toggle the LED on PICO
        LED_state = LED_state? false : true ;
        gpio_put(25, LED_state);
        //
        // NEVER exit while
      } // END WHILE(1)
  PT_END(pt);
} // blink thread


// ==================================================
// === toggle gpio 4 thread -- RUNNING on core 1
// ==================================================
// toggle gpio 4 
static PT_THREAD (protothread_toggle_gpio4(struct pt *pt))
{
    PT_BEGIN(pt);
    static bool LED_state = false ;
    //
     // set up LED gpio 4 to blink
     gpio_init(4) ;	
     gpio_set_dir(4, GPIO_OUT) ;
     gpio_put(4, true);
     // data structure for interval timer
     PT_INTERVAL_INIT() ;

      while(1) {
        //
        PT_YIELD_INTERVAL(20) ;
        // toggle gpio 4
        LED_state = !LED_state ;
        gpio_put(4, LED_state);
        //
        // NEVER exit while
      } // END WHILE(1)
  PT_END(pt);
} // blink thread

// ==================================================
// === user's serial input thread on core 1
// ==================================================
// serial_read an serial_write do not block any thread
// except this one
static PT_THREAD (protothread_serial(struct pt *pt))
{
    PT_BEGIN(pt);
      static int test_in1, test_in2, sum ;
      //
      while(1) {
        // print prompt
        sprintf(pt_serial_out_buffer, "input x y: ");
        // spawn a thread to do the non-blocking write
        serial_write ;

        // spawn a thread to do the non-blocking serial read
         serial_read ;
        
        // convert input string to number
        //sscanf(pt_serial_in_buffer,"%d %d ",  &col1, &col2) ;
        //printf("color=%d\n\r", readPixel(col1, col2)) ;

        // NEVER exit while
      } // END WHILE(1)
  PT_END(pt);
} // serial thread

// ========================================
// === core 1 main -- started in main below
// ========================================
void core1_main(){ 
  //
  //  === add threads  ====================
  // for core 1
  pt_add_thread(protothread_toggle_gpio4) ;
  //pt_add_thread(protothread_serial) ;
  //
  // === initalize the scheduler ==========
  pt_schedule_start ;
  // NEVER exits
  // ======================================
}

// ========================================
// === core 0 main
// ========================================
int main(){
  // set the clock
  //set_sys_clock_khz(250000, true); // 171us
  // start the serial i/o
  stdio_init_all() ;
  // announce the threader version on system reset
  printf("\n\rProtothreads RP2040 v1.11 two-core\n\r");

  // Initialize the VGA screen
  initVGA() ;
     
  // start core 1 threads
  multicore_reset_core1();
  multicore_launch_core1(&core1_main);

  // === config threads ========================
  // for core 0
  pt_add_thread(protothread_graphics);
  pt_add_thread(protothread_toggle25);
  //
  // === initalize the scheduler ===============
  pt_schedule_start ;
  // NEVER exits
  // ===========================================
} // end main