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

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;

  PetscInt nev = 6;

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

  ierr = MatCreate(PETSC_COMM_WORLD,&A);CHKERRQ(ierr);
  
  PetscViewer viewer;
  PetscViewerBinaryOpen(PETSC_COMM_WORLD,"MatA.gz",FILE_MODE_READ,&viewer);
  MatLoad(A,viewer);
  PetscViewerDestroy(&viewer);

  PetscInt N;
  MatGetSize(A, &N, NULL);

  ierr = PetscPrintf(PETSC_COMM_WORLD,"\n Acoustic Modes for matrix read from MatA.gz of size, N=%D \n\n", N);CHKERRQ(ierr);

  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);

  PetscInt usejd = 0;
  ierr = PetscOptionsGetInt(NULL,NULL,"-usejd",&usejd,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);
  }
  else 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_TARGET_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;
}