package electrosphere;

import java.io.File;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.FloatBuffer;
import java.nio.file.Files;
import java.util.LinkedList;
import java.util.List;
import java.util.Random;
import java.util.Set;

import org.joml.Vector2i;
import org.lwjgl.BufferUtils;
import org.lwjgl.PointerBuffer;
import org.lwjgl.system.MemoryUtil;

/**
 * Simulates a fluid via opencl
 */
public class FluidSim {

    static {
        System.out.println(System.getProperty("user.dir"));
        System.load(System.getProperty("user.dir") + "/shared-folder/libfluidsim.dll");
    }

    public static final int DIM = 18;

    ByteBuffer x = ByteBuffer.allocateDirect(DIM * DIM * DIM * 4);
    ByteBuffer x0 = ByteBuffer.allocateDirect(DIM * DIM * DIM * 4);
    ByteBuffer u = ByteBuffer.allocateDirect(DIM * DIM * DIM * 4);
    ByteBuffer v = ByteBuffer.allocateDirect(DIM * DIM * DIM * 4);
    ByteBuffer w = ByteBuffer.allocateDirect(DIM * DIM * DIM * 4);
    ByteBuffer u0 = ByteBuffer.allocateDirect(DIM * DIM * DIM * 4);
    ByteBuffer v0 = ByteBuffer.allocateDirect(DIM * DIM * DIM * 4);
    ByteBuffer w0 = ByteBuffer.allocateDirect(DIM * DIM * DIM * 4);

    //The densities for every voxel for the current frame
    float[] densityArrayView = new float[DIM * DIM * DIM];
    //Should be set to add water to the current frame
    float[] density0ArrayView = new float[DIM * DIM * DIM];
    //The force vector for each voxel for the current frame
    float[] uArrayView = new float[DIM * DIM * DIM];
    float[] vArrayView = new float[DIM * DIM * DIM];
    float[] wArrayView = new float[DIM * DIM * DIM];
    //these should be set to the 
    float[] u0ArrayView = new float[DIM * DIM * DIM];
    float[] v0ArrayView = new float[DIM * DIM * DIM];
    float[] w0ArrayView = new float[DIM * DIM * DIM];
    

    static final float DIFFUSION_CONSTANT = 0.0f;
    static final float VISCOSITY_CONSTANT = 0.0f;

    static final int LINEARSOLVERTIMES = 20;

    static final float GRAVITY = -1000f;
    
    public void setup(){
        x.order(ByteOrder.LITTLE_ENDIAN);
        x0.order(ByteOrder.LITTLE_ENDIAN);
        u.order(ByteOrder.LITTLE_ENDIAN);
        v.order(ByteOrder.LITTLE_ENDIAN);
        w.order(ByteOrder.LITTLE_ENDIAN);
        u0.order(ByteOrder.LITTLE_ENDIAN);
        v0.order(ByteOrder.LITTLE_ENDIAN);
        w0.order(ByteOrder.LITTLE_ENDIAN);
        FloatBuffer xf = x.asFloatBuffer();
        FloatBuffer uf = u.asFloatBuffer();
        FloatBuffer vf = v.asFloatBuffer();
        FloatBuffer wf = w.asFloatBuffer();

        Random rand = new Random(1);
        //make a cube of water in the center
        for(int i = 0; i < DIM; i++){
            for(int j = 0; j < DIM; j++){
                for(int k = 0; k < DIM; k++){
                    if(
                        Math.abs(16 - i) < 4 &&
                        Math.abs(8 - j) < 4 &&
                        Math.abs(10 - k) < 4
                    ){
                        xf.put(1);
                        uf.put(rand.nextFloat() * 0.1f);
                        vf.put(-1f);
                        wf.put(rand.nextFloat() * 0.1f);
                    } else {
                        xf.put(0);
                        uf.put(0);
                        vf.put(0);
                        wf.put(0);
                    }
                }
            }
        }
    }

