-- **************************************************
-- Provide Moho with the name of this script object
-- **************************************************

ScriptName = "AE_MergeSkeletons"

-- **************************************************
-- General information about this script
-- **************************************************

AE_MergeSkeletons = {}

function AE_MergeSkeletons:Name()
	return "Merge Skeletons"
end

function AE_MergeSkeletons:Version()
	return "2.0"
end

function AE_MergeSkeletons:UILabel()
	return "Merge Skeletons"
end

function AE_MergeSkeletons:Creator()
	return "Alexandra Evseeva"
end

function AE_MergeSkeletons:Description()
	return "Merge current skeleton with parent one"
end




-- **************************************************
-- Recurring values
-- **************************************************

-- AE_MergeSkeletons.value1 = false

-- **************************************************
-- Is Enabled
-- **************************************************

function AE_MergeSkeletons:IsEnabled(moho)
	return true
end


-- **************************************************
-- The guts of this script
-- **************************************************

function AE_MergeSkeletons:Run(moho)

	moho.document:PrepUndo(moho.layer)
	moho.document:SetDirty()
	
	local sourceBoneLayer = moho.layer
	local sourceSkeleton = moho:Skeleton()
   
	local targetSkeleton = sourceBoneLayer:ControllingSkeleton()
	local targetBoneLayer = sourceBoneLayer:Parent()
	
	if sourceSkeleton and not targetSkeleton then 
		-- if no parent skeleton, perform unmerge instead of merge
		moho:SetCurFrame(0)
		local bones2detach, layers2detach, actions2detach = self:Bones2LayersAndActions(moho)
		self:UnMergeSkeleton(moho, bones2detach, layers2detach, actions2detach)
		return 
	end

	nextBoneLayer = moho:LayerAsBone(targetBoneLayer)
	while not nextBoneLayer or not nextBoneLayer:Skeleton() == targetSkeleton do 
		targetBoneLayer = targetBoneLayer:Parent()
		nextBoneLayer = moho:LayerAsBone(targetBoneLayer)
	end	
		
	moho:SetSelLayer(targetBoneLayer)

	local boneMap = {}
	for boneID=0, sourceSkeleton:CountBones()-1 do
		local sourceBone = sourceSkeleton:Bone(boneID)
		local targetBone = targetSkeleton:AddBone(0)
		boneMap[sourceBone] = targetBone
		self:CopyBoneProps(sourceBone, targetBone)
		targetBone:SetName(sourceBone:Name())
	end
	for k, v in pairs(boneMap) do
		if(k.fParent>-1) then
			v.fParent = targetSkeleton:BoneID(boneMap[sourceSkeleton:Bone(k.fParent)])
		end
		if(k.fAngleControlParent>-1) then
			v.fAngleControlParent = targetSkeleton:BoneID(boneMap[sourceSkeleton:Bone(k.fAngleControlParent)])
			v.fAngleControlScale = k.fAngleControlScale
		end
		if(k.fPosControlParent>-1) then
			v.fPosControlParent = targetSkeleton:BoneID(boneMap[sourceSkeleton:Bone(k.fPosControlParent)])
			v.fPosControlScale = k.fPosControlScale
		end
		if(k.fScaleControlParent>-1) then
			v.fScaleControlParent = targetSkeleton:BoneID(boneMap[sourceSkeleton:Bone(k.fScaleControlParent)])
			v.fScaleControlScale = k.fScaleControlScale
		end
		if(k.fTargetBone:GetValue(0)>-1) then
			v.fTargetBone:SetValue(0, targetSkeleton:BoneID(boneMap[sourceSkeleton:Bone(k.fTargetBone:GetValue(0))]))
		end	  
	end

	local layersMap = {}
	local oldGroup = moho:LayerAsGroup(sourceBoneLayer)
	local newGroup = moho:LayerAsGroup(targetBoneLayer)
	self:StoreBindings(moho, layersMap, oldGroup, sourceSkeleton)
	while oldGroup:CountLayers()>0 do
		local childLayer = oldGroup:Layer(0)
		moho:PlaceLayerInGroup(childLayer, newGroup, false)
	end
	local boneIDMap = {}
	for k, v in pairs(boneMap) do
	  boneIDMap[sourceSkeleton:BoneID(k)] = targetSkeleton:BoneID(v)
	end
	self:RestoreBindings(moho, layersMap, boneIDMap) 

	for k, v in pairs(boneMap) do
       local nextActionName = k:Name() 
       if sourceBoneLayer:HasAction(nextActionName) then 
          sourceBoneLayer:ActivateAction(nextActionName)
          local animAngle = k.fAnimAngle
          print(nextActionName, " keys: ", animAngle:CountKeys())
          targetBoneLayer:ActivateAction(nextActionName)
          for f=0, animAngle:Duration() do
             if animAngle:HasKey(f) then
                v.fAnimAngle:SetValue(f,animAngle:GetValue(f))
             end
            for otherBoneID = 0, sourceSkeleton:CountBones()-1 do 
                local otherBone = sourceSkeleton:Bone(otherBoneID)
                local otherTargetBone = boneMap[otherBone]
                if otherBone.fAnimAngle:HasKey(f) then
                   otherTargetBone.fAnimAngle:SetValue(f, otherBone.fAnimAngle:GetValue(f))
                end
                if otherBone.fAnimPos:HasKey(f) then
                   otherTargetBone.fAnimPos:SetValue(f, otherBone.fAnimPos:GetValue(f))
                end 
                if otherBone.fAnimScale:HasKey(f) then
                   otherTargetBone.fAnimScale:SetValue(f, otherBone.fAnimScale:GetValue(f))
                end
                if otherBone.fFlipH:HasKey(f) then
                   otherTargetBone.fFlipH:SetValue(f, otherBone.fFlipH:GetValue(f))
                end 
                if otherBone.fFlipV:HasKey(f) then
                   otherTargetBone.fFlipV:SetValue(f, otherBone.fFlipV:GetValue(f))
                end              
            end 
          end
       end
       nextActionName = k:Name() .. " 2"
       if sourceBoneLayer:HasAction(nextActionName) then 
          sourceBoneLayer:ActivateAction(nextActionName)
          local animAngle = k.fAnimAngle
          print(nextActionName, " keys: ", animAngle:CountKeys())
          targetBoneLayer:ActivateAction(nextActionName)
          for f=0, animAngle:Duration() do
             if animAngle:HasKey(f) then
                v.fAnimAngle:SetValue(f,animAngle:GetValue(f))
             end
            for otherBoneID = 0, sourceSkeleton:CountBones()-1 do 
                local otherBone = sourceSkeleton:Bone(otherBoneID)
                local otherTargetBone = boneMap[otherBone]
                if otherBone.fAnimAngle:HasKey(f) then
                   otherTargetBone.fAnimAngle:SetValue(f, otherBone.fAnimAngle:GetValue(f))
                end
                if otherBone.fAnimPos:HasKey(f) then
                   otherTargetBone.fAnimPos:SetValue(f, otherBone.fAnimPos:GetValue(f))
                end 
                if otherBone.fAnimScale:HasKey(f) then
                   otherTargetBone.fAnimScale:SetValue(f, otherBone.fAnimScale:GetValue(f))
                end
                if otherBone.fFlipH:HasKey(f) then
                   otherTargetBone.fFlipH:SetValue(f, otherBone.fFlipH:GetValue(f))
                end 
                if otherBone.fFlipV:HasKey(f) then
                   otherTargetBone.fFlipV:SetValue(f, otherBone.fFlipV:GetValue(f))
                end              
            end 
          end
       end       
   end
   
   sourceBoneLayer:ActivateAction("")
   
   while sourceSkeleton:CountBones()>0 do
      sourceSkeleton:DeleteBone(0,0)
   end
   moho.layer:ActivateAction("")
	
