//********************************************************************************
//*
//*  C++ finite element method for heat equation
//*  James sandham
//*  15 April 2015
//*
//********************************************************************************

//********************************************************************************
//
// HeatFE is free software; you can redistribute it and/or modify it under the
// terms of the GNU Lesser General Public License (as published by the Free
// Software Foundation) version 2.1 dated February 1999.
//
//********************************************************************************
#include<iostream>
#include"FeModel.h"
#include"Solver.h"
#include<GL/glut.h>
 

// global pointers needed for opengl rendering 
// (the fact that I have to use globals for opengl annoys me greatly)
int *elemType, *numOfElem;
int **connArray;
double *x_pts, *y_pts, *z_pts, *soln;

GLfloat x_angle = 0.0;
GLfloat y_angle = 0.0;
GLfloat scale = 20.0;
GLfloat x_position = 0.0;
GLfloat y_position = 0.0;
GLfloat z_position = 0.0;

int mx, my;
double xmax, xmin;

void opengl_main(int argc, char** argv);
void render_mesh();
void render_solution();
void drawWireLine1D(int i);
void drawWireTri2D(int i);
void drawWireQuad2D(int i);
void drawWireTetra3D(int i);
void drawSolidLine1D(int i);
void drawSolidTri2D(int i);
void drawSolidQuad2D(int i);
void drawSolidTetra3D(int i);
double jetColorRed(double s);
double jetColorGreen(double s);
double jetColorBlue(double s);
void keyboard(unsigned char key, int x, int y);
void keyboard_arrows(int key, int x, int y);
void mouse(int button, int state, int x, int y);
void mouse_moved(int x, int y);




//********************************************************************************
//
// Main Program 
//
//********************************************************************************
int main(int argc, char *argv[])
{
  int itr = 0;

  FeModel model(argv[1]); //create model (read mesh into constructor)

  model.readModel();  //read model data

  //for(int i=0;i<model.N;i++){
  //  std::cout<<model.xpoints[i]<<std::endl;
  //}
  //for(int i=0;i<model.Ne;i++){
  //  for(int j=0;j<model.Npe;j++){
  //   std::cout<<model.connect[i][j]<<" ";
  //  }
  //  std::cout<<""<<std::endl;
  //}
  //for(int i=0;i<model.Nb;i++){
  //  for(int j=0;j<model.Nbpe+1;j++){
  //   std::cout<<model.bconnect[i][j]<<" ";
  //  }
  //  std::cout<<""<<std::endl;
  //}
  //for(int i=0;i<4;i++){
  //  std::cout<<model.dirichletGrps[i]<<" ";
  //  std::cout<<model.dirichletVals[i]<<" ";
  //}
  //std::cout<<""<<std::endl;
  //std::cout<<model.nDirGrps<<std::endl;
  //std::cout<<model.nNeuGrps<<std::endl;


  SolverPCG solver(model);  //create solver object and pass model to constructor

  itr = solver.solve();     //solve Global system

  std::cout<<"iterations: "<<itr<<std::endl;

  //set constant pointers to model arrays
  elemType = &model.Type;
  numOfElem = &model.Ne;
  connArray = model.connect;
  x_pts = model.xpoints;
  y_pts = model.ypoints;
  z_pts = model.zpoints;
  soln = solver.x;
  xmax = soln[0];
  xmin = soln[0];
  for(int i=1;i<model.N;i++){
    if(soln[i]>xmax){xmax = soln[i];}
    else if(soln[i]<xmin){xmin = soln[i];}
  }

  opengl_main(argc,argv);  

  solver.deleteSolver();    //delete solver
  model.deleteModel();      //delete model
}










//********************************************************************************
//
// OpenGL functions
//
//********************************************************************************


//----------------------------------------------------------------------------
// Main opengl loop
//----------------------------------------------------------------------------
void opengl_main(int argc, char** argv)
{
  //initialize window
  glutInit(&argc,argv);
  glutInitDisplayMode(GLUT_DOUBLE);
  glutInitWindowPosition(100,100);
  glutInitWindowSize(800,800);
  glutCreateWindow("finite element mesh");
  glutSetCursor(GLUT_CURSOR_CROSSHAIR);

  //keyboard and mouse functions
  glutKeyboardFunc(keyboard);
  glutSpecialFunc(keyboard_arrows);
  glutMouseFunc(mouse);
  glutMotionFunc(mouse_moved);

  //render function
  glutDisplayFunc(render_mesh);

  //main loop
  glutMainLoop();
}



