#include <jni.h>
#include <stdio.h>
#include <immintrin.h>
#include <stdint.h>

#include "includes/utilities.h"
#include "includes/chunkmask.h"



#define LINEARSOLVERTIMES 10



void diffuse(JNIEnv * env, uint32_t chunk_mask, int N, int b, float * x, float * x0, float diff, float dt);
void advect(JNIEnv * env, uint32_t chunk_mask, int N, int b, jobjectArray jrd, float * d0, float * u, float * v, float * w, float dt);
void project(JNIEnv * env, uint32_t chunk_mask, int N, jobjectArray jru, jobjectArray jrv, jobjectArray jrw, float * p, float * div);
void set_bnd(JNIEnv * env, uint32_t chunk_mask, int N, int b, float * x);
void dens_step(JNIEnv * env, uint32_t chunk_mask, int N, jobjectArray jrx, float * x0, jobjectArray jru, jobjectArray jrv, jobjectArray jrw, float diff, float dt);
void vel_step(JNIEnv * env, uint32_t chunk_mask, int N, jobjectArray jru, jobjectArray jrv, jobjectArray jrw, float * u0, float * v0, float * w0, float visc, float dt);
void lin_solve(JNIEnv * env, uint32_t chunk_mask, int N, int b, float* x, float* x0, float a, float c);

/**
 * The core simulation function
*/
JNIEXPORT void JNICALL Java_electrosphere_FluidSim_simulate(
    JNIEnv * env,
    jobject this,
    jint DIM_X,
    jint chunk_mask,
    jobject jx,
    jobject jx0,
    jobject ju,
    jobject jv,
    jobject jw,
    jobject ju0,
    jobject jv0,
    jobject jw0,
    jfloat DIFFUSION_RATE,
    jfloat VISCOSITY_RATE,
    jfloat timestep){
    jboolean isCopy;
    // float * x = (*env)->GetDirectBufferAddress(env,jx);
    float * x0 = (*env)->GetDirectBufferAddress(env,jx0);
    // float * u = (*env)->GetDirectBufferAddress(env,ju);
    // float * v = (*env)->GetDirectBufferAddress(env,jv);
    // float * w = (*env)->GetDirectBufferAddress(env,jw);
    float * u0 = (*env)->GetDirectBufferAddress(env,ju0);
    float * v0 = (*env)->GetDirectBufferAddress(env,jv0);
    float * w0 = (*env)->GetDirectBufferAddress(env,jw0);
    int N = DIM_X;
    int i,j,k;
    vel_step(env, chunk_mask, DIM_X, ju, jv, jw, u0, v0, w0, VISCOSITY_RATE, timestep);
    dens_step(env, chunk_mask, DIM_X, jx, x0, ju, jv, jw, DIFFUSION_RATE, timestep);
}

/**
 * Adds values from a source array to a current frame array (eg more density to the main density array)
*/
void add_source(int N, float * x, float * s, float dt){
	int i;
    int size=N*N*N;
	for(i=0; i<size; i++){
        x[i] += dt*s[i];
    }
}

/**
 * Diffuses a given array by a diffusion constant
*/
void diffuse(JNIEnv * env, uint32_t chunk_mask, int N, int b, float * x, float * x0, float diff, float dt){
    float a=dt*diff*N*N*N;
    lin_solve(env, chunk_mask, N, b, x, x0, a, 1+6*a);
}

