package electrosphere;

import java.io.File;
import java.io.IOException;
import java.io.OutputStream;
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.joml.Vector3i;
import org.lwjgl.BufferUtils;
import org.lwjgl.PointerBuffer;
import org.lwjgl.glfw.GLFW;
import org.lwjgl.system.MemoryUtil;

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

    /**
     * Load fluid sim library
     */
    static {
        String osName = System.getProperty("os.name").toLowerCase();
        System.out.println(osName);
        if(osName.contains("win")){
            System.load(new File("./shared-folder/libfluidsim.dll").toPath().toAbsolutePath().toString());
        } else {
            System.load(new File("./shared-folder/libfluidsim.so").toPath().toAbsolutePath().toString());
        }
    }

    public static final int DIM = 18;

    //                           
    //              +-------------+             ^ Y                
    //            / |           / |             |    _             
    //          /   |         /   |             |    /\  Z         
    //(0,2,0) +-----+-------+     |             |   /              
    //        |     +-------+-----+             |  /               
    //        |   /         |   /               | /                
    //        | /           | /                 |/                 
    //        +-------------+ (2,0,0)           +---------------> X

    //Buffers that contain density for current frame
    public ByteBuffer[] density = new ByteBuffer[27];
    //Buffers that contain new density to add to the simulation
    public ByteBuffer[] densityAddition = new ByteBuffer[27];
    //Buffers that contain u vector directions
    public ByteBuffer[] uVector = new ByteBuffer[27];
    //Buffers that contain v vector directions
    public ByteBuffer[] vVector = new ByteBuffer[27];
    //Buffers that contain w vector directions
    public ByteBuffer[] wVector = new ByteBuffer[27];
    //Buffers that contain u vector directions to add to the simulation
    public ByteBuffer[] uAdditionVector = new ByteBuffer[27];
    //Buffers that contain v vector directions to add to the simulation
    public ByteBuffer[] vAdditionVector = new ByteBuffer[27];
    //Buffers that contain w vector directions to add to the simulation
    public ByteBuffer[] wAdditionVector = new ByteBuffer[27];

    //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];
    public float[] v0ArrayView = new float[DIM * DIM * DIM];
    float[] w0ArrayView = new float[DIM * DIM * DIM];

    public int chunkMask = 0;
    

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

    static final int LINEARSOLVERTIMES = 20;

    static final float GRAVITY = -1000f;
    
    public void setup(Vector3i offset){
        //allocate buffers for this chunk
        density[13] = ByteBuffer.allocateDirect(DIM * DIM * DIM * 4);
        densityAddition[13] = ByteBuffer.allocateDirect(DIM * DIM * DIM * 4);
        uVector[13] = ByteBuffer.allocateDirect(DIM * DIM * DIM * 4);
        vVector[13] = ByteBuffer.allocateDirect(DIM * DIM * DIM * 4);
        wVector[13] = ByteBuffer.allocateDirect(DIM * DIM * DIM * 4);
        uAdditionVector[13] = ByteBuffer.allocateDirect(DIM * DIM * DIM * 4);
        vAdditionVector[13] = ByteBuffer.allocateDirect(DIM * DIM * DIM * 4);
        wAdditionVector[13] = ByteBuffer.allocateDirect(DIM * DIM * DIM * 4);
        //order endian-ness
        density[13].order(ByteOrder.LITTLE_ENDIAN);
        densityAddition[13].order(ByteOrder.LITTLE_ENDIAN);
        uVector[13].order(ByteOrder.LITTLE_ENDIAN);
        vVector[13].order(ByteOrder.LITTLE_ENDIAN);
        wVector[13].order(ByteOrder.LITTLE_ENDIAN);
        uAdditionVector[13].order(ByteOrder.LITTLE_ENDIAN);
        vAdditionVector[13].order(ByteOrder.LITTLE_ENDIAN);
        wAdditionVector[13].order(ByteOrder.LITTLE_ENDIAN);

        //set initial values for chunk
        FloatBuffer xf = density[13].asFloatBuffer();
        FloatBuffer uf = uVector[13].asFloatBuffer();
        FloatBuffer vf = vVector[13].asFloatBuffer();
        FloatBuffer wf = wVector[13].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(offset.x == 0 && offset.y == 0 && offset.z == 0){
                        if(
                            Math.abs(16 - i) < 4 &&
                            Math.abs(8 - j) < 4 &&
                            Math.abs(10 - k) < 4
                            // &&
                            // i < 17 && i > 0 &&
                            // j < 17 && j > 0 &&
                            // k < 17 && k > 0
                        ){
                            xf.put(1);
                            uf.put(0.1f);
                            vf.put(-1);
                            wf.put(0.1f);
                        }
                         else {
                            xf.put(0);
                            uf.put(0);
                            vf.put(0);
                            wf.put(0);
                        }
                    // } else {
                    //     if(
                    //         Math.abs(0 - i) < 5 &&
                    //         Math.abs(j) < 5 &&
                    //         Math.abs(0 - k) < 5 &&
                    //         i < 17 && i > 0 &&
                    //         j < 17 && j > 0 &&
                    //         k < 17 && k > 0
                    //     ){
                    //         // xf.put(1);
                    //         // uf.put(50);
                    //         // vf.put(0);
                    //         // wf.put(rand.nextFloat() * 0.1f);
                    //     } else {
                    //         xf.put(0);
                    //         uf.put(0);
                    //         vf.put(0);
                    //         wf.put(0);
                    //     }
                    // }
                }
            }
        }
    }

    static int i = 0;
    static double time = 0;
    static double lastTime = 0;
    public static void simChunks(FluidSim[][][] simArray, int step, float timestep){

        // simArray[0][1][2].density0ArrayView[IX(10,10,3)] = 3.0f;
        // simArray[2][1][2].density0ArrayView[IX(10,10,3)] = 3.0f;
        // simArray[2][1][0].density0ArrayView[IX(10,10,3)] = 3.0f;

        List<FluidSim> chunksToSim = new LinkedList<FluidSim>();
        //
        //init data for upcoming frame
        for(int x = 0; x < simArray.length; x++){
            for(int y = 0; y < simArray[0].length; y++){
                for(int z = 0; z < simArray[0][0].length; z++){
                    //
                    // Add forces and density here
                    //
                    simArray[x][y][z].addGravity();
                    //
                    //Performs main fluid simulation logic
                    //
                    simArray[x][y][z].writeNewStateIntoBuffers();
                    //
                    // add to queue
                    //
                    chunksToSim.add(simArray[x][y][z]);
                }
            }
        }

        lastTime = GLFW.glfwGetTime();
        //
        //simulate
        simulateWrapper(chunksToSim,0.01f);
        //clock
        time = time + (GLFW.glfwGetTime() - lastTime);
        i++;
        if(i == 100){
            System.out.println(time / 100.0 * 1000.0);
        }


        //
        //Read out of buffers back into accessible arrays
        for(int x = 0; x < simArray.length; x++){
            for(int y = 0; y < simArray[0].length; y++){
                for(int z = 0; z < simArray[0][0].length; z++){
                    //
                    //Reads out the results of the fluid sim
                    //
                    simArray[x][y][z].readDataIntoArrays();
                }
            }
        }


        if(Main.endStep == 0){
            System.exit(0);
        }

        

    }



    //
    //
    //Management functions
    //
    //
    /**
     * Creates a fluid simulation array based on a given dimension set
     * @param dimx the x dimension
     * @param dimy the y dimension
     * @param dimz the z dimension
     * @return The array of fluid sim chunks
     */
    public static FluidSim[][][] initFluidSim(int dimx, int dimy, int dimz){
        FluidSim[][][] simArray = new FluidSim[dimx][dimy][dimz];
        for(int x = 0; x < simArray.length; x++){
            for(int y = 0; y < simArray[0].length; y++){
                for(int z = 0; z < simArray[0][0].length; z++){
                    simArray[x][y][z] = new FluidSim();
                    simArray[x][y][z].setup(new Vector3i(x,y,z));
                }
            }
        }
        //set sim adjacencies
        for(int x = 0; x < simArray.length; x++){
            for(int y = 0; y < simArray[0].length; y++){
                for(int z = 0; z < simArray[0][0].length; z++){
                    FluidSim current = simArray[x][y][z];

                    for(int i = -1; i < 2; i++){
                        for(int j = -1; j < 2; j++){
                            for(int k = -1; k < 2; k++){
                                if(i == j && j == k && k == 0){
                                    continue;
                                }
                                if(
                                    0 <= x + i && x + i < simArray.length &&
                                    0 <= y + j && y + j < simArray[0].length &&
                                    0 <= z + k && z + k < simArray[0][0].length
                                ){
                                    current.setNeighbor(i+1,j+1,k+1,simArray[x+i][y+j][z+k]);
                                }
                            }
                        }
                    }

                }
            }
        }
        return simArray;
    }























    static void printLayer(FluidSim[][][] simArray, int step, int layer){
        if(step == 0){
            for(int x = 0; x < DIM; x++){
                for(int y = 0; y < DIM; y++){
                    System.out.printf("%.2f\t",simArray[0][0][0].uArrayView[IX(x,layer,y)]);
                }
                System.out.println();
            }
        }
    }

    private static double sumAllDensity(FluidSim[][][] simArray){
        double rVal = 0;
        for(int x = 0; x < simArray.length; x++){
            for(int y = 0; y < simArray[0].length; y++){
                for(int z = 0; z < simArray[0][0].length; z++){
                    rVal = rVal + simArray[x][y][z].sumDensity();
                }
            }
        }
        return rVal;
    }

    private double sumDensity(){
        double rVal = 0;
        for(int x = 1; x < DIM - 1; x++){
            for(int y = 1; y < DIM - 1; y++){
                for(int z = 1; z < DIM - 1; z++){
                    rVal = rVal + densityArrayView[IX(x,y,z)];
                }
            }
        }
        return rVal;
    }

    private double sumU(){
        double rVal = 0;
        for(int x = 1; x < DIM - 1; x++){
            for(int y = 1; y < DIM - 1; y++){
                for(int z = 1; z < DIM - 1; z++){
                    rVal = rVal + Math.abs(uArrayView[IX(x,y,z)]);
                }
            }
        }
        return rVal;
    }

    private static double sumAllU(FluidSim[][][] simArray){
        double rVal = 0;
        for(int x = 0; x < simArray.length; x++){
            for(int y = 0; y < simArray[0].length; y++){
                for(int z = 0; z < simArray[0][0].length; z++){
                    simArray[x][y][z].readDataIntoArrays();
                    rVal = rVal + simArray[x][y][z].sumU();
                }
            }
        }
        return rVal;
    }

    private double sumU0(){
        double rVal = 0;
        for(int x = 1; x < DIM - 1; x++){
            for(int y = 1; y < DIM - 1; y++){
                for(int z = 1; z < DIM - 1; z++){
                    rVal = rVal + Math.abs(u0ArrayView[IX(x,y,z)]);
                }
            }
        }
        return rVal;
    }

    private static double sumAllU0(FluidSim[][][] simArray){
        double rVal = 0;
        for(int x = 0; x < simArray.length; x++){
            for(int y = 0; y < simArray[0].length; y++){
                for(int z = 0; z < simArray[0][0].length; z++){
                    simArray[x][y][z].readDataIntoArrays();
                    rVal = rVal + simArray[x][y][z].sumU0();
                }
            }
        }
        return rVal;
    }

    private double sumV(){
        double rVal = 0;
        for(int x = 1; x < DIM - 1; x++){
            for(int y = 1; y < DIM - 1; y++){
                for(int z = 1; z < DIM - 1; z++){
                    rVal = rVal + Math.abs(vArrayView[IX(x,y,z)]);
                }
            }
        }
        return rVal;
    }

    private static double sumAllV(FluidSim[][][] simArray){
        double rVal = 0;
        for(int x = 0; x < simArray.length; x++){
            for(int y = 0; y < simArray[0].length; y++){
                for(int z = 0; z < simArray[0][0].length; z++){
                    simArray[x][y][z].readDataIntoArrays();
                    rVal = rVal + simArray[x][y][z].sumV();
                }
            }
        }
        return rVal;
    }

    private double sumV0(){
        double rVal = 0;
        for(int x = 1; x < DIM - 1; x++){
            for(int y = 1; y < DIM - 1; y++){
                for(int z = 1; z < DIM - 1; z++){
                    rVal = rVal + Math.abs(v0ArrayView[IX(x,y,z)]);
                }
            }
        }
        return rVal;
    }

    private static double sumAllV0(FluidSim[][][] simArray){
        double rVal = 0;
        for(int x = 0; x < simArray.length; x++){
            for(int y = 0; y < simArray[0].length; y++){
                for(int z = 0; z < simArray[0][0].length; z++){
                    simArray[x][y][z].readDataIntoArrays();
                    rVal = rVal + simArray[x][y][z].sumV0();
                }
            }
        }
        return rVal;
    }



    /**
     * Main simulation function
     * @param timestep
     */
    private static void simulateWrapper(List<FluidSim> chunks, float timestep){
        simulate(chunks,timestep);
    }
    private static native void simulate(List<FluidSim> chunks, float timestep);



    /**
     * Dumps this chunk's density values to disk
     * @param path the path to dump to
     */
    public void dumpToDisk(String path){
        OutputStream fileOS = null;
        try {
            fileOS = Files.newOutputStream(new File(path).toPath());
            ByteBuffer density = this.density[13];
            while(density.hasRemaining()){
                fileOS.write(density.get());
            }
            fileOS.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

















































    /**
     * Used post-simulation call to read data from the simulation from buffers back into arrays
     */
    private void readDataIntoArrays(){
        //
        //Read data back into arrays
        //
        if(density[13].position() > 0){
            density[13].position(0);
        }
        if(uVector[13].position() > 0){
            uVector[13].position(0);
        }
        if(vVector[13].position() > 0){
            vVector[13].position(0);
        }
        if(wVector[13].position() > 0){
            wVector[13].position(0);
        }
        if(uAdditionVector[13].position() > 0){
            uAdditionVector[13].position(0);
        }
        if(vAdditionVector[13].position() > 0){
            vAdditionVector[13].position(0);
        }
        if(wAdditionVector[13].position() > 0){
            wAdditionVector[13].position(0);
        }
        FloatBuffer xFloatView = density[13].asFloatBuffer();
        FloatBuffer uFloatView = uVector[13].asFloatBuffer();
        FloatBuffer vFloatView = vVector[13].asFloatBuffer();
        FloatBuffer wFloatView = wVector[13].asFloatBuffer();
        FloatBuffer u0FloatView = uAdditionVector[13].asFloatBuffer();
        FloatBuffer v0FloatView = vAdditionVector[13].asFloatBuffer();
        FloatBuffer w0FloatView = wAdditionVector[13].asFloatBuffer();
        int index = 0;
        for(int i = 0; i < DIM; i++){
            for(int j = 0; j < DIM; j++){
                for(int k = 0; k < DIM; k++){
                    index = ((i)+(DIM)*(j) + (DIM)*(DIM)*(k));
                    densityArrayView[index] = xFloatView.get();
                    uArrayView[index] = uFloatView.get();
                    vArrayView[index] = vFloatView.get();
                    wArrayView[index] = wFloatView.get();
                    u0ArrayView[index] = u0FloatView.get();
                    v0ArrayView[index] = v0FloatView.get();
                    w0ArrayView[index] = w0FloatView.get();
                }
            }
        }
    }

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

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












































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

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

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

    public ByteBuffer getDensityBuffer(){
        return density[getNeighborIndex(1,1,1)];
    }

    public ByteBuffer getDensityAdditionBuffer(){
        return densityAddition[getNeighborIndex(1,1,1)];
    }

    public ByteBuffer getUBuffer(){
        return uVector[getNeighborIndex(1,1,1)];
    }

    public ByteBuffer getVBuffer(){
        return vVector[getNeighborIndex(1,1,1)];
    }

    public ByteBuffer getWBuffer(){
        return wVector[getNeighborIndex(1,1,1)];
    }

    public ByteBuffer getUAddBuffer(){
        return uAdditionVector[getNeighborIndex(1,1,1)];
    }

    public ByteBuffer getVAddBuffer(){
        return vAdditionVector[getNeighborIndex(1,1,1)];
    }

    public ByteBuffer getWAddBuffer(){
        return wAdditionVector[getNeighborIndex(1,1,1)];
    }

    public void setNeighbor(int x, int y, int z, FluidSim neighbor){
        density[getNeighborIndex(x,y,z)] = neighbor.getDensityBuffer();
        densityAddition[getNeighborIndex(x,y,z)] = neighbor.getDensityAdditionBuffer();
        uVector[getNeighborIndex(x,y,z)] = neighbor.getUBuffer();
        vVector[getNeighborIndex(x,y,z)] = neighbor.getVBuffer();
        wVector[getNeighborIndex(x,y,z)] = neighbor.getWBuffer();
        uAdditionVector[getNeighborIndex(x,y,z)] = neighbor.getUAddBuffer();
        vAdditionVector[getNeighborIndex(x,y,z)] = neighbor.getVAddBuffer();
        wAdditionVector[getNeighborIndex(x,y,z)] = neighbor.getWAddBuffer();
    }

}
