#pragma once
#include "lbm.h"

template< typename MM>
struct D2Q9 : LBM<MM>
{
	static constexpr const char* id = "d2q9";
	virtual void setC(matrix &C)
	{
		C = matrix(2, 9);
		int maxima_matrix[2][9] = {
			{0,	1,	0,	-1,	0,	1,	-1,	-1,	1},
			{0,	0,	1,	0,	-1,	1,	1,	-1,	-1}
		};
		for (int i=0;i<2;i++)
		for (int j=0;j<9;j++)
			C(i,j) = (ex)(maxima_matrix[i][j]);
	}
	D2Q9() : LBM<MM>(9) {}
};

struct D2Q9_M1
{
	static constexpr const char* id = "M1";
	static void setM(matrix &M, matrix &C)
	{
		const int Q=9,D=2;
		M = matrix(Q, Q);
		for (int j=0;j<Q;j++)
		{
			// zeroth moment
			 M(0,j) = (ex)1;
			// first moments
			for (int d=0;d<D;d++) M(d+1,j) = C(d,j);
			// the rest of moments
			M(3,j) = C(0,j)*C(1,j);
			M(4,j) = pow(C(0,j), 2);
			M(5,j) = pow(C(1,j), 2);
			M(6,j) = pow(C(0,j), 2)*C(1,j);
			M(7,j) = pow(C(1,j), 2)*C(0,j);
			M(8,j) = pow(C(0,j), 2)*pow(C(1,j), 2);
		}
	}

	static ex get_moments_eq(int q, FUN& fun, ex sym_cs, matrix &C)
	{
		// equilibrium Asinari_PRE_2008bR.pdf P7 Eq 16
		if (q==0) return fun.rho;
		// beware of correct order_eq of moments - Asinari has different 4th, 5th and 6th row!
		if (q==1) return fun.rho*fun.vel[0];
		if (q==2) return fun.rho*fun.vel[1];
 		if (q==3) return fun.rho*(fun.vel[0]*fun.vel[1]);
 		if (q==4) return fun.rho*(fun.vel[0]*fun.vel[0] + sym_cs*sym_cs);
 		if (q==5) return fun.rho*(fun.vel[1]*fun.vel[1] + sym_cs*sym_cs);
 		if (q==6) return fun.rho*(fun.vel[1]*sym_cs*sym_cs + fun.vel[0]*fun.vel[0]*fun.vel[1]);
 		if (q==7) return fun.rho*(fun.vel[0]*sym_cs*sym_cs + fun.vel[0]*fun.vel[1]*fun.vel[1]);
		if (q==8) return fun.rho*(sym_cs*sym_cs*sym_cs*sym_cs + sym_cs*sym_cs*(fun.vel[0]*fun.vel[0] + fun.vel[1]*fun.vel[1]) + fun.vel[0]*fun.vel[0]*fun.vel[1]*fun.vel[1]);
		return 0;
	}

};

struct D2Q9_M2
{
	static constexpr const char* id = "M2";
	static void setM(matrix &M, matrix &C)
	{
		const int Q=9,D=2;
		M = matrix(Q, Q);
		for (int j=0;j<Q;j++)
		{
			// zeroth moment
			 M(0,j) = (ex)1;
			// first moments
			for (int d=0;d<D;d++) M(d+1,j) = C(d,j);
			// dev to see whether we can get the same result as in maxima: 
			// central moment basis k00,k10,k01,k10*k01,k20+k02,k20-k02,k21,k12,k22
			M(3,j) = C(0,j)*C(1,j);
			M(4,j) = pow(C(0,j), 2) + pow(C(1,j), 2);
			M(5,j) = pow(C(0,j), 2) - pow(C(1,j), 2);
			M(6,j) = pow(C(0,j), 2)*C(1,j);
			M(7,j) = pow(C(1,j), 2)*C(0,j);
			M(8,j) = pow(C(0,j), 2)*pow(C(1,j), 2);
		}
	}

	static ex get_moments_eq(int q, FUN&fun, ex sym_cs, matrix &C)
	{
		const int Q=9;
		// workaround:
		// mu_eq = 
		matrix mu1_eq = matrix(Q,1);
		for (int i=0;i<Q;i++)
			mu1_eq(i,0) = D2Q9_M1::get_moments_eq(i,fun,sym_cs,C);
		matrix M1,M2;
		D2Q9_M1::setM(M1,C);
		setM(M2,C);
		matrix M1i = inverse(M1);
		matrix mu2_eq = M2.mul(M1i).mul(mu1_eq);
		return mu2_eq(q,0);
	}
};

struct D2Q9_M3dev
{
	static constexpr const char* id = "M3dev";
	static void setM(matrix &M, matrix &C)
	{
		const int Q=9,D=2;
		M = matrix(Q, Q);
		for (int j=0;j<Q;j++)
		{
			// zeroth moment
			 M(0,j) = (ex)1;
			// first moments
			for (int d=0;d<D;d++) M(d+1,j) = C(d,j);
			for (int d=D+1;d<Q;d++) M(d,j) = 0;
			if (j>D) M(j,j) = (ex)1;
		}
	}

	static ex get_moments_eq(int q, FUN&fun, ex sym_cs, matrix &C)
	{
		const int Q=9;
		// workaround:
		// mu_eq = 
		matrix mu1_eq = matrix(Q,1);
		for (int i=0;i<Q;i++)
			mu1_eq(i,0) = D2Q9_M1::get_moments_eq(i,fun,sym_cs,C);
		matrix M1,M2;
		D2Q9_M1::setM(M1,C);
		setM(M2,C);
		matrix M1i = inverse(M1);
		matrix mu2_eq = M2.mul(M1i).mul(mu1_eq);
		return mu2_eq(q,0);
	}
};

template< typename MM>
struct D2Q9_SRT : D2Q9<MM>
{
	static constexpr const char* id_col = "srt";
	using LBM<MM>::Q;

	virtual void setA(matrix &A, ex& u, ex&v, ex&w)
	{
		A = matrix(Q, Q);
		for (int i=0;i<Q;i++)
		for (int j=0;j<Q;j++)
			A(i,j) = (i==j) ? (ex)LBM<MM>::sym_o : (ex)0;
	}

};

template< typename MM>
struct D2Q9_MRT1 : D2Q9<MM>
{
	static constexpr const char* id_col = "mrt1";
	virtual void setM(matrix &M, matrix &C)
	{
		D2Q9_M1::setM(M,C);
	}
	virtual void setA(matrix &A, ex& u, ex&v, ex&w)
	{
		matrix S,M;
		LBM<MM>::setS(S);
	        LBM<MM>::setM(M);
	        matrix Mi = inverse(M);
	        A = Mi.mul(S).mul(M);
	}
};

template< typename MM>
struct D2Q9_MRT2 : D2Q9_MRT1<MM>
{
	static constexpr const char* id_col = "mrt2";
	virtual void setM(matrix &M, matrix &C)
	{
		D2Q9_M2::setM(M,C);
	}
};

template< typename MM>
struct D2Q9_CLBM1 : D2Q9_MRT1<MM>
{
	static constexpr const char* id_col = "clbm1";
	virtual void setA(matrix &A, ex& u, ex&v, ex&w)
	{
		matrix S,K;
		LBM<MM>::setS(S);
		LBM<MM>::setK(K, u, v, w);
		matrix Ki = inverse(K);
		A = Ki.mul(S).mul(K);
	}
};

template< typename MM>
struct D2Q9_CLBM2 : D2Q9_CLBM1<MM>
{
	static constexpr const char* id_col = "clbm2";
	virtual void setM(matrix &M, matrix &C)
	{
		D2Q9_M2::setM(M,C);
	}
};

template< typename MM>
struct D2Q9_CuLBM1 : D2Q9<MM>
{
    static constexpr const char* id_col = "culbm1";
    static constexpr const bool is_matrix_lbm = false;
    using LBM<MM>::sym_os;
    using LBM<MM>::sym_cs;
    virtual void collision(matrix &rhs, matrix &dfs, ex& u, ex&v, ex&w, ex&rho)
    {
	// converted from lbm: lbm_col_cum.h version 2021.08.21
	// simplified way of using fractions in the CuLBM codes
	const ex no1 = (ex)1;
	const ex n2o3 = 2/(ex)3;
	const ex n1o2 = 1/(ex)2;
    const ex n1o3 = 1/(ex)3;
	// following the defintion of C for D3Q27
	//      C = matrix(D, Q);
	//     int maxima_matrix[2][9] = {
	//         {0,	1,	0,	-1,	0,	1,	-1,	-1,	1},
	//         {0,	0,	1,	0,	-1,	1,	1,	-1,	-1}
	//     };

	enum{zz,pz,zp,mz,zm,pp,mp,mm,pm};
	//Robert code indexation
	//double cx[F]={0,-1, 0,1,0,1,-1,-1, 1};
	//double cy[F]=  {0, 0,-1,0,1,1, 1,-1,-1};

	ex ff0 = dfs(zz, 0);
	ex ff1 = dfs(pz, 0);
	ex ff2 = dfs(zp, 0);
	ex ff3 = dfs(mz, 0);
	ex ff4 = dfs(zm, 0);
	ex ff5 = dfs(pp, 0);
	ex ff6 = dfs(mp, 0);
	ex ff7 = dfs(mm, 0);
	ex ff8 = dfs(pm, 0);

	ff0 = ff3 + ff1 + ff0; 
	ff1 = -ff3 + ff1;
	ff3 = ff1 + 2*ff3; 
	ff2 = ff6 + ff5 + ff2;
	ff5 = -ff6 + ff5;
	ff6 = ff5 + 2*ff6; 
	ff4 = ff7 + ff8 + ff4;
	ff8 = -ff7 + ff8;
	ff7 = ff8 + 2*ff7; 
	ff0 = ff4 + ff2 + ff0;
	ff2 = -ff4 + ff2;
	ff4 = ff2 + 2*ff4; 
	ff1 = ff8 + ff5 + ff1;
	ff5 = -ff8 + ff5;
	ff8 = ff5 + 2*ff8; 
	ff3 = ff7 + ff6 + ff3;
	ff6 = -ff7 + ff6;
	ff7 = ff6 + 2*ff7; 

	//Cumulant calculation from moments
	ex cu0 = ff0;
	ex cu1 = ff1/ff0;
	ex cu3 = ( -cu1*ff1 + ff3 )/ff0;
	ex cu2 = ff2/ff0;
	ex cu5 = ( -cu1*ff2 + ff5 )/ff0;
	ex cu6 = ( -cu5*ff1 - cu3*ff2 - cu1*ff5 + ff6 )/ff0;
	ex cu4 = ( -cu2*ff2 + ff4 )/ff0;
	ex cu8 = ( -cu1*ff4 + ff8 - 2*cu5*ff2)/ff0;
	ex cu7 = ( -cu8*ff1 - cu3*ff4 - cu1*ff8 + ff7 + 2*( -cu6*ff2 - cu5*ff5 ))/ff0;

	//Cumulant relaxation:
	ex a = (cu3 + cu4);
	ex b = (cu3 - cu4);

	ex w0 = sym_os[0];
	ex w1 = sym_os[1];
	ex w2 = sym_os[2];//((ex)2 - w0)/((ex)1 - w0/(ex)8);
	ex w3 = sym_os[3];
    ex w4 = sym_os[4];
    ex w5 = sym_os[5];

	cu3 = (no1 - w0)*cu3 + w0*sym_cs*sym_cs;
	cu4 = (no1 - w1)*cu4 + w1*sym_cs*sym_cs;
	cu5 = (no1 - w2)*cu5;
	cu6 = (no1 - w3)*cu6;
	cu7 = (no1 - w4)*cu7;
	cu8 = (no1 - w5)*cu8; 

	// Moment calculation from cummulants

	ff0 = ff0;
	ff1 = cu1*ff0;
	ff3 = cu3*ff0 + cu1*ff1;
	ff2 = cu2*ff0;
	ff5 = cu5*ff0 + cu1*ff2;
	ff6 = cu6*ff0 + cu5*ff1 + cu3*ff2 + cu1*ff5;
	ff4 = cu4*ff0 + cu2*ff2;
	ff8 = cu8*ff0 + cu1*ff4 + 2*cu5*ff2;
	ff7 = cu7*ff0 + cu8*ff1 + cu3*ff4 + cu1*ff8 + 2*( cu6*ff2 + cu5*ff5 );

	//Transformation from moment to density distribution function

	ff0 = -ff3 + ff0;
	ff1 = ( ff3 + ff1 )*n1o2;
	ff3 = ff3 - ff1; 
	ff2 = -ff6 + ff2;
	ff5 = ( ff6 + ff5 )*n1o2;
	ff6 = ff6 - ff5; 
	ff4 = -ff7 + ff4;
	ff8 = ( ff7 + ff8 )*n1o2;
	ff7 = ff7 - ff8; 
	ff0 = -ff4 + ff0;
	ff2 = ( ff4 + ff2 )*n1o2;
	ff4 = ff4 - ff2; 
	ff1 = -ff8 + ff1;
	ff5 = ( ff8 + ff5 )*n1o2;
	ff8 = ff8 - ff5; 
	ff3 = -ff7 + ff3;
	ff6 = ( ff7 + ff6 )*n1o2;
	ff7 = ff7 - ff6;

	rhs(zz, 0) = ff0;
	rhs(pz, 0) = ff1;
	rhs(zp, 0) = ff2;
	rhs(mz, 0) = ff3;
	rhs(zm, 0) = ff4;
	rhs(pp, 0) = ff5;
	rhs(mp, 0) = ff6;
	rhs(mm, 0) = ff7;
	rhs(pm, 0) = ff8;

    }
};

template< typename MM>
struct D2Q9_CuLBM2 : D2Q9<MM>
{
    static constexpr const char* id_col = "culbm2";
    static constexpr const bool is_matrix_lbm = false;
    using LBM<MM>::sym_os;
    using LBM<MM>::sym_cs;
    virtual void collision(matrix &rhs, matrix &dfs, ex& u, ex&v, ex&w, ex&rho)
    {
	// converted from lbm: lbm_col_cum.h version 2021.08.21
	// simplified way of using fractions in the CuLBM codes
	const ex no1 = (ex)1;
	const ex n2o3 = 2/(ex)3;
	const ex n1o2 = 1/(ex)2;
	// following the defintion of C for D3Q27
	//      C = matrix(D, Q);
	//     int maxima_matrix[2][9] = {
	//         {0,	1,	0,	-1,	0,	1,	-1,	-1,	1},
	//         {0,	0,	1,	0,	-1,	1,	1,	-1,	-1}
	//     };

	enum{zz,pz,zp,mz,zm,pp,mp,mm,pm};
	//Robert code indexation
	//double cx[F]={0,-1, 0,1,0,1,-1,-1, 1};
	//double cy[F]=  {0, 0,-1,0,1,1, 1,-1,-1};

	ex ff0 = dfs(zz, 0);
	ex ff1 = dfs(pz, 0);
	ex ff2 = dfs(zp, 0);
	ex ff3 = dfs(mz, 0);
	ex ff4 = dfs(zm, 0);
	ex ff5 = dfs(pp, 0);
	ex ff6 = dfs(mp, 0);
	ex ff7 = dfs(mm, 0);
	ex ff8 = dfs(pm, 0);

	ff0 = ff3 + ff1 + ff0; 
	ff1 = -ff3 + ff1;
	ff3 = ff1 + 2*ff3; 
	ff2 = ff6 + ff5 + ff2;
	ff5 = -ff6 + ff5;
	ff6 = ff5 + 2*ff6; 
	ff4 = ff7 + ff8 + ff4;
	ff8 = -ff7 + ff8;
	ff7 = ff8 + 2*ff7; 
	ff0 = ff4 + ff2 + ff0;
	ff2 = -ff4 + ff2;
	ff4 = ff2 + 2*ff4; 
	ff1 = ff8 + ff5 + ff1;
	ff5 = -ff8 + ff5;
	ff8 = ff5 + 2*ff8; 
	ff3 = ff7 + ff6 + ff3;
	ff6 = -ff7 + ff6;
	ff7 = ff6 + 2*ff7; 

	//Cumulant calculation from moments
	ex cu0 = ff0;
	ex cu1 = ff1/ff0;
	ex cu3 = ( -cu1*ff1 + ff3 )/ff0;
	ex cu2 = ff2/ff0;
	ex cu5 = ( -cu1*ff2 + ff5 )/ff0;
	ex cu6 = ( -cu5*ff1 - cu3*ff2 - cu1*ff5 + ff6 )/ff0;
	ex cu4 = ( -cu2*ff2 + ff4 )/ff0;
	ex cu8 = ( -cu1*ff4 + ff8 - 2*cu5*ff2)/ff0;
	ex cu7 = ( -cu8*ff1 - cu3*ff4 - cu1*ff8 + ff7 + 2*( -cu6*ff2 - cu5*ff5 ))/ff0;

	//Cumulant relaxation:
	ex a = (cu3 + cu4);
	ex b = (cu3 - cu4);

	ex w0 = sym_os[0];
	ex w1 = sym_os[1];
	ex w2 = sym_os[2];//((ex)2 - w0)/((ex)1 - w0/(ex)8);
	ex w3 = sym_os[3];

	cu3 = ((no1 - w1)*a + w1*2*sym_cs*sym_cs + (1 - w0)*b)*n1o2;
	cu4 = ((no1 - w1)*a + w1*2*sym_cs*sym_cs - (1 - w0)*b)*n1o2;
	cu5 = (no1 - w0)*cu5;
	cu6 = (no1 - w2)*cu6;
	cu7 = (no1 - w3)*cu7;
	cu8 = (no1 - w2)*cu8; 

	// Moment calculation from cummulants

	ff0 = ff0;
	ff1 = cu1*ff0;
	ff3 = cu3*ff0 + cu1*ff1;
	ff2 = cu2*ff0;
	ff5 = cu5*ff0 + cu1*ff2;
	ff6 = cu6*ff0 + cu5*ff1 + cu3*ff2 + cu1*ff5;
	ff4 = cu4*ff0 + cu2*ff2;
	ff8 = cu8*ff0 + cu1*ff4 + 2*cu5*ff2;
	ff7 = cu7*ff0 + cu8*ff1 + cu3*ff4 + cu1*ff8 + 2*( cu6*ff2 + cu5*ff5 );

	//Transformation from moment to density distribution function

	ff0 = -ff3 + ff0;
	ff1 = ( ff3 + ff1 )*n1o2;
	ff3 = ff3 - ff1; 
	ff2 = -ff6 + ff2;
	ff5 = ( ff6 + ff5 )*n1o2;
	ff6 = ff6 - ff5; 
	ff4 = -ff7 + ff4;
	ff8 = ( ff7 + ff8 )*n1o2;
	ff7 = ff7 - ff8; 
	ff0 = -ff4 + ff0;
	ff2 = ( ff4 + ff2 )*n1o2;
	ff4 = ff4 - ff2; 
	ff1 = -ff8 + ff1;
	ff5 = ( ff8 + ff5 )*n1o2;
	ff8 = ff8 - ff5; 
	ff3 = -ff7 + ff3;
	ff6 = ( ff7 + ff6 )*n1o2;
	ff7 = ff7 - ff6;

	rhs(zz, 0) = ff0;
	rhs(pz, 0) = ff1;
	rhs(zp, 0) = ff2;
	rhs(mz, 0) = ff3;
	rhs(zm, 0) = ff4;
	rhs(pp, 0) = ff5;
	rhs(mp, 0) = ff6;
	rhs(mm, 0) = ff7;
	rhs(pm, 0) = ff8;

    }
};
