/*
A GLUT program using OpenGL, GLUI, and GLSL.

It uses the Newton physics API:
	http://www.newtondynamics.com/
	http://walaber.com/newton_wiki/index.php?title=Main_Page


It simulates some simple square-wheeled robots.

Orion Sky Lawlor, olawlor@acm.org, 2008/09/28 (Public Domain)
*/
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <math.h>

#include <GL/glew.h> /* GL Extensions Wrangler, http://glew.sourceforge.net/ */
#include <GL/glut.h> /* GL Utilities Toolkit, http://freeglut.sourceforge.net/ */
#include <GL/glui.h> /* GL User Interface, http://glui.sourceforge.net/ */
#include "osl/vec4.h" /* GLSL-style C++ types, by Orion Lawlor */
#include "osl/mat4.h"
using osl::mat4;

#include "osl/random.h"
#include "osl/random_vector.h"
#include "osl/random.cpp"

/* Dr. Lawlor's OpenGL utilities */
#include "ogl/util.cpp"
#include "ogl/fast_mipmaps.c"

/* Newton Physics Engine: http://www.newtondynamics.com/ */
#define _NEWTON_USE_LIB 1 /* I want the statically linked library, no DLLs! */
#include "Newton.h"


/* Variables updated by the GUI */
float distance=2.5; /* robot-to-camera distance */
int win_w, win_h; /**< Width and height (in pixels) of our window at last resize */

float gravity=9.8; /* gravitational accelleration, m/s^2 */
float torque=5.0; /* overall torque to apply to robot wheels, newton-meters */

int slowmo=0; /* slow motion enabled? */


/* Array indicating which keyboard keys are currently held down */
int keyDown[256]={0};
void key_down(unsigned char c,int x,int y) {keyDown[c]=1;}
void key_up(unsigned char c,int x,int y) {keyDown[c]=0;}

/* These are rotation matrices, updated by the GLUI_Rotation spheres */
mat4 modelMat(1.0), viewMat(1.0);

/// Scale by this offset
inline void glScalefv(const float *s) {
	glScalef(s[0],s[1],s[2]);
}
/// Translate by this offset
inline void glTranslatefv(const float *s) {
	glTranslatef(s[0],s[1],s[2]);
}
/** Transpose this mat4.  If the mat4 was a rotation matrix, 
  this actually inverts the matrix. */
mat4 transpose(const mat4 &m) {
	return mat4(
		m[0][0], m[1][0], m[2][0], m[3][0],
		m[0][1], m[1][1], m[2][1], m[3][1],
		m[0][2], m[1][2], m[2][2], m[3][2],
		m[0][3], m[1][3], m[2][3], m[3][3]);
}


/********* Robots ***************/

class RobotPhysics {
public:
	RobotPhysics(NewtonWorld *world_,vec3 location);
	void draw(void);
	vec3 color;
	/* List of simulated bodies.  First four are wheels */
	enum {lb=0,lf=1,rb=2,rf=3,chassis=4,nbodies=5};
	/* List of net torques that motors apply to each wheel */
	float wheelTorques[4];
	
	vec3 getLocation(void) const {
		mat4 m;
		NewtonBodyGetMatrix(bodies[chassis],&m[0][0]);
		return vec3(m[3]);
	}
	
	int selected; /* UI state--is checkbox checked? */
	
private:
	NewtonWorld *world;
	NewtonBody *bodies[nbodies];
	vec3 bodysize[nbodies];
	
	struct wheelInfo {
		RobotPhysics *robot;
		int b; /* body number of our wheel */
		wheelInfo(RobotPhysics *r_,int b_) :robot(r_), b(b_) {}
	};
	
	static void wheelForceAndTorque (const NewtonBody* body);
	static void chassisForceAndTorque (const NewtonBody* body);
};

