Wind creator

//Forum thread : https://forum.unity.com/threads/how-to-add-wind-to-custom-trees-grass-in-unity-using-tree-creator-shaders.871126/
using System;
using System.Collections.Generic;
using UnityEngine;

public class WindCreator : MonoBehaviour
{
    private Mesh treeMesh;
    private List<Material> mats = new List<Material>();
    private bool barkMatFirst;

    [SerializeField]
    [Tooltip("If set to true, all the vertex colors and uv2 already set in this mesh will be replaced with new ones")]
    private bool overwriteExistingChannels = false;

    //Vertex colors variables
    [Header("Vertex Colors")]
    [SerializeField]
    [Tooltip("Do you want to use vertex colors on your mesh ?")]
    private bool useVertexColors = true;

    [SerializeField]
    [Tooltip("The type of gradient that will be apllied to the vertex colors")]
    private GradientType colorGradientType = GradientType.Linear;

    [SerializeField]
    [Tooltip("Put the bark material of your tree here")]
    private Material barkMat;

    [SerializeField]
    [Tooltip("Put the leaves material of your tree here")]
    private Material leafMat;

    [SerializeField]
    [Tooltip("How much the red value of the vertex colors should be painted ?")]
    [Range(0, 1)]
    private float vertexPercent = 0.5f;

    //Uv2 variables
    [Header("UV2 Channel")]
    [SerializeField]
    [Tooltip("Do you want to use uv2 on your mesh ?")]
    private bool useUv2 = true;

    [SerializeField]
    [Tooltip("The type of gradient that will be apllied to the uv2")]
    private GradientType uv2GradientType = GradientType.Circular;

    [SerializeField]
    [Tooltip("How much the values of the uv2 should be painted ?")]
    [Range(0, 1)]
    private float uv2Percent = 0.3f;

    [SerializeField]
    [Tooltip("How much the X value of the uv2 for the bark should be painted ?")]
    [Range(0, 1)]
    private float barkMultiplierX = 1f;

    [SerializeField]
    [Tooltip("How much the Y value of the uv2 for the bark should be painted ?")]
    [Range(0, 1)]
    private float barkMultiplierY = 0.5f;

    [SerializeField]
    [Tooltip("How much the X value of the uv2 for the leaves should be painted ?")]
    [Range(0, 1)]
    private float leafMutliplierX = 1;

    [SerializeField]
    [Tooltip("How much the Y value of the uv2 for the leaves should be painted ?")]
    [Range(0, 1)]
    private float leafMutliplierY = 1;


