/**********************************************************
      Grid searching for L-G-F-T by fitting P-to-S phase (Ps) harmonic feature in receiver functions.
	    V1 (converted from bash version), By Mei Feng, 2024/07/03
	    V0 (bash version) By Mei Feng

***********************************************************/

#ifdef _OPENMP
#include <omp.h>
#endif

#include <cstdio>
#include <string>
#include <cmath>
#include "mystds.h"     //General functions
#include "Sac_head.h"

using namespace std;

struct Arguement
{
	string input_filename;
	string fepsi;
	double mo, ka, vpc, vpm; // a priori moho and kappa 
	double L, G, F, T;           //solution for Lamda, Gama, Fai, Tau.
	bool no_openmp; //close openmp
};

struct sacRF
{
	int npts;
	double b, dt, baz;
	float *time;
	float *amp;
};


int get_argument(struct Arguement &argu, int argc, char**argv);

int main(int argc, char **argv)
{
	struct Arguement argu;
	get_argument(argu, argc, argv);

#ifdef _OPENMP
	if (argu.no_openmp == false)
		set_openmp_parallel_process_num();
#endif


	double rayp = 0.0576;
	double tPs = argu.mo * (sqrt((argu.ka/argu.vpc)*(argu.ka/argu.vpc) - rayp*rayp) - sqrt((1.0/argu.vpc)*(1.0/argu.vpc) - rayp*rayp));

	// read R-RFs into memory
	int nbaz = 0;
	int nbaz_bin[360] = {0};
	struct sacRF * sacsR = new struct sacRF[360];
	memset(sacsR, 0, sizeof(struct sacRF)*360);

	string strFile=string("r")+argu.input_filename;
	FILE *fRF = openfile(strFile.c_str(), "r");
	char line[1024], strTmp[8];
	int ibaz;
	double dBaz;

	if(fRF == NULL) exit(0);
	while(feof(fRF)==false)
	{
		if(fgets(line, 1024, fRF) == NULL) continue;
		if(line[0]!='>') continue;

		if (sscanf(line, "%s%d%lf%d%lf%lf%lf%d", strTmp, &ibaz, &dBaz, &sacsR[nbaz].npts, &sacsR[nbaz].b, &sacsR[nbaz].dt, &sacsR[nbaz].baz, &nbaz_bin[nbaz] ) < 7) continue;

		sacsR[nbaz].time = new float[sacsR[nbaz].npts];
		sacsR[nbaz].amp  = new float[sacsR[nbaz].npts];
		for(int i=0; i<sacsR[nbaz].npts; i++)
		{
			if(fgets(line, 1024, fRF) == NULL) continue;
			if(sscanf(line, "%f%f", sacsR[nbaz].time+i, sacsR[nbaz].amp+i) != 2) continue;
		}
		nbaz++;
	}
	fclose(fRF);

	// read T-RFs into memory
	struct sacRF * sacsT = new struct sacRF[nbaz];
	memset(sacsT, 0, sizeof(struct sacRF)*nbaz);

	strFile=string("t")+argu.input_filename;
	fRF = openfile(strFile.c_str(), "r");
	nbaz=0;
	
	if(fRF == NULL) exit(0);
	while(feof(fRF)==false)
	{
		if(fgets(line, 1024, fRF) == NULL) continue;
		if(line[0]!='>') continue;

		if (sscanf(line, "%s%d%lf%d%lf%lf%lf%d", strTmp, &ibaz, &dBaz, &sacsT[nbaz].npts, &sacsT[nbaz].b, &sacsT[nbaz].dt, &sacsT[nbaz].baz, &nbaz_bin[nbaz] ) < 7) continue;

		sacsT[nbaz].time = new float[sacsT[nbaz].npts];
		sacsT[nbaz].amp  = new float[sacsT[nbaz].npts];
		for(int i=0; i<sacsT[nbaz].npts; i++)
		{
			if(fgets(line, 1024, fRF) == NULL) continue;
			if (sscanf(line, "%f%f", sacsT[nbaz].time+i, sacsT[nbaz].amp+i) != 2) continue;
		}
		nbaz++;
	}
	fclose(fRF);

	// read epsilon into memory
	float *epsi = new float[nbaz];
	memset(epsi, 0, sizeof(float)*nbaz);
	float *baztmp = new float[nbaz];
	memset(baztmp, 0, sizeof(float)*nbaz);

	FILE *fepsi = openfile(argu.fepsi.c_str(), "r");
	nbaz=0;
	if(fepsi == NULL) exit(0);
	while(feof(fepsi)==false)
	{
		if(fgets(line, 1024, fepsi) == NULL) continue;
		if (sscanf(line, "%f%f", &baztmp[nbaz], &epsi[nbaz] ) < 2) continue;
		epsi[nbaz]*=pai/180;
		nbaz++;
	}
	fclose(fepsi);

	double lamda = argu.L;
	double lam = lamda * pai/180;
	// double stk = lam - 0.5*pai;
	double gama = argu.G;
	double gam = gama * pai/180;
	double az = argu.F;
	double tau = argu.T;

	strFile = string("ACorrected")+argu.input_filename;
	FILE *fARF = openfile(strFile.c_str(), "w");
	strFile = string("ADCorrected")+argu.input_filename;
	FILE *fADRF = openfile(strFile.c_str(), "w");
	for(int ibaz=0; ibaz<nbaz; ibaz++)
	{
		double bazi=sacsR[ibaz].baz*pai/180;

		// calculate tPs for dipping interface, according to TanPing, ChenYun et al., 2018. CJG
		double PI1=-rayp*cos(bazi);
		double PI2=-rayp*sin(bazi);
		double PI3=-sqrt(1/(argu.vpm*argu.vpm)-(rayp*rayp));          
		// double P1I1=sin(lam)*PI1-cos(lam)*PI2;
		// double P1I2=cos(gam)*cos(lam)*PI1+cos(gam)*sin(lam)*PI2+sin(gam)*PI3;
		// double P1I3=-sin(gam)*cos(lam)*PI1-sin(gam)*sin(lam)*PI2+cos(gam)*PI3;           
		// double P2I1=P1I1;
		// double P2I2=P1I2;
		// double P2I3=-sqrt(1/(argu.vpc*argu.vpc)-P1I1*P1I1-P1I2*P1I2);              
		double n1=sin(gam)*cos(lam);
		double n2=sin(gam)*sin(lam);
		double n3=-cos(gam);
		double nPI=n1*PI1+n2*PI2+n3*PI3;
		double tPsDip = argu.mo * cos(gam) * (sqrt((argu.ka/argu.vpc)*(argu.ka/argu.vpc)-1/(argu.vpm*argu.vpm)+nPI*nPI)-sqrt(1/(argu.vpc*argu.vpc)-1/(argu.vpm*argu.vpm)+nPI*nPI));
		
		// polarization angle, epsilon, according to Langston, 1977, BSSA
		// double Nj = -sin(bazi-0.5*pai)*argu.vpm*rayp;
		// double Ej =  cos(bazi-0.5*pai)*argu.vpm*rayp;
		// double Zj = -sqrt(1-(argu.vpm*rayp)*(argu.vpm*rayp));
		// double qb1 =  cos(stk) * Nj - cos(gam) * sin(stk) * Ej + sin(gam) * sin(stk) * Zj;
		// double qb2 =  sin(stk) * Nj + cos(gam) * cos(stk) * Ej - sin(gam) * cos(stk) * Zj;
		// double qb3 =                  sin(gam) * sin(stk) * Ej + cos(gam) * Zj;
		// double sinib = sqrt(qb1*qb1 + qb2*qb2)/sqrt(qb1*qb1 + qb2*qb2 + qb3*qb3);
		// double sinia = sinib/argu.vpm*argu.vpc;
		// double cosia = sqrt(1-sinia*sinia); 
		// double alpha = sinia/sqrt(qb1*qb1 + qb2*qb2);
		// double qa1 = alpha * qb1;
		// double qa2 = alpha * qb2;
		// double sign = (qb3>0) ? 1 : -1;
		// double qa3 = sign*cosia;
		// double N =  cos(stk) * qa1 + sin(stk) * qa2;
		// double E = -cos(gam) * sin(stk) * qa1 + cos(gam) * cos(stk) * qa2 + sin(gam) * qa3;
		// // double Z =  sin(gam) * sin(stk) * qa1 - sin(gam) * cos(stk) * qa2 + cos(gam) * qa3;
		// double epsi = atan2(E, N)-(bazi-pai);
		// // if(epsi>=2*pai)
		// // 	epsi-=2*pai;
		// // else if (epsi<=0)
		// // 	epsi+=2*pai;
		
		// fprintf(stderr, "%f\t%lg\t%f\n", baztmp[ibaz], sacsR[ibaz].baz, -epsi[ibaz]);

		// rotate from R & T to F & S 
		double fai = (az - sacsR[ibaz].baz)*pai/180;

		double *ampFdt = new double[sacsR[0].npts];
		double *ampSdt = new double[sacsR[0].npts];
		double *timef = new double[sacsR[0].npts];
		memset(ampFdt, 0, sizeof(double) * sacsR[0].npts);
		memset(ampSdt, 0, sizeof(double) * sacsR[0].npts);
		memset(timef,  0, sizeof(double) * sacsR[0].npts);


		int iptk=0;
		int iptm=0;
		for(int j=0; j<sacsR[0].npts; j++)
		{
			double ttmp=sacsR[ibaz].time[j] + 0.5 * tau * cos(gam);
			if(ttmp>=-3 && ttmp<=43) 
			{
				timef[iptk]=ttmp;
				ampFdt[iptk]= sacsR[ibaz].amp[j] * cos(fai) + sacsT[ibaz].amp[j] * sin(fai);
				// fprintf(stderr, "%lg\t%lg\n", sacsR[ibaz].time[j], ampFdt[iptk]);
				iptk++;
			}
			ttmp=sacsR[ibaz].time[j] - 0.5 * tau * cos(gam);
			if(ttmp>=-3 && ttmp<=43) 
			{
				ampSdt[iptm] = -1*sacsR[ibaz].amp[j] * sin(fai) + sacsT[ibaz].amp[j] * cos(fai);
				// fprintf(stderr, "%lg\t%lg\n", ttmp, ampSdt[iptm]);
				iptm++;
			}
		}

		int km = (iptm>iptk) ? iptk : iptm;
		double *ampRcP = new double[km];
		double *ampTcP = new double[km];
		double *timet = new double[km];
		memset(ampRcP, 0, sizeof(double) * km);
		memset(ampTcP, 0, sizeof(double) * km);
		memset(timet,  0, sizeof(double) * km);

		iptk=0;
		iptm=0;
		fprintf(fARF, "> %d %g %d %g %g %g %d\n", ibaz, dBaz, km, sacsR[ibaz].b, sacsR[ibaz].dt, sacsR[ibaz].baz, nbaz_bin[ibaz]);
		for(int k=0; k<km; k++)
		{
			double rc =  ampFdt[k] * cos(-fai) + ampSdt[k] * sin(-fai);
			double tc = -ampFdt[k] * sin(-fai) + ampSdt[k] * cos(-fai);
			// fprintf(stderr, "%g\t%g\t%g\n", timef[k], rc, tc);
			fprintf(fARF, "%g\t%g\t%g\n", timef[k], rc, tc);

			double ttmp = timef[k] + (tPs-tPsDip);
			if(ttmp>=-1 && ttmp<=41)
			{
				timet[iptm]=ttmp;
				ampRcP[iptm] = rc * cos(-epsi[ibaz]) + tc * sin(-epsi[ibaz]);
				// fprintf(stderr, "%g\t%g\n", timet[iptm], ampTcP[iptm]);
				iptm++;
			}
			ttmp = timef[k];
			if(ttmp>=-1 && ttmp<=41)
			{
				ampTcP[iptk] = -rc * sin(-epsi[ibaz]) + tc * cos(-epsi[ibaz]);
				// fprintf(stderr, "%g\t%g\n", timet[iptk], ampTcP[iptk]);
				iptk++;
			}
		}

		km = (iptm>iptk) ? iptk : iptm;
		fprintf(fADRF, "> %d %g %d %g %g %g %d\n", ibaz, dBaz, km, sacsR[ibaz].b, sacsR[ibaz].dt, sacsR[ibaz].baz, nbaz_bin[ibaz]);
		for(int k=0; k<km; k++)
		{
			fprintf(fADRF, "%g\t%g\t%g\n", timet[k], ampRcP[k], ampTcP[k]);
		}
		
		delete ampFdt; ampFdt=NULL;
		delete ampSdt; ampSdt=NULL;
		delete timef; timef=NULL;
		delete ampRcP; ampRcP=NULL;
		delete ampTcP; ampTcP=NULL;
		delete timet; timet=NULL;
	}

	fclose(fARF);
	fclose(fADRF);

	delete sacsR;
	delete sacsT;

	return 0;
}

int show_help(char* argv0 /*=argv[0]*/)
{
	printf("%s (Version 2.1) by Meijian An, 2023/11/14\n\n", argv0);
	printf("Info:\n"
		"   Make L-G-F-T (model space) grid by fitting P-to-S phase (Ps) in receiver functions\n\n");
	printf("arguments:\n"
		"     -h:        Show this help information.\n"
		"     files [rt]RFbin_Ps.out stored the binned Receiver functions should be prepared. \n"
		"  -input=x:   postfix of the input RFbin file name (RFbin_x.out).\n"
		"  -fepsi=x:   filename for epsilon, calculated by Raysum. Default is epsilon.lst.\n"
		"  -moho=x:      preliminary Moho depth. Default is 40 km \n"
		"  -kappa=x:     preliminary Kappa. Default is 1.732\n"
		"  -vpc=x:     vp in the crust. Default is 6.4 km/s \n"
		"  -vpm=x:     vp in the mantle. Default is 8.1 km/s \n"
		"  -L=x:      determined dipping direction Lamda. \n"
		"  -G=x:      determined dipping angle     Gama.  \n"
		"  -F=x:      determined fast-wave direction Fai. \n"
		"  -T=x:      determined delay time Tau. \n"
		"-no_openmp      No parallel computing.\n"
		);

	return 1;
}

int get_argument(struct Arguement &argu, int argc, char**argv)
{
	if(get_switch("-h", false, argc, argv) == true || argc < 3)
	{
		show_help(argv[0]);
		exit(0);
	}

 	//Get All arguement values.
	argu.input_filename   = get_Parameter("-input=",    string("RFbin.out"), argc, argv);
	argu.fepsi   = get_Parameter("-fepsi=",    string("epsilon.lst"), argc, argv);

	argu.mo    = get_Parameter("-moho=", double(40.0),   argc, argv);
	argu.ka    = get_Parameter("-kappa=", double(1.732), argc, argv);
	argu.vpc    = get_Parameter("-vpc=", double(6.4), argc, argv);
	argu.vpm    = get_Parameter("-vpm=", double(8.1), argc, argv);

	argu.L = get_Parameter("-L=", double(10.0), argc, argv);
	argu.G = get_Parameter("-G=", double(2.5), argc, argv);
	argu.F = get_Parameter("-F=", double(5.0), argc, argv);
	argu.T = get_Parameter("-T=", double(0.1), argc, argv);

#ifdef _OPENMP
	argu.no_openmp = get_switch("-no_openmp", false, argc, argv);
#else
	argu.no_openmp = true;
#endif

	SHOWerr1("%s options are:\n", argv[0]);
	SHOWerr1("\t         input = %s\n", argu.input_filename.c_str());
	SHOWerr1("\t         fepsi = %s\n", argu.fepsi.c_str());
	SHOWerr1("\t         moho = %g\n", argu.mo);
	SHOWerr1("\t         kappa = %g\n", argu.ka);
	SHOWerr1("\t         vpc = %g\n", argu.vpc);
	SHOWerr1("\t         vpm = %g\n", argu.vpm);
	SHOWerr1("\t       Lsol = %g\n", argu.L);
	SHOWerr1("\t       Gsol = %g\n", argu.G);
	SHOWerr1("\t       Fsol = %g\n", argu.F);
	SHOWerr1("\t       Tsol = %g\n", argu.T);

	return 1;
}