/* Set up the geometry of a robot simulation */
RobotPhysics::RobotPhysics(NewtonWorld *world_,vec3 location)
	:selected(1), world(world_)
{
	for (int w=0;w<4;w++) wheelTorques[w]=0.1;
	color=vec3(0.8,0.3,0.8);
	struct bodyInfo {
		float loc[3];
		float sz[3];
	};
	const static float halftrack=0.3; /* x distance between tires (track) / 2 */
	const static float halfbase=0.35; /* y distance between tires (wheelbase) /2 */
	const static float wheelsize=0.2; /* diameter of drive wheels */
	const static float wheelthick=0.08; /* thickness of drive wheels */
	const static float wheelht=-0.1; /* location of wheel axles, relative to chassis COM (z) */
	const static float chassislen=2.4*halfbase; /* chassis along drive axis (y) */
	const static float chassiswid=2*(halftrack-0.5*wheelthick-0.01); /* chassis along axle axis (x) */
	const static float chassisht=0.25; /* chassis along vertical axis (z) */
	const static bodyInfo bi[nbodies]={
		{{-halftrack,-0.5*chassislen,wheelht}, {wheelthick,wheelsize,wheelsize}}, /* lb wheel */
		{{-halftrack,+0.5*chassislen,wheelht}, {wheelthick,wheelsize,wheelsize}}, /* lf wheel */
		{{+halftrack,-0.5*chassislen,wheelht}, {wheelthick,wheelsize,wheelsize}}, /* rb wheel */
		{{+halftrack,+0.5*chassislen,wheelht}, {wheelthick,wheelsize,wheelsize}}, /* lf wheel */
		
		{{0,0,0}, {chassiswid,chassislen,chassisht}} /* main chassis */
	};

#define ARR_TO_XYZ(arr) (arr)[0],(arr)[1],(arr)[2]
	for (int b=nbodies-1;b>=0;b--) {
		bodysize[b]=bi[b].sz;
		NewtonCollision *c=NewtonCreateBox(world,ARR_TO_XYZ(bi[b].sz),NULL);
		NewtonBody *body=NewtonCreateBody(world,c);
		bodies[b]=body;
		NewtonReleaseCollision(world,c);
		NewtonBodySetUserData(body,new wheelInfo(this,b));
		
		double mass=40.0, rotNertia=1.0;
		if (b<chassis) {mass=1.0; rotNertia=1.0e-3;} /* wheels have lower mass and inertia */
		/* FIXME: figure out real masses and rotational inertia */
		NewtonBodySetMassMatrix(body,1,rotNertia,rotNertia,rotNertia);
		NewtonBodySetForceAndTorqueCallback(body,
			(b==chassis)?chassisForceAndTorque:wheelForceAndTorque);
		
		vec3 loc=location+bi[b].loc; /* center of our body: robot plus local origin */
		mat4 m(1.0);
		m[3][0]=loc[0];
		m[3][1]=loc[1];
		m[3][2]=loc[2];
		NewtonBodySetMatrix(body,&m[0][0]);
		
		if (b<chassis) { // Each wheel is stuck to the chassis with a hinge joint (axle)
			vec3 pinDir(1,0,0); /* axles rotate along the X axis */
			NewtonConstraintCreateHinge(world,
				loc,pinDir,body,bodies[chassis]);
		}
	}
}

/* Render the simulated robot */
void RobotPhysics::draw(void)
{
	for (int b=0;b<nbodies;b++) {
		glPushMatrix();

		mat4 m(1.0);
		NewtonBodyGetMatrix(bodies[b],&m[0][0]);
		glMultMatrixf(&m[0][0]);
		glScalef(ARR_TO_XYZ(bodysize[b]));
		
		if (b!=chassis) glColor3fv(color);  /* wheels */
		else { /* chassis */
			if (selected)
				glColor3f(0,1,1); /* bright green */
			else
				glColor3f(0.4,0.2,0.2); /* dark red */
		}
		glutSolidCube(1.0);
		if (b==chassis) { /* fake a LIDAR scanner on front end of chassis */
			glTranslatef(0.0f,0.4f,0.5f);
			glScalef(0.8,0.3,0.6);
			glutSolidCube(1.0);
		}

		glPopMatrix();
	}
}

