Vertex color baker for Stylized Water 2

using Unity.Mathematics;
using UnityEngine;
#if UNITY_EDITOR
using UnityEditor;
#endif

[RequireComponent(typeof(MeshFilter))]
[ExecuteInEditMode]
public class WaterVertexColorBaker : MonoBehaviour
{
    [Header("Geometry")]
    public MeshFilter meshFilter;
    public Mesh sourceMesh;

    [Space]

    [Header("Raycasting")]
    [Tooltip("Only accept raycast hits from terrain geometry")]
    public bool terrainOnly = true;
    [Tooltip("Instead of raycasting precesly from the vertex's position, move the origin of the ray up this much. " +
             "\n\nThis is particularly useful when a vertex sits just underneath the terrain at a shoreline, but you need it to register a hit")]
    public float heightOffset = 0f;

    [Space]

    [Header("Effect information")]
    [Min(0f)]
    public float intersectionLength = 1f;
    [Min(0f)]
    public float depthExponent = 1f;
    [Min(0f)]
    public float waveAttenuation = 1f;

    [Space]

    [Header("Procedural surface foam painting")]
    [Range(0f,1f)]
    public float foamNoiseAmount;
    public float foamNoiseFrequency = 0.1f;

    private void Reset()
    {
        meshFilter = GetComponent<MeshFilter>();
        sourceMesh = meshFilter.sharedMesh;
    }

    [ContextMenu("BAKE")]
    public void Bake()
    {
        if(!meshFilter || !sourceMesh) return;

        meshFilter.sharedMesh = BakeVertexColors(transform, heightOffset, sourceMesh, intersectionLength, depthExponent, waveAttenuation, foamNoiseAmount, foamNoiseFrequency, terrainOnly);
    }

    public static Mesh BakeVertexColors(Transform transform, float heightOffset, Mesh sourceMesh, float intersectionLength, float depthExponent, float waveAttenuation, float foamAmount, float foamFrequency, bool terrainOnly)
    {
        Vector3[] verts = sourceMesh.vertices;
        Mesh newMesh = new Mesh();
        newMesh.name = sourceMesh.name + " (Baked vertex colors)";

        Color[] colors = new Color[verts.Length];

        RaycastHit hit;

        float dist = 0;
        bool hasHitBelow = false;
        bool hasHitAbove = false;

        for (int i = 0; i < verts.Length; i++)
        {
            Vector3 position = transform.TransformPoint(verts[i]);
            Vector3 rayOrigin = position;

            //Default value for missed hits
            dist = 0;
            hasHitAbove = false;
            hasHitBelow = false;

            rayOrigin.y += heightOffset;

            if (Physics.Raycast(rayOrigin, Vector3.down, out hit, Mathf.Infinity, -1, QueryTriggerInteraction.Ignore))
            {
                hasHitBelow = hit.point.y < position.y;
                hasHitAbove = hit.point.y >= position.y;

                dist = hit.point.y - position.y;

                if (terrainOnly && hit.collider.GetType() != typeof(TerrainCollider))
                {
                    dist = 0;
                    hasHitAbove = false;
                    hasHitBelow = false;
                }
            }

            float intersection = 1-Mathf.Abs(dist / intersectionLength);
            float depthVal = Mathf.Exp(dist * depthExponent);
            float waveVal = Mathf.Clamp01(1-Mathf.Abs(dist / waveAttenuation));
            float foam = (noise.cnoise(new float2(position.x * foamFrequency * 0.1f, position.z * foamFrequency * 0.1f)) * 0.5f + 0.5f) * foamAmount;

            intersection = hasHitBelow ? intersection : 0f;
            intersection = hasHitAbove ? 1f : intersection;

            depthVal = hasHitBelow ? depthVal : 0f;
            depthVal = hasHitAbove ? 1f : depthVal;

            waveVal = hasHitBelow ? waveVal : 0f;
            waveVal = hasHitAbove ? 1f : waveVal;

            colors[i] = new Color(intersection, depthVal, waveVal, foamAmount > 0 ? foam : colors[i].a);
        }

        newMesh.vertices = verts;
        newMesh.triangles = sourceMesh.triangles;
        newMesh.normals = sourceMesh.normals;
        newMesh.tangents = sourceMesh.tangents;
        newMesh.uv = sourceMesh.uv;

        newMesh.colors = colors;

        newMesh.RecalculateBounds();

        return newMesh;
    }
}

#if UNITY_EDITOR
[CustomEditor(typeof(WaterVertexColorBaker))]
[CanEditMultipleObjects]
class WaterVertexColorBakerEditor : Editor
{
    public override void OnInspectorGUI()
    {
        EditorGUI.BeginChangeCheck();

        base.OnInspectorGUI();

        if (EditorGUI.EndChangeCheck())
        {
            EditorApplication.delayCall += () =>
            {
                foreach (var target in targets)
                    ((WaterVertexColorBaker)target).Bake();
            };
        }
    }
}
#endif