#include <jni.h>

//library includes
//include stb ds
#define STB_DS_IMPLEMENTATION
#include "../lib/stb/stb_ds.h"

//local includes
#include "../includes/chunk.h"
#include "../includes/chunkmask.h"
#include "../includes/utilities.h"
#include "../includes/simulation.h"


//defines


//function defines
#define getChunk(i) (*env)->CallObjectMethod(env,chunkList,jListGet,i)
#define getBuffArr(buffId) (*env)->GetObjectField(env,chunkJRaw,buffId)
#define setBuffArr(buffId,value) (*env)->SetObjectField(env,chunkJRaw,buffId,value)
#define GET_ARR(env,src,i) (*env)->GetDirectBufferAddress(env,(*env)->GetObjectArrayElement(env,src,i))



//declarations
void readInChunks(JNIEnv * env, jobject chunkList);
int calculateChunkMask(JNIEnv * env, jobjectArray jrx);
uint32_t matrix_transform(JNIEnv * env, jobjectArray jrx);



//the list of chunks
Chunk ** javaChunkView = NULL;
//the number of chunks
int numChunks = 0;


JNIEXPORT void JNICALL Java_electrosphere_FluidSim_simulate(
    JNIEnv * env,
    jclass fluidSimClass,
    jobject chunkList,
    jfloat dt
    ){
    readInChunks(env,chunkList);
    simulate(numChunks,javaChunkView,dt);
}


/**
 * Reads chunks into the dynamic array
*/
void readInChunks(JNIEnv * env, jobject chunkList){
    jclass listClass = (*env)->FindClass(env,"java/util/List");
    jclass fluidSimClass = (*env)->FindClass(env,"electrosphere/FluidSim");
    //JNIEnv *env, jclass clazz, const char *name, const char *sig
    jmethodID jListSize = (*env)->GetMethodID(env, listClass, "size", "()I");
    jmethodID jListGet  = (*env)->GetMethodID(env, listClass, "get", "(I)Ljava/lang/Object;");
    jmethodID jListAdd  = (*env)->GetMethodID(env, listClass, "add", "(Ljava/lang/Object;)Z");

    //ByteBuffer[]
    jfieldID dJId = (*env)->GetFieldID(env,fluidSimClass,"density","[Ljava/nio/ByteBuffer;");
    jfieldID d0JId = (*env)->GetFieldID(env,fluidSimClass,"densityAddition","[Ljava/nio/ByteBuffer;");
    jfieldID uJId = (*env)->GetFieldID(env,fluidSimClass,"uVector","[Ljava/nio/ByteBuffer;");
    jfieldID vJId = (*env)->GetFieldID(env,fluidSimClass,"vVector","[Ljava/nio/ByteBuffer;");
    jfieldID wJId = (*env)->GetFieldID(env,fluidSimClass,"wVector","[Ljava/nio/ByteBuffer;");
    jfieldID u0JId = (*env)->GetFieldID(env,fluidSimClass,"uAdditionVector","[Ljava/nio/ByteBuffer;");
    jfieldID v0JId = (*env)->GetFieldID(env,fluidSimClass,"vAdditionVector","[Ljava/nio/ByteBuffer;");
    jfieldID w0JId = (*env)->GetFieldID(env,fluidSimClass,"wAdditionVector","[Ljava/nio/ByteBuffer;");
    jfieldID chunkmaskJId = (*env)->GetFieldID(env,fluidSimClass,"chunkMask","I");

    //the number of chunks
    numChunks = (*env)->CallIntMethod(env,chunkList,jListSize);

    //current chunk (this)
    jobject chunkJRaw;
    //current chunk fields
    jobjectArray jd;
    jobjectArray jd0;
    jobjectArray u;
    jobjectArray v;
    jobjectArray w;
    jobjectArray u0;
    jobjectArray v0;
    jobjectArray w0;
    int chunkMask;

    //solve chunk mask
    for(int i = 0; i < numChunks; i++){
        chunkJRaw = getChunk(i);
        chunkMask = calculateChunkMask(env,getBuffArr(dJId));
        (*env)->SetIntField(env,chunkJRaw,chunkmaskJId,chunkMask);

        Chunk * newChunk;
        if(i >= stbds_arrlen(javaChunkView)){
            // printf("allocate chunk %d\n",i);
            // fflush(stdout);
            newChunk = (Chunk *)malloc(sizeof(Chunk));
            // printf("new chunk %p\n",newChunk);
            // fflush(stdout);
            stbds_arrput(javaChunkView,newChunk);
            // printf("new chunk %p\n",chunks[i]);
            // fflush(stdout);
        } else {
            newChunk = javaChunkView[i];
            // printf("get chunk %d: %p\n",i,newChunk);
            // fflush(stdout);
        }
        jd = (*env)->GetObjectField(env,chunkJRaw,dJId);
        jd0 = (*env)->GetObjectField(env,chunkJRaw,d0JId);
        u = (*env)->GetObjectField(env,chunkJRaw,uJId);
        v = (*env)->GetObjectField(env,chunkJRaw,vJId);
        w = (*env)->GetObjectField(env,chunkJRaw,wJId);
        u0 = (*env)->GetObjectField(env,chunkJRaw,u0JId);
        v0 = (*env)->GetObjectField(env,chunkJRaw,v0JId);
        w0 = (*env)->GetObjectField(env,chunkJRaw,w0JId);
        newChunk->chunkMask = chunkMask;
        for(int j = 0; j < 27; j++){
            if((chunkMask & CHUNK_INDEX_ARR[j]) > 0){
                newChunk->d[j] = GET_ARR(env,jd,j);
                newChunk->d0[j] = GET_ARR(env,jd0,j);
                newChunk->u[j] = GET_ARR(env,u,j);
                newChunk->v[j] = GET_ARR(env,v,j);
                newChunk->w[j] = GET_ARR(env,w,j);
                newChunk->u0[j] = GET_ARR(env,u0,j);
                newChunk->v0[j] = GET_ARR(env,v0,j);
                newChunk->w0[j] = GET_ARR(env,w0,j);
            }
        }
    }
}


/**
 * Calculates the bitmask for available chunks for the provided chunk's neighbor array
*/
int calculateChunkMask(JNIEnv * env, jobjectArray jrx){
    return matrix_transform(env,jrx);
}

/**
 * Calculates a mask that represents all nearby chunks that are actually accessible and exist
*/
uint32_t matrix_transform(JNIEnv * env, jobjectArray jrx){

    //The returned value, an availability mask that contains the availability of each neighbor chunk
    uint32_t rVal = 0;

    //Add to maks for initial chunks
    for(int i = 0; i < CENTER_LOC; i++){
        if((*env)->GetObjectArrayElement(env,jrx,i)!=NULL){
            rVal = rVal + 1;
        }
        rVal = rVal << 1;
    }
    //add 1 for center chunk because we already have that
    rVal = rVal + 1;
    rVal = rVal << 1;
    //continue on for remaining chunks
    for(int i = CENTER_LOC+1; i < 27; i++){
        if((*env)->GetObjectArrayElement(env,jrx,i)!=NULL){
            rVal = rVal + 1;
        }
        if(i < 26){
            rVal = rVal << 1;
        }
    }

    return rVal;
}