/******************************************************************************
 *
 * ALBERTA:  an Adaptive multi Level finite element toolbox  using
 *           Bisectioning refinement and Error control by Residual
 *           Techniques
 *
 * File: oem_sp_solve.c
 *
 * Interface to the OEM saddle-point CG-method oem_spcg. NOTE: this
 * only for symmetric saddle-point problems.
 *
 ******************************************************************************
 *
 *  author:     Claus-Justus Heine
 *              Abteilung fuer Angewandte Mathematik
 *              Albert-Ludwigs-Universitaet Freiburg
 *              Hermann-Herder-Str. 10
 *              79104 Freiburg
 *              Germany
 *              claus@mathematik.uni-freiburg.de
 *
 *  http://www.mathematik.uni-freiburg.de/IAM/ALBERTA
 *
 *  (c) by C.-J. Heine (2006)
 *
 ******************************************************************************/

#include "alberta_intern.h"
#include "alberta.h"

/* In the case of multiple constraints we need projections matrices
 * between the different constraint spaces to implement a sensible
 * precondioner.
 *
 * The first constraint has no projections attached, the second holds
 * the projection from the first to the second space, the third
 * constraint data holds the projections from the first to the third
 * and the second to the third and so on:
 *
 * C1: --
 * C2: P12
 * C3: P13 P23
 * C4: P14 P24 P34
 *
 * The projections are matrices; the dual projection is the action of
 * the transpose matrix.
 */
typedef struct proj_chain {
  const DOF_MATRIX *P;
  MatrixTranspose  trans;
  MatrixTranspose  notrans;
  DBL_LIST_NODE    chain;
} PROJ_CHAIN;

typedef struct constraint_chain
{
  const SP_CONSTRAINT *constraint;
  const DOF_REAL_VEC  *rhs;
  DOF_REAL_VEC        *unknown;

  int                 x_dim;
  int                 y_dim;
  
  DOF_REAL_VEC_D      *x_skel;
  DOF_REAL_VEC        *y_skel;

  DBL_LIST_NODE       chain;
  PROJ_CHAIN          *proj_chain;
  DOF_REAL_VEC        *Cr_tmp;   /* storage for rhs and solution*/
  DOF_REAL_VEC        *g_Btu;    /* skeleton to access g_Btu */
  DOF_REAL_VEC        *r;        /* skeleton for projected residual */
  DOF_REAL_VEC        *Cr;       /* Cr skeleton */
} CONSTRAINT_CHAIN;

typedef struct oem_sp_solve_ds_data
{
  OEM_SP_DATA      ospd;
  CONSTRAINT_CHAIN *first_constraint;
  REAL             *Cr_tmp;
} OEM_SP_SOLVE_DS_DATA;

/******************************************************************************
 *
 * Matrix-vector routines to compute (alpha B q -> v) and (alpha Bt v -> q)
 *
 * We support multiple constraint (e.g. divergence + slip b.c. using
 * Verfuerth's saddle-point approach. We use those "chain"-macros,
 * although the constraints do not form a "horizontal" direct sum.
 *
 ******************************************************************************/

static inline
void __Bp_add(CONSTRAINT_CHAIN *data,
	      REAL factor, int dim_y, const REAL *y, int dim_x, REAL *x)
{
  FUNCNAME("__Bp_add");
  DOF_REAL_VEC_D *dof_x = data->x_skel;
  DOF_REAL_VEC   *dof_y = data->y_skel;

  DEBUG_TEST_EXIT(dim_x == data->x_dim,
		  "arguemnt dim_x %d != data->x_dim %d.\n", dim_x, data->x_dim);
  DEBUG_TEST_EXIT(dim_y == data->y_dim,
		  "arguemnt dim_y %d != data->y_dim %d.\n", dim_y, data->y_dim);

  distribute_to_dof_real_vec_d_skel(data->x_skel, x);
  distribute_to_dof_real_vec_skel(data->y_skel, y);

  dof_gemv_dow_scl(NoTranspose,
		   factor, data->constraint->B, data->constraint->bound,
		   dof_y, 1.0, dof_x);
}

static void Bp_add(void *ud,
		   REAL factor, int dim_y, const REAL *y, int dim_x, REAL *x)
{
  CONSTRAINT_CHAIN *data = (CONSTRAINT_CHAIN *)ud;
  
  CHAIN_DO(data, CONSTRAINT_CHAIN) {
    __Bp_add(data, factor, data->y_dim, y, data->x_dim, x);
    y += data->y_dim;
  } CHAIN_WHILE(data, CONSTRAINT_CHAIN);
}

