#include "drtp.h"

void DRTP::ReadConfig(std::string filename)
{
	gopt_.add_options({GRIDIN}, {true});
	gopt_.add_options({GRIDOUT, GRIDPARA, USEMEAN, TAYLOR}, {false, false, false, false});
	gopt_.add_options({CALLGEOMAG, GEOMAGEXE, GEOMAGPARA, GEOMAGOUT, IGRFMODEL}, {false, false, false, false, false});
	gopt_.read_options(filename);
	gopt_.show_options();
	gopt_.check_mandatory();
	return;
}

void DRTP::ReadGrid()
{
	if (gopt_.has_value(GRIDPARA))
	{
		std::stringstream tmpss(gopt_.get_value(GRIDPARA));
		gctl::parse_string_to_value(gopt_.get_value(GRIDPARA), ',', false, netcdf_lon_name_, netcdf_lat_name_, netcdf_data_name_);
	}

	ingrid_name_ = gopt_.get_value(GRIDIN);
	gctl::read_netcdf_grid(ingrid_name_, in_lon_, in_lat_, grid_, netcdf_lon_name_, netcdf_lat_name_, netcdf_data_name_);

	rtp_.resize(grid_.row_size(), grid_.col_size(), NAN);
	rsize_ = in_lat_.size();
	csize_ = in_lon_.size();
	dlon_ = (in_lon_[in_lon_.size()-1] - in_lon_[0])/(in_lon_.size()-1);
	dlat_ = (in_lat_[in_lat_.size()-1] - in_lat_[0])/(in_lat_.size()-1);
	return;
}

void DRTP::ReadGEOMAG()
{
	std::string geomag_output = gopt_.get_value(GEOMAGOUT);

	std::ifstream infile;
	gctl::open_infile(infile, geomag_output, "");

	geomag_dec_.resize(rsize_, csize_, NAN);
	geomag_inc_.resize(rsize_, csize_, NAN);

	std::string err_str, tmp_str;
	std::stringstream tmp_ss;
	std::string date, cs_sys, alti_str, dd_str, dm_str, id_str, im_str;
	double lat, lon, H, X, Y, Z, F, dD, dI, dH, dX, dY, dZ, dF;
	int m, n, d_deg, d_min, i_deg, i_min;
	while (getline(infile, tmp_str))
	{
		if (tmp_str[0] == 'D') // 如果第一行是头信息就跳过
			continue;

		//Date Coord-System Altitude Latitude Longitude D_deg D_min I_deg I_min H_nT X_nT Y_nT Z_nT F_nT dD_min dI_min dH_nT dX_nT dY_nT dZ_nT dF_nT
		//2017,6,15 D M4000.00 30.0167 105.017 -2d 16m 46d 52m 34358.9 34332.2 -1354.1 36667.2 50249.6 -3.8 5.6 -13.2 -14.8 -37.9 105.0 67.6
		tmp_ss.clear();
		tmp_ss.str(tmp_str);
		tmp_ss >> date >> cs_sys >> alti_str >> lat >> lon 
			>> dd_str >> dm_str >> id_str >> im_str
			>> H >> X >> Y >> Z >> F >> dD >> dI >> dH 
			>> dX >> dY >> dZ >> dF;

		if (1 != sscanf(dd_str.c_str(), "%dd", &d_deg) || 
			1 != sscanf(dm_str.c_str(), "%dm", &d_min) ||
			1 != sscanf(id_str.c_str(), "%dd", &i_deg) ||
			1 != sscanf(im_str.c_str(), "%dm", &i_min))
		{
			throw gctl::invalid_argument("Wrong format: " + tmp_str);
		}

		m = round((lat - in_lat_[0])/dlat_);
		n = round((lon - in_lon_[0])/dlon_);

		geomag_inc_[m][n] = fabs(i_deg) + i_min/60.0;
		geomag_dec_[m][n] = fabs(d_deg) + d_min/60.0;
		if (i_deg < 0) geomag_inc_[m][n] *= -1;
		if (d_deg < 0) geomag_dec_[m][n] *= -1;
	}
	infile.close();

	mean_inc_ = mean_dec_ = 0.0;
	for (int i = 0; i < rsize_; i++)
	{
		for (int j = 0; j < csize_; j++)
		{
			if (std::isnan(geomag_dec_[i][j]) || std::isnan(geomag_inc_[i][j]))
			{
				throw gctl::runtime_error("Invalid geomagnetic declination or inclination found.");
			}

			mean_dec_ += geomag_dec_[i][j];
			mean_inc_ += geomag_inc_[i][j];
		}
	}
	mean_inc_ /= (geomag_dec_.row_size()*geomag_dec_.col_size());
	mean_dec_ /= (geomag_inc_.row_size()*geomag_inc_.col_size());
	return;
}

void DRTP::SaveGrid()
{
	if (gopt_.has_value(GRIDOUT)) outgrid_name_ = gopt_.get_value(GRIDOUT);
	else if (gopt_.has_value(USEMEAN)) outgrid_name_ = gopt_.get_value(GRIDIN) + "_RTP";
	else outgrid_name_ = gopt_.get_value(GRIDIN) + "_DRTP";

	gctl::save_netcdf_grid(outgrid_name_, rtp_, in_lon_, in_lat_, netcdf_lon_name_, netcdf_lat_name_, netcdf_data_name_);
	return;
}

void DRTP::CalNormalField()
{
	std::string geomag_exe = gopt_.get_value(GEOMAGEXE);
	std::string geomag_para = gopt_.get_value(GEOMAGPARA);
	std::string geomag_output = gopt_.get_value(GEOMAGOUT);
	std::string IGRF = gopt_.get_value(IGRFMODEL);

	// 生成临时文件
	std::ofstream outfile;
	gctl::open_outfile(outfile, "drtp_tmp", ".txt");
	for (int i = 0; i < in_lat_.size(); i++)
	{
		for (int j = 0; j < in_lon_.size(); j++)
		{
			outfile << geomag_para << " " << in_lat_[i] << " " << in_lon_[j] << std::endl;
		}
	}
	outfile.close();

	std::string geomag_cmd = "./" + geomag_exe + " " + IGRF + " f drtp_tmp.txt " + geomag_output; 
	int ret = system(geomag_cmd.c_str());
	ret = system("rm drtp_tmp.txt");
	return;
}

void DRTP::RTP(double inc, double dec)
{
	int ori_row = 0, ori_col = 0;
	gctl::matrix<double> ingrid_ex;
	gctl::cosine_extrapolate_2d(grid_, ingrid_ex, ori_row, ori_col);

	gctl::save_netcdf_grid(gopt_.get_value(GRIDIN) + "_expanded", ingrid_ex, 0, 1, 0, 1);

	int M = ingrid_ex.row_size();
	int N = ingrid_ex.col_size();
	fftw_complex *in = (fftw_complex*) fftw_malloc(sizeof(fftw_complex)*M*N);
	fftw_complex *out= (fftw_complex*) fftw_malloc(sizeof(fftw_complex)*M*N);

	for (int i = 0; i < M; i++)
	{
		for (int j = 0; j < N; j++)
		{
			in[i*N+j][0] = ingrid_ex[i][j];
			in[i*N+j][1] = 0.0;
		}
	}

	fftw_plan p = fftw_plan_dft_2d(M, N, in, out, FFTW_FORWARD, FFTW_ESTIMATE);
	fftw_execute(p);

	int half_M = ceil(0.5*M);
	int half_N = ceil(0.5*N);
	double du = 2.0*GCTL_Pi/((M-1)*dlat_);
	double dv = 2.0*GCTL_Pi/((N-1)*dlon_);

	double cc = cosd(dec)*cosd(inc);
	double sc = sind(dec)*cosd(inc);
	double s  = sind(inc);

	int idx;
	double U, V, Rpart, Ipart;
	double temp, temp1, temp2, temp3, temp4, base;
	for (int i = 0; i < half_M; i++)
	{
		for (int j = 0; j < half_N; j++)
		{
			U = (i+1)*du; V = (j+1)*dv;
			temp = U*U+V*V;

			idx = i*N+j;
			temp1 = cc*U+sc*V;
			temp2 = s*s*temp-temp1*temp1;
			base = temp2*temp2+4.0*s*s*temp1*temp1*temp;
			Rpart = temp*temp2/base;
			Ipart = -2.0*s*temp1*pow(temp,1.5)/base;
			temp3 = out[idx][0];
			temp4 = out[idx][1];
			in[idx][0] = Rpart*temp3-temp4*Ipart;
			in[idx][1] = Rpart*temp4+Ipart*temp3;

			idx = (M-1-i)*N+j;
			temp1 = -cc*U+sc*V;
			temp2 = s*s*temp-temp1*temp1;
			base = temp2*temp2+4.0*s*s*temp1*temp1*temp;
			Rpart = temp*temp2/base;
			Ipart = -2.0*s*temp1*pow(temp,1.5)/base;
			temp3 = out[idx][0];
			temp4 = out[idx][1];
			in[idx][0] = Rpart*temp3-temp4*Ipart;
			in[idx][1] = Rpart*temp4+Ipart*temp3;

			idx = i*N+N-1-j;
			temp1 = cc*U-sc*V;
			temp2 = s*s*temp-temp1*temp1;
			base = temp2*temp2+4.0*s*s*temp1*temp1*temp;
			Rpart = temp*temp2/base;
			Ipart = -2.0*s*temp1*pow(temp,1.5)/base;
			temp3 = out[idx][0];
			temp4 = out[idx][1];
			in[idx][0] = Rpart*temp3-temp4*Ipart;
			in[idx][1] = Rpart*temp4+Ipart*temp3;

			idx = (M-1-i)*N+N-1-j;
			temp1 = -cc*U-sc*V;
			temp2 = s*s*temp-temp1*temp1;
			base = temp2*temp2+4.0*s*s*temp1*temp1*temp;
			Rpart = temp*temp2/base;
			Ipart = -2.0*s*temp1*pow(temp,1.5)/base;
			temp3 = out[idx][0];
			temp4 = out[idx][1];
			in[idx][0] = Rpart*temp3-temp4*Ipart;
			in[idx][1] = Rpart*temp4+Ipart*temp3;
		}
	}

	p = fftw_plan_dft_2d(M, N, in, out, FFTW_BACKWARD, FFTW_ESTIMATE);
	fftw_execute(p);

	for (int i = 0; i < M; i++)
	{
		for (int j = 0; j < N; j++)
		{
			ingrid_ex[i][j] = out[i*N+j][0]/(M*N);
		}
	}
	

	for (int i = 0; i < rsize_; i++)
	{
		for (int j = 0; j < csize_; j++)
		{
			rtp_[i][j] = ingrid_ex[i+ori_row][j+ori_col];
		}
	}

	fftw_destroy_plan(p);
	fftw_free(in); 
	fftw_free(out);
	return;
}

void DRTP::DiffRTP(int order)
{
	// 使用微小扰动计算RTP相对与倾角与偏角的导数
	double dinc = 0.01, ddec = 0.01;
	// 拷贝平均化极磁异常到临时数组
	gctl::matrix<double> drtp(rtp_), tmp(rtp_);
	gctl::matrix<double> diffinc(rtp_.row_size(), rtp_.col_size());
	gctl::matrix<double> diffdec(rtp_.row_size(), rtp_.col_size());
	RTP(mean_inc_+dinc, mean_dec_);
	for (int i = 0; i < rsize_; i++)
	{
		for (int j = 0; j < csize_; j++)
		{
			diffinc[i][j] = (geomag_inc_[i][j] - mean_inc_)*(rtp_[i][j] - tmp[i][j])/dinc;
			drtp[i][j] += diffinc[i][j];
		}
	}

	RTP(mean_inc_, mean_dec_+ddec);
	for (int i = 0; i < rsize_; i++)
	{
		for (int j = 0; j < csize_; j++)
		{
			diffdec[i][j] = (geomag_dec_[i][j] - mean_dec_)*(rtp_[i][j] - tmp[i][j])/ddec;
			drtp[i][j] += diffdec[i][j];
		}
	}

	double factor;
	for (int o = 2; o <= order; o++)
	{
		factor = 1.0;
		for (int n = o; n > 1; n--)
		{
			factor *= n;
		}
		factor = 1.0/factor;

		grid_ = diffinc;
		RTP(mean_inc_, mean_dec_);

		tmp = rtp_;
		RTP(mean_inc_+dinc, mean_dec_);
		for (int i = 0; i < rsize_; i++)
		{
			for (int j = 0; j < csize_; j++)
			{
				diffinc[i][j] = factor*pow(geomag_inc_[i][j] - mean_inc_, o)*(rtp_[i][j] - tmp[i][j])/dinc;
				drtp[i][j] += diffinc[i][j];
			}
		}

		grid_ = diffdec;
		RTP(mean_inc_, mean_dec_);

		tmp = rtp_;
		RTP(mean_inc_, mean_dec_+ddec);
		for (int i = 0; i < rsize_; i++)
		{
			for (int j = 0; j < csize_; j++)
			{
				diffdec[i][j] = factor*pow(geomag_dec_[i][j] - mean_dec_, o)*(rtp_[i][j] - tmp[i][j])/ddec;
				drtp[i][j] += diffdec[i][j];
			}
		}		
	}

	rtp_ = drtp;
	return;
}

void DRTP::Process()
{
	std::string err_str;
	// 读入网格文件
	ReadGrid();
	// 使用输入的平均参数
	std::string geomag_str = gopt_.get_value(USEMEAN);
	if (gopt_.has_value(USEMEAN))
	{
		double inc, dec;
		gctl::parse_string_to_value(geomag_str, '/', false, inc, dec);
		if (inc < -90 || inc > 90 || dec < -90 || dec > 90)
		{
			gctl::invalid_argument("Invalid geomagnetic parameters: " + geomag_str);
		}
		RTP(inc, dec);
	}
	else
	{
		// 计算正常场参数
		if (gopt_.get_value(CALLGEOMAG) == "yes")
			CalNormalField();
		// 读入正常场参数
		ReadGEOMAG();
		// 读入差分阶数
		if (gopt_.has_value(TAYLOR))
		{
			int order = atoi(gopt_.get_value(TAYLOR).c_str());
			if (order < 1 || order > 5)
			{
				gctl::invalid_argument("Invalid taylor order: " + gopt_.get_value(TAYLOR));
			}
			taylor_order_ = order;
		}
		// 化极
		RTP(mean_inc_, mean_dec_);
		DiffRTP(taylor_order_);
	}
	// 输出网格文件
	SaveGrid();
	return;
}