/**
 * Advects a given array based on the force vectors in the simulation
*/
void advect(JNIEnv * env, uint32_t chunk_mask, int N, int b, jobjectArray jrd, float * d0, float * u, float * v, float * w, float dt){
    int i, j, k, i0, j0, k0, i1, j1, k1;
    float x, y, z, s0, t0, s1, t1, u1, u0, dtx,dty,dtz;
    
    dtx=dty=dtz=dt*N;

    float * d = GET_ARR(env,jrd,CENTER_LOC);

    for(k=1; k<N-1; k++){
        for(j=1; j<N-1; j++){
            for(i=1; i<N-1; i++){
                x = i-dtx*u[IX(i,j,k)]; y = j-dty*v[IX(i,j,k)]; z = k-dtz*w[IX(i,j,k)];
                if (x<0.5f) x=0.5f; if (x>N+0.5f) x=N+0.5f; i0=(int)x; i1=i0+1;
                if (y<0.5f) y=0.5f; if (y>N+0.5f) y=N+0.5f; j0=(int)y; j1=j0+1;
                if (z<0.5f) z=0.5f; if (z>N+0.5f) z=N+0.5f; k0=(int)z; k1=k0+1;

                s1 = x-i0; s0 = 1-s1; t1 = y-j0; t0 = 1-t1; u1 = z-k0; u0 = 1-u1;
                if(i0 >= N){
                    i0 = N - 1;
                }
                // if(i0 < 0){
                //     i0 = 0;
                // }
                if(j0 >= N){
                    j0 = N - 1;
                }
                // if(j0 < 0){
                //     j0 = 0;
                // }
                if(k0 >= N){
                    k0 = N - 1;
                }
                // if(k0 < 0){
                //     k0 = 0;
                // }
                if(i1 >= N){
                    i1 = N - 1;
                }
                // if(i1 < 0){
                //     i1 = 0;
                // }
                if(j1 >= N){
                    j1 = N - 1;
                }
                // if(j1 < 0){
                //     j1 = 0;
                // }
                if(k1 >= N){
                    k1 = N - 1;
                }
                // if(k1 < 0){
                //     k1 = 0;
                // }
                d[IX(i,j,k)] = s0*(t0*u0*d0[IX(i0,j0,k0)]+t1*u0*d0[IX(i0,j1,k0)]+t0*u1*d0[IX(i0,j0,k1)]+t1*u1*d0[IX(i0,j1,k1)])+
                    s1*(t0*u0*d0[IX(i1,j0,k0)]+t1*u0*d0[IX(i1,j1,k0)]+t0*u1*d0[IX(i1,j0,k1)]+t1*u1*d0[IX(i1,j1,k1)]);
            }
        }
    }
    set_bnd(env, chunk_mask, N, b, d);
}

/**
 * The main density step function
*/
void dens_step(JNIEnv * env, uint32_t chunk_mask, int N, jobjectArray jrx, float * x0, jobjectArray jru, jobjectArray jrv, jobjectArray jrw, float diff, float dt){
    float * x = GET_ARR(env,jrx,CENTER_LOC);
    float * u = GET_ARR(env,jru,CENTER_LOC);
    float * v = GET_ARR(env,jrv,CENTER_LOC);
    float * w = GET_ARR(env,jrw,CENTER_LOC);
    add_source(N, x, x0, dt);
    SWAP(x0, x);
    diffuse(env, chunk_mask, N, 0, x, x0, diff, dt);
    SWAP(x0, x);
    advect(env, chunk_mask, N, 0, jrx, x0, u, v, w, dt);
}

/**
 * The main velocity step function
*/
void vel_step(JNIEnv * env, uint32_t chunk_mask, int N, jobjectArray jru, jobjectArray jrv, jobjectArray jrw, float * u0, float * v0, float * w0, float visc, float dt){
    float * u = GET_ARR(env,jru,CENTER_LOC);
    float * v = GET_ARR(env,jrv,CENTER_LOC);
    float * w = GET_ARR(env,jrw,CENTER_LOC);
    add_source(N, u, u0, dt);
    add_source(N, v, v0, dt);
    add_source(N, w, w0, dt);
    SWAP(u0, u);
    diffuse(env, chunk_mask, N, 1, u, u0, visc, dt);
    SWAP(v0, v);
    diffuse(env, chunk_mask, N, 2, v, v0, visc, dt);
    SWAP(w0, w);
    diffuse(env, chunk_mask, N, 3, w, w0, visc, dt);
    project(env, chunk_mask, N, jru, jrv, jrw, u0, v0);
    SWAP(u0, u);
    SWAP(v0, v);
    SWAP(w0, w);
    advect(env, chunk_mask, N, 1, jru, u0, u0, v0, w0, dt);
    advect(env, chunk_mask, N, 2, jrv, v0, u0, v0, w0, dt);
    advect(env, chunk_mask, N, 3, jrw, w0, u0, v0, w0, dt);
    project(env, chunk_mask, N, jru, jrv, jrw, u0, v0);
}

//used for temporary vector storage when appropriate
float container[16];

