// Copyright (c) 2017, Lawrence Livermore National Security, LLC. Produced at
// the Lawrence Livermore National Laboratory. LLNL-CODE-734707. All Rights
// reserved. See files LICENSE and NOTICE for details.
//
// This file is part of CEED, a collection of benchmarks, miniapps, software
// libraries and APIs for efficient high-order finite element and spectral
// element discretizations for exascale applications. For more information and
// source code availability see http://github.com/ceed.
//
// The CEED research is supported by the Exascale Computing Project 17-SC-20-SC,
// a collaborative effort of two U.S. Department of Energy organizations (Office
// of Science and the National Nuclear Security Administration) responsible for
// the planning and preparation of a capable exascale ecosystem, including
// software, applications, hardware, advanced system engineering and early
// testbed platforms, in support of the nation's exascale computing imperative.

#ifndef setup_h
#define setup_h

#include <stdbool.h>
#include <string.h>
#include <petsc.h>
#include <petscdmplex.h>
#include <petscfe.h>
#include <ceed.h>
#include "qfunctions/bps/common.h"
#include "qfunctions/bps/bp1.h"
#include "qfunctions/bps/bp2.h"
#include "qfunctions/bps/bp3.h"
#include "qfunctions/bps/bp4.h"

// -----------------------------------------------------------------------------
// PETSc Operator Structs
// -----------------------------------------------------------------------------

// Data for PETSc Matshell
typedef struct UserO_ *UserO;
struct UserO_ {
  MPI_Comm comm;
  DM dm;
  Vec Xloc, Yloc, diag;
  CeedVector xceed, yceed;
  CeedOperator op;
  Ceed ceed;
};

// Data for PETSc Interp/Restrict Matshells
typedef struct UserIR_ *UserIR;
struct UserIR_ {
  MPI_Comm comm;
  DM dmc, dmf;
  Vec Xloc, Yloc, mult;
  CeedVector ceedvecc, ceedvecf;
  CeedOperator op;
  Ceed ceed;
};

// -----------------------------------------------------------------------------
// libCEED Data Struct
// -----------------------------------------------------------------------------

// libCEED data struct for level
typedef struct CeedData_ *CeedData;
struct CeedData_ {
  Ceed ceed;
  CeedBasis basisx, basisu, basisctof;
  CeedElemRestriction Erestrictx, Erestrictu, Erestrictxi, Erestrictui,
                      Erestrictqdi;
  CeedQFunction qf_apply;
  CeedOperator op_apply, op_restrict, op_interp;
  CeedVector qdata, xceed, yceed;
};

// -----------------------------------------------------------------------------
// Command Line Options
// -----------------------------------------------------------------------------

// Coarsening options
typedef enum {
  COARSEN_UNIFORM = 0, COARSEN_LOGARITHMIC = 1
} coarsenType;
static const char *const coarsenTypes [] = {"uniform","logarithmic",
                                            "coarsenType","COARSEN",0
                                           };

// -----------------------------------------------------------------------------
// Boundary Conditions
// -----------------------------------------------------------------------------

// Diff boundary condition function
PetscErrorCode BCsDiff(PetscInt dim, PetscReal time, const PetscReal x[],
                       PetscInt ncompu, PetscScalar *u, void *ctx) {
  // *INDENT-OFF*
  #ifndef M_PI
  #define M_PI    3.14159265358979323846
  #endif
  // *INDENT-ON*
  const CeedScalar c[3] = { 0, 1., 2. };
  const CeedScalar k[3] = { 1., 2., 3. };

  PetscFunctionBeginUser;

  for (PetscInt i = 0; i < ncompu; i++)
    u[i] = sin(M_PI*(c[0] + k[0]*x[0])) *
           sin(M_PI*(c[1] + k[1]*x[1])) *
           sin(M_PI*(c[2] + k[2]*x[2]));

  PetscFunctionReturn(0);
}

// Mass boundary condition function
PetscErrorCode BCsMass(PetscInt dim, PetscReal time, const PetscReal x[],
                       PetscInt ncompu, PetscScalar *u, void *ctx) {
  PetscFunctionBeginUser;

  for (PetscInt i = 0; i < ncompu; i++)
    u[i] = PetscSqrtScalar(PetscSqr(x[0]) + PetscSqr(x[1]) +
                           PetscSqr(x[2]));

  PetscFunctionReturn(0);
}

// Create BC label
static PetscErrorCode CreateBCLabel(DM dm, const char name[]) {
  int ierr;
  DMLabel label;

  PetscFunctionBeginUser;

  ierr = DMCreateLabel(dm, name); CHKERRQ(ierr);
  ierr = DMGetLabel(dm, name, &label); CHKERRQ(ierr);
  ierr = DMPlexMarkBoundaryFaces(dm, 1, label); CHKERRQ(ierr);
  ierr = DMPlexLabelComplete(dm, label); CHKERRQ(ierr);

  PetscFunctionReturn(0);
}

// -----------------------------------------------------------------------------
// BP Option Data
// -----------------------------------------------------------------------------

// BP options
typedef enum {
  CEED_BP1 = 0, CEED_BP2 = 1, CEED_BP3 = 2,
  CEED_BP4 = 3, CEED_BP5 = 4, CEED_BP6 = 5
} bpType;
static const char *const bpTypes[] = {"bp1","bp2","bp3","bp4","bp5","bp6",
                                      "bpType","CEED_BP",0
                                     };

// BP specific data
typedef struct {
  CeedInt ncompu, qdatasize, qextra;
  CeedQFunctionUser setupgeo, setuprhs, apply, error;
  const char *setupgeofname, *setuprhsfname, *applyfname, *errorfname;
  CeedEvalMode inmode, outmode;
  CeedQuadMode qmode;
  PetscBool enforce_bc;
  PetscErrorCode (*bcs_func)(PetscInt, PetscReal, const PetscReal *,
                             PetscInt, PetscScalar *, void *);
} bpData;

bpData bpOptions[6] = {
  [CEED_BP1] = {
    .ncompu = 1,
    .qdatasize = 1,
    .qextra = 1,
    .setupgeo = SetupMassGeo,
    .setuprhs = SetupMassRhs,
    .apply = Mass,
    .error = Error,
    .setupgeofname = SetupMassGeo_loc,
    .setuprhsfname = SetupMassRhs_loc,
    .applyfname = Mass_loc,
    .errorfname = Error_loc,
    .inmode = CEED_EVAL_INTERP,
    .outmode = CEED_EVAL_INTERP,
    .qmode = CEED_GAUSS,
    .enforce_bc = false,
    .bcs_func = BCsMass
  },
  [CEED_BP2] = {
    .ncompu = 3,
    .qdatasize = 1,
    .qextra = 1,
    .setupgeo = SetupMassGeo,
    .setuprhs = SetupMassRhs3,
    .apply = Mass3,
    .error = Error3,
    .setupgeofname = SetupMassGeo_loc,
    .setuprhsfname = SetupMassRhs3_loc,
    .applyfname = Mass3_loc,
    .errorfname = Error3_loc,
    .inmode = CEED_EVAL_INTERP,
    .outmode = CEED_EVAL_INTERP,
    .qmode = CEED_GAUSS,
    .enforce_bc = false,
    .bcs_func = BCsMass
  },
  [CEED_BP3] = {
    .ncompu = 1,
    .qdatasize = 6,
    .qextra = 1,
    .setupgeo = SetupDiffGeo,
    .setuprhs = SetupDiffRhs,
    .apply = Diff,
    .error = Error,
    .setupgeofname = SetupDiffGeo_loc,
    .setuprhsfname = SetupDiffRhs_loc,
    .applyfname = Diff_loc,
    .errorfname = Error_loc,
    .inmode = CEED_EVAL_GRAD,
    .outmode = CEED_EVAL_GRAD,
    .qmode = CEED_GAUSS,
    .enforce_bc = true,
    .bcs_func = BCsDiff
  },
  [CEED_BP4] = {
    .ncompu = 3,
    .qdatasize = 6,
    .qextra = 1,
    .setupgeo = SetupDiffGeo,
    .setuprhs = SetupDiffRhs3,
    .apply = Diff3,
    .error = Error3,
    .setupgeofname = SetupDiffGeo_loc,
    .setuprhsfname = SetupDiffRhs3_loc,
    .applyfname = Diff_loc,
    .errorfname = Error3_loc,
    .inmode = CEED_EVAL_GRAD,
    .outmode = CEED_EVAL_GRAD,
    .qmode = CEED_GAUSS,
    .enforce_bc = true,
    .bcs_func = BCsDiff
  },
  [CEED_BP5] = {
    .ncompu = 1,
    .qdatasize = 6,
    .qextra = 0,
    .setupgeo = SetupDiffGeo,
    .setuprhs = SetupDiffRhs,
    .apply = Diff,
    .error = Error,
    .setupgeofname = SetupDiffGeo_loc,
    .setuprhsfname = SetupDiffRhs_loc,
    .applyfname = Diff_loc,
    .errorfname = Error_loc,
    .inmode = CEED_EVAL_GRAD,
    .outmode = CEED_EVAL_GRAD,
    .qmode = CEED_GAUSS_LOBATTO,
    .enforce_bc = true,
    .bcs_func = BCsDiff
  },
  [CEED_BP6] = {
    .ncompu = 3,
    .qdatasize = 6,
    .qextra = 0,
    .setupgeo = SetupDiffGeo,
    .setuprhs = SetupDiffRhs3,
    .apply = Diff3,
    .error = Error3,
    .setupgeofname = SetupDiffGeo_loc,
    .setuprhsfname = SetupDiffRhs3_loc,
    .applyfname = Diff_loc,
    .errorfname = Error3_loc,
    .inmode = CEED_EVAL_GRAD,
    .outmode = CEED_EVAL_GRAD,
    .qmode = CEED_GAUSS_LOBATTO,
    .enforce_bc = true,
    .bcs_func = BCsDiff
  }
};

// -----------------------------------------------------------------------------
// PETSc FE Boilerplate
// -----------------------------------------------------------------------------

// Create FE by degree
static int PetscFECreateByDegree(DM dm, PetscInt dim, PetscInt Nc,
                                 PetscBool isSimplex, const char prefix[],
                                 PetscInt order, PetscFE *fem) {
  PetscQuadrature q, fq;
  DM              K;
  PetscSpace      P;
  PetscDualSpace  Q;
  PetscInt        quadPointsPerEdge;
  PetscBool       tensor = isSimplex ? PETSC_FALSE : PETSC_TRUE;
  PetscErrorCode  ierr;

  PetscFunctionBeginUser;
  /* Create space */
  ierr = PetscSpaceCreate(PetscObjectComm((PetscObject) dm), &P); CHKERRQ(ierr);
  ierr = PetscObjectSetOptionsPrefix((PetscObject) P, prefix); CHKERRQ(ierr);
  ierr = PetscSpacePolynomialSetTensor(P, tensor); CHKERRQ(ierr);
  ierr = PetscSpaceSetFromOptions(P); CHKERRQ(ierr);
  ierr = PetscSpaceSetNumComponents(P, Nc); CHKERRQ(ierr);
  ierr = PetscSpaceSetNumVariables(P, dim); CHKERRQ(ierr);
  ierr = PetscSpaceSetDegree(P, order, order); CHKERRQ(ierr);
  ierr = PetscSpaceSetUp(P); CHKERRQ(ierr);
  ierr = PetscSpacePolynomialGetTensor(P, &tensor); CHKERRQ(ierr);
  /* Create dual space */
  ierr = PetscDualSpaceCreate(PetscObjectComm((PetscObject) dm), &Q);
  CHKERRQ(ierr);
  ierr = PetscDualSpaceSetType(Q,PETSCDUALSPACELAGRANGE); CHKERRQ(ierr);
  ierr = PetscObjectSetOptionsPrefix((PetscObject) Q, prefix); CHKERRQ(ierr);
  ierr = PetscDualSpaceCreateReferenceCell(Q, dim, isSimplex, &K); CHKERRQ(ierr);
  ierr = PetscDualSpaceSetDM(Q, K); CHKERRQ(ierr);
  ierr = DMDestroy(&K); CHKERRQ(ierr);
  ierr = PetscDualSpaceSetNumComponents(Q, Nc); CHKERRQ(ierr);
  ierr = PetscDualSpaceSetOrder(Q, order); CHKERRQ(ierr);
  ierr = PetscDualSpaceLagrangeSetTensor(Q, tensor); CHKERRQ(ierr);
  ierr = PetscDualSpaceSetFromOptions(Q); CHKERRQ(ierr);
  ierr = PetscDualSpaceSetUp(Q); CHKERRQ(ierr);
  /* Create element */
  ierr = PetscFECreate(PetscObjectComm((PetscObject) dm), fem); CHKERRQ(ierr);
  ierr = PetscObjectSetOptionsPrefix((PetscObject) *fem, prefix); CHKERRQ(ierr);
  ierr = PetscFESetFromOptions(*fem); CHKERRQ(ierr);
  ierr = PetscFESetBasisSpace(*fem, P); CHKERRQ(ierr);
  ierr = PetscFESetDualSpace(*fem, Q); CHKERRQ(ierr);
  ierr = PetscFESetNumComponents(*fem, Nc); CHKERRQ(ierr);
  ierr = PetscFESetUp(*fem); CHKERRQ(ierr);
  ierr = PetscSpaceDestroy(&P); CHKERRQ(ierr);
  ierr = PetscDualSpaceDestroy(&Q); CHKERRQ(ierr);
  /* Create quadrature */
  quadPointsPerEdge = PetscMax(order + 1,1);
  if (isSimplex) {
    ierr = PetscDTGaussJacobiQuadrature(dim,   1, quadPointsPerEdge, -1.0, 1.0,
                                        &q); CHKERRQ(ierr);
    ierr = PetscDTGaussJacobiQuadrature(dim-1, 1, quadPointsPerEdge, -1.0, 1.0,
                                        &fq); CHKERRQ(ierr);
  } else {
    ierr = PetscDTGaussTensorQuadrature(dim,   1, quadPointsPerEdge, -1.0, 1.0,
                                        &q); CHKERRQ(ierr);
    ierr = PetscDTGaussTensorQuadrature(dim-1, 1, quadPointsPerEdge, -1.0, 1.0,
                                        &fq); CHKERRQ(ierr);
  }
  ierr = PetscFESetQuadrature(*fem, q); CHKERRQ(ierr);
  ierr = PetscFESetFaceQuadrature(*fem, fq); CHKERRQ(ierr);
  ierr = PetscQuadratureDestroy(&q); CHKERRQ(ierr);
  ierr = PetscQuadratureDestroy(&fq); CHKERRQ(ierr);

  PetscFunctionReturn(0);
}

