/*--------------------------------------------------------------------------*/
/*  Solving a linear saddle point problem with a preconditioned conjugate   */
/*  gradient method                                                         */
/*  Matrix:         A symmetric and positve definite                        */
/*                  X and Y satisfy LBB condition with respect to B         */
/*                                                                          */
/*                                                                          */
/*        [A   B] [u]   [f]                                                 */
/*        [     ] [ ] = [ ]  in  X x Y                                      */
/*        [B^* 0] [p]   [g]                                                 */
/*                                                                          */
/*  Reformulation using the Schur-Complement operator T := B^* A^{-1} B     */
/*                                                                          */
/*                                                                          */
/*       <=>    T p = B^* A^{-1} B p = B^* A^{-1} f - g   in Y              */
/*                                                                          */
/*  u is then given as  u = A^{-1} f - A^{-1} B p                           */
/*--------------------------------------------------------------------------*/
/*  Workspace:      3*dimY + 2*dimX                                         */
/*                                                                          */
/*  author: Kunibert G. Siebert                                             */
/*          Institut fuer Mathematik                                        */
/*          Universitaet Augsburg                                           */
/*          Universitaetsstr. 14                                            */
/*          D-86159 Augsburg, Germany                                       */
/*                                                                          */
/*          http://scicomp.math.uni-augsburg.de/Siebert/                    */
/*                                                                          */
/* (c) by K.G. Siebert (2000-2003)                                          */
/*--------------------------------------------------------------------------*/

#include "alberta_util.h"
#include "alberta_util_intern.h"

/*--------------------------------------------------------------------------*/
/*---  information output for solver  --------------------------------------*/
/*--------------------------------------------------------------------------*/

static void spcg_start_info(const char *funcName, OEM_SP_DATA *data, int nAuf,
			    int npj, int npc)
{
  data->info = data->info > 6 ? 6 : data->info;

  INFO(data->info, 6, "with tolerance %le", data->tolerance);
  if (data->restart > 0)
    PRINT_INFO(data->info, 6, " and restart %d\n", data->restart);
  else
    PRINT_INFO(data->info, 6, "\n");
  INFO(data->info, 2, "    |    residual |     red. |");

  if (nAuf >= 0) 
  {
    PRINT_INFO(data->info, 2, "   Au=f |");
  }
  if (npj >= 0)  
  {
    PRINT_INFO(data->info, 2, "   B^*u |");
  }
  if (npc >= 0)
  {
    PRINT_INFO(data->info, 2, "     Cr |");
  }

  PRINT_INFO(data->info, 2, "\n");

  fflush(stdout);
  return;
}

static const int step[7] = {0, 0, 10, 5, 2, 1, 1};

static void spcg_break_info(const char *funcName, OEM_SP_DATA *data, 
			    const char *reason,
			    int iter, int nAuf, int npj, int npc,
			    REAL res2, REAL *ores, WORKSPACE *ws)
{
  REAL  res = sqrt(res2);
  int iter_step  = iter == 0 ? 1 : step[data->info];
  int step_mod   = iter_step == 0 ? 1 : (iter % iter_step);
  double divisor = step_mod == 0 ? iter_step : step_mod;

  if (*ores && *ores > 0) {
    INFO(data->info, 2, "%3d | %11.5le | %8.2le |", iter, res, res/(*ores));
  } else {
    INFO(data->info, 2, "%3d | %11.5le |         |", iter);
  }

  if (nAuf >= 0) {
    PRINT_INFO(data->info, 2, " %6.1f |", (double)nAuf/divisor);
  }
  if (npj >= 0) {
    PRINT_INFO(data->info, 2, " %6.1f |", (double)npj/divisor);
  }
  if (npc >= 0) {
    PRINT_INFO(data->info, 2, " ------ |");
  }
  PRINT_INFO(data->info, 2, "\n");
  INFO(data->info, 2, "stop due to: %s\n", reason);

  data->residual = res;
  if (ws != data->ws)
    FREE_WORKSPACE(ws);

  fflush(stdout);
  return;
}