end

function AE_MergeSkeletons:CopyBoneProps(src, trg)
   trg.fLength = src.fLength  
   trg.fAnimPos:SetValue(0, src.fPos)        
   trg.fAnimAngle:SetValue(0, src.fAngle)    
   trg.fAnimScale:SetValue(0, src.fScale) 
   trg:ShowLabel(src:IsLabelShowing())
   trg.fShy = src.fShy
   trg.fStrength = src.fStrength
   trg.fFixedAngle = src.fFixedAngle
   trg.fMinConstraint = src.fMinConstraint
   trg.fMaxConstraint = src.fMaxConstraint
   trg.fConstraints = src.fConstraints
   trg.fIgnoredByIK = src.fIgnoredByIK

end

function AE_MergeSkeletons:StoreBindings(moho, storeArray, rootGroup, sourceSkeleton)
  for childID = 0, rootGroup:CountLayers()-1 do
    local childLayer = rootGroup:Layer(childID)
	if childLayer:ControllingSkeleton() == sourceSkeleton then
		local storeObject = {}
		storeObject["layer"] = childLayer:LayerParentBone()
		if childLayer:LayerType() == MOHO.LT_VECTOR then 
		  local vertices = {}
		  moho:SetSelLayer(childLayer)
		  local mesh = moho:Mesh()
		  for pointID = 0, mesh:CountPoints()-1 do
			vertices[pointID] = mesh:Point(pointID).fParent
		  end
		  storeObject["vertices"] = vertices
		end  
		storeArray[childLayer] = storeObject
		if childLayer:IsGroupType() then
		  local childGroup = moho:LayerAsGroup(childLayer)
		  self:StoreBindings(moho, storeArray, childGroup, sourceSkeleton)
		end
	end
  end