// -----------------------------------------------------------------------------
// PETSc Setup for Level
// -----------------------------------------------------------------------------

// This function sets up a DM for a given degree
static int SetupDMByDegree(DM dm, PetscInt degree, PetscInt ncompu,
                           bpType bpChoice) {
  PetscInt ierr, dim, marker_ids[1] = {1};
  PetscFE fe;

  PetscFunctionBeginUser;

  // Setup FE
  ierr = DMGetDimension(dm, &dim); CHKERRQ(ierr);
  ierr = PetscFECreateByDegree(dm, dim, ncompu, PETSC_FALSE, NULL, degree, &fe);
  CHKERRQ(ierr);
  ierr = DMSetFromOptions(dm); CHKERRQ(ierr);
  ierr = DMAddField(dm, NULL, (PetscObject)fe); CHKERRQ(ierr);

  // Setup DM
  ierr = DMCreateDS(dm); CHKERRQ(ierr);
  if (bpOptions[bpChoice].enforce_bc) {
    PetscBool hasLabel;
    DMHasLabel(dm, "marker", &hasLabel);
    if (!hasLabel) {CreateBCLabel(dm, "marker");}
    ierr = DMAddBoundary(dm, DM_BC_ESSENTIAL, "wall", "marker", 0, 0, NULL,
                         (void(*)(void))bpOptions[bpChoice].bcs_func,
                         1, marker_ids, NULL);
    CHKERRQ(ierr);
  }
  ierr = DMPlexSetClosurePermutationTensor(dm, PETSC_DETERMINE, NULL);
  CHKERRQ(ierr);
  ierr = PetscFEDestroy(&fe); CHKERRQ(ierr);

  PetscFunctionReturn(0);
}

// -----------------------------------------------------------------------------
// libCEED Setup for Level
// -----------------------------------------------------------------------------

// Destroy libCEED operator objects
static PetscErrorCode CeedDataDestroy(CeedInt i, CeedData data) {
  PetscInt ierr;

  CeedVectorDestroy(&data->qdata);
  CeedVectorDestroy(&data->xceed);
  CeedVectorDestroy(&data->yceed);
  CeedBasisDestroy(&data->basisx);
  CeedBasisDestroy(&data->basisu);
  CeedElemRestrictionDestroy(&data->Erestrictu);
  CeedElemRestrictionDestroy(&data->Erestrictx);
  CeedElemRestrictionDestroy(&data->Erestrictui);
  CeedElemRestrictionDestroy(&data->Erestrictxi);
  CeedElemRestrictionDestroy(&data->Erestrictqdi);
  CeedQFunctionDestroy(&data->qf_apply);
  CeedOperatorDestroy(&data->op_apply);
  if (i > 0) {
    CeedOperatorDestroy(&data->op_interp);
    CeedBasisDestroy(&data->basisctof);
    CeedOperatorDestroy(&data->op_restrict);
  }
  ierr = PetscFree(data); CHKERRQ(ierr);

  PetscFunctionReturn(0);
}

// Get CEED restriction data from DMPlex
static int CreateRestrictionPlex(Ceed ceed, CeedInt P, CeedInt ncomp,
                                 CeedElemRestriction *Erestrict, DM dm) {
  PetscInt ierr;
  PetscInt c, cStart, cEnd, nelem, nnodes, *erestrict, eoffset;
  PetscSection section;
  Vec Uloc;

  PetscFunctionBeginUser;

  // Get Nelem
  ierr = DMGetSection(dm, &section); CHKERRQ(ierr);
  ierr = DMPlexGetHeightStratum(dm, 0, &cStart,& cEnd); CHKERRQ(ierr);
  nelem = cEnd - cStart;

  // Get indices
  ierr = PetscMalloc1(nelem*P*P*P, &erestrict); CHKERRQ(ierr);
  for (c=cStart, eoffset=0; c<cEnd; c++) {
    PetscInt numindices, *indices, i;
    ierr = DMPlexGetClosureIndices(dm, section, section, c, &numindices,
                                   &indices, NULL); CHKERRQ(ierr);
    for (i=0; i<numindices; i+=ncomp) {
      for (PetscInt j=0; j<ncomp; j++) {
        if (indices[i+j] != indices[i] + (PetscInt)(copysign(j, indices[i])))
          SETERRQ1(PETSC_COMM_SELF, PETSC_ERR_ARG_INCOMP,
                   "Cell %D closure indices not interlaced", c);
      }
      // Essential boundary conditions are encoded as -(loc+1)
      PetscInt loc = indices[i] >= 0 ? indices[i] : -(indices[i] + 1);
      erestrict[eoffset++] = loc/ncomp;
    }
    ierr = DMPlexRestoreClosureIndices(dm, section, section, c, &numindices,
                                       &indices, NULL); CHKERRQ(ierr);
  }

  // Setup CEED restriction
  ierr = DMGetLocalVector(dm, &Uloc); CHKERRQ(ierr);
  ierr = VecGetLocalSize(Uloc, &nnodes); CHKERRQ(ierr);

  ierr = DMRestoreLocalVector(dm, &Uloc); CHKERRQ(ierr);
  CeedElemRestrictionCreate(ceed, nelem, P*P*P, nnodes/ncomp, ncomp,
                            CEED_MEM_HOST, CEED_COPY_VALUES, erestrict,
                            Erestrict);
  ierr = PetscFree(erestrict); CHKERRQ(ierr);

  PetscFunctionReturn(0);
}

// Set up libCEED for a given degree
static int SetupLibceedByDegree(DM dm, Ceed ceed, CeedInt degree, CeedInt dim,
                                CeedInt qextra, PetscInt ncompu, PetscInt gsize,
                                PetscInt xlsize, bpType bpChoice, CeedData data,
                                PetscBool setup_rhs, CeedVector rhsceed,
                                CeedVector *target) {
  int ierr;
  DM dmcoord;
  PetscSection section;
  Vec coords;
  const PetscScalar *coordArray;
  CeedBasis basisx, basisu;
  CeedElemRestriction Erestrictx, Erestrictu, Erestrictxi,
                      Erestrictui, Erestrictqdi;
  CeedQFunction qf_setupgeo, qf_apply;
  CeedOperator op_setupgeo, op_apply;
  CeedVector xcoord, qdata, xceed, yceed;
  CeedInt qdatasize = bpOptions[bpChoice].qdatasize, ncompx = dim, P, Q,
          cStart, cEnd, nelem;

  // CEED bases
  P = degree + 1;
  Q = P + qextra;
  CeedBasisCreateTensorH1Lagrange(ceed, dim, ncompu, P, Q,
                                  bpOptions[bpChoice].qmode, &basisu);
  CeedBasisCreateTensorH1Lagrange(ceed, dim, ncompx, 2, Q,
                                  bpOptions[bpChoice].qmode, &basisx);

  // CEED restrictions
  ierr = DMGetCoordinateDM(dm, &dmcoord); CHKERRQ(ierr);
  ierr = DMPlexSetClosurePermutationTensor(dmcoord, PETSC_DETERMINE, NULL);
  CHKERRQ(ierr);

  CreateRestrictionPlex(ceed, 2, ncompx, &Erestrictx, dmcoord);
  CreateRestrictionPlex(ceed, P, ncompu, &Erestrictu, dm);

  ierr = DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd); CHKERRQ(ierr);
  nelem = cEnd - cStart;

  CeedElemRestrictionCreateIdentity(ceed, nelem, Q*Q*Q, nelem*Q*Q*Q, ncompu,
                                    &Erestrictui); CHKERRQ(ierr);
  CeedElemRestrictionCreateIdentity(ceed, nelem, Q*Q*Q, nelem*Q*Q*Q,
                                    qdatasize, &Erestrictqdi); CHKERRQ(ierr);
  CeedElemRestrictionCreateIdentity(ceed, nelem, Q*Q*Q, nelem*Q*Q*Q, ncompx,
                                    &Erestrictxi); CHKERRQ(ierr);

  // Element coordinates
  ierr = DMGetCoordinatesLocal(dm, &coords); CHKERRQ(ierr);
  ierr = VecGetArrayRead(coords, &coordArray); CHKERRQ(ierr);
  ierr = DMGetSection(dmcoord, &section); CHKERRQ(ierr);

  CeedElemRestrictionCreateVector(Erestrictx, &xcoord, NULL);
  CeedVectorSetArray(xcoord, CEED_MEM_HOST, CEED_COPY_VALUES,
                     (PetscScalar *)coordArray);
  ierr = VecRestoreArrayRead(coords, &coordArray); CHKERRQ(ierr);

  // Create the persistent vectors that will be needed in setup and apply
  CeedInt nqpts;
  CeedBasisGetNumQuadraturePoints(basisu, &nqpts);
  CeedVectorCreate(ceed, qdatasize*nelem*nqpts, &qdata);
  CeedVectorCreate(ceed, xlsize, &xceed);
  CeedVectorCreate(ceed, xlsize, &yceed);

  // Create the Q-function that builds the operator (i.e. computes its
  // quadrature data) and set its context data
  CeedQFunctionCreateInterior(ceed, 1, bpOptions[bpChoice].setupgeo,
                              bpOptions[bpChoice].setupgeofname, &qf_setupgeo);
  CeedQFunctionAddInput(qf_setupgeo, "dx", ncompx*dim, CEED_EVAL_GRAD);
  CeedQFunctionAddInput(qf_setupgeo, "weight", 1, CEED_EVAL_WEIGHT);
  CeedQFunctionAddOutput(qf_setupgeo, "qdata", qdatasize, CEED_EVAL_NONE);

  // Set up PDE operator
  CeedInt inscale = bpOptions[bpChoice].inmode==CEED_EVAL_GRAD ? dim : 1;
  CeedInt outscale = bpOptions[bpChoice].outmode==CEED_EVAL_GRAD ? dim : 1;
  CeedQFunctionCreateInterior(ceed, 1, bpOptions[bpChoice].apply,
                              bpOptions[bpChoice].applyfname, &qf_apply);
  CeedQFunctionAddInput(qf_apply, "u", ncompu*inscale,
                        bpOptions[bpChoice].inmode);
  CeedQFunctionAddInput(qf_apply, "qdata", qdatasize, CEED_EVAL_NONE);
  CeedQFunctionAddOutput(qf_apply, "v", ncompu*outscale,
                         bpOptions[bpChoice].outmode);

  // Create the operator that builds the quadrature data for the operator
  CeedOperatorCreate(ceed, qf_setupgeo, CEED_QFUNCTION_NONE,
                     CEED_QFUNCTION_NONE, &op_setupgeo);
  CeedOperatorSetField(op_setupgeo, "dx", Erestrictx, CEED_TRANSPOSE,
                       basisx, CEED_VECTOR_ACTIVE);
  CeedOperatorSetField(op_setupgeo, "weight", Erestrictxi, CEED_NOTRANSPOSE,
                       basisx, CEED_VECTOR_NONE);
  CeedOperatorSetField(op_setupgeo, "qdata", Erestrictqdi, CEED_NOTRANSPOSE,
                       CEED_BASIS_COLLOCATED, CEED_VECTOR_ACTIVE);

  // Create the operator
  CeedOperatorCreate(ceed, qf_apply, CEED_QFUNCTION_NONE, CEED_QFUNCTION_NONE,
                     &op_apply);
  CeedOperatorSetField(op_apply, "u", Erestrictu, CEED_TRANSPOSE,
                       basisu, CEED_VECTOR_ACTIVE);
  CeedOperatorSetField(op_apply, "qdata", Erestrictqdi, CEED_NOTRANSPOSE,
                       CEED_BASIS_COLLOCATED, qdata);
  CeedOperatorSetField(op_apply, "v", Erestrictu, CEED_TRANSPOSE,
                       basisu, CEED_VECTOR_ACTIVE);

  // Setup qdata
  CeedOperatorApply(op_setupgeo, xcoord, qdata, CEED_REQUEST_IMMEDIATE);

  // Set up RHS if needed
  if (setup_rhs) {
    CeedQFunction qf_setuprhs;
    CeedOperator op_setuprhs;
    CeedVectorCreate(ceed, nelem*nqpts*ncompu, target);

    // Create the q-function that sets up the RHS and true solution
    CeedQFunctionCreateInterior(ceed, 1, bpOptions[bpChoice].setuprhs,
                                bpOptions[bpChoice].setuprhsfname, &qf_setuprhs);
    CeedQFunctionAddInput(qf_setuprhs, "x", dim, CEED_EVAL_INTERP);
    CeedQFunctionAddInput(qf_setuprhs, "dx", ncompx*dim, CEED_EVAL_GRAD);
    CeedQFunctionAddInput(qf_setuprhs, "weight", 1, CEED_EVAL_WEIGHT);
    CeedQFunctionAddOutput(qf_setuprhs, "true_soln", ncompu, CEED_EVAL_NONE);
    CeedQFunctionAddOutput(qf_setuprhs, "rhs", ncompu, CEED_EVAL_INTERP);

    // Create the operator that builds the RHS and true solution
    CeedOperatorCreate(ceed, qf_setuprhs, CEED_QFUNCTION_NONE,
                       CEED_QFUNCTION_NONE, &op_setuprhs);
    CeedOperatorSetField(op_setuprhs, "x", Erestrictx, CEED_TRANSPOSE,
                         basisx, CEED_VECTOR_ACTIVE);
    CeedOperatorSetField(op_setuprhs, "dx", Erestrictx, CEED_TRANSPOSE,
                         basisx, CEED_VECTOR_ACTIVE);
    CeedOperatorSetField(op_setuprhs, "weight", Erestrictxi, CEED_NOTRANSPOSE,
                         basisx, CEED_VECTOR_NONE);
    CeedOperatorSetField(op_setuprhs, "true_soln", Erestrictui, CEED_NOTRANSPOSE,
                         CEED_BASIS_COLLOCATED, *target);
    CeedOperatorSetField(op_setuprhs, "rhs", Erestrictu, CEED_TRANSPOSE,
                         basisu, CEED_VECTOR_ACTIVE);

    // Setup RHS and target
    CeedOperatorApply(op_setuprhs, xcoord, rhsceed, CEED_REQUEST_IMMEDIATE);
    CeedVectorSyncArray(rhsceed, CEED_MEM_HOST);

    // Cleanup
    CeedQFunctionDestroy(&qf_setuprhs);
    CeedOperatorDestroy(&op_setuprhs);
  }

  // Cleanup
  CeedQFunctionDestroy(&qf_setupgeo);
  CeedOperatorDestroy(&op_setupgeo);
  CeedVectorDestroy(&xcoord);

  // Save libCEED data required for level
  data->basisx = basisx; data->basisu = basisu;
  data->Erestrictx = Erestrictx;
  data->Erestrictu = Erestrictu;
  data->Erestrictxi = Erestrictxi;
  data->Erestrictui = Erestrictui;
  data->Erestrictqdi = Erestrictqdi;
  data->qf_apply = qf_apply;
  data->op_apply = op_apply;
  data->qdata = qdata;
  data->xceed = xceed;
  data->yceed = yceed;

  PetscFunctionReturn(0);
}