/* Return "true" if the iteration should be stopped, false otherwise. */
static bool spcg_solve_info(const char *funcName, OEM_SP_DATA *data, int iter, 
			    int *nAuf, int *npj, int *npc,
			    REAL res2, REAL *ores, 
			    WORKSPACE *ws)
{
  REAL res = sqrt(res2);
  int iter_step  = iter == 0 ? 1 : step[data->info];
  int step_mod   = iter_step == 0 ? 1 : (iter % iter_step);
  double divisor = step_mod == 0 ? iter_step : step_mod;

  if (res2 < 0.0 && fabs(res2) < SQR(1e-2*data->tolerance)) {
    /* This should be due to solver tolerance, round-off errors etc. */
    res = sqrt(fabs(res2));
  }

  if (res <= data->tolerance ||
      ((data->info > 1) && step_mod == 0) ||
      iter == data->max_iter) {
    if (ores) {
      if (*ores > 0.0) {
	REAL  red = res/(*ores);
	INFO(data->info, 2, "%3d | %11.5le | %8.2le |", iter, res, red);
      } else {
	INFO(data->info, 2, "%3d | %11.5le | -------- |", iter, res);
      }
      *ores = res;
    } else {
      INFO(data->info, 2, "%2d | %11.5le |", iter, res);
    }

    if (*nAuf >= 0) {
      PRINT_INFO(data->info, 2, " %6.1f |", (double)*nAuf/divisor);
      *nAuf = 0;
    }
    if (*npj >= 0) {
      PRINT_INFO(data->info, 2, " %6.1f |", (double)*npj/divisor);
      *npj = 0;
    }
    if (*npc >= 0) {
      PRINT_INFO(data->info, 2, " %6.1f |", (double)*npc/divisor);
      *npc = 0;
    }
    PRINT_INFO(data->info, 2, "\n");
  }
 
  data->residual = res;

  if (iter == data->max_iter && res > data->tolerance) {
    INFO(data->info, 1, "tolerance %le not reached after %d iterations\n", 
		      data->tolerance, iter);
    if (ws != data->ws) FREE_WORKSPACE(ws);

    fflush(stdout);
    return true;
  }

  if (res <= data->tolerance) {
    INFO(data->info, 6, "finished successfully with %d iterations\n", iter);
    if (ws != data->ws) FREE_WORKSPACE(ws);

    fflush(stdout);
    return true;
  }

  fflush(stdout);
  return false;
}

#define SPCG_START_INFO(d, n, k, l)		\
  spcg_start_info(funcName, d, n, k, l)
#define SPCG_BREAK_INFO(d, s, i, n, k, l, r, o, w) \
 spcg_break_info(funcName, d, s, i, n, k, l, r, o, w)
#define SPCG_SOLVE_INFO(data, iter, nAuf, npj, npc, res2, ores, ws) \
 spcg_solve_info(funcName, data, iter, &nAuf, &npj, &npc, res2, ores, ws)

