#include "PajaTypes.h"
#include "rfftw.h"
#include "FluidSolverC.h"


using namespace PajaTypes;




bool			FluidSolverC::m_bInitialised = false;
rfftwnd_plan	FluidSolverC::m_rPlanRC;
rfftwnd_plan	FluidSolverC::m_rPlanCR;


FluidSolverC::FluidSolverC() :
	m_f32Time( 0 ),
	m_i32Dim( 0 ),
	m_pU( 0 ),
	m_pV( 0 ),
	m_pU0( 0 ),
	m_pV0( 0 ),
	m_pRho( 0 ),
	m_pRhoSmooth( 0 ),
	m_pRho0( 0 ),
	m_pU_U0( 0 ),
	m_pU_V0( 0 )
{
	// empty
}

FluidSolverC::~FluidSolverC()
{
	delete [] m_pU;
	delete [] m_pV;
	delete [] m_pU0;
	delete [] m_pV0;
	delete [] m_pRho;
	delete [] m_pRhoSmooth;
	delete [] m_pRho0;
	delete [] m_pU_U0;
	delete [] m_pU_V0;
}

void
FluidSolverC::init( int32 i32Dim )
{
	m_i32Dim = i32Dim;

	m_pU = new fftw_real[m_i32Dim * 2 * (m_i32Dim / 2 + 1)];
	m_pV = new fftw_real[m_i32Dim * 2 * (m_i32Dim / 2 + 1)];
	m_pU0 = new fftw_real[m_i32Dim * 2 * (m_i32Dim / 2 + 1)];
	m_pV0 = new fftw_real[m_i32Dim * 2 * (m_i32Dim / 2 + 1)];

	m_pRho = new fftw_real[m_i32Dim * m_i32Dim];
	m_pRhoSmooth = new fftw_real[m_i32Dim * m_i32Dim];
	m_pRho0 = new fftw_real[m_i32Dim * m_i32Dim];

	m_pU_U0 = new fftw_real[m_i32Dim * m_i32Dim];
	m_pU_V0 = new fftw_real[m_i32Dim * m_i32Dim];


	if( !m_bInitialised ) {
		m_rPlanRC = rfftw2d_create_plan( m_i32Dim, m_i32Dim, FFTW_REAL_TO_COMPLEX, FFTW_IN_PLACE);
		m_rPlanCR = rfftw2d_create_plan( m_i32Dim, m_i32Dim, FFTW_COMPLEX_TO_REAL, FFTW_IN_PLACE);
		m_bInitialised = true;
	}

	reset();
}

void
FluidSolverC::reset()
{
	for( uint32 i = 0; i < m_i32Dim * m_i32Dim; i++ )
		m_pU[i] = m_pV[i] = m_pU0[i] = m_pV0[i] = m_pRho[i] = m_pRho0[i] = m_pU_U0[i] = m_pU_V0[i] = 0.0f;

	m_f32Time = 0;
}

void
FluidSolverC::solve( float32 f32DT )
{
	// apply force
	set_forces();
	stable_solve( m_i32Dim, m_pU, m_pV, m_pU0, m_pV0, 0.001, f32DT );
//	diffuse_matter( m_i32Dim, m_pU, m_pV, m_pRho, m_pRho0, f32DT );
	m_f32Time += f32DT;
}

