/*
 * Created by jiadong chen
 * http://www.chenjd.me
 */

using System.IO;
using UnityEditor;
using UnityEngine;

public class AnimMapBakerWindow : EditorWindow
{
    private enum SaveStrategy
    {
        AnimMap,//only anim map
        Mat,//with shader
        Prefab//prefab with mat
    }

    #region FIELDS

    // 目标对象
    private static GameObject _targetGo;

    // 烘焙器
    private static AnimMapBaker _baker;

    // 路径
    private static string _path = "DefaultPath";

    // 路径补充
    private static string _subPath = "SubPath";

    // 保存策略
    private static SaveStrategy _stratege = SaveStrategy.AnimMap;

    // 材质Shader
    private static Shader _animMapShader;

    #endregion FIELDS

    #region METHODS

    [MenuItem("Tool/AnimMapBaker")]
    public static void ShowWindow()
    {
        EditorWindow.GetWindow(typeof(AnimMapBakerWindow));
        _baker = new AnimMapBaker();
        _animMapShader = Shader.Find("chenjd/AnimMapShader");
    }

    private void OnGUI()
    {
        _targetGo = (GameObject)EditorGUILayout.ObjectField(_targetGo, typeof(GameObject), true);
        _subPath = _targetGo == null ? _subPath : _targetGo.name;
        EditorGUILayout.LabelField(string.Format($"output path:{Path.Combine(_path, _subPath)}"));
        _path = EditorGUILayout.TextField(_path);
        _subPath = EditorGUILayout.TextField(_subPath);

        _stratege = (SaveStrategy)EditorGUILayout.EnumPopup("output type:", _stratege);

        if (!GUILayout.Button("Bake")) return;
        if (_targetGo == null)
        {
            EditorUtility.DisplayDialog("错误", "目标对象为空", "OK");
            return;
        }

        if (_baker == null)
        {
            _baker = new AnimMapBaker();
        }
        // 写入动画对象
        _baker.SetAnimData(_targetGo);
        // 烘焙
        var list = _baker.Bake();

        if (list == null) return;
        foreach (var t in list)
        {
            var data = t;
            // 保存烘焙信息
            Save(ref data);
        }
    }

    private void Save(ref BakedData data)
    {
        switch (_stratege)
        {
            case SaveStrategy.AnimMap:
                SaveAsAsset(ref data);
                break;

            case SaveStrategy.Mat:
                SaveAsMat(ref data);
                break;

            case SaveStrategy.Prefab:
                SaveAsPrefab(ref data);
                break;
        }
        AssetDatabase.SaveAssets();
        AssetDatabase.Refresh();
    }

    /// <summary>
    /// 保存为贴图
    /// </summary>
    /// <param name="data"></param>
    /// <returns></returns>
    private Texture2D SaveAsAsset(ref BakedData data)
    {
        var folderPath = CreateFolder();
        // 创建贴图
        var animMap = new Texture2D(data.AnimMapWidth, data.AnimMapHeight, TextureFormat.RGBAHalf, false);
        // 加载数据到贴图
        animMap.LoadRawTextureData(data.RawAnimMap);
        // 创建对象
        AssetDatabase.CreateAsset(animMap, Path.Combine(folderPath, data.Name + ".asset"));
        return animMap;
    }

    /// <summary>
    /// 保存为材质
    /// </summary>
    /// <param name="data"></param>
    /// <returns></returns>
    private Material SaveAsMat(ref BakedData data)
    {
        if (_animMapShader == null)
        {
            EditorUtility.DisplayDialog("err", "shader is null!!", "OK");
            return null;
        }

        if (_targetGo == null || !_targetGo.GetComponentInChildren<SkinnedMeshRenderer>())
        {
            EditorUtility.DisplayDialog("err", "SkinnedMeshRender is null!!", "OK");
            return null;
        }

        var smr = _targetGo.GetComponentInChildren<SkinnedMeshRenderer>();
        var mat = new Material(_animMapShader);
        var animMap = SaveAsAsset(ref data);
        mat.SetTexture("_MainTex", smr.sharedMaterial.mainTexture);
        mat.SetTexture("_AnimMap", animMap);
        mat.SetFloat("_AnimLen", data.AnimLen);

        var folderPath = CreateFolder();
        AssetDatabase.CreateAsset(mat, Path.Combine(folderPath, data.Name + ".mat"));

        return mat;
    }

