#pragma once

#include "lbm.h"

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

	D3Q7() : LBM<MM>(7) {}
};

struct D3Q7_M1
{
	static constexpr const char* id = "M1";
	static void setM(matrix &M, matrix &C)
	{
		const int Q=7,D=3;
		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(4,j) = pow(C(0,j),2); 	// c[0]^2
			M(5,j) = pow(C(1,j),2);	 	// c[1]^2
			M(6,j) = pow(C(2,j),2); 	// c[2]^2
		}
	}

	static ex get_moments_eq(int q, FUN& fun, ex sym_cs, matrix &C)
	{
		if (q==0) return fun.rho;
		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[2];
 		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[2]*fun.vel[2] + sym_cs*sym_cs);
 		return 0;
	}

};

struct D3Q7_M2
{
	static constexpr const char* id = "M2";
	static void setM(matrix &M, matrix &C)
	{
		const int Q=7,D=3;
		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(4,j) = pow(C(0,j),2) + pow(C(1,j),2) + pow(C(2,j),2); // m200 + m020 + m002
			M(5,j) = pow(C(0,j),2) - pow(C(1,j),2); 	// m200 - m020
			M(6,j) = pow(C(0,j),2) - pow(C(2,j),2);	 	// m200 - m002
		}
	}

	static ex get_moments_eq(int q, FUN&fun, ex sym_cs, matrix &C)
	{
		const int Q=7;
		// workaround:
		// mu_eq = 
		matrix mu1_eq = matrix(Q,1);
		for (int i=0;i<Q;i++)
			mu1_eq(i,0) = D3Q7_M1::get_moments_eq(i,fun,sym_cs,C);
		matrix M1,M2;
		D3Q7_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 D3Q7_SRT : D3Q7<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 D3Q7_MRT1 : D3Q7<MM>
{
	static constexpr const char* id_col = "mrt1";
	virtual void setM(matrix &M, matrix &C)
	{
		D3Q7_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 D3Q7_MRT2 : D3Q7_MRT1<MM>
{
	static constexpr const char* id_col = "mrt2";
	virtual void setM(matrix &M, matrix &C)
	{
		D3Q7_M2::setM(M,C);
	}
};

template< typename MM>
struct D3Q7_CLBM1 : D3Q7_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 D3Q7_CLBM2 : D3Q7_CLBM1<MM>
{
	static constexpr const char* id_col = "clbm2";
	virtual void setM(matrix &M, matrix &C)
	{
		D3Q7_M2::setM(M,C);
	}
};