static inline
void __Btu_add(CONSTRAINT_CHAIN *data,
	       REAL factor, int dim_x, const REAL *x, int dim_y, REAL *y)
{
  FUNCNAME("__Btu_add");
  DOF_REAL_VEC_D *dof_x = data->x_skel;
  DOF_REAL_VEC   *dof_y = data->y_skel;

  DEBUG_TEST_EXIT(dim_x == data->x_dim,
		  "arguemnt dim_x %d != data->x_dim %d.\n", dim_x, data->x_dim);
  DEBUG_TEST_EXIT(dim_y == data->y_dim,
		  "arguemnt dim_y %d != data->y_dim %d.\n", dim_y, data->y_dim);

  distribute_to_dof_real_vec_d_skel(data->x_skel, x);
  distribute_to_dof_real_vec_skel(data->y_skel, y);

  if (data->constraint->Bt != NULL) {
    dof_gemv_scl_dow(NoTranspose,
		     factor, data->constraint->Bt, NULL, dof_x, 1.0, dof_y);
  } else {
    dof_gemv_scl_dow(Transpose,
		     factor, data->constraint->B, NULL, dof_x, 1.0, dof_y);
  }
}

static void Btu_add(void *ud,
		    REAL factor, int dim_x, const REAL *x, int dim_y, REAL *y)
{
  CONSTRAINT_CHAIN *data = (CONSTRAINT_CHAIN *)ud;

  CHAIN_DO(data, CONSTRAINT_CHAIN) {
    __Btu_add(data, factor, data->x_dim, x, data->y_dim, y);
    y += data->y_dim;
  } CHAIN_WHILE(data, CONSTRAINT_CHAIN);
}

static inline
CONSTRAINT_CHAIN *init_constraint_chain(
  const SP_CONSTRAINT *constr, const DOF_REAL_VEC *g, DOF_REAL_VEC *p)
{
  CONSTRAINT_CHAIN *data;
  const FE_SPACE *y_space, *x_space;
  
  data             = MEM_CALLOC(1, CONSTRAINT_CHAIN);
  data->constraint = constr;
  data->rhs        = g;
  data->unknown    = p;
  x_space          = constr->B->row_fe_space;
  y_space          = constr->B->col_fe_space;

  data->x_skel = init_dof_real_vec_d_skel(
    MEM_ALLOC(CHAIN_LENGTH(x_space), DOF_REAL_VEC_D), "x skel", x_space);
  data->y_skel = init_dof_real_vec_skel(
    MEM_ALLOC(CHAIN_LENGTH(y_space), DOF_REAL_VEC), "y_skel", y_space);

  data->x_dim = dof_real_vec_d_length(x_space);
  data->y_dim = dof_real_vec_length(y_space);

  CHAIN_INIT(data);

  data->proj_chain = NULL;

  data->Cr_tmp = init_dof_real_vec_skel(
    MEM_ALLOC(CHAIN_LENGTH(y_space), DOF_REAL_VEC), "Cr_tmp", y_space);
  data->Cr = init_dof_real_vec_skel(
    MEM_ALLOC(CHAIN_LENGTH(y_space), DOF_REAL_VEC), "Cr", y_space);
  data->r = init_dof_real_vec_skel(
    MEM_ALLOC(CHAIN_LENGTH(y_space), DOF_REAL_VEC), "r", y_space);
  data->g_Btu = init_dof_real_vec_skel(
    MEM_ALLOC(CHAIN_LENGTH(y_space), DOF_REAL_VEC), "g_Btu", y_space);

  return data;
}

static inline
void __free_constraint_chain(CONSTRAINT_CHAIN *data)
{
  CHAIN_DEL(data);
  if (data->proj_chain) {
    DBL_LIST_NODE *next;
    PROJ_CHAIN *chain;
    CHAIN_FOREACH_SAFE(chain, next, data->proj_chain, PROJ_CHAIN) {
      CHAIN_DEL(chain);
      MEM_FREE(chain, 1, PROJ_CHAIN);
    }
    MEM_FREE(data->proj_chain, 1, PROJ_CHAIN);
  }

  MEM_FREE(data->x_skel, CHAIN_LENGTH(data->x_skel->fe_space), DOF_REAL_VEC_D);
  MEM_FREE(data->y_skel, CHAIN_LENGTH(data->y_skel->fe_space), DOF_REAL_VEC);
  MEM_FREE(data->Cr_tmp, CHAIN_LENGTH(data->Cr_tmp->fe_space), DOF_REAL_VEC);
  MEM_FREE(data->Cr, CHAIN_LENGTH(data->Cr->fe_space), DOF_REAL_VEC);
  MEM_FREE(data->r, CHAIN_LENGTH(data->r->fe_space), DOF_REAL_VEC);
  MEM_FREE(data->g_Btu, CHAIN_LENGTH(data->g_Btu->fe_space), DOF_REAL_VEC);

  MEM_FREE(data, 1, CONSTRAINT_CHAIN);
}

static
void free_constraint_chain(CONSTRAINT_CHAIN *data)
{
  DBL_LIST_NODE *next;
  CONSTRAINT_CHAIN *chain;

  CHAIN_FOREACH_SAFE(chain, next, data, CONSTRAINT_CHAIN) {
    __free_constraint_chain(chain);
  }
  __free_constraint_chain(data);
}