void
FluidSolverC::drag( float32 f32X, float32 f32Y, Vector2C rForce, PajaTypes::float32 f32Matter )
{
	f32X *= (float32)m_i32Dim;
	f32Y *= (float32)m_i32Dim;

	int32	i32X = floor( f32X );
	int32	i32Y = floor( f32Y );
	int32	i32NextX, i32NextY;

	float32	f32AX = f32X - (float32)i32X;
	float32	f32AY = f32Y - (float32)i32Y;

	if( i32X < 0 )
		i32X = 0;
	if( i32X >= m_i32Dim )
		i32X = m_i32Dim - 1;

	if( i32Y < 0 )
		i32Y = 0;
	if( i32Y >= m_i32Dim )
		i32Y = m_i32Dim - 1;

	i32NextX = i32X + 1;
	i32NextY = i32Y + 1;

	if( i32NextX >= m_i32Dim )
		i32NextX = m_i32Dim - 1;
	if( i32NextY >= m_i32Dim )
		i32NextY = m_i32Dim - 1;

	uint32	ui32Idx = i32X + i32Y * m_i32Dim;
	float32	f32Alpha;

	ui32Idx = i32X + i32Y * m_i32Dim;
	f32Alpha = (1.0f - f32AX) * (1.0f - f32AY);
	m_pU_U0[ui32Idx] += rForce[0] * f32Alpha;
	m_pU_V0[ui32Idx] += rForce[1] * f32Alpha;
	m_pRho[ui32Idx] += f32Matter * f32Alpha;

	ui32Idx = i32NextX + i32Y * m_i32Dim;
	f32Alpha = f32AX * (1.0f - f32AY);
	m_pU_U0[ui32Idx] += rForce[0] * f32Alpha;
	m_pU_V0[ui32Idx] += rForce[1] * f32Alpha;
	m_pRho[ui32Idx] += f32Matter * f32Alpha;

	ui32Idx = i32X + i32NextY * m_i32Dim;
	f32Alpha = (1.0f - f32AX) * f32AY;
	m_pU_U0[ui32Idx] += rForce[0] * f32Alpha;
	m_pU_V0[ui32Idx] += rForce[1] * f32Alpha;
	m_pRho[ui32Idx] += f32Matter * f32Alpha;

	ui32Idx = i32NextX + i32NextY * m_i32Dim;
	f32Alpha = f32AX * f32AY;
	m_pU_U0[ui32Idx] += rForce[0] * f32Alpha;
	m_pU_V0[ui32Idx] += rForce[1] * f32Alpha;
	m_pRho[ui32Idx] += f32Matter * f32Alpha;
}

void
FluidSolverC::push( float32 f32X, float32 f32Y, Vector2C rForce, PajaTypes::float32 f32Amount )
{
	f32X *= (float32)m_i32Dim;
	f32Y *= (float32)m_i32Dim;

	int32	i32X = floor( f32X );
	int32	i32Y = floor( f32Y );
	int32	i32NextX, i32NextY;

	float32	f32AX = f32X - (float32)i32X;
	float32	f32AY = f32Y - (float32)i32Y;

	if( i32X < 0 )
		i32X = 0;
	if( i32X >= m_i32Dim )
		i32X = m_i32Dim - 1;

	if( i32Y < 0 )
		i32Y = 0;
	if( i32Y >= m_i32Dim )
		i32Y = m_i32Dim - 1;

	i32NextX = i32X + 1;
	i32NextY = i32Y + 1;

	if( i32NextX >= m_i32Dim )
		i32NextX = m_i32Dim - 1;
	if( i32NextY >= m_i32Dim )
		i32NextY = m_i32Dim - 1;

	uint32	ui32Idx = i32X + i32Y * m_i32Dim;
	float32	f32Alpha;;

	ui32Idx = i32X + i32Y * m_i32Dim;
	f32Alpha = (1.0f - f32AX) * (1.0f - f32AY) * f32Amount;
	m_pU_U0[ui32Idx] += rForce[0] * f32Alpha;
	m_pU_V0[ui32Idx] += rForce[1] * f32Alpha;

	ui32Idx = i32NextX + i32Y * m_i32Dim;
	f32Alpha = f32AX * (1.0f - f32AY) * f32Amount;
	m_pU_U0[ui32Idx] += rForce[0] * f32Alpha;
	m_pU_V0[ui32Idx] += rForce[1] * f32Alpha;

	ui32Idx = i32X + i32NextY * m_i32Dim;
	f32Alpha = (1.0f - f32AX) * f32AY * f32Amount;
	m_pU_U0[ui32Idx] += rForce[0] * f32Alpha;
	m_pU_V0[ui32Idx] += rForce[1] * f32Alpha;

	ui32Idx = i32NextX + i32NextY * m_i32Dim;
	f32Alpha = f32AX * f32AY * f32Amount;
	m_pU_U0[ui32Idx] += rForce[0] * f32Alpha;
	m_pU_V0[ui32Idx] += rForce[1] * f32Alpha;
}

