#include "Sundance.h"
#include <math.h>

extern "C" {
#include <pnm.h>                          // for pnm images
}

//
// image problem using Galerkin approximation
//

bool img_invert = false;                         // flag to invert image colors


void readTargetImage( Expr & ustar, string img_name, int nx, int ny )
{
  // the image in img_name gives ustar. by convention, black is the 0 pixel and
  // white is the highest pixel value - in our case, 255.  the image pixels
  // should range from black to white (a gray picture).
  
  int format, row, col, plane;
  int num_row, num_col;
  xelval maxval;
  xel * xelrow;
  FILE * file;

  // get vector associated with expression ustar
  TSFVector vector;
  ustar.getVector( vector );

  file = fopen( img_name.c_str(), "r" );
  if ( file == NULL )
    TSFError::raise("Error: could not open image file for reading");

  // read an image in RPGM_FORMAT format
  pnm_readpnminit( file, &num_col, &num_row, &maxval, &format );

  if ( format != RPGM_FORMAT )
    TSFError::raise("Error: image is not in proper PNM/PGM format");

  // image and mesh sizes must match
  if ( num_col != nx+1 || num_row != ny+1 )
    TSFError::raise("Error: image has different size than mesh");

  // allocate memory for a row of data
  xelrow = pnm_allocrow( num_col );

  for ( row = 0; row < num_row; row++ ) {
    // each row in the picture runs from left to right and it starts from top
    // left corner. thus the picture is vertically flipped w.r.t. the mesh
    
    pnm_readpnmrow( file, xelrow, num_col, maxval, format );
    int dof_index = ny - row;             // column position in the current row

    if ( row == 0 ) {
      // check if image need to be inverted. the code works for black (0) as
      // background color and white (255 = maxval) the target color
      if ( xelrow[0].b == maxval ) img_invert = true;
    }
    
    for ( col = 0; col < num_col; ++col ) {
      // only blue pixel should be set
      int pixel_value = xelrow[col].b;
      if ( img_invert ) pixel_value = maxval - pixel_value; // invert pixel
      vector[dof_index] = pixel_value/(double)maxval;
      
      dof_index += num_row;
    }
  }
  
  fclose( file );
  pnm_freerow( xelrow );
}

void saveAsImage( Expr & u, string filename, int nx, int ny, bool isGray = true)
{
  // save discrete function 'u' as a pnm image in file 'filename'. the image
  // size is given by the size of the mesh 'u' is discretized into.  if
  // 'isGray' is true, the picture is a gray one (this is the default).
  // otherwise it is a binary B&W image.
  
  int num_row = ny + 1;
  int num_col = nx + 1;
  int white = 0;
  int black = 1;
  int format;                             // PBM or PGM format (PPM later)
  xelval maxval;

  int row, col, pixel_value;
  double pix_value;
  xel * xelrow;
  FILE * file;

  file = fopen( filename.c_str(), "w" );
  if ( file == NULL )
    TSFError::raise("Error: could not open image file for writing");

  if ( isGray ) {
    // raw pgm format
    format = RPGM_FORMAT;
    maxval = PNM_MAXMAXVAL;
  }
  else {
    // pbm (binary) file
    format = PNM_FORMAT_TYPE(4);
    maxval = 1;                           // black
  }

  // write header in image file
  pnm_writepnminit( file, num_col, num_row, maxval, format, 0 );
  
  // array for a row of pixels
  xelrow = pnm_allocrow( num_col );

  // get vector associated with expression u
  TSFVector vector;
  u.getVector( vector );
  
  for ( row = 0; row < num_row; row++ ) {
    // each row in the picture runs from left to right and it starts from top
    // left corner. thus the picture is vertically flipped w.r.t. the mesh
    
    int dof_index = ny - row;             // column position in the current row
    for ( col = 0; col < num_col; ++col ) {
      pix_value = ( img_invert ) ? 1.0 - vector[dof_index] : vector[dof_index];
      if ( isGray )
        pixel_value = (xelval) round( maxval * pix_value );
      else
        pixel_value = ( pix_value < 0.5 ) ? 0 : 1;

      PNM_ASSIGN1( xelrow[col], pixel_value );
      dof_index += num_row;
    }
    // write out row
    pnm_writepnmrow( file, xelrow, num_col, maxval, format, 0 );
  }

  fclose( file );
  pnm_freerow( xelrow );
}