/* Apply traction to the robot's wheels */
void RobotPhysics::wheelForceAndTorque (const NewtonBody* body)
{	
	wheelInfo *w=(wheelInfo *)NewtonBodyGetUserData(body);
	
	float t=w->robot->wheelTorques[w->b];

	if (1) { /* Scale down the applied torque at high velocity;
		this keeps the wheel speeds reasonable */
		vec3 vel; /* angular velocity vector (omega) */
		NewtonBodyGetOmega(body,vel);
		float maxvel=3.0; /* maximum angular velocity to allow (radians/sec) */
		float tdown=(1.0-length(vel)/maxvel); /* stop applying torque once wheel's spinning fast */
		if (tdown<0.0) tdown=0.0; /* never reverse applied torque */
		t*=tdown; 
	}
	
	/* 3D torque vector points along our local axle axis, x */
	vec3 torque(-t,0.0,0.0);
	
	/* apply body's rotation matrix to our torque vector to get world coordinates */
	mat4 m;
	NewtonBodyGetMatrix(body,&m[0][0]);
	torque=vec3(m*vec4(torque,0.0));
	NewtonBodySetTorque(body,torque);
}



void standardPhysics(const NewtonBody *body)
{
	float mass,Ixx,Iyy,Izz;

	NewtonBodyGetMassMatrix (body, &mass, &Ixx, &Iyy, &Izz);
	vec3 force (0.0f, 0.0f, -mass * gravity); /* gravity points downward, force proportional to mass */
	NewtonBodySetForce (body, force);
}
void RobotPhysics::chassisForceAndTorque (const NewtonBody* body)
{
	standardPhysics(body);
}



/******************* Heightfield ************************/

class HeightField {
public:
	NewtonWorld *world;
	NewtonBody *body;
	HeightField(NewtonWorld *world, 
		int w_,int h_,
		vec3 scale_,vec3 offset_);
	
	void draw(void);
	
	double terrainHeight(const vec3 &loc) {
		int x=(int)((loc.x-offset.x)/scale.x+0.5);
		int y=(int)((loc.y-offset.y)/scale.y+0.5);
		int i=y*w+x;
		if (i<0 || i>w*h) return 0.0; /* out of bounds */
		else return heights[i];
	}
private:
	/* Size of area, in pixels */
	int w,h;
	vec3 scale, offset;
	/* Array of w x h height values */
	std::vector<float> heights;
	/* Return 3D location of this 2D pixel. */
	vec3 loc(int x,int y) {
		return vec3(x,y,heights[y*w+x])*scale+offset;
	}
	
	/* Add a face to this collision list */
	void addFace(NewtonCollision *col,const vec3 &a,const vec3 &b,const vec3 &c)
	{
		vec3 va[3]; va[0]=a; va[1]=b; va[2]=c;
		NewtonTreeCollisionAddFace(col,3,(const dFloat*)va,sizeof(vec3),0);
	}
	void drawFace(const vec3 &a,const vec3 &b,const vec3 &c)
	{
		glVertex3fv(a); glVertex3fv(b); glVertex3fv(c); 
	}
};