static inline
int __Yproject(CONSTRAINT_CHAIN *data, int dimY, const REAL *rhs, REAL *p)
{
  const SP_CONSTRAINT *constr = data->constraint;
  int iter = 0;

  if (constr->project) {
    iter = constr->project(constr->project_data, dimY, rhs, p);
  } else {
    dcopy(dimY, rhs, 1, p, 1);
  }

  return iter;
}

static
int Yproject(void *ud, int dimY, const REAL *rhs, REAL *p)
{
  OEM_SP_SOLVE_DS_DATA *data  = (OEM_SP_SOLVE_DS_DATA *)ud;
  CONSTRAINT_CHAIN     *chain = data->first_constraint;
  int iter = 0;

  CHAIN_DO(chain, CONSTRAINT_CHAIN) {
    int sub_iter;
    
    sub_iter = __Yproject(chain, chain->y_dim, rhs, p);
    
    rhs += chain->y_dim;
    p   += chain->y_dim;
    
    iter = MAX(iter, sub_iter);

  } CHAIN_WHILE(chain, CONSTRAINT_CHAIN);

  return iter;
}

static inline
int __Yprecon(CONSTRAINT_CHAIN *data,
	      int dimY, const REAL *g_Btu, const REAL *r, REAL *Cr)
{
  const SP_CONSTRAINT *constr = data->constraint;
  int iter = 0, i;

  if (constr->precon) {
    iter = constr->precon(constr->precon_data, dimY, g_Btu, Cr);
  
    for (i = 0; i < dimY; i++) {
      Cr[i] = constr->proj_factor * r[i] + constr->prec_factor * Cr[i];
    }
  } else {
    dcopy(dimY, r, 1, Cr, 1);
  }

  return iter;
}

static
int Yprecon(void *ud, int dimY, const REAL *g_Btu, const REAL *r, REAL *Cr)
{
  OEM_SP_SOLVE_DS_DATA *data  = (OEM_SP_SOLVE_DS_DATA *)ud;
  CONSTRAINT_CHAIN     *chain = data->first_constraint;
  const REAL *g_Btu_saved = g_Btu;
  const REAL *r_saved     = r;
  REAL *Cr_saved          = Cr;
  int iter = 0;

  if (CHAIN_SINGLE(chain)) {
    iter = __Yprecon(chain, chain->y_dim, g_Btu, r, Cr);
  } else {
    CHAIN_DO(chain, CONSTRAINT_CHAIN) {
      distribute_to_dof_real_vec_skel(chain->g_Btu, g_Btu);
      distribute_to_dof_real_vec_skel(chain->Cr, Cr);
      distribute_to_dof_real_vec_skel(chain->r, r);
      g_Btu += chain->y_dim;
      Cr    += chain->y_dim;
      r     += chain->y_dim;
    } CHAIN_WHILE(chain, CONSTRAINT_CHAIN);

    g_Btu = g_Btu_saved;
    r     = r_saved;
    Cr    = Cr_saved;

    CHAIN_DO(chain, CONSTRAINT_CHAIN) {
      CONSTRAINT_CHAIN    *other;
      PROJ_CHAIN          *pr     = chain->proj_chain;
      const SP_CONSTRAINT *constr = chain->constraint;
      int sub_iter = 0;
      int y_dim = chain->y_dim;

      if (constr->prec_factor != 0.0 && constr->precon) {
	dcopy(chain->y_dim, g_Btu, 1, chain->Cr_tmp->vec, 1);
	/* add the interdependencies from the other constraints */
	other = data->first_constraint;
	CHAIN_DO(pr, PROJ_CHAIN) {
	  if (other == chain) {
	    other = CHAIN_NEXT(other, CONSTRAINT_CHAIN);
	  }
	  if (pr->P) {
	    const SP_CONSTRAINT *otherconstr = other->constraint;
	    dof_gemv(pr->notrans,
		     otherconstr->prec_factor / constr->prec_factor,
		     pr->P, NULL, /*other->g_Btu*/ other->r,
		     1.0, chain->Cr_tmp);
	  }
	  other = CHAIN_NEXT(other, CONSTRAINT_CHAIN);
	} CHAIN_WHILE(pr, PROJ_CHAIN);
	
	/* solve */
	sub_iter =
	  constr->precon(
	    constr->precon_data, y_dim, chain->Cr_tmp->vec, chain->Cr_tmp->vec);
	/* prolongate the solution to the other constraint spaces, for
	 * later back-projection.
	 */
	other = data->first_constraint;
	CHAIN_DO(pr, PROJ_CHAIN) {
	  if (other == chain) {
	    other = CHAIN_NEXT(other, CONSTRAINT_CHAIN);
	  }
	  if (pr->P) {
	    const SP_CONSTRAINT *otherconstr = other->constraint;
	    TEST_EXIT(fabs(otherconstr->prec_factor)
		      ==
		      fabs(constr->prec_factor),
		      "Sorry, the absolute value of the factors for "
		      "the precondtioners has to be the same.");
	    
	    dof_gemv(pr->trans,
		     otherconstr->prec_factor, pr->P, NULL, chain->Cr_tmp,
		     1.0, other->Cr);
	  }
	  other = CHAIN_NEXT(other, CONSTRAINT_CHAIN);
	} CHAIN_WHILE(pr, PROJ_CHAIN);
      }

#if 0
      /* Add "non-preconditioned" contributions from the other constraints */
      other = data->first_constraint;
      CHAIN_DO(pr, PROJ_CHAIN) {
	if (other == chain) {
	    other = CHAIN_NEXT(other, CONSTRAINT_CHAIN);
	}
	if (pr->P) {
	  const SP_CONSTRAINT *otherconstr = other->constraint;
	  dof_gemv(pr->notrans,
		   otherconstr->proj_factor, pr->P, NULL,
		   other->g_Btu,
		   1.0, chain->Cr);
	}
	other = CHAIN_NEXT(other, CONSTRAINT_CHAIN);
      } CHAIN_WHILE(pr, PROJ_CHAIN);
#endif

      g_Btu += chain->y_dim;
      Cr    += chain->y_dim;
    
      iter = MAX(iter, sub_iter);

    } CHAIN_WHILE(chain, CONSTRAINT_CHAIN);
    /* Now we have the contributions from the projected residuals of
     * the other constraints in chain->Cr_skel, the solution from the
     * "preconditioner" in chain->tmp. We now have to project
     * chain->Cr_skel one more time.
     */
    g_Btu = g_Btu_saved;
    r     = r_saved;
    Cr    = Cr_saved;

    CHAIN_DO(chain, CONSTRAINT_CHAIN) {
      const SP_CONSTRAINT *constr = chain->constraint;
      int i, sub_iter;
      
      /* solve */
      sub_iter = __Yproject(chain, chain->y_dim, Cr, Cr);

      if (constr->prec_factor != 0.0 && constr->precon) {
	for (i = 0; i < chain->y_dim; i++) {
	  Cr[i] +=
	    constr->prec_factor * chain->Cr_tmp->vec[i]
	    +
	    constr->proj_factor * r[i];
	}
      } else {
	for (i = 0; i < chain->y_dim; i++) {
	  Cr[i] += constr->proj_factor * r[i];
	}
      }
      r  += chain->y_dim;
      Cr += chain->y_dim;
    
      iter = MAX(iter, sub_iter);

    } CHAIN_WHILE(chain, CONSTRAINT_CHAIN);
  }

  return iter;
}

