//-----------------------------------------------------------------------bl-
//--------------------------------------------------------------------------
//
// Copyright (C) 2014 Paul T. Bauman
//
// This library is free software; you can redistribute it and/or
// modify it under the terms of the Version 2.1 GNU Lesser General
// Public License as published by the Free Software Foundation.
//
// This library is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc. 51 Franklin Street, Fifth Floor,
// Boston, MA  02110-1301  USA
//
//-----------------------------------------------------------------------el-

#include <petsc.h> // petsc.h will have everything we need for PETSc
#include <cmath> // Math functions for C++
#include <iostream>

// Stubs for functions that we need to write
PetscErrorCode FormFunction(SNES,Vec,Vec,void*);
PetscErrorCode FormJacobian(SNES,Vec,Mat,Mat,void*);

/* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
   Our user-defined context. Use this to pass auxillary information
   to our FormFunction and FormJacobian functions.
   - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
struct MyContext
{
  DM dm;
  DM dmf;
};

/* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
   The purpose of this code is to illustrate the PETSc DMDA
   framework for solving nonlinear equations posed on structured
   grids. The DMDA framework will take care of things such as
   partitioning, ghost points, etc.
   This program was run using PETSc 3.4.3. Your mileage may vary
   with other PETSc versions.
   - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
int main(int argc,char **args)
{
  PetscErrorCode ierr;
  SNES           snes;         /* nonlinear solver context */
  Vec            u,r;          /* solution, residual vectors */
  Mat            J,interp;            /* Jacobian matrix */

  MyContext      ctx; /* Our context */

  /* Only call this *once* in a program */
  PetscInitialize(&argc,&args,PETSC_NULL,PETSC_NULL);

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
     Create nonlinear solver context
     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  ierr = SNESCreate(PETSC_COMM_WORLD,&snes);CHKERRQ(ierr);


  const int N = 101; /* Default to 101 points, h = 1/100 */

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
     Create our DM object.
      - Operate on the the COMM_WORLD communicator
      - This is a GHOSTED DMDA (as opposed to periodic)
      - Number of points. We use a negative value, because then we
        can set the size from the command line using -da_grid_x <nx> 
      - 1 dof per node
      - Stencil width is 1 (standard central difference scheme)
      - Supply NULL to let PETSc decide node layout over the processors
      - DM object where this info will be stashed


      PetscErrorCode  DMDACreate2d(MPI_Comm comm,DMDABoundaryType bx,DMDABoundaryType by,DMDAStencilType stencil_type,
                          PetscInt M,PetscInt N,PetscInt m,PetscInt n,PetscInt dof,PetscInt s,const PetscInt lx[],const PetscInt ly[],DM *da)

      comm	- MPI communicator
      bx,by	- type of ghost nodes the array have. Use one of DMDA_BOUNDARY_NONE, DMDA_BOUNDARY_GHOSTED, DMDA_BOUNDARY_PERIODIC.
      stencil_type	- stencil type. Use either DMDA_STENCIL_BOX or DMDA_STENCIL_STAR.
      M,N	- global dimension in each direction of the array (use -M and or -N to indicate that it may be set to a different value
                  from the command line with -da_grid_x <M> -da_grid_y <N>)
      m,n	- corresponding number of processors in each dimension (or PETSC_DECIDE to have calculated)
      dof	- number of degrees of freedom per node
      s	- stencil width
      lx, ly	- arrays containing the number of nodes in each cell along the x and y coordinates, or NULL. If non-null,
                  these must be of length as m and n, and the corresponding m and n cannot be PETSC_DECIDE. The sum of the lx[]
                  entries must be M, and the sum of the ly[] entries must be N.
     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  ierr = DMDACreate2d( PETSC_COMM_WORLD, DM_BOUNDARY_GHOSTED, DM_BOUNDARY_GHOSTED, DMDA_STENCIL_STAR, -N, -N, PETSC_DECIDE, PETSC_DECIDE, 1, 1, NULL, NULL, &ctx.dm );CHKERRQ(ierr);



  std::cout<<std::endl<<"testing proj matrix" << std::endl; 
  // proj mtx testing
  ierr = DMRefine(ctx.dm, PETSC_COMM_WORLD, &ctx.dmf);CHKERRQ(ierr);
  ierr = DMCreateInterpolation(ctx.dm, ctx.dmf, &interp, 0 );
  ierr = MatView(interp, PETSC_VIEWER_STDOUT_SELF);

  PetscInt nn,mm;
  ierr = MatGetSize(interp, &mm,&nn);
  std::cout<<mm<< "x" << nn << std::endl; 





  ierr = SNESSetDM(snes,ctx.dm);CHKERRQ(ierr);

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
     Create matrix and vector data structures; set corresponding routines
     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  ierr = DMCreateGlobalVector(ctx.dm,&u);CHKERRQ(ierr); // Solution vector
  ierr = VecDuplicate(u,&r);CHKERRQ(ierr); // Residual vector

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
     Use DM call to create matrix as it will take care of preallocating
     the matrix entries for us as well as the parallel layout for the
     given DA. Note that SNES doesn't care about the DM, per se, we can
     just hand it solution, residual, and Jacobian and it does the right
     thing.
     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  ierr = DMCreateMatrix(ctx.dm,&J);CHKERRQ(ierr); // Jacobian matrix
  ierr = MatSetFromOptions(J);CHKERRQ(ierr); // Allows setting options from command line

  // Tell SNES what the residual function is
  ierr = SNESSetFunction(snes,r,FormFunction,&ctx);CHKERRQ(ierr);

  // Tell SNES what the Jabobian function is
  ierr = SNESSetJacobian(snes,J,J,FormJacobian,&ctx);CHKERRQ(ierr);

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
     Set SNES/KSP/KSP/PC runtime options, e.g.,
         -snes_view -snes_monitor -ksp_type <ksp> -pc_type <pc>
     These options will override those specified above as long as
     SNESSetFromOptions() is called _after_ any other customization
     routines.
     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  ierr = SNESSetFromOptions(snes);CHKERRQ(ierr);

  // Solve our nonlinear problem
  ierr = SNESSolve(snes,NULL,u);CHKERRQ(ierr);

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
     PETSc has some built-in capability for outputting vectors and
     matrices to nice formats. Here, we dump out the solution vector
     in an ASCII format that can be read by Octave/Matlab.
   - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  PetscViewer viewer;
  ierr = PetscViewerASCIIOpen(PETSC_COMM_WORLD, "p1_solution.vtk", &viewer);CHKERRQ(ierr);
  ierr = PetscViewerSetFormat(viewer, PETSC_VIEWER_ASCII_VTK);CHKERRQ(ierr);
  ierr = VecView(u,viewer);CHKERRQ(ierr);
  ierr = PetscViewerDestroy(&viewer);CHKERRQ(ierr);

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
     Free work space.  All PETSc objects should be destroyed when they
     are no longer needed.
     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  ierr = VecDestroy(&u);CHKERRQ(ierr);
  ierr = VecDestroy(&r);CHKERRQ(ierr);
  ierr = MatDestroy(&J);CHKERRQ(ierr);
  ierr = SNESDestroy(&snes);CHKERRQ(ierr);
  ierr = DMDestroy(&ctx.dm);CHKERRQ(ierr);

  /* Only call this *once* in a program */
  ierr = PetscFinalize();

  return 0;
}