    /// <summary>
    /// 保存为预制体
    /// </summary>
    /// <param name="data"></param>
    private void SaveAsPrefab(ref BakedData data)
    {
        var mat = SaveAsMat(ref data);

        if (mat == null)
        {
            EditorUtility.DisplayDialog("err", "mat is null!!", "OK");
            return;
        }

        var go = new GameObject();
        go.AddComponent<MeshRenderer>().sharedMaterial = mat;
        go.AddComponent<MeshFilter>().sharedMesh = _targetGo.GetComponentInChildren<SkinnedMeshRenderer>().sharedMesh;

        var folderPath = CreateFolder();
        PrefabUtility.SaveAsPrefabAsset(go, Path.Combine(folderPath, data.Name + ".prefab")
            .Replace("\\", "/"));
    }

    /// <summary>
    /// 创建文件夹
    /// </summary>
    /// <returns></returns>
    private static string CreateFolder()
    {
        var folderPath = Path.Combine("Assets/" + _path, _subPath);
        if (!AssetDatabase.IsValidFolder(folderPath))
        {
            AssetDatabase.CreateFolder("Assets/" + _path, _subPath);
        }
        return folderPath;
    }

    #endregion METHODS
}

/*
 * Created by jiadong chen
 * http://www.chenjd.me
 *
 * 用来烘焙动作贴图。烘焙对象使用animation组件，并且在导入时设置Rig为Legacy
 */

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

/// <summary>
/// 保存需要烘焙的动画的相关数据
/// </summary>
public struct AnimData
{
    #region FIELDS

    // 顶点数
    private int _vertexCount;

    // 贴图宽度
    private int _mapWidth;

    // 动画Clip
    private readonly List<AnimationState> _animClips;

    private string _name;

    private Animation _animation;
    private SkinnedMeshRenderer _skin;

    public List<AnimationState> AnimationClips => _animClips;
    public int MapWidth => _mapWidth;
    public string Name => _name;

    #endregion FIELDS

    public AnimData(Animation anim, SkinnedMeshRenderer smr, string goName)
    {
        _vertexCount = smr.sharedMesh.vertexCount;
        _mapWidth = Mathf.NextPowerOfTwo(_vertexCount);
        _animClips = new List<AnimationState>(anim.Cast<AnimationState>());
        _animation = anim;
        _skin = smr;
        _name = goName;
    }

    #region METHODS

    public void AnimationPlay(string animName)
    {
        _animation.Play(animName);
    }

    /// <summary>
    /// 采样动画并且烘焙贴图
    /// </summary>
    /// <param name="m"></param>
    public void SampleAnimAndBakeMesh(ref Mesh m)
    {
        SampleAnim();
        BakeMesh(ref m);
    }

    /// <summary>
    /// 采样动画
    /// </summary>
    private void SampleAnim()
    {
        if (_animation == null)
        {
            Debug.LogError("animation is null!!");
            return;
        }

        _animation.Sample();
    }

    /// <summary>
    /// 烘焙材质
    /// </summary>
    /// <param name="m"></param>
    private void BakeMesh(ref Mesh m)
    {
        if (_skin == null)
        {
            Debug.LogError("skin is null!!");
            return;
        }

        _skin.BakeMesh(m);
    }

    #endregion METHODS
}

/// <summary>
/// 烘焙后的数据
/// </summary>
public struct BakedData
{
    #region FIELDS

    private readonly string _name;
    private readonly float _animLen;
    private readonly byte[] _rawAnimMap;
    private readonly int _animMapWidth;
    private readonly int _animMapHeight;

    #endregion FIELDS

    public BakedData(string name, float animLen, Texture2D animMap)
    {
        _name = name;
        _animLen = animLen;
        _animMapHeight = animMap.height;
        _animMapWidth = animMap.width;
        _rawAnimMap = animMap.GetRawTextureData();
    }

    public int AnimMapWidth => _animMapWidth;

    public string Name => _name;

    public float AnimLen => _animLen;

    public byte[] RawAnimMap => _rawAnimMap;

    public int AnimMapHeight => _animMapHeight;
}

/// <summary>
/// 烘焙器
/// </summary>
public class AnimMapBaker
{
    #region FIELDS

    private AnimData? _animData = null;
    private Mesh _bakedMesh;
    private readonly List<Vector3> _vertices = new List<Vector3>();
    private readonly List<BakedData> _bakedDataList = new List<BakedData>();

    #endregion FIELDS

    #region METHODS

    /// <summary>
    /// 设置烘焙对象
    /// </summary>
    /// <param name="go"></param>
    public void SetAnimData(GameObject go)
    {
        if (go == null)
        {
            Debug.LogError("go is null!!");
            return;
        }
        // 获取动画组件与蒙皮网格组件
        var anim = go.GetComponent<Animation>();
        var smr = go.GetComponentInChildren<SkinnedMeshRenderer>();

        if (anim == null || smr == null)
        {
            Debug.LogError("anim or smr is null!!");
            return;
        }
        // 创建网格
        _bakedMesh = new Mesh();
        // 创建动画信息
        _animData = new AnimData(anim, smr, go.name);
    }