static inline
int extract_vecs_for_oem_dow(REAL **fvecp, REAL **uvecp,
			     const DOF_REAL_VEC_D *f, const DOF_REAL_VEC_D *u)
{
  int dimX = dof_real_vec_d_length(u->fe_space);

  if (!CHAIN_SINGLE(u)) {
    *uvecp = MEM_ALLOC(dimX, REAL);
    copy_from_dof_real_vec_d(*uvecp, u);
    if (f != NULL) {
      *fvecp = MEM_ALLOC(dimX, REAL);
      copy_from_dof_real_vec_d(*fvecp, f);
    }
  } else {
    FOR_ALL_USED_FREE_DOFS(u->fe_space->admin,
			   if (u->stride != 1) {
			     SET_DOW(0.0, ((REAL_D *)u->vec)[dof]);
			     if (f) {
			       SET_DOW(0.0, ((REAL_D *)f->vec)[dof]);
			     }
			   } else {
			     u->vec[dof] = 0.0;
			     if (f) {
			       f->vec[dof] = 0.0;
			     }
			   });
    *uvecp = u->vec;
    *fvecp = f ? f->vec : NULL;
  }
  return dimX;
}

static inline
int extract_vecs_for_oem_scl(REAL **gvecp, REAL **pvecp,
			     const DOF_REAL_VEC *g, const DOF_REAL_VEC *p)
{
  return extract_vecs_for_oem_dow(
    gvecp, pvecp, (const DOF_REAL_VEC_D *)g, (const DOF_REAL_VEC_D *)p);
}

static inline int
extract_constraint_vecs(REAL **gvecp, REAL **pvecp, CONSTRAINT_CHAIN *chain)
{
  int dimY = 0;
  
  if (CHAIN_SINGLE(chain)) {
    return extract_vecs_for_oem_scl(gvecp, pvecp, chain->rhs, chain->unknown);
  } else {
    REAL *gvec, *pvec;

    /* determine the total length of the constraint space */
    CHAIN_DO(chain, CONSTRAINT_CHAIN) {
      dimY += chain->y_dim;
    } CHAIN_WHILE(chain, CONSTRAINT_CHAIN);

    pvec = *pvecp = MEM_ALLOC(dimY, REAL);
    gvec = *gvecp = MEM_ALLOC(dimY, REAL);
    
    CHAIN_DO(chain, CONSTRAINT_CHAIN) {
      copy_from_dof_real_vec(pvec, chain->unknown);
      if (chain->rhs != NULL) {
	copy_from_dof_real_vec(gvec, chain->rhs);
      } else {
	memset(gvec, 0, chain->y_dim * sizeof(REAL));
      }

      pvec += chain->y_dim;
      gvec += chain->y_dim;

    } CHAIN_WHILE(chain, CONSTRAINT_CHAIN);

    return dimY;
  }
}