// Function for evaluating residual
PetscErrorCode FormFunction(SNES snes, Vec u, Vec r, void* ctx)
{
  PetscErrorCode     ierr;
  // Cast ctx pointer to correct type
  MyContext* user = (MyContext*) ctx;
  
  // Extract DM data structure
  DM dm = user->dm;

  PetscInt N, M; // n_nodes
  PetscScalar h; // mesh size

  Vec u_local;

  // Use DM to extract local part of the solution vector
  DMGetLocalVector(dm,&u_local);

  // Populate local vector from global vector, including ghost points
  DMGlobalToLocalBegin(dm,u,INSERT_VALUES,u_local);
  DMGlobalToLocalEnd(dm,u,INSERT_VALUES,u_local);

  // Pull out pointer to solution vector data
  PetscScalar** u_ptr;
  DMDAVecGetArray(dm,u_local,&u_ptr);

  // Pull out pointer to residual vector data
  PetscScalar** r_ptr;
  DMDAVecGetArray(dm,r,&r_ptr);

  PetscInt xs, ys, xw, yw;

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
     Retrieving the starting index at the corner and the width
     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  DMDAGetCorners(dm,&xs,&ys,NULL,&xw,&yw,NULL);

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
     Retrieving number of points in the domain
     Note this can change at runtime through the command line
     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  DMDAGetInfo(dm,NULL,&N,&M,NULL,NULL,NULL,NULL,NULL,
              NULL,NULL,NULL,NULL,NULL);

  if( N != M )
    {
      std::cerr << "Error: Currently assuming N = M!" << std::endl
                << "       Found N = " << N << ", M = " << M << std::endl;
    }

  // Assuming domain width is 1.0
  PetscScalar domain_size = 1.0;
  h = domain_size/(N-1);

  // Variables to cache starting and stopping indices
  PetscInt start_x = xs;
  PetscInt stop_x = xs+xw;

  PetscInt start_y = ys;
  PetscInt stop_y = ys+yw;

  // Loop over the nodes on the current processor and assemble the residual
  for( PetscInt j = start_y; j < stop_y; j++ )
    {
      PetscScalar y = j*h;

      for( PetscInt i = start_x; i < stop_x; i++ )
        {
          PetscScalar x = i*h;

          PetscScalar forcing =  2*( (1-6*x*x)*y*y*(1-y*y) + (1-6*y*y)*x*x*(1-x*x) );
            
          if(i == 0 || j == 0 || i == N-1 || j == M-1)
            {
              r_ptr[j][i] = 0.0;
            }
          else
            {
              // FD approximation of: -u'' = f
              // Nonlinear function is -u'' - f = 0
              r_ptr[j][i] = (-u_ptr[j][i-1] -u_ptr[j-1][i]+ 4*u_ptr[j][i] - u_ptr[j][i+1] - u_ptr[j+1][i])/(h*h) - forcing;
            }
        }
    }

  /* Restore vectors */
  ierr = DMDAVecRestoreArray(dm,u_local,&u_ptr);CHKERRQ(ierr);
  ierr = DMDAVecRestoreArray(dm,r,&r_ptr);CHKERRQ(ierr);
  ierr = DMRestoreLocalVector(dm,&u_local);CHKERRQ(ierr);

  return 0;
}