end

function AE_MergeSkeletons:StoreBindings2(moho, storeArray, layer, sourceSkeleton)
	if layer:ControllingSkeleton() ~= sourceSkeleton then return end
	local storeObject = {}
	storeObject["layer"] = layer:LayerParentBone()
	if layer:LayerType() == MOHO.LT_VECTOR then 
		local vertices = {}
		local mesh = moho:LayerAsVector(layer):Mesh()
		for pointID = 0, mesh:CountPoints()-1 do
		vertices[pointID] = mesh:Point(pointID).fParent
		end
		storeObject["vertices"] = vertices
	end  
	storeArray[layer] = storeObject
	if layer:IsGroupType() then
		local childGroup = moho:LayerAsGroup(layer)
		for nch = 0, childGroup:CountLayers() - 1 do 
			self:StoreBindings2(moho, storeArray, childGroup:Layer(nch), sourceSkeleton)
		end
	end
end 

function AE_MergeSkeletons:RestoreBindings(moho, storeArray, boneIDMap)
  for k, v in pairs(storeArray) do
    if v.layer > -1 then
      k:SetLayerParentBone(boneIDMap[v.layer]) 
    end
    if k:LayerType() == MOHO.LT_VECTOR then
      moho:SetSelLayer(k)
      local mesh = moho:Mesh()
      for pointID = 0, mesh:CountPoints()-1 do
        if v.vertices[pointID] > -1 then
          local newBone = boneIDMap[v.vertices[pointID]]
          mesh:Point(pointID).fParent = newBone
        end
      end
    end 
  end
end