int oem_spcg(OEM_SP_DATA *data, int dimX, const REAL *f, REAL *u, int dimY, 
	     const REAL *g, REAL *p)
{
  FUNCNAME("oem_spcg");
  REAL   rho, Bdchi, delta, gamma;
  REAL   *d, *r, *Cr, *chi, *Bp, *Bd, *g_Btu;
  REAL   ores = -1.0;
  int    iter, nAuf = -1, npj = -1, npc = -1;

  int  (*Auf)(void *, int, const REAL *, REAL *) = data->solve_Auf;
  void *Auf_data = data->solve_Auf_data;
  void (*B)(void *, REAL, int, const REAL *, int, REAL *) = data->B;
  void *B_data = data->B_data;
  void (*Bt)(void *, REAL, int, const REAL *, int, REAL *) = data->Bt;
  void *Bt_data = data->Bt_data;
  int  (*project)(void *, int, const REAL *, REAL *) = data->project;
  void *project_data = data->project_data;
  int  (*precon)(void *, int, const REAL *, const REAL *, REAL *) =
    data->precon;
  void *precon_data = data->precon_data;
  int  info = data->info;

  const REAL TOL = 1.e-30;

  size_t size = precon?3*dimY+dimX+MAX(dimX, dimY):2*dimY+dimX+MAX(dimX, dimY);
  WORKSPACE *ws = CHECK_WORKSPACE(size, data->ws);

/*--------------------------------------------------------------------------*/
/*---  partitioning of the workspace  --------------------------------------*/
/*--------------------------------------------------------------------------*/

  d = (REAL *)ws->work;
  r = d + dimY;
  Cr = precon ? r+dimY : r;

  chi = Cr + dimY;

  Bp = Bd = chi + dimX;       /*---  vectors storing B p, B d and B^t u  ---*/
  g_Btu = project ? Bp : r;

  delta = g ? dnrm2(dimY, g, 1) : 0.0;
  if (delta + dnrm2(dimX, f, 1) < TOL) {
    INFO(info, 2, "f, g = 0 => u, p = 0 is the solution\n");
    dset(dimX, 0.0, u, 1);
    dset(dimY, 0.0, p, 1);
    data->initial_residual = data->residual = 0.0;
    return 0;
  }

/*--------------------------------------------------------------------------*/
/*--- Initialization  ------------------------------------------------------*/
/*--------------------------------------------------------------------------*/

/*--- Bd = f - B p  --------------------------------------------------------*/
  dcopy(dimX, f, 1, Bp, 1);
  (*B)(B_data, -1.0, dimY, p, dimX, Bp);

/*--- u = A^{-1} (f - Bp)  -------------------------------------------------*/
  nAuf = (*Auf)(Auf_data, dimX, Bp, u);

/*--- r = -B^*u + g  -------------------------------------------------------*/
  if (g) {
    dcopy(dimY, g, 1, g_Btu, 1);
  } else {
    dset(dimY, 0.0, g_Btu, 1);
  }
  (*Bt)(Bt_data, -1.0, dimX, u, dimY, g_Btu);
  if (project) {
    dset(dimY, 0.0, r, 1);
    npj = (*project)(project_data, dimY, g_Btu, r);
  }

/*--- preconditioning  -----------------------------------------------------*/
  if (precon) {
    dset(dimY, 0.0, Cr, 1);
    npc = (*precon)(precon_data, dimY, g_Btu, r, Cr);
  }

/*---  check initial residual  (r, r)_C = (r, Cr)  = (g-B^*u, Cr)  ---------*/
  delta = ddot(dimY, Cr, 1, g_Btu, 1);

  data->initial_residual = sqrt(delta);

  SPCG_START_INFO(data, nAuf, npj, npc);
  if (SPCG_SOLVE_INFO(data, 0, nAuf, npj, npc, delta, &ores, ws))
    return 0;

/*---  decent direction  ---------------------------------------------------*/
  dcopy(dimY, Cr, 1, d, 1);

/*--------------------------------------------------------------------------*/
/*---  Iteration  ----------------------------------------------------------*/
/*--------------------------------------------------------------------------*/

  for (iter = 1; iter <= data->max_iter; iter++)
  {
/*---  solve A chi = B d  --------------------------------------------------*/
    dset(dimX, 0.0, Bd, 1);
    (*B)(B_data, 1.0, dimY, d, dimX, Bd);

    dset(dimX, 0.0, chi, 1);
    nAuf += (*Auf)(Auf_data, dimX, Bd, chi);

/*---  decent parameter rho = delta/(h, d)    -------------------------------*/
    Bdchi = ddot(dimX, Bd, 1, chi, 1);
    if (ABS(Bdchi) < TOL) {
      SPCG_BREAK_INFO(data, "(Td, d) = 0", iter, nAuf, npj, npc, delta, 
		      &ores, ws);
      return iter;
    }

/*---  update p and u  -----------------------------------------------------*/
    rho = delta/Bdchi;
    daxpy(dimY, -rho, d, 1, p, 1);
    daxpy(dimX, rho, chi, 1, u, 1);

/*--- r = -B^*u + g  -------------------------------------------------------*/
    if (g)
      dcopy(dimY, g, 1, g_Btu, 1);
    else
      dset(dimY, 0.0, g_Btu, 1);
    (*Bt)(Bt_data, -1.0, dimX, u, dimY, g_Btu);
    if (project) {
      dset(dimY, 0.0, r, 1);
      npj += (*project)(project_data, dimY, g_Btu, r);
    }

/*--- preconditioning  -----------------------------------------------------*/
    if (precon) {
      dset(dimY, 0.0, Cr, 1);
      npc += (*precon)(precon_data, dimY, g_Btu, r, Cr);
    }

/*---  check initial residual  (r, r)_C = (r, Cr)  = (g-B^*u, Cr)  ----------*/
    gamma = 1.0/delta;
    delta = ddot(dimY, Cr, 1, g_Btu, 1);

    if (SPCG_SOLVE_INFO(data, iter, nAuf, npj, npc, delta, &ores, ws)) 
      return iter;

/*---  update of the decent direction  -------------------------------------*/
    gamma *= delta;
    dxpay(dimY, Cr, 1, gamma, d, 1);
  }
  return 0;    /*  statement never reached!!!                              */
}