// Function for evaluating Jacobian
PetscErrorCode FormJacobian(SNES snes, Vec u, Mat J, Mat Jprec, void* ctx)
{
  PetscErrorCode     ierr;
  // Cast ctx pointer to correct type
  MyContext* user = (MyContext*) ctx;

  // Extract DM data structure
  DM dm = user->dm;

  PetscScalar h;
  PetscInt N, M;

  Vec u_local;

  // Use DM to extract local part of the solution vector
  DMGetLocalVector(dm,&u_local);

  // Populate local vector from global vector, including ghost points
  DMGlobalToLocalBegin(dm,u,INSERT_VALUES,u_local);
  DMGlobalToLocalEnd(dm,u,INSERT_VALUES,u_local);

  // Pull out pointer to solution vector data
  PetscScalar* u_ptr;
  DMDAVecGetArray(dm,u_local,&u_ptr);

  PetscInt xs, xw;
  PetscInt ys, yw;

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
     Retrieving the starting index at the corner and the width
     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  DMDAGetCorners(dm,&xs,&ys,NULL,&xw,&yw,NULL);

  /* - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
     Retrieving number of points in the domain
     Note this can change at runtime through the command line
     - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - */
  DMDAGetInfo(dm,NULL,&N,&M,NULL,NULL,NULL,NULL,NULL,
              NULL,NULL,NULL,NULL,NULL);

  if( N != M )
    {
      std::cerr << "Error: Currently assuming N = M!" << std::endl
                << "       Found N = " << N << ", M = " << M << std::endl;
    }

  // Assuming domain width is 1.0
  PetscScalar domain_size = 1.0;
  h = domain_size/(N-1);

  // Variables to cache starting and stopping indices
  PetscInt start_x = xs;
  PetscInt stop_x = xs+xw;

  PetscInt start_y = ys;
  PetscInt stop_y = ys+yw;

  MatStencil row;
  MatStencil columns[5];
  PetscScalar values[5];
  
  MatStencil* boundary_rows;
  PetscInt nrows = 0;

  ierr = PetscMalloc(xw*yw*sizeof(MatStencil),&boundary_rows);CHKERRQ(ierr);

  // Loop over the nodes on the current processor and assemble the Jacobian
  for( PetscInt j = start_y; j < stop_y; j++ )
    {
      for( PetscInt i = start_x; i < stop_x; i++ )
        {
          if(i == 0 || j == 0 || i == N-1 || j == M-1)
            {
              boundary_rows[nrows].i = i;
              boundary_rows[nrows].j = j;
              nrows++;
            }

          row.i = i;
          row.j = j;

          columns[0].i = i-1; columns[0].j = j;
          values[0] = -1.0/(h*h);

          columns[1].i = i; columns[1].j = j-1;
          values[1] = -1.0/(h*h);

          columns[2].i = i; columns[2].j = j;
          values[2] = 4.0/(h*h);

          columns[3].i = i+1; columns[3].j = j;
          values[3] = -1.0/(h*h);

          columns[4].i = i; columns[4].j = j+1;
          values[4] = -1.0/(h*h);

          ierr = MatSetValuesStencil(J,1,&row,5,columns,values,INSERT_VALUES);CHKERRQ(ierr);
        }
    }

  /* Restore vectors */
  ierr = DMDAVecRestoreArray(dm,u_local,&u_ptr);CHKERRQ(ierr);
  ierr = DMRestoreLocalVector(dm,&u_local);CHKERRQ(ierr);

  // Tell PETSc you're done putting values into the matrix
  ierr = MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
  ierr = MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);

  ierr = MatZeroRowsColumnsStencil(J,nrows,boundary_rows,1.0,NULL,NULL);CHKERRQ(ierr);
  ierr = PetscFree(boundary_rows);CHKERRQ(ierr);

  ierr = MatSetOption(J,MAT_NEW_NONZERO_LOCATION_ERR,PETSC_TRUE);CHKERRQ(ierr);
  ierr = MatView(J, PETSC_VIEWER_STDOUT_SELF);

  return 0;
}