// this code is under MIT License, by Robert Yang + others (credits in comments)
// a lot of this is based on http://wiki.unity3d.com/index.php?title=SkinnedMeshCombiner
// but I removed the atlasing stuff because I don't need it

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

// please don't ask me for help with this... this is more for learning purposes, and it's not really an easy-to-use Asset Store thing
// I have a lot of hacks specific for my uses... again, this is for devs doing similar things in their own code
public class SkinnedMeshCombiner : MonoBehaviour {
	public SkinnedMeshRenderer[] smRenderers; // you can leave this empty if you want it autodetect SMRs in children

	public LODGroup lod; // leave this undefined, legacy behavior... I used to bake to different LODs.
	public int bakeToLOD = 0;

	public Bounds boundsOverride; // this is for stuff like Mixamo hair, where I find I often have to type in my own Bounds or else the hair disappears randomly because of weird bounds that don't calculate well

	public Transform rootBoneCombine; // leave this undefined... this is for legacy behavior
	bool runInAwakeInstead = true; // I like having it run in Awake() so that my other scripts can do all their init stuff in Start() and they're guaranteed to be OK
	public bool ignoreBlendshapes = false;

	[SerializeField] Material[] randomMaterials; // leave empty if you don't want it to randomly pick a material for the final combined SMR

	[SerializeField] bool doRestoreBindPose = false;

	void Awake () {
		if ( runInAwakeInstead && this.enabled )
			Combine();
	}

	void Start () {
		if ( !runInAwakeInstead )
			Combine();
		transform.localPosition = Vector3.zero;
		transform.localRotation = Quaternion.identity;
	}

	// based on some code by JoeStrout https://forum.unity3d.com/threads/mesh-bindposes.383752/
	public void RestoreBindPose () {
		// if you're using Mixamo models, every SkinnedMeshRenderer has a fraction of the full bind pose data...
		// so first we have to search through all of the SMRs and combine all the bindpose data into a dictionary
		var	smRenderers2 = transform.parent.GetComponentsInChildren<SkinnedMeshRenderer>().OrderBy( smr => smr.bones.Length ).ToArray();

		Dictionary<Transform, Matrix4x4> bindPoseMap = new Dictionary<Transform, Matrix4x4>();

		foreach ( var smr in smRenderers2 ) {
			for( int i=0; i<smr.bones.Length; i++) {
				if ( !bindPoseMap.ContainsKey( smr.bones[i] ) ) {
					bindPoseMap.Add( smr.bones[i], smr.sharedMesh.bindposes[i] );
				}
			}
		}
	
		// based on data, now move the bones based on the bindPoseMap
		foreach( var kvp in bindPoseMap ) {
			Transform boneTrans = kvp.Key;
			Matrix4x4 bindPose = kvp.Value;

			// Recreate the local transform matrix of the bone
			Matrix4x4 localMatrix = bindPoseMap.ContainsKey(boneTrans.parent) ? (bindPose * bindPoseMap[boneTrans.parent].inverse).inverse : bindPose.inverse;
			// Recreate local transform from that matrix
			boneTrans.localPosition = localMatrix.MultiplyPoint (Vector3.zero);
			boneTrans.localRotation = Quaternion.LookRotation (localMatrix.GetColumn (2), localMatrix.GetColumn (1));
			boneTrans.localScale = new Vector3 (localMatrix.GetColumn (0).magnitude, localMatrix.GetColumn (1).magnitude, localMatrix.GetColumn (2).magnitude);
		}
		Debug.Log("Reset " + bindPoseMap.Count + " bones to bind pose");
	}

	// internal vars for blendshapes and boneweights etc
	class BlendShapeFrame {
		public string name;
		public float weight;
		public List<Vector3> deltaVertices = new List<Vector3>();
		public List<Vector3> deltaNormals = new List<Vector3>();
		public List<Vector3> deltaTangents = new List<Vector3>();
	}
		
