#include "odeimpeul.h"
#include <cassert>

// Same calling conventions as VL
static double BiCGSTAB(const SparseMatd &A, Vecd &x, const Vecd &b,
                double eps, int *steps = 0);

void OdeImpEul1::step(Vecd &xout, const OdeFunc1 &f, const Vecd &x, double h) const
{
    int n = x.Elts();
    
    SparseMatd M(n, n);
    M.MakeDiag(1/h);
    M -= f.dx(x);
    
    Vecd b  = f(x);

    Vecd dx = h * b;

    int i=100;
    BiCGSTAB(M, dx, b, 1e-4, &i);
    // double resid = BiCGSTAB(M, dx, b, 1e-4, &i);
  //  double resid = SolveConjGrad(M, dx, b, 1e-8, &i);
    //cout << "iters: " << i << "residual: " << ret << endl;
 
    xout = x + dx;
}

void OdeImpEul2::step(Vecd &xout, Vecd &vout, const OdeFunc2 &f, const Vecd &x, const Vecd &v, double h) const
{
    int n = v.Elts();

    SparseMatd M(n,n);
    M.MakeDiag(vl_one);
    M -= (f.dv(x,v) * h);
    SparseMatd fdx = f.dx(x,v);
    M -= (fdx*h*h);

    Vecd fxv = f(x,v);

    Vecd b = h*( fxv + h*(fdx*v) );

    Vecd dv(h*fxv);

    int i=100;
    BiCGSTAB(M, dv, b, 1e-4, &i);
   // double resid = SolveConjGrad(M, dv, b, 1e-8, &i);
    //cout << "iters: " << i << "residual: " << resid << endl;

    vout = v + dv;
    xout = x + h*vout;
}

static double BiCGSTAB(const SparseMatd &A, Vecd &x, const Vecd &b,
                double eps, int *steps)
{
    int max_iter = (steps) ? *steps : 1000; // new
    int n = x.Elts();
    double resid;
    // initialize to avoid compiler warnings
    double rho_1=1.0, rho_2=1.0, alpha=1.0, beta=1.0, omega=1.0;
    Vecd p(n), phat(n), s(n), shat(n), t(n), v(n);
    
    double normb = len(b);
    Vecd r = b - A * x;
    Vecd rtilde = r;
    

    // Eugene: potential floating point roundoff error gives NaN!
    if (normb < 1e-154)
    {
        normb = 1;
    }
    
    if ((resid = len(r) / normb) <= eps) {
        eps = resid;
        if (steps) *steps = 0;
        return eps;
    }
    
    for (int i = 1; i <= max_iter; i++) {
       // cout << "RLEN: " << len(r) << endl;
    
        rho_1 = dot(rtilde, r);
        if (rho_1 == 0) {
            assert(!"Something is wrong in BiCGSTAB");
        }
        if (i == 1)
            p = r;
        else {
            beta = (rho_1/rho_2) * (alpha/omega);
            p = r + beta * (p - omega * v);
        }
        phat = p; 
        // phat should be M\p where M is preconditioner,
        // but I deleted it.
        v = A * phat;
        alpha = rho_1 / dot(rtilde, v);
        s = r - alpha * v;
        if ((resid = len(s)/normb) < eps) {
            x += alpha * phat;
            eps = resid;
            if (steps) *steps = i;
            return eps;
        }
        shat = s;
        // shat should be M\s where M is preconditioner,
        // but I deleted it.
        t = A * shat;
        omega = dot(t,s) / dot(t,t);
        x += alpha * phat + omega * shat;

        Vecd tmp = s - omega * t;
        if (resid < len(tmp)/normb)
        {
           // cout << "RESID: " << resid << endl;
            eps = resid;
            if (steps) *steps = i;
            return eps;
        }
        else
        {
         //   cout << "RESID: " << resid << endl;
            r = tmp;
        }


       // r = b - A*x;
        
        rho_2 = rho_1;
        if ((resid = len(r) / normb) < eps) {
            eps = resid;
            if (steps) *steps = i;
            return eps;
        }
        if (omega == 0) {
            assert(!"Something is wrong in BiCGSTAB");
        }
    }
    
    eps = resid;
    return eps;
}