/**
 * Projects a given array based on force vectors
*/
void project(JNIEnv * env, uint32_t chunk_mask, int N, jobjectArray jru, jobjectArray jrv, jobjectArray jrw, float * p, float * div){
    int i, j, k;

    __m256 nVector = _mm256_set1_ps(N);
    __m256 constScalar = _mm256_set1_ps(-1.0/3.0);
    __m256 zeroVec = _mm256_set1_ps(0);
    __m256 vector, vector2, vector3;

    float * u = GET_ARR(env,jru,CENTER_LOC);
    float * v = GET_ARR(env,jrv,CENTER_LOC);
    float * w = GET_ARR(env,jrw,CENTER_LOC);

    for(k=1; k<N-1; k++){
        for(j=1; j<N-1; j++){
            i = 1;
            //
            //lower
            //
            //first part
            vector = _mm256_loadu_ps(&u[IX(i+1,j,k)]);
            vector = _mm256_sub_ps(vector,_mm256_loadu_ps(&u[IX(i-1,j,k)]));
            vector = _mm256_div_ps(vector,nVector);
            //second part
            vector2 = _mm256_loadu_ps(&v[IX(i,j+1,k)]);
            vector2 = _mm256_sub_ps(vector2,_mm256_loadu_ps(&v[IX(i,j-1,k)]));
            vector2 = _mm256_div_ps(vector2,nVector);
            //third part
            vector3 = _mm256_loadu_ps(&w[IX(i,j,k+1)]);
            vector3 = _mm256_sub_ps(vector3,_mm256_loadu_ps(&w[IX(i,j,k-1)]));
            vector3 = _mm256_div_ps(vector3,nVector);
            //multiply and finalize
            vector = _mm256_add_ps(vector,_mm256_add_ps(vector2,vector3));
            vector = _mm256_mul_ps(vector,constScalar);
            //store
            _mm256_storeu_ps(&div[IX(i,j,k)],vector);
            _mm256_storeu_ps(&p[IX(i,j,k)],zeroVec);
            i = 9;
            //
            //upper
            //
            //first part
            vector = _mm256_loadu_ps(&u[IX(i+1,j,k)]);
            vector = _mm256_sub_ps(vector,_mm256_loadu_ps(&u[IX(i-1,j,k)]));
            vector = _mm256_div_ps(vector,nVector);
            //second part
            vector2 = _mm256_loadu_ps(&v[IX(i,j+1,k)]);
            vector2 = _mm256_sub_ps(vector2,_mm256_loadu_ps(&v[IX(i,j-1,k)]));
            vector2 = _mm256_div_ps(vector2,nVector);
            //third part
            vector3 = _mm256_loadu_ps(&w[IX(i,j,k+1)]);
            vector3 = _mm256_sub_ps(vector3,_mm256_loadu_ps(&w[IX(i,j,k-1)]));
            vector3 = _mm256_div_ps(vector3,nVector);
            //multiply and finalize
            vector = _mm256_add_ps(vector,_mm256_add_ps(vector2,vector3));
            vector = _mm256_mul_ps(vector,constScalar);
            //store
            _mm256_storeu_ps(&div[IX(i,j,k)],vector);
            _mm256_storeu_ps(&p[IX(i,j,k)],zeroVec);
        }
    }
    
    set_bnd(env, chunk_mask, N, 0, div);
    set_bnd(env, chunk_mask, N, 0, p);

    lin_solve(env, chunk_mask, N, 0, p, div, 1, 6);


    constScalar = _mm256_set1_ps(0.5f*N);
    for ( k=1 ; k<N-1 ; k++ ) {
        for ( j=1 ; j<N-1 ; j++ ) {
            //
            //v
            //
            //lower
            vector =  _mm256_loadu_ps(&p[IX(1+1,j,k)]);
            vector2 = _mm256_loadu_ps(&p[IX(1-1,j,k)]);
            vector =  _mm256_sub_ps(vector,vector2);
            vector = _mm256_mul_ps(vector,constScalar);
            vector = _mm256_sub_ps(_mm256_loadu_ps(&u[IX(1,j,k)]),vector);
            _mm256_storeu_ps(&u[IX(1,j,k)],vector);
            //upper
            vector =  _mm256_loadu_ps(&p[IX(9+1,j,k)]);
            vector2 = _mm256_loadu_ps(&p[IX(9-1,j,k)]);
            vector =  _mm256_sub_ps(vector,vector2);
            vector = _mm256_mul_ps(vector,constScalar);
            vector = _mm256_sub_ps(_mm256_loadu_ps(&u[IX(9,j,k)]),vector);
            _mm256_storeu_ps(&u[IX(9,j,k)],vector);
            //
            //v
            //
            //lower
            vector =  _mm256_loadu_ps(&p[IX(1,j+1,k)]);
            vector2 = _mm256_loadu_ps(&p[IX(1,j-1,k)]);
            vector =  _mm256_sub_ps(vector,vector2);
            vector = _mm256_mul_ps(vector,constScalar);
            vector = _mm256_sub_ps(_mm256_loadu_ps(&v[IX(1,j,k)]),vector);
            _mm256_storeu_ps(&v[IX(1,j,k)],vector);
            //upper
            vector =  _mm256_loadu_ps(&p[IX(9,j+1,k)]);
            vector2 = _mm256_loadu_ps(&p[IX(9,j-1,k)]);
            vector =  _mm256_sub_ps(vector,vector2);
            vector = _mm256_mul_ps(vector,constScalar);
            vector = _mm256_sub_ps(_mm256_loadu_ps(&v[IX(9,j,k)]),vector);
            _mm256_storeu_ps(&v[IX(9,j,k)],vector);
            //
            //w
            //
            //lower
            vector =  _mm256_loadu_ps(&p[IX(1,j,k+1)]);
            vector2 = _mm256_loadu_ps(&p[IX(1,j,k-1)]);
            vector =  _mm256_sub_ps(vector,vector2);
            vector = _mm256_mul_ps(vector,constScalar);
            vector = _mm256_sub_ps(_mm256_loadu_ps(&w[IX(1,j,k)]),vector);
            _mm256_storeu_ps(&w[IX(1,j,k)],vector);
            //upper
            vector =  _mm256_loadu_ps(&p[IX(9,j,k+1)]);
            vector2 = _mm256_loadu_ps(&p[IX(9,j,k-1)]);
            vector =  _mm256_sub_ps(vector,vector2);
            vector = _mm256_mul_ps(vector,constScalar);
            vector = _mm256_sub_ps(_mm256_loadu_ps(&w[IX(9,j,k)]),vector);
            _mm256_storeu_ps(&w[IX(9,j,k)],vector);
        }
    }
    
    set_bnd(env, chunk_mask, N, 1, u);
    set_bnd(env, chunk_mask, N, 2, v);
    set_bnd(env, chunk_mask, N, 3, w);
}

