#define PETSC_USE_COMPLEX 1
#include <slepceps.h>
#include <iostream>
#include <chrono>

inline
int
fillMatA(PetscInt nx, PetscInt ny, PetscReal Dx, PetscReal Dy, Mat &A)
{
  PetscErrorCode ierr;
  PetscInt i, j, II, Istart, Iend;
  ierr = MatGetOwnershipRange(A,&Istart,&Iend);CHKERRQ(ierr);

  for (II=Istart; II<Iend; II++) 
  {
    i = II/nx; 
    j = II-i*nx;
    
    if (i>0) { ierr = MatSetValue(A,II,II-nx,-Dy,INSERT_VALUES);CHKERRQ(ierr); }
    if (i<ny-1) { ierr = MatSetValue(A,II,II+nx,-Dy,INSERT_VALUES);CHKERRQ(ierr); }
    if (j>0) { ierr = MatSetValue(A,II,II-1,-Dx,INSERT_VALUES);CHKERRQ(ierr); }
    if (j<nx-1) { ierr = MatSetValue(A,II,II+1,-Dx,INSERT_VALUES);CHKERRQ(ierr); }
    
    ierr = MatSetValue(A, II, II, Dx+Dy, INSERT_VALUES);CHKERRQ(ierr);
  }

  ierr = MatAssemblyBegin(A,MAT_FLUSH_ASSEMBLY);CHKERRQ(ierr);
  ierr = MatAssemblyEnd(A,MAT_FLUSH_ASSEMBLY);CHKERRQ(ierr);

  for (II=Istart; II<Iend; II++) 
  {
    i = II/nx; 
    j = II-i*nx;

    if ((i == 0 || i == (ny-1)) && j != 0 && j != (nx-1))
      {
	ierr = MatSetValue(A, II,II, Dx, ADD_VALUES);CHKERRQ(ierr);
      }
    else if ((j == 0 || j == (nx-1)) && i != 0 && i != (ny-1))
      {
	ierr = MatSetValue(A,II,II, Dy, ADD_VALUES);CHKERRQ(ierr);
      }
    else if (i > 0 && i < (ny-1) && j > 0 && j < (nx-1))
      {
	ierr = MatSetValue(A,II,II, Dx+Dy, ADD_VALUES);CHKERRQ(ierr);
      }
  }

  ierr = MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
  ierr = MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
}

struct Matrix
{
  Mat A;
};

static
PetscErrorCode
matMultA(Mat A, Vec x, Vec y)
{
  PetscFunctionBeginUser;

  void *ctx = NULL;

  MatShellGetContext(A, &ctx);
  Matrix &matrix = *static_cast<Matrix*>(ctx);

  MatMult(matrix.A, x, y);

  PetscFunctionReturn(0);
  return 0;
};