function AE_MergeSkeletons:Bones2LayersAndActions(moho)
	-- get current skeleton selected bones and collect:
	-- layers, driven with selected bones
	-- actions, driving only these layers and nothing else
	-- additional bones, driving only these layers and nothing else
	-- additional smartbones, driving only these layers and bones and nothing else
	
	local skel = moho:Skeleton()
	local bones2detach = {}
	for b = 0, skel:CountBones() - 1 do
		if skel:Bone(b).fSelected then table.insert(bones2detach, skel:Bone(b)) end
	end
	
	local layers2detach = {}
	local mainGroup = moho:LayerAsGroup(moho.layer)
	for nchld = 0, mainGroup:CountLayers() - 1 do
		local nextChild = mainGroup:Layer(nchld)
		local doesBelong, additionalBones = self:CheckBoneBelong(moho, skel, bones2detach, nextChild)
		if doesBelong then 
			table.insert(layers2detach, nextChild)
			for i, bone in pairs(additionalBones) do
				if not bone.fSelected then
					bone.fSelected = true
					table.insert(bones2detach, bone)
				end
			end
		end
	end
	
	for nchld = 0, mainGroup:CountLayers() - 1 do
		local nextChild = mainGroup:Layer(nchld)
		local doesBelong, additionalBones = self:CheckBoneBelong(moho, skel, bones2detach, nextChild)
		if doesBelong then 
			if #additionalBones > 0 then 
				local boneList = ""
				for j, bone in pairs(additionalBones) do
					boneList = boneList .. " " .. bone:Name()
				end
				local warning = "These bones do not drive some separate layers to detach"
				local waring2 = "additional bones used are, for example: "
				LM.GUI.Alert(LM.GUI.ALERT_WARNING, warning, warning2, boneList)
				return bones2detach, layers2detach, {}
			end
			local newLayer = true
			for i, layer in pairs(layers2detach) do
				if nextChild == layer then 
					newLayer = false
					break
				end
			end
			if newLayer then table.insert(layers2detach, nextChild) end
		end
	end

	-- collect actions used only in collected layers
	local actions2detach = {}
	for a = 0, moho.layer:CountActions() - 1 do
		local actionName = moho.layer:ActionName(a)
		local foundInCollected = false
		local foundInOthers = false
		for nchld = 0, mainGroup:CountLayers() - 1 do
			local layer = mainGroup:Layer(nchld)
			local isLayerCollected = false
			for j, colLayer in pairs(layers2detach) do
				if colLayer == layer then 
					isLayerCollected = true
					break
				end
			end	
			if layer:ActionDuration(actionName) > 0 then 
				if isLayerCollected then foundInCollected = true 
				else foundInOthers = true 
				end
			end
			if foundInCollected and foundInOthers then break end
		end
		if foundInCollected and not foundInOthers then table.insert(actions2detach, actionName) end
	end
	-- collect smartbones driving collected actions (inserting into bones array too)
	for i, actionName in pairs(actions2detach) do
		if moho.layer:IsSmartBoneAction(actionName) then
			local smartBoneName = actionName
			if string.sub(actionName, -2) == " 2" then 
				smartBoneName = string.sub(actionName, 1, -3)
			end
			local smartbone = skel:BoneByName(smartBoneName)
			if smartbone then 
				local newBone = true
				for j, bone in pairs(bones2detach) do
					if bone == smartbone then 
						newBone = false
						break
					end
				end
				if newBone then table.insert(bones2detach, smartbone) end
			end
		end
	end
	
	return bones2detach, layers2detach, actions2detach
end