	static List< List<BlendShapeFrame> > blendFrames = new List< List<BlendShapeFrame> >();
	static string[] blendshapeNames = new string[] { "Blink_Left", "Blink_Right", "BrowsDown_Left", 
		"BrowsDown_Right", "BrowsIn_Left", "BrowsIn_Right", "BrowsOuterLower_Left", "BrowsOuterLower_Right", 
		"BrowsUp_Left", "BrowsUp_Right", "CheekPuff_Left", "CheekPuff_Right", "EyesWide_Left", 
		"EyesWide_Right", "Frown_Left", "Frown_Right", "JawBackward", "JawForward", 
		"JawRotateY_Left", "JawRotateY_Right", "JawRotateZ_Left", "JawRotateZ_Right", "Jaw_Down", 
		"Jaw_Left", "Jaw_Right", "Jaw_Up", "LowerLipDown_Left", "LowerLipDown_Right", 
		"LowerLipIn", "LowerLipOut", "Midmouth_Left", "Midmouth_Right", "MouthDown", 
		"MouthNarrow_Left", "MouthNarrow_Right", "MouthOpen", "MouthUp", "MouthWhistle_NarrowAdjust_Left", 
		"MouthWhistle_NarrowAdjust_Right", "NoseScrunch_Left", "NoseScrunch_Right", "Smile_Left", "Smile_Right", 
		"Squint_Left", "Squint_Right", "TongueUp", "UpperLipIn", "UpperLipOut", 
		"UpperLipUp_Left", "UpperLipUp_Right" };
	static Vector3[] deltaVertices, deltaNormals, deltaTangents;
	static List<BoneWeight> boneWeights = new List<BoneWeight>();
	static List<CombineInstance> combineInstances = new List<CombineInstance>();
	static BoneWeight[] meshBoneweight;

