/* Author: Jonathan Brumberg
   Date  : April 6, 2006
   File  : p3.cpp
*/

/* Main file for application of Photon Map and Ray Tracer to rendering
   a Cornell Box via global illumination */


#include<cstdio>
#include<cstdlib>
#include<cmath>
#include<ctime>
#include<GL/glut.h>
#include<GL/glu.h>
#include<GL/glext.h>

#include "RtObject.h"
#include "Intersection.h"
#include "Box.h"
#include "BoundingBox.h"
#include "Sphere.h"
#include "BoxIntersection.h"
#include "SphereIntersection.h"
#include "Vector.h"
#include "Ray.h"
#include "PhotonMap.h"
#include "jpegUI.h"

#define WIDTH 600
#define HEIGHT 600

#define IM_WIDTH 400
#define IM_HEIGHT 400
#define NUM_PHOTONS 50000
#define DRAW_PHOT_MAP 1
#define MAX_PHOTON_DEPTH 5
#define MAX_PHOTONS NUM_PHOTONS*MAX_PHOTON_DEPTH
#define DEBUG 0

#define NUM_RAYS IM_WIDTH*IM_HEIGHT

float rtImg[IM_HEIGHT][IM_WIDTH][3];   // Ray trace rendering buffer
float photImg[IM_HEIGHT][IM_WIDTH][3]; // Photon map rendering buffer

int rendWindow;
int photWindow;
int rtWindow;

Vector Isource(0.6,0.6,0.6); // light source for just Ray Tracing
Vector light(0.5,1.0,-1.5);  // light source center position

Ray eyeRays[NUM_RAYS];       // Array of eye rays for ray tracing
BoundingBox boundVol;        // Bounding volume
Box boxObj;                  // Box scene object
Sphere sphereObj;            // Sphere scene object
PhotonMap pm(MAX_PHOTONS);   // Photon map data structure
Ray photRays[NUM_PHOTONS];   // Photon rays used to trace photon paths

void tracePhoton(Ray pr, Vector power, int depth);
void traceRay(Ray r);

/*--------------------------------------------------*/
void initPhotonRays(void) {
  int i=0;
  float x,y,z;
  Vector power(1.0,1.0,1.0);
  Vector offset;
 
  printf("Tracing Photons\n");
  for(i=0;i<NUM_PHOTONS;i++) {
    /* Uniformly distribute photon origins within confines of square
       light source */
    x=0.2*rand()/RAND_MAX-0.1;
    z=0.2*rand()/RAND_MAX-0.1;
    y=0;
    offset.setVector(x,y,z);
    photRays[i].setOrigin(light+offset);
    
    /* determine initial direction: photons always in -y direction,
       remaining directions obtained via rejection sampling */
    do {
      x=(1.0*rand()/RAND_MAX-0.5)*2;
      y=-1.0*rand()/RAND_MAX;
      z=(1.0*rand()/RAND_MAX-0.5)*2;
    } while(x*x+y*y+z*z>1 || y==0);
    
    photRays[i].setDirection(x,y,z);
    photRays[i].normalize();
    /* store photons are origin for light source rendering */
    pm.store(power.getVector(),
	     photRays[i].getOrigin().getVector(),
	     photRays[i].getDirection().getVector());    
    /* trace photons through scene */
    tracePhoton(photRays[i],power,0);
  }
  /* photon map, kd-tree must be balanced before use */
  pm.balance();
  pm.scale_photon_power(1.0/NUM_PHOTONS);
  if(DEBUG) {
    printf("\n");
    printf("Balanced!!!\n");
  }
}

/*--------------------------------------------------*/
Vector diffuseReflection(Vector N) {
  float x,y,z;
  Vector dir;
  
  /* reject any direction that is opposite to the surface normal */
  do {
    /* determine new direction by rejection sampling */
    do {
      x=1.0*rand()/RAND_MAX;
      y=1.0*rand()/RAND_MAX;
      z=1.0*rand()/RAND_MAX;
    } while(x*x+y*y+z*z>1);

    x=2*(x-0.5);
    y=2*(y-0.5);
    z=2*(z-0.5);
    dir.setVector(x,y,z);
  } while(dir*N>=0);
  return dir;
}

