#pragma once
#include "lbm.h"

template< typename MM>
struct D2Q5 : LBM<MM>
{
	static constexpr const char* id = "d2q5";

	virtual void setC(matrix &C)
	{
		C = matrix(2, 5);
		int maxima_matrix[2][5] = {
			{0,	1,	0,	-1,	0},
			{0,	0,	1,	0,	-1}
		};
		for (int i=0;i<2;i++)
		for (int j=0;j<5;j++)
			C(i,j) = (ex)(maxima_matrix[i][j]);
	}
	D2Q5() : LBM<MM>(5) {}
};

struct D2Q5_M1
{
	static constexpr const char* id = "M1";
	static void setM(matrix &M, matrix &C)
	{
		const int Q=5,D=2;
		M = matrix(Q, Q);
		for (int j=0;j<Q;j++)
		{
			M(0,j) = (ex)1;
			for (int d=0;d<D;d++) M(d+1,j) = C(d,j);
			M(3,j) = pow(C(0,j), 2);
			M(4,j) = 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[0] + sym_cs*sym_cs);
 		if (q==4) return fun.rho*(fun.vel[1]*fun.vel[1] + sym_cs*sym_cs);
		return 0;
	}

};

struct D2Q5_M2
{
	static constexpr const char* id = "M2";
	static void setM(matrix &M, matrix &C)
	{
		const int Q=5,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);
			M(3,j) = pow(C(0,j), 2) + pow(C(1,j), 2);
			M(4,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=5;
		matrix mu1_eq = matrix(Q,1);
		for (int i=0;i<Q;i++)
			mu1_eq(i,0) = D2Q5_M1::get_moments_eq(i,fun,sym_cs,C);
		matrix M1,M2;
		D2Q5_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 D2Q5_SRT : D2Q5<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 D2Q5_MRT1 : D2Q5<MM>
{
	static constexpr const char* id_col = "mrt1";

	virtual void setM(matrix &M, matrix &C)
	{
		D2Q5_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 D2Q5_MRT2 : D2Q5_MRT1<MM>
{
	static constexpr const char* id_col = "mrt2";

	virtual void setM(matrix &M, matrix &C)
	{
		D2Q5_M2::setM(M,C);
	}
};

template< typename MM>
struct D2Q5_CLBM1 : D2Q5_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 D2Q5_CLBM2 : D2Q5_CLBM1<MM>
{
	static constexpr const char* id_col = "clbm2";

	virtual void setM(matrix &M, matrix &C)
	{
		D2Q5_M2::setM(M,C);
	}
};
