/*
   - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
   SLEPc - Scalable Library for Eigenvalue Problem Computations
   Copyright (c) 2002-, Universitat Politecnica de Valencia, Spain

   This file is part of SLEPc.
   SLEPc is distributed under a 2-clause BSD license (see LICENSE).
   - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
*/
/*
   LME routines related to problem setup
*/

#include <slepc/private/lmeimpl.h>       /*I "slepclme.h" I*/

static inline PetscErrorCode LMESetUp_Lyapunov(LME lme)
{
  Mat            C1,C2,X1,X2;
  Vec            dc,dx;

  PetscFunctionBegin;
  PetscCall(MatLRCGetMats(lme->C,NULL,&C1,&dc,&C2));
  PetscCheck(C1==C2,PetscObjectComm((PetscObject)lme),PETSC_ERR_ARG_WRONGSTATE,"Lyapunov matrix equation requires symmetric right-hand side C");
  PetscCheck(!dc,PetscObjectComm((PetscObject)lme),PETSC_ERR_ARG_WRONGSTATE,"Lyapunov solvers currently require positive-definite right-hand side C");
  if (lme->X) {
    PetscCall(MatLRCGetMats(lme->X,NULL,&X1,&dx,&X2));
    PetscCheck(X1==X2,PetscObjectComm((PetscObject)lme),PETSC_ERR_ARG_WRONGSTATE,"Lyapunov matrix equation requires symmetric solution X");
    PetscCheck(!dx,PetscObjectComm((PetscObject)lme),PETSC_ERR_ARG_WRONGSTATE,"Lyapunov solvers currently assume a positive-definite solution X");
  }
  PetscFunctionReturn(PETSC_SUCCESS);
}

/*@
   LMESetUp - Sets up all the internal data structures necessary for the
   execution of the linear matrix equation solver.

   Collective

   Input Parameter:
.  lme - the linear matrix equation solver context

   Notes:
   This function need not be called explicitly in most cases, since `LMESolve()`
   calls it. It can be useful when one wants to measure the set-up time
   separately from the solve time.

   Level: developer

.seealso: [](ch:lme), `LMECreate()`, `LMESolve()`, `LMEDestroy()`
@*/
PetscErrorCode LMESetUp(LME lme)
{
  PetscInt       N;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(lme,LME_CLASSID,1);

  /* reset the convergence flag from the previous solves */
  lme->reason = LME_CONVERGED_ITERATING;

  if (lme->setupcalled) PetscFunctionReturn(PETSC_SUCCESS);
  PetscCall(PetscLogEventBegin(LME_SetUp,lme,0,0,0));

  /* Set default solver type (LMESetFromOptions was not called) */
  if (!((PetscObject)lme)->type_name) PetscCall(LMESetType(lme,LMEKRYLOV));

  /* Check problem dimensions */
  PetscCheck(lme->A,PetscObjectComm((PetscObject)lme),PETSC_ERR_ARG_WRONGSTATE,"LMESetCoefficients must be called first");
  PetscCheck(lme->C,PetscObjectComm((PetscObject)lme),PETSC_ERR_ARG_WRONGSTATE,"LMESetRHS must be called first");
  PetscCall(MatGetSize(lme->A,&N,NULL));
  if (lme->ncv > N) lme->ncv = N;

  /* setup options for the particular equation type */
  switch (lme->problem_type) {
    case LME_LYAPUNOV:
      PetscCall(LMESetUp_Lyapunov(lme));
      break;
    case LME_SYLVESTER:
      LMECheckCoeff(lme,lme->B,"B","Sylvester");
      break;
    case LME_GEN_LYAPUNOV:
      LMECheckCoeff(lme,lme->D,"D","Generalized Lyapunov");
      break;
    case LME_GEN_SYLVESTER:
      LMECheckCoeff(lme,lme->B,"B","Generalized Sylvester");
      LMECheckCoeff(lme,lme->D,"D","Generalized Sylvester");
      LMECheckCoeff(lme,lme->E,"E","Generalized Sylvester");
      break;
    case LME_DT_LYAPUNOV:
      break;
    case LME_STEIN:
      LMECheckCoeff(lme,lme->E,"E","Stein");
      break;
  }
  PetscCheck(lme->problem_type==LME_LYAPUNOV,PetscObjectComm((PetscObject)lme),PETSC_ERR_SUP,"There is no solver yet for this matrix equation type");

  /* call specific solver setup */
  PetscUseTypeMethod(lme,setup);

  /* set tolerance if not yet set */
  if (lme->tol==(PetscReal)PETSC_DETERMINE) lme->tol = SLEPC_DEFAULT_TOL;

  PetscCall(PetscLogEventEnd(LME_SetUp,lme,0,0,0));
  lme->setupcalled = 1;
  PetscFunctionReturn(PETSC_SUCCESS);
}

static inline PetscErrorCode LMESetCoefficients_Private(LME lme,Mat A,Mat *lmeA)
{
  PetscInt       m,n;

  PetscFunctionBegin;
  PetscCall(MatGetSize(A,&m,&n));
  PetscCheck(m==n,PetscObjectComm((PetscObject)lme),PETSC_ERR_ARG_WRONG,"Matrix is non-square");
  if (!lme->setupcalled) PetscCall(MatDestroy(lmeA));
  PetscCall(PetscObjectReference((PetscObject)A));
  *lmeA = A;
  PetscFunctionReturn(PETSC_SUCCESS);
}

/*@
   LMESetCoefficients - Sets the coefficient matrices that define the linear matrix
   equation to be solved.

   Collective

   Input Parameters:
+  lme - the linear matrix equation solver context
.  A   - first coefficient matrix
.  B   - second coefficient matrix
.  D   - third coefficient matrix
-  E   - fourth coefficient matrix

   Notes:
   The matrix equation takes the general form $AXE+DXB=C$, where matrix $C$ is not
   provided here but with `LMESetRHS()`. Not all four matrices must be passed, some
   can be `NULL` instead, see `LMESetProblemType()` for details.

   It must be called before `LMESetUp()`. If it is called again after `LMESetUp()` then
   the `LME` object is reset.

   In order to delete a previously set matrix, pass `NULL` in the corresponding
   argument.

   Level: beginner

.seealso: [](ch:lme), `LMESolve()`, `LMESetUp()`, `LMESetRHS()`, `LMESetProblemType()`
@*/
PetscErrorCode LMESetCoefficients(LME lme,Mat A,Mat B,Mat D,Mat E)
{
  PetscFunctionBegin;
  PetscValidHeaderSpecific(lme,LME_CLASSID,1);
  PetscValidHeaderSpecific(A,MAT_CLASSID,2);
  PetscCheckSameComm(lme,1,A,2);
  if (B) {
    PetscValidHeaderSpecific(B,MAT_CLASSID,3);
    PetscCheckSameComm(lme,1,B,3);
  }
  if (D) {
    PetscValidHeaderSpecific(D,MAT_CLASSID,4);
    PetscCheckSameComm(lme,1,D,4);
  }
  if (E) {
    PetscValidHeaderSpecific(E,MAT_CLASSID,5);
    PetscCheckSameComm(lme,1,E,5);
  }

  if (lme->setupcalled) PetscCall(LMEReset(lme));

  PetscCall(LMESetCoefficients_Private(lme,A,&lme->A));
  if (B) PetscCall(LMESetCoefficients_Private(lme,B,&lme->B));
  else if (!lme->setupcalled) PetscCall(MatDestroy(&lme->B));
  if (D) PetscCall(LMESetCoefficients_Private(lme,D,&lme->D));
  else if (!lme->setupcalled) PetscCall(MatDestroy(&lme->D));
  if (E) PetscCall(LMESetCoefficients_Private(lme,E,&lme->E));
  else if (!lme->setupcalled) PetscCall(MatDestroy(&lme->E));

  lme->setupcalled = 0;
  PetscFunctionReturn(PETSC_SUCCESS);
}

/*@
   LMEGetCoefficients - Gets the coefficient matrices of the matrix equation.

   Collective

   Input Parameter:
.  lme - the linear matrix equation solver context

   Output Parameters:
+  A   - first coefficient matrix
.  B   - second coefficient matrix
.  D   - third coefficient matrix
-  E   - fourth coefficient matrix

   Level: intermediate

.seealso: [](ch:lme), `LMESolve()`, `LMESetCoefficients()`
@*/
PetscErrorCode LMEGetCoefficients(LME lme,Mat *A,Mat *B,Mat *D,Mat *E)
{
  PetscFunctionBegin;
  PetscValidHeaderSpecific(lme,LME_CLASSID,1);
  if (A) *A = lme->A;
  if (B) *B = lme->B;
  if (D) *D = lme->D;
  if (E) *E = lme->E;
  PetscFunctionReturn(PETSC_SUCCESS);
}

/*@
   LMESetRHS - Sets the right-hand side of the matrix equation, as a low-rank
   matrix.

   Collective

   Input Parameters:
+  lme - the linear matrix equation solver context
-  C   - the right-hand side matrix

   Notes:
   The matrix equation takes the general form $AXE+DXB=C$, where matrix $C$ is
   given with this function. $C$ must be a low-rank matrix of type `MATLRC`, that is,
   $C = UDV^*$ where $D$ is diagonal of order $k$, and $U$, $V$ are dense tall-skinny
   matrices with $k$ columns. No sparse matrix must be provided when creating the
   `MATLRC` matrix.

   In equation types that require $C$ to be symmetric, such as Lyapunov, `C` must be
   created with $V=U$ (or `V=NULL`).

   Level: beginner

.seealso: [](ch:lme), `LMESetSolution()`, `LMESetProblemType()`, `LMEGetRHS()`
@*/
PetscErrorCode LMESetRHS(LME lme,Mat C)
{
  Mat            A;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(lme,LME_CLASSID,1);
  PetscValidHeaderSpecific(C,MAT_CLASSID,2);
  PetscCheckSameComm(lme,1,C,2);
  PetscCheckTypeName(C,MATLRC);

  PetscCall(MatLRCGetMats(C,&A,NULL,NULL,NULL));
  PetscCheck(!A,PetscObjectComm((PetscObject)C),PETSC_ERR_SUP,"The MatLRC must not have a sparse matrix term");

  PetscCall(PetscObjectReference((PetscObject)C));
  PetscCall(MatDestroy(&lme->C));
  lme->C = C;
  PetscFunctionReturn(PETSC_SUCCESS);
}

/*@
   LMEGetRHS - Gets the right-hand side of the matrix equation.

   Collective

   Input Parameter:
.  lme - the linear matrix equation solver context

   Output Parameter:
.  C   - the right-hand side matrix

   Level: intermediate

.seealso: [](ch:lme), `LMESolve()`, `LMESetRHS()`
@*/
PetscErrorCode LMEGetRHS(LME lme,Mat *C)
{
  PetscFunctionBegin;
  PetscValidHeaderSpecific(lme,LME_CLASSID,1);
  PetscAssertPointer(C,2);
  *C = lme->C;
  PetscFunctionReturn(PETSC_SUCCESS);
}

/*@
   LMESetSolution - Sets the placeholder for the solution of the matrix
   equation, as a low-rank matrix.

   Collective

   Input Parameters:
+  lme - the linear matrix equation solver context
-  X   - the solution matrix

   Notes:
   The matrix equation takes the general form $AXE+DXB=C$, where the solution
   matrix is of low rank and is written in factored form $X = UDV^*$. This function
   provides a `Mat` object of type `MATLRC` that stores $U$, $V$ and (optionally) $D$.
   These factors will be computed during `LMESolve()`.

   In equation types whose solution $X$ is symmetric, such as Lyapunov, $X$ must be
   created with $V=U$ (or `V=NULL`).

   If the user provides `X` with this function, then the solver will
   return a solution with rank at most the number of columns of $U$. Alternatively,
   it is possible to let the solver choose the rank of the solution, by
   setting `X` to `NULL` and then calling `LMEGetSolution()` after `LMESolve()`.

   Level: intermediate

.seealso: [](ch:lme), `LMEGetSolution()`, `LMESetRHS()`, `LMESetProblemType()`, `LMESolve()`
@*/
PetscErrorCode LMESetSolution(LME lme,Mat X)
{
  Mat            A;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(lme,LME_CLASSID,1);
  if (X) {
    PetscValidHeaderSpecific(X,MAT_CLASSID,2);
    PetscCheckSameComm(lme,1,X,2);
    PetscCheckTypeName(X,MATLRC);
    PetscCall(MatLRCGetMats(X,&A,NULL,NULL,NULL));
    PetscCheck(!A,PetscObjectComm((PetscObject)X),PETSC_ERR_SUP,"The MatLRC must not have a sparse matrix term");
    PetscCall(PetscObjectReference((PetscObject)X));
  }
  PetscCall(MatDestroy(&lme->X));
  lme->X = X;
  PetscFunctionReturn(PETSC_SUCCESS);
}

/*@
   LMEGetSolution - Gets the solution of the matrix equation.

   Collective

   Input Parameter:
.  lme - the linear matrix equation solver context

   Output Parameter:
.  X   - the low-rank matrix

   Notes:
   If called after `LMESolve()`, `X` will contain the solution of the equation.

   The matrix `X` may have been passed by the user via `LMESetSolution()`,
   although this is not required.

   Level: beginner

.seealso: [](ch:lme), `LMESolve()`, `LMESetSolution()`
@*/
PetscErrorCode LMEGetSolution(LME lme,Mat *X)
{
  PetscFunctionBegin;
  PetscValidHeaderSpecific(lme,LME_CLASSID,1);
  PetscAssertPointer(X,2);
  *X = lme->X;
  PetscFunctionReturn(PETSC_SUCCESS);
}

/*@
   LMEAllocateSolution - Allocate memory storage for common variables such
   as the basis vectors.

   Collective

   Input Parameters:
+  lme   - the linear matrix equation solver context
-  extra - number of additional positions, used for methods that require a
           working basis slightly larger than `ncv`

   Developer Notes:
   This is `SLEPC_EXTERN` because it may be required by user plugin `LME`
   implementations.

   Level: developer

.seealso: [](ch:lme), `LMESetUp()`, `LMESetDimensions()`
@*/
PetscErrorCode LMEAllocateSolution(LME lme,PetscInt extra)
{
  PetscInt       oldsize,requested;
  Vec            t;

  PetscFunctionBegin;
  requested = lme->ncv + extra;

  /* oldsize is zero if this is the first time setup is called */
  PetscCall(BVGetSizes(lme->V,NULL,NULL,&oldsize));

  /* allocate basis vectors */
  if (!lme->V) PetscCall(LMEGetBV(lme,&lme->V));
  if (!oldsize) {
    if (!((PetscObject)lme->V)->type_name) PetscCall(BVSetType(lme->V,BVMAT));
    PetscCall(MatCreateVecsEmpty(lme->A,&t,NULL));
    PetscCall(BVSetSizesFromVec(lme->V,t,requested));
    PetscCall(VecDestroy(&t));
  } else PetscCall(BVResize(lme->V,requested,PETSC_FALSE));
  PetscFunctionReturn(PETSC_SUCCESS);
}
