/********************************************************
 *  ██████╗  ██████╗████████╗██╗
 * ██╔════╝ ██╔════╝╚══██╔══╝██║
 * ██║  ███╗██║        ██║   ██║
 * ██║   ██║██║        ██║   ██║
 * ╚██████╔╝╚██████╗   ██║   ███████╗
 *  ╚═════╝  ╚═════╝   ╚═╝   ╚══════╝
 * 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 "gctl/core.h"
#include "gctl/algorithm.h"
#include "../lib/optimization.h"

#define M 1000
#define N 900

// get random floating points
double random_double(double l, double t)
{
	return (t-l)*rand()*1.0/RAND_MAX + l;
}

// get random integral numbers
int random_int(int small, int big)
{
	return (rand() % (big - small)) + small;
}

double max_diff(const gctl::_1d_array &a, const gctl::_1d_array &b)
{
	double max = -1.0;
	for (size_t i = 0; i < a.size(); i++)
	{
		max = std::max(fabs(a[i] - b[i]), max);
	}
	return max;
}

class ex8 : public gctl::lbfgs_solver, public gctl::grad_norm
{
public:
    ex8();
    virtual ~ex8();
    virtual double LBFGS_Evaluate(const gctl::_1d_array &x, gctl::_1d_array &g);
    virtual int LBFGS_Progress(const gctl::_1d_array &x, const gctl::_1d_array &g, const double fx, 
        const double converge, const double rate, const gctl::lbfgs_para param, int k, int ls, std::ostream &ss);

    void CalTarget(const gctl::_1d_array &x);

private:
    gctl::_1d_array obs1, obs2, obs3, tmp, grad;
    gctl::_2d_matrix k1, k2, k3;
};

ex8::ex8()
{
    srand(time(0));

    tmp.resize(M);
	grad.resize(N);

    k1.resize(M, N);
    obs1.resize(M);
	// 添加一些大数
	int tmp_id, tmp_size;
	double tmp_val;
	for (int i = 0; i < M; i++)
	{
		tmp_size = random_int(25, 35);
		for (int j = 0; j < tmp_size; j++)
		{
			tmp_id = random_int(0, N);
			tmp_val = random_double(-1.0, 1.0);

			k1[i][tmp_id] = tmp_val;
		}
	}

    k2.resize(M, N);
    obs2.resize(M);
	// 添加一些大数
	for (int i = 0; i < M; i++)
	{
		tmp_size = random_int(25, 35);
		for (int j = 0; j < tmp_size; j++)
		{
			tmp_id = random_int(0, N);
			tmp_val = random_double(-200.0, 200.0);

			k2[i][tmp_id] = tmp_val;
		}
	}

    k3.resize(M, N);
    obs3.resize(M);
	// 添加一些大数
	for (int i = 0; i < M; i++)
	{
		tmp_size = random_int(25, 35);
		for (int j = 0; j < tmp_size; j++)
		{
			tmp_id = random_int(0, N);
			tmp_val = random_double(-0.01, 0.01);

			k3[i][tmp_id] = tmp_val;
		}
	}
}

ex8::~ex8(){}

double ex8::LBFGS_Evaluate(const gctl::_1d_array &x, gctl::_1d_array &g)
{
    gctl::matvec(tmp, k1, x);
    tmp -= obs1;

    gctl::matvec(grad, k1, tmp, gctl::Trans);
    gctl::scale(grad, 2.0/M);

    AddSingleLoss(gctl::power2(gctl::module(tmp, gctl::L2))/M, grad);

    gctl::matvec(tmp, k2, x);
    tmp -= obs2;

    gctl::matvec(grad, k2, tmp, gctl::Trans);
    gctl::scale(grad, 2.0/M);

   	AddSingleLoss(gctl::power2(gctl::module(tmp, gctl::L2))/M, grad);

    gctl::matvec(tmp, k3, x);
    tmp -= obs3;

    gctl::matvec(grad, k3, tmp, gctl::Trans);
    gctl::scale(grad, 2.0/M);

    AddSingleLoss(gctl::power2(gctl::module(tmp, gctl::L2))/M, grad);

    return GradNormLoss(g);
}

int ex8::LBFGS_Progress(const gctl::_1d_array &x, const gctl::_1d_array &g, const double fx, 
    const double converge, const double rate, const gctl::lbfgs_para param, int k, int ls, std::ostream &ss)
{
    UpdateWeights();

    return gctl::lbfgs_solver::LBFGS_Progress(x, g, fx, converge, rate, param, k, ls, ss);
}

void ex8::CalTarget(const gctl::_1d_array &x)
{
    // 计算正演值
    gctl::matvec(obs1, k1, x);
	for (int i = 0; i < M; i++)
	{
		// 添加噪声
		obs1[i] += random_double(-1e-3, 1e-3);
	}

    gctl::matvec(obs2, k2, x);
	for (int i = 0; i < M; i++)
	{
		// 添加噪声
		obs2[i] += random_double(-1e-3, 1e-3);
	}

    gctl::matvec(obs3, k3, x);
	for (int i = 0; i < M; i++)
	{
		// 添加噪声
		obs3[i] += random_double(-1e-3, 1e-3);
	}
    return;
}

int main(int argc, char const *argv[])
{
	// 生成一组正演解
    gctl::_1d_array fm(N);
    random(fm, 1.0, 2.0, gctl::RdUniform);

    ex8 test;

	// 计算拟合目标项
    test.CalTarget(fm);

	// 声明一组解
	gctl::_1d_array m(N, 0.0);

    gctl::lbfgs_para self_para = test.default_lbfgs_para();
    self_para.linesearch = gctl::LBFGS_LINESEARCH_BACKTRACKING_STRONG_WOLFE;
	self_para.epsilon = 1e-6;

	test.set_lbfgs_para(self_para);
	test.show_lbfgs_para();

    test.InitGradNorm(3, N);
    test.set_control_weight(1.0);
    test.set_weight_step(0.00001);

	double fx = test.LBFGS_Minimize(m);

    std::clog << "maximal difference: " << max_diff(fm, m) << std::endl;

    gctl::_1d_array records;
    test.get_records(records);
    for (size_t i = 0; i < records.size(); i++)
    {
        if ((i+1)%3 == 0)
        {
            std::cout << records[i] << "\n";
        }
        else std::cout << records[i] << " ";
    }
	return 0;
}