HeightField::HeightField(NewtonWorld *world_, 
		int w_,int h_,
		vec3 scale_,vec3 offset_)
	:world(world_), w(w_), h(h_), scale(scale_), offset(offset_)
{
	int x,y;
	//heights.resize(w*h,0.0);
	for (y=0;y<h;y++) 
	for (x=0;x<w;x++) {
		double z=0.4*sin(x*0.2)*sin(y*0.1);
		heights.push_back(z);
	}
	
	double start=oglTimer();
	printf("Building collision tree.\n");
	NewtonCollision *c=NewtonCreateTreeCollision(world,NULL);
	NewtonTreeCollisionBeginBuild(c);
	
	for (y=0;y<h-1;y++) 
	for (x=0;x<w-1;x++) {
		addFace(c,loc(x,y),loc(x+1,y),loc(x,y+1));
		addFace(c,loc(x+1,y),loc(x+1,y+1),loc(x,y+1));
	}
	
	printf("Done adding faces to collision tree; optimizing\n");
	NewtonTreeCollisionEndBuild(c,1);
	double elapsed=oglTimer()-start;
	printf("Done: %.3f s\n",elapsed);
	
	body=NewtonCreateBody(world,c);
}
	
void HeightField::draw(void)
{
	int x,y;
	glBegin(GL_TRIANGLES);
	for (y=0;y<h-1;y++) 
	for (x=0;x<w-1;x++) {
		vec3 a(loc(x,y)), b(loc(x+1,y)), c(loc(x,y+1));
		glNormal3fv((b-a).cross(c-a).dir());
		
		float g=0.7, d=0.3;
		if ((1&(int)(a.x))^(1&(int)(a.y)) == 0) /* meter-size checkerboard pattern */
			glColor3f(g,g,g);
		else
			glColor3f(d,d,d);
		
		drawFace(a,b,c);
		drawFace(b,loc(x+1,y+1),c);
	}
	glEnd();
}

/*************** Rocks **************/
class Rock {
public:
	Rock(NewtonWorld *world_,osl::Random &r,HeightField *hf);
	void draw(void);
	NewtonWorld *world;
	NewtonBody *body;
	double sz; /* diameter of our object */
	bool isSphere;
	vec3 color;
	
	// add our forces and torques to rigid body
	static void PhysicsForceAndTorque (const NewtonBody* body);
};

Rock::Rock(NewtonWorld *world_,osl::Random &r,HeightField *hf)
	:world(world_)
{
	sz=0.01+0.2*r.nextFloat();
	NewtonCollision *c;
	if (r.nextInt(3)!=0) {
		c=NewtonCreateBox(world,sz,sz,sz,NULL);
		isSphere=false;
	} else {
		c=NewtonCreateSphere(world,sz,sz,sz,NULL);
		isSphere=true;
	}
	body=NewtonCreateBody(world,c);
	NewtonReleaseCollision(world,c);
	
	double rotNertia=0.01;
	NewtonBodySetMassMatrix(body,1,rotNertia,rotNertia,rotNertia);
	NewtonBodySetForceAndTorqueCallback(body,standardPhysics);
	
	mat4 m(1.0);
	vec3 v(r.nextFloat(),r.nextFloat(),r.nextFloat());
	v=(v-vec3(0.5f))*6+vec3(0,0,1);
	v.z=1.3*sz+hf->terrainHeight(v); /* don't exist below ground */
	m[3][0]=v.x; /* spread objects apart */
	m[3][1]=v.y;
	m[3][2]=v.z;
	NewtonBodySetMatrix(body,&m[0][0]);
	
	// Initial angular velocity
	vec3 angVel(r.nextFloat(),r.nextFloat(),r.nextFloat());
	NewtonBodySetOmega(body,3*(angVel-vec3(0.5)));
	
	color=vec3(r.nextFloat(),r.nextFloat(),r.nextFloat());
}
void Rock::draw(void)
{
	glPushMatrix();
	
	mat4 m(1.0);
	NewtonBodyGetMatrix(body,&m[0][0]);
	glMultMatrixf(&m[0][0]);
	
	glColor3fv(color);
	if (isSphere) glutSolidSphere(sz,6,5);
	else glutSolidCube(sz);
	
	glPopMatrix();
}

/* A set of interacting RobotPhysics and Rock objects */
class PhysicsWorld {
public:
	NewtonWorld *world;
	HeightField *hf;
	std::vector<RobotPhysics *> robots;
	std::vector<Rock *> rocks;
	
	
	PhysicsWorld();
	void draw(void);
};