//----------------------------------------------------------------------------
// render the mesh
//----------------------------------------------------------------------------
void render_mesh()
{
  void (*drawElement)(int i);  

  glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
  gluLookAt(0.0,0.0,1.0,0.0,0.0,0.0,0.0,1.0,0.0);
  glMatrixMode(GL_MODELVIEW);
  glLoadIdentity();
  glOrtho(-320,320,-240,240,-200.0,200.0);
  glEnable(GL_DEPTH_TEST);

  //translations, rotations, and scalings
  glTranslatef(x_position,y_position,z_position);
  glRotatef(x_angle,0,1,0);
  glRotatef(y_angle,1,0,0);
  glScalef(scale,scale,scale);

  //set function pointer to correct element to draw
  switch(*elemType)
  {
    case(1):
      drawElement = &drawSolidLine1D;
      break;
    case(2):
      drawElement = &drawSolidTri2D;
      break;
    case(3):
      drawElement = &drawWireQuad2D;
      break;
    case(4):
      drawElement = &drawWireTetra3D;
      break;
    case(8):
      drawElement = &drawSolidLine1D;
      break;
    case(9):
      drawElement = &drawSolidTri2D;
      break;
    case(10):
      drawElement = &drawWireQuad2D;
      break;
    case(11):
      drawElement = &drawWireTetra3D;
      break;
  }

  //draw mesh
  glNewList(1,GL_COMPILE);
    for(int i=0;i<*numOfElem;i++){
      drawElement(i);
    }
  glEndList();
  glCallList(1);

  glutSwapBuffers();
}



//----------------------------------------------------------------------------
// render the solution
//----------------------------------------------------------------------------
void render_solution()
{
  void (*drawElement)(int i);

  glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
  gluLookAt(0.0,0.0,1.0,0.0,0.0,0.0,0.0,1.0,0.0);
  glMatrixMode(GL_MODELVIEW);
  glLoadIdentity();
  glOrtho(-320,320,-240,240,-200.0,200.0);
  glEnable(GL_DEPTH_TEST);

  //translations, rotations, and scalings
  glTranslatef(x_position,y_position,z_position);
  glRotatef(x_angle,0,1,0);
  glRotatef(y_angle,1,0,0);
  glScalef(scale,scale,scale);

  //set function pointer to correct element to draw
  switch(*elemType)
  {
    case(1):
      drawElement = &drawSolidLine1D;
      break;
    case(2):
      drawElement = &drawSolidTri2D;
      break;
    case(3):
      drawElement = &drawSolidQuad2D;
      break;
    case(4):
      drawElement = &drawSolidTetra3D;
      break;
    case(8):
      drawElement = &drawSolidLine1D;
      break;
    case(9):
      drawElement = &drawSolidTri2D;
      break;
    case(10):
      drawElement = &drawSolidQuad2D;
      break;
    case(11):
      drawElement = &drawSolidTetra3D;
      break;
  }

  //draw mesh
  glNewList(1,GL_COMPILE);
    for(int i=0;i<*numOfElem;i++){
      drawElement(i);
    }
  glEndList();
  glCallList(1);

  glutSwapBuffers();
}



//----------------------------------------------------------------------------
// draw wire functions for different elements
//----------------------------------------------------------------------------
void drawWireLine1D(int i)
{
  glBegin(GL_LINE_LOOP);
    for(int j=0;j<2;j++){
      glVertex3f(x_pts[connArray[i][j]-1],
                 y_pts[connArray[i][j]-1],
                 z_pts[connArray[i][j]-1]);
    }
  glEnd();
}


void drawWireTri2D(int i)
{
  glBegin(GL_LINE_LOOP);
    for(int j=0;j<3;j++){
      glVertex3f(x_pts[connArray[i][j]-1],
                 y_pts[connArray[i][j]-1],
                 z_pts[connArray[i][j]-1]);
    }
  glEnd();
}


void drawWireQuad2D(int i)
{
  glBegin(GL_LINE_LOOP);
    for(int j=0;j<4;j++){
      glVertex3f(x_pts[connArray[i][j]-1],
                 y_pts[connArray[i][j]-1],
                 z_pts[connArray[i][j]-1]);
    }
  glEnd();
}