function AE_MergeSkeletons:CheckBoneBelong(moho, skel, bones, layer)
	-- flexi binding without restrictions
	if layer:LayerParentBone() == -2 then return false, {} end
	-- layer belongs to a single bone	
	if layer:LayerParentBone() > -1 then
		for i, bone in pairs(bones) do
			if skel:BoneID(bone) == layer:LayerParentBone() then return true, {} end
		end
		return false, {}
	end
	-- layer belongs to multiple bones 
	if layer:LayerParentBone() == -3 then
		local flexiBones = {}
		local doesBelong = false
		for b = 0, skel:CountBones() - 1 do
			if layer:IsIncludedInFlexiBoneSubset(b) then 
				local thisBone = false
				for i, bone in pairs(bones) do
					if skel:BoneID(bone) == b then 
						doesBelong = true
						thisBone = true
						break
					end
				end
				if not thisBone then table.insert(flexiBones, skel:Bone(b)) end
			end
		end
		if doesBelong then return true, flexiBones end
		return false, {}
	end
	-- layer is vector
	if moho:LayerAsVector(layer) then
		local otherBones = {}
		local doesBelong = false
		local mesh = moho:LayerAsVector(layer):Mesh()
		for p = 0, mesh:CountPoints()-1 do
			local boneID = mesh:Point(p).fParent			
			if boneID > -1 then
				local thisBone = false
				for i, bone in pairs(bones) do
					if skel:BoneID(bone) == boneID then 
						doesBelong = true
						thisBone = true
						break
					end
				end
				if not thisBone then table.insert(otherBones, skel:Bone(boneID)) end
			end
		end
		if doesBelong then return true, otherBones end
		return false, {}
	end
	-- layer is group
	if moho:LayerAsGroup(layer) then
		local doesBelong = false
		local otherBones = {}
		local aGroup = moho:LayerAsGroup(layer)
		for nch = 0, aGroup:CountLayers() - 1 do
			local nextChild = aGroup:Layer(nch)
			if nextChild:ControllingSkeleton() == skel then
				local nextBelong, nextBones = self:CheckBoneBelong(moho, skel, bones, nextChild)
				--print("Checking ", nextChild:Name(), "... ", tostring(nextBelong))
				if nextBelong then doesBelong = true end
				for i, bone in pairs(nextBones) do
					local allwaysCollected = false
					for j, clbone in pairs(otherBones) do
						if clbone == bone then 
							allwaysCollected = true
							break
						end
					end
					if not allwaysCollected then table.insert(otherBones, bone) end
				end
			end
		end
		if doesBelong then return true, otherBones end
		return false, {}
	end
	return false, {}