/*--------------------------------------------------*/
void tracePhoton(Ray pr, Vector power, int depth) {
  Vector reflect_power;          // reflected photon power
  float kd_avg, k_rand;          // average diffuse reflectance
  Vector kd, N;                  /* local copies of normal vectors and
				    reflectance */
  BoxIntersection boundInt;      /* bounding volume intersection
				    storage */
  BoxIntersection boxInt;        /* box object intersection storage */
  SphereIntersection sphereInt;  /* sphere object intersection storage */

  float t;
  Vector Rt;

  if(DEBUG) {
    printf("depth: %d\n",depth);
  }

  /* Test for maximum reflection depth */
  if(depth<MAX_PHOTON_DEPTH) {
    /* if photon intersects either scene objects (not bounding volume,
       later... */
    if(pr.intersect(&boxObj, &boxInt) || pr.intersect(&sphereObj, &sphereInt)) {
      /* if box intersection, copy important information */
      if(boxInt.getStatus()==HIT) {
	kd=boxObj.getKDiffuse();
	t=boxInt.getTHit();
	N=boxObj.getNormal(boxInt.getFace());
      }
      /* if sphere,... */
      else if(sphereInt.getStatus()==HIT) {
	kd=sphereObj.getKDiffuse();
	t=sphereInt.getTHit();
	N=sphereInt.getNormal();
      }
    }
    /* If the photon does not intersect either scene object, it must
       intersect the bounding volume */
    else {
      pr.intersect(&boundVol,&boundInt);
      kd=boundVol.getKDiffuse(boundInt.getFace());
      t=boundInt.getTExit();
      N=boundVol.getNormal(boundInt.getFace());
    }
      /* calculate average reflectance and random variable for use in
	 determine absorbtion or reflection */
    kd_avg=(kd.x[0]+kd.x[1]+kd.x[2])/3.0;
    k_rand=1.0*rand()/RAND_MAX;
    if(DEBUG) {
      printf("k_rand <? kd_avg:\t%1.2f <? %1.2f\n",k_rand,kd_avg);
    }
    
    if(k_rand<kd_avg) {
      /* store and reflect */
      Rt=pr.cast(t);                    /* cast ray to intersection
					   point */
      Vector dir=diffuseReflection(N);  /* calculate new direction */
      if(DEBUG) {
	pr.print();
	power.print();
      }
      /* store photon location, direction and power at intersection */
      pm.store(power.getVector(),
	       pr.getOrigin().getVector(),
	       pr.getDirection().getVector());
      if(DEBUG) {
	printf("Store photon\n");
	printf("----- ------\n");
      }
      /* initialize reflected photon ray to new origin and direction */
      pr.setOrigin(Rt);
      pr.setDirection(dir);
      pr.normalize();
      if(DEBUG) {
	pr.print();
      }
      /* calculate reflected power */
      reflect_power=power.multScalar(kd*(1.0/kd_avg));
      /* recursively trace photon */
      tracePhoton(pr,reflect_power,depth+1);
    }
    else {
      /*absorb*/
      return;
    }
  }
}

/*--------------------------------------------------*/
void rendPhotMap(void) {
  /* Direct rendering of the photon map, no tree lookup */
  int i=0;
  float *x,*c;
  Photon* photons;

  glMatrixMode(GL_PROJECTION);
  glLoadIdentity();
  gluPerspective(50,1.0,1,100);
  glMatrixMode(GL_MODELVIEW);
  glLoadIdentity();
  gluLookAt(0.5,0.5,0.0,
	    0.5,0.5,-1.0,
	    0.0,1.0,0.0);

  glClearColor(0.0,0.0,0.0,0.0);
  glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

  photons=pm.getPhotons();

  if(DRAW_PHOT_MAP) {
    glBegin(GL_POINTS);
    for(i=0;i<pm.get_stored_photons();i++) {
      x=photons[i].pos;
      c=photons[i].power;
      glColor4f(c[0]*NUM_PHOTONS,c[1]*NUM_PHOTONS,c[2]*NUM_PHOTONS,1.0);
      glVertex3f(x[0],x[1],x[2]);
    }
    glEnd();
  }
  glutSwapBuffers();
}