void drawWireTetra3D(int i)
{
  glBegin(GL_LINE_LOOP);
    glVertex3f(x_pts[connArray[i][0]-1],
               y_pts[connArray[i][0]-1],
               z_pts[connArray[i][0]-1]);
    glVertex3f(x_pts[connArray[i][1]-1],
               y_pts[connArray[i][1]-1],
               z_pts[connArray[i][1]-1]);
    glVertex3f(x_pts[connArray[i][2]-1],
               y_pts[connArray[i][2]-1],
               z_pts[connArray[i][2]-1]);
  glEnd();
  glBegin(GL_LINE_LOOP);
    glVertex3f(x_pts[connArray[i][0]-1],
               y_pts[connArray[i][0]-1],
               z_pts[connArray[i][0]-1]);
    glVertex3f(x_pts[connArray[i][1]-1],
               y_pts[connArray[i][1]-1],
               z_pts[connArray[i][1]-1]);
    glVertex3f(x_pts[connArray[i][3]-1],
               y_pts[connArray[i][3]-1],
               z_pts[connArray[i][3]-1]);
  glEnd();
  glBegin(GL_LINE_LOOP);
    glVertex3f(x_pts[connArray[i][0]-1],
               y_pts[connArray[i][0]-1],
               z_pts[connArray[i][0]-1]);
    glVertex3f(x_pts[connArray[i][2]-1],
               y_pts[connArray[i][2]-1],
               z_pts[connArray[i][2]-1]);
    glVertex3f(x_pts[connArray[i][3]-1],
               y_pts[connArray[i][3]-1],
               z_pts[connArray[i][3]-1]);
  glEnd();
  glBegin(GL_LINE_LOOP);
    glVertex3f(x_pts[connArray[i][1]-1],
               y_pts[connArray[i][1]-1],
               z_pts[connArray[i][1]-1]);
    glVertex3f(x_pts[connArray[i][2]-1],
               y_pts[connArray[i][2]-1],
               z_pts[connArray[i][2]-1]);
    glVertex3f(x_pts[connArray[i][3]-1],
               y_pts[connArray[i][3]-1],
               z_pts[connArray[i][3]-1]);
  glEnd();
}



//----------------------------------------------------------------------------
// draw solid functions for different elements
//----------------------------------------------------------------------------
void drawSolidLine1D(int i)
{
  //jet colormap
  double s = soln[connArray[i][0]-1];
  double red = jetColorRed(s);
  double green = jetColorGreen(s);
  double blue = jetColorBlue(s);

  glColor3f(red,green,blue);

  glBegin(GL_LINE_LOOP);
    for(int j=0;j<2;j++){
      glVertex3f(x_pts[connArray[i][j]-1],
                 y_pts[connArray[i][j]-1],
                 z_pts[connArray[i][j]-1]);
    }
  glEnd();
}


void drawSolidTri2D(int i)
{
  //jet colormap
  double s = soln[connArray[i][0]-1];
  double red = jetColorRed(s);
  double green = jetColorGreen(s);
  double blue = jetColorBlue(s);

  glColor3f(red,green,blue);

  glBegin(GL_TRIANGLES);
    for(int j=0;j<3;j++){
      glVertex3f(x_pts[connArray[i][j]-1],
                 y_pts[connArray[i][j]-1],
                 z_pts[connArray[i][j]-1]);
    }
  glEnd();
}


void drawSolidQuad2D(int i)
{
  //jet colormap
  double s = soln[connArray[i][0]-1];
  double red = jetColorRed(s);
  double green = jetColorGreen(s);
  double blue = jetColorBlue(s);

  glColor3f(red,green,blue);

  glBegin(GL_QUADS);
    for(int j=0;j<4;j++){
      glVertex3f(x_pts[connArray[i][j]-1],
                 y_pts[connArray[i][j]-1],
                 z_pts[connArray[i][j]-1]);
    }
  glEnd();
}