end
function AE_MergeSkeletons:UnMergeSkeleton(moho, bones2detach, layers2detach, actions2detach)
	-- duplicate current skeleton layer
	local sourceLayer = moho.layer
	local sourceSkeleton = moho:LayerAsBone(sourceLayer):Skeleton()
	local targetLayer = moho:DuplicateLayer(sourceLayer)
	local targetSkeleton = moho:LayerAsBone(targetLayer):Skeleton()
	if #layers2detach > 0 then targetLayer:SetName(layers2detach[1]:Name()) end
	
	-- delete layers from target	
	local otherLayers = {}
	local sourceGroup = moho:LayerAsGroup(sourceLayer)
	local targetGroup = moho:LayerAsGroup(targetLayer)
	for i = 0, sourceGroup:CountLayers() - 1 do
		local toCollect = true
		for j, layer in pairs(layers2detach) do
			if layer == sourceGroup:Layer(i) then
				toCollect = false
				break
			end
		end
		if toCollect then table.insert(otherLayers, targetGroup:Layer(i)) end
	end
	for i, layer in pairs(otherLayers) do moho:DeleteLayer(layer) end


	-- delete bones from target
	local otherBones = {}
	local boneMap = {}
	for b = 0, sourceSkeleton:CountBones() - 1 do
		local toCollect = true
		for i, bone in pairs(bones2detach) do
			if bone == sourceSkeleton:Bone(b) then
				toCollect = false
				break
			end
		end
		if toCollect then 
			table.insert(otherBones, targetSkeleton:Bone(b)) 
		else 
			boneMap[targetSkeleton:Bone(b)] = sourceSkeleton:Bone(b)
		end
	end

	for i, bone in pairs(otherBones) do 
		local boneID = targetSkeleton:BoneID(bone)
		for j = 0, targetGroup:CountLayers() - 1 do targetGroup:Layer(j):DeleteParentBone(boneID) end
		targetSkeleton:DeleteBone(boneID)
	end
	
	-- delete NOT USED actions from target
	local otherActions = {}
	for a = 0, targetLayer:CountActions() - 1 do
		local actionName = targetLayer:ActionName(a)
		targetLayer:ActivateAction(actionName)
		targetLayer:ActivateAction(nil)
		if targetLayer:ActionDuration(actionName) < 1 then
			table.insert(otherActions, actionName)
		end
	end
	for i, actionName in pairs(otherActions) do targetLayer:DeleteAction(actionName) end
	
	-- copy animations into root target bones from their source copies (for all actions)
	for b = 0, targetSkeleton:CountBones() - 1 do
		local nextBone = targetSkeleton:Bone(b)
		if nextBone.fAnimParent:GetValue(0) == -1 then
			self:FixBoneAnimation(moho, nextBone, boneMap[nextBone], sourceSkeleton, moho.document:EndFrame() + targetLayer:TotalTimingOffset())
			for a = 0, targetLayer:CountActions() - 1 do
				local actionName = targetLayer:ActionName(a)
				self:FixBoneAnimation(moho, nextBone, boneMap[nextBone], sourceSkeleton, targetLayer:ActionDuration(actionName), actionName)
			end
		end
	end
	
	-- delete layers, bones and actions from old skeleton
	local topNumber = -1
	for i = 0, sourceGroup:CountLayers()-1 do
		for j, layer in pairs(layers2detach) do
			if sourceGroup:Layer(i) == layer then topNumber = i + 1 end
		end
	end
	local topLayer = nil
	if topNumber < sourceGroup:CountLayers() then topLayer = sourceGroup:Layer(topNumber) end
	
	for i, layer in pairs(layers2detach) do moho:DeleteLayer(layer) end	
	for i, actionName in pairs(actions2detach) do sourceLayer:DeleteAction(actionName) end	
	for i, bone in pairs(bones2detach) do 
		local boneID = sourceSkeleton:BoneID(bone)
		for nch = 0, sourceGroup:CountLayers() - 1 do
			sourceGroup:Layer(nch):DeleteParentBone(boneID)
		end
		sourceSkeleton:DeleteBone(boneID)
	end
	
	-- place new into old (and reset new layer transformation)
	targetLayer.fTranslation:Clear()
	targetLayer.fScale:Clear()
	targetLayer.fRotationX:Clear()
	targetLayer.fRotationY:Clear()
	targetLayer.fRotationZ:Clear()
	targetLayer.fShear:Clear()
	targetLayer.fFlipH:Clear()
	targetLayer.fFlipV:Clear()
	local zeroVec = LM.Vector3:new_local()
	targetLayer.fTranslation:SetValue(0, zeroVec)
	targetLayer.fShear:SetValue(0, zeroVec)
	zeroVec:Set(1,1,1)
	targetLayer.fScale:SetValue(0, zeroVec)
	targetLayer.fRotationX:SetValue(0,0)
	targetLayer.fRotationY:SetValue(0,0)
	targetLayer.fRotationZ:SetValue(0,0)
	targetLayer.fFlipH:SetValue(0, false)
	targetLayer.fFlipV:SetValue(0, false)
	
	moho:PlaceLayerInGroup(targetLayer, sourceGroup, true)
	if topLayer then moho:PlaceLayerBehindAnother(targetLayer, topLayer) end
	moho:SetSelLayer(targetLayer)
end

function AE_MergeSkeletons:FixBoneAnimation(moho, trgBone, srcBone, srcSkel, duration, actionName)
	-- iterate every frame with pos or angle key
	local posChannel = trgBone.fAnimPos
	local angleChannel = trgBone.fAnimAngle
	if actionName then
		posChannel = posChannel:ActionByName(actionName)
		if posChannel then posChannel = moho:ChannelAsAnimVec2(posChannel) end
		angleChannel = angleChannel:ActionByName(actionName)
		if angleChannel then angleChannel = moho:ChannelAsAnimVal(angleChannel) end
	end
	for f = 1, duration do
		if (posChannel and posChannel:HasKey(f)) or (angleChannel and angleChannel:HasKey(f)) then
			--local theMatrix = AE_Utilities:GetGlobalBoneMatrix(moho, srcSkel, srcBone, f, actionName)
			--local pos, angle = AE_Utilities:Matrix2transform(theMatrix)
			local pos, angle = AE_Utilities:GetGlobalBonePRS(moho, srcSkel, srcBone, f, actionName)
			if posChannel then posChannel:SetValue(f, pos) end
			if angleChannel then angleChannel:SetValue(f, angle) end
		end
	end
end