/*--------------------------------------------------*/
void initImage(void) {
  /* initialize photon map rendering buffer, simple ray trace
     rendering buffer */
  int i,j;

  for(i=0;i<IM_HEIGHT;i++) {
    for(j=0;j<IM_WIDTH;j++) {
      rtImg[i][j][0]=0.0;
      rtImg[i][j][1]=0.0;
      rtImg[i][j][2]=0.0;

      photImg[i][j][0]=0.0;
      photImg[i][j][1]=0.0;
      photImg[i][j][2]=0.0;
    }
  }
}

/*--------------------------------------------------*/
void initVars(void) {
  int i,j;
  float pixSize[2];

  srand(time(NULL));   /* initialize random number generator */
  initImage();         /* initialize rendering buffers */
  pixSize[0]=1.0/IM_WIDTH;
  pixSize[1]=1.0/IM_HEIGHT;

  /* initialize eye-ray origin and direction to completely sample
     viewing plane */
  for(i=0;i<IM_HEIGHT;i++) {
    for(j=0;j<IM_WIDTH;j++) {
      eyeRays[i*IM_HEIGHT+j].setOrigin(0.5,0.5,0);
      eyeRays[i*IM_HEIGHT+j].setDirection(1.0*i*pixSize[0]-0.5+pixSize[0]/2,
					  1.0*j*pixSize[1]-0.5+pixSize[1]/2,
					  -1.0);
      eyeRays[i*IM_HEIGHT+j].normalize();
    }
  }

  /* set bounding volume characteristics */
  boundVol.setOrigin(0.0,0.0,-1.0);
  boundVol.setExtant(1.0,1.0,-2.0);
  for(i=0;i<6;i++) {
    if(i==0) {
      boundVol.setKAmbient(0.5,0.05,0.05,i);
      boundVol.setKDiffuse(0.8,0.2,0.2,i);
    }
    else if(i==1) {
      boundVol.setKAmbient(0.05,0.5,0.05,i);
      boundVol.setKDiffuse(0.2,0.8,0.2,i);
    }
    else if(i==5) {
      boundVol.setKAmbient(0,0,0,i);
      boundVol.setKDiffuse(0,0,0,i);
    }
    else {
      boundVol.setKAmbient(0.5,0.5,0.5,i);
      boundVol.setKDiffuse(238.0f/255,232.0f/255,170.0f/255,i);
    }
  }

  /* Box and Sphere */
  boxObj.setOrigin(0.7,0.0,-1.1);
  boxObj.setExtant(0.9,0.3,-1.3);
  boxObj.setKAmbient(0.2,0.2,1.0);
  boxObj.setKDiffuse(0.7,0.7,0.4);

  sphereObj.setRadius(0.1);
  sphereObj.setCenter(0.3,0.1,-1.4);
  sphereObj.setKAmbient(0.7,0.7,0.2);
  sphereObj.setKDiffuse(0.3,0.9,0.3);

  /* initialize photon rays, photon map, and trace the photons */
  initPhotonRays();
  
  /* trace eye rays through scene...backwards ray tracing */
  printf("Tracing Rays (calculating radiance estimates\n");
  for(i=0;i<NUM_RAYS;i++) {
    traceRay(eyeRays[i]);
  }
}

/*------------------------------------------------*/
Vector calcDiffuseIndirect(Ray L, Vector N) {
  /* photon map radiance calculation */
  Vector c;
  float irrad[3];

  /* calculate radiance estimate via photon map */
  pm.irradiance_estimate(irrad,L.getOrigin().getVector(),
			 N.getVector(),0.5,100);
  c.setVector(irrad);
  return c;
}