#ifdef multigrid
// Setup libCEED level transfer operator objects
static PetscErrorCode CeedLevelTransferSetup(Ceed ceed, CeedInt numlevels,
    CeedInt ncompu, bpType bpChoice, CeedData *data, CeedInt *leveldegrees,
    CeedQFunction qf_restrict, CeedQFunction qf_prolong) {
  // Return early if numlevels=1
  if (numlevels==1)
    PetscFunctionReturn(0);

  // Set up each level
  for (CeedInt i=1; i<numlevels; i++) {
    // P coarse and P fine
    CeedInt Pc = leveldegrees[i-1] + 1;
    CeedInt Pf = leveldegrees[i] + 1;

    // Restriction - Fine to corse
    CeedBasis basisctof;
    CeedOperator op_restrict;

    // Basis
    CeedBasisCreateTensorH1Lagrange(ceed, 3, ncompu, Pc, Pf,
                                    CEED_GAUSS_LOBATTO, &basisctof);

    // Create the restriction operator
    CeedOperatorCreate(ceed, qf_restrict, CEED_QFUNCTION_NONE,
                       CEED_QFUNCTION_NONE, &op_restrict);
    CeedOperatorSetField(op_restrict, "input", data[i]->Erestrictu,
                         CEED_NOTRANSPOSE, CEED_BASIS_COLLOCATED,
                         CEED_VECTOR_ACTIVE);
    CeedOperatorSetField(op_restrict, "output", data[i-1]->Erestrictu,
                         CEED_TRANSPOSE, basisctof, CEED_VECTOR_ACTIVE);

    // Save libCEED data required for level
    data[i]->basisctof = basisctof;
    data[i]->op_restrict = op_restrict;

    // Interpolation - Corse to fine
    CeedOperator op_interp;

    // Create the prolongation operator
    CeedOperatorCreate(ceed, qf_prolong, CEED_QFUNCTION_NONE,
                       CEED_QFUNCTION_NONE, &op_interp);
    CeedOperatorSetField(op_interp, "input", data[i-1]->Erestrictu,
                         CEED_NOTRANSPOSE, basisctof, CEED_VECTOR_ACTIVE);
    CeedOperatorSetField(op_interp, "output", data[i]->Erestrictu,
                         CEED_TRANSPOSE, CEED_BASIS_COLLOCATED,
                         CEED_VECTOR_ACTIVE);

    // Save libCEED data required for level
    data[i]->op_interp = op_interp;
  }

  PetscFunctionReturn(0);
}
#endif

// -----------------------------------------------------------------------------
// Mat Shell Functions
// -----------------------------------------------------------------------------

#ifdef multigrid
// This function returns the computed diagonal of the operator
static PetscErrorCode MatGetDiag(Mat A, Vec D) {
  PetscErrorCode ierr;
  UserO user;

  PetscFunctionBeginUser;
  ierr = MatShellGetContext(A, &user); CHKERRQ(ierr);

  ierr = VecCopy(user->diag, D); CHKERRQ(ierr);

  PetscFunctionReturn(0);
}
#endif

// This function uses libCEED to compute the action of the Laplacian with
// Dirichlet boundary conditions
static PetscErrorCode MatMult_Ceed(Mat A, Vec X, Vec Y) {
  PetscErrorCode ierr;
  UserO user;
  PetscScalar *x, *y;

  PetscFunctionBeginUser;
  ierr = MatShellGetContext(A, &user); CHKERRQ(ierr);

  // Global-to-local
  ierr = DMGlobalToLocalBegin(user->dm, X, INSERT_VALUES, user->Xloc);
  CHKERRQ(ierr);
  ierr = DMGlobalToLocalEnd(user->dm, X, INSERT_VALUES, user->Xloc);
  CHKERRQ(ierr);
  ierr = VecZeroEntries(user->Yloc); CHKERRQ(ierr);

  // Setup CEED vectors
  ierr = VecGetArrayRead(user->Xloc, (const PetscScalar **)&x); CHKERRQ(ierr);
  ierr = VecGetArray(user->Yloc, &y); CHKERRQ(ierr);
  CeedVectorSetArray(user->xceed, CEED_MEM_HOST, CEED_USE_POINTER, x);
  CeedVectorSetArray(user->yceed, CEED_MEM_HOST, CEED_USE_POINTER, y);

  // Apply CEED operator
  CeedOperatorApply(user->op, user->xceed, user->yceed, CEED_REQUEST_IMMEDIATE);
  CeedVectorSyncArray(user->yceed, CEED_MEM_HOST);

  // Restore PETSc vectors
  ierr = VecRestoreArrayRead(user->Xloc, (const PetscScalar **)&x);
  CHKERRQ(ierr);
  ierr = VecRestoreArray(user->Yloc, &y); CHKERRQ(ierr);

  // Local-to-global
  ierr = VecZeroEntries(Y); CHKERRQ(ierr);
  ierr = DMLocalToGlobalBegin(user->dm, user->Yloc, ADD_VALUES, Y);
  CHKERRQ(ierr);
  ierr = DMLocalToGlobalEnd(user->dm, user->Yloc, ADD_VALUES, Y);
  CHKERRQ(ierr);

  PetscFunctionReturn(0);
}

#ifdef multigrid
// This function uses libCEED to compute the action of the interp operator
static PetscErrorCode MatMult_Interp(Mat A, Vec X, Vec Y) {
  PetscErrorCode ierr;
  UserIR user;
  PetscScalar *x, *y;

  PetscFunctionBeginUser;
  ierr = MatShellGetContext(A, &user); CHKERRQ(ierr);

  // Global-to-local
  ierr = VecZeroEntries(user->Xloc); CHKERRQ(ierr);
  ierr = DMGlobalToLocalBegin(user->dmc, X, INSERT_VALUES, user->Xloc);
  CHKERRQ(ierr);
  ierr = DMGlobalToLocalEnd(user->dmc, X, INSERT_VALUES, user->Xloc);
  CHKERRQ(ierr);
  ierr = VecZeroEntries(user->Yloc); CHKERRQ(ierr);

  // Setup CEED vectors
  ierr = VecGetArrayRead(user->Xloc, (const PetscScalar **)&x); CHKERRQ(ierr);
  ierr = VecGetArray(user->Yloc, &y); CHKERRQ(ierr);
  CeedVectorSetArray(user->ceedvecc, CEED_MEM_HOST, CEED_USE_POINTER, x);
  CeedVectorSetArray(user->ceedvecf, CEED_MEM_HOST, CEED_USE_POINTER, y);

  // Apply CEED operator
  CeedOperatorApply(user->op, user->ceedvecc, user->ceedvecf,
                    CEED_REQUEST_IMMEDIATE);
  CeedVectorSyncArray(user->ceedvecf, CEED_MEM_HOST);

  // Restore PETSc vectors
  ierr = VecRestoreArrayRead(user->Xloc, (const PetscScalar **)&x);
  CHKERRQ(ierr);
  ierr = VecRestoreArray(user->Yloc, &y); CHKERRQ(ierr);

  // Multiplicity
  ierr = VecPointwiseMult(user->Yloc, user->Yloc, user->mult);

  // Local-to-global
  ierr = VecZeroEntries(Y); CHKERRQ(ierr);
  ierr = DMLocalToGlobalBegin(user->dmf, user->Yloc, ADD_VALUES, Y);
  CHKERRQ(ierr);
  ierr = DMLocalToGlobalEnd(user->dmf, user->Yloc, ADD_VALUES, Y);
  CHKERRQ(ierr);

  PetscFunctionReturn(0);
}

// This function uses libCEED to compute the action of the restriction operator
static PetscErrorCode MatMult_Restrict(Mat A, Vec X, Vec Y) {
  PetscErrorCode ierr;
  UserIR user;
  PetscScalar *x, *y;

  PetscFunctionBeginUser;
  ierr = MatShellGetContext(A, &user); CHKERRQ(ierr);

  // Global-to-local
  ierr = VecZeroEntries(user->Xloc); CHKERRQ(ierr);
  ierr = DMGlobalToLocalBegin(user->dmf, X, INSERT_VALUES, user->Xloc);
  CHKERRQ(ierr);
  ierr = DMGlobalToLocalEnd(user->dmf, X, INSERT_VALUES, user->Xloc);
  CHKERRQ(ierr);
  ierr = VecZeroEntries(user->Yloc); CHKERRQ(ierr);

  // Multiplicity
  ierr = VecPointwiseMult(user->Xloc, user->Xloc, user->mult); CHKERRQ(ierr);

  // Setup CEED vectors
  ierr = VecGetArrayRead(user->Xloc, (const PetscScalar **)&x); CHKERRQ(ierr);
  ierr = VecGetArray(user->Yloc, &y); CHKERRQ(ierr);
  CeedVectorSetArray(user->ceedvecf, CEED_MEM_HOST, CEED_USE_POINTER, x);
  CeedVectorSetArray(user->ceedvecc, CEED_MEM_HOST, CEED_USE_POINTER, y);

  // Apply CEED operator
  CeedOperatorApply(user->op, user->ceedvecf, user->ceedvecc,
                    CEED_REQUEST_IMMEDIATE);
  CeedVectorSyncArray(user->ceedvecc, CEED_MEM_HOST);

  // Restore PETSc vectors
  ierr = VecRestoreArrayRead(user->Xloc, (const PetscScalar **)&x);
  CHKERRQ(ierr);
  ierr = VecRestoreArray(user->Yloc, &y); CHKERRQ(ierr);

  // Local-to-global
  ierr = VecZeroEntries(Y); CHKERRQ(ierr);
  ierr = DMLocalToGlobalBegin(user->dmc, user->Yloc, ADD_VALUES, Y);
  CHKERRQ(ierr);
  ierr = DMLocalToGlobalEnd(user->dmc, user->Yloc, ADD_VALUES, Y);
  CHKERRQ(ierr);

  PetscFunctionReturn(0);
}
#endif

// This function calculates the error in the final solution
static PetscErrorCode ComputeErrorMax(UserO user, CeedOperator op_error,
                                      Vec X, CeedVector target,
                                      PetscReal *maxerror) {
  PetscErrorCode ierr;
  PetscScalar *x;
  CeedVector collocated_error;
  CeedInt length;

  PetscFunctionBeginUser;
  CeedVectorGetLength(target, &length);
  CeedVectorCreate(user->ceed, length, &collocated_error);

  // Global-to-local
  ierr = DMGlobalToLocal(user->dm, X, INSERT_VALUES, user->Xloc); CHKERRQ(ierr);

  // Setup CEED vector
  ierr = VecGetArrayRead(user->Xloc, (const PetscScalar **)&x); CHKERRQ(ierr);
  CeedVectorSetArray(user->xceed, CEED_MEM_HOST, CEED_USE_POINTER, x);

  // Apply CEED operator
  CeedOperatorApply(op_error, user->xceed, collocated_error,
                    CEED_REQUEST_IMMEDIATE);

  // Restore PETSc vector
  VecRestoreArrayRead(user->Xloc, (const PetscScalar **)&x); CHKERRQ(ierr);

  // Reduce max error
  *maxerror = 0;
  const CeedScalar *e;
  CeedVectorGetArrayRead(collocated_error, CEED_MEM_HOST, &e);
  for (CeedInt i=0; i<length; i++) {
    *maxerror = PetscMax(*maxerror, PetscAbsScalar(e[i]));
  }
  CeedVectorRestoreArrayRead(collocated_error, &e);
  ierr = MPI_Allreduce(MPI_IN_PLACE, maxerror,
                       1, MPIU_REAL, MPIU_MAX, user->comm); CHKERRQ(ierr);

  // Cleanup
  CeedVectorDestroy(&collocated_error);

  PetscFunctionReturn(0);
}
#endif


/* Function in common for mlfit1, mlfitN, etc ...
 */

#include <stdlib.h>
#include <stdio.h>
#include <assert.h>
#include <stdint.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_sort.h>
#include <gsl/gsl_statistics.h>

#include "mlfit.h"
#include "blit3.h"

#ifndef verbose
#define verbose 0
#endif

double estimateNphot(double *V, size_t Vm)
  // V is a Vm x Vm matrix
  // The number of photons locally as sum(V-min(V))
{
  double s = 0;
  double minV = INFINITY;

  for(size_t kk=0; kk<Vm*Vm; kk++)
    s+=V[kk];

  for(size_t kk=0; kk<Vm*Vm; kk++)
    minV = GSL_MIN(V[kk], minV);

  s = s-minV*Vm*Vm;

  //printf("estimateNphot: %f\n", s);
  return s;
}

int double_cmp(const void * a, const void * b)
{
  if (*(const double*)a > *(const double*)b) {
    return 1;
  } else {
    if (*(const double*)a < *(const double*)b) {
      return -1;
    } else {
    return 0;
  }
}
}

double estimateBGV(double *V, size_t Vm, size_t Vn, size_t Vp, double * D)
  // estimate the background level in V [VmxVnxVp] around the point D
  // (x,y,z)
{

  int32_t radius = 7;
  double * ePixels = malloc(sizeof(double)*6*(2*radius+1)*(2*radius+1)); // a few more than we need

  int32_t x = nearbyint(D[0]);
  int32_t y = nearbyint(D[1]);
  int32_t z = nearbyint(D[2]);

  // Loop over the box, and copy the edge pixels to ePixels
  size_t nPixels = 0;
  for(int32_t zz = z-radius; zz<=z+radius; zz++)
    for(int32_t yy = y-radius; yy<=y+radius; yy++)
     for(int32_t xx = x-radius; xx<=x+radius; xx++)
       if((abs(zz-z)+abs(yy-y)+abs(xx-x)) == radius)
         if(xx>=0 && yy>=0 && zz>=0 && xx< (int32_t) Vm && yy<(int32_t) Vn && zz<(int32_t) Vp)
           ePixels[nPixels++] = V[xx + yy*Vm + zz*Vm*Vn];

#if verbose > 0
  printf("%lu ePixels\n", nPixels);
#endif

  double median = 0;
  if(nPixels>0)
  {
  qsort(ePixels, nPixels, sizeof(double),  double_cmp);
  gsl_sort(ePixels, 1, nPixels);
  median = gsl_stats_median_from_sorted_data(ePixels, 1, nPixels);
  //median = (double) nPixels;
#if verbose > 0
  printf("Median: %f\n", median);
#endif
  }
  else
  { median = 0;}

  free(ePixels);
  return(median);
  //return((double) q);
}