//
// main function
//

int main (int argc, void ** argv)
{
  try
  {
    string img_name = "image.pnm";
    string out_name = "galerkin";
    int nx = 64, ny = 64;
    int degree = 1;
    double beta = 1.0;
    
    // init Sundance and Trilinos command line reader
    Sundance::init( &argc, &argv );
    TSFCommandLine::init( argc, argv );
    TSFCommandLine::verbose() = false;
    
    // read command line switches
    TSFCommandLine::findInt("-d", degree);
    TSFCommandLine::findInt("-nx", nx );
    TSFCommandLine::findInt("-ny", ny );
    TSFCommandLine::findDouble("-beta", beta );
    TSFCommandLine::findString( "-img", img_name);
    TSFCommandLine::findString( "-out", out_name);
    
    // Create a mesh object. In this example, we will use a built-in method to
    // create a uniform mesh on the unit square. In more realistic problems we
    // would use a mesher to create a mesh, and then read the mesh using a
    // MeshReader object.

    MeshGenerator mesher = new RectangleMesher( 0., 1., nx, 0., 1., ny );
    Mesh mesh = mesher.getMesh();

    // define symbolic objects to represent x and y coordinate functions
    // and derivatives w.r.t. x and y
      // x, y and gradient operator
    Expr x = new CoordExpr(0, "x");
    Expr y = new CoordExpr(1, "y");
    Expr dx = new Derivative(0,1);
    Expr dy = new Derivative(1,1);
    Expr grad = List( dx, dy );

    // define a cell set that contains all boundary cells 
    CellSet boundary = new BoundaryCellSet();

    // create a discrete space 'space' on the mesh.
    // we will need 'space' to assign an image to a discrete function.
    TSFVectorType petra  = new PetraVectorType();
    BasisFamily lagr     = new Lagrange(degree);
    TSFVectorSpace space = new SundanceVectorSpace( mesh, lagr, petra );

    // define an unknown function and its variation. The constructor argument
    // is the basis family with which the function will be represented, in this
    // case Lagrange (nodal) polynomials.
    Expr u = new UnknownFunction ( lagr, "u" );
    Expr var_u = new TestFunction( lagr, "v" );
    
    // setup the target image
    Expr ustar = DiscreteFunction::discretize( space, Expr(0.0) );
    readTargetImage( ustar, img_name, nx, ny );
    
    // having constructed the unknown, variation, and differentiation operator,
    // we can now define the weak form of the PDE
    Expr weak_form = Integral( (u - ustar) * var_u +
                               beta * (grad * u) * (grad * var_u) );
    
    // boundary condition also in a weak form sense
    EssentialBC bc = EssentialBC( boundary, u * var_u );

    // Create a solver object: stablized biconjugate gradient solver
    TSFPreconditionerFactory prec = new ILUKPreconditionerFactory( 1 );
    TSFLinearSolver solver = new BICGSTABSolver( prec, 1.0e-12, 300 );

    // combine the geometry, the variational form, the BCs, and the solver to
    // form a complete problem.
    StaticLinearProblem prob( mesh, weak_form, bc, var_u, u );

    // solve the problem, obtaining the solution as a (discrete) Expr object
    Expr soln = prob.solve( solver );

    // write the solution in a form readable by matlab
    ofstream outfile( (out_name + ".dat").c_str() );
    outfile.precision( 12 );
    soln.matlabDump( outfile );

    // write out soln as an image
    saveAsImage( soln, out_name + ".pnm", nx, ny );
    
  }
  catch( std::exception & e ) {
    TSFOut::println( e.what() );
  }

  Sundance::finalize ();
}