/**
 * Solves a linear system of equations in a vectorized manner
*/
void lin_solve(JNIEnv * env, uint32_t chunk_mask, int N, int b, float* x, float* x0, float a, float c){
    int i, j, k, l, m;
    __m256 aScalar = _mm256_set1_ps(a);
    __m256 cScalar = _mm256_set1_ps(c);
    // iterate the solver
    for ( l=0 ; l<LINEARSOLVERTIMES ; l++ ) {
        // update for each cell
        for(k=1; k<N-1; k++){
            for(j=1; j<N-1; j++){
                int n = 0;
                //solve as much as possible vectorized
                for(i = 1; i < N-1; i=i+8){
                    __m256 vector = _mm256_loadu_ps(&x[IX(i-1,j,k)]);
                    vector = _mm256_add_ps(vector,_mm256_loadu_ps(&x[IX(i+1,j,k)]));
                    vector = _mm256_add_ps(vector,_mm256_loadu_ps(&x[IX(i,j-1,k)]));
                    vector = _mm256_add_ps(vector,_mm256_loadu_ps(&x[IX(i,j+1,k)]));
                    vector = _mm256_add_ps(vector,_mm256_loadu_ps(&x[IX(i,j,k-1)]));
                    vector = _mm256_add_ps(vector,_mm256_loadu_ps(&x[IX(i,j,k+1)]));
                    vector = _mm256_mul_ps(vector,aScalar);
                    vector = _mm256_add_ps(vector,_mm256_loadu_ps(&x0[IX(i,j,k)]));
                    vector = _mm256_div_ps(vector,cScalar);
                    _mm256_storeu_ps(&x[IX(i,j,k)],vector);
                }
                //If there is any leftover, perform manual solving
                if(i>N-1){
                    for(i=i-8; i < N-1; i++){
                        x[IX(i,j,k)] = (x0[IX(i,j,k)] + a*(x[IX(i-1,j,k)]+x[IX(i+1,j,k)]+x[IX(i,j-1,k)]+x[IX(i,j+1,k)]+x[IX(i,j,k-1)]+x[IX(i,j,k+1)]))/c;
                    }
                }
            }
        }
        set_bnd(env, chunk_mask, N, b, x);
    }
}