double estimateBG(double *V, size_t Vm)
  // Estimate the background level as the median of the edge pixels
{

  size_t nEdge = 4*Vm-4;
  double * ePixels = malloc(nEdge*sizeof(double));

  size_t pos = 0;
  for(uint32_t kk = 0; kk<Vm; kk++, pos++) // Left
    ePixels[pos] = V[kk];
  for(uint32_t kk = 0; kk<Vm; kk++, pos++) // Right
    ePixels[pos] = V[kk+(Vm-1)*Vm];
  for(uint32_t kk = 1; kk<Vm-1; kk++, pos++) // Top
    ePixels[pos] = V[Vm*kk];
  for(uint32_t kk = 1; kk<Vm-1; kk++, pos++) // Bottom
    ePixels[pos] = (double) V[Vm-1+Vm*kk];

  assert(pos == 4*Vm-4);

  if(0) {
    for(uint32_t kk = 0; kk<nEdge; kk++)
      printf("%2d :%f\n", kk, ePixels[kk]);
    printf("\n");
  }

//  qsort(ePixels, nEdge, sizeof(double),  double_cmp);

  if(0) {
    for(uint32_t kk = 0; kk<nEdge; kk++)
      printf("%2d %f\n", kk, ePixels[kk]);
    printf("\n");
  }

  gsl_sort(ePixels, 1, nEdge);
  double median = gsl_stats_median_from_sorted_data(ePixels, 1, nEdge);
  // double mean = gsl_stats_mean(ePixels, 1, nEdge);
#if verbose > 0
  printf("Median: %f\n", median);
#endif
  free(ePixels);
  return(median);
}

int getZLine(double *W, size_t Ws,
    double * V, size_t Vm, size_t Vn, size_t Vp,
    double *D)
  /* Copy a line from V into W.
   * The line will be D[x, y, z-hWs:z+hWs]
   * where hWs = (Ws-1)/2 and D = [x,y,z]
   *
   * returns 1 on failure (i.e. line is out of bounds)
   * returns 0 if ok.
   *
   * See also: getRegion
   */
{

  int64_t x = nearbyint(D[0]);
  int64_t y = nearbyint(D[1]);
  int64_t z = nearbyint(D[2]);

  size_t Ws2 = (Ws-1)/2;

  if(z+Ws2 >= Vp)
    return 1;
  if(z<(int64_t) Ws2)
    return 1;

  size_t pos = 0;
  for(size_t zz = z-Ws2; zz<=z+Ws2; zz++)
  {
    W[pos++] = V[x + y*Vm + zz*Vm*Vn];
  }
  return 0;
}

int getRegion(double * W, size_t Ws,
    double * V, size_t Vm, size_t Vn, size_t Vp,
    double * D)
  // Get a 2D region for constant z or returns 0 if out of bounds
  // Vm, Vn, Vp is the centre of the region.
{
  size_t Ws2 = (Ws-1)/2;
  size_t pos = 0;

  int64_t x = nearbyint(D[0]);
  int64_t y = nearbyint(D[1]);
  int64_t z = nearbyint(D[2]);

#if verbose > 0
  printf("getRegion round(D): (%lu %lu %lu)\n", x, y, z);
  printf("getRegion size(V)   (%lu %lu %lu)\n", Vm, Vn, Vp);
#endif

  if(x-(int64_t) Ws2 < 0)
    return 1;
  if(y-(int64_t) Ws2 < 0)
    return 1;
  if(x+(uint64_t) Ws2  >= Vm)
    return 1;
  if(y+(uint64_t) Ws2  >= Vn)
    return 1;
  if(z<0)
    return 1;
  if((size_t) z>=Vp)
    return 1;

#if verbose > 1
  printf("Region valid\n");
  printf("%lu %lu %lu\n", Vm, Vn, Vp);
#endif

  uint32_t zz = z;
  for(uint32_t yy = y-Ws2; yy<=y+Ws2; yy++) {
    for(uint32_t xx = x-Ws2; xx<=x+Ws2; xx++) {
      assert(xx<Vm); assert(yy<Vn); assert(zz<Vp);
    // printf("x %lu y %lu z %lu\n", x, y, z);

      W[pos++] = V[xx + yy*Vm + zz*Vm*Vn];
    }
  }
  //  showRegion(W, Ws);
 return 0;
}

void showRegion(double * W, size_t Ws)
  // Dump some quadratic region to the terminal
{
printf("Show region ...\n");
  for(size_t xx = 0 ; xx<Ws; xx++) {
    for(size_t yy = 0; yy<Ws; yy++) {
      printf("%f ", W[yy*Ws+xx]);
    }
    printf("\n");
  }
}


///
/// @file
///
/// @author Mirko Myllykoski (mirkom@cs.umu.se), Umeå University
///
/// @internal LICENSE
///
/// Copyright (c) 2019-2020, Umeå Universitet
///
/// Redistribution and use in source and binary forms, with or without
/// modification, are permitted provided that the following conditions are met:
///
/// 1. Redistributions of source code must retain the above copyright notice,
///    this list of conditions and the following disclaimer.
///
/// 2. Redistributions in binary form must reproduce the above copyright notice,
///    this list of conditions and the following disclaimer in the documentation
///    and/or other materials provided with the distribution.
///
/// 3. Neither the name of the copyright holder nor the names of its
///    contributors may be used to endorse or promote products derived from this
///    software without specific prior written permission.
///
/// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
/// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
/// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
/// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
/// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
/// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
/// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
/// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
/// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
/// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
/// POSSIBILITY OF SUCH DAMAGE.
///

#include <starneig_test_config.h>
#include <starneig/configuration.h>
#include "solvers.h"
#include "../common/common.h"
#include "../common/parse.h"
#include "../common/local_pencil.h"
#ifdef STARNEIG_ENABLE_MPI
#include "../common/starneig_pencil.h"
#endif
#include "../common/threads.h"
#include <starneig/starneig.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <cblas.h>
#include <omp.h>

#ifdef MAGMA_FOUND
#include <cuda_runtime_api.h>
#include <cuda.h>
#include <magma_auxiliary.h>
#include <magma_d.h>
#endif

static hook_solver_state_t lapack_prepare(
    int argc, char * const *argv, struct hook_data_env *env)
{
    return (hook_solver_state_t) env->data;
}

static int lapack_finalize(hook_solver_state_t state, struct hook_data_env *env)
{
    return 0;
}

static int lapack_run(hook_solver_state_t state)
{
    pencil_t data = (pencil_t) state;

    extern void dgehrd_(
        int const *,    // the order of the matrix A
        int const *,    // left bound
        int const *,    // right bound
        double *,       // input/output matrix
        int const *,    // input/output matrix leading dimension
        double *,       // tau (scalar factors)
        double *,       // work space
        int const *,    // work space size
        int *);         // info

    extern void dormhr_(
        char const *,   // side
        char const *,   // transpose
        int const *,    // row count
        int const *,    // column count
        int const *,    // left bound
        int const *,    // right bound
        double const *, // elementary reflectors
        int const *,    // elementary reflector leading dimension
        double const *, // scalar factors (tau)
        double *,       // input/output matrix
        int const *,    // input/output matrix leading dimension
        double *,       // work space
        int const *,    // work space size
        int *);         // info

    extern void dgeqrf_(int const *, int const *, double *, int const *,
        double *, double *, int const *, int *);

    extern void dormqr_(char const *, char const *, int const *, int const *,
        int const *, double const *, int const *, double const *, double *,
        int const *, double*, const int *, int *);

    extern void dgghd3_(char const *, char const *, int const *, int const *,
        int const *, double *, int const *, double *, int const *, double *,
        int const *, double *, int const *, double *, int const *, int *);

    int n = LOCAL_MATRIX_N(data->mat_a);
    double *A = LOCAL_MATRIX_PTR(data->mat_a);
    int ldA = LOCAL_MATRIX_LD(data->mat_a);
    double *Q = LOCAL_MATRIX_PTR(data->mat_q);
    int ldQ = LOCAL_MATRIX_LD(data->mat_q);

    double *B = NULL;
    int ldB = 0;
    double *Z = NULL;
    int ldZ = 0;

    double *tau = NULL;
    double *work = NULL;
    int info, ilo = 1, ihi = n;

    threads_set_mode(THREADS_MODE_LAPACK);

    if (data->mat_b != NULL) {
        B = LOCAL_MATRIX_PTR(data->mat_b);
        ldB = LOCAL_MATRIX_LD(data->mat_b);
        Z = LOCAL_MATRIX_PTR(data->mat_z);
        ldZ = LOCAL_MATRIX_LD(data->mat_z);

        //
        // allocate workspace
        //

        int lwork = 0;

        {
            int _lwork = -1;
            double dlwork;

            dgeqrf_(&n, &n, B, &ldB, tau, &dlwork, &_lwork, &info);
            if (info != 0)
                goto cleanup;

            lwork = MAX(lwork, dlwork);
        }

        {
            int _lwork = -1;
            double dlwork;

            dormqr_("L", "T", &n, &n, &n,
                B, &ldB, tau, A, &ldA, &dlwork, &_lwork, &info);
            if (info != 0)
                goto cleanup;

            lwork = MAX(lwork, dlwork);
        }

        {
            int _lwork = -1;
            double dlwork;

            dormqr_("R", "N", &n, &n, &n,
                B, &ldB, tau, Q, &ldQ, &dlwork, &_lwork, &info);
            if (info != 0)
                goto cleanup;

            lwork = MAX(lwork, dlwork);
        }

        {
            int _lwork = -1;
            double dlwork;

            dgghd3_("V", "V", &n, &ilo, &ihi,
                A, &ldA, B, &ldB, Q, &ldQ, Z, &ldZ, &dlwork, &_lwork, &info);
            if (info != 0)
                goto cleanup;

            lwork = MAX(lwork, dlwork);
        }

        tau = malloc(n*sizeof(double));
        work = malloc(lwork*sizeof(double));

        //
        // reduce
        //

        // form B = ~Q * R
        dgeqrf_(&n, &n, B, &ldB, tau, work, &lwork, &info);
        if (info != 0)
            goto cleanup;

        // A <- ~Q^T * A
        dormqr_("L", "T", &n, &n, &n, B, &ldB, tau, A, &ldA, work, &lwork,
            &info);
        if (info != 0)
            goto cleanup;

        // Q <- Q * ~Q
        dormqr_("R", "N", &n, &n, &n, B, &ldB, tau, Q, &ldQ, work, &lwork,
            &info);
        if (info != 0)
            goto cleanup;

        // clean B (B <- R)
        for (int i = 0; i < n; i++)
            for (int j = i+1; j < n; j++)
                B[(size_t)i*ldB+j] = 0.0;

        // reduce (A,B) to Hessenberg-triangular form
        dgghd3_("V", "V", &n, &ilo, &ihi,
            A, &ldA, B, &ldB, Q, &ldQ, Z, &ldZ, work, &lwork, &info);
        if (info != 0)
            goto cleanup;
    }
    else {

        int lwork = -1;
        double dlwork;

        // request optimal work space size
        dgehrd_(&n, &ilo, &ihi, A, &ldA,
            tau, &dlwork, &lwork, &info);

        if (info != 0)
            goto cleanup;

        lwork = dlwork;
        work = malloc(lwork*sizeof(double));
        tau = malloc(n*sizeof(double));

        // reduce
        dgehrd_(&n, &ilo, &ihi, A, &ldA,
            tau, work, &lwork, &info);

        if (info != 0)
            goto cleanup;

        free(work);
        work = NULL;

        // request optimal work space size
        lwork = -1;
        dormhr_("Right", "No transpose", &n, &n, &ilo, &ihi, A, &ldA, tau,
            Q, &ldQ, &dlwork, &lwork, &info);

        if (info != 0)
            goto cleanup;

        lwork = dlwork;
        work = malloc(lwork*sizeof(double));

        // form Q
        dormhr_("Right", "No transpose", &n, &n, &ilo, &ihi, A, &ldA, tau,
            Q, &ldQ, work, &lwork, &info);

        if (info != 0)
            goto cleanup;

        for (int i = 0; i < n; i++)
            for (int j = i+2; j < n; j++)
                A[i*ldA+j] = 0.0;
    }

cleanup:

    threads_set_mode(THREADS_MODE_DEFAULT);

    free(work);
    free(tau);
    return info;
}

const struct hook_solver hessenberg_lapack_solver = {
    .name = "lapack",
    .desc = "LAPACK's dgehrd/dgghrd subroutine",
    .formats = (hook_data_format_t[]) { HOOK_DATA_FORMAT_PENCIL_LOCAL, 0 },
    .prepare = &lapack_prepare,
    .finalize = &lapack_finalize,
    .run = &lapack_run
};

////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////

#if defined(PDGEHRD_FOUND) && defined(PDORMHR_FOUND) && defined(PDLASET_FOUND)

static hook_solver_state_t scalapack_prepare(
    int argc, char * const *argv, struct hook_data_env *env)
{
    return env;
}

static int scalapack_finalize(
    hook_solver_state_t state, struct hook_data_env *env)
{
    return 0;
}

static int has_valid_descr(
    int matrix_size, int section_size, const starneig_blacs_descr_t *descr)
{
    if (descr->m != matrix_size || descr->n != matrix_size)
        return 0;
    if (descr->sm != section_size || descr->sn != section_size)
        return 0;
    return 1;
}

