// ---------------------------------------------------------------------------
// - Krylov.cpp                                                              -
// - afnix:mth module - krylov based algorithm class implementation          -
// ---------------------------------------------------------------------------
// - This program is free software;  you can redistribute it  and/or  modify -
// - it provided that this copyright notice is kept intact.                  -
// -                                                                         -
// - 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.  In no event shall -
// - the copyright holder be liable for any  direct, indirect, incidental or -
// - special damages arising in any way out of the use of this software.     -
// ---------------------------------------------------------------------------
// - copyright (c) 1999-2013 amaury darsch                                   -
// ---------------------------------------------------------------------------

#include "Real.hpp"
#include "Math.hpp"
#include "Krylov.hpp"
#include "Algebra.hpp"
#include "Exception.hpp"
#include "Utility.hpp"
#include "OutputTerm.hpp"

namespace afnix {
  static OutputTerm terr (OutputTerm::ERROR);
  // -------------------------------------------------------------------------
  // - public section                                                        -
  // -------------------------------------------------------------------------

  // solve a system with conjugate gradient squared method

  bool Krylov::cgs (Rvi& x, const Rmi& a, const Rvi& b, const long ni) {
    // check for square matrix and vector consistency
    t_long size = x.getsize ();
    if ((a.getrsiz () != size) || (a.getcsiz () != size) || 
	(b.getsize () != size)) {
      throw Exception ("krylov-error", 
		       "incompatible matrix/vector in cgs solver");
    }
    // compute operating norms
    t_real an = a.norm ();
    t_real bn = b.norm ();
    // check for null solution
    if (bn == 0.0) {
      x.clear ();
      return true;
    }
    // check for null matrix
    if (an == 0.0) {
      throw Exception ("krylov-error", "invalid null matrix in cgs solver");
    }
    // operating vector list
    Rvi* ri = nilp;
    Rvi* rn = nilp;
    Rvi*  q = nilp;
    Rvi*  p = nilp;
    Rvi*  u = nilp;
    Rvi* qh = nilp;
    Rvi* uh = nilp;
    Rvi* vh = nilp;
    // protected block
    try {
      // clear target and clone
      x.clear ();
      ri = dynamic_cast <Rvi*> (x.clone ());
      rn = dynamic_cast <Rvi*> (x.clone ());
      q  = dynamic_cast <Rvi*> (x.clone ());
      p  = dynamic_cast <Rvi*> (x.clone ());
      u  = dynamic_cast <Rvi*> (x.clone ());
      qh = dynamic_cast <Rvi*> (x.clone ());
      uh = dynamic_cast <Rvi*> (x.clone ());
      vh = dynamic_cast <Rvi*> (x.clone ());
      // set the initial vector
      x.set (Math::d_aeps);
      // compute ri = b - Ax
      a.mul (*ri, x, -1.0); (*ri) += b;
      // check convergence - the initial x might be the solution
      if (ri->norm () < Math::d_reps * (an * x.norm () + bn)) {
	delete ri; delete rn;
	delete  q; delete  p; delete  u; 
	delete uh; delete qh; delete vh;
	return true;
      }
      // copy ri into rn
      rn->cpy (*ri);
      // initialize factors
      t_real rhop = 1.0;
      // initialize status
      bool status = false;
      // main loop
      for (long i = 0; i < ni; i++) {
	// compute next rho
	t_real rhon = (*ri) ^ (*rn);
	if (rhon == 0.0) {
	  status = false;
	  break;
	}
	if (i == 0) {
	  u->cpy (*rn);
	  p->cpy (*rn);
	} else {
	  // compute beta
	  t_real beta = rhon / rhop;
	  // compute u = rn + beta*q
	  u->add (*rn, *q, beta);
	  // compute p = u + beta (q + beta*p)
	  p->req (*q, beta); p->req (*u, beta);
	}
	// compute vh = A*p
	a.mul (*vh, *p, 1.0);
	// compute alfa = rhon / <ri, v>
	t_real alfa = rhon / ((*ri) ^ (*vh));
	if (Math::isinf (alfa) == true) {
	  status = false;
	  break;
	}
	// compute q = u -alfa*vh
	q->add (*u, *vh, -alfa);
	// solve uh = u + q
	uh->add (*u, *q);
	// compute x = x + alfa*uh
	x.aeq (*uh, alfa);
	// compute qh = A*uh
	a.mul (*qh, *uh, 1.0);
	// compute rn = rn - alfa*qh
	rn->aeq (*qh, -alfa);
	// set previous rho
	rhop = rhon;
	// check convergence
	t_real nr = rn->norm ();
	if ((nr < Math::d_aeps) && 
	    (rn->norm () < Math::d_reps * (an * x.norm () + bn))) {
	  status = true;
	  break;
	}
      }
      // clean and return
      delete ri; delete rn;
      delete  q; delete  p; delete u; delete uh; delete qh; delete vh;
      return status;
    } catch (...) {
      delete ri; delete rn;
      delete  q; delete  p; delete u; delete uh; delete qh; delete vh;
      throw;
    }
  }