/**
 * Sets the bounds of the simulation
*/
void set_bnd(JNIEnv * env, uint32_t chunk_mask, int N, int b, float * target){
    int DIM = N;
    for(int x=1; x < DIM-1; x++){
        for(int y = 1; y < DIM-1; y++){
            //((x)+(DIM)*(y) + (DIM)*(DIM)*(z))
            target[0 + DIM * x + DIM * DIM * y] =         b==1 ? -target[1 + DIM * x + DIM * DIM * y] : target[1 + DIM * x + DIM * DIM * y];
            target[IX(DIM-1,x,y)] =     b==1 ? -target[IX(DIM-2,x,y)] : target[IX(DIM-2,x,y)];
            target[IX(x,0,y)] =         b==2 ? -target[IX(x,1,y)] : target[IX(x,1,y)];
            target[IX(x,DIM-1,y)] =     b==2 ? -target[IX(x,DIM-2,y)] : target[IX(x,DIM-2,y)];
            target[IX(x,y,0)] =         b==3 ? -target[IX(x,y,1)] : target[IX(x,y,1)];
            target[IX(x,y,DIM-1)] =     b==3 ? -target[IX(x,y,DIM-2)] : target[IX(x,y,DIM-2)];
        }
    }
    for(int x = 1; x < DIM-1; x++){
        target[IX(x,0,0)]         = (float)(0.5f * (target[IX(x,1,0)] + target[IX(x,0,1)]));
        target[IX(x,DIM-1,0)]     = (float)(0.5f * (target[IX(x,DIM-2,0)] + target[IX(x,DIM-1,1)]));
        target[IX(x,0,DIM-1)]     = (float)(0.5f * (target[IX(x,1,DIM-1)] + target[IX(x,0,DIM-2)]));
        target[IX(x,DIM-1,DIM-1)] = (float)(0.5f * (target[IX(x,DIM-2,DIM-1)] + target[IX(x,DIM-1,DIM-2)]));

        target[IX(0,x,0)] = (float)(0.5f * (target[IX(1,x,0)] + target[IX(0,x,1)]));
        target[IX(DIM-1,x,0)] = (float)(0.5f * (target[IX(DIM-2,x,0)] + target[IX(DIM-1,x,1)]));
        target[IX(0,x,DIM-1)] = (float)(0.5f * (target[IX(1,x,DIM-1)] + target[IX(0,x,DIM-2)]));
        target[IX(DIM-1,x,DIM-1)] = (float)(0.5f * (target[IX(DIM-2,x,DIM-1)] + target[IX(DIM-1,x,DIM-2)]));


        target[IX(0,0,x)] = (float)(0.5f * (target[IX(1,0,x)] + target[IX(0,1,x)]));
        target[IX(DIM-1,0,x)] = (float)(0.5f * (target[IX(DIM-2,0,x)] + target[IX(DIM-1,1,x)]));
        target[IX(0,DIM-1,x)] = (float)(0.5f * (target[IX(1,DIM-1,x)] + target[IX(0,DIM-2,x)]));
        target[IX(DIM-1,DIM-1,x)] = (float)(0.5f * (target[IX(DIM-2,DIM-1,x)] + target[IX(DIM-1,DIM-2,x)]));

    }
    target[IX(0,0,0)]             = (float)((target[IX(1,0,0)]+target[IX(0,1,0)]+target[IX(0,0,1)])/3.0);
    target[IX(DIM-1,0,0)]         = (float)((target[IX(DIM-2,0,0)]+target[IX(DIM-1,1,0)]+target[IX(DIM-1,0,1)])/3.0);
    target[IX(0,DIM-1,0)]         = (float)((target[IX(1,DIM-1,0)]+target[IX(0,DIM-2,0)]+target[IX(0,DIM-1,1)])/3.0);
    target[IX(0,0,DIM-1)]         = (float)((target[IX(0,0,DIM-2)]+target[IX(1,0,DIM-1)]+target[IX(0,1,DIM-1)])/3.0);
    target[IX(DIM-1,DIM-1,0)]     = (float)((target[IX(DIM-2,DIM-1,0)]+target[IX(DIM-1,DIM-2,0)]+target[IX(DIM-1,DIM-1,1)])/3.0);
    target[IX(0,DIM-1,DIM-1)]     = (float)((target[IX(1,DIM-1,DIM-1)]+target[IX(0,DIM-2,DIM-1)]+target[IX(0,DIM-1,DIM-2)])/3.0);
    target[IX(DIM-1,0,DIM-1)]     = (float)((target[IX(DIM-1,0,DIM-2)]+target[IX(DIM-2,0,DIM-1)]+target[IX(DIM-1,1,DIM-1)])/3.0);
    target[IX(DIM-1,DIM-1,DIM-1)] = (float)((target[IX(DIM-1,DIM-1,DIM-2)]+target[IX(DIM-1,DIM-2,DIM-1)]+target[IX(DIM-1,DIM-1,DIM-2)])/3.0);
}