/*------------------------------------------------*/
Vector calcDiffuseDirect(Vector kd, Ray L, Vector N) {
  /* ray trace radiance calculation */
  Vector c;

  c=kd.multScalar(Isource*(N*L.getDirection()));
  
  return c;
}

// /*------------------------------------------------*/
// Vector calcAmbientI(Vector ka) {
//   Vector c;

//   c=ka.multScalar(Ia);
//   return c;
// }

// /*------------------------------------------------*/
// Vector shadePixel(Vector ka, Vector kd, Ray L, Vector N) {
//   Vector c,ca,cd;

//   //  ca=calcAmbientI(ka);
//   cd=calcDiffuseDirect(kd,L,N);

//   c=ca+cd;
//   return c;
// }

/*------------------------------------------------*/
Ray calcLightRay(Ray r, float t) {
  /* calculate ray from intersection point to light source for ray
     tracing */
  Vector Rt, D;
  Ray L;

  Rt=r.cast(t);
  D=light-Rt;
  L.setOrigin(Rt);
  L.setDirection(D);
  L.normalize();
  return L;
}

/*------------------------------------------------*/
void traceRay(Ray r) {
  /* backwards ray tracing, from eye origin to light source */
  int i=0;
  Ray L;               /* light ray */
  Vector Rt; 
  Vector C,C_dd,C_di;  /* color values from ray tracing/photon mapping */
  int row,col;
  Vector N,ka,kd;      /* local copy of surface information */
  float t_view;

  BoxIntersection boundInt;
  BoxIntersection boxInt;
  SphereIntersection sphereInt;

  BoxIntersection ShadowBoxInt;        /* intersection information for */
  SphereIntersection ShadowSphereInt;  /* shadow feelers from
					  intersection point to light
					  source */

  /* Cast ray at each object */
  /* If ray hits object,
     - Determine object
     - Color pixel
     else
     - Determine bounding box face
     - Color pixel
  */

  if(r.intersect(&boundVol, &boundInt)) {
    /* determine view plane position, for use in rendering buffer */
    t_view=boundInt.getTHit();
  }
  else {
    return;
  }
  
  if(r.intersect(&boxObj, &boxInt) || r.intersect(&sphereObj, &sphereInt)) {
    if(boxInt.getStatus()==HIT) {
      /* Ray hits box */
      L=calcLightRay(r,boxInt.getTHit());
      C_di=calcDiffuseIndirect(L,boxObj.getNormal(boxInt.getFace()));
      C_dd=calcDiffuseDirect(boxObj.getKDiffuse(),
			     L,boxObj.getNormal(boxInt.getFace()));
    }
    else if(sphereInt.getStatus()==HIT) {
      /* Ray hits sphere */
      L=calcLightRay(r,sphereInt.getTHit());
      C_di=calcDiffuseIndirect(L,sphereInt.getNormal());
      C_dd=calcDiffuseDirect(sphereObj.getKDiffuse(),
			     L,sphereInt.getNormal());
    }
  }
  else {
    /* Ray hits background */
    L=calcLightRay(r,boundInt.getTExit());
    C_di=calcDiffuseIndirect(L,boundVol.getNormal(boundInt.getFace()));
    if(L.intersect(&boxObj, &ShadowBoxInt) || 
       L.intersect(&sphereObj, &ShadowSphereInt)) {
      C_dd.setVector(0,0,0);
    }
    else {
      C_dd=calcDiffuseDirect(boundVol.getKDiffuse(boundInt.getFace()),
			     L,boundVol.getNormal(boundInt.getFace()));
    }
  }
  C=C_di;
  Rt=r.cast(t_view);
  row=(int)floor(Rt.x[1]*IM_HEIGHT);
  col=(int)floor(Rt.x[0]*IM_WIDTH);
  for(i=0;i<3;i++) {
    photImg[row][col][i]+=C.x[i];
    rtImg[row][col][i]+=C_dd.x[i];
  } 
}