PhysicsWorld::PhysicsWorld()
{
	/* Fire up the physics */
	world = NewtonCreate(NULL,NULL);
	
	// Make the terrain
	int w=100, h=100;  vec3 scale(10.0,10.0,1.0);
	scale*=(1.0/w); scale.z=1.0;
	hf=new HeightField(world,w,h,scale,-vec3(w,h,0)*scale*0.5);

	/* Add some robots and rocks */
	osl::Random48 r(1);
	int i;
	for (i=0;i<3;i++) {
		vec3 loc(4*r.nextFloat()-2,4*r.nextFloat()-4.0,0.0);
		loc.z=0.3+hf->terrainHeight(loc);
		robots.push_back(new RobotPhysics(world,loc));
	}
	for (i=0;i<10;i++) {
		rocks.push_back(new Rock(world,r,hf));
	}
}
void PhysicsWorld::draw(void)
{
	hf->draw();
	for (unsigned int i=0;i<rocks.size();i++)
		rocks[i]->draw();
	for (unsigned int r=0;r<robots.size();r++)
		robots[r]->draw();
}


PhysicsWorld *physicsworld=0; /*<- silly global variable */
GLUI *glui;  /* user interface */

/* GLUT calls this routine every time our window needs to be redrawn */
void main_display(void) {
	/* Make sure there's a world out there */
	if (physicsworld==0) {
		physicsworld=new PhysicsWorld();
		for (unsigned int r=0;r<physicsworld->robots.size();r++)
		{ /* stupid way to control multiple robots at once */
			char title[100]; sprintf(title,"Control Robot %d",r);
			new GLUI_Checkbox(glui,title,&physicsworld->robots[r]->selected);
		}
	}

	static double last_elapsed=0.0;
	glFinish(); double elapsed=oglTimer();
	// Clear the color and depth buffers
	glClearColor(0.2,0.4,0.8,0.0); /* sky blue */
	glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT); /* actually clear color and depth */
	
	// Load all needed matrices into OpenGL
	glMatrixMode(GL_PROJECTION);
	glLoadIdentity(); /*<- clean out any old leftover matrices */
	gluPerspective(90.0, /* <- Y camera field-of-view, in degrees */
		win_w/(float)win_h, /* viewport's aspect ratio: always width over height */
		0.01, /* Near clipping plane depth */
		100.0 /* Far clipping plane depth */
	);
	glTranslatef(0,0,-distance); /* push world origin down the *unrotated* Z axis */
	glMultMatrixf(&viewMat[0][0]); /* camera orientation: could also go under GL_MODELVIEW */
	
	glMatrixMode(GL_MODELVIEW);
	glLoadIdentity();
	glMultMatrixf(&modelMat[0][0]);
	glTranslatef(ARR_TO_XYZ(-physicsworld->robots[0]->getLocation())); // center view on robot 0

	oglDirectionLight(GL_LIGHT0,oglVector3d(0.3,-0.2,1.0),
		oglVector3d(1,1,1),0.1,0.9,0.1);
	glEnable(GL_LIGHTING);
	glEnable(GL_DEPTH_TEST);
	glEnable(GL_COLOR_MATERIAL);
	
	
	
	/* User interface for robot driving */
	float forward=0.0, right=0.0;
	if (keyDown[(int)'w']) forward+=1.0;
	if (keyDown[(int)'s']) forward-=1.0;
	if (keyDown[(int)'a']) right-=1.0;
	if (keyDown[(int)'d']) right+=1.0;
	
	for (unsigned int r=0;r<physicsworld->robots.size();r++)
	{
		RobotPhysics *robot=physicsworld->robots[r];
		if (robot->selected) {
			robot->wheelTorques[0]=torque*(forward+right);
			robot->wheelTorques[1]=torque*(forward+right);
			robot->wheelTorques[2]=torque*(forward-right); // to turn right, right wheels go backwards
			robot->wheelTorques[3]=torque*(forward-right);
		}
	}
	
	physicsworld->draw();
	
	
	glFinish(); elapsed=oglTimer()-elapsed; if (elapsed<=1.0e-3) elapsed=1.0e-3;
	last_elapsed=elapsed;
	
	double slowmo_scale=1.0;
	if (slowmo) slowmo_scale=0.001;
	static double newtime_last=0.0;
	double newtime=oglTimer();
	NewtonUpdate(physicsworld->world, (elapsed+newtime_last)*slowmo_scale);
	newtime=oglTimer()-newtime;
	newtime_last=newtime;
	
	if (1) /* print frames per second */
	{
		/* overwrite all matrices with identity-- back to screen coordinates */
		glMatrixMode(GL_PROJECTION); glLoadIdentity();
		glMatrixMode(GL_MODELVIEW); glLoadIdentity(); 
		glColor4f(1,1,1,1);
		oglPrintf(oglVector3d(-0.9,-0.9,0.0),"%.2f ms/frame draw; %.2f ms/frame physics; %4.2f fps\n",
			1.0e3*elapsed,1.0e3*newtime,1.0/(elapsed+newtime));
	}
	// Display the stuff drawn so far to the screen:
	glutSwapBuffers();
	// Repeatedly ask to be redrawn
	glutPostRedisplay();
}

