/********************************************************
 *  ██████╗  ██████╗████████╗██╗
 * ██╔════╝ ██╔════╝╚══██╔══╝██║
 * ██║  ███╗██║        ██║   ██║
 * ██║   ██║██║        ██║   ██║
 * ╚██████╔╝╚██████╗   ██║   ███████╗
 *  ╚═════╝  ╚═════╝   ╚═╝   ╚══════╝
 * Geophysical Computational Tools & Library (GCTL)
 *
 * Copyright (c) 2022  Yi Zhang (yizhang-geo@zju.edu.cn)
 *
 * GCTL is distributed under a dual licensing scheme. You can redistribute 
 * it and/or modify it under the terms of the GNU Lesser General Public 
 * License as published by the Free Software Foundation, either version 2 
 * of the License, or (at your option) any later version. You should have 
 * received a copy of the GNU Lesser General Public License along with this 
 * program. If not, see <http://www.gnu.org/licenses/>.
 * 
 * If the terms and conditions of the LGPL v.2. would prevent you from using 
 * the GCTL, please consider the option to obtain a commercial license for a 
 * fee. These licenses are offered by the GCTL's original author. As a rule, 
 * licenses are provided "as-is", unlimited in time for a one time fee. Please 
 * send corresponding requests to: yizhang-geo@zju.edu.cn. Please do not forget 
 * to include some description of your company and the realm of its activities. 
 * Also add information on how to contact you by electronic and paper mail.
 ******************************************************/

#include "clcg.h"

/**
 * Default parameter for conjugate gradient methods
 */
static const gctl::clcg_para clcg_defparam = {0, 1e-8, 0};

int gctl::clcg_solver::CLCG_Progress(const array<std::complex<double> > &m, const double converge, const clcg_para &param, size_t t)
{
    if (converge <= param.epsilon)
    {
        std::clog << GCTL_CLEARLINE << "\rIteration-times: " << t << "\tconvergence: " << converge;
        return 0;
    }

    if (clcg_inter_ > 0 && t%clcg_inter_ == 0)
    {
        std::clog << GCTL_CLEARLINE << "\rIteration-times: " << t << "\tconvergence: " << converge;
    }
    return 0;
}

gctl::clcg_solver::clcg_solver()
{
    clcg_param_ = clcg_defparam;
    clcg_inter_ = 1;
    clcg_silent_ = false;
}

gctl::clcg_solver::~clcg_solver(){}

void gctl::clcg_solver::clcg_silent()
{
    clcg_silent_ = true;
    return;
}

void gctl::clcg_solver::set_clcg_report_interval(size_t inter)
{
    clcg_inter_ = inter;
    return;
}

void gctl::clcg_solver::set_clcg_para(const clcg_para &in_param)
{
    clcg_param_ = in_param;
    return;
}

void gctl::clcg_solver::set_clcg_para(const toml::value &toml_data)
{
    clcg_param_ = clcg_defparam;

    std::string CLCG = "clcg";
	if (toml_data.contains(CLCG))
	{
		if (toml_data.at(CLCG).contains("max_iterations")) clcg_param_.max_iterations = toml::find<int>(toml_data, CLCG, "max_iterations");
		if (toml_data.at(CLCG).contains("epsilon"))        clcg_param_.epsilon = toml::find<double>(toml_data, CLCG, "epsilon");
        if (toml_data.at(CLCG).contains("abs_diff"))       clcg_param_.abs_diff = toml::find<int>(toml_data, CLCG, "abs_diff");
	}
    return;
}

void gctl::clcg_solver::clcg_error_str(clcg_return_code err_code, std::ostream &ss, bool err_throw)
{
#if defined _WINDOWS || __WIN32__
    if (!er_throw)
    {
        if (err_code >= 0)
        {
            SetConsoleTextAttribute(GetStdHandle(STD_ERROR_HANDLE), FOREGROUND_INTENSITY | FOREGROUND_GREEN);
            ss << "Success! ";
        }
        else
        {
            SetConsoleTextAttribute(GetStdHandle(STD_ERROR_HANDLE), FOREGROUND_INTENSITY | FOREGROUND_RED);
            ss << "Fail! ";
        }
    }
#else
    if (!err_throw)
    {
        if (err_code >= 0)
            ss << "\033[1m\033[32mCLCG Success! ";
        else
            ss << "\033[1m\033[31mCLCG Fail! ";
    }
#endif

    std::string err_str;
    switch (err_code)
	{
		case CLCG_SUCCESS:
			err_str = "Iteration reached convergence."; break;
		case CLCG_STOP:
			err_str = "Iteration is stopped by the progress evaluation function."; break;
		case CLCG_ALREADY_OPTIMIZIED:
			err_str = "The variables are already optimized."; break;
		case CLCG_UNKNOWN_ERROR:
			err_str = "Unknown error."; break;
		case CLCG_INVILAD_VARIABLE_SIZE:
			err_str = "The size of the variables is negative."; break;
		case CLCG_INVILAD_MAX_ITERATIONS:
			err_str = "The maximal iteration times is negative."; break;
		case CLCG_INVILAD_EPSILON:
			err_str = "The epsilon is not in the range (0, 1)."; break;
		case CLCG_REACHED_MAX_ITERATIONS:
			err_str = "The maximal iteration has been reached."; break;
		case CLCG_NAN_VALUE:
			err_str = "The model values are NaN."; break;
		case CLCG_INVALID_POINTER:
			err_str = "Invalid pointer."; break;
		case CLCG_SIZE_NOT_MATCH:
			err_str = "The sizes of the solution and target do not match."; break;
		case CLCG_UNKNOWN_SOLVER:
			err_str = "Unknown solver."; break;
		default:
			err_str = "Unknown error."; break;
	}

    if (err_throw && err_code < 0) throw std::runtime_error(err_str.c_str());
    else ss << err_str;

#if defined _WINDOWS || __WIN32__
    if (!er_throw)
    {
        if (err_code >= 0)
        {
            SetConsoleTextAttribute(GetStdHandle(STD_ERROR_HANDLE), 7);
            ss << std::endl;
        }
        else
        {
            SetConsoleTextAttribute(GetStdHandle(STD_ERROR_HANDLE), 7);
            ss << std::endl;
        }	
    }
#else
    if (!err_throw)
    {
        if (err_code >= 0)
            ss << "\033[0m" << std::endl;
        else
            ss << "\033[0m" << std::endl;	
    }
#endif

    return;
}

gctl::clcg_para gctl::clcg_solver::default_clcg_para()
{
    clcg_para dp = clcg_defparam;
    return dp;
}

void gctl::clcg_solver::CLCG_Minimize(array<std::complex<double> > &m, const array<std::complex<double> > &B, 
    clcg_solver_type solver_id, std::ostream &ss, bool verbose, bool er_throw)
{
	if (clcg_silent_)
    {
        clcg_return_code ret;
        if (solver_id == CLCG_BICG) ret = clbicg(m, B);
        else if (solver_id == CLCG_BICG_SYM) ret = clbicg_symmetric(m, B);
        else if (solver_id == CLCG_CGS) ret = clcgs(m, B);
        else if (solver_id == CLCG_BICGSTAB) ret = clbicgstab(m, B);
        else if (solver_id == CLCG_TFQMR) ret = cltfqmr(m, B);
        else throw std::invalid_argument("Invalid solver type. gctl::clcg_solver<T>::Minimize(...)");

        if (ret < 0) clcg_error_str(ret, ss, true);
        return;
    }

#ifdef GCTL_OPENMP
    double start = omp_get_wtime();
    clcg_return_code ret;
    if (solver_id == CLCG_BICG) ret = clbicg(m, B);
    else if (solver_id == CLCG_BICG_SYM) ret = clbicg_symmetric(m, B);
    else if (solver_id == CLCG_CGS) ret = clcgs(m, B);
    else if (solver_id == CLCG_BICGSTAB) ret = clbicgstab(m, B);
    else if (solver_id == CLCG_TFQMR) ret = cltfqmr(m, B);
    else throw std::invalid_argument("Invalid solver type. gctl::clcg_solver<T>::Minimize(...)");
    double end = omp_get_wtime();

    double costime = 1000*(end-start);
#else
    clock_t start = clock();
    clcg_return_code ret;
    if (solver_id == CLCG_BICG) ret = clbicg(m, B);
    else if (solver_id == CLCG_BICG_SYM) ret = clbicg_symmetric(m, B);
    else if (solver_id == CLCG_CGS) ret = clcgs(m, B);
    else if (solver_id == CLCG_BICGSTAB) ret = clbicgstab(m, B);
    else if (solver_id == CLCG_TFQMR) ret = cltfqmr(m, B);
    else throw std::invalid_argument("Invalid solver type. gctl::clcg_solver<T>::Minimize(...)");
    clock_t end = clock();

    double costime = 1000*(end-start)/(double)CLOCKS_PER_SEC;
#endif
    
    if (!er_throw)
    {
        ss << std::endl;
        switch (solver_id)
		{
			case CLCG_BICG:
				std::clog << "Solver: Bi-CG. Times cost: " << costime << " ms" << std::endl;
				break;
			case CLCG_BICG_SYM:
				std::clog << "Solver: Bi-CG (symmetrically accelerated). Times cost: " << costime << " ms" << std::endl;
				break;
			case CLCG_CGS:
				std::clog << "Solver: CGS. Times cost: " << costime << " ms" << std::endl;
				break;
            case CLCG_BICGSTAB:
				std::clog << "Solver: CGS. Times cost: " << costime << " ms" << std::endl;
				break;
			case CLCG_TFQMR:
				std::clog << "Solver: TFQMR. Times cost: " << costime << " ms" << std::endl;
				break;
			default:
				std::clog << "Solver: Unknown. Times cost: " << costime << " ms" << std::endl;
				break;
		}
    }

    if (verbose) clcg_error_str(ret, ss, er_throw);
    else if (ret < 0) clcg_error_str(ret, ss, er_throw);
    return;
}

gctl::clcg_return_code gctl::clcg_solver::clbicg(array<std::complex<double> > &m, const array<std::complex<double> > &B)
{
    clcg_return_code ret;
    return ret;
}

gctl::clcg_return_code gctl::clcg_solver::clbicg_symmetric(array<std::complex<double> > &m, const array<std::complex<double> > &B)
{
	size_t n_size = B.size();
	//check parameters
	if (n_size <= 0) return CLCG_INVILAD_VARIABLE_SIZE;
	if (clcg_param_.max_iterations < 0) return CLCG_INVILAD_MAX_ITERATIONS;
	if (clcg_param_.epsilon <= 0.0 || clcg_param_.epsilon >= 1.0) return CLCG_INVILAD_EPSILON;

	r1k.resize(n_size);
    d1k.resize(n_size);
	Ax.resize(n_size);

	CLCG_Ax(m, Ax, gctl::NoTrans, gctl::NoConj);

    std::complex<double> one_z(1.0, 0.0);
    vecdiff(r1k, B, Ax, one_z, one_z);
    veccpy(d1k, r1k, one_z);

	std::complex<double> rkrk = vecdot(r1k, r1k);

	double r0_square, rk_square;
	std::complex<double> r0_mod, rk_mod;
	rk_mod = vecinner(r1k, r1k);
	r0_square = rk_square = std::norm(rk_mod);

	if (r0_square < 1.0) r0_square = 1.0;

	clcg_return_code ret;
	if (clcg_param_.abs_diff && sqrt(rk_square)/n_size <= clcg_param_.epsilon)
	{
		ret = CLCG_ALREADY_OPTIMIZIED;
        CLCG_Progress(m, sqrt(rk_square)/n_size, clcg_param_, 0);
		return ret;
	}	
	else if (rk_square/r0_square <= clcg_param_.epsilon)
	{
		ret = CLCG_ALREADY_OPTIMIZIED;
        CLCG_Progress(m, rk_square/r0_square, clcg_param_, 0);
		return ret;
	}

	double residual;
    std::complex<double> ak, rkrk2, betak, dkAx;

    size_t t = 0;
	while(1)
	{
		if (clcg_param_.abs_diff) residual = sqrt(rk_square)/n_size;
		else residual = rk_square/r0_square;

        if (CLCG_Progress(m, residual, clcg_param_, t))
        {
            ret = CLCG_STOP; return ret;
        }

		if (residual <= clcg_param_.epsilon)
		{
			ret = CLCG_CONVERGENCE; return ret;
		}

		if (clcg_param_.max_iterations > 0 && t+1 > clcg_param_.max_iterations)
		{
			ret = CLCG_REACHED_MAX_ITERATIONS;
			break;
		}
		
		t++;

		CLCG_Ax(d1k, Ax, gctl::NoTrans, gctl::NoConj);
		dkAx = vecdot(d1k, Ax);
		ak = rkrk/dkAx;

        vecapp(m, d1k, ak);
        vecsub(r1k, Ax, ak);

		rk_mod = vecdot(r1k, r1k);
		rk_square = std::norm(rk_mod);

        if (!vecvalid(m))
        {
            ret = CLCG_NAN_VALUE; return ret;
        }

		rkrk2 = vecdot(r1k, r1k);
		betak = rkrk2/rkrk;
		rkrk = rkrk2;

        vecadd(d1k, d1k, r1k, betak, one_z);
	}

	return ret;
}

gctl::clcg_return_code gctl::clcg_solver::clcgs(array<std::complex<double> > &m, const array<std::complex<double> > &B)
{
	clcg_return_code ret;
    return ret;
}

gctl::clcg_return_code gctl::clcg_solver::clbicgstab(array<std::complex<double> > &m, const array<std::complex<double> > &B)
{
	clcg_return_code ret;
    return ret;
}

gctl::clcg_return_code gctl::clcg_solver::cltfqmr(array<std::complex<double> > &m, const array<std::complex<double> > &B)
{
	clcg_return_code ret;
    return ret;
}