static inline
void distribute_vecs_from_oem_dow(DOF_REAL_VEC_D *u, REAL *uvec, REAL *fvec,
				  int dimX)
{
  if (!CHAIN_SINGLE(u)) {
    if (fvec) {
      MEM_FREE(fvec, dimX, REAL);
    }
    copy_to_dof_real_vec_d(u, uvec);
    MEM_FREE(uvec, dimX, REAL);  
  }
}

static inline
void distribute_vecs_from_oem_scl(DOF_REAL_VEC *p, REAL *pvec, REAL *gvec,
				  int dimY)
{
  distribute_vecs_from_oem_dow((DOF_REAL_VEC_D *)p, pvec, gvec, dimY);
}

static inline
void distribute_constraint_vecs(CONSTRAINT_CHAIN *chain, REAL *pvec, REAL *gvec,
				int dimY)
{
  if (CHAIN_SINGLE(chain)) {
    distribute_vecs_from_oem_scl(chain->unknown, pvec, gvec, chain->y_dim);
  } else {
    REAL *chain_vec = pvec;
    CHAIN_DO(chain, CONSTRAINT_CHAIN) {
      copy_to_dof_real_vec(chain->unknown, chain_vec);
      chain_vec += chain->y_dim;
    } CHAIN_WHILE(chain, CONSTRAINT_CHAIN);
    MEM_FREE(pvec, dimY, REAL);
    MEM_FREE(gvec, dimY, REAL);
  }
}

/* Solve a saddle point problem by applying an iterative method to the
 * Schur-complement. Currently, only a PCG method is supported.
 *
 * If we have multiple constraints and one of the constraints
 * implements a Quasi-Stokes like preconditioner, then we also require
 * a projection matrix from Y1 -> Y_2 where Y_i are the fe-spaces for
 * the constraints.
 */
int oem_sp_schur_solve(OEM_SOLVER sp_solver,
		       REAL sp_tol, int sp_max_iter, int sp_info,
		       OEM_MV_FCT principal_inverse,
		       OEM_DATA *principal_data,
		       const DOF_REAL_VEC_D *f,
		       DOF_REAL_VEC_D *u,
		       SP_CONSTRAINT *constraint,
		       const DOF_REAL_VEC *g,
		       DOF_REAL_VEC *p,
		       ...)
{
  OEM_SP_SOLVE_DS_DATA data;
  CONSTRAINT_CHAIN *chain;
  int dimX, dimY, iter, ccnt;
  REAL *uvec, *fvec, *pvec, *gvec = NULL;
  va_list ap;
  bool have_precon;

  memset(&data, 0, sizeof(data));

  TEST_EXIT(sp_solver == CG, "Only implemented for solver == CG.\n");
  TEST_EXIT(FE_SPACE_EQ_P(u->fe_space, f->fe_space),
	    "`Velocity' row and column FE_SPACEs don't match!\n");
  TEST_EXIT(g == NULL || FE_SPACE_EQ_P(p->fe_space, g->fe_space),
	    "`Pressure' row and column FE_SPACEs don't match!\n");

  chain = init_constraint_chain(constraint, g, p);
  have_precon = constraint->precon != NULL;

  ccnt = 1;
  va_start(ap, p);
  while ((constraint = va_arg(ap, SP_CONSTRAINT *)) != NULL) {
    CONSTRAINT_CHAIN *constr;
    g = va_arg(ap, DOF_REAL_VEC *);
    p = va_arg(ap, DOF_REAL_VEC *);
    constr = init_constraint_chain(constraint, g, p);
    CHAIN_ADD_TAIL(chain, constr);
    if (constraint->precon) {
      have_precon = true;
    }
    /* we require ccnt many projection matrices between the
     * individual constraint spaces.
     */
    if (ccnt > 0) {
      CONSTRAINT_CHAIN *other;
      PROJ_CHAIN *pr, *transpr;
      int i;
      
      other = chain;
      constr->proj_chain = pr = MEM_ALLOC(1, PROJ_CHAIN);
      CHAIN_INIT(pr);
      pr->P     = va_arg(ap, const DOF_MATRIX *);
      pr->notrans = NoTranspose;
      pr->trans   = Transpose;
      transpr = MEM_ALLOC(1, PROJ_CHAIN);
      CHAIN_INIT(transpr);
      transpr->P     = pr->P;
      transpr->notrans = Transpose;
      transpr->trans   = NoTranspose;
      if (other->proj_chain == NULL) {
	other->proj_chain = transpr;
      } else {
	CHAIN_ADD_TAIL(other->proj_chain, transpr);
      }
      for (i = 1; i < ccnt; i++) {
	other = CHAIN_NEXT(other, CONSTRAINT_CHAIN);
	pr = MEM_ALLOC(1, PROJ_CHAIN);
	CHAIN_INIT(pr);
	pr->P     = va_arg(ap, const DOF_MATRIX *);
	pr->notrans = NoTranspose;
	pr->trans   = Transpose;
	CHAIN_ADD_TAIL(constr->proj_chain, pr);

	transpr = MEM_ALLOC(1, PROJ_CHAIN);
	CHAIN_INIT(transpr);
	transpr->P     = pr->P;
	transpr->notrans = Transpose;
	transpr->trans   = NoTranspose;
	CHAIN_ADD_TAIL(other->proj_chain, transpr);
      }
    }
    ++ccnt;
  }
  va_end(ap);

  data.first_constraint    = chain;
  data.ospd.info           = sp_info;
  data.ospd.tolerance      = sp_tol;
  data.ospd.max_iter       = sp_max_iter;
  data.ospd.info           = sp_info;
  data.ospd.solve_Auf      = principal_inverse;
  data.ospd.solve_Auf_data = principal_data;
  data.ospd.B              = Bp_add;
  data.ospd.B_data         = chain;
  data.ospd.Bt             = Btu_add;
  data.ospd.Bt_data        = chain;
  data.ospd.project        = Yproject;
  data.ospd.project_data   = &data;

  dimX = extract_vecs_for_oem_dow(&fvec, &uvec, f, u);
  dimY = extract_constraint_vecs(&gvec, &pvec, data.first_constraint);

  if (have_precon) {
    data.ospd.precon      = Yprecon;
    data.ospd.precon_data = &data;
    if (ccnt > 1) {
      REAL *Cr_tmp = data.Cr_tmp = MEM_ALLOC(dimY, REAL);
      CHAIN_DO(chain, CONSTRAINT_CHAIN) {
	distribute_to_dof_real_vec_skel(chain->Cr_tmp, Cr_tmp);
	Cr_tmp += chain->y_dim;
      } CHAIN_WHILE(chain, CONSTRAINT_CHAIN);
    }
  }

  iter = oem_spcg(&data.ospd, dimX, fvec, uvec, dimY, gvec, pvec);

  distribute_constraint_vecs(data.first_constraint, pvec, gvec, dimY);
  distribute_vecs_from_oem_dow(u, uvec, fvec, dimX);

  free_constraint_chain(data.first_constraint);

  if (ccnt > 1 && have_precon) {
    MEM_FREE(data.Cr_tmp, dimY, REAL);
  }

  return iter;
}