/* GLUT calls this routine every time our window changes size */
void main_reshape(int w, int h) {
	win_w=w; win_h=h;
	glViewport(0, 0, (GLsizei)win_w, (GLsizei)(win_h));
}

int main(int argc, char** argv)
{
	glutInitWindowSize(800,600);
	glutInit(&argc,argv);

	// Ask for a back buffer, color buffer, and Z buffer in a window.
	int mode=GLUT_DOUBLE | GLUT_RGBA | GLUT_DEPTH;
	glutInitDisplayMode (mode);
	int main_window = glutCreateWindow("UAF RobotSim 2008-09-28: Try the WASD keys!");

	// Turn on OpenGL extensions, like programmable shaders
	//   WARNING: you MUST call glutCreateWindow before doing this!
	//   If you don't call glewInit, or call it wrong, you'll segfault!
	glewInit();
	
	// Ask GLUT to call us when we need to redraw
	glutDisplayFunc (main_display);
	glutReshapeFunc (main_reshape);
	glutKeyboardFunc(key_down);
	glutKeyboardUpFunc(key_up);
	
	/* Set up view matrix so we're looking down Z axis properly */
	vec3 up=normalize(vec3(0.0,0.3,-1.0));
	viewMat[1]=vec4(up,0.0);
	viewMat[2]=vec4(vec3(up.x,-up.z,up.y),0.0);
	
	if (1) {
		// Make some sliders (this is the only GLUI code you'll need!)
		glui = GLUI_Master.create_glui_subwindow(main_window, 
			GLUI_SUBWINDOW_TOP);
		/* This causes the main window to be redrawn when a GLUI control changes */
		glui->set_main_gfx_window( main_window );

		(new GLUI_Spinner( glui, "Distance:", &distance))->set_float_limits(0.0, 10.0);
		new GLUI_Column(glui);
		(new GLUI_Spinner( glui, "Gravity:", &gravity))->set_float_limits(-0.1,100.0);
		(new GLUI_Spinner( glui, "Torque:", &torque))->set_float_limits(-1.0e1,1.0e1);
		new GLUI_Checkbox(glui,"Slow-mo",&slowmo);
		
		new GLUI_Column(glui);
		new GLUI_Rotation(glui,"Model:",&modelMat[0][0]);
		new GLUI_Column(glui,false);
		new GLUI_Rotation(glui,"View:",&viewMat[0][0]);
		new GLUI_Column(glui,true);
	}
	
	// Process events until user closes our window...
	glutMainLoop();
	
	return 0;
}