    /**
     * Runs a frame of the fluid simulation
     */
    public void simulate(int step, float timestep){

        //
        // Add forces and density here
        //
        addGravity();

        //
        //Performs main fluid simulation logic
        //
        writeNewStateIntoBuffers();

        if(x.position() > 0){
            x.position(0);
        }
        if(x0.position() > 0){
            x0.position(0);
        }
        if(u.position() > 0){
            u.position(0);
        }
        if(v.position() > 0){
            v.position(0);
        }
        if(w.position() > 0){
            w.position(0);
        }
        if(u0.position() > 0){
            u0.position(0);
        }
        if(v0.position() > 0){
            v0.position(0);
        }
        if(w0.position() > 0){
            w0.position(0);
        }
        simulate(DIM, DIM, DIM, x, x0, u, v, w, u0, v0, w0, DIFFUSION_CONSTANT, VISCOSITY_CONSTANT, timestep);

        //
        //Reads out the results of the fluid sim
        //
        readDataIntoArrays();

    }


    /**
     * Used post-simulation call to read data from the simulation from buffers back into arrays
     */
    private void readDataIntoArrays(){
        //
        //Read data back into arrays
        //
        if(x.position() > 0){
            x.position(0);
        }
        if(u.position() > 0){
            u.position(0);
        }
        if(v.position() > 0){
            v.position(0);
        }
        if(w.position() > 0){
            w.position(0);
        }
        for(int i = 0; i < DIM; i++){
            for(int j = 0; j < DIM; j++){
                for(int k = 0; k < DIM; k++){
                    densityArrayView[IX(i,j,k)] = x.getFloat();
                    uArrayView[IX(i,j,k)] = u.getFloat();
                    vArrayView[IX(i,j,k)] = v.getFloat();
                    wArrayView[IX(i,j,k)] = w.getFloat();
                }
            }
        }
    }

    /**
     * Writes data from the java-side arrays into buffers that get passed into c-side
     */
    private void writeNewStateIntoBuffers(){
        if(x0.position() > 0){
            x0.position(0);
        }
        if(u0.position() > 0){
            u0.position(0);
        }
        if(v0.position() > 0){
            v0.position(0);
        }
        if(w0.position() > 0){
            w0.position(0);
        }
        FloatBuffer x0FloatView = x0.asFloatBuffer();
        FloatBuffer u0FloatView = u0.asFloatBuffer();
        FloatBuffer v0FloatView = v0.asFloatBuffer();
        FloatBuffer w0FloatView = w0.asFloatBuffer();
        for(int i = 0; i < DIM; i++){
            for(int j = 0; j < DIM; j++){
                for(int k = 0; k < DIM; k++){
                    x0FloatView.put(density0ArrayView[IX(i,j,k)]);
                    u0FloatView.put(u0ArrayView[IX(i,j,k)]);
                    v0FloatView.put(v0ArrayView[IX(i,j,k)]);
                    w0FloatView.put(w0ArrayView[IX(i,j,k)]);
                }
            }
        }
    }

    /**
     * Adds gravity to the simulation
     */
    private void addGravity(){
        for(int i = 0; i < DIM; i++){
            for(int j = 0; j < DIM; j++){
                for(int k = 0; k < DIM; k++){
                    v0ArrayView[IX(i,j,k)] = densityArrayView[IX(i,j,k)] * GRAVITY;
                }
            }
        }
    }

    /**
     * The native function call to simulate a frame of fluid
     * @param DIM_X
     * @param DIM_Y
     * @param DIM_Z
     * @param x
     * @param x0
     * @param u
     * @param v
     * @param w
     * @param u0
     * @param v0
     * @param w0
     * @param DIFFUSION_CONSTANT
     * @param VISCOSITY_CONSTANT
     * @param timestep
     */
    private native void simulate(
        int DIM_X,
        int DIM_Y,
        int DIM_Z,
        ByteBuffer x,
        ByteBuffer x0,
        ByteBuffer u,
        ByteBuffer v,
        ByteBuffer w,
        ByteBuffer u0,
        ByteBuffer v0,
        ByteBuffer w0,
        float DIFFUSION_CONSTANT,
        float VISCOSITY_CONSTANT,
        float timestep
    );

    public float[] getData(){
        return densityArrayView;
    }

    public static final int IX(int x, int y, int z){
        return ((x)+(DIM)*(y) + (DIM)*(DIM)*(z));
    }

}