/**********************************************************
      Make H-K (model space) grid by stacking harmonic corrected P-to-S phase (Ps) 
      and multiples (PpPs, PsPs+PpSs) in receiver functions.
	    V1 (converted from awk version), By Mei Feng, 2024/07/03
	    V0 (awk 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 overlays_filename, harmonic_filename, output_filename;
	double Lsol, Gsol, Fsol, Tsol; // Moho dipping and crustal anisotropy
	double Ws[3];   //weights for the three phases
	double vpc, vpm;
	struct data_range<double> Hmimx; 	//minimum and maximum of thickness (H)
	struct data_range<double> Kmimx; 	//minimum and maximum of velocity ratio (K)
	double dH, dK;
	double dBaz;

	bool no_openmp; //close openmp
	bool normax;   // only normalize the maximium, but not consider the range (minimium)
	double minPsM1;  //not stack when Ps or M1 is small than this value
};

struct sacRF
{
	int npts, ibaz;
	double b, dt, rayp, rayp2, baz;
	double tPs, tM1, tM2; // delay time after anisotropic and dipping correction
	float *amp;
};

struct layer
{
	double thick, kappa, Vp;
};


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

	// back azimuth bins
	int nbaz = 1;
	if(argu.dBaz > 0.0 && argu.dBaz <= 180)
		nbaz = (int)(360.0/argu.dBaz) + 1;
	int *nbaz_bin = new int[nbaz];
	memset(nbaz_bin, 0, sizeof(int)*nbaz);

	// read RFs into memory
	int maxpt = 1000000;
	float  *RFTmp = new float [maxpt];
	struct SacHead sacHdr;

	int nfile = 0;
	struct sacRF * sacs = new struct sacRF[argc];
	for(int i=1; i<argc; i++)
	{
		sacs[nfile].npts = 0;
		if(argv[i][0] == '-') continue;

		FILE *fRF = openfile(argv[i], "rb");
		if(fRF == NULL) continue;

		//read sac head
		if(fread((char*)&sacHdr, 1, sizeof(sacHdr), fRF) == sizeof(sacHdr) && sacHdr.npts > 10)
		{
			sacs[nfile].npts = (sacHdr.npts > maxpt) ? maxpt:sacHdr.npts;
			sacs[nfile].dt   = sacHdr.delta;
			sacs[nfile].b    = sacHdr.b;
			if(sacHdr.baz < 0.0) sacHdr.baz += 360;
			sacs[nfile].baz    = sacHdr.baz;
			sacs[nfile].ibaz  = (nbaz == 1 || sacHdr.baz < -360.0)?0:((int)(sacHdr.baz/argu.dBaz)) ;
			sacs[nfile].rayp = sacHdr.user4;
			sacs[nfile].rayp2 = pow(sacHdr.user4, 2);

			//read sac data
			int nret= fread(RFTmp, sizeof(float), sacs[nfile].npts, fRF);
			if(nret < sacs[nfile].npts)
				sacs[nfile].npts = nret;
		}
		fclose(fRF);
		
		if (sacs[nfile].npts < 10)
			continue;

		nbaz_bin[ sacs[nfile].ibaz ] ++;
		sacs[nfile].amp = new float[sacs[nfile].npts];
		memcpy(sacs[nfile].amp, RFTmp, sizeof(float) * sacs[nfile].npts);
		nfile ++;
	}
	SHOWerr1("files= %d\n", nfile);

	// create array with grid points
	int nH = (int) ((argu.Hmimx.max - argu.Hmimx.min) / argu.dH + 1.499);
	int nK = (int) ((argu.Kmimx.max - argu.Kmimx.min) / argu.dK + 1.499);
	float** rastacksum = new float*[nH];
	for (int i = 0; i < nH; i++)
	{
		rastacksum[i] = new float[nK];
		memset(rastacksum[i], 0, sizeof(float)* nK);
	}

	double lam = D2A(argu.Lsol);
	double gam = D2A(argu.Gsol);
	double az = argu.Fsol;
	double tau = argu.Tsol;

	// make H-K stacking grid for all RF files
	for(int iH=0; iH < nH; iH++)
	{
		double h = argu.Hmimx.min + iH * argu.dH;
		for(int iK=0; iK < nK; iK++)
		{
			double k = argu.Kmimx.min + iK * argu.dK;
			double sumi = 0;
			for(int ifile=0; ifile<nfile; ifile++)
			{
				double dRFw = pow(nbaz_bin[sacs[ifile].ibaz], 0.9);
				
				// calculate tPs, tM1 (PpPs), tM2 (PpSs or PsPs) for flat Moho model, for debug
				// double dtS = sqrt(pow((k/argu.vpc),2) - sacs[ifile].rayp2);
				// double dtP = sqrt(pow((1.0/argu.vpc),2) - sacs[ifile].rayp2);
				// double tPs0 = h * (dtS - dtP);
				// double tPpPs0 = h * (dtS + dtP);
				// double tPpSs0 = h * 2.0 * dtS;		
				
				// correct tPs, tPpPs and tPpSs and take their Amplitude
				// calculate tPs for dipping interface, according to TanPing, ChenYun et al., 2018. CJG
				double bazi=D2A(sacs[ifile].baz);
				double fai=D2A(az)-bazi;

				double PI1=-sacs[ifile].rayp*cos(bazi);
				double PI2=-sacs[ifile].rayp*sin(bazi);
				double PI3=-sqrt(1/(argu.vpm*argu.vpm)-(sacs[ifile].rayp2));          
				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 P3I1=sin(lam)*P2I1+cos(gam)*cos(lam)*P2I2-sin(gam)*cos(lam)*P2I3;
				double P3I2=-cos(lam)*P2I1+cos(gam)*sin(lam)*P2I2-sin(gam)*sin(lam)*P2I3;
				// double P3I3=sin(gam)*P2I2+cos(gam)*P2I3;
				double PR1=P3I1;
				double PR2=P3I2;
				double PR3=sqrt(1/(argu.vpc*argu.vpc)-P3I1*P3I1-P3I2*P3I2);
				double nPR=n1*PR1+n2*PR2+n3*PR3;
				double SR1=P3I1;
				double SR2=P3I2;
				double SR3=sqrt(1/((argu.vpc/k)*(argu.vpc/k))-P3I1*P3I1-P3I2*P3I2);
				double nSR=n1*SR1+n2*SR2+n3*SR3;
				double tPs = h * cos(gam) * (sqrt((k/argu.vpc)*(k/argu.vpc)-1/(argu.vpm*argu.vpm)+nPI*nPI)-sqrt(1/(argu.vpc*argu.vpc)-1/(argu.vpm*argu.vpm)+nPI*nPI));
				double tPpPs = h * cos(gam) * (sqrt(nPR*nPR) + sqrt((k*k)/(argu.vpc*argu.vpc)-1/(argu.vpc*argu.vpc)+nPR*nPR));
				double tPpSs = 2 * h * cos(gam) * sqrt(nSR*nSR);
				// printf("%lg\t%lg\t%lg\t%lg\t%lg\t%lg\t%lg\n", sacs[ifile].baz, tPs0, tPs, tPpPs0, tPpPs, tPpSs0, tPpSs);
				
				tPs -= 0.5 * tau * cos(gam) * cos(2*fai);
				int jpt=int((tPs - sacs[ifile].b)/sacs[ifile].dt + 0.4999);
				double raPs  = sacs[ifile].amp[jpt];
				// fprintf(stderr, "%lg\t%lg\t%lg\n", sacs[ifile].baz, tPs, raPs);  //debug

				tPpPs -= 0.5 * tau * cos(gam) * cos(2*fai);
				jpt=int((tPpPs - sacs[ifile].b)/sacs[ifile].dt + 0.4999);
				double raPpPs  = sacs[ifile].amp[jpt];
				// fprintf(stderr, "%lg\t%lg\t%lg\n", sacs[ifile].baz, tPpPs, raPpPs);  //debug


				tPpSs -= tau * cos(gam) * cos(2*fai);
				jpt=int((tPpSs - sacs[ifile].b)/sacs[ifile].dt + 0.4999);
				double raPpSs  = sacs[ifile].amp[jpt];
				// fprintf(stderr, "%lg\t%lg\t%lg\n", sacs[ifile].baz, tPpSs, raPpSs);  //debug

				sumi += (raPs * argu.Ws[0] + raPpPs * argu.Ws[1] + raPpSs * argu.Ws[2])/dRFw;
			}
			rastacksum[iH][iK] = sumi;
		}
	}

	// get min & max values for normalization
	float rmax, rmin;
	int iHmax=0, iKmax=0;
	for (int iH = 0; iH < nH; iH++)
	{
		for (int iK = 0; iK < nK; iK++)
		{
			if (iH == 0 && iK == 0)
			{
				rmax = rmin = rastacksum[0][0];
			}
			else if (rastacksum[iH][iK] > rmax)
			{
				rmax = rastacksum[iH][iK];
				iHmax = iH;
				iKmax = iK;
			}
			else if (rastacksum[iH][iK] < rmin)
				rmin = rastacksum[iH][iK];
		}
	}

	// normalize and output the array 
	double dr = (rmax - rmin);
	double thick = argu.Hmimx.min;
	if(argu.normax == true) rmin = 0.0;

	FILE* fout = (argu.output_filename.empty())?stdout:openfile(argu.output_filename.c_str(), "w");
	fprintf(fout, "> Hmi,Hmx,dH,Kmi,Kmx,dK,Vp,Hsol,Ksol= %g %g %g %g %g %g %g %g %g\n", argu.Hmimx.min, argu.Hmimx.max, argu.dH, argu.Kmimx.min, argu.Kmimx.max, argu.dK, argu.vpc, argu.Hmimx.min+iHmax*argu.dH, argu.Kmimx.min+iKmax*argu.dK);
	for (int iH = 0; iH < nH; iH++)
	{
		double kappa = argu.Kmimx.min;
		for (int iK = 0; iK < nK; iK++)
		{
			fprintf(fout, "%g\t%g\t%g\t%g\n", thick, kappa, rastacksum[iH][iK], (rastacksum[iH][iK] - rmin) / dr);
			kappa += argu.dK;
		}
		thick += argu.dH;
	}

	if (fout != stdout) fclose(fout);

	// free memory
	for (int iH = 0; iH < nH; iH++)
		delete rastacksum[iH];
	delete rastacksum;

	for (int i = 0; i < nfile; i++)
		delete sacs[i].amp;
	delete sacs;

	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 H-K (model space) grid by stacking P-to-S phase (Ps) and multiples (PpPs, PsPs+PpSs) in receiver functions\n\n");
	printf("arguments:\n"
		"     -h:        Show this help information.\n"
		"     Receiver function files in sac format\n"
		"                Sac headers of USER4 (= Rayp) and Baz are necessary.\n"
		"-overlays=x:    Filename or stdin (\"-\") for structure of overlying layers. Column format is: thick xx vp/vs yy vp\n"
		"                Used to calculate time delays.\n"
		"     -o=x:      Output model-space file name.\n"
		" -Hmimx=x1,x2:  Thickness range.\n"
		" -Kmimx=x1,x2:  Velocity ratio range.\n"
		"    -dH=x:      Thickness interval. Default: 0.1 km\n"
		"    -dK=x:      Velocity ratio interval. Default: 0.001\n"
		"    -vpc=x:      crustal Vp. Default: 6.4 km/s\n"
		"    -vpm=x:      mantle Vp. Default: 8.1 km/s\n"
		"  -dBaz=[5]     Width of Back-azimuth bins for RF normlization. Baz <=0 or >180, all RFs are equally stacked. Default bin is 5 degree.\n"
		"    -Ws=\"x1 x2 x3\": Weight for Ps, PpPs, PsPs+PpSs. Default: 0.6. 0.3 -0.1\n"
		"  -Lsol=x:      dipping direction of Moho\n"
		"  -Gsol=x:      dipping angle of Moho\n"
		"  -Fsol=x:      Fast-wave direction\n"
		"  -Tsol=x:      delay time of slow-wave\n"
		"   -normax      Normalize using maximum value, otherwise using the difference between maximum and minimum\n"
		"-minPsM1=x      Not stack points when r(Ps) or r(M1) < minPsM1\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.overlays_filename = get_Parameter("-overlays=", string(""), argc, argv);
	argu.output_filename   = get_Parameter("-o=",        string(""), argc, argv);

	struct data_range<double> default_range = { -1, -1 };
	argu.dH    = get_Parameter("-dH=", double(0.1),   argc, argv);
	argu.dK    = get_Parameter("-dK=", double(0.001), argc, argv);
	argu.Hmimx = get_Parameter_range("-Hmimx=", default_range, argc, argv);
	argu.Kmimx = get_Parameter_range("-Kmimx=", default_range, argc, argv);
	argu.Hmimx.max = argu.Hmimx.min + argu.dH * int((argu.Hmimx.max - argu.Hmimx.min)/argu.dH + 0.5);
	argu.Kmimx.max = argu.Kmimx.min + argu.dK * int((argu.Kmimx.max - argu.Kmimx.min)/argu.dK + 0.5);

	string str_ws = get_Parameter("-Ws=", string(""), argc, argv);
	int n = sscanf(str_ws.c_str(), "%lf%lf%lf", argu.Ws, argu.Ws+1, argu.Ws+2);
	if(n < 3) argu.Ws[2] = -0.1;
	if(n < 2) argu.Ws[1] = 0.3;
	if(n < 1) argu.Ws[0] = 0.6;

	argu.vpc   = get_Parameter("-vpc=",   double(6.4), argc, argv);
	argu.vpm   = get_Parameter("-vpm=",   double(8.1), argc, argv);
	argu.dBaz = get_Parameter("-dBaz=", double(5.0), argc, argv);

	argu.Lsol = get_Parameter("-Lsol=", double(-1.0), argc, argv);
	argu.Gsol = get_Parameter("-Gsol=", double(-1.0), argc, argv);
	argu.Fsol = get_Parameter("-Fsol=", double(-1.0), argc, argv);
	argu.Tsol = get_Parameter("-Tsol=", double(-1.0), argc, argv);

	argu.minPsM1 = get_Parameter("-minPsM1=", double(-99.0), argc, argv);
	argu.normax  = get_switch("-normax",   false, 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]);
	SHOWerr2("\t      Hmimx = %g,%g\n", argu.Hmimx.min, argu.Hmimx.max);
	SHOWerr1("\t         dH = %g\n", argu.dH);
	SHOWerr2("\t      Kmimx = %g,%g\n", argu.Kmimx.min, argu.Kmimx.max);
	SHOWerr1("\t         dK = %g\n", argu.dK);
	SHOWerr1("\t         vpc = %g\n", argu.vpc);
	SHOWerr1("\t         vpm = %g\n", argu.vpm);
	SHOWerr1("\t       dBaz = %g\n", argu.dBaz);
	SHOWerr3("\t         ws = %g %g %g\n", argu.Ws[0], argu.Ws[1], argu.Ws[2]);
	SHOWerr1("\t overlayers = %s\n", argu.overlays_filename.c_str());
	SHOWerr1("\t       Lsol = %g\n", argu.Lsol);
	SHOWerr1("\t       Gsol = %g\n", argu.Gsol);
	SHOWerr1("\t       Fsol = %g\n", argu.Fsol);
	SHOWerr1("\t       Tsol = %g\n", argu.Tsol);
	SHOWerr1("\t    minPsM1 = %g\n", argu.minPsM1);
	SHOWerr1("\t normal. max: %s\n", (argu.normax == true)?"Yes":"No");

	return 1;
}