/*------------------------------------------------*/
void saveImages(void) {
  /* save resulting images to jpeg format */
  FILE *jfile=NULL;
  image jpegIMG;
  char *fname;

  jpegIMG.buffer=(GLubyte *)calloc((size_t)IM_WIDTH*IM_HEIGHT*3,sizeof(GLubyte));
  jpegIMG.w=IM_WIDTH;
  jpegIMG.h=IM_HEIGHT;
  jpegIMG.c=3;
  fname = "../FIGURES/pmapIMG_new.jpg";


  printf("Saving Photon Mapped Image\n");
  glutSetWindow(rendWindow);
  glReadPixels(0,0,(GLsizei)IM_WIDTH,(GLsizei)IM_HEIGHT,GL_RGB,GL_UNSIGNED_BYTE,
	       (GLubyte *)jpegIMG.buffer);
  flipJPEG(&jpegIMG);
  jfile=fopen(fname,"w");
  writeJPEG(&jpegIMG,jfile);
  fclose(jfile);

  fname = "../FIGURES/pmapRAW_new.jpg";

  printf("Saving Raw Photon Map Visualization\n");
  glutSetWindow(photWindow);
  glReadPixels(0,0,(GLsizei)IM_WIDTH,(GLsizei)IM_HEIGHT,GL_RGB,GL_UNSIGNED_BYTE,
	       (GLubyte *)jpegIMG.buffer);
  flipJPEG(&jpegIMG);
  jfile=fopen(fname,"w");
  writeJPEG(&jpegIMG,jfile);
  fclose(jfile);

  fname = "../FIGURES/rtImage.jpg";

  printf("Saving Simple Ray Traced image\n");
  glutSetWindow(rtWindow);
  glReadPixels(0,0,(GLsizei)IM_WIDTH,(GLsizei)IM_HEIGHT,GL_RGB,GL_UNSIGNED_BYTE,
	       (GLubyte *)jpegIMG.buffer);
  flipJPEG(&jpegIMG);
  jfile=fopen(fname,"w");
  writeJPEG(&jpegIMG,jfile);
  fclose(jfile);

}

/*------------------------------------------------------------*/
void keyboard(unsigned char key, int x, int y) {
  if(key==27 || key == 'q')
    exit(0);
  if(key=='s') {
    saveImages();
  }
}

/*------------------------------------------------------------*/
 void renderImg(void) {

  glClearColor(0.0,0.0,0.0,0.0);
  glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

  glDrawPixels(IM_WIDTH,IM_HEIGHT,GL_RGB,GL_FLOAT,photImg);
  glutSwapBuffers();
}

/*------------------------------------------------------------*/
void renderRTImg(void) {
  glClearColor(0.0,0.0,0.0,0.0);
  glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

  glDrawPixels(IM_WIDTH,IM_HEIGHT,GL_RGB,GL_FLOAT,rtImg);
  glutSwapBuffers();
}

/*------------------------------------------------------------*/
int main(int argc, char **argv) {
  initVars();
  
  glutInit(&argc,argv);
  glutInitDisplayMode(GLUT_DEPTH | GLUT_DOUBLE | GLUT_RGBA);
  
  glutInitWindowPosition(100,100);
  glutInitWindowSize(IM_WIDTH,IM_HEIGHT);
  rendWindow=glutCreateWindow("Photon Mapped Image");
  glutDisplayFunc(renderImg);
  glutKeyboardFunc(keyboard);

  glutInitWindowPosition(650,100);
  glutInitWindowSize(IM_WIDTH,IM_HEIGHT);
  photWindow=glutCreateWindow("Photon map visualization");
  glutDisplayFunc(rendPhotMap);
  glutKeyboardFunc(keyboard);

  glutInitWindowPosition(1350,100);
  glutInitWindowSize(IM_WIDTH,IM_HEIGHT);
  rtWindow=glutCreateWindow("Ray Tracer Image");
  glutDisplayFunc(renderRTImg);
  glutKeyboardFunc(keyboard);

  glutMainLoop();
  return 0;
}