ZMWSLI0-SL2021-GR11/cw7/GrassController.cs

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

struct DifferentialEquationData
{
    public float water;
    public float vegetation;
}


public class GrassController : MonoBehaviour
{
    int _structSize()
    {
        return sizeof(float) + sizeof(float);
    }
    private GameObject cube;

    public float updateInterval;
    float timePassed;
    Terrain terrain;
    public ComputeShader computeShader;
    public ComputeBuffer oldVegetation;
    public ComputeBuffer vegetation;

    public ComputeBuffer oldWater;
    public ComputeBuffer water;

    public ComputeBuffer feedModifier;
    public ComputeBuffer killModifier;

    [Header("equation coeffs")]
    public float DWater = 1.0f;
    public float DVegetation = 0.5f;
    public float feedRate = 0.055f;
    public float killRate = 0.0462f;
    public float dT = 0.10f;
    public int stepsNumber = 1;


    int sizeX;
    int sizeY;
    // Start is called before the first frame update
    void Start()
    {
        terrain = gameObject.GetComponent<Terrain>();

        sizeX = terrain.terrainData.detailWidth;
        sizeY = terrain.terrainData.detailWidth;

        initComputeShader();
    }


    bool doUpdate()
    {
        timePassed += Time.deltaTime;
        if (timePassed < updateInterval)
        {
            return false;
        }
        else
        {
            timePassed = 0;
            return true;
        }

    }

    float transition(float start, float end, float value)
    {
        if (start > end)
        {
            return 1 - transition(end, start, value);
        }
        if (value < start)
        {
            return 0;
        }
        if (value >= end)
        {
            return 1;
        }
        var scaledValue = (value - start) / (end - start);
        return scaledValue;
    }

    void initKillFeedBuffers()
    {
        feedModifier = new ComputeBuffer(sizeX * sizeY, sizeof(float));
        killModifier = new ComputeBuffer(sizeX * sizeY, sizeof(float));
        computeShader.SetBuffer(0, "feedModifier", feedModifier);
        computeShader.SetBuffer(0, "killModifier", killModifier);

        var feedValues = new float[sizeX * sizeY];
        var killValues = new float[sizeX * sizeY];

        var width = terrain.terrainData.alphamapWidth;
        var height = terrain.terrainData.alphamapHeight;

        var splatMap = terrain.terrainData.GetAlphamaps(0, 0, width, height);

        for (var i = 0; i < sizeX; i++)
        {
            for (var j = 0; j < sizeY; j++)
            {
                var scale = height / (float)sizeX;
                var x = Mathf.FloorToInt(scale * i);
                var y = Mathf.FloorToInt(scale * j);
                var terrainHeight = terrain.terrainData.GetInterpolatedHeight(x, y);
                feedValues[i + j * sizeX] = transition(i, j, terrainHeight);

                var snow = splatMap[x, y, 5];
                if (snow > 0.5)
                {
                    killValues[i + j * sizeX] = 2;
                }
            }
        }

        feedModifier.SetData(killValues);
        killModifier.SetData(killValues);
    }

    public void initBufferValues()
    {
        var initValues = new float[sizeX * sizeY];

        for (var i = 0; i < sizeX * sizeY; i++)
        {
            initValues[i] = Random.Range(0.0f, 1.0f);
        }

        oldVegetation.SetData(initValues);
        vegetation.SetData(initValues);

        oldWater.SetData(initValues);
        water.SetData(initValues);
    }

    public void resetGrass()
    {
        initBufferValues();
        var detal = new int[sizeX, sizeY];

        for (var i = 0; i < sizeX; i++)
        {
            detal[i, i] = 16;
        }

        for (var i = 0; i < sizeX; i++)
        {
            detal[i, sizeX - i] = 16;
        }

        terrain.terrainData.SetDetailLayer(0, 0, 1, detal);
    }
    void initParameters()
    {
        computeShader.SetInt("sizeX", sizeX);
        computeShader.SetInt("sizeY", sizeY);
        computeShader.SetFloat("DWater", DWater);
        computeShader.SetFloat("DVegetation", DVegetation);
        computeShader.SetFloat("feedRate", feedRate);
        computeShader.SetFloat("killRate", killRate);
        computeShader.SetFloat("dT", dT);
    }

    void initComputeShader()
    {
        oldVegetation = new ComputeBuffer(sizeX * sizeY, sizeof(float));
        vegetation = new ComputeBuffer(sizeX * sizeY, sizeof(float));

        computeShader.SetBuffer(0, "oldVegetation", oldVegetation);
        computeShader.SetBuffer(0, "vegetation", vegetation);

        oldWater = new ComputeBuffer(sizeX * sizeY, sizeof(float));
        water = new ComputeBuffer(sizeX * sizeY, sizeof(float));

        computeShader.SetBuffer(0, "oldWater", oldWater);
        computeShader.SetBuffer(0, "water", water);

        initKillFeedBuffers();
        initBufferValues();
        initParameters();
    }


    void swapBuffers()
    {
        var temp = vegetation;
        vegetation = oldVegetation;
        oldVegetation = temp;
        computeShader.SetBuffer(0, "oldVegetation", oldVegetation);
        computeShader.SetBuffer(0, "vegetation", vegetation);

        var temp2 = water;
        water = oldWater;
        oldWater = temp2;
        computeShader.SetBuffer(0, "oldWater", oldWater);
        computeShader.SetBuffer(0, "water", water);
    }

    void doSimulationSteps(int steps)
    {
        for (int i = 0; i < steps; i++)
        {
            swapBuffers();
            computeShader.Dispatch(0, sizeX, sizeY, 1);
        }
    }

    void recoverData()
    {
        var result = new float[sizeX * sizeY];
        vegetation.GetData(result);
        var detailLayer = new int[sizeX, sizeY];

        for (var i = 0; i < sizeX; i++)
        {
            for (var j = 0; j < sizeX; j++)
            {
                detailLayer[i, j] = Mathf.FloorToInt(result[i + j * sizeX] * 10);
            }
        }

        terrain.terrainData.SetDetailLayer(0, 0, 1, detailLayer);
    }
    // Update is called once per frame
    void Update()
    {
        if (doUpdate())
        {
            terrain = gameObject.GetComponent<Terrain>();
            initParameters();
            doSimulationSteps(stepsNumber);
            recoverData();
        }
    }
}