    /// <summary>
    /// 开始烘焙
    /// </summary>
    /// <returns></returns>
    public List<BakedData> Bake()
    {
        if (_animData == null)
        {
            Debug.LogError("bake data is null!!");
            return _bakedDataList;
        }

        //遍历动画动作，每一个动作都生成一个动作图
        foreach (var t in _animData.Value.AnimationClips)
        {
            if (!t.clip.legacy)
            {
                Debug.LogError(string.Format($"{t.clip.name} is not legacy!!"));
                continue;
            }
            BakePerAnimClip(t);
        }

        return _bakedDataList;
    }

    // 烘焙动画剪辑
    private void BakePerAnimClip(AnimationState curAnim)
    {
        var curClipFrame = 0;
        float sampleTime = 0;
        float perFrameTime = 0;
        // 获取离值最近的二次方数（动画帧率*动画长度）
        curClipFrame = Mathf.ClosestPowerOfTwo((int)(curAnim.clip.frameRate * curAnim.length));
        perFrameTime = curAnim.length / curClipFrame; ;
        // 创建贴图
        var animMap = new Texture2D(_animData.Value.MapWidth, curClipFrame, TextureFormat.RGBAHalf, true);
        // 设置贴图名称
        animMap.name = string.Format($"{_animData.Value.Name}_{curAnim.name}.animMap");
        // 播放动画
        _animData.Value.AnimationPlay(curAnim.name);

        for (var i = 0; i < curClipFrame; i++)
        {
            curAnim.time = sampleTime;
            // 采样动画并且烘焙贴图
            _animData.Value.SampleAnimAndBakeMesh(ref _bakedMesh);

            for (var j = 0; j < _bakedMesh.vertexCount; j++)
            {
                var vertex = _bakedMesh.vertices[j];

                // 将动画顶点坐标写入贴图
                animMap.SetPixel(j, i, new Color(vertex.x, vertex.y, vertex.z));
            }

            sampleTime += perFrameTime;
        }
        animMap.Apply();

        _bakedDataList.Add(new BakedData(animMap.name, curAnim.clip.length, animMap));
    }

    #endregion METHODS
}

/*
Created by jiadong chen
http://www.chenjd.me
*/

Shader "chenjd/AnimMapShader"
{
    Properties
    {
        _MainTex ("Texture", 2D) = "white" { }
        // 动画贴图
        _AnimMap ("AnimMap", 2D) = "white" { }
        // 动画长度
        _AnimLen ("Anim Length", Float) = 0
    }
    SubShader
    {
        Tags { "RenderType" = "Opaque" }
        LOD 100
        Cull off

        Pass
        {
            CGPROGRAM
            
            #pragma vertex vert
            #pragma fragment frag
            //开启gpu instancing
            #pragma multi_compile_instancing


            #include "UnityCG.cginc"

            struct appdata
            {
                float2 uv: TEXCOORD0;
                float4 pos: POSITION;
                UNITY_VERTEX_INPUT_INSTANCE_ID
            };

            struct v2f
            {
                float2 uv: TEXCOORD0;
                float4 vertex: SV_POSITION;
                UNITY_VERTEX_INPUT_INSTANCE_ID
            };

            sampler2D _MainTex;
            float4 _MainTex_ST;

            // 动画贴图
            sampler2D _AnimMap;
            float4 _AnimMap_TexelSize;//x == 1/width
            // 动画长度
            float _AnimLen;

            
            v2f vert(appdata v, uint vid: SV_VertexID)
            {
                UNITY_SETUP_INSTANCE_ID(v);

                float f = _Time.y / _AnimLen;
                // fmod返回余数
                fmod(f, 1.0);

                float animMap_x = (vid + 0.5) * _AnimMap_TexelSize.x;
                float animMap_y = f;

                // 返回贴图中坐标位置
                float4 pos = tex2Dlod(_AnimMap, float4(animMap_x, animMap_y, 0, 0));

                v2f o;
                o.uv = TRANSFORM_TEX(v.uv, _MainTex);
                o.vertex = UnityObjectToClipPos(pos);
                return o;
            }
            
            fixed4 frag(v2f i): SV_Target
            {
                fixed4 col = tex2D(_MainTex, i.uv);
                return col;
            }
            ENDCG
            
        }
    }
}