void drawSolidTetra3D(int i)
{
  //jet colormap
  double s = soln[connArray[i][0]-1];
  double red = jetColorRed(s);
  double green = jetColorGreen(s);
  double blue = jetColorBlue(s);

  glColor3f(red,green,blue);

  glBegin(GL_TRIANGLES);
    glVertex3f(x_pts[connArray[i][0]-1],
               y_pts[connArray[i][0]-1],
               z_pts[connArray[i][0]-1]);
    glVertex3f(x_pts[connArray[i][1]-1],
               y_pts[connArray[i][1]-1],
               z_pts[connArray[i][1]-1]);
    glVertex3f(x_pts[connArray[i][2]-1],
               y_pts[connArray[i][2]-1],
               z_pts[connArray[i][2]-1]);
  glEnd();
  glBegin(GL_TRIANGLES);
    glVertex3f(x_pts[connArray[i][0]-1],
               y_pts[connArray[i][0]-1],
               z_pts[connArray[i][0]-1]);
    glVertex3f(x_pts[connArray[i][1]-1],
               y_pts[connArray[i][1]-1],
               z_pts[connArray[i][1]-1]);
    glVertex3f(x_pts[connArray[i][3]-1],
               y_pts[connArray[i][3]-1],
               z_pts[connArray[i][3]-1]);
  glEnd();
  glBegin(GL_TRIANGLES);
    glVertex3f(x_pts[connArray[i][0]-1],
               y_pts[connArray[i][0]-1],
               z_pts[connArray[i][0]-1]);
    glVertex3f(x_pts[connArray[i][2]-1],
               y_pts[connArray[i][2]-1],
               z_pts[connArray[i][2]-1]);
    glVertex3f(x_pts[connArray[i][3]-1],
               y_pts[connArray[i][3]-1],
               z_pts[connArray[i][3]-1]);
  glEnd();
  glBegin(GL_TRIANGLES);
    glVertex3f(x_pts[connArray[i][1]-1],
               y_pts[connArray[i][1]-1],
               z_pts[connArray[i][1]-1]);
    glVertex3f(x_pts[connArray[i][2]-1],
               y_pts[connArray[i][2]-1],
               z_pts[connArray[i][2]-1]);
    glVertex3f(x_pts[connArray[i][3]-1],
               y_pts[connArray[i][3]-1],
               z_pts[connArray[i][3]-1]);
  glEnd();
}




//----------------------------------------------------------------------------
// colormaps
//----------------------------------------------------------------------------
double jetColorRed(double s)
{
  //set red
  double red;
  if((s-xmin)/(xmax-xmin)<0.4){red = 0.0;}
  else if((s-xmin)/(xmax-xmin)<0.6){red = 5.0*(s-xmin)/(xmax-xmin) - 2.0;}
  else if((s-xmin)/(xmax-xmin)<0.8){red = 1.0;}
  else{red = -5/2*(s-xmin)/(xmax-xmin)+7/2;}

  return red;
}


double jetColorGreen(double s)
{
  //green
  double green;
  if((s-xmin)/(xmax-xmin)<0.1){green = 0.0;}
  else if((s-xmin)/(xmax-xmin)<0.4){green = 10/3*(s-xmin)/(xmax-xmin)-1/3;}
  else if((s-xmin)/(xmax-xmin)<0.6){green = 1.0;}
  else if((s-xmin)/(xmax-xmin)<0.8){green = -5.0*(s-xmin)/(xmax-xmin)+4.0;}
  else{green = 0.0;}

  return green;
}


double jetColorBlue(double s)
{
  //blue
  double blue;
  if((s-xmin)/(xmax-xmin)<0.1){blue = 5*(s-xmin)/(xmax-xmin)+0.5;}
  else if((s-xmin)/(xmax-xmin)<0.4){blue = 1.0;}
  else if((s-xmin)/(xmax-xmin)<0.8){blue = -5*(s-xmin)/(xmax-xmin)+3.0;}
  else{blue = 0.0;}
  return blue;
}



//----------------------------------------------------------------------------
// keyboard and mouse interactions
//----------------------------------------------------------------------------
void keyboard(unsigned char key, int x, int y)
{
  if(key==27){exit(0);}
  if(key=='z'){scale += 1.0; glutPostRedisplay();}
  if(key=='x'){scale -= 1.0; glutPostRedisplay();}
  if(key=='a'){x_angle -= 1.0; glutPostRedisplay();}
}



void keyboard_arrows(int key, int x, int y)
{
  if(key==GLUT_KEY_LEFT){x_position += 5; glutPostRedisplay();}
  if(key==GLUT_KEY_RIGHT){x_position -= 5; glutPostRedisplay();}
  if(key==GLUT_KEY_UP){y_position -= 5; glutPostRedisplay();}
  if(key==GLUT_KEY_DOWN){y_position += 5; glutPostRedisplay();}
}



void mouse(int button, int state, int x, int y)
{
  mx = x;
  my = y;
}


void mouse_moved(int x, int y)
{
  x_position += (mx-x);
  y_position += (my-y);
  mx = x;
  my = y;
  glutPostRedisplay();
}
