// ===========================================================================
//
//       Filename:  bicgstab.h
//
//    Description:
//
//        Version:  0.0
//        Created:  10/27/2009 03:15:06 PM
//       Revision:  none
//       Compiler:  g++ (c++)
//
//         Author:  Trevor Irons (ti)
//
//   Organisation:  Colorado School of Mines (CSM)
//                  United States Geological Survey (USGS)
//
//          Email:  tirons@mines.edu, tirons@usgs.gov
//
//  This program is free software: you can redistribute it and/or modify
//  it under the terms of the GNU General Public License as published by
//  the Free Software Foundation, either version 3 of the License, or
//  (at your option) any later version.
//
//  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.  See the
//  GNU General Public License for more details.
//
//  You should have received a copy of the GNU General Public License
//  along with this program.  If not, see <http://www.gnu.org/licenses/>.
//
// ===========================================================================
#pragma once

#ifndef  BICGSTAB_INC
#define  BICGSTAB_INC

#include <Eigen/Core>
#include <Eigen/Sparse>

#ifdef CHOLMODPRECONDITION
#include <Eigen/CholmodSupport>
#endif // CHOLMODPRECONDITION

//#include "unsupported/Eigen/IterativeSolvers"
//#include <unsupported/Eigen/SuperLUSupport>

#include <iostream>
#include <string>
#include <complex>
#include <fstream>
#include <LemmaCore>
#include "timer.h"

using namespace Eigen;
using namespace Lemma;

typedef Eigen::SparseMatrix<Complex> SparseMat;



// Computes implicit calculation
template < typename Solver >
inline VectorXcr implicitDCInvBPhi3 (
                        const SparseMat& D,
                        const Solver& solver,
                        const Ref<VectorXcr const> ioms,
                        const Ref<VectorXcr const> Phi,
                        Real& tol,   // not used
                        int& max_it  // not used
                ) {
    VectorXcr b = (ioms).asDiagonal() * (D.transpose()*Phi);
    VectorXcr y = solver.solve(b);
    //max_it = solver.iterations(); // actualy no need to pass this
    return  D*y;
    //return  y;
}


// // Simple extraction of indices in idx into reduceed array x1
// void vmap( const Ref<VectorXcr const> x0, Ref<VectorXcr> x1, const std::vector<int>& idx ) {
//     for (unsigned int ii=0; ii<idx.size(); ++ii) {
//         x1(ii) = x0(idx[ii]);
//     }
// }

// Simple extraction of indices in idx into reduceed array x1
VectorXcr vmap( const Ref<VectorXcr const> x0, const std::vector<int>& idx ) {
    VectorXcr x1 = VectorXcr::Zero( idx.size() );
    for (unsigned int ii=0; ii<idx.size(); ++ii) {
        x1(ii) = x0(idx[ii]);
    }
    return x1;
}

// reverse of above
void ivmap( Ref<VectorXcr > x0, const Ref<VectorXcr const> x1, const std::vector<int>& idx ) {
    for (unsigned int ii=0; ii<idx.size(); ++ii) {
        x0(idx[ii]) = x1(ii);
    }
}


// On Input
// A = Matrix
// B = Right hand side
// X = initial guess, and solution
// maxit = maximum Number of iterations
// tol = error tolerance
// On Output
// X real solution vector
// errorn = Real error norm
template < typename CSolver >
int implicitbicgstab(//const SparseMat& D,
                     //const SparseMat& C,
                     const Ref< Eigen::SparseMatrix<Complex> const > D,
                     const std::vector<int>& idx,
                     const Ref< VectorXcr const > ioms,
                     const Ref< VectorXcr const > rhs,
                     Ref <VectorXcr> phi,
                     CSolver& solver,
                     int &max_it, const Real &tol, Real &errorn, int &iter_done, ofstream& logio) {

    logio << "using the preconditioned implicit solver" << std::endl;

    Complex omega, rho, rho_1, alpha, beta;
    Real    tol2;
    int     iter, max_it2, max_it1;

    // Look at reduced problem
    VectorXcr rhs2 = vmap(rhs, idx);
    VectorXcr phi2 = vmap(phi, idx);

    // Determine size of system and init vectors
    int n = idx.size();        // was phi.size();

    std::cout << "BiCGStab SIZES  " << n << "\t" << phi.size() << "\t" << ioms.size() << std::endl;

    VectorXcr r(n);
    VectorXcr r_tld(n);
    VectorXcr p(n);
    VectorXcr s(n);

    VectorXcr v = VectorXcr::Zero(ioms.size());
    VectorXcr t = VectorXcr::Zero(ioms.size());

//    VectorXcr vm1 = VectorXcr::Zero(ioms.size());
//    VectorXcr tm1 = VectorXcr::Zero(ioms.size());

//     TODO, refigure for implicit large system
//     std::cout << "Start BiCGStab, memory needed: "
//               <<  (sizeof(Complex)*(9+2)*n/(1024.*1024*1024)) << " [Gb]\n";

    // Initialise
    iter_done = 0;
    Real eps = 1e-100;

    Real bnrm2 = rhs.norm();
    if (bnrm2 == 0) {
        phi.setConstant(0.0);
        errorn = 0;
        std::cerr << "Trivial case of Ax = b, where b is 0\n";
        return (0);
    }

    // If there is an initial guess
    if ( phi.norm() ) {
        tol2 = tol;
        max_it2 = 50000;
        //r = rhs - implicitDCInvBPhi3(D, solver, ioms, phi, tol2, max_it2);
        //r = rhs - implicitDCInvBPhi3(D, solver, ioms, phi, tol2, max_it2);
        r = rhs2 - vmap( implicitDCInvBPhi3(D, solver, ioms, phi, tol2, max_it2), idx );
    } else {
        r = vmap(rhs, idx);
    }

    jsw_timer timer;

    errorn = r.norm() / bnrm2;
    omega = 1.;
    r_tld = r;
    Real errornold = 1e14;
    // Get down to business
    for (iter=0; iter<max_it; ++iter) {

        timer.begin();

        rho = r_tld.dot(r);
        if (abs(rho) < eps) {
            ivmap( phi, phi2, idx );
            return (0);
        }

        if (iter > 0) {
            beta = (rho/rho_1) * (alpha/omega);
            p = r.array() + beta*(p.array()-omega*v.array()).array();
        } else {
            p = r;
        }

        tol2 = tol;

        max_it2 = 500000;
        ivmap(phi, p, idx);
        v = vmap(implicitDCInvBPhi3(D, solver, ioms, phi, tol2, max_it2), idx);

        alpha = rho / r_tld.dot(v);
        s = r.array() - alpha*v.array();
        errorn = s.norm()/bnrm2;

        if (errorn < tol && iter > 1) {
            phi2.array() += alpha*p.array();
            ivmap( phi, phi2, idx );
            return (0);
        }

        tol2 = tol;

        max_it1 = 500000;
        //t = implicitDCInvBPhi2(D, C, ioms, solver, s, tol2, max_it1);
        //t = implicitDCInvBPhi3(D, solver, ioms, s, tol2, max_it1);
        ivmap(phi, s, idx);
        t = vmap(implicitDCInvBPhi3(D, solver, ioms, phi, tol2, max_it1), idx);
        omega = t.dot(s)  / t.dot(t);

        r = s.array() - omega*t.array();
        errorn = r.norm() / bnrm2;
        iter_done = iter;

        if (errorn <= tol) {
            ivmap( phi, phi2, idx );
            return (0);
        }

        if (abs(omega) < eps) {
            ivmap( phi, phi2, idx );
            return (0);
        }

        rho_1 = rho;

        logio << "iteration " << std::setw(3) << iter
              << "  errorn " << std::setw(6) << std::setprecision(4) << std::scientific << errorn
              //<< "\timplicit iterations " << std::setw(5) << max_it1+max_it2
              << "  time " << std::setw(6) << std::fixed << std::setprecision(2) << timer.end()/60. << " [m]  "
              << max_it2+max_it2 << " iterations " << std::endl;

        // Check to see how progress is going
        if (errornold - errorn < 0) {
            logio << "Irregular non-monotonic (negative) convergence. Recommend restart. \n";
            ivmap( phi, phi2, idx );
            return (2);
        }

        /*
        if (errornold - errorn < 1e-14) {
            logio << "not making any progress. Giving up\n";
            return (1);
        }
        */

        //std::cout << "|| p-s ||" << (alpha*p - omega*s).norm() << std::endl;

        // only update phi if good things are happening
        phi2.array() += alpha*p.array() + omega*s.array();
        errornold = errorn;

    }
    ivmap( phi, phi2, idx );
    return (0);
}

#endif   // ----- #ifndef BICGSTAB_INC  -----