static int scalapack_run(hook_solver_state_t state)
{
    extern void pdgehrd_(int const *, int const *, int const *, double *,
        int const *, int const *, const starneig_blacs_descr_t *, double *,
        double *, int const *, int *);

    extern void pdormhr_(char const *, char const *, int const *, int const *,
        int const *, int const *, double *, int const *, int const *,
        const starneig_blacs_descr_t *, double *, double *, int const *,
        int const *, const starneig_blacs_descr_t *, double *, int const *,
        int *);

    extern void pdlaset_(char const *, int const *, int const *,
        double const *, double const *, double *, int const *, int const *,
        const starneig_blacs_descr_t *);

    threads_set_mode(THREADS_MODE_SCALAPACK);

    struct hook_data_env *env = state;
    pencil_t pencil = (pencil_t) env->data;

    if (pencil->mat_a == NULL) {
        fprintf(stderr, "Missing matrix A.\n");
        return -1;
    }

    if (pencil->mat_b != NULL) {
        fprintf(stderr, "Solver does not support generalized cases.\n");
        return -1;
    }

    int n = STARNEIG_MATRIX_N(pencil->mat_a);
    int sn = STARNEIG_MATRIX_BN(pencil->mat_a);

    starneig_distr_t distr = STARNEIG_MATRIX_DISTR(pencil->mat_a);
    starneig_blacs_context_t context = starneig_distr_to_blacs_context(distr);

    starneig_blacs_descr_t desc_a, desc_q;
    double *local_a, *local_q;
    STARNEIG_BLACS_MATRIX_DESCR_LOCAL(
        pencil->mat_a, context, &desc_a, (void **)&local_a);
    STARNEIG_BLACS_MATRIX_DESCR_LOCAL(
        pencil->mat_q, context, &desc_q, (void **)&local_q);

    if (!has_valid_descr(n, sn, &desc_a)) {
        fprintf(stderr, "Matrix A has invalid dimensions.\n");
        return -1;
    }

    if (pencil->mat_q != NULL && !has_valid_descr(n, sn, &desc_q)) {
        fprintf(stderr, "Matrix Q has invalid dimension.\n");
        return -1;
    }

    int ilo = 1, ihi = n, ia = 1, ja = 1, ic = 1, jc = 1, lwork, info;
    double *work = NULL, *tau = NULL, _work;

    lwork = -1;
    pdgehrd_(&n, &ilo, &ihi, NULL, &ia, &ja, &desc_a, NULL,
        &_work, &lwork, &info);

    if (info)
        goto cleanup;

    lwork = _work;
    work = malloc(lwork*sizeof(double));
    tau = malloc(n*sizeof(double));

    pdgehrd_(&n, &ilo, &ihi, local_a, &ia, &ja, &desc_a, tau,
        work, &lwork, &info);

    if (info)
        goto cleanup;

    lwork = -1;
    pdormhr_("Right", "No transpose", &n, &n, &ilo, &ihi, NULL,
        &ia, &ja, &desc_a, NULL, NULL, &ic, &jc,
        &desc_q, &_work, &lwork, &info);

    if (info)
        goto cleanup;

    free(work);
    lwork = _work;
    work = malloc(lwork*sizeof(double));

    pdormhr_("Right", "No transpose", &n, &n, &ilo, &ihi, local_a,
        &ia, &ja, &desc_a, tau, local_q, &ic, &jc,
        &desc_q, work, &lwork, &info);

    {
        int nm1 = n-2, one = 1, three = 3;
        double dzero = 0.0;
        pdlaset_("Lower", &nm1, &nm1, &dzero, &dzero, local_a, &three,
            &one, &desc_a);
    }

cleanup:

    threads_set_mode(THREADS_MODE_DEFAULT);
    starneig_blacs_gridexit(context);

    free(work);
    free(tau);

    return info;
}

const struct hook_solver hessenberg_scalapack_solver = {
    .name = "scalapack",
    .desc = "pdgehrd subroutine from scaLAPACK",
    .formats = (hook_data_format_t[]) {
#ifdef STARNEIG_ENABLE_BLACS
        HOOK_DATA_FORMAT_PENCIL_BLACS,
#endif
        0 },
    .prepare = &scalapack_prepare,
    .finalize = &scalapack_finalize,
    .run = &scalapack_run
};

#endif // PDGEHRD_FOUND

////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////

struct starpu_state {
    int argc;
    char * const *argv;
    struct hook_data_env *env;
};

static void starpu_print_usage(int argc, char * const *argv)
{
    printf(
        "  --cores [default,(num)} -- Number of CPU cores\n"
        "  --gpus [default,(num)} -- Number of GPUS\n"
        "  --tile-size [default,(num)] -- tile size\n"
        "  --panel-width [default,(num)] -- Panel width\n"
    );
}

static int starpu_check_args(int argc, char * const *argv, int *argr)
{
    struct multiarg_t arg_cores = read_multiarg(
        "--cores", argc, argv, argr, "default", NULL);
    struct multiarg_t arg_gpus = read_multiarg(
        "--gpus", argc, argv, argr, "default", NULL);
    struct multiarg_t tile_size = read_multiarg(
        "--tile-size", argc, argv, argr, "default", NULL);
    struct multiarg_t panel_width = read_multiarg(
        "--panel-width", argc, argv, argr, "default", NULL);

    if (arg_cores.type == MULTIARG_INVALID)
        return -1;

    if (arg_gpus.type == MULTIARG_INVALID)
        return -1;

    if (tile_size.type == MULTIARG_INVALID ||
    (tile_size.type == MULTIARG_INT && tile_size.int_value < 1)) {
        fprintf(stderr, "Invalid tile size.\n");
        return -1;
    }

    if (panel_width.type == MULTIARG_INVALID ||
    (panel_width.type == MULTIARG_INT && panel_width.int_value < 1)) {
        fprintf(stderr, "Invalid panel width.\n");
        return -1;
    }

    return 0;
}

static void starpu_print_args(int argc, char * const *argv)
{
    print_multiarg("--cores", argc, argv, "default", NULL);
    print_multiarg("--gpus", argc, argv, "default", NULL);
    print_multiarg("--tile-size", argc, argv, "default", NULL);
    print_multiarg("--panel-width", argc, argv, "default", NULL);
}

static hook_solver_state_t starpu_prepare(
    int argc, char * const *argv, struct hook_data_env *env)
{
    struct starpu_state *state = malloc(sizeof(struct starpu_state));

    state->argc = argc;
    state->argv = argv;
    state->env = env;

    struct multiarg_t arg_cores = read_multiarg(
        "--cores", argc, argv, NULL, "default", NULL);
    struct multiarg_t arg_gpus = read_multiarg(
        "--gpus", argc, argv, NULL, "default", NULL);

    int cores = STARNEIG_USE_ALL;
    if (arg_cores.type == MULTIARG_INT)
        cores = arg_cores.int_value;

    int gpus = STARNEIG_USE_ALL;
    if (arg_gpus.type == MULTIARG_INT)
        gpus = arg_gpus.int_value;

#ifdef STARNEIG_ENABLE_MPI
    if (env->format == HOOK_DATA_FORMAT_PENCIL_STARNEIG ||
    env->format == HOOK_DATA_FORMAT_PENCIL_BLACS)
        starneig_node_init(cores, gpus, STARNEIG_FAST_DM);
    else
#endif
        starneig_node_init(
            cores, gpus, STARNEIG_HINT_SM | STARNEIG_AWAKE_WORKERS);

    return state;
}

static int starpu_finalize(hook_solver_state_t state, struct hook_data_env *env)
{
    if (state == NULL)
        return 0;

    starneig_node_finalize();

    free(state);
    return 0;
}

static int starpu_run(hook_solver_state_t state)
{
    int argc = ((struct starpu_state *) state)->argc;
    char * const *argv = ((struct starpu_state *) state)->argv;
    struct hook_data_env *env = ((struct starpu_state *) state)->env;

    struct starneig_hessenberg_conf conf;
    starneig_hessenberg_init_conf(&conf);

    struct multiarg_t tile_size = read_multiarg(
        "--tile-size", argc, argv, NULL, "default", NULL);
    struct multiarg_t panel_width = read_multiarg(
        "--panel-width", argc, argv, NULL, "default", NULL);

    if (tile_size.type == MULTIARG_INT)
        conf.tile_size = tile_size.int_value;
    if (panel_width.type == MULTIARG_INT)
        conf.panel_width = panel_width.int_value;

    int ret = 0;

    if (env->format == HOOK_DATA_FORMAT_PENCIL_LOCAL) {
        pencil_t pencil = (pencil_t) env->data;
        if (pencil->mat_b != NULL) {
            ret = starneig_GEP_SM_HessenbergTriangular(
                LOCAL_MATRIX_N(pencil->mat_a),
                LOCAL_MATRIX_PTR(pencil->mat_a), LOCAL_MATRIX_LD(pencil->mat_a),
                LOCAL_MATRIX_PTR(pencil->mat_b), LOCAL_MATRIX_LD(pencil->mat_b),
                LOCAL_MATRIX_PTR(pencil->mat_q), LOCAL_MATRIX_LD(pencil->mat_q),
                LOCAL_MATRIX_PTR(pencil->mat_z), LOCAL_MATRIX_LD(pencil->mat_z)
            );
        }
        else {
            ret = starneig_SEP_SM_Hessenberg_expert(&conf,
                LOCAL_MATRIX_N(pencil->mat_a), 0, LOCAL_MATRIX_N(pencil->mat_a),
                LOCAL_MATRIX_PTR(pencil->mat_a), LOCAL_MATRIX_LD(pencil->mat_a),
                LOCAL_MATRIX_PTR(pencil->mat_q), LOCAL_MATRIX_LD(pencil->mat_q)
            );
        }
    }
#ifdef STARNEIG_ENABLE_MPI
    if (env->format == HOOK_DATA_FORMAT_PENCIL_BLACS) {
        pencil_t pencil = (pencil_t) env->data;
        if (pencil->mat_b != NULL) {
#ifdef STARNEIG_GEP_DM_HESSENBERGTRIANGULAR
            ret = starneig_GEP_DM_HessenbergTriangular(
                STARNEIG_MATRIX_HANDLE(pencil->mat_a),
                STARNEIG_MATRIX_HANDLE(pencil->mat_b),
                STARNEIG_MATRIX_HANDLE(pencil->mat_q),
                STARNEIG_MATRIX_HANDLE(pencil->mat_z));
#else
            fprintf(stderr,
                "Solver does not support generalized cases in distributed "
                "memory.\n");
            return -1;
#endif
        }
        else {
            ret = starneig_SEP_DM_Hessenberg(
                STARNEIG_MATRIX_HANDLE(pencil->mat_a),
                STARNEIG_MATRIX_HANDLE(pencil->mat_q));
        }
    }
#endif

    return ret;
}

const struct hook_solver hessenberg_starpu_solver = {
    .name = "starneig",
    .desc = "StarPU based subroutine",
    .formats = (hook_data_format_t[]) {
        HOOK_DATA_FORMAT_PENCIL_LOCAL,
#ifdef STARNEIG_ENABLE_BLACS
        HOOK_DATA_FORMAT_PENCIL_BLACS,
#endif
        0 },
    .print_usage = &starpu_print_usage,
    .print_args = &starpu_print_args,
    .check_args = &starpu_check_args,
    .prepare = &starpu_prepare,
    .finalize = &starpu_finalize,
    .run = &starpu_run
};

////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////

static void starpu_simple_print_usage(int argc, char * const *argv)
{
    printf(
        "  --cores [default,(num)} -- Number of CPU cores\n"
        "  --gpus [default,(num)} -- Number of GPUS\n"
    );
}

static void starpu_simple_print_args(int argc, char * const *argv)
{
    print_multiarg("--cores", argc, argv, "default", NULL);
    print_multiarg("--gpus", argc, argv, "default", NULL);
}

static int starpu_simple_check_args(int argc, char * const *argv, int *argr)
{
    struct multiarg_t arg_cores = read_multiarg(
        "--cores", argc, argv, argr, "default", NULL);
    struct multiarg_t arg_gpus = read_multiarg(
        "--gpus", argc, argv, argr, "default", NULL);

    if (arg_cores.type == MULTIARG_INVALID)
        return -1;

    if (arg_gpus.type == MULTIARG_INVALID)
        return -1;

    return 0;
}

static hook_solver_state_t starpu_simple_prepare(
    int argc, char * const *argv, struct hook_data_env *env)
{
    struct starpu_state *state = malloc(sizeof(struct starpu_state));

    state->argc = argc;
    state->argv = argv;
    state->env = env;

    struct multiarg_t arg_cores = read_multiarg(
        "--cores", argc, argv, NULL, "default", NULL);
    struct multiarg_t arg_gpus = read_multiarg(
        "--gpus", argc, argv, NULL, "default", NULL);

    int cores = STARNEIG_USE_ALL;
    if (arg_cores.type == MULTIARG_INT)
        cores = arg_cores.int_value;

    int gpus = STARNEIG_USE_ALL;
    if (arg_gpus.type == MULTIARG_INT)
        gpus = arg_gpus.int_value;

#ifdef STARNEIG_ENABLE_MPI
    if (env->format == HOOK_DATA_FORMAT_PENCIL_STARNEIG ||
    env->format == HOOK_DATA_FORMAT_PENCIL_BLACS)
        starneig_node_init(cores, gpus, STARNEIG_FAST_DM);
    else
#endif
        starneig_node_init(
            cores, gpus, STARNEIG_HINT_SM | STARNEIG_AWAKE_WORKERS);

    return state;
}

static int starpu_simple_finalize(
    hook_solver_state_t state, struct hook_data_env *env)
{
    if (state == NULL)
        return 0;

    starneig_node_finalize();

    free(state);
    return 0;
}

static int starpu_simple_run(hook_solver_state_t state)
{
    struct hook_data_env *env = ((struct starpu_state *) state)->env;

    int ret = 0;

    if (env->format == HOOK_DATA_FORMAT_PENCIL_LOCAL) {
        pencil_t pencil = (pencil_t) env->data;

        if (pencil->mat_b != NULL) {
            ret = starneig_GEP_SM_HessenbergTriangular(
                LOCAL_MATRIX_N(pencil->mat_a),
                LOCAL_MATRIX_PTR(pencil->mat_a), LOCAL_MATRIX_LD(pencil->mat_a),
                LOCAL_MATRIX_PTR(pencil->mat_b), LOCAL_MATRIX_LD(pencil->mat_b),
                LOCAL_MATRIX_PTR(pencil->mat_q), LOCAL_MATRIX_LD(pencil->mat_q),
                LOCAL_MATRIX_PTR(pencil->mat_z), LOCAL_MATRIX_LD(pencil->mat_z)
            );
        }
        else {
            ret = starneig_SEP_SM_Hessenberg(LOCAL_MATRIX_N(pencil->mat_a),
                LOCAL_MATRIX_PTR(pencil->mat_a), LOCAL_MATRIX_LD(pencil->mat_a),
                LOCAL_MATRIX_PTR(pencil->mat_q), LOCAL_MATRIX_LD(pencil->mat_q)
            );
        }
    }
#ifdef STARNEIG_ENABLE_MPI
    if (env->format == HOOK_DATA_FORMAT_PENCIL_BLACS) {
        pencil_t pencil = (pencil_t) env->data;

        if (pencil->mat_b != NULL) {
#ifdef STARNEIG_GEP_DM_HESSENBERGTRIANGULAR
            ret = starneig_GEP_DM_HessenbergTriangular(
                STARNEIG_MATRIX_HANDLE(pencil->mat_a),
                STARNEIG_MATRIX_HANDLE(pencil->mat_b),
                STARNEIG_MATRIX_HANDLE(pencil->mat_q),
                STARNEIG_MATRIX_HANDLE(pencil->mat_z));
#else
            fprintf(stderr,
                "Solver does not support distributed memory in generalized "
                "cases.\n");
            return -1;
#endif
        }
        else {
            ret = starneig_SEP_DM_Hessenberg(
                STARNEIG_MATRIX_HANDLE(pencil->mat_a),
                STARNEIG_MATRIX_HANDLE(pencil->mat_q));
        }
    }
#endif

    return ret;
}

const struct hook_solver hessenberg_starpu_simple_solver = {
    .name = "starneig-simple",
    .desc = "StarPU based subroutine (simplified interface)",
    .formats = (hook_data_format_t[]) {
        HOOK_DATA_FORMAT_PENCIL_LOCAL,
#ifdef STARNEIG_ENABLE_BLACS
        HOOK_DATA_FORMAT_PENCIL_BLACS,
#endif
        0 },
    .print_usage = &starpu_simple_print_usage,
    .print_args = &starpu_simple_print_args,
    .check_args = &starpu_simple_check_args,
    .prepare = &starpu_simple_prepare,
    .finalize = &starpu_simple_finalize,
    .run = &starpu_simple_run
};

////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////

#ifdef MAGMA_FOUND

static hook_solver_state_t magma_dgehrd_prepare(
    int argc, char * const *argv, struct hook_data_env *env)
{
    if (magma_init() != MAGMA_SUCCESS)
        return NULL;
    return (hook_solver_state_t) env->data;
}

static int magma_dgehrd_finalize(
    hook_solver_state_t state, struct hook_data_env *env)
{
    magma_finalize();
    return 0;
}

static int magma_dgehrd_run(hook_solver_state_t state)
{
    pencil_t data = (pencil_t) state;

    int n = LOCAL_MATRIX_N(data->mat_a);
    double *A = LOCAL_MATRIX_PTR(data->mat_a);
    double *Q = LOCAL_MATRIX_PTR(data->mat_q);
    int ldA = LOCAL_MATRIX_LD(data->mat_a);
    int ldQ = LOCAL_MATRIX_LD(data->mat_q);

    int info;

    double *tau = NULL;
    double *work = NULL;
    double *dT = NULL;

    threads_set_mode(THREADS_MODE_LAPACK);

    // request optimal work space size
    double _work;
    magma_dgehrd(n, 1, n, A, ldA, tau, &_work, -1, dT, &info);

    if (info != 0)
        goto finalize;

    tau = malloc(LOCAL_MATRIX_N(data->mat_a)*sizeof(double));

    int lwork = _work;
    work = malloc(lwork*sizeof(double));

    int nb = magma_get_dgehrd_nb(LOCAL_MATRIX_N(data->mat_a));
    cudaMalloc((void**)&dT, nb*n*sizeof(double));

    // reduce
    magma_dgehrd(n, 1, n, A, ldA, tau, work, lwork, dT, &info);

    if (info != 0)
        goto finalize;

    free(work);
    work = NULL;

    // copy A -> Q
    for (int i = 0; i < n; i++)
        memcpy(Q+i*ldQ, A+i*ldA, n*sizeof(double));

    // form Q
    magma_dorghr(n, 1, n, Q, ldQ, tau, dT, nb, &info);

    if (info != 0)
        goto finalize;

    // zero entries below the first sub-diagonal
    for (int i = 0; i < n-1; i++)
        memset(A+i*ldA+i+2, 0, (n-i-2)*sizeof(double));

finalize:

    threads_set_mode(THREADS_MODE_DEFAULT);

    cudaFree(dT);
    free(work);
    free(tau);
    return info;
}

const struct hook_solver hessenberg_magma_solver = {
    .name = "magma",
    .desc = "MAGMA's dgehrd subroutine",
    .formats = (hook_data_format_t[]) { HOOK_DATA_FORMAT_PENCIL_LOCAL, 0 },
    .prepare = &magma_dgehrd_prepare,
    .finalize = &magma_dgehrd_finalize,
    .run = &magma_dgehrd_run
};

#endif // MAGMA_FOUND


/*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*
**                                                                                **
** This file forms part of the Underworld geophysics modelling application.         **
**                                                                                **
** For full license and copyright information, please refer to the LICENSE.md file  **
** located at the project root, or contact the authors.                             **
**                                                                                **
**~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*~*/
#include <mpi.h>
#include <StGermain/StGermain.h>
#include <StgDomain/StgDomain.h>
#include <StgFEM/StgFEM.h>
#include "types.h"
#include <petsc.h>
#include <petscvec.h>
#include <petscmat.h>
#include <petscksp.h>
#include <petscpc.h>
#include <petscsnes.h>
#include <petscis.h>
#include <petscviewer.h>
#include <petscsys.h>
#include <petscversion.h>

#if ( (PETSC_VERSION_MAJOR >= 3) && (PETSC_VERSION_MINOR >=3) )
  #if (PETSC_VERSION_MINOR >=6)
     #include "petsc/private/kspimpl.h"
  #else
     #include "petsc-private/kspimpl.h"   /*I "petscksp.h" I*/
  #endif
#else
  #include "private/kspimpl.h"   /*I "petscksp.h" I*/
#endif

//#include "ksptypes.h"
#include "ksp-register.h"
#include "StokesBlockKSPInterface.h"

#include <stdio.h>
#include <stdlib.h>
#include <assert.h>
#include <string.h>

/* Macro for checking number integrity - i.e. checks if number is infinite or "not a number" */
#define SBKSP_isGoodNumber( number ) ( (! isnan( number ) ) && ( ! isinf( number ) ) )
#define SBKSP_GetPetscMatrix( matrix ) ( (Mat)(matrix) )
#define SBKSP_GetPetscVector( vector ) ( (Vec)(vector) )
#define SBKSP_GetPetscKSP( solver ) ( (KSP)(solver)  )
PetscErrorCode _BlockSolve( void* solver, void* _stokesSLE );

const Type StokesBlockKSPInterface_Type = "StokesBlockKSPInterface";

void* _StokesBlockKSPInterface_DefaultNew( Name name ) {
    SizeT                                              _sizeOfSelf = sizeof(StokesBlockKSPInterface);
    Type                                                      type = StokesBlockKSPInterface_Type;
    Stg_Class_DeleteFunction*                              _delete = _SLE_Solver_Delete;
    Stg_Class_PrintFunction*                                _print = _SLE_Solver_Print;
    Stg_Class_CopyFunction*                                  _copy = _SLE_Solver_Copy;
    Stg_Component_BuildFunction*                            _build = _StokesBlockKSPInterface_Build;
    Stg_Component_InitialiseFunction*                  _initialise = _StokesBlockKSPInterface_Initialise;
    Stg_Component_ExecuteFunction*                        _execute = _SLE_Solver_Execute;
    Stg_Component_DestroyFunction*                        _destroy = _SLE_Solver_Destroy;
    SLE_Solver_GetResidualFunc*                       _getResidual = NULL;
    Stg_Component_DefaultConstructorFunction*  _defaultConstructor = _StokesBlockKSPInterface_DefaultNew;
    Stg_Component_ConstructFunction*                    _construct = _StokesBlockKSPInterface_AssignFromXML;
    SLE_Solver_SolverSetupFunction*                   _solverSetup = _StokesBlockKSPInterface_SolverSetup;
    SLE_Solver_SolveFunction*                               _solve = _StokesBlockKSPInterface_Solve;

    AllocationType  nameAllocationType = NON_GLOBAL /* default value NON_GLOBAL */;
    return (void*) _StokesBlockKSPInterface_New( STOKESBLOCKKSPINTERFACE_PASSARGS );
}

/* Creation implementation / Virtual constructor */
/* Set up function pointers */
StokesBlockKSPInterface* _StokesBlockKSPInterface_New( STOKESBLOCKKSPINTERFACE_DEFARGS )
{
    StokesBlockKSPInterface* self;
    /* Allocate memory */
    assert( _sizeOfSelf >= sizeof(StokesBlockKSPInterface) );

    self = (StokesBlockKSPInterface*) _SLE_Solver_New( SLE_SOLVER_PASSARGS );

    /* Virtual info */
    return self;
}

void _StokesBlockKSPInterface_Init(
		StokesBlockKSPInterface*      self,
		StiffnessMatrix*   preconditioner,
		Stokes_SLE *       st_sle,
		PETScMGSolver *    mg,
        Name   filename,
        char * string,
		StiffnessMatrix*  k2StiffMat,
		StiffnessMatrix*  mStiffMat,
		ForceVector*	  f2ForceVec,
		ForceVector*	  jForceVec,
		double            penaltyNumber,
		double            hFactor,
		StiffnessMatrix*  vmStiffMat,
		ForceVector*	  vmForceVec  )
{
	self->preconditioner = preconditioner;
	self->st_sle = st_sle;
	self->mg     = mg;
    self->optionsFile = filename;
    self->optionsString = string;
	self->k2StiffMat = k2StiffMat;
	self->f2ForceVec = f2ForceVec;
	self->penaltyNumber = penaltyNumber;
	self->hFactor       = hFactor;
	self->mStiffMat = mStiffMat;
	self->jForceVec = jForceVec;
	self->vmStiffMat = vmStiffMat;
	self->vmForceVec = vmForceVec;
	/* add the vecs and matrices to the Base SLE class's dynamic lists, so they can be
	initialised and built properly */

	/*
	if (k2StiffMat )
	    SystemLinearEquations_AddStiffnessMatrix( st_sle, k2StiffMat );

	if (f2ForceVec )
	    SystemLinearEquations_AddForceVector( st_sle, f2ForceVec );

	if (mStiffMat )
	    SystemLinearEquations_AddStiffnessMatrix( st_sle, mStiffMat );

	if (jForceVec )
	    SystemLinearEquations_AddForceVector( st_sle, jForceVec );

	if (vmStiffMat )
	    SystemLinearEquations_AddStiffnessMatrix( st_sle, vmStiffMat );

	if (vmForceVec )
	    SystemLinearEquations_AddForceVector( st_sle, vmForceVec );
    */
}

void _StokesBlockKSPInterface_Build( void* solver, void* sle ) {/* it is the sle here being passed in*/
	StokesBlockKSPInterface*	self  = (StokesBlockKSPInterface*)solver;

	Stream_IndentBranch( StgFEM_Debug );

	/* Build Preconditioner */
	if ( self->preconditioner ) {
		Stg_Component_Build( self->preconditioner, sle, False );
		SystemLinearEquations_AddStiffnessMatrix( self->st_sle, self->preconditioner );

	}
	if( self->mg ){
	    Stg_Component_Build( self->mg, sle, False );
	}

	Stream_UnIndentBranch( StgFEM_Debug );
}

void _StokesBlockKSPInterface_AssignFromXML( void* solver, Stg_ComponentFactory* cf, void* data ) {
	StokesBlockKSPInterface* self         = (StokesBlockKSPInterface*) solver;
	//double                  tolerance;
	//Iteration_Index         maxUzawaIterations, minUzawaIterations;
	StiffnessMatrix*  preconditioner;
	StiffnessMatrix*  k2StiffMat;
	ForceVector*	  f2ForceVec;
	StiffnessMatrix*  mStiffMat;
	ForceVector*	  jForceVec;
	double            penaltyNumber;
	double            hFactor;
	StiffnessMatrix*  vmStiffMat;
	ForceVector*	  vmForceVec;
	//Bool                    useAbsoluteTolerance;
	//Bool                    monitor;
	Stokes_SLE *            st_sle;
	PETScMGSolver *         mg;
	//Name                filename;
	//char* 		        string;

	_SLE_Solver_AssignFromXML( self, cf, data );

	preconditioner = Stg_ComponentFactory_ConstructByKey( cf, self->name, (Dictionary_Entry_Key)"Preconditioner", StiffnessMatrix, False, data  );
	st_sle  = Stg_ComponentFactory_ConstructByKey( cf, self->name, (Dictionary_Entry_Key)"stokesEqn", Stokes_SLE, True, data  );
	mg      = Stg_ComponentFactory_ConstructByKey( cf, self->name, (Dictionary_Entry_Key)"mgSolver", PETScMGSolver, False, data);
	k2StiffMat = Stg_ComponentFactory_ConstructByKey( cf, self->name, (Dictionary_Entry_Key)"2ndStressTensorMatrix", StiffnessMatrix, False, data  );
	f2ForceVec = Stg_ComponentFactory_ConstructByKey( cf, self->name, (Dictionary_Entry_Key)"2ndForceVector", ForceVector, False, data  );
	penaltyNumber   = Stg_ComponentFactory_GetDouble( cf, self->name, (Dictionary_Entry_Key)"penaltyNumber", 0.0  );
	hFactor         = Stg_ComponentFactory_GetDouble( cf, self->name, (Dictionary_Entry_Key)"hFactor", 0.0  );
	mStiffMat = Stg_ComponentFactory_ConstructByKey( cf, self->name, (Dictionary_Entry_Key)"MassMatrix", StiffnessMatrix, False, data  );
	jForceVec = Stg_ComponentFactory_ConstructByKey( cf, self->name, (Dictionary_Entry_Key)"JunkForceVector", ForceVector, False, data  );
	vmStiffMat = Stg_ComponentFactory_ConstructByKey( cf, self->name, (Dictionary_Entry_Key)"VelocityMassMatrix", StiffnessMatrix, False, data  );
	vmForceVec = Stg_ComponentFactory_ConstructByKey( cf, self->name, (Dictionary_Entry_Key)"VMassForceVector", ForceVector, False, data  );

	_StokesBlockKSPInterface_Init( self, preconditioner, st_sle, mg, NULL, NULL, k2StiffMat, mStiffMat,
                                   f2ForceVec, jForceVec, penaltyNumber, hFactor, vmStiffMat, vmForceVec);

}

void _StokesBlockKSPInterface_Initialise( void* solver, void* data ) {
	StokesBlockKSPInterface* self = (StokesBlockKSPInterface*) solver;
	Stokes_SLE*             sle  = (Stokes_SLE*)  self->st_sle;

	/* Initialise Parent */
	_SLE_Solver_Initialise( self, sle );

	KSPRegisterAllKSP("Solvers/KSPSolvers/src");
}

/* SolverSetup */

void _StokesBlockKSPInterface_SolverSetup( void* solver, void* stokesSLE ) {
	StokesBlockKSPInterface* self = (StokesBlockKSPInterface*) solver;
	//Stokes_SLE*             sle  = (Stokes_SLE*)             stokesSLE;

 	Journal_DPrintf( self->debug, "In %s:\n", __func__ );
	Stream_IndentBranch( StgFEM_Debug );

	Stream_UnIndentBranch( StgFEM_Debug );
}
void SBKSP_SetSolver( void* solver, void* stokesSLE ) {
	SLE_Solver* self = (SLE_Solver*) solver;
	Stokes_SLE*              sle  = (Stokes_SLE*) stokesSLE;

    sle->solver=self;

}
void SBKSP_SetPenalty( void* solver, double penalty ) {
  StokesBlockKSPInterface* self = (StokesBlockKSPInterface*) solver;
  self->penaltyNumber=penalty;
}

int SBKSP_GetPressureIts(void *solver){
  StokesBlockKSPInterface* self = (StokesBlockKSPInterface*) solver;
  return self->stats.pressure_its;
}

/***********************************************************************************************************/
/***********************************************************************************************************/
/***********************************************************************************************************/

void SBKSP_GetStokesOperators(
		Stokes_SLE *stokesSLE,
		Mat *K,Mat *G,Mat *D,Mat *C,Mat *approxS,
		Vec *f,Vec *h,Vec *u,Vec *p )
{

	*K = *G = *D = *C = PETSC_NULL;
	if (stokesSLE->kStiffMat){      *K = SBKSP_GetPetscMatrix( stokesSLE->kStiffMat->matrix );     }
	if (stokesSLE->gStiffMat){      *G = SBKSP_GetPetscMatrix( stokesSLE->gStiffMat->matrix );     }
	if (stokesSLE->dStiffMat){      *D = SBKSP_GetPetscMatrix( stokesSLE->dStiffMat->matrix );     }
	if (stokesSLE->cStiffMat){      *C = SBKSP_GetPetscMatrix( stokesSLE->cStiffMat->matrix );     }

	/* preconditioner */
	*approxS = PETSC_NULL;
	if( ((StokesBlockKSPInterface*)stokesSLE->solver)->preconditioner ) {
		StiffnessMatrix *preconditioner;

		preconditioner = ((StokesBlockKSPInterface*)stokesSLE->solver)->preconditioner;
		*approxS = SBKSP_GetPetscMatrix( preconditioner->matrix );
	}

	*f = *h = PETSC_NULL;
	if (stokesSLE->fForceVec){      *f = SBKSP_GetPetscVector( stokesSLE->fForceVec->vector );     }
	if (stokesSLE->hForceVec){      *h = SBKSP_GetPetscVector( stokesSLE->hForceVec->vector );     }

	*u = *p = PETSC_NULL;
	if (stokesSLE->uSolnVec){       *u = SBKSP_GetPetscVector( stokesSLE->uSolnVec->vector );      }
	if (stokesSLE->pSolnVec){       *p = SBKSP_GetPetscVector( stokesSLE->pSolnVec->vector );      }

}
/***********************************************************************************************************/
/***********************************************************************************************************/
/***********************************************************************************************************/
/* Sets up Solver to be a custom ksp (KSP_BSSCR) solve by default: */
void _StokesBlockKSPInterface_Solve( void* solver, void* _stokesSLE ) {
  StokesBlockKSPInterface* self    = (StokesBlockKSPInterface*)solver;
  PetscLogDouble flopsA,flopsB;
  PetscTruth found, get_flops;

  found = PETSC_FALSE;
  get_flops = PETSC_FALSE;
  PetscOptionsGetTruth( PETSC_NULL, "-get_flops", &get_flops, &found);
  if(get_flops){
    PetscGetFlops(&flopsA); }

  _BlockSolve(solver, _stokesSLE);

  if(get_flops){
    PetscGetFlops(&flopsB);
    self->stats.total_flops=(double)(flopsB-flopsA); }
}
PetscErrorCode _BlockSolve( void* solver, void* _stokesSLE ) {
  Stokes_SLE*  stokesSLE  = (Stokes_SLE*)_stokesSLE;
  StokesBlockKSPInterface* Solver    = (StokesBlockKSPInterface*)solver;

  /* Create shortcuts to stuff needed on sle */
  Mat       K;
  Mat       G;
  Mat       Gt;
  Mat       D;
  Mat       C;
  Mat       approxS;
  Vec       u;
  Vec       p;
  Vec       f;
  Vec       h;
  Mat stokes_P;
  Mat stokes_A;
  Vec stokes_x;
  Vec stokes_b;
  Mat a[2][2];
  Vec x[2];
  Vec b[2];
  KSP stokes_ksp;
  PC  stokes_pc;
  PetscTruth sym,flg;
  PetscErrorCode ierr;

  PetscInt   N,n;

  SBKSP_GetStokesOperators( stokesSLE, &K,&G,&D,&C, &approxS, &f,&h, &u,&p );

  /* create Gt */
  if( !D ) {
    ierr = MatTranspose( G, MAT_INITIAL_MATRIX, &Gt);CHKERRQ(ierr);
    sym = PETSC_TRUE;
    Solver->DIsSym = sym;
  }
  else {
    Gt = D;
    sym = PETSC_FALSE;
    Solver->DIsSym = sym;
  }
  flg=PETSC_FALSE;
  PetscOptionsHasName(PETSC_NULL,"-use_petsc_ksp",&flg);
  if (flg) {
    if( !C ) {
      /* Everything in this bracket, dependent on !C, is to build
         a matrix with diagonals of 0 for C the previous comment ways

      need a 'zero' matrix to keep fieldsplit happy in petsc? */
      MatType mtype;
      Vec V;
      //MatGetSize( G, &M, &N );
      VecGetSize(p, &N);
      VecGetLocalSize( p, &n );
      MatCreate( PetscObjectComm((PetscObject) K), &C );
      MatSetSizes( C, PETSC_DECIDE ,PETSC_DECIDE, N, N );
#if (((PETSC_VERSION_MAJOR==3) && (PETSC_VERSION_MINOR>=3)) || (PETSC_VERSION_MAJOR>3) )
      MatSetUp(C);
#endif
      MatGetType( G, &mtype );
      MatSetType( C, mtype );
      MatGetVecs( G, &V, PETSC_NULL );
      VecSet(V, 0.0);
      //VecSet(h, 1.0);
      ierr = VecAssemblyBegin( V );CHKERRQ(ierr);
      ierr = VecAssemblyEnd  ( V );CHKERRQ(ierr);
      ierr = MatDiagonalSet(C,V,INSERT_VALUES);CHKERRQ(ierr);
      ierr = MatAssemblyBegin( C, MAT_FINAL_ASSEMBLY );CHKERRQ(ierr);
      ierr = MatAssemblyEnd  ( C, MAT_FINAL_ASSEMBLY );CHKERRQ(ierr);
    }
  }
  a[0][0]=K;  a[0][1]=G;
  a[1][0]=Gt; a[1][1]=C;
  ierr = MatCreateNest(PetscObjectComm((PetscObject) K), 2, NULL, 2, NULL, (Mat *)a, &stokes_A);CHKERRQ(ierr);
  ierr = MatAssemblyBegin( stokes_A, MAT_FINAL_ASSEMBLY );CHKERRQ(ierr);
  ierr = MatAssemblyEnd( stokes_A, MAT_FINAL_ASSEMBLY );CHKERRQ(ierr);

  x[0]=u;
  x[1]=p;
  ierr = VecCreateNest(PetscObjectComm((PetscObject) u), 2, NULL, x, &stokes_x);CHKERRQ(ierr);
  ierr = VecAssemblyBegin( stokes_x );CHKERRQ(ierr);
  ierr = VecAssemblyEnd( stokes_x);CHKERRQ(ierr);

  b[0]=f;
  b[1]=h;
  ierr = VecCreateNest(PetscObjectComm((PetscObject) f), 2, NULL, b, &stokes_b);CHKERRQ(ierr);
  ierr = VecAssemblyBegin( stokes_b );CHKERRQ(ierr);
  ierr = VecAssemblyEnd( stokes_b);CHKERRQ(ierr);

  /* if( approxS ) { */
  /*   a[0][0]=K;    a[0][1]=G; */
  /*   a[1][0]=NULL; a[1][1]=approxS; */
  /*   ierr = MatCreateNest(PetscObjectComm((PetscObject) K), 2, NULL, 2, NULL, (Mat *)a, &stokes_P);CHKERRQ(ierr); */
  /*   ierr = MatAssemblyBegin( stokes_P, MAT_FINAL_ASSEMBLY );CHKERRQ(ierr); */
  /*   ierr = MatAssemblyEnd( stokes_P, MAT_FINAL_ASSEMBLY );CHKERRQ(ierr); */
  /* } */
  /* else { */
    stokes_P = stokes_A;
  /* } */

  /* probably should make a Destroy function for these two */
  /* Update options from file and/or string here so we can change things on the fly */
  //PetscOptionsInsertFile(PETSC_COMM_WORLD, Solver->optionsFile, PETSC_FALSE);
  //PetscOptionsInsertString(Solver->optionsString);

  ierr = KSPCreate( PETSC_COMM_WORLD, &stokes_ksp );CHKERRQ(ierr);
  Stg_KSPSetOperators( stokes_ksp, stokes_A, stokes_P, SAME_NONZERO_PATTERN );
  ierr = KSPSetType( stokes_ksp, "bsscr" );/* i.e. making this the default solver : calls KSPCreate_XXX */CHKERRQ(ierr);

  ierr = KSPGetPC( stokes_ksp, &stokes_pc );CHKERRQ(ierr);
  ierr = PCSetType( stokes_pc, PCNONE );CHKERRQ(ierr);
  ierr = KSPSetInitialGuessNonzero( stokes_ksp, PETSC_TRUE );CHKERRQ(ierr);
  ierr = KSPSetFromOptions( stokes_ksp );CHKERRQ(ierr);

  /*
    Doing this so the KSP Solver has access to the StgFEM Multigrid struct (PETScMGSolver).
    As well as any custom stuff on the Stokes_SLE struct
  */
  if( stokes_ksp->data ){/* then ksp->data has been created in a KSpSetUp_XXX function */
    /* testing for our KSP types that need the data that is on Solver... */
    /* for the moment then, this function not completely agnostic about our KSPs */
    //if(!strcmp("bsscr",stokes_ksp->type_name)){/* if is bsscr then set up the data on the ksp */
    flg=PETSC_FALSE;
    PetscOptionsHasName(PETSC_NULL,"-use_petsc_ksp",&flg);
    if (!flg) {
      ((KSP_COMMON*)(stokes_ksp->data))->st_sle         = Solver->st_sle;
      ((KSP_COMMON*)(stokes_ksp->data))->mg             = Solver->mg;
      ((KSP_COMMON*)(stokes_ksp->data))->DIsSym         = Solver->DIsSym;
      ((KSP_COMMON*)(stokes_ksp->data))->preconditioner = Solver->preconditioner;
      ((KSP_COMMON*)(stokes_ksp->data))->solver         = Solver;
    }
  }

  ierr = KSPSolve( stokes_ksp, stokes_b, stokes_x );CHKERRQ(ierr);

  Stg_KSPDestroy(&stokes_ksp );
  //if( ((StokesBlockKSPInterface*)stokesSLE->solver)->preconditioner )
  if(stokes_P != stokes_A) { Stg_MatDestroy(&stokes_P ); }

  Stg_MatDestroy(&stokes_A );

  Stg_VecDestroy(&stokes_x);
  Stg_VecDestroy(&stokes_b);

  if(!D){ Stg_MatDestroy(&Gt); }
  if(C && (stokesSLE->cStiffMat->matrix != C) ){ Stg_MatDestroy(&C); }

  PetscFunctionReturn(0);
}


/* integration/workspace.c
 *
 * Copyright (C) 1996, 1997, 1998, 1999, 2000 Brian Gough
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or (at
 * your option) any later version.
 *
 * This program 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
 * General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
 */

#include <config.h>
#include <stdlib.h>
#include <gsl/gsl_integration.h>
#include <gsl/gsl_errno.h>

gsl_integration_workspace *
gsl_integration_workspace_alloc (const size_t n)
{
  gsl_integration_workspace * w ;

  if (n == 0)
    {
      GSL_ERROR_VAL ("workspace length n must be positive integer",
			GSL_EDOM, 0);
    }

  w = (gsl_integration_workspace *)
    malloc (sizeof (gsl_integration_workspace));

  if (w == 0)
    {
      GSL_ERROR_VAL ("failed to allocate space for workspace struct",
			GSL_ENOMEM, 0);
    }

  w->alist = (double *) malloc (n * sizeof (double));

  if (w->alist == 0)
    {
      free (w);		/* exception in constructor, avoid memory leak */

      GSL_ERROR_VAL ("failed to allocate space for alist ranges",
			GSL_ENOMEM, 0);
    }

  w->blist = (double *) malloc (n * sizeof (double));

  if (w->blist == 0)
    {
      free (w->alist);
      free (w);		/* exception in constructor, avoid memory leak */

      GSL_ERROR_VAL ("failed to allocate space for blist ranges",
			GSL_ENOMEM, 0);
    }

  w->rlist = (double *) malloc (n * sizeof (double));

  if (w->rlist == 0)
    {
      free (w->blist);
      free (w->alist);
      free (w);		/* exception in constructor, avoid memory leak */

      GSL_ERROR_VAL ("failed to allocate space for rlist ranges",
			GSL_ENOMEM, 0);
    }

  w->elist = (double *) malloc (n * sizeof (double));

  if (w->elist == 0)
    {
      free (w->rlist);
      free (w->blist);
      free (w->alist);
      free (w);		/* exception in constructor, avoid memory leak */

      GSL_ERROR_VAL ("failed to allocate space for elist ranges",
			GSL_ENOMEM, 0);
    }

  w->order = (size_t *) malloc (n * sizeof (size_t));

  if (w->order == 0)
    {
      free (w->elist);
      free (w->rlist);
      free (w->blist);
      free (w->alist);
      free (w);		/* exception in constructor, avoid memory leak */

      GSL_ERROR_VAL ("failed to allocate space for order ranges",
			GSL_ENOMEM, 0);
    }

  w->level = (size_t *) malloc (n * sizeof (size_t));

  if (w->level == 0)
    {
      free (w->order);
      free (w->elist);
      free (w->rlist);
      free (w->blist);
      free (w->alist);
      free (w);		/* exception in constructor, avoid memory leak */

      GSL_ERROR_VAL ("failed to allocate space for order ranges",
			GSL_ENOMEM, 0);
    }

  w->size = 0 ;
  w->limit = n ;
  w->maximum_level = 0 ;

  return w ;
}

void
gsl_integration_workspace_free (gsl_integration_workspace * w)
{
  free (w->level) ;
  free (w->order) ;
  free (w->elist) ;
  free (w->rlist) ;
  free (w->blist) ;
  free (w->alist) ;
  free (w) ;
}

/*
size_t
gsl_integration_workspace_limit (gsl_integration_workspace * w)
{
  return w->limit ;
}

size_t
gsl_integration_workspace_size (gsl_integration_workspace * w)
{
  return w->size ;
}
*/


#include <stdlib.h>

#include <gsl/gsl_statistics_double.h>
#include <hdf5.h>
#include <hdf5_hl.h>

#include "qdm.h"

qdm_ijk *
qdm_ijk_alloc(
    size_t size1,
    size_t size2,
    size_t size3
)
{
  qdm_ijk *t = malloc(sizeof(qdm_ijk));

  t->size1 = size1;
  t->size2 = size2;
  t->size3 = size3;

  t->data = malloc(sizeof(double) * size1 * size2 * size3);
  t->owner = true;

  return t;
}

qdm_ijk *
qdm_ijk_calloc(
    size_t size1,
    size_t size2,
    size_t size3
)
{
  qdm_ijk *t = malloc(sizeof(qdm_ijk));

  t->size1 = size1;
  t->size2 = size2;
  t->size3 = size3;

  t->data = calloc(sizeof(double), size1 * size2 * size3);
  t->owner = true;

  return t;
}

void
qdm_ijk_free(
    qdm_ijk *t
)
{
  if (t == NULL) {
    return;
  }

  if (t->owner) {
    free(t->data);
  }

  free(t);
}

qdm_ijk_view
qdm_ijk_view_array(
    double *base,
    size_t size1,
    size_t size2,
    size_t size3
)
{
  qdm_ijk_view v = {
    .ijk = {
      .size1 = size1,
      .size2 = size2,
      .size3 = size3,

      .data = base,
      .owner = false,
    },
  };

  return v;
}

gsl_vector_view
qdm_ijk_get_k(
    qdm_ijk *t,
    size_t i,
    size_t j
)
{
  return gsl_vector_view_array_with_stride(
      &t->data[i * t->size2 + j],
      t->size1 * t->size2,
      t->size3
  );
}

gsl_matrix_view
qdm_ijk_get_ij(
    qdm_ijk *t,
    size_t k
)
{
  return gsl_matrix_view_array(
      &t->data[k * t->size1 * t->size2],
      t->size1,
      t->size2
  );
}

double
qdm_ijk_get(
    qdm_ijk *t,
    size_t i,
    size_t j,
    size_t k
)
{
  gsl_matrix_view ij = qdm_ijk_get_ij(t, k);

  return gsl_matrix_get(&ij.matrix, i, j);
}

gsl_matrix *
qdm_ijk_cov(
    qdm_ijk *t
)
{
  size_t d = t->size1 * t->size2;
  size_t dr = 0;

  gsl_matrix *r = gsl_matrix_calloc(d, d);

  for (size_t i0 = 0; i0 < t->size1; i0++) {
    for (size_t j0 = 0; j0 < t->size2; j0++) {
      gsl_vector_view k0 = qdm_ijk_get_k(t, i0, j0);

      /* TODO: When we have time (and if it matters), this is symmetric and we
       * don't need to do a bunch of these if we get a little clever about
       * copying the mirror'd results.
       */
      for (size_t ip = 0; ip < t->size1; ip++) {
        for (size_t jp = 0; jp < t->size2; jp++) {
          gsl_vector_view kp = qdm_ijk_get_k(t, ip, jp);

          double cov = gsl_stats_covariance(
              k0.vector.data,
              k0.vector.stride,
              kp.vector.data,
              kp.vector.stride,
              kp.vector.size
          );

          gsl_matrix_set(r, dr, ip * t->size2 + jp, cov);
        }
      }

      dr++;
    }
  }

  return r;
}

int
qdm_ijk_write(
    hid_t id,
    const char *name,
    const qdm_ijk *t
)
{
  int status = 0;

  hid_t datatype  = -1;
  hid_t dataspace = -1;
  hid_t dataset   = -1;

  if (t == NULL) {
    goto cleanup;
  }

  datatype = H5Tcopy(H5T_NATIVE_DOUBLE);

  status = H5Tset_order(datatype, H5T_ORDER_LE);
  if (status != 0) {
    goto cleanup;
  }

  hsize_t dims[3] = {
    t->size3,
    t->size1,
    t->size2,
  };
  dataspace = H5Screate_simple(3, dims, NULL);
  if (dataspace < 0) {
    status = dataspace;
    goto cleanup;
  }

  dataset = H5Dcreate(id, name, datatype, dataspace, H5P_DEFAULT, H5P_DEFAULT, H5P_DEFAULT);
  if (dataset < 0) {
    status = dataset;
    goto cleanup;
  }

  status = H5Dwrite(dataset, H5T_NATIVE_DOUBLE, H5S_ALL, H5S_ALL, H5P_DEFAULT, t->data);
  if (status != 0) {
    goto cleanup;
  }

cleanup:
  if (dataset >= 0) {
    H5Dclose(dataset);
  }

  if (dataspace >= 0) {
    H5Sclose(dataspace);
  }

  if (datatype >= 0) {
    H5Tclose(datatype);
  }

  H5Oflush(id);

  return status;
}

int
qdm_ijk_read(
    hid_t id,
    const char *name,
    qdm_ijk **t
)
{
  int status = 0;

  int rank = 0;

  status = H5LTget_dataset_ndims(id, name, &rank);
  if (status < 0) {
    return status;
  }

  if (rank != 3) {
    return -1;
  }

  hsize_t dims[3];

  status = H5LTget_dataset_info(id, name, dims, NULL, NULL);
  if (status < 0) {
    return status;
  }

  size_t size3 = dims[0];
  size_t size1 = dims[1];
  size_t size2 = dims[2];

  qdm_ijk *tmp = qdm_ijk_alloc(size1, size2, size3);

  status = H5LTread_dataset_double(id, name, tmp->data);
  if (status < 0) {
    qdm_ijk_free(tmp);

    return status;
  }

  *t = tmp;

  return status;
}


//setzt constanten auf
#define EXTERN_VARIABLES
#include <stdio.h>
#include <math.h>
#include <gsl/gsl_sf_bessel.h>
#include "global.h"
#include "utility.h"
#include "steppmethods.h"
#include <gsl/gsl_cblas.h>
#include <time.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#include <omp.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>

void Set_Target_Length(double * TARGET_LENGTH, double volfrac,  void (*Stepper_Method)(int N,  double  x[N/2] , double v[N/2], double a[N/2], double *t,
					void (*deriv) (double *y, double *ans, double t,int N)),
					void (*Poti_Handle) (double *y, double *ans, double t,int N))
{	int FLAG = 1;
	*TARGET_LENGTH = 0.0;
	if (Stepper_Method == VVerlet_Step_deriv)
	{
		 *TARGET_LENGTH = 0.0;
		printf(" No  Targets\n");
		FLAG = 0;
	}
	//if (Stepper_Method == VVerlet_Step_Target_Circle)
	//{
	//	*TARGET_LENGTH = LATTICE_SPACING * sqrt(volfrac/ M_PI);
	//	printf(" Targets set to hard circles\n");
	//}

	if (Stepper_Method == VVerlet_Step_hardsphere_reflect)
	{
		 *TARGET_LENGTH = LATTICE_SPACING * sqrt(volfrac/ M_PI);
		printf(" Targets set to hard circles\n");
		FLAG = 0;
	}

	if (Stepper_Method == VVerlet_Step_hardsphere_reflect_Kupf)
	{
		 *TARGET_LENGTH = LATTICE_SPACING * sqrt(volfrac/ M_PI);
		printf(" Targets set to hard circles Kupferman \n");
		FLAG = 0;
	}

	if (Stepper_Method == VVerlet_Step_hardsphere_reflect_all)
	{
		 *TARGET_LENGTH = LATTICE_SPACING * sqrt(volfrac/ M_PI);
		printf(" Targets set to hard circles, reflect ALL particles\n");
		FLAG = 0;
	}

	if (Stepper_Method == VVerlet_Step_deriv_Kupf)
	{
		*TARGET_LENGTH = 0.0;
		printf(" No  Targets\n, Kupfermann bath");
		FLAG = 0;
	}

	if (Stepper_Method == VVerlet_Step_deriv_Box)
	{
		 //*TARGET_LENGTH = LATTICE_SPACING * sqrt(volfrac/ M_PI/5.0);
		//printf(" Targets set to soft Yukawa circles in 3 per pore corner\n");
		*TARGET_LENGTH = 0;
		FLAG = 0;
		if (DIM > 1)
		{
			printf("wrong Dimension, please switch to 1d");
			exit(1);
		}
		if (volfrac > 0.0)
		{
			printf("no need for volume fractions >0 \n");
			exit(1);
		}
		printf("box set up with hard wall \n");
	}
	if(FLAG)
	{
		printf("Kein bekannter Stepper in global, cannot set tarhet length!\n");
		exit(1);
	}
}

void Constants(){
// call to setup the bathh parameters and give rudimentary output for the run
	int i;
	OUTPUT_FLAG = 1;
	printf ("OSSZI = %d  DIM =%d  ORDER =%d  \n", OSSZI, DIM, ORDER);
	printf ("KBOLTZ = %e \nTIME_STEPS =%f  \nTIME_END =%f  \n", KBOLTZ, TIME_STEPS, TIME_END);
	for(i=0;i<OSSZI;i++){
		coupling[i] = 0.0;
	}
	for(i=0;i<OSSZI;i++){
		ommega[i] = 0.0;
	}
	for(i=0;i<OSSZI;i++){
		massq[i] = 1.0;
	}
	for(i=0;i<ORDER;i++){
		y[i] = 1.0;
	}
	printf("ALPHA = %1.2E  GAMMA = %1.2E  mass  = %1.2E \n", ALPHA, GAMMA, mass);
	printf("NUMBER_TO_INF = %d  \nTIME_ARRIVAL =%1.2E  \n ", NUMBER_TO_INF, TIME_ARRIVAL);
	printf("Startbedingungen = ");
	printf("%s",LABEL);
	printf("\n");
	if (!(OMMEGA_READ_BINARY))
	{
		// random Number Preparation for Taus GSL genrator
		gsl_rng * r;
		const gsl_rng_type * T;
		gsl_rng_env_setup();
		T = gsl_rng_taus2;
		r = gsl_rng_alloc (T);
		gsl_rng_set(r, time(NULL)); // Seed with time
		// Random number end //
	    Ommega_Setup(r);
	    gsl_rng_free (r);
	}else	// read in from ommega.dat
	{
		FILE *fp;
		fp = fopen("ommega.dat", "r");
		for (i=0;i<OSSZI;i++)
		{
			if (fscanf(fp, "%lf", &ommega[i]) != 1)
			{
	            printf("ERROT READING ommega.dat");
	            break;
        	}
		}
		fclose(fp);
		printf("read in ommega.dat\n");
	}
	if(VIRTUAL_FLAG) printf("Virtuelles Teilchen benutzt, ohne Badrückkopplung");
  	Gamma_Setup(coupling);

  	double ommega_max = ommega[cblas_idamax(OSSZI, ommega, 1)];
  	double coupling_max = coupling[cblas_idamax(OSSZI, coupling, 1)];
  	printf("MAX coupling =  %1.2e\n",coupling_max);
  	if(!(MAXSTEPSIZE_BINARY))
	{
  		MAXSTEPSIZE = (2*M_PI/ommega[cblas_idamax(OSSZI, ommega, 1)] * MAXSTEPSIZEMULT);
  		printf("MAXSTEPSIZE =  %1.2e set according to smallest period times %3.3e\n",MAXSTEPSIZE, MAXSTEPSIZEMULT);
  	}else
  	{
  		MAXSTEPSIZE = (MAXSTEPSIZEMULT / (double) OSSZI);
  		printf("MAXSTEPSIZE =  %1.2e set according to OSSZInr times Stepsize dt = %3.3e\n",MAXSTEPSIZE, MAXSTEPSIZEMULT);
  	}
  	Sum_Method =  SUMMATION_METHOD;
  	Stepper_Method = STEPPER_METHOD;
  	Poti_Handle = POTI_HANDLE;

  	int m = ipow(NUMBER_TO_INF,DIM); int n = DIM;
  	// Setze Gitter auf auf globale Matrix POSITIONS
	POSITIONS = (double **) malloc(m * sizeof(int *));
	POSITIONS[0] = (double *) malloc(m* n * sizeof(int));
	for (i=1; i<m; i++) POSITIONS[i] = POSITIONS[0] + n *i;
	vec_zero_i(lattice_position,DIM);
	printf("lattice allocated, current lattice cell set to zero\n");
	TARGET_CONTACT = 0;
	OUTPUT_FLAG = 0; // subsequent setup calls should not display text
	printf("Running on %d threads, optimally DIM = %d\n", THREADS, DIM);

	// set masses for chosen type of coupling
	if( !(KUPF_BINARY) )
	{
		for (int osi = 0; osi < OSSZI; osi ++)
		{
			massq[osi] 		= 1.0;
		}
		printf("Set all masses équal\n");
	}else
	{
		for (int osi = 0; osi < OSSZI; osi ++)
		{
			massq[osi] 		= coupling[osi] / ommega[osi] / ommega[osi];

		}
		printf("Set all masses according to k = m w^2l\n");
	}
	// Setup masses of bath for simpsons rule, split ends, odd and even parts

	if (SIMPSON_BINARY)
	{
		for (int osi = 1; osi < OSSZI/2-1; osi ++)
		{
			massq[osi * 2] 		*= 2.0/3.0;
			massq[osi * 2 + 1] 	*= 4.0/3.0;
		}
		massq[0] 		*= 1.0/3.0;;
		massq[OSSZI-1] 	*= 1.0/3.0;

		if (KUPF_BINARY)
		{
			for (int osi = 0; osi < OSSZI; osi ++)
			{
				coupling[osi] = massq[osi] * ommega[osi] * ommega[osi];
			}
		}
	}
}