  // solve a system with preconditioned conjugate gradient squared method

  bool Krylov::cgs (Rvi& x, const Rmi& a, const Rvi& m, const Rvi& b,
		    const long ni) {
    // check for square matrix and vector consistency
    t_long size = x.getsize ();
    if ((a.getrsiz () != size) || (a.getcsiz () != size) || 
	(b.getsize () != size)) {
      throw Exception ("krylov-error", 
		       "incompatible matrix/vector in cgs solver");
    }
    // compute operating norms
    t_real an = a.norm ();
    t_real bn = b.norm ();
    // check for null solution
    if (bn == 0.0) {
      x.clear ();
      return true;
    }
    // check for null matrix
    if (an == 0.0) {
      throw Exception ("krylov-error", "invalid null matrix in cgs solver");
    }
    // operating vector list
    Rvi* ri = nilp;
    Rvi* rn = nilp;
    Rvi*  q = nilp;
    Rvi*  p = nilp;
    Rvi*  u = nilp;
    Rvi* ph = nilp;
    Rvi* qh = nilp;
    Rvi* uh = nilp;
    Rvi* vh = nilp;
    // protected block
    try {
      // clear target and clone
      x.clear ();
      ri = dynamic_cast <Rvi*> (x.clone ());
      rn = dynamic_cast <Rvi*> (x.clone ());
      q  = dynamic_cast <Rvi*> (x.clone ());
      p  = dynamic_cast <Rvi*> (x.clone ());
      u  = dynamic_cast <Rvi*> (x.clone ());
      ph = dynamic_cast <Rvi*> (x.clone ());
      qh = dynamic_cast <Rvi*> (x.clone ());
      uh = dynamic_cast <Rvi*> (x.clone ());
      vh = dynamic_cast <Rvi*> (x.clone ());
      // set the initial vector
      x.set (Math::d_aeps);
      // compute ri = b - Ax
      a.mul (*ri, x, -1.0); (*ri) += b;
      // check convergence - the initial x might be the solution
      if (ri->norm () < Math::d_reps * (an * x.norm () + bn)) {
	delete ri; delete rn;
	delete  q; delete  p; delete u; 
	delete ph; delete qh; delete uh; delete vh;
	return true;
      }
      // copy ri into rn
      rn->cpy (*ri);
      // initialize factors
      t_real rhop = 1.0;
      // initialize status
      bool status = false;
      // main loop
      for (long i = 0; i < ni; i++) {
	// compute next rho
	t_real rhon = (*ri) ^ (*rn);
	if (rhon == 0.0) {
	  status = false;
	  break;
	}
	if (i == 0) {
	  u->cpy (*rn);
	  p->cpy (*rn);
	} else {
	  // compute beta
	  t_real beta = rhon / rhop;
	  // compute u = rn + beta*q
	  u->add (*rn, *q, beta);
	  // compute p = u + beta (q + beta*p)
	  p->req (*q, beta); p->req (*u, beta);
	}
	// solve M*ph = p
	ph->mul (m, *p);
	// compute vh = A*ph
	a.mul (*vh, *ph, 1.0);
	// compute alfa = rhon / <ri, v>
	t_real alfa = rhon / ((*ri) ^ (*vh));
	if (Math::isinf (alfa) == true) {
	  status = false;
	  break;
	}
	// compute q = u -alfa*vh
	q->add (*u, *vh, -alfa);
	// solve M*uh = u + q
	uh->add (*u, *q);
	uh->mul (m, *uh);
	// compute x = x + alfa*uh
	x.aeq (*uh, alfa);
	// compute qh = A*uh
	a.mul (*qh, *uh, 1.0);
	// compute rn = rn - alfa*qh
	rn->aeq (*qh, -alfa);
	// set previous rho
	rhop = rhon;
	// check convergence
	t_real nr = rn->norm ();
	if ((nr < Math::d_aeps) && 
	    (rn->norm () < Math::d_reps * (an * x.norm () + bn))) {
	  status = true;
	  break;
	}
      }
      // clean and return
      delete ri; delete rn;
      delete  q; delete  p; delete u; 
      delete ph; delete uh; delete qh; delete vh;
      return status;
    } catch (...) {
      delete ri; delete rn;
      delete  q; delete  p; delete u; 
      delete ph; delete uh; delete qh; delete vh;
      throw;
    }
  }

  // solve a system with bi-conjugate stabilized method

  bool Krylov::bcs (Rvi& x, const Rmi& a, const Rvi& b, const long ni) {
    // check for square matrix and vector consistency
    t_long size = x.getsize ();
    if ((a.getrsiz () != size) || (a.getcsiz () != size) || 
	(b.getsize () != size)) {
      throw Exception ("krylov-error", 
		       "incompatible matrix/vector in bcs solver");
    }
    // compute operating norms
    t_real an = a.norm ();
    t_real bn = b.norm ();
    // check for null solution
    if (bn == 0.0) {
      x.clear ();
      return true;
    }
    // check for null matrix
    if (an == 0.0) {
      throw Exception ("krylov-error", "invalid null matrix in bcs solver");
    }
    // operating vector list
    Rvi* ri = nilp;
    Rvi* rn = nilp;
    Rvi*  p = nilp;
    Rvi*  v = nilp;
    Rvi*  s = nilp;
    Rvi*  t = nilp;
    // protected block
    try {
      // clear target and clone
      x.clear ();
      ri = dynamic_cast <Rvi*> (x.clone ());
      rn = dynamic_cast <Rvi*> (x.clone ());
      p  = dynamic_cast <Rvi*> (x.clone ());
      v  = dynamic_cast <Rvi*> (x.clone ());
      s  = dynamic_cast <Rvi*> (x.clone ());
      t  = dynamic_cast <Rvi*> (x.clone ());
      // set the initial vector
      x.set (Math::d_aeps);      
      // compute ri = b - Ax
      a.mul (*ri, x, -1.0); (*ri) += b;
      // check convergence - the initial x might be the solution
      if (ri->norm () < Math::d_reps * (an * x.norm () + bn)) {
	delete ri; delete rn;
	delete  p; delete  v; delete s; delete t;
	return true;
      }
      // copy ri into rn
      rn->cpy (*ri);
      // initialize factors
      t_real rhop = 1.0;
      t_real alfa = 1.0;
      t_real omga = 1.0;
      // initialize status
      bool status = false;
      // main loop
      for (long i = 0; i < ni; i++) {
	// compute next rho
	t_real rhon = (*ri) ^ (*rn);
	if (rhon == 0.0) {
	  status = false;
	  break;
	}
	if (i == 0) {
	  p->cpy (*rn);
	} else {
	  // compute beta
	  t_real beta = (rhon / rhop) * (alfa / omga);
	  // compute next p = rn + beta (p - omga*v)
	  p->aeq (*v, -omga); p->req (*rn, beta);
	}
	// compute v = A*p
	a.mul (*v, *p, 1.0);
	// compute alfa = rhon / <ri, v>
	alfa = rhon / ((*ri) ^ (*v));
	if (Math::isinf (alfa) == true) {
	  status = false;
	  break;
	}
	// compute s = rn - alfa*v
	s->add (*rn, *v, -alfa);
	// compute t = A*s
	a.mul (*t, *s, 1.0);
	// compute omga = <t, s> / <t, t>
	omga = ((*t) ^ (*s)) / ((*t) ^ (*t));
	// compute x = x + alfa*p + omga*s
	x.aeq (*p, alfa); x.aeq (*s, omga);
	// compute rn = s - omga*t
	rn->add (*s, *t, -omga);
	// set previous rho
	rhop = rhon;
	// check convergence
	t_real nr = rn->norm ();
	if ((nr < Math::d_aeps) && 
	    (rn->norm () < Math::d_reps * (an * x.norm () + bn))) {
	  status = true;
	  break;
	}
	// check omga
	if ((omga == 0.0) || (Math::isinf (omga) == true)){
	  status = false;
	  break;
	}
      }
      // clean and return
      delete ri; delete rn;
      delete  p; delete  v; delete s; delete t;
      return status;
    } catch (...) {
      delete ri; delete rn;
      delete  p; delete  v; delete s; delete t;
      throw;
    }
  }

  // solve a system with a precontioned bi-conjugate stabilized method

  bool Krylov::bcs (Rvi& x, const Rmi& a, const Rvi& m, const Rvi& b,
		    const long ni) {
    // check for square matrix and vector consistency
    t_long size = x.getsize ();
    if ((a.getrsiz () != size) || (a.getcsiz () != size) || 
	(b.getsize () != size)) {
      throw Exception ("krylov-error", 
		       "incompatible matrix/vector in bcs solver");
    }
    // compute operating norms
    t_real an = a.norm ();
    t_real bn = b.norm ();
    // check for null solution
    if (bn == 0.0) {
      x.clear ();
      return true;
    }
    // check for null matrix
    if (an == 0.0) {
      throw Exception ("krylov-error", "invalid null matrix in bcs solver");
    }
    // operating vector list
    Rvi* ri = nilp;
    Rvi* rn = nilp;
    Rvi*  p = nilp;
    Rvi*  v = nilp;
    Rvi*  s = nilp;
    Rvi*  t = nilp;
    Rvi* ph = nilp;
    Rvi* sh = nilp;
    // protected block
    try {
      // clear target and clone
      x.clear ();
      ri = dynamic_cast <Rvi*> (x.clone ());
      rn = dynamic_cast <Rvi*> (x.clone ());
      p  = dynamic_cast <Rvi*> (x.clone ());
      v  = dynamic_cast <Rvi*> (x.clone ());
      s  = dynamic_cast <Rvi*> (x.clone ());
      t  = dynamic_cast <Rvi*> (x.clone ());
      ph = dynamic_cast <Rvi*> (x.clone ());
      sh = dynamic_cast <Rvi*> (x.clone ());
      // set the initial vector
      x.set (Math::d_aeps);      
      // compute ri = b - Ax
      a.mul (*ri, x, -1.0); (*ri) += b;
      // check convergence - the initial x might be the solution
      if (ri->norm () < Math::d_reps * (an * x.norm () + bn)) {
	delete ri; delete rn;
	delete  p; delete  v; delete s; delete t;
	delete ph; delete sh;
	return true;
      }
      // copy ri into rn
      rn->cpy (*ri);
      // initialize factors
      t_real rhop = 1.0;
      t_real alfa = 1.0;
      t_real omga = 1.0;
      // initialize status
      bool status = false;
      // main loop
      for (long i = 0; i < ni; i++) {
	// compute next rho
	t_real rhon = (*ri) ^ (*rn);
	if (rhon == 0.0) {
	  status = false;
	  break;
	}
	if (i == 0) {
	  p->cpy (*rn);
	} else {
	  // compute beta
	  t_real beta = (rhon / rhop) * (alfa / omga);
	  // compute next p = rn + beta (p - omga*v)
	  p->aeq (*v, -omga); p->req (*rn, beta);
	}
	// solve M*ph = p
	ph->mul (m, *p);
	// compute v = A*ph
	a.mul (*v, *ph, 1.0);
	// compute alfa = rhon / <ri, v>
	alfa = rhon / ((*ri) ^ (*v));
	if (Math::isinf (alfa) == true) {
	  status = false;
	  break;
	}
	// compute s = rn - alfa*v
	s->add (*rn, *v, -alfa);
	// solve M*sh = s
	sh->mul (m, *s);
	// compute t = A*sh
	a.mul (*t, *sh, 1.0);
	// compute omga = <t, s> / <t, t>
	omga = ((*t) ^ (*s)) / ((*t) ^ (*t));
	// compute x = x + alfa*ph + omga*sh
	x.aeq (*ph, alfa); x.aeq (*sh, omga);
	// compute rn = s - omga*t
	rn->add (*s, *t, -omga);
	// set previous rho
	rhop = rhon;
	// check convergence
	t_real nr = rn->norm ();
	if ((nr < Math::d_aeps) && 
	    (rn->norm () < Math::d_reps * (an * x.norm () + bn))) {
	  status = true;
	  break;
	}
	// check omga
	if ((omga == 0.0) || (Math::isinf (omga) == true)){
	  status = false;
	  break;
	}
      }
      // clean and return
      delete ri; delete rn;
      delete  p; delete  v; delete s; delete t;
      delete ph; delete sh;
      return status;
    } catch (...) {
      delete ri; delete rn;
      delete  p; delete  v; delete s; delete t;
      delete ph; delete sh;
      throw;
    }
  }
}