void
FluidSolverC::add_force( uint32 ui32X, uint32 ui32Y, Vector2C rForce )
{
	uint32	ui32Idx = ui32X + ui32Y * m_i32Dim;
	m_pU_U0[ui32Idx] += rForce[0];
	m_pU_V0[ui32Idx] += rForce[1];
}

#define FFT(s,u)\
	if( s == 1 ) rfftwnd_one_real_to_complex( m_rPlanRC, (fftw_real *)u, (fftw_complex*)u );\
	else rfftwnd_one_complex_to_real( m_rPlanCR, (fftw_complex *)u, (fftw_real *)u )

#define floor( x ) ((x) >= 0.0 ? ((int)(x)) : (-((int)(1-(x)))))

void
FluidSolverC::stable_solve( int32 n, fftw_real* u, fftw_real* v, fftw_real* u0, fftw_real* v0, fftw_real visc, fftw_real dt )
{
	fftw_real	x, y, x0, y0, f, r, U[2], V[2], s, t;
	int			i, j, i0, j0, i1, j1;
	int			mask = n - 1;
	float		dt_visc = dt * visc;

	for( i = 0; i < n * n ; i++ ) {
		u[i] += dt * u0[i]; 
		u0[i] = u[i];
		
		v[i] += dt * v0[i]; 
		v0[i] = v[i];
	}    

	for( x = 0.5f / n, i = 0; i < n; i++, x += 1.0f / n ) {
		for( y = 0.5f / n, j = 0; j < n; j++, y += 1.0f / n ) {
			x0 = n * (x - dt * u0[i + n * j]) - 0.5f; 
			y0 = n * (y - dt * v0[i + n * j]) - 0.5f;
			i0 = floor( x0 );
			s = x0 - i0;
			i0 = (n + (i0 & mask)) & mask;
			i1 = (i0 + 1) & mask;
			j0 = floor( y0 );
			t = y0  -j0;
			j0 = (n + (j0 & mask)) & mask;
			j1 = (j0 + 1) & mask;

			u[i + n * j] = (1 - s) * ((1 - t) * u0[i0 + n * j0] + t * u0[i0 + n * j1]) +                        
				s * ((1 - t) * u0[i1 + n * j0] + t * u0[i1 + n * j1]);
			v[i + n * j] = (1 - s) * ((1 - t) * v0[i0 + n * j0] + t * v0[i0 + n * j1]) +
				s * ((1 - t) * v0[i1 + n * j0] + t * v0[i1 + n * j1]);

			m_pRho[i + n * j] = (1 - s) * ((1 - t) * m_pRho0[i0 + n * j0] + t * m_pRho0[i0 + n * j1]) +                        
				s * ((1 - t) * m_pRho0[i1 + n * j0] + t * m_pRho0[i1 + n * j1]);
		}    
	} 
	
	for( i = 0; i < n ;i++ ) {
		for( j = 0; j < n ; j++ ) { 
			u0[i + (n + 2) * j] = u[i + n * j]; 
			v0[i + (n + 2) * j] = v[i + n * j];
		}
	}

	FFT( 1, u0 );
	FFT( 1, v0 );

	for( i = 0; i <= n; i += 2 ) {
		x = 0.5f * i;
		int offset = (n + 2) * j;
		int y2 = y * y;

		for( j = 0; j < n; j++ ) {
			y = j <= n / 2 ? (fftw_real)j : (fftw_real)j - n;
			r = x * x + y * y;
			if( r == 0.0f ) continue;

			int	offset = i + (n + 2) * j;

			float	x_inv_r = x / r;
			float	y_inv_r = y / r;

			f = (fftw_real)exp( -r * dt * visc );
//			f = 0.0001 * -r * dt_visc + 1;
			
			U[0] = u0[offset    ]; V[0] = v0[offset    ];
			U[1] = u0[offset + 1]; V[1] = v0[offset + 1];

			u0[offset    ] = f * ((1 - x * x_inv_r) * U[0]       -x * y_inv_r  * V[0]);
			u0[offset + 1] = f * ((1 - x * x_inv_r) * U[1]       -x * y_inv_r  * V[1]);
			v0[offset    ] = f * (    -y * x_inv_r  * U[0] + (1 - y * y_inv_r) * V[0]);
			v0[offset + 1] = f * (    -y * x_inv_r  * U[1] + (1 - y * y_inv_r) * V[1]);
		}    
	}

	FFT( -1, u0 ); 
	FFT( -1, v0 );

	f = 1.0 / (n * n);
	for( i = 0; i < n; i++ ) {
		for( j = 0; j < n ; j++ ) {
			u[i + n * j] = f * u0[i + (n + 2) * j]; 
			v[i + n * j] = f * v0[i + (n + 2) * j]; 
		}
	}
}

