// based on https://github.com/m-schuetz/compute_rasterizer

using UnityEngine;

public struct Particle
{
    public Vector4 position;
    public Vector4 color;
}

// PUT THIS ON THE CAMERA! Otherwise OnRenderImage will not be called..
public class RenderPC : MonoBehaviour
{
    public ComputeShader splatShader;

    public Vector2Int size = new Vector2Int(2560, 1440);
    public float epsilonColor = .01f;
    public float epsilonExpand = .05f;
    public int expandIters = 1;

    public Transform target;
    public Mesh mesh;

    private GraphicsBuffer points;
    private GraphicsBuffer depth;
    private GraphicsBuffer color;

    private GraphicsBuffer depth2;
    private GraphicsBuffer color2;


    void Awake()
    {
        var particleCount = mesh.vertexCount;
        var pointData = new Particle[particleCount];
        var verts = mesh.vertices;
        var colors = mesh.colors;
        for (int i = 0; i < particleCount; i++) {
            pointData[i] = new Particle() {
                position = verts[i],
                color = colors[i],
            };
        }
        points = new GraphicsBuffer(GraphicsBuffer.Target.Structured, particleCount, sizeof(float) * 8);
        points.SetData(pointData);

        depth = new GraphicsBuffer(GraphicsBuffer.Target.Structured, size.x * size.y, sizeof(uint));
        color = new GraphicsBuffer(GraphicsBuffer.Target.Structured, size.x * size.y, sizeof(uint) * 4);

        depth2 = new GraphicsBuffer(GraphicsBuffer.Target.Structured, size.x * size.y, sizeof(uint));
        color2 = new GraphicsBuffer(GraphicsBuffer.Target.Structured, size.x * size.y, sizeof(uint) * 4);
    }



    void OnRenderImage(RenderTexture source, RenderTexture destination)
    {

        splatShader.SetInts("size", size.x, size.y);
        splatShader.SetFloat("epsilonColor", epsilonColor);
        splatShader.SetFloat("epsilonExpand", epsilonExpand);
        var mainCam = Camera.main;
        splatShader.SetMatrix("CameraVP", mainCam.projectionMatrix * mainCam.worldToCameraMatrix * target.localToWorldMatrix);

        {
            var kernel = splatShader.FindKernel("clear");
            splatShader.SetTextureFromGlobal(kernel, "CameraDepth", "_CameraDepthTexture");
            splatShader.SetBuffer(kernel, "DepthBuffer", depth);
            splatShader.SetBuffer(kernel, "ColorBuffer", color);
            splatShader.Dispatch(kernel, Mathf.CeilToInt(size.x / (float)32), Mathf.CeilToInt(size.y / (float)32), 1);
        }

        {
            var kernel = splatShader.FindKernel("splat_depth");
            splatShader.SetBuffer(kernel, "points", points);
            splatShader.SetBuffer(kernel, "DepthBuffer", depth);
            splatShader.Dispatch(kernel, Mathf.CeilToInt(points.count / (float)1024), 1, 1);
        }

        {
            var kernel = splatShader.FindKernel("splat_color");
            splatShader.SetBuffer(kernel, "points", points);
            splatShader.SetBuffer(kernel, "DepthBuffer", depth);
            splatShader.SetBuffer(kernel, "ColorBuffer", color);
            splatShader.Dispatch(kernel, Mathf.CeilToInt(points.count / (float)1024), 1, 1);
        }

        for (int i = 0; i < expandIters; i++) {
            var kernel = splatShader.FindKernel("expand");
            splatShader.SetBuffer(kernel, "DepthBuffer", depth);
            splatShader.SetBuffer(kernel, "ColorBuffer", color);
            splatShader.SetBuffer(kernel, "DepthBuffer2", depth2);
            splatShader.SetBuffer(kernel, "ColorBuffer2", color2);
            splatShader.Dispatch(kernel, Mathf.CeilToInt(size.x / (float)32), Mathf.CeilToInt(size.y / (float)32), 1);

            (depth, depth2) = (depth2, depth);
            (color, color2) = (color2, color);
        }

        {
            var desc = source.descriptor;
            desc.enableRandomWrite = true;
            desc.msaaSamples = 1;
            var tmp = RenderTexture.GetTemporary(desc);

            var kernel = splatShader.FindKernel("blit");
            splatShader.SetBuffer(kernel, "ColorBuffer", color);
            splatShader.SetTexture(kernel, "source", source);
            splatShader.SetTexture(kernel, "res", tmp);
            splatShader.SetTextureFromGlobal(kernel, "CameraDepth", "_CameraDepthTexture");
            splatShader.Dispatch(kernel, Mathf.CeilToInt(size.x / (float)32), Mathf.CeilToInt(size.y / (float)32), 1);

            Graphics.Blit(tmp, destination);
            RenderTexture.ReleaseTemporary(tmp);
        }
    }

    void OnDestroy()
    {
        points.Release();
        depth.Release();
        color.Release();
        depth2.Release();
        color2.Release();
    }
}