using System.Collections;
using System.Collections.Generic;
using UnityEngine;

public class FluidSimulation : MonoBehaviour
{

    public RenderTexture readA;
    public RenderTexture readB;
    public RenderTexture readC;
    public RenderTexture readD;

    public RenderTexture writeA;
    public RenderTexture writeB;
    public RenderTexture writeC;
    public RenderTexture writeD;

    public ComputeShader Compute;

    public Material OutputMaterial;

    [Range(8, 1024)] public int resolution = 8;

    private int stepKernel;

    private int nthreads = 1;

    private int threadresolution => (resolution / nthreads);

    private Camera cam;

    [Range(1,50)]public int stepMod = 1;

    [Range(0, 50)] public int stepsPerFrame = 0;


    // Start is called before the first frame update
    void Start()
    {
        cam = Camera.main;
        Reset();
    }

    private void Reset()
    {
        readA = CreateTexture(RenderTextureFormat.ARGBFloat);
        readB = CreateTexture(RenderTextureFormat.ARGBFloat);
        readC = CreateTexture(RenderTextureFormat.ARGBFloat);
        readD = CreateTexture(RenderTextureFormat.ARGBFloat);

        writeA = CreateTexture(RenderTextureFormat.ARGBFloat);
        writeB = CreateTexture(RenderTextureFormat.ARGBFloat);
        writeC = CreateTexture(RenderTextureFormat.ARGBFloat);
        writeD = CreateTexture(RenderTextureFormat.ARGBFloat);

        stepKernel = Compute.FindKernel("StepKernel");

        GPUResetKernel();
    }

    void GPUResetKernel()
    {
        int k = Compute.FindKernel("ResetKernel");

        Compute.SetTexture(k, "readA", readA);
        Compute.SetTexture(k, "readB", readB);
        Compute.SetTexture(k, "readC", readC);
        Compute.SetTexture(k, "readD", readD);

        Compute.SetTexture(k, "writeA", writeA);
        Compute.SetTexture(k, "writeB", writeB);
        Compute.SetTexture(k, "writeC", writeC);
        Compute.SetTexture(k, "writeD", writeD);

        Compute.SetInt("nthreads", nthreads);

        Compute.SetInt("resolution", resolution);

        Compute.Dispatch(k, threadresolution, threadresolution, 1);
    }

    // Update is called once per frame
    void Update()
    {
        if (Time.frameCount % stepMod == 0)
        {
            for (int i = 0; i < stepsPerFrame; i++)
            {
                Step();
            }
        }
    }

    void Step()
    {
        if (readA == null || readB == null || readC == null || readD == null)
        {
            Reset();
        }

        Compute.SetTexture(stepKernel, "readA", readA);
        Compute.SetTexture(stepKernel, "readB", readB);
        Compute.SetTexture(stepKernel, "readC", readC);
        Compute.SetTexture(stepKernel, "readD", readD);

        Compute.SetTexture(stepKernel, "writeA", writeA);
        Compute.SetTexture(stepKernel, "writeB", writeB);
        Compute.SetTexture(stepKernel, "writeC", writeC);
        Compute.SetTexture(stepKernel, "writeD", writeD);

        Compute.Dispatch(stepKernel, resolution/nthreads, resolution/nthreads, 1);

        SwapTex();

        OutputMaterial.SetTexture("_UnlitColorMap", readD);
    }

    protected RenderTexture CreateTexture(RenderTextureFormat format)
    {
        RenderTexture tex = new RenderTexture(resolution, resolution, 0, format);

        //IMPORTANT FOR GPU SHADERS, allows random access (like gpus will do)
        tex.enableRandomWrite = true;
        tex.dimension = UnityEngine.Rendering.TextureDimension.Tex2D;
        tex.volumeDepth = 1;
        tex.filterMode = FilterMode.Point;
        tex.wrapMode = TextureWrapMode.Repeat;
        tex.useMipMap = false;
        tex.Create();

        return tex;
    }

    void SwapTex()
    {
        RenderTexture tmp = readA;
        readA = writeA;
        writeA = tmp;

        tmp = readB;
        readB = writeB;
        writeB = tmp;

        tmp = readC;
        readC = writeC;
        writeC = tmp;

        tmp = readD;
        readD = writeD;
        writeD = tmp;
    }
}