	// this is the main function
	// you could probably optimize it if you really had to
	void Combine() {        
		// the initial state of the model matters a lot
		// so let's turn off the animator (just in case) and ideally try to reset the model back to bindpose
		var anim = GetComponentInParent<Animator>();
		anim.enabled = false;

		if (smRenderers.Length == 0)
			smRenderers = GetComponentsInChildren<SkinnedMeshRenderer>();

		if ( doRestoreBindPose ) 
			{ RestoreBindPose(); }

		// I use CP_SSSSS on GitHub for my skin subsurface scattering 
		bool hasSSS = OnlyHighQuality.usingHighQuality && smRenderers[0].GetComponent<CP_SSSSS_Object>() != null;
		Debug.Log( "debug: " + smRenderers[0].name + " Sss is " + hasSSS.ToString() );
		Color sssColor = hasSSS ? smRenderers[0].GetComponent<CP_SSSSS_Object>().subsurfaceColor : Color.clear;

		bool doBlendshapes = !ignoreBlendshapes && smRenderers[0].sharedMesh.blendShapeCount > 0;

		Vector3 oldPos = transform.position;
		Quaternion oldRot = transform.rotation;

		// reset the model to 0,0,0 (temporarily) so we can work easier with it
		transform.position = Vector3.zero;
		transform.rotation = Quaternion.identity;

		List<Transform> bones = new List<Transform>();        
		boneWeights.Clear();      
		combineInstances.Clear();
	
		Material mat = smRenderers[0].material;
		int numSubs = 0;

		// destroy and delete unused mesh renderers when we're done with them
		foreach(SkinnedMeshRenderer smr in smRenderers) {
			if ( !smr.gameObject.activeSelf ) {
				Destroy( smr.gameObject );
				continue;
			}
			numSubs += smr.sharedMesh.subMeshCount;
		}
		smRenderers = smRenderers.Where( x => x.gameObject.activeSelf ).ToArray(); // THIS IS WASTEFUL, I DON'T CARE, I DON'T
		
		int[] meshIndex = new int[numSubs];
		int boneOffset = 0;

		// bone map stuff
		var boneMap = new Dictionary<string, int>();
		var allBones = new List<Transform>();
		if (rootBoneCombine != null ) {
				allBones = rootBoneCombine.GetComponentsInChildren<Transform>().ToList();
			for ( int i=0; i<allBones.Count; i++ ) {
				boneMap.Add( allBones[i].name, i );
			}
		}

		blendFrames.Clear();

		var castShadows = smRenderers[0].shadowCastingMode;
		var receiveShadows = smRenderers[0].receiveShadows;
		// "BODY" MUST BE THE FIRST SMRENDERER in your smRenderer array, BECAUSE IT HAS *EVERY* BLENDSHAPE AND BLEND FRAME
		for( int s = 0; s < smRenderers.Length; s++ ) {
			SkinnedMeshRenderer smr = smRenderers[s];  

			var sMesh = smr.sharedMesh;

			// blend shape stuff
			// mixamo meshes sometimes have different blendshape counts and indices (WTF???)
			bool doNameSearch = s >= 1 && sMesh.blendShapeCount > 0 && sMesh.blendShapeCount < 50;

			for ( int shapeIndex=0; shapeIndex<50 && doBlendshapes; shapeIndex++ ) { // hardcoded to Mixamo's 50 because even meshes without blendshapes will need zeroed deltas when they get merged

				// runs only on first mesh, allocates space and sets up structure
				if ( blendFrames.Count == shapeIndex ) {
					var shapeName = "Facial_Blends." + blendshapeNames[shapeIndex];
					var frameCount = sMesh.GetBlendShapeFrameCount( shapeIndex );
					blendFrames.Add( new List<BlendShapeFrame>() );
					for ( int frameIndex = 0; frameIndex<frameCount; frameIndex++) {
						var newFrame = new BlendShapeFrame();
						newFrame.name = shapeName;
						newFrame.weight = sMesh.GetBlendShapeFrameWeight(shapeIndex,frameIndex);
						blendFrames[shapeIndex].Add( newFrame );
					}

				}

				// runs on every mesh, adds the deltas for each frame... if there are no deltas, then it adds empty deltas
				for ( int frameIndex = 0; frameIndex < blendFrames[shapeIndex].Count; frameIndex++ ) {
					deltaVertices = new Vector3[ sMesh.vertexCount ];
					deltaNormals = new Vector3[ sMesh.vertexCount ];
					deltaTangents = new Vector3[ sMesh.vertexCount ];

					// search by blendshape name sometimes, because sometimes some meshes don't have all the blendshapes
					int actualShapeIndex = shapeIndex;
					if ( doNameSearch ) {
						actualShapeIndex = -1;
						if ( sMesh.blendShapeCount > shapeIndex ) {
							string groupName = sMesh.GetBlendShapeName(shapeIndex);
							for( int i=0; i<blendFrames.Count; i++) {
								if ( blendFrames[i][0].name == groupName ) { // does this iterated frameGroup's shapeName equal this current frameGroup's shapeName?
									actualShapeIndex = i;
									break;
								}
							}
						}
					}
					// only fetch deltas if there's a blendShapeFrame there
					int actualFrameCount = actualShapeIndex > -1 && sMesh.blendShapeCount > actualShapeIndex ? sMesh.GetBlendShapeFrameCount( actualShapeIndex ) : 0;

					if ( actualShapeIndex >= 0 && frameIndex < actualFrameCount ) {
						sMesh.GetBlendShapeFrameVertices( shapeIndex, frameIndex, deltaVertices, deltaNormals, deltaTangents );
					}

					var frame = blendFrames[shapeIndex][frameIndex];
					frame.deltaVertices.AddRange( deltaVertices );
					frame.deltaNormals.AddRange( deltaNormals );
					frame.deltaTangents.AddRange( deltaTangents );
				}
			} // end blendshape merging

			// bone weight stuff
			meshBoneweight = smr.sharedMesh.boneWeights;
			var smrBones = smr.bones;
			
			// May want to modify this if the renderer shares bones as unnecessary bones will get added.
			foreach( BoneWeight bw in meshBoneweight ) {
				BoneWeight bWeight = bw;

				if (boneMap.Count == 0 ) {
					bWeight.boneIndex0 += boneOffset;
					bWeight.boneIndex1 += boneOffset;
					bWeight.boneIndex2 += boneOffset;
					bWeight.boneIndex3 += boneOffset;   
				} else {
					bWeight.boneIndex0 = boneMap[ smrBones[bWeight.boneIndex0].name ];
					bWeight.boneIndex1 = boneMap[ smrBones[bWeight.boneIndex1].name ];
					bWeight.boneIndex2 = boneMap[ smrBones[bWeight.boneIndex2].name ];
					bWeight.boneIndex3 = boneMap[ smrBones[bWeight.boneIndex3].name ];
				}
				
				boneWeights.Add( bWeight );
			}
			boneOffset += smr.bones.Length;

			if ( boneMap.Count == 0 ) {
				Transform[] meshBones = smr.bones;
				foreach( Transform bone in meshBones )
					bones.Add( bone );
			}

			CombineInstance ci = new CombineInstance();
			ci.mesh = smr.sharedMesh;
			meshIndex[s] = ci.mesh.vertexCount;
			ci.transform = smr.transform.localToWorldMatrix;
			combineInstances.Add( ci );
			
			Object.Destroy( smr.gameObject );
		}
		
		List<Matrix4x4> bindposes = new List<Matrix4x4>();

		if (allBones.Count > 0)
			bones = allBones;
		
		for( int b = 0; b < bones.Count; b++ ) {
			bindposes.Add( bones[b].worldToLocalMatrix * transform.worldToLocalMatrix );
		}

		// begin constructing new mesh
		var newMesh = new Mesh();
		newMesh.CombineMeshes( combineInstances.ToArray(), true, true );

		// add blendshape frames
		for( int frameGroup=0; frameGroup<blendFrames.Count; frameGroup++) {
			for ( int frame=0; frame<blendFrames[frameGroup].Count; frame++ ) {
				var thisFrame = blendFrames[frameGroup][frame];
				newMesh.AddBlendShapeFrame( thisFrame.name, thisFrame.weight, thisFrame.deltaVertices.ToArray(), thisFrame.deltaNormals.ToArray(), thisFrame.deltaTangents.ToArray() );
			}
		}

		SkinnedMeshRenderer r = GetComponent<SkinnedMeshRenderer>();
		if ( r == null ) { r = gameObject.AddComponent<SkinnedMeshRenderer>(); }

		r.sharedMesh = newMesh;

		// workaround for Unity bug 824384 where blendshapes would not update properly at runtime... thanks Julius!
		SkinnedMeshRenderer smrr = r;
		Mesh m = smrr.sharedMesh;
		r.sharedMesh = m;
		// Debug.Log( "debug blendshape check: " + r.sharedMesh.blendShapeCount.ToString() );


		r.shadowCastingMode = castShadows;
		r.receiveShadows = receiveShadows;

		r.sharedMaterial = mat;
		
		r.bones = bones.ToArray();
		r.sharedMesh.boneWeights = boneWeights.ToArray();
		r.sharedMesh.bindposes = bindposes.ToArray();

		r.sharedMesh.RecalculateBounds();
//		if ( boundsOverride.center.sqrMagnitude > 0.001f ) {
//			//r.sharedMesh.bounds = boundsOverride;
//			r.localBounds = boundsOverride;
//		} else {
//			transform.position = oldPos;
//			transform.rotation = oldRot;
//		}

		// commit to LODs, if we're using them (more or less disabled since august 2016 though)
		if ( lod != null ) {
			var lods = lod.GetLODs();
			var rends = lods[bakeToLOD].renderers.Where( x => x != null).ToList();
			rends.Add( r );
			lods[bakeToLOD].renderers = rends.ToArray();
			lod.SetLODs( lods );
			//r.sharedMesh.RecalculateBounds();
			lod.RecalculateBounds();
		}

		// if defined, pick a random material
		if ( randomMaterials != null && randomMaterials.Length > 0 ) {
			r.sharedMaterial = randomMaterials[ Random.Range(0, randomMaterials.Length) ];
		}

		// did we have SSS? if so, assign it
		if ( hasSSS ) {
			var newSSS = r.gameObject.AddComponent<CP_SSSSS_Object>();
			newSSS.maskSource = CP_SSSSS_MaskSource.wholeObject;
			newSSS.subsurfaceColor = sssColor;
		}

		// turn animator back on
		anim.enabled = true;
	}
}