#include <iostream>
#include <petscksp.h>

using namespace std;


typedef struct {
  int val;
} MyCtx;



class ShellClass {
  Mat matShell;
  KSP ksp;
  PC pc;
  Vec x;
  Vec b;
  
public:
  void userMult(Mat Amat, Vec x, Vec y) {
    cout << "Inside userMult" << endl;
    
    MyCtx *ctx;
    MatShellGetContext(Amat, (void *) ctx);
    
    cout << "End userMult" << endl;
  }

  void solveShell() {
    // context
    MyCtx *ctx = new MyCtx;
    ctx->val = 42;

    // pc
    PCCreate(PETSC_COMM_WORLD, &pc);
    PCSetType(pc, PCNONE);

    // ksp
    KSPCreate(PETSC_COMM_WORLD, &ksp);
    KSPSetType(ksp, KSPCG);
    KSPSetPC(ksp, pc);
    KSPSetFromOptions(ksp);

    // matshell
    int m = 10;
    int n = 10;
    MatCreateShell(PETSC_COMM_WORLD, m, n, PETSC_DETERMINE, PETSC_DETERMINE, ctx, &matShell);
    MatShellSetOperation(matShell, MATOP_MULT, (void(*)(void))&ShellClass::userMult);


    // create vectors
    MatCreateVecs(matShell, &x, 0);
    VecDuplicate(x, &b);
    VecSet(b, 1.);

    // set operators
    KSPSetOperators(ksp, matShell, matShell);

    // solve (call to userMult)
    KSPSolve(ksp, b, x);
  }
};



int main(int argc, char** argv) {
  PetscInitialize(&argc, &argv, NULL, NULL);
    
  ShellClass foo;
  foo.solveShell();
  
  PetscFinalize();
  return 0;
}