    /// <summary>
    /// This is the main function
    /// </summary>
    [ContextMenu("Create wind")]
    private void CreateWind()
    {
        try
        {
            treeMesh = CopyMesh(GetComponent<MeshFilter>().sharedMesh);
            mats.Clear();
            mats.AddRange(GetComponent<MeshRenderer>().sharedMaterials);
            if (mats[0].Equals(barkMat)) barkMatFirst = true;
            else if (mats[0].Equals(leafMat)) barkMatFirst = false;

            if (useVertexColors)
            {
                if (treeMesh.colors.Length > 0 && !overwriteExistingChannels)
                    Debug.LogWarning("Existing vertex colors channel detected, wind wont be set to this channel, consider enabling overwriteExistingChannels.");
                else if ((treeMesh.colors.Length <= 0 && !overwriteExistingChannels) || (treeMesh.colors.Length > 0 && overwriteExistingChannels))
                    PaintColors();
            }

            if (useUv2)
            {
                if (treeMesh.uv2.Length > 0 && !overwriteExistingChannels)
                    Debug.LogWarning("Existing uv2 channel detected, wind wont be set to this channel, consider enabling overwriteExistingChannels.");
                else if ((treeMesh.uv2.Length <= 0 && !overwriteExistingChannels) || (treeMesh.uv2.Length > 0 && overwriteExistingChannels))
                    PaintUV2();
            }

            if (useVertexColors || useUv2) CreateNewMesh();
            else Debug.LogWarning("Nothing to add to the mesh, please set useVertexColors or useUv2 to true.");
        }
        catch (Exception ex)
        {
            Debug.LogError(ex.Message);
        }
    }
    /// <summary>
    /// This function will add the colors channel to the mesh
    /// </summary>
    private void PaintColors()
    {
        try
        {
            Vector3 range = treeMesh.bounds.max - treeMesh.bounds.min;
            List<Color> colors = new List<Color>(new Color[treeMesh.vertexCount]);
            HashSet<int> tris = new HashSet<int>();
            for (int i = 0; i < treeMesh.subMeshCount; i++)
            {
                int start = tris.Count;
                tris.UnionWith(new HashSet<int>(treeMesh.GetTriangles(i)));
                for (int a = start; a < tris.Count; a++)
                {
                    float power = 0;
                    //Linear gradient
                    if (colorGradientType == GradientType.Linear)
                        power = treeMesh.vertices[a].y / range.y * vertexPercent;
                    //Circular gradient
                    if (colorGradientType == GradientType.Circular)
                        power = Mathf.Sqrt(Mathf.Pow(treeMesh.vertices[a].x / range.x, 2) + Mathf.Pow(treeMesh.vertices[a].y / range.y, 2) + Mathf.Pow(treeMesh.vertices[a].z / range.z, 2)) / Mathf.Sqrt(3) * vertexPercent;

                    //Bark mat
                    if ((barkMatFirst && i == 0) || (!barkMatFirst && i != 0))
                        colors[a] = new Color(power, 0, 0);

                    //Leaf mat
                    else if ((!barkMatFirst && i == 0) || (barkMatFirst && i != 0))
                        colors[a] = new Color(power, 1, 0);
                }
            }
            treeMesh.colors = colors.ToArray();
        }
        catch (Exception ex)
        {
            Debug.LogError(ex.Message);
        }
    }
    /// <summary>
    /// This function will add the second uv channel to the mesh
    /// </summary>
    private void PaintUV2()
    {
        try
        {
            Vector3 range = treeMesh.bounds.max - treeMesh.bounds.min;
            List<Vector2> uv2 = new List<Vector2>(new Vector2[treeMesh.vertexCount]);
            HashSet<int> tris = new HashSet<int>();
            for (int i = 0; i < treeMesh.subMeshCount; i++)
            {
                int start = tris.Count;
                tris.UnionWith(new HashSet<int>(treeMesh.GetTriangles(i)));
                for (int a = start; a < tris.Count; a++)
                {
                    float power = 0;
                    //Linear gradient
                    if (uv2GradientType == GradientType.Linear)
                        power = treeMesh.vertices[a].y / range.y * uv2Percent;
                    //Circular gradient
                    if (uv2GradientType == GradientType.Circular)
                        power = Mathf.Sqrt(Mathf.Pow(treeMesh.vertices[a].x / range.x, 2) + Mathf.Pow(treeMesh.vertices[a].y / range.y, 2) + Mathf.Pow(treeMesh.vertices[a].z / range.z, 2))  / Mathf.Sqrt(3) * uv2Percent;

                    //Bark mat
                    if ((barkMatFirst && i == 0) || (!barkMatFirst && i != 0))
                        uv2[a] = new Vector2(power * barkMultiplierX, power * barkMultiplierY);

                    //Leaf mat
                    else if ((!barkMatFirst && i == 0) || (barkMatFirst && i != 0))
                        uv2[a] = new Vector2(power * leafMutliplierX, power * leafMutliplierY);
                }
            }
            treeMesh.uv2 = uv2.ToArray();
        }
        catch (Exception ex)
        {
            Debug.LogError(ex.Message);
        }
    }
    /// <summary>
    /// This function will create the new gameobject with all the needed components
    /// </summary>
    private void CreateNewMesh()
    {
        GameObject tree = new GameObject(gameObject.name + "_Wind");
        tree.AddComponent<MeshFilter>().sharedMesh = treeMesh;
        tree.AddComponent<MeshRenderer>().sharedMaterials = mats.ToArray();
        tree.AddComponent<WindMultipleMat>();
        Debug.Log("Wind successfully added to : " + gameObject.name, tree);
        gameObject.SetActive(false);
    }
    /// <summary>
    /// This function makes a copy of a Mesh, it breaks the reference
    /// </summary>
    /// <param name="mesh">The Mesh to copy</param>
    /// <returns>The copied Mesh</returns>
    private Mesh CopyMesh(Mesh mesh)
    {
        return Instantiate(mesh);
    }
    /// <summary>
    /// The gradient that will be applied
    /// </summary>
    private enum GradientType
    {
        Linear,
        Circular
    }

}