// Lattice Boltzmann Method Analysis Tool
// https://mmg-gitlab.fjfi.cvut.cz/gitlab/lbm/lbmat

#pragma once

#include "base.h"

template< typename MM >
struct LBM
{
	static constexpr const bool is_matrix_lbm = true; // default
	using T_MM = MM;

	int pde_type = NSE; // default=NSE

	int Q=0;
	int D=0;

	matrix M, Mi, E[MQ], I;

	symbol sym_o=symbol("o","\\omega");
	symbol sym_os[MQ];
	symbol sym_cs=symbol("c_s","c_s");
	symbol x=symbol("x","\\x"),y=symbol("y","\\y"),z=symbol("z","\\z"),t=symbol("t");
	symbol dl=symbol("dl","\\dl"), dt=symbol("dt","\\dt");

	// these needs to be specified in subsequent DdQq models
	virtual void setC(matrix &iC) {}
	// virtualize for non-matrix-LBM
	virtual void collision(matrix &rhs, matrix &dfs, ex&u, ex&v, ex&w, ex&rho) {}

	// matrix that defines EPDEs moments - should be independent of the actual LBM model used! virtualize if needed
	virtual void setM(matrix &M)
	{
		matrix C0;
		setC(C0);
		MM::setM(M,C0);
	}

	virtual ex get_moments_eq(int q, FUN& fun)
	{
		matrix C0;
		setC(C0);
		return MM::get_moments_eq(q,fun,sym_cs,C0); // a bit complicated but this is needed for the transform M2.M1i.mu1_eq
	}

	virtual void setA(matrix &A, ex& u, ex&v, ex&w) {}
	virtual void setS(matrix &S)
	{
		S = matrix(Q, Q);
		for (int i=0;i<Q;i++)
		for (int j=0;j<Q;j++)
			S(i,j) = (i==j) ? (ex)sym_os[i] : (ex)0;
	}

	// set CLBM matrix K = N M 
	void setK(matrix &K, ex &u, ex &v, ex &w)
	{
		matrix C0;
		setC(C0);
		// shifted matrix C
		matrix C = matrix(D, Q);
		for (int j=0;j<Q;j++)
		{
			C(0,j) = C0(0,j) - u;
			if (D>=2) C(1,j) = C0(1,j) - v;
			if (D>=3) C(2,j) = C0(2,j) - w;
		}
		MM::setM(K, C);
	}
/*
	void setN(matrix &N, ex &u, ex &v, ex &w)
	{
		matrix M = matrix(Q,Q);
		matrix K = matrix(Q,Q);
		matrix C0;
		setC(C0);
		MM::setM(M,C0);
		setK(K,u,v,w);
		N = K.mul( inverse(M) );
		simplify_matrix(N);
	}
*/
	inline void simplify_matrix(matrix &A, int Nrows, int Ncols) { for (int i=0;i<Nrows;i++) for (int j=0;j<Ncols;j++) A(i,j) = expand(A(i,j)); }
	inline void simplify_matrix(matrix &A) { simplify_matrix(A, Q, Q); }
	inline void empty_matrix(matrix &A) { empty_matrix(A,Q,Q); }
	inline void empty_matrix(matrix &A, int Nrows, int Ncols) { for (int i=0;i<Nrows;i++) for (int j=0;j<Ncols;j++) A(i,j)=0;; }

	int getDim()
	{
		if (Q==3) return 1;
		if (Q==5 || Q==9) return 2;
		return 3;
	}

	virtual ex get_moments(int q, FUN& fun)
	{
		if (q==0) return fun.rho;
		if (pde_type==NSE)
		{
			switch (D)
			{
				case 3:	if (q==3) return fun.vel[2]*fun.rho;
				case 2:	if (q==2) return fun.vel[1]*fun.rho;
				case 1:	if (q==1) return fun.vel[0]*fun.rho;
			}
		}
		return fun.mu[q];
	}

	bool allocated=false;
	void allocate()
	{
		if (Q<=0) printf("error @ lbm_common: Q = %d\n",Q);
		if (allocated) return;
		// velocity vector matrix
		matrix C;
		setC(C);
		// raw moments defining matrix Mtilde
		MM::setM(M,C);
		Mi = inverse(M);
		// row/column selector matrices E
		for (int m=0;m<Q;m++)
		{
			E[m] = matrix(Q,Q);
			for (int i=0;i<Q;i++)
			for (int j=0;j<Q;j++)
				E[m](i,j) = (i==j && j==m) ? 1 : 0;
		}
		// identity matrix
		I = matrix(Q,Q);
		empty_matrix(I);
		for (int i=0;i<Q;i++) I(i,i) = (ex)1;
		allocated=true;
	}

	LBM(int iQ)
	{
		Q=iQ;
		D=getDim();
		char desc[100], desc_tex[100];
		for (int q=0;q<MQ;q++)
		{
			sprintf(desc,"o_%d",q+1);
			sprintf(desc_tex,"\\omega_{%d}",q+1);
			sym_os[q]=symbol(desc,desc_tex);
		}

	}
};