SP_CONSTRAINT *init_sp_constraint(
  const DOF_MATRIX *B,
  const DOF_MATRIX *Bt,
  const DOF_SCHAR_VEC *bound,
  REAL tol, int info,
  const DOF_MATRIX *Yproj,
  OEM_SOLVER Yproj_solver, int Yproj_max_iter, const PRECON *Yproj_prec,
  const DOF_MATRIX *Yprec,
  OEM_SOLVER Yprec_solver, int Yprec_max_iter, const PRECON *Yprec_prec,
  REAL Yproj_frac, REAL Yprec_frac)
{
  SP_CONSTRAINT *constr = MEM_CALLOC(1, SP_CONSTRAINT);
  
  constr->B  = B;
  constr->Bt = Bt;
  constr->bound = bound;

  constr->project = get_oem_solver(Yproj_solver);
  constr->project_data = init_oem_solve(
    Yproj, NULL, tol, Yproj_prec, -1 /* restart */, Yproj_max_iter, info);

  if (Yprec) {
    constr->precon = get_oem_solver(Yprec_solver);
    constr->precon_data = init_oem_solve(
      Yprec, NULL, tol, Yprec_prec, -1 /* restart */, Yprec_max_iter, info);
  }

  constr->proj_factor = Yproj_frac;
  constr->prec_factor = Yprec_frac;
  
  return constr;
}

void release_sp_constraint(SP_CONSTRAINT *constr) 
{
  if (constr->precon) {
    release_oem_solve(constr->precon_data);
  }
  release_oem_solve(constr->project_data);
  MEM_FREE(constr, 1, SP_CONSTRAINT);
}

