#include "sgm_esl.h"

void SGM_ESL::OptimizeCG_Precondition(_1dArray obs_array){
	int i,j;
	double data_object_func, dTAd, zTr, z1Tr1, ak, beta_k;
	_1dArray r_k, r_k1, z_k, z_k1, d_k;
	_1dArray Adk, Adk_part;
	_1dArray predict_field;

	r_k.resize(mod_blocks_num_); r_k1.resize(mod_blocks_num_);
	z_k.resize(mod_blocks_num_); z_k1.resize(mod_blocks_num_);
	d_k.resize(mod_blocks_num_);
	Adk.resize(mod_blocks_num_); Adk_part.resize(obs_points_num_);
	predict_field.resize(obs_points_num_);

	//初始化反演密度向量
	for (i = 0; i < mod_blocks_num_; i++)
		invert_model_[i] = init_model_[i];

#pragma omp parallel for private(i,j) schedule(guided)
	for (i = 0; i < obs_points_num_ ; i++){
		Adk_part[i] = 0.0;
		for (j = 0; j < mod_blocks_num_; j++){
			Adk_part[i] += GM_kernel_[i][j]*invert_model_[j];
		}
	}

#pragma omp parallel for private(i,j) schedule(guided)
	for (j = 0; j < mod_blocks_num_; j++){
		Adk[j] = 0.0;
		for (i = 0; i < obs_points_num_ ; i++){
			Adk[j] += GM_kernel_[i][j]*WdTWd[i]*Adk_part[i];
		}
	}

#pragma omp parallel for private(i) schedule(guided)
	for (j = 0; j < mod_blocks_num_; j++)
		r_k[j] = PartB[j] - Adk[j];

#pragma omp parallel for private(i) schedule(guided)
	for (j = 0; j < mod_blocks_num_; j++){
		z_k[j] = Wp[j]*r_k[j];
		d_k[j] = z_k[j];
	}

	for (int t = 0; t < iter_times_; t++){
		//计算数据拟合差函数
#pragma omp parallel for private(i,j) schedule(guided)
		for (i = 0; i < obs_points_num_; i++){
			predict_field[i] = 0.0;
			for (j = 0; j < mod_blocks_num_; j++){
				predict_field[i] += GM_kernel_[i][j] * invert_model_[j];
			}
		}

		data_object_func = 0.0;
		for (i = 0; i < obs_points_num_; i++){
			data_object_func += pow((obs_array[i] - predict_field[i])/obs_points_[i].dev,2);
		}
		data_object_func /= obs_points_num_;

		if (data_object_func <= 1.0){
			cout << "Iteration limit reached!\tconvergence value = " << data_object_func << endl;
			cout << "===================================" << endl;
			break;
		}
		else{
			cout << "Iteration times: " << t << "\tconvergence value = " << data_object_func << endl;
			MOVEUP(1);
			cout << CLEARLINE;
		}

		//计算r_0
#pragma omp parallel for private(i,j) schedule(guided)
		for (i = 0; i < obs_points_num_ ; i++){
			Adk_part[i] = 0.0;
			for (j = 0; j < mod_blocks_num_; j++){
				Adk_part[i] += GM_kernel_[i][j]*d_k[j];
			}
		}

#pragma omp parallel for private(i,j) schedule(guided)
		for (j = 0; j < mod_blocks_num_; j++){
			Adk[j] = 0.0;
			for (i = 0; i < obs_points_num_ ; i++){
				Adk[j] += GM_kernel_[i][j]*WdTWd[i]*Adk_part[i];
			}
		}

		zTr = 0.0; dTAd = 0.0;
		for (j = 0; j < mod_blocks_num_; j++){
			zTr += z_k[j]*r_k[j];
			dTAd += d_k[j]*Adk[j];
		}
		ak = zTr/dTAd;

		for (j = 0; j < mod_blocks_num_; j++){
			invert_model_[j] += ak*d_k[j];
			r_k1[j] = r_k[j] - ak*Adk[j];
			z_k1[j] = Wp[j]*r_k1[j];
		}

		z1Tr1 = 0.0;
		for (j = 0; j < mod_blocks_num_; j++)
			z1Tr1 += z_k1[j]*r_k1[j];
		beta_k = z1Tr1/zTr;

		for (j = 0; j < mod_blocks_num_; j++){
			d_k[j] = z_k1[j] + beta_k*d_k[j];
			r_k[j] = r_k1[j];
			z_k[j] = z_k1[j];
		}
	}
	return;
}