int main(int argc, char **argv)
{
  Mat A;
  EPS eps;
  ST st;
  KSP ksp;
  PC pc;
  EPSType type;
  PetscErrorCode ierr;

  // Box dimension and speed of sound (all fixed)
  PetscReal Lx = 0.4;
  PetscReal Ly = 0.2;
  PetscReal sos = 347.0;

  // Mesh size (can be changed from the command line)
  PetscInt nx = 1600;
  PetscInt ny = 400;
  PetscInt nev = 6;

  ierr = SlepcInitialize(&argc,&argv,(char*)0,NULL);if (ierr) return ierr;

  ierr = PetscOptionsGetInt(NULL,NULL,"-nx",&nx,NULL);CHKERRQ(ierr);
  ierr = PetscOptionsGetInt(NULL,NULL,"-ny",&ny,NULL);CHKERRQ(ierr);

  PetscInt N = nx*ny;

  ierr = PetscPrintf(PETSC_COMM_WORLD,"\n Acoustic Modes in a 2D Rectangular Box, N=%D (%Dx%D grid)\n\n",N,nx,ny);CHKERRQ(ierr);

  ierr = MatCreate(PETSC_COMM_WORLD,&A);CHKERRQ(ierr);
  ierr = MatSetSizes(A,PETSC_DECIDE,PETSC_DECIDE,N,N);CHKERRQ(ierr);
  ierr = MatSetFromOptions(A);CHKERRQ(ierr);
  ierr = MatSetUp(A);CHKERRQ(ierr);

  PetscReal dx = Lx/nx;
  PetscReal dy = Ly/ny;
  PetscReal Dx = -(sos*sos)/(dx*dx);
  PetscReal Dy = -(sos*sos)/(dy*dy);

  fillMatA(nx, ny, Dx, Dy, A);

  std::chrono::high_resolution_clock::time_point ts, te;
  std::chrono::duration<double> time;
  ts = std::chrono::high_resolution_clock::now();

  Mat Ashell;
  Matrix ctx;
  ctx.A = A;

  PetscInt m,n;
  MatGetLocalSize(A, &m, &n);
  MatCreateShell(PETSC_COMM_WORLD, m, n, N, N, &ctx, &Ashell);
  MatShellSetOperation(Ashell, MATOP_MULT, reinterpret_cast<void(*)(void)>(matMultA));

  PetscInt shell = 1;
  ierr = PetscOptionsGetInt(NULL,NULL,"-shell",&shell,NULL);CHKERRQ(ierr);

  ierr = EPSCreate(PETSC_COMM_WORLD,&eps);CHKERRQ(ierr);

  if (shell == 1)
  {
    ierr = EPSSetOperators(eps,Ashell,NULL);CHKERRQ(ierr);
  }
  else
  {
    ierr = EPSSetOperators(eps,A,NULL);CHKERRQ(ierr);
  }

  ierr = EPSSetProblemType(eps,EPS_NHEP);CHKERRQ(ierr);

  PetscInt sinvert = 0;
  ierr = PetscOptionsGetInt(NULL,NULL,"-sinvert",&sinvert,NULL);CHKERRQ(ierr);
  if (sinvert == 1)
  {
    ierr = EPSGetST(eps,&st);CHKERRQ(ierr);
    ierr = STSetType(st,STSINVERT);
    ierr = STSetPreconditionerMat(st,A); CHKERRQ(ierr);

    ierr = STGetKSP(st,&ksp);CHKERRQ(ierr);
    ierr = KSPSetType(ksp,KSPBCGSL);CHKERRQ(ierr);
    ierr = KSPSetTolerances(ksp,1e-4,PETSC_DEFAULT,PETSC_DEFAULT,PETSC_DEFAULT);CHKERRQ(ierr);

    ierr = KSPGetPC(ksp, &pc);CHKERRQ(ierr);
    ierr = PCSetType(pc, PCBJACOBI);CHKERRQ(ierr);
  }

  PetscInt usejd = 0;
  ierr = PetscOptionsGetInt(NULL,NULL,"-usejd",&usejd,NULL);CHKERRQ(ierr);
  if (usejd == 1)
  {
    ierr = EPSSetType(eps,EPSJD);CHKERRQ(ierr);

    ierr = EPSGetST(eps,&st);CHKERRQ(ierr);
    ierr = STSetType(st,STPRECOND);
    ierr = STSetPreconditionerMat(st,A); CHKERRQ(ierr);

    ierr = STGetKSP(st,&ksp);CHKERRQ(ierr);
    ierr = KSPSetType(ksp,KSPBCGSL);CHKERRQ(ierr);
    ierr = KSPSetTolerances(ksp,1e-3,PETSC_DEFAULT,PETSC_DEFAULT,90);CHKERRQ(ierr);

    ierr = KSPGetPC(ksp, &pc);CHKERRQ(ierr);
    ierr = PCSetType(pc, PCBJACOBI);CHKERRQ(ierr);

    ierr = EPSSetExtraction(eps,EPS_HARMONIC);CHKERRQ(ierr);
  }

  //ierr = EPSSetWhichEigenpairs(eps,EPS_SMALLEST_MAGNITUDE);CHKERRQ(ierr);
  ierr = EPSSetTarget(eps,0.0);CHKERRQ(ierr);
    
  ierr = PetscOptionsGetInt(NULL,NULL,"-nev",&nev,NULL);CHKERRQ(ierr);
  ierr = EPSSetDimensions(eps, nev, PETSC_DEFAULT, PETSC_DEFAULT);CHKERRQ(ierr);
  ierr = EPSSetFromOptions(eps);CHKERRQ(ierr);

  PetscInt deflate = 0;
  ierr = PetscOptionsGetInt(NULL,NULL,"-deflate",&deflate,NULL);CHKERRQ(ierr);
  if (deflate == 1)
  {
    Vec x;
    ierr = MatCreateVecs(A,&x,NULL);CHKERRQ(ierr);
    ierr = VecSet(x,std::complex<double>(1.0,1.0));CHKERRQ(ierr);
    ierr = EPSSetDeflationSpace(eps,1,&x);CHKERRQ(ierr);
    ierr = VecDestroy(&x);CHKERRQ(ierr);
  }

  ierr = EPSSolve(eps);CHKERRQ(ierr);

  ierr = EPSGetType(eps,&type);CHKERRQ(ierr);
  ierr = PetscPrintf(PETSC_COMM_WORLD," Solution method: %s\n\n",type);CHKERRQ(ierr);
  ierr = EPSGetConverged(eps,&nev);CHKERRQ(ierr);
  ierr = PetscPrintf(PETSC_COMM_WORLD," Number of converged eigenvalues: %D\n", nev);CHKERRQ(ierr);

  ierr = PetscViewerPushFormat(PETSC_VIEWER_STDOUT_WORLD,PETSC_VIEWER_ASCII_INFO_DETAIL);CHKERRQ(ierr);
  ierr = EPSErrorView(eps,EPS_ERROR_RELATIVE,PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
  ierr = PetscViewerPopFormat(PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);

  PetscScalar kr, ki;
  Vec xr, xi;
  if (shell == 1)
  {
    MatCreateVecs(Ashell,NULL,&xr);
    MatCreateVecs(Ashell,NULL,&xi);
  }
  else
  {
    MatCreateVecs(A,NULL,&xr);
    MatCreateVecs(A,NULL,&xi);
  }

  PetscPrintf(PETSC_COMM_WORLD," -------------------------------------------------------------------\n");
  PetscPrintf(PETSC_COMM_WORLD," %10s %20s %20s\n", "Mode", "Frequency, Real", "Frequency, Imag");
  PetscPrintf(PETSC_COMM_WORLD," -------------------------------------------------------------------\n");
  for (int i = 0; i < nev; ++i)
    {
      EPSGetEigenpair(eps,i, &kr, &ki, xr, xi);
      PetscScalar omega = sqrt(-kr);
      PetscScalar f = omega/2.0/3.1415;
      PetscPrintf(PETSC_COMM_WORLD," %10D %20.3f %20.3f\n", i, std::real(f), std::imag(f));
    }
  PetscPrintf(PETSC_COMM_WORLD," -------------------------------------------------------------------\n");

  ierr = EPSDestroy(&eps);CHKERRQ(ierr);
  ierr = MatDestroy(&A);CHKERRQ(ierr);
  ierr = VecDestroy(&xr);CHKERRQ(ierr);
  ierr = VecDestroy(&xi);CHKERRQ(ierr);
  ierr = SlepcFinalize();

  te = std::chrono::high_resolution_clock::now();
  time = te - ts;

  std::cout<<"\n Computation finished in " << time.count() << " seconds" << std::endl;

  return ierr;
}