int oem_sp_solve_dow_scl(
  OEM_SOLVER sp_solver,
  REAL sp_tol, REAL tol_incr,
  int sp_max_iter, int sp_info,
  const DOF_MATRIX *A, const DOF_SCHAR_VEC *bound,
  OEM_SOLVER A_solver, int A_max_iter, const PRECON *A_prec,
  DOF_MATRIX *B,
  DOF_MATRIX *Bt,
  DOF_MATRIX *Yproj,
  OEM_SOLVER Yproj_solver, int Yproj_max_iter, const PRECON *Yproj_prec,
  DOF_MATRIX *Yprec,
  OEM_SOLVER Yprec_solver, int Yprec_max_iter, const PRECON *Yprec_prec,
  REAL Yproj_frac, REAL Yprec_frac,
  const DOF_REAL_VEC_D *f,
  const DOF_REAL_VEC *g,
  DOF_REAL_VEC_D *u,
  DOF_REAL_VEC *p)
{
  OEM_SP_SOLVE_DS_DATA data;
  REAL A_tol = sp_tol / tol_incr;
  int  A_info = MAX(0, sp_info - 3);
  int dimX, dimY, iter;
  REAL *uvec, *fvec, *pvec, *gvec = NULL;
  SP_CONSTRAINT *constr;

  TEST_EXIT(sp_solver == CG, "Only implemented for solver == CG.\n");

  TEST_EXIT(B != NULL || Bt != NULL, "Either B or Bt must be != NULL\n");

  TEST_EXIT(FE_SPACE_EQ_P(u->fe_space, f->fe_space),
	    "`Velocity' row and column FE_SPACEs don't match!\n");
  TEST_EXIT(g == NULL || FE_SPACE_EQ_P(p->fe_space, g->fe_space),
	    "`Pressure' row and column FE_SPACEs don't match!\n");

  memset(&data, 0, sizeof(data));

  data.ospd.ws        = NULL;
  data.ospd.tolerance = sp_tol;
  data.ospd.max_iter  = sp_max_iter;
  data.ospd.info      = MAX(0, sp_info);
  data.ospd.solve_Auf = get_oem_solver(A_solver);
  data.ospd.solve_Auf_data = init_oem_solve(
    A, NULL, A_tol, A_prec, -1 /* restart */, A_max_iter, A_info);

  constr = init_sp_constraint(
    B, Bt, bound,
    sp_tol / tol_incr, MAX(0, sp_info-3),
    Yproj, Yproj_solver, Yproj_max_iter, Yproj_prec,
    Yprec, Yprec_solver, Yprec_max_iter, Yprec_prec,
    Yproj_frac, Yprec_frac);

  data.ospd.project      = constr->project;
  data.ospd.project_data = constr->project_data;
  if (constr->precon) {
    data.ospd.precon = Yprecon;
    data.ospd.precon_data = &data;
  }
  
  data.first_constraint = init_constraint_chain(constr, g, p);
  data.ospd.B       = Bp_add;
  data.ospd.B_data  = data.first_constraint;
  data.ospd.Bt      = Btu_add;
  data.ospd.Bt_data = data.first_constraint;

  dimX = extract_vecs_for_oem_dow(&fvec, &uvec, f, u);
  dimY = extract_constraint_vecs(&gvec, &pvec, data.first_constraint);

  iter = oem_spcg(&data.ospd, dimX, fvec, uvec, dimY, gvec, pvec);

  distribute_constraint_vecs(data.first_constraint, pvec, gvec, dimY);
  distribute_vecs_from_oem_dow(u, uvec, fvec, dimX);

  free_constraint_chain(data.first_constraint);

  /* Don't forget to release all the resources afterwards: */
  release_sp_constraint(constr);
  release_oem_solve((const OEM_DATA *)data.ospd.solve_Auf_data);
  
  return iter;
}

/* Compute a modification of the RHS in Y-space to assure the
 * solvability of the linear system in the presence of Dirichlet
 * boundary conditions. This can be done on the linear algebra level.
 *
 * NOTE: the modification of g_h is performed ONLY when no non-NoSlip
 * boundary segments are encountered. In any case, the flux over the
 * no-slip boundary segments is returned.
 *
 * An initial flux already stored in g_h is taken into account.
 */
REAL sp_dirichlet_bound_dow_scl(MatrixTranspose transpose,
				const DOF_MATRIX *Bt,
				const DOF_SCHAR_VEC *bound,
				const DOF_REAL_VEC_D *u_h,
				DOF_REAL_VEC *g_h)
{
  return sp_flux_adjust_dow_scl(transpose, Bt, bound, u_h, g_h, 0.0, false);
}

REAL sp_flux_adjust_dow_scl(MatrixTranspose transpose,
			    const DOF_MATRIX *Bt,
			    const DOF_SCHAR_VEC *bound,
			    const DOF_REAL_VEC_D *u_h,
			    DOF_REAL_VEC *g_h,
			    REAL flux,
			    bool do_adjust)
{
  int nnzero, dof;
  bool adjust;

  /* First compute the total excess flux generated by the boundary
   * values.
   */
  nnzero = 0;
  adjust = true;
  if (transpose == NoTranspose) {
    bool first = true;
    ROW_CHAIN_DO(Bt, const DOF_MATRIX) {
      COL_CHAIN_DO(Bt, const DOF_MATRIX) {
	if (u_h->stride != 1) {
	  for (dof = 0; dof < Bt->row_fe_space->admin->size_used; dof++) {
	    if (Bt->matrix_row[dof]) {
	      if (first) {
		nnzero ++;
		flux -= g_h->vec[dof];
	      }
	      FOR_ALL_MAT_COLS(REAL_D, Bt->matrix_row[dof],
			       if (bound->vec[col_dof] <= NEUMANN) {
				 adjust = false;
				 /* no need to adjust anything */
			       } else if (bound->vec[col_dof] >= DIRICHLET) {
				 flux += SCP_DOW(row->entry[col_idx],
						 ((REAL_D *)u_h->vec)[col_dof]);
			       });
	    }
	  }
	} else {
	  for (dof = 0; dof < Bt->row_fe_space->admin->size_used; dof++) {
	    if (Bt->matrix_row[dof]) {
	      if (first) {
		nnzero ++;
		flux -= g_h->vec[dof];
	      }
	      FOR_ALL_MAT_COLS(REAL, Bt->matrix_row[dof],
			       if (bound->vec[col_dof] <= NEUMANN) {
				 adjust = false;
				 /* no need to adjust anything */
			       } else if (bound->vec[col_dof] >= DIRICHLET) {
				 flux += 
				   row->entry[col_idx] * u_h->vec[col_dof];
			       });
	    }
	  }
	}
	first = false;
	g_h = CHAIN_NEXT(g_h, DOF_REAL_VEC);
      } COL_CHAIN_WHILE(Bt, const DOF_MATRIX);
      u_h = CHAIN_NEXT(u_h, DOF_REAL_VEC_D);
      bound = CHAIN_NEXT(bound, DOF_SCHAR_VEC);
    } ROW_CHAIN_WHILE(Bt, const DOF_MATRIX);
    
    if (adjust || do_adjust) {
      /* Now distribute the noise across all pressure DOFs */
      REAL dofflux = flux / (REAL)nnzero;
      COL_CHAIN_DO(Bt, const DOF_MATRIX) {
	for (dof = 0; dof < Bt->row_fe_space->admin->size_used; dof++) {
	  if (Bt->matrix_row[dof]) {
	    g_h->vec[dof] += dofflux;
	  }
	}
	g_h = CHAIN_NEXT(g_h, DOF_REAL_VEC);
      } COL_CHAIN_WHILE(Bt, const DOF_MATRIX);
    }
  } else {
    /* So Bt is actually B */
    DOF_SCHAR_VEC *flags;
    
    flags = get_dof_schar_vec("flags", g_h->fe_space);
    FOREACH_DOF(g_h->fe_space,
		flags->vec[dof] = false,
		flags = CHAIN_NEXT(flags, DOF_SCHAR_VEC));
    COL_CHAIN_DO(Bt, const DOF_MATRIX) {
      ROW_CHAIN_DO(Bt, const DOF_MATRIX) {
	if (u_h->stride != 1) {
	  for (dof = 0; dof < Bt->row_fe_space->admin->size_used; dof++) {
	    if (Bt->matrix_row[dof] == NULL) {
	      continue;
	    }
	    if (bound->vec[dof] <= NEUMANN) {
	      adjust = false; /* no need to adjust anything */
	    }
	    if (bound->vec[dof] >= DIRICHLET) {
	      FOR_ALL_MAT_COLS(REAL_D, Bt->matrix_row[dof],
			       flux += SCP_DOW(row->entry[col_idx],
					       ((REAL_D *)u_h->vec)[dof]);
			       if (!flags->vec[col_dof]) {
				 flags->vec[col_dof] = true;
				 nnzero ++;
				 flux -= g_h->vec[col_dof];
			       });
	    }
	  }
	} else {
	  for (dof = 0; dof < Bt->row_fe_space->admin->size_used; dof++) {
	    if (Bt->matrix_row[dof] == NULL) {
	      continue;
	    }
	    if (bound->vec[dof] <= NEUMANN) {
	      adjust = false; /* no need to adjust anything */
	    } else if (bound->vec[dof] >= DIRICHLET) {
	      FOR_ALL_MAT_COLS(REAL, Bt->matrix_row[dof],
			       flux += row->entry[col_idx] * u_h->vec[dof];
			       if (!flags->vec[col_dof]) {
				 flags->vec[col_dof] = true;
				 nnzero ++;
				 flux -= g_h->vec[col_dof];
			       });
	    }
	  }
	}
	g_h = CHAIN_NEXT(g_h, DOF_REAL_VEC);
      } ROW_CHAIN_WHILE(Bt, const DOF_MATRIX);
      u_h = CHAIN_NEXT(u_h, const DOF_REAL_VEC_D);
      bound = CHAIN_NEXT(bound, const DOF_SCHAR_VEC);
    } COL_CHAIN_WHILE(Bt, const DOF_MATRIX);

    if (adjust || do_adjust) {
      /* Now distribute the noise across all pressure DOFs */
      REAL dofflux = flux / (REAL)nnzero;
      FOREACH_DOF(g_h->fe_space,
		  if (flags->vec[dof]) {
		    g_h->vec[dof] += dofflux;
		  },
		  g_h = CHAIN_NEXT(g_h, DOF_REAL_VEC);
		  flags = CHAIN_NEXT(flags, DOF_SCHAR_VEC));
      free_dof_schar_vec(flags);
    }
  }
  return flux;
}

