# -*- Pyrex -*-
#
# ======================================================================
#
#                           Brad T. Aagaard
#                        U.S. Geological Survey
#
# {LicenseText}
#
# ======================================================================
#

#header{
#include <petscmesh.hh>
#include <petscksp.h>

#include <assert.h>
#include <stdexcept>
#include <Python.h>
#}header

# ----------------------------------------------------------------------
cdef extern from "Python.h":
  object PyCObject_FromVoidPtr(void*, void (*destruct)(void*))
  void* PyCObject_AsVoidPtr(object)

cdef void* ptrFromHandle(obj):
  """Extract pointer from PyCObject."""
  return PyCObject_AsVoidPtr(obj.handle)

cdef extern from "stdlib.h":
  ctypedef unsigned long size_t
  void* malloc(size_t size)
  void free(void* mem)

cdef void KSP_destructor(void* obj):
  """
  Destroy KSP object.
  """
  # create shim for destructor
  #embed{ void KSP_destructor_cpp(void* objVptr)
  try {
    KSP* ksp = (KSP*) objVptr;
    PetscErrorCode err = KSPDestroy(*ksp);
    if (err) {
      PetscError(__LINE__,__FUNCT__,__FILE__,__SDIR__,err,0," ");
      throw std::runtime_error("Could not destroy KSP object.");
    } // if
    delete ksp;
  } catch (const std::exception& err) {
    PyErr_SetString(PyExc_RuntimeError,
                    const_cast<char*>(err.what()));
  } catch (const ALE::Exception& err) {
    PyErr_SetString(PyExc_RuntimeError,
                    const_cast<char*>(err.msg().c_str()));
  } catch (...) {
    PyErr_SetString(PyExc_RuntimeError,
                    "Caught unknown C++ exception.");
  } // try/catch
  #}embed
  KSP_destructor_cpp(obj)
  return


# ----------------------------------------------------------------------
cdef class Solver:

  cdef void* thisptr # Pointer to C++ object
  cdef readonly object handle # PyCObject holding pointer to C++ object
  cdef readonly object name # Identifier for object base type
  cdef void* vecScatterVptr # Handle to VecScatter
  cdef void* vecInVptr # Handle to solver rhs
  cdef void* vecOutVptr # Handle to solver solution

  def __init__(self):
    """
    Constructor.
    """
    self.handle = None
    self.thisptr = NULL
    self.name = "pylith_solver_Solver"
    self.vecScatterVptr = NULL
    self.vecInVptr = NULL
    self.vecOutVptr = NULL

    # create shim for constructor
    #embed{ void* KSP_create()
    void* result = 0;
    try {
      KSP* ksp = new KSP;
      PetscErrorCode err = KSPCreate(PETSC_COMM_WORLD, ksp);
      if (err) {
        PetscError(__LINE__,__FUNCT__,__FILE__,__SDIR__,err,0," ");
        throw std::runtime_error("Could not create KSP object.");
      } // if
      err = KSPSetFromOptions(*ksp);
      if (err) {
        PetscError(__LINE__,__FUNCT__,__FILE__,__SDIR__,err,0," ");
        throw std::runtime_error("Could not set KSP options.");
      } // if
      result = (void*) ksp;
    } catch (const std::exception& err) {
      PyErr_SetString(PyExc_RuntimeError,
                      const_cast<char*>(err.what()));
    } catch (const ALE::Exception& err) {
      PyErr_SetString(PyExc_RuntimeError,
                      const_cast<char*>(err.msg().c_str()));
    } catch (...) {
      PyErr_SetString(PyExc_RuntimeError,
                      "Caught unknown C++ exception.");
    } // try/catch
    return result;
    #}embed
    self.thisptr = KSP_create()
    self.handle = self._createHandle()
    return


  def initialize(self, mesh, field):
    """
    Initialzie solver.
    """
    # create shim for method 'initialize'
    #embed{ void Solver_initialize(void* objVptr, void** scatterVptr, void** vecInVptr, void** vecOutVptr, void* meshVptr, void* fieldVptr)
    try {
      assert(0 != objVptr);
      assert(0 != scatterVptr);
      assert(0 != vecInVptr);
      assert(0 != vecOutVptr);
      assert(0 != meshVptr);
      assert(0 != fieldVptr);
      ALE::Obj<ALE::Mesh>* mesh =
        (ALE::Obj<ALE::Mesh>*) meshVptr;
      ALE::Obj<ALE::Mesh::real_section_type>* field =
        (ALE::Obj<ALE::Mesh::real_section_type>*) fieldVptr;
      VecScatter scatter;
      PetscErrorCode err = MeshCreateGlobalScatter(*mesh, *field, &scatter);
      if (err) {
        PetscError(__LINE__,__FUNCT__,__FILE__,__SDIR__,err,0," ");
        throw std::runtime_error("Could not create scatter.");
      } // if
      if (0 != *scatterVptr) {
        err = VecScatterDestroy((VecScatter) *scatterVptr);
        if (err) {
          PetscError(__LINE__,__FUNCT__,__FILE__,__SDIR__,err,0," ");
          throw std::runtime_error("Could not destroy scatter.");
        } // if
      } // if
      *scatterVptr = (void *) scatter;
      Vec in;
      const ALE::Obj<ALE::Mesh::order_type>& order = (*mesh)->getFactory()->getGlobalOrder((*mesh), "default", (*field));

      err = VecCreate((*mesh)->comm(), &in);
      err = VecSetSizes(in, order->getLocalSize(), order->getGlobalSize());
      err = VecSetFromOptions(in);
      if (err) {
        PetscError(__LINE__,__FUNCT__,__FILE__,__SDIR__,err,0," ");
        throw std::runtime_error("Could not create vector.");
      } // if
      if (0 != *vecInVptr) {
        err = VecDestroy((Vec) *vecInVptr);
        if (err) {
          PetscError(__LINE__,__FUNCT__,__FILE__,__SDIR__,err,0," ");
          throw std::runtime_error("Could not destroy vector.");
        } // if
      } // if
      *vecInVptr = (void *) in;
      Vec out;
      err = VecDuplicate(in, &out);
      if (err) {
        PetscError(__LINE__,__FUNCT__,__FILE__,__SDIR__,err,0," ");
        throw std::runtime_error("Could not create vector.");
      } // if
      if (0 != *vecOutVptr) {
        err = VecDestroy((Vec) *vecOutVptr);
        if (err) {
          PetscError(__LINE__,__FUNCT__,__FILE__,__SDIR__,err,0," ");
          throw std::runtime_error("Could not destroy vector.");
        } // if
      } // if
      *vecOutVptr = (void *) out;
    } catch (const std::exception& err) {
      PyErr_SetString(PyExc_RuntimeError,
                      const_cast<char*>(err.what()));
    } catch (const ALE::Exception& err) {
      PyErr_SetString(PyExc_RuntimeError,
                      const_cast<char*>(err.msg().c_str()));
    } catch (...) {
      PyErr_SetString(PyExc_RuntimeError,
                      "Caught unknown C++ exception.");
    } // try/catch      
    #}embed

    if not mesh.name == "pylith_topology_Mesh":
      raise TypeError, \
            "Argument must be extension module type " \
            "'pylith::topology::Mesh'."
    cdef void* fieldVptr
    fieldVptr = PyCObject_AsVoidPtr(field)
    Solver_initialize(self.thisptr, &self.vecScatterVptr,
                      &self.vecInVptr, &self.vecOutVptr,
                      ptrFromHandle(mesh), fieldVptr)
    return


  def _createHandle(self):
    """
    Wrap pointer to C++ object in PyCObject.
    """
    return PyCObject_FromVoidPtr(self.thisptr, KSP_destructor)


# ----------------------------------------------------------------------
cdef class SolverLinear(Solver):

  def __init__(self):
    """
    Constructor.
    """
    Solver.__init__(self)
    return


  def solve(self, fieldOut, jacobian, fieldIn):
    """
    Solve linear system.
    """
    # create shim for method 'solve'
    #embed{ int SolverLinear_solve(void* objVptr, void* fieldOutVptr, void* jacobianVptr, void* fieldInVptr, void* scatterVptr, void* vecInVptr, void* vecOutVptr)
    typedef ALE::Mesh::real_section_type real_section_type;
    try {
      assert(0 != objVptr);
      assert(0 != fieldOutVptr);
      assert(0 != jacobianVptr);
      assert(0 != fieldInVptr);
      assert(0 != scatterVptr);
      assert(0 != vecInVptr);
      assert(0 != vecOutVptr);

      KSP* ksp = (KSP*) objVptr;
      ALE::Obj<real_section_type>* fieldOut =
        (ALE::Obj<real_section_type>*) fieldOutVptr;
      Mat* jacobian = (Mat*) jacobianVptr;
      ALE::Obj<real_section_type>* fieldIn =
        (ALE::Obj<real_section_type>*) fieldInVptr;
      VecScatter scatter = (VecScatter) scatterVptr;
      Vec vecIn = (Vec) vecInVptr;
      Vec vecOut = (Vec) vecOutVptr;

      PetscErrorCode err = 0;
      Vec localVec;

      err = VecCreateSeqWithArray(PETSC_COMM_SELF, (*fieldIn)->sizeWithBC(),
                                  (*fieldIn)->restrict(), &localVec);CHKERRQ(err);
      err = VecScatterBegin(scatter, localVec, vecIn, INSERT_VALUES, SCATTER_FORWARD
                            );CHKERRQ(err);
      err = VecScatterEnd(scatter, localVec, vecIn, INSERT_VALUES, SCATTER_FORWARD
                          ); CHKERRQ(err);
      err = VecDestroy(localVec); CHKERRQ(err);
      err = KSPSetOperators(*ksp, *jacobian, *jacobian,
                            DIFFERENT_NONZERO_PATTERN); CHKERRQ(err);
      err = KSPSolve(*ksp, vecIn, vecOut); CHKERRQ(err);
      err = VecCreateSeqWithArray(PETSC_COMM_SELF, (*fieldOut)->sizeWithBC(),
                                (*fieldOut)->restrict(), &localVec);CHKERRQ(err);
      err = VecScatterBegin(scatter, vecOut, localVec, INSERT_VALUES, SCATTER_REVERSE
                            ); CHKERRQ(err);
      err = VecScatterEnd(scatter, vecOut, localVec, INSERT_VALUES, SCATTER_REVERSE
                          ); CHKERRQ(err);
      err = VecDestroy(localVec); CHKERRQ(err);
    } catch (const std::exception& err) {
      PyErr_SetString(PyExc_RuntimeError,
                      const_cast<char*>(err.what()));
    } catch (const ALE::Exception& err) {
      PyErr_SetString(PyExc_RuntimeError,
                      const_cast<char*>(err.msg().c_str()));
    } catch (...) {
      PyErr_SetString(PyExc_RuntimeError,
                      "Caught unknown C++ exception.");
    } // try/catch      
    #}embed

    cdef void* fieldOutVptr
    cdef void* jacobianVptr
    cdef void* fieldInVptr
    fieldOutVptr = PyCObject_AsVoidPtr(fieldOut)
    jacobianVptr = PyCObject_AsVoidPtr(jacobian)
    fieldInVptr = PyCObject_AsVoidPtr(fieldIn)
    SolverLinear_solve(self.thisptr, fieldOutVptr, jacobianVptr,
                       fieldInVptr, self.vecScatterVptr,
                       self.vecInVptr, self.vecOutVptr)
    return


# End of file