#include "base.h"
#include "core.h"
#include "d1q3.h"

struct LBM_local : LBM<D1Q3_M1> // any MM matrix works here
{
	FUN fun;
	LBM_local() : LBM(27)
	{
	        fun.rho = frho(t, x, y, z);
	        fun.vel[0]=fu(t, x, y, z);
	        fun.vel[1]=fv(t, x, y, z);
	        fun.vel[2]=fw(t, x, y, z);
	}
};

void export_latex_compare(const char *dirname_prefix, int MO=4, int MFCO=2, bool dohyper=true)
{
	// MO and MFCO needs to be specified as well as basis functions
	LBM_local lbm;
	BASE B;
	char subdir[200];
	char dirname[1000];
	symtab table;
	table[lbm.sym_o.get_name()] = lbm.sym_o;
	table[lbm.dt.get_name()] = lbm.dt;
	table[lbm.dl.get_name()] = lbm.dl;
	table[lbm.sym_cs.get_name()] = lbm.sym_cs;
//	table[lbm.sym_cs.get_name()] = sqrt((ex)1/(ex)3);//lbm.sym_cs;
	table[lbm.t.get_name()] = lbm.t;
	table[lbm.x.get_name()] = lbm.x;
	table[lbm.y.get_name()] = lbm.y;
	table[lbm.z.get_name()] = lbm.z;
	// velocity and density
	table["rho"] = lbm.fun.rho;
	table["u"] = lbm.fun.vel[0];
	table["v"] = lbm.fun.vel[1];
	table["w"] = lbm.fun.vel[2];
	table["D11"] = symbol("D11","{\\color{red}{D_{11}}}");
	table["D12"] = symbol("D12","{\\color{red}{D_{12}}}");
	table["D13"] = symbol("D13","{\\color{red}{D_{13}}}");
	table["D21"] = symbol("D21","{\\color{red}{D_{21}}}");
	table["D22"] = symbol("D22","{\\color{red}{D_{22}}}");
	table["D23"] = symbol("D23","{\\color{red}{D_{23}}}");
	table["D31"] = symbol("D31","{\\color{red}{D_{31}}}");
	table["D32"] = symbol("D32","{\\color{red}{D_{32}}}");
	table["D33"] = symbol("D33","{\\color{red}{D_{33}}}");
	table["zeta"] = symbol("zeta","{\\color{red}{\\zeta}}");
	table["nu"] = symbol("nu","{\\color{red}{\\nu}}");

	for (int i=0;i<MQ;i++)
		table[lbm.sym_os[i].get_name()] = lbm.sym_os[i];

	vct bas;
	bas.push_back( lbm.fun.rho );
	bas.push_back( lbm.fun.vel[0] );
	bas.push_back( lbm.fun.vel[1] );
	bas.push_back( lbm.fun.vel[2] );
	int Nbas=bas.size();

	// fill names
	const int ModMAX=9;
	char modname[ModMAX][20];
	char modlatex[ModMAX][20];
	for (int model=0;model<ModMAX;model++)
	switch (model)
	{
		case 0:	sprintf(modname[model],"srt"); sprintf(modlatex[model],"SRT"); break; // ignore SRT
		case 1:	sprintf(modname[model],"mrt1"); sprintf(modlatex[model],"MRT1"); break;
		case 2:	sprintf(modname[model],"mrt2"); sprintf(modlatex[model],"MRT2"); break;
		case 3:	sprintf(modname[model],"clbm1"); sprintf(modlatex[model],"CLBM1"); break;
		case 4:	sprintf(modname[model],"clbm2"); sprintf(modlatex[model],"CLBM2"); break;
		case 5:	sprintf(modname[model],"culbm1"); sprintf(modlatex[model],"CuLBM1"); break;
		case 6:	sprintf(modname[model],"culbm2"); sprintf(modlatex[model],"CuLBM2"); break;
		case 7:	sprintf(modname[model],"culbm"); sprintf(modlatex[model],"CuLBM"); break;
		case 8:	sprintf(modname[model],"special"); sprintf(modlatex[model],"\\RAWMODEL"); break;
		default: sprintf(modname[model],"nothing");
	}

	// browse pdes
	for (int pi=0;pi<4;pi++) // maximum is D+1 for D=3
	{
		TPDE S[ModMAX]; // max pdes
		bool exists[ModMAX];
		int firstmod=-1;	// the index of the first model that exists


		for (int model=0;model<ModMAX;model++)
		{
			exists[model]=false;
			// symbolic
			sprintf(dirname,"%s_%s_symbolic", dirname_prefix, modname[model]);
			if (B.file_exists(dirname))
			{
				sprintf(subdir,"%s/pde_%02d",dirname,pi);
				if (B.file_exists(subdir))
				{
					exists[model]=true;
					if (firstmod<0) firstmod=model;
					cout << "pde "<<pi<<" found "<<subdir<<endl;
					S[model].allocate(bas, MO, MFCO, lbm.t, lbm.x, lbm.y, lbm.z);
					S[model].load(subdir, table);
//					S[model].save_tex2(dirname, "spatial_EPDE","pde_%02d", pi);
//					S[model].save_raw(dirname, "spatial_EPDE","pde_%02d", pi);
				}
			}
			// raw
			sprintf(dirname,"%s_%s", dirname_prefix, modname[model]);
			if (!exists[model])
			if (B.file_exists(dirname))
			{
				sprintf(subdir,"%s/pde_%02d",dirname,pi);
				if (B.file_exists(subdir))
				{
					exists[model]=true;
					if (firstmod<0) firstmod=model;
					cout << "pde "<<pi<<" found "<<subdir<<endl;
					S[model].allocate(bas, MO, MFCO, lbm.t, lbm.x, lbm.y, lbm.z);
					S[model].load(subdir, table);
				}
			}
		}

		if (firstmod<0) { continue; } // no directory or pde found 

		/**
		// determine dim if needed
		int DIM=1;
		for (int ok=1;ok<=2*MO;ok++)
		for (int oi=1;oi<=MO;oi++) for (int i3=0;i3<=oi;i3++)  for (int i2=0;i2<=oi;i2++) for (int i1=0;i1<=oi;i1++) for (int i0=0;i0<=oi;i0++) if (i0+i1+i2+i3==oi)
		{
			// 1D terms
			if (oi==ok)
			for (int bi=0;bi<Nbas;bi++)
			{
				if (S[firstmod].data_r(bi,i0,i1,i2,i3)!=0) 
				{
//					if (DIM<1 && i1>0) DIM=1;
					if (DIM<2 && i2>0) DIM=2;
					if (DIM<3 && i3>0) DIM=3;
				}
			}
			// 2D terms
			for (int oj=1;oj<=MO;oj++) for (int j3=0;j3<=oj;j3++) for (int j2=0;j2<=oj;j2++) for (int j1=0;j1<=oj;j1++) for (int j0=0;j0<=oj;j0++) if (j0+j1+j2+j3==oj)
			if (oi+oj==ok)
			if (oj+oi<=MFCO)
			for (int bj=0;bj<Nbas;bj++)
			for (int bi=0;bi<Nbas;bi++)
			{
				if (S[firstmod].data_r(bi,i0,i1,i2,i3, bj,j0,j1,j2,j3)!=0)
				{
//					if (DIM<1 && i1>0) DIM=1;
					if (DIM<2 && i2>0) DIM=2;
					if (DIM<3 && i3>0) DIM=3;
//					if (DIM<1 && j1>0) DIM=1;
					if (DIM<2 && j2>0) DIM=2;
					if (DIM<3 && j3>0) DIM=3;
				}
			}
		}
		// what a stupid :) way of determining the dimension of the problem ... DIM=1, 2, 3 ... sorry
		/**/

		// prepare compare tex file
		char expdir[1000];
//		sprintf(expdir,"compare_%s_%s_%d_%d",dirname_prefix,dirname_suffix,MO,MFCO);
		sprintf(expdir,"compare_%s",dirname_prefix);
		mkdir(expdir,0777);

		char coefid[1000];
		sprintf(coefid,"compare_%d_%s_",pi,expdir);

		char name[1024];
		sprintf(name,"%s/pde_%02d.tex",expdir,pi);
		ostringstream sout1, sout2; // sout1 stores the main equation, sout2 is to store extra substitutions
		sout1 << latex;
		sout2 << latex;
		// header
		sout1<< "\\begin{flushleft}" << endl;
		sout1<< "\\noindent\n";

		ostringstream sout_coef_exp, sout_coef_exp_c;
		sout_coef_exp << latex; // exponent to distinguish coefficients
		sout_coef_exp << "("<<(pi)<<")\\modelq";
		sout_coef_exp_c << latex; // exponent to distinguish coefficients
		sout_coef_exp_c << "("<<(pi)<<"),\\modelqc";

		bool first=true;
		int counter=1;
		for (int ok=1;ok<=2*MO;ok++)
		for (int oi=1;oi<=MO;oi++) for (int i3=0;i3<=oi;i3++)  for (int i2=0;i2<=oi;i2++) for (int i1=0;i1<=oi;i1++) for (int i0=0;i0<=oi;i0++) if (i0+i1+i2+i3==oi)
		{

			// 1D terms
			if (oi==ok)
			for (int bi=0;bi<Nbas;bi++)
			{
				ex co = S[firstmod].data_r(bi,i0,i1,i2,i3);
				// compare coefficients
				bool thesame=true;
				for (int mod=firstmod+1;mod<ModMAX;mod++) 
					if (exists[mod])
						if (co != S[mod].data_r(bi,i0,i1,i2,i3)) thesame=false;

				ex dldt = pow(lbm.dt, i0-1)*pow(lbm.dl,i1+i2+i3);
				if (!thesame)
				{
					ostringstream sout_coef, sout_expr;
					sout_coef << latex;
					sout_expr << latex;
//					sout_coef << "C_{"<<counter<<"}";
					sout_coef << "C_{"<<S[firstmod].derivative_tex_short(bi, i0, i1, i2, i3)<<"}";
//					sout_coef << "C_{"<<counter<<","<<S[firstmod].derivative_tex_short(bi, i0, i1, i2, i3)<<"}";
//					sout_coef << "C_{"<<S[firstmod].derivative_tex_multiindex(bi, i0, i1, i2, i3, DIM)<<"}";
//					if (DIM==1) sout_coef << "C_{\\D_{("<<i0<<","<<i1<<")}"<<bas[bi]<<"}";
//					if (DIM==2) sout_coef << "C_{\\D_{("<<i0<<","<<i1<<","<<i2<<")}"<<bas[bi]<<"}";
//					if (DIM==3) sout_coef << "C_{\\D_{("<<i0<<","<<i1<<","<<i2<<","<<i3<<")}"<<bas[bi]<<"}";

					if (first) first=false; else sout1 << "$+$\n";
					sout1 << "${\\color{darkgreen}";
					if (dohyper) sout1 << "\\hyperref[C"<<coefid<<counter<<"]";
					sout1 << "{"<<sout_coef.str()<<"^{"<<sout_coef_exp.str()<<"}}}";
					if (dldt!=1) sout1 << dldt;
					sout_expr << "{\\color{blue}"<< S[firstmod].derivative_tex(bi, i0, i1, i2, i3) << "}";
					sout1 << sout_expr.str() << "$\n";

					if (dohyper) sout1 << "{\\phantomsection\\label{E"<<coefid<<counter<<"}} \n";

					if (dohyper) sout2 << "\\phantomsection\\label{C"<<coefid<<counter<<"}\n";
					sout2 << "\n\\noindent\\textbf{coefficient ${\\color{darkgreen}"<<sout_coef.str()<<"^{"<<sout_coef_exp.str()<<"}}$ at ";
					if (dohyper) sout2 << "\\hyperref[E"<<coefid<<counter<<"]";
					sout2 << "{$"<<sout_expr.str()<<"$}:";
					sout2 << "}\n\\vskip 1em\n\n";
					for (int mod=firstmod;mod<ModMAX;mod++) if (exists[mod])
					{
						ex co2 = S[mod].data_r(bi,i0,i1,i2,i3);
						sout2 << "${\\color{darkgreen}"<<sout_coef.str()<<"^{"<<sout_coef_exp_c.str()<<"\\mathrm{"<< modlatex[mod] <<"}}} = \n ";
						// compare to previous 
						bool prevmodfound=false;
						if (co2!=0)
						for (int prevmod=firstmod;prevmod<mod;prevmod++)
						{
							if (exists[prevmod])
							if (co2 == S[prevmod].data_r(bi,i0,i1,i2,i3))
							{
								prevmodfound=true;
								sout2 << sout_coef.str()<< "^{"<<sout_coef_exp_c.str()<<"\\mathrm{"<< modlatex[prevmod] <<"}} \n ";
								break;
							}
						}

						if (!prevmodfound)
						{
							co2 = normal(expand(co2/dldt));
							if (co2!=1) S[firstmod].factorize_latex_simple(sout2,co2);
						}
						sout2 << "$\n\\vskip 1em\n\n";
					}
					counter++;
				} else
				{
					if (co!=0)
					{
						if (!is_a<GiNaC::add>(co) && B.isnegative(co))
						{
							co*=-1;
							sout1 << "$-$\n";
							if (first) first=false;
						} else
						{
							if (first) first=false; else sout1 << "$+$\n";
						}
						if (dldt!=1) co = normal(expand(co/dldt));
						sout1 << "$";
						if (co!=1) S[firstmod].factorize_latex_simple(sout1,co);
						if (dldt!=1) sout1 << dldt;
						sout1 << "{\\color{blue}"<< S[firstmod].derivative_tex(bi, i0, i1, i2, i3) << "}$\n";
					}
				}
			}

			// 2D terms
			for (int oj=1;oj<=MO;oj++) for (int j3=0;j3<=oj;j3++) for (int j2=0;j2<=oj;j2++) for (int j1=0;j1<=oj;j1++) for (int j0=0;j0<=oj;j0++) if (j0+j1+j2+j3==oj)
			if (oi+oj==ok)
			if (oj+oi<=MFCO)
			for (int bj=0;bj<Nbas;bj++)
			for (int bi=0;bi<Nbas;bi++)
			{
				ex co = S[firstmod].data_r(bi,i0,i1,i2,i3, bj,j0,j1,j2,j3);
				// compare coefficients
				bool thesame=true;
				for (int mod=firstmod+1;mod<ModMAX;mod++) if (exists[mod])
					if (co != S[mod].data_r(bi,i0,i1,i2,i3, bj,j0,j1,j2,j3)) thesame=false;

				ex dldt = pow(lbm.dt, j0+i0-1)*pow(lbm.dl,j1+j2+j3+i1+i2+i3);

				if (!thesame)
				{
					ostringstream sout_coef, sout_expr;
					sout_coef << latex;
					sout_expr << latex;
//					sout_coef << "C_{"<<counter<<"}";
					sout_coef << "C_{"<<S[firstmod].derivative_tex_short(bi, i0, i1, i2, i3)<<","<<S[firstmod].derivative_tex_short(bj, j0, j1, j2, j3)<<"}";
//					sout_coef << "C_{"<<counter<<","<<S[firstmod].derivative_tex_short(bi, i0, i1, i2, i3)<<","<<S[firstmod].derivative_tex_short(bj, j0, j1, j2, j3)<<"}";
//					sout_coef << "C_{"<<S[firstmod].derivative_tex_multiindex(bi, i0, i1, i2, i3, DIM)<<","<<S[firstmod].derivative_tex_multiindex(bj, j0, j1, j2, j3, DIM)<<"}";
//					if (DIM==1) sout_coef << "C_{\\D_{("<<i0<<","<<i1<<")}"<<bas[bi]<<","<< "\\D_{("<<j0<<","<<j1<<")}"<<bas[bj]<<"}";
//					if (DIM==2) sout_coef << "C_{\\D_{("<<i0<<","<<i1<<","<<i2<<")}"<<bas[bi]<<","<< "\\D_{("<<j0<<","<<j1<<","<<j2<<")}"<<bas[bj]<<"}";
//					if (DIM==3) sout_coef << "C_{\\D_{("<<i0<<","<<i1<<","<<i2<<","<<i3<<")}"<<bas[bi]<<","<< "\\D_{("<<j0<<","<<j1<<","<<j2<<","<<j3<<")}"<<bas[bj]<<"}";

					if (first) first=false; else sout1 << "$+$\n";
					sout1 << "${\\color{darkgreen}";
					if (dohyper) sout1 << "\\hyperref[C"<<coefid<<counter<<"]";
					sout1 << "{"<<sout_coef.str()<<"^{"<<sout_coef_exp.str()<<"}}}";
					if (dldt!=1) sout1 << dldt;

					if (S[firstmod].ind(bi,i0,i1,i2,i3)==S[firstmod].ind(bj,j0,j1,j2,j3))
						sout_expr << "\\left({\\color{blue}"<<S[firstmod].derivative_tex(bj, j0, j1, j2, j3)<< "}\\right)^2";
					else
						sout_expr << "{\\color{blue}"<< S[firstmod].derivative_tex(bi, i0, i1, i2, i3) << " "<< S[firstmod].derivative_tex(bj, j0, j1, j2, j3)<< "}";

					sout1 << sout_expr.str() << "$\n";

					if (dohyper) sout1 << "{\\phantomsection\\label{E"<<coefid<<counter<<"}} \n";

					if (dohyper) sout2 << "\\phantomsection\\label{C"<<coefid<<counter<<"}\n";
					sout2 << "\n\\noindent\\textbf{coefficient ${\\color{darkgreen}"<<sout_coef.str()<<"^{"<<sout_coef_exp.str()<<"}}$ at ";
					if (dohyper) sout2 << "\\hyperref[E"<<coefid<<counter<<"]";
					sout2 << "{$"<<sout_expr.str()<<"$}:";
					sout2 << "}\n\\vskip 1em\n\n";

					for (int mod=firstmod;mod<ModMAX;mod++) if (exists[mod])
					{
						ex co2 = S[mod].data_r(bi,i0,i1,i2,i3, bj,j0,j1,j2,j3);
						sout2 << "${\\color{darkgreen}"<<sout_coef.str()<<"^{"<<sout_coef_exp_c.str()<<"\\mathrm{"<< modlatex[mod] <<"}}} = \n ";

						// compare to previous 
						bool prevmodfound=false;
						if (co2!=0)
						for (int prevmod=firstmod;prevmod<mod;prevmod++)
							if (exists[prevmod])
							if (co2 == S[prevmod].data_r(bi,i0,i1,i2,i3, bj,j0,j1,j2,j3))
							{
								prevmodfound=true;
								sout2 << sout_coef.str()<<"^{"<<sout_coef_exp_c.str()<<"\\mathrm{"<< modlatex[prevmod] <<"}} \n ";
								break;
							}

						if (!prevmodfound)
						{
							co2 = normal(expand(co2/dldt));
							if (co2!=1) S[firstmod].factorize_latex_simple(sout2,co2);
						}
						sout2 << "$\n\\vskip 1em\n\n";
					}
					counter++;
				} else
				{
					if (co!=0)
					{
						if (!is_a<GiNaC::add>(co) && B.isnegative(co))
						{
							co*=-1;
							sout1 << "$-$\n";
							if (first) first=false;
						} else
						{
							if (first) first=false; else sout1 << "$+$\n";
						}

						// extract dl/dt from co
						// ex dldt = (co!=0) ? normal(expand(co / subs( subs( co, lbm.dt == 1 ), lbm.dl == 1) )) : (ex)1;
						if (dldt!=1) co=normal(expand(co/dldt));

//						if (co!=1) S[firstmod].factorize_latex(sout1,sout2,co,counter);
						sout1 << "$";
						if (co!=1) S[firstmod].factorize_latex_simple(sout1,co);
						if (dldt!=1) sout1 << dldt;
						if (S[firstmod].ind(bi,i0,i1,i2,i3)==S[firstmod].ind(bj,j0,j1,j2,j3))
							sout1 << "\\left({\\color{blue}"<<S[firstmod].derivative_tex(bj, j0, j1, j2, j3)<< "}\\right)^2";
						else
							sout1 << "{\\color{blue}"<< S[firstmod].derivative_tex(bi, i0, i1, i2, i3) << " "<< S[firstmod].derivative_tex(bj, j0, j1, j2, j3)<< "}";
						sout1 << "$\n";
					}
				}
			}
		}

		if (sout2.str().length()==0)
			sout1<< "$=0.$\n";
		else
			sout1<< "$=0,$\n\\vskip 1em\nwhere:";
		sout1<<"\\end{flushleft}" << endl;

		ofstream fout(name);
		fout << latex;
		fout << sout1.str() << endl;
		// put coefs definitions if any
		if (sout2.str().length()>0)
		{
			fout << "\\begin{flushleft}" << endl;
			fout << "\\scriptsize"<<endl;
			fout << "\\noindent\n";
			fout << sout2.str();
			fout << "\\end{flushleft}" << endl;
		}
		fout.close();
	}

}


int main(int argc, char* argv[])
{
	auto start = chrono::steady_clock::now();

	if (argc<2) { cerr << "directory prefix missing" << endl; return 0; }

	export_latex_compare(argv[1]);

 	auto end = chrono::steady_clock::now();
	cout<<endl;
	cout<<chrono::duration_cast<chrono::milliseconds>(end-start).count()<<" milliseconds"<<endl;
}