// This function diffuses matter that has been placed
// in the velocity field. It's almost identical to the
// velocity diffusion step in the function above. The
// input matter densities are in rho0 and the result
// is written into rho.
void
FluidSolverC::diffuse_matter( int32 n, fftw_real* u, fftw_real* v, fftw_real* rho, fftw_real* rho0, fftw_real dt )
{
/*	fftw_real	x, y, x0, y0, s, t;
	int			mask = n - 1;
	int			i, j, i0, j0, i1, j1;

	for( x = 0.5f / n, i = 0; i < n; i++, x += 1.0f / n ) {
		for( y = 0.5f / n, j = 0; j < n; j++, y += 1.0f / n ) {
			x0 = n * (x - dt * u[i + n * j]) - 0.5f; 
			y0 = n * (y - dt * v[i + n * j]) - 0.5f;
			i0 = floor( x0 );
			s = x0 - i0;
			i0 = (n + (i0 & mask)) & mask;
			i1 = (i0 + 1) & mask;
			j0 = floor( y0 );
			t = y0 - j0;
			j0 = (n + (j0 & mask)) & mask;
			j1 = (j0 + 1) & mask;

			rho[i + n * j] = (1 - s) * ((1 - t) * rho0[i0 + n * j0] + t * rho0[i0 + n * j1]) +                        
				s * ((1 - t) * rho0[i1 + n * j0] + t * rho0[i1 + n * j1]);
		}    
	} */
}

void
FluidSolverC::set_forces( void )
{
	int i;
	for( i = 0; i < m_i32Dim * m_i32Dim; i++ ) {
		m_pRho0[i] = 0.995 * m_pRho[i];

		m_pU_U0[i] *= 0.85;
		m_pU_V0[i] *= 0.85;

		m_pU0[i] = m_pU_U0[i];
		m_pV0[i] = m_pU_V0[i];
	}
}

int32
FluidSolverC::get_dimension() const
{
	return m_i32Dim;
}

fftw_real*
FluidSolverC::get_density_ptr()
{
	return m_pRho;
}

fftw_real*
FluidSolverC::get_smooth_density_ptr()
{
	return m_pRhoSmooth;
}

fftw_real*
FluidSolverC::get_velocity_u_ptr()
{
	return m_pU;
}

fftw_real*
FluidSolverC::get_velocity_v_ptr()
{
	return m_pV;
}

void
FluidSolverC::smooth()
{
	int32	i, j;
	int32	i32Mask = m_i32Dim - 1;

	int32	i32PrevX, i32NextX;
	int32	i32PrevY, i32NextY;
	int32	i32X, i32Y;

	i32PrevY = 0;

	for( i = 0; i < m_i32Dim; i++ ) {

		i32Y = i * m_i32Dim;
		i32NextY = __min( i + 1, m_i32Dim - 1 ) * m_i32Dim;
		
		i32PrevX = 0;
		for( j = 0; j < m_i32Dim; j++ ) {
			i32X = j;
			i32NextX = __min( j + 1, m_i32Dim - 1 );

			m_pRhoSmooth[i32X + i32Y] = (2.0f * m_pRho[i32X + i32Y] + m_pRho[i32PrevX + i32Y] + m_pRho[i32NextX + i32Y] +
												 m_pRho[i32X + i32PrevY] + m_pRho[i32X + i32NextY]) / 6.0f;

			i32PrevX = j;
		}

		i32PrevY = i32Y;
	}
}
