/*
	Project example - spread of disease in a region

*/
#include "glut.h"
#include <stdlib.h>
#include <stdio.h>
#include <math.h>

#define ROWS 30			//	area is assumed to be ROWSxCOLS (unitless) squares
#define COLS 30
#define P    10000.		//	total population and maximum infections assumed
#define MAXINFECT P/5.
#define MULTIPLIER 10.

#define BOTH
#undef  GEOM
#undef  COLOR

typedef struct population {
		float s, i, r;			// susceptible, immune, recovered
		float aN, aE, aS, aW;	// inwards contact rates to region
	} population;

GLfloat angle = 0.0;
population regions[ROWS][COLS], back[ROWS+2][COLS+2];
GLfloat theta = 0.0, vp = 24.0;
float alpha = 0.0001, beta = 0.15;	// infection and recovery rates, resp

int onceonly = 0;
int winwide = 700, winheight = 500;

void myinit(void);
void cube(void);
void display(void);
void doRasterString( float, float, float, char *);
void reshape(int, int);
void colorRamp(float, float *, float *, float *); 
void animate(void);

void myinit(void)
{
	int i,j;
	
	glEnable (GL_DEPTH_TEST);

	glClearColor(0.6, 0.6, 0.6, 1.0);

	//	set up initial states
	for (i=0; i<ROWS; i++) {
		for (j=0; j < COLS; j++) {
			regions[i][j].s = P;
			regions[i][j].i = regions[i][j].r = 0.;
			regions[i][j].aN = regions[i][j].aE = 
				regions[i][j].aS = regions[i][j].aW = 0.01;
		}
	}
	//	exceptions -- lake in the center of the region and 
	//	initial infection in one region
	for (i = 13; i<18; i++ )
		for (j=13; j<18; j++ ) {
			regions[i][j].s = 0.;
			regions[i][j].aN = regions[i][j].aE = 
				regions[i][j].aS = regions[i][j].aW = 0.0;
	}
	regions[10][10].i = 200.;
	regions[10][10].s = P - 200;
}

//	Unit cube in first octant
void cube (void)
{
	typedef GLfloat point [3];

	point v[8] = {
		{0.0, 0.0, 0.0},
		{0.0, 0.0, 1.0},
		{0.0, 1.0, 0.0},
		{0.0, 1.0, 1.0},
		{1.0, 0.0, 0.0},
		{1.0, 0.0, 1.0},
		{1.0, 1.0, 0.0},
		{1.0, 1.0, 1.0} };

	glBegin (GL_QUAD_STRIP);
		glVertex3fv(v[4]);
		glVertex3fv(v[5]);
		glVertex3fv(v[0]);
		glVertex3fv(v[1]);
		glVertex3fv(v[2]);
		glVertex3fv(v[3]);
		glVertex3fv(v[6]);
		glVertex3fv(v[7]);
	glEnd();

	glBegin (GL_QUAD_STRIP);
		glVertex3fv(v[1]);
		glVertex3fv(v[3]);
		glVertex3fv(v[5]);
		glVertex3fv(v[7]);
		glVertex3fv(v[4]);
		glVertex3fv(v[6]);
		glVertex3fv(v[0]);
		glVertex3fv(v[2]);
	glEnd();
	
#ifdef GEOM
	glColor3f(1.0,1.0,1.0);
	glBegin (GL_LINE_STRIP);
		glVertex3fv(v[1]), glVertex3fv(v[3]);
		glVertex3fv(v[7]), glVertex3fv(v[5]);
		glVertex3fv(v[1]);
	glEnd();
#endif
}

void colorRamp(float yval, float *r, float *g, float *b)
{
	if (yval < 0.30) { *r = yval/0.3; *g = 0.0; *b = 0.0; return; }
	if ((yval >= 0.30) && (yval < 0.89))
		{ *r = 1.0; *g = (yval-0.3)/0.59; *b = 0.0; return; }
	if (yval >= 0.89) { *r = 1.0; *g = 1.0; *b = (yval-0.89)/0.11; }
	return;
}

void display( void )
{
	int i,j;
	float r, g, b, scaledValue;
	char *s;

	glViewport(0,0,(int)(5.*(float)winwide/7.),winheight);
	
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);  

	glMatrixMode(GL_PROJECTION);
	glLoadIdentity();
	gluPerspective(60.0, (5.*(float)winwide/7.)/(float)winheight, 1.0, 300.0);

	glMatrixMode(GL_MODELVIEW);
	glLoadIdentity();
	gluLookAt(vp, vp, vp/2.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0);
    
	glPushMatrix();
    glRotatef(theta, 0.0, 0.0, 1.0);
	//	Draw the bars
	for (i = 0; i < ROWS; i++) {
		for (j = 0; j < COLS; j++) {
#ifndef GEOM
			scaledValue = MULTIPLIER*regions[i][j].i/MAXINFECT;
			colorRamp(scaledValue, &r, &g, &b);
			if ((i > 12)&&(i < 18)&&(j > 12)&&(j < 18)) glColor3f(0.,0.,1.);
			else glColor3f(r, g, b);
#endif
#ifdef GEOM
			glColor3f( 1.0, 0.3, 0.3 );
#endif
			glPushMatrix();
			glTranslatef((float)i-(float)ROWS/2.0,(float)j-(float)COLS/2.0,0.0);
#ifndef COLOR
			glScalef(1.0, 1.0, 0.1+4.0*scaledValue);
#endif
			cube();
			glPopMatrix();
			}
		}
    glPopMatrix();  // undo last rotation you set

//	draw the legend in its own viewport
	glViewport((int)(5.*(float)winwide/7.),0,(int)(2.*(float)winwide/7.),winheight);
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);  
	glMatrixMode(GL_PROJECTION);
	glLoadIdentity();
//	glOrtho(0.,2.,0.,5.,-1.,1.);
	gluPerspective(30.0, (2.*(float)winwide/7.)/(float)winheight, 1.0, 300.0);
	
	glMatrixMode(GL_MODELVIEW);
	glLoadIdentity();
	gluLookAt(1., 2.5, 10., 1., 2.5, 0.0, 0.0, 1.0, 0.0);

	glPushMatrix();
	glEnable (GL_SMOOTH);
	glColor3f(1.,1.,1.);
	doRasterString(0.1, 4.8, 0., "Number Infected");
	sprintf(s,"%5.0f",MAXINFECT/MULTIPLIER);
	doRasterString(0.,4.4,0.,s);
//	color cutoffs are at 0.3 and 0.89
	glBegin(GL_QUADS);
		glColor3f(0.,0.,0.);
		glVertex3f(0.7, 0.1, 0.);
		glVertex3f(1.7, 0.1, 0.);
		colorRamp(0.3, &r, &g, &b);
		glColor3f(r,g,b);
		glVertex3f(1.7, 1.36, 0.);
		glVertex3f(0.7, 1.36, 0.);
		
		glVertex3f(0.7, 1.36, 0.);
		glVertex3f(1.7, 1.36, 0.);
		colorRamp(0.89, &r, &g, &b);
		glColor3f(r,g,b);
		glVertex3f(1.7, 4.105, 0.);
		glVertex3f(0.7, 4.105, 0.);
		
		glVertex3f(0.7, 4.105, 0.);
		glVertex3f(1.7, 4.105, 0.);
		glColor3f(1.,1.,1.);
		glVertex3f(1.7, 4.6, 0.);
		glVertex3f(0.7, 4.6, 0.);
	glEnd();
	sprintf(s,"%5.0f",0.0);
	doRasterString(.1,.1,0.,s);
	glPopMatrix();
	glDisable(GL_SMOOTH);
    
	glutSwapBuffers();
}

void doRasterString( float x, float y, float z, char *s)
{
	char c;
	
	glRasterPos3f(x,y,z);
	for ( ; (c = *s) != '\0'; s++ )
		glutBitmapCharacter(GLUT_BITMAP_TIMES_ROMAN_24, c);
}

void reshape(int w,int h)
{
	winwide = w; winheight = h;
}

void animate(void)
{
	int i, j, ii, jj;
	float newi;
	
	//	update angle for rotating display
	theta += 1.0;
	if (theta > 360.0) theta = theta - 360.0;

	// back up regions[][] to back[][] up so you can recreate regions
	for (i=0; i<ROWS; i++)
		for (j=0; j<COLS; j++)
			back[i+1][j+1] = regions[i][j];
	for (i=1; i<ROWS+2; i++) {
		back[i][0]=back[i][1];
		back[i][COLS+1]=back[i][COLS];
		}
	for (j=0; j<COLS+2; j++) {
		back[0][j] = back[1][j];
		back[ROWS+1][j]=back[ROWS][j];
		}
		
	// use diffusion based on model using back values to compute temp
	for (ii=0; ii<ROWS; ii++)
		for (jj=0; jj<COLS; jj++) {
			i = ii+1;
			j = jj+1;
			newi = alpha*back[i][j].i*back[i][j].s;
			if (i != 0) newi += alpha*(back[i][j].aN*back[i-1][j].i*back[i][j].s);
			if (j != COLS) newi += alpha*(back[i][j].aE*back[i][j+1].i*back[i][j].s);
			if (i != ROWS) newi += alpha*(back[i][j].aS*back[i+1][j].i*back[i][j].s);
			if (j != 0) newi += alpha*(back[i][j].aW*back[i][j-1].i*back[i][j].s);
			regions[ii][jj].r += beta * back[i][j].i;
			regions[ii][jj].i = (1.-beta)*back[i][j].i + newi;
			regions[ii][jj].s = back[i][j].s - newi;
		}

	//	finished updating; now do the display
	glutPostRedisplay(); /* perform display again */
}

void main(int argc, char** argv)
{
/* Standard GLUT initialization */
	glutInit(&argc,argv);
	glutInitDisplayMode (GLUT_DOUBLE | GLUT_RGB | GLUT_DEPTH); 
	glutInitWindowSize(winwide, winheight);
	glutInitWindowPosition(50,50);
	glutCreateWindow("Spread of disease in geographic region");
	glutDisplayFunc(display);
	glutReshapeFunc(reshape);
	glutIdleFunc(animate);

	myinit();
	glutMainLoop(); /* enter event loop */
}