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

ScriptName = "MR_SmartBoneFixer"

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

MR_SmartBoneFixer = {}

function MR_SmartBoneFixer:Name()
	return self:Localize('UILabel')
end

function MR_SmartBoneFixer:Version()
	return '1.0'
end

function MR_SmartBoneFixer:UILabel()
	return self:Localize('UILabel')
end

function MR_SmartBoneFixer:Creator()
	return 'Eugene Babich'
end

function MR_SmartBoneFixer:Description()
	return self:Localize('Description')
end

-- **************************************************
-- Is Relevant / Is Enabled
-- **************************************************

function MR_SmartBoneFixer:IsRelevant(moho)
	return true
end

function MR_SmartBoneFixer:IsEnabled(moho)
	return true
end

-- **************************************************
-- Keyboard/Mouse Control
-- **************************************************

function MR_SmartBoneFixer:OnMouseDown(moho, mouseEvent)
	
end

function MR_SmartBoneFixer:LoadPrefs(prefs)
	self.considerNewParentRotation = prefs:GetBool('MR_SmartBoneFixer.considerNewParentRotation', true)
	self.considerOldParentRotation = prefs:GetBool('MR_SmartBoneFixer.considerOldParentRotation', true)
end

function MR_SmartBoneFixer:SavePrefs(prefs)
	prefs:SetBool('MR_SmartBoneFixer.considerNewParentRotation', self.considerNewParentRotation)
	prefs:SetBool('MR_SmartBoneFixer.considerOldParentRotation', self.considerOldParentRotation)
end

function MR_SmartBoneFixer:ResetPrefs()
	MR_SmartBoneFixer.considerNewParentRotation = false
	MR_SmartBoneFixer.considerOldParentRotation = false
end

-- **************************************************
-- Recurring Values
-- **************************************************

MR_SmartBoneFixer.considerNewParentRotation = false
MR_SmartBoneFixer.considerOldParentRotation = false
MR_SmartBoneFixer.boneList = {}
MR_SmartBoneFixer.boneRefList = {}
MR_SmartBoneFixer.boneOrigPosList = {}
MR_SmartBoneFixer.boneOrigGlobalPosList = {}
MR_SmartBoneFixer.boneAngleList = {}
MR_SmartBoneFixer.boneAngleExtraList = {}
MR_SmartBoneFixer.boneAngleOffsetList = {}
MR_SmartBoneFixer.boneParentAngleList = {}
MR_SmartBoneFixer.boneParentsPosList = {}
MR_SmartBoneFixer.boneParentList = {}
MR_SmartBoneFixer.boneParentMatrixList = {}
MR_SmartBoneFixer.boneLayerId = nil
MR_SmartBoneFixer.status = 'No bones transformation collected yet.'
MR_SmartBoneFixer.fixedBonesList = {}
MR_SmartBoneFixer.fixedActionsList = {}
MR_SmartBoneFixer.totalBones = 0
MR_SmartBoneFixer.refKey = 1

-- **************************************************
-- settingsDialog
-- **************************************************

local settingsDialog = {}

settingsDialog.CONSIDER_NEW_PARENT_ROTATION = MOHO.MSG_BASE
settingsDialog.CONSIDER_OLD_PARENT_ROTATION = MOHO.MSG_BASE + 1

function settingsDialog:new()
    local d = LM.GUI.SimpleDialog(MR_SmartBoneFixer:Localize('UILabel'), settingsDialog)
    local l = d:GetLayout()

    d.considerNewParentRotationCheckbox = LM.GUI.CheckBox(MR_SmartBoneFixer:Localize('Consider new parent rotation'), d.CONSIDER_NEW_PARENT_ROTATION)
    l:AddChild(d.considerNewParentRotationCheckbox, LM.GUI.ALIGN_LEFT, 0)

    d.considerOldParentRotationCheckbox = LM.GUI.CheckBox(MR_SmartBoneFixer:Localize('Consider old parent rotation'), d.CONSIDER_OLD_PARENT_ROTATION)
    l:AddChild(d.considerOldParentRotationCheckbox, LM.GUI.ALIGN_LEFT, 0)
    return d
end

function settingsDialog:UpdateWidgets(moho)
    self.considerNewParentRotationCheckbox:SetValue(MR_SmartBoneFixer.considerNewParentRotation)
    self.considerOldParentRotationCheckbox:SetValue(MR_SmartBoneFixer.considerOldParentRotation)
end

function settingsDialog:OnOK(moho)
    MR_SmartBoneFixer.considerNewParentRotation = self.considerNewParentRotationCheckbox:Value()
    MR_SmartBoneFixer.considerOldParentRotation = self.considerOldParentRotationCheckbox:Value()
end

function settingsDialog:HandleMessage(msg)
    if msg == self.CONSIDER_NEW_PARENT_ROTATION then
		self.considerNewParentRotation = self.considerNewParentRotationCheckbox:Value()
    elseif msg == self.CONSIDER_OLD_PARENT_ROTATION then
        self.considerOldParentRotation = self.considerOldParentRotationCheckbox:Value()
    end
end

-- **************************************************
-- Tool Panel Layout
-- **************************************************

MR_SmartBoneFixer.GET_OLD_TRANSFORMATION = MOHO.MSG_BASE
MR_SmartBoneFixer.FIX_ACTIONS = MOHO.MSG_BASE + 1
MR_SmartBoneFixer.CLEAR = MOHO.MSG_BASE + 2

function MR_SmartBoneFixer:DoLayout(moho, layout)
	self.dlog = settingsDialog:new()
    self.settingsPopup = LM.GUI.PopupDialog(self:Localize('Settings'), false, 0)
    self.settingsPopup:SetDialog(self.dlog)
    layout:AddChild(self.settingsPopup, LM.GUI.ALIGN_LEFT, 0)
	
	self.getOldTransformationBtn = LM.GUI.Button(self:Localize('Get original transformation'), self.GET_OLD_TRANSFORMATION)
	layout:AddChild(self.getOldTransformationBtn, LM.GUI.ALIGN_LEFT, 0)
	
	self.fixActionsBtn = LM.GUI.Button(self:Localize('Fix Actions'), self.FIX_ACTIONS)
	layout:AddChild(self.fixActionsBtn, LM.GUI.ALIGN_LEFT, 0)
	
	self.clearBtn = LM.GUI.Button(self:Localize('Clear'), self.CLEAR)
	layout:AddChild(self.clearBtn, LM.GUI.ALIGN_LEFT, 0)
	
	layout:AddPadding(10)

    layout:AddChild(LM.GUI.Divider(true), LM.GUI.ALIGN_FILL)

    layout:AddPadding(10)
	
	self.statusText = LM.GUI.DynamicText(self:Localize('Status:'), 500)
    layout:AddChild(self.statusText, LM.GUI.ALIGN_LEFT, 0)
end

function MR_SmartBoneFixer:HandleMessage(moho, view, msg)
	if msg == self.GET_OLD_TRANSFORMATION then
		self:GetBoneTransformation(moho)
	elseif msg == self.FIX_ACTIONS then
		self:ApplyBoneTransformation(moho)
	elseif msg == self.CLEAR then
		self:Cancel(moho)	
	end
end

function MR_SmartBoneFixer:UpdateWidgets(moho)
	local isMarker = self:CheckMarker(moho)
	local curLayerId = moho.layer:UUID()
	if self.boneLayerId ~= curLayerId or not isMarker then
		self.fixActionsBtn:Enable(false)
	else 
		self.fixActionsBtn:Enable(true)
	end
	if moho.frame ~= 0 then
		self.getOldTransformationBtn:Enable(false)
	else
		self.getOldTransformationBtn:Enable(true)
	end
	if isMarker then
		self.clearBtn:Enable(true)
		self.getOldTransformationBtn:Enable(false)
	else
		self.clearBtn:Enable(false)
	end	
	if moho.layer:LayerType() ~= MOHO.LT_BONE or moho.document:CurrentDocAction() ~= "" then
		self.getOldTransformationBtn:Enable(false)
		self.fixActionsBtn:Enable(false)
	end
	self.statusText:SetValue('Status: '.. self.status)
end

function MR_SmartBoneFixer:Cancel(moho)
	moho.document:PrepUndo(moho.layer)
	moho.document:SetDirty()
	local markerChannels = moho.layer.fTimelineMarkers
	local framesToDel = {}
	local isNeedUpdate = false
	if markerChannels:Duration() > 0 then
		for i=1, markerChannels:CountKeys()-1 do
			local markerTime = markerChannels:GetKeyWhen(i)
			local markerText = markerChannels:GetValue(markerTime)
			if markerText == 'Do not delete or edit!' then
				self.refKey = markerTime
				table.insert(framesToDel, markerTime)
				isNeedUpdate = true
			end
		end
		for y in pairs(framesToDel) do
			markerChannels:DeleteKey(framesToDel[y])
			self:ClearFrame(moho, framesToDel[y])
			if framesToDel[y] > 1 then
				self:ClearFrame(moho, framesToDel[y]-1)
			end
		end
		self.status = 'Collected data cleared.'
		self:UpdateWidgets(moho)	
		if isNeedUpdate then
			self:HardUpdate(moho)
		end
	end
end

function MR_SmartBoneFixer:CheckMarker(moho)
	local found = false
	local markerChannels = moho.layer.fTimelineMarkers
	if markerChannels:Duration() > 0 then
		for i=1, markerChannels:CountKeys()-1 do
			local markerTime = markerChannels:GetKeyWhen(i)
			local markerText = markerChannels:GetValue(markerTime)
			if markerText == 'Do not delete or edit!' then
				found = true
				break
			end
		end	
	end		
	return found
end

function MR_SmartBoneFixer:FindLastKey(moho)
	local lastKey = 0
	local skel = moho:Skeleton()
	for i=0, skel:CountBones()-1 do 
		local myBone = skel:Bone(i)
		local channelAngle = myBone.fAnimAngle
		local channelPos = myBone.fAnimPos
		local channelScale = myBone.fAnimScale
		if channelAngle:Duration() > lastKey then
			lastKey = channelAngle:Duration()
		end
		if channelPos:Duration() > lastKey then
			lastKey = channelPos:Duration()
		end
		if channelScale:Duration() > lastKey then
			lastKey = channelScale:Duration()
		end
	end
	return lastKey
end

function MR_SmartBoneFixer:RejoinDimensionsInSelectedBones(moho)
	local skel = moho:Skeleton()
	local isAllowRejoin = false
	local isSuccessful = true
	local isIgnoreMainLine = false
	for i=0, skel:CountBones()-1 do 
		local myBone = skel:Bone(i)
		if myBone.fSelected then
			for act=0, myBone.fAnimPos:CountActions()-1 do
				local actName = myBone.fAnimPos:ActionName(act)
				local isSmartBone = moho.layer:IsSmartBoneAction(actName)
				if isSmartBone then
					local actionChannel = moho:ChannelAsAnimVec2(myBone.fAnimPos:Action(act))
					if actionChannel and actionChannel.AreDimensionsSplit then
						if actionChannel:AreDimensionsSplit() then
							if not isAllowRejoin then
								local ans = LM.GUI.Alert(LM.GUI.ALERT_WARNING, "Channels with separated dimensions were found.",
								"You can Rejoin Dimensions to continue or cancel operation.", "",  "Rejoin Dimensions", "Cancel","")
								if ans == 0 then
									isAllowRejoin = true
								elseif ans == 1 then
									isSuccessful = false
									return isSuccessful
								end	
							end	
							if isAllowRejoin then
								actionChannel:SplitDimensions(false)
							end	
						end
					end
				end	
			end
		end
	end
	return isSuccessful
end

function MR_SmartBoneFixer:GetBoneTransformation(moho)
	for k in pairs(self.boneList) do
		self.boneList[k] = nil
		self.boneOrigPosList[k] = nil
		self.boneOrigGlobalPosList[k] = nil
		self.boneAngleList[k] = nil
		self.boneAngleExtraList[k] = nil
		self.boneAngleOffsetList[k] = nil
		self.boneParentAngleList[k] = nil
		self.boneParentsPosList[k] = nil
		self.boneParentList[k] = nil
		self.boneParentMatrixList[k] = nil
	end
	
	for k in pairs(self.boneRefList) do
		self.boneRefList[k] = nil
	end
	
	local skel = moho:Skeleton()
	
	self.totalBones = 0
	local bonesFound = moho:CountSelectedBones(true)
	
	if bonesFound > 0 then
		moho.document:PrepUndo(moho.layer)
		moho.document:SetDirty()
		
		local curLayerId = moho.layer:UUID()
		if self.boneLayerId ~= curLayerId then
			self.refKey = 1
		end
		
		local markerChannels = moho.layer.fTimelineMarkers
		local framesToDel = {}
		local isNeedUpdate = false
		if markerChannels:Duration() > 0 then
			for i=1, markerChannels:CountKeys()-1 do
				local markerTime = markerChannels:GetKeyWhen(i)
				local markerText = markerChannels:GetValue(markerTime)
				if markerText == 'Do not delete or edit!' then
					self.refKey = markerTime
					table.insert(framesToDel, markerTime)
					isNeedUpdate = true
				end
			end
			for y in pairs(framesToDel) do
				markerChannels:DeleteKey(framesToDel[y])
				self:ClearFrame(moho, framesToDel[y])
				if framesToDel[y] > 1 then
					self:ClearFrame(moho, framesToDel[y]-1)
				end
			end	
			if isNeedUpdate then
				self:HardUpdate(moho)
			end
		end

		local lastFrame = self:FindLastKey(moho)
		
		if lastFrame >= 1 then
			self.refKey = lastFrame + 2
		else
			self.refKey = 1
		end
		
		for i=0, skel:CountBones()-1 do 
			table.insert(self.boneRefList,i)
		end
		
		if self.refKey > 1 then
			self:SetKey(moho, self.refKey -1, self.refKey -1)
		end
		self:SetKey(moho, self.refKey, 0)
		local keyInterp = MOHO.InterpSetting:new_local()
		keyInterp.tags = 1
		moho.layer.fTimelineMarkers:SetValue(self.refKey,'Do not delete or edit!')
		moho.layer.fTimelineMarkers:SetKeyInterp(self.refKey, keyInterp)
	end
	
	
	local isSuccessful = self:RejoinDimensionsInSelectedBones(moho)
	
	if not isSuccessful then
		self.status = 'Canceled.'
		self:UpdateWidgets(moho)
		self:ClearFrame(moho, self.refKey)
	
		if self.refKey > 1 then
			self:ClearFrame(moho, self.refKey-1)
		end
		self:HardUpdate(moho)
		return
	end
	
	for i=0, skel:CountBones()-1 do 
		local myBone = skel:Bone(i)
		self.totalBones = self.totalBones + 1
		if myBone.fSelected then
			table.insert(self.boneList, i)
			local parentBone = skel:Bone(myBone.fParent)
			
			if myBone.fParent >= 0 then
				local parentMatrix = LM.Matrix:new_local()
				parentMatrix:Set(parentBone.fRestMatrix)
				table.insert(self.boneParentList, myBone.fParent)
				table.insert(self.boneParentMatrixList, parentMatrix)
			else
				local emtyMatrix = LM.Matrix:new_local()
				emtyMatrix:Identity()
				table.insert(self.boneParentMatrixList, emtyMatrix)
				table.insert(self.boneParentList, -1)
			end
			
			local bonePos = LM.Vector2:new_local()		
			local parentBonePos = LM.Vector2:new_local()
			parentBonePos:Set(0,0)
			bonePos = myBone.fAnimPos:GetValue(0)
			local boneAngle = myBone.fAnimAngle:GetValue(0)
			local boneAngleExtra = myBone.fAnimAngle:GetValue(self.refKey)
			table.insert(self.boneOrigPosList, myBone.fAnimPos:GetValue(0))
			local boneParentAngle = 0
				
			if myBone.fParent >= 0 then
				boneParentAngle = parentBone.fAnimAngle:GetValue(0)
				parentBonePos:Set(parentBone.fAnimPos:GetValue(0))
				if parentBone.fParent >= 0 then
					local parentForParent = skel:Bone(parentBone.fParent)
					parentForParent.fRestMatrix:Transform(parentBonePos)
				end
				parentBone.fRestMatrix:Transform(bonePos)
			end
			
			table.insert(self.boneParentsPosList, parentBonePos)
			table.insert(self.boneOrigGlobalPosList, bonePos)
			table.insert(self.boneAngleList, boneAngle)
			table.insert(self.boneAngleExtraList, boneAngleExtra)
			table.insert(self.boneParentAngleList, boneParentAngle)
		end	
	end
	
	local layerName = moho.layer:Name()
	
	if bonesFound > 0 then
		self.boneLayerId = moho.layer:UUID()
		if bonesFound == 1 then
			self.status = 'Original transformation collected for 1 bone on layer ' .. layerName..'.'
		else
			self.status = 'Original transformation collected for ' .. bonesFound.. ' bones on layer ' .. layerName..'.'
		end	
	else
		self.boneLayerId = nil
		self.status = 'No bones was selected. Please select bones first.'
	end
	
	moho.view:DrawMe()
	moho:UpdateUI()
	moho.layer:UpdateCurFrame()
	self:UpdateWidgets(moho)
end

function MR_SmartBoneFixer:SwapTwoArrayKeys(array, key1, key2)
	local keyTmp1 = array[key1]
	local keyTmp2 = array[key2]
	return keyTmp2, keyTmp1
end

function MR_SmartBoneFixer:ApplyBoneTransformation(moho)
	self.status = 'Correction in progress. Please wait...'
	self:UpdateWidgets(moho)
	
	for q in pairs(self.fixedBonesList) do
		self.fixedBonesList[q] = nil
	end

	for m in pairs(self.fixedActionsList) do
		self.fixedActionsList[m] = nil
	end
	
	self.fixedBones = 0
	self.fixedActions = 0
	
	local curFrame = moho.frame
	local skel = moho:Skeleton()
	
	if (skel == nil) then
		return
	elseif (self.boneList[1] == nil) then
		return
	end
	
	for b in pairs(self.boneList) do
		local bone = skel:Bone(self.boneList[b])
		if bone == nil then
			self.status = 'Correction was canceled.'
			self:UpdateWidgets(moho)
			moho.document:PrepUndo(moho.layer)
			moho.document:SetDirty()
			self:ClearFrame(moho, self.refKey)
			if self.refKey > 1 then
				self:ClearFrame(moho, self.refKey-1)
			end
			local markerChannels = moho.layer.fTimelineMarkers
			markerChannels:DeleteKey(self.refKey)
			self:HardUpdate(moho)
			return LM.GUI.Alert(LM.GUI.ALERT_WARNING, "Some bones are missing. The skeleton structure was changed.", "", "", "EXIT")
		end
	end
		
	local totalCurBones = 0
	for i=0, skel:CountBones()-1 do 
		totalCurBones = totalCurBones + 1
		local myBone = skel:Bone(i)
		local channelAngle = myBone.fAnimAngle
		local channelPos = myBone.fAnimPos
		local channelScale = myBone.fAnimScale
		local isRefKeysOk = true
		
		local isPosChannelSplit = channelPos:AreDimensionsSplit()
		if isPosChannelSplit then
			local subChannel1 = channelPos:DimensionChannel(0)
			local subChannel2 = channelPos:DimensionChannel(1)
			if not subChannel1:HasKey(self.refKey) or not subChannel2:HasKey(self.refKey) then
				isRefKeysOk = false
			end
		elseif not channelPos:HasKey(self.refKey) then
			isRefKeysOk = false
		end
		if not channelAngle:HasKey(self.refKey) or not channelScale:HasKey(self.refKey) then
			isRefKeysOk = false
		end
		
		if not isRefKeysOk and self.boneRefList[i + 1] ~= nil then
			self.status = 'Correction was canceled.'
			self:UpdateWidgets(moho)
			moho.document:PrepUndo(moho.layer)
			moho.document:SetDirty()
			self:ClearFrame(moho, self.refKey)
			if self.refKey > 1 then
				self:ClearFrame(moho, self.refKey-1)
			end
			local markerChannels = moho.layer.fTimelineMarkers
			markerChannels:DeleteKey(self.refKey)
			self:HardUpdate(moho)
			return LM.GUI.Alert(LM.GUI.ALERT_WARNING, "Some reference keys are missing. Please do not remove or edit any reference keys.", "", "", "EXIT")
		end
	end
	
	if totalCurBones < self.totalBones then
		self.status = 'Correction was canceled.'
		self:UpdateWidgets(moho)
		moho.document:PrepUndo(moho.layer)
		moho.document:SetDirty()
		self:ClearFrame(moho, self.refKey)
		if self.refKey > 1 then
			self:ClearFrame(moho, self.refKey-1)
		end
		local markerChannels = moho.layer.fTimelineMarkers
		markerChannels:DeleteKey(self.refKey)
		self:HardUpdate(moho)
		return LM.GUI.Alert(LM.GUI.ALERT_WARNING, "The original number of bones is different.", "", "", "EXIT")
	end
	
	local isChanges = false
	
	for r in pairs(self.boneList) do
		local myBone = skel:Bone(self.boneList[r])
		local extraBoneAngle = self.boneAngleExtraList[r]
		local extraBoneAngleTmp = myBone.fAnimAngle:GetValue(self.refKey)
		if myBone.fParent > -1 then
			local origParentMatrix = LM.Matrix:new_local()
			origParentMatrix:Set(self.boneParentMatrixList[r])
			extraBoneAngle = self:GetGlobalBoneAngle(moho, origParentMatrix, extraBoneAngle, 0, false)
			extraBoneAngleTmp = self:GetGlobalBoneAngle(moho, origParentMatrix, extraBoneAngleTmp, self.refKey, false)
		end	
		local reparentOffset = extraBoneAngle - extraBoneAngleTmp
		table.insert(self.boneAngleOffsetList, reparentOffset)
	end
	
	repeat
		isChanges = false
		for y in pairs(self.boneList) do
			local myBone = skel:Bone(self.boneList[y])
			if myBone.fParent > -1 then
				local newParentId = myBone.fParent
				for t in pairs(self.boneList) do
					if newParentId == self.boneList[t] then
						if y < t then
							isChanges = true
							local key2 = self.boneList[t]
							self.boneList[y], self.boneList[t] = self:SwapTwoArrayKeys(self.boneList, y, t)
							self.boneOrigPosList[y], self.boneOrigPosList[t] = self:SwapTwoArrayKeys(self.boneOrigPosList, y, t)
							self.boneOrigGlobalPosList[y], self.boneOrigGlobalPosList[t] = self:SwapTwoArrayKeys(self.boneOrigGlobalPosList, y, t)
							self.boneAngleList[y], self.boneAngleList[t] = self:SwapTwoArrayKeys(self.boneAngleList, y, t)
							self.boneAngleExtraList[y], self.boneAngleExtraList[t] = self:SwapTwoArrayKeys(self.boneAngleExtraList, y, t)
							self.boneAngleOffsetList[y], self.boneAngleOffsetList[t] = self:SwapTwoArrayKeys(self.boneAngleOffsetList, y, t)
							self.boneParentAngleList[y], self.boneParentAngleList[t] = self:SwapTwoArrayKeys(self.boneParentAngleList, y, t)
							self.boneParentsPosList[y], self.boneParentsPosList[t] = self:SwapTwoArrayKeys(self.boneParentsPosList, y, t)
							self.boneParentList[y], self.boneParentList[t] = self:SwapTwoArrayKeys(self.boneParentList, y, t)
							self.boneParentMatrixList[y], self.boneParentMatrixList[t] = self:SwapTwoArrayKeys(self.boneParentMatrixList, y, t)
							break
						end	
					else
						local nextBone = skel:Bone(newParentId)
						repeat 
							local prevBone = nextBone
							for g in pairs(self.boneList) do
								if prevBone.fParent == self.boneList[g] then
									if y < g then
										isChanges = true
										local key2 = self.boneList[g]
										self.boneList[y], self.boneList[g] = self:SwapTwoArrayKeys(self.boneList, y, g)
										self.boneOrigPosList[y], self.boneOrigPosList[g] = self:SwapTwoArrayKeys(self.boneOrigPosList, y, g)
										self.boneOrigGlobalPosList[y], self.boneOrigGlobalPosList[g] = self:SwapTwoArrayKeys(self.boneOrigGlobalPosList, y, g)
										self.boneAngleList[y], self.boneAngleList[g] = self:SwapTwoArrayKeys(self.boneAngleList, y, g)
										self.boneAngleExtraList[y], self.boneAngleExtraList[g] = self:SwapTwoArrayKeys(self.boneAngleExtraList, y, g)
										self.boneAngleOffsetList[y], self.boneAngleOffsetList[g] = self:SwapTwoArrayKeys(self.boneAngleOffsetList, y, g)
										self.boneParentAngleList[y], self.boneParentAngleList[g] = self:SwapTwoArrayKeys(self.boneParentAngleList, y, g)
										self.boneParentsPosList[y], self.boneParentsPosList[g] = self:SwapTwoArrayKeys(self.boneParentsPosList, y, g)
										self.boneParentList[y], self.boneParentList[g] = self:SwapTwoArrayKeys(self.boneParentList, y, g)
										self.boneParentMatrixList[y], self.boneParentMatrixList[g] = self:SwapTwoArrayKeys(self.boneParentMatrixList, y, g)
									end	
								end		
							end
							if nextBone.fParent > -1 then
								nextBone = skel:Bone(nextBone.fParent)
							end
						until nextBone == prevBone
					end
				end
			end	
		end	
	until isChanges == false
		
	moho.document:PrepUndo(moho.layer)
	moho.document:SetDirty()
	
	for i in pairs(self.boneList) do
		local myBone = skel:Bone(self.boneList[i])
		for act=0, moho.layer:CountActions()-1 do
			local actName = moho.layer:ActionName(act)
			local actionChannel = myBone.fAnimAngle:ActionByName(actName)
			local found = false
			local isSmartBone = moho.layer:IsSmartBoneAction(actName)
			
			if actionChannel and actionChannel:Duration()>0 and isSmartBone then
				self:ScanAction(moho, skel, myBone, actName, i)
				found = true
			end
			if found == false and isSmartBone then
				actionChannel = myBone.fAnimPos:ActionByName(actName)
				if actionChannel and actionChannel:Duration()>0 then 
					self:ScanAction(moho, skel, myBone, actName, i)
					found = true
				end
			end
			
			if found == true then
				self.fixedActions = self.fixedActions + 1
			end	
		end
	end
	
	moho.layer:ActivateAction(nil)
	
	self:ClearFrame(moho, self.refKey)
	local markerChannels = moho.layer.fTimelineMarkers
	markerChannels:DeleteKey(self.refKey)
	if self.refKey > 1 then
		self:ClearFrame(moho, self.refKey-1)
	end
	
	local totalBones = #self.fixedBonesList
	local totalActions = #self.fixedActionsList
	local actionsStr = ' actions.'
	local bonesStr = ' bones were'
	
	if totalBones == 1 then
		bonesStr = ' bone was'
	end
	if totalActions == 1 then
		actionsStr = ' action.'
	end
	
	self.status = 'Correction completed! ' .. totalBones.. bonesStr..' fixed in '.. totalActions.. actionsStr
	
	if totalBones == 0 or totalActions == 0 then
		self.status = 'Could not find any actions that need to be fixed.'
	end
	
	self:UpdateWidgets(moho)
	moho:SetCurFrame(curFrame)
	self:HardUpdate(moho)
end

function MR_SmartBoneFixer:HardUpdate(moho)
	moho.document:PrepUndo(moho.layer)
	moho.document:SetDirty()
	moho:UpdateSelectedChannels()
	moho.document:Undo()
	moho.view:DrawMe()
	moho:UpdateUI()
	moho.layer:UpdateCurFrame()
end

function MR_SmartBoneFixer:ClearFrame(moho, frame)
	local skel = moho:Skeleton()
	for i=0, skel:CountBones()-1 do 
		local myBone = skel:Bone(i)
		myBone.fAnimPos:DeleteKey(frame)
		myBone.fAnimAngle:DeleteKey(frame)
		myBone.fAnimScale:DeleteKey(frame)
	end
end

function MR_SmartBoneFixer:SetKey(moho, frame, valueFromFrame)
	local keyInterp = MOHO.InterpSetting:new_local()
	keyInterp.tags = 1
	local skel = moho:Skeleton()
	for i=0, skel:CountBones()-1 do 
		local myBone = skel:Bone(i)
		local channelPos = myBone.fAnimPos
		local pos = channelPos:GetValue(valueFromFrame)
		local angle = myBone.fAnimAngle:GetValue(valueFromFrame)
		local scale = myBone.fAnimScale:GetValue(valueFromFrame)
		local isPosChannelSplit = channelPos:AreDimensionsSplit()
		if isPosChannelSplit then
			local subChannel1 = channelPos:DimensionChannel(0)
			local subChannel2 = channelPos:DimensionChannel(1)
			if subChannel1 then
				subChannel1:SetValue(frame, pos.x)
				subChannel1:SetKeyInterp(frame, keyInterp)
			end	
			if subChannel2 then
				subChannel2:SetValue(frame, pos.y)
				subChannel2:SetKeyInterp(frame, keyInterp)
			end
		else	
			myBone.fAnimPos:SetValue(frame, pos)
			myBone.fAnimPos:SetKeyInterp(frame, keyInterp)
		end
		myBone.fAnimAngle:SetValue(frame, angle)
		myBone.fAnimAngle:SetKeyInterp(frame, keyInterp)
		if valueFromFrame == 0 then
			myBone.fAnimScale:SetValue(frame, 1)
			myBone.fAnimScale:SetKeyInterp(frame, keyInterp)
		else
			myBone.fAnimScale:SetValue(frame, scale)
			myBone.fAnimScale:SetKeyInterp(frame, keyInterp)
		end
	end
end

function MR_SmartBoneFixer:ScanAction(moho, skel, bone, actionName, boneNum)
	moho.layer:ActivateAction(actionName)

	local actionChannelPos =  bone.fAnimPos:ActionByName(actionName)
	local actionChannelAngle =  bone.fAnimAngle:ActionByName(actionName)

	for i=1, actionChannelPos:CountKeys()-1 do
		local keyTime = actionChannelPos:GetKeyWhen(i)
		if keyTime > 0 then
			self:ApplyPosition(moho, skel, bone, keyTime, boneNum, actionName)
		end
	end
	for i=1, actionChannelAngle:CountKeys()-1 do
		local keyTime = actionChannelAngle:GetKeyWhen(i)
		if keyTime > 0 then
			self:ApplyAngle(moho, skel, bone, keyTime, boneNum, actionName)
		end	
	end
end

function MR_SmartBoneFixer:ApplyPosition(moho, skel, bone, frame, boneNum, actionName)
	moho:SetCurFrame(frame)
	local myBone = bone
	local parentBone = myBone.fParent
	local oldParentBone = self.boneParentList[boneNum]
	local origBonePos = LM.Vector2:new_local()
	origBonePos:Set(self.boneOrigPosList[boneNum])
	local localOrigBonePos = LM.Vector2:new_local()
	localOrigBonePos:Set(self.boneOrigPosList[boneNum])
	local origBoneGlobalPos = LM.Vector2:new_local()
	origBoneGlobalPos:Set(self.boneOrigGlobalPosList[boneNum])
	local bonePosFrameZero = LM.Vector2:new_local()	
	local oldMatrix = LM.Matrix:new_local()
	oldMatrix:Set(self.boneParentMatrixList[boneNum])
	local parent = skel:Bone(parentBone)
	local parentOffset = LM.Vector2:new_local()
	local parentOffsetZero = LM.Vector2:new_local()
	parentOffset:Set(0,0)
	parentOffsetZero:Set(0,0)
	
	local oldParentBonePos = self.boneParentsPosList[boneNum]

	if parent ~= nil then
		parent.fMovedMatrix:Transform(parentOffset)  
		parent.fRestMatrix:Transform(parentOffsetZero)
	end

	parentOffset = parentOffsetZero - parentOffset

	local curBonePos = LM.Vector2:new_local()
	curBonePos:Set(myBone.fAnimPos:GetValue(frame))
	
	oldMatrix:Transform(origBonePos)
	bonePosFrameZero:Set(myBone.fAnimPos:GetValue(0))
	local oldParent = skel:Bone(oldParentBone)
	
	if parent ~= nil then
		parent.fRestMatrix:Transform(bonePosFrameZero)  
	end
	
	local bonePos = LM.Vector2:new_local()	
	bonePos:Set(myBone.fAnimPos:GetValue(frame))
	
	local bonePosDifFrameZero = LM.Vector2:new_local()
	bonePosDifFrameZero = origBoneGlobalPos - bonePosFrameZero
	local localBonePosFrameZero = LM.Vector2:new_local()	
	localBonePosFrameZero:Set(myBone.fAnimPos:GetValue(0))
	oldMatrix:Transform(bonePos)   
	
	bonePos = bonePos - bonePosDifFrameZero  - parentOffset

	if parent ~= nil then
		local inverseM = LM.Matrix:new_local()
		inverseM:Set(parent.fMovedMatrix) 
		inverseM:Invert()
		inverseM:Transform(bonePos)
	end
	
	if localOrigBonePos.x ~= localBonePosFrameZero.x
	or localOrigBonePos.y ~= localBonePosFrameZero.y
	or oldParentBone ~= parentBone then
		myBone.fAnimPos:SetValue(frame, bonePos)
		self:CollectLog(myBone:Name(), actionName)
	end
end

function MR_SmartBoneFixer:CollectLog(boneName, actionName)
	local isBoneNew = true
	local isActionNew = true
	
	for z in pairs(self.fixedBonesList) do
		if self.fixedBonesList[z] == boneName then
			isBoneNew = false
			break
		end	
	end
	if isBoneNew then
		table.insert(self.fixedBonesList, boneName)
	end

	if actionName ~= '' then
		for a in pairs(self.fixedActionsList) do
			if self.fixedActionsList[a] == actionName then
				isActionNew = false
				break
			end	
		end
		if isActionNew then
			table.insert(self.fixedActionsList, actionName)
		end
	end
end

function MR_SmartBoneFixer:GetGlobalBoneAngle(moho, matrix, boneAngle, frame, round)
	local v1 = LM.Vector2:new_local()
	local v2 = LM.Vector2:new_local()
	v1:Set(0,0)
	v2:Set(1,0)
	local newAngle = 0
	local invMatrix = LM.Matrix:new_local()	
										
	invMatrix:Set(matrix)
	invMatrix:Invert()
	invMatrix:Transform(v1)
	invMatrix:Transform(v2)
	v2 = v2-v1
	newAngle = math.atan2(v2.y, v2.x)
	
	if round == true then
		while newAngle > 2 * math.pi do
			newAngle = newAngle - 2 * math.pi
		end
		while newAngle < - 2 * math.pi do
			newAngle = newAngle + 2 * math.pi
		end	
	end
	newAngle = boneAngle - newAngle
	return newAngle
end

function MR_SmartBoneFixer:ApplyAngle(moho, skel, bone, frame, boneNum, actionName)
	moho:SetCurFrame(frame)
	local myBone = bone
	local oldBoneAngle = self.boneAngleList[boneNum]
	local newBoneAngle = myBone.fAnimAngle:GetValue(frame)
	local newBoneAngleFrameZero = myBone.fAnimAngle:GetValue(0)
	local newParentBoneId = myBone.fParent
	local oldParentBoneId = self.boneParentList[boneNum]
	local boneAngleDifFrameZero = 0
	local parentAngleOffset = 0
	local oldParentAngleOffset = 0
	
	if oldBoneAngle ~= newBoneAngleFrameZero and oldParentBoneId == newParentBoneId then
		local oldParentBone = skel:Bone(oldParentBoneId)
		boneAngleDifFrameZero = oldBoneAngle - newBoneAngleFrameZero
		newBoneAngle = (newBoneAngle - boneAngleDifFrameZero) + self.boneAngleOffsetList[boneNum]
		self:CollectLog(myBone:Name(), actionName) 
	elseif oldBoneAngle ~= newBoneAngleFrameZero and oldParentBoneId ~= newParentBoneId and oldParentBoneId >= 0 then
		local oldParentBone = skel:Bone(oldParentBoneId)
		local oldMatrix = LM.Matrix:new_local()
		oldMatrix = oldParentBone.fMovedMatrix
		
		if newParentBoneId > -1 then
			local parentBone = skel:Bone(newParentBoneId)
			if self.considerNewParentRotation then
				local parentAngleZero = parentBone.fAnimAngle:GetValue(0)
				local parentAngle = parentBone.fAnimAngle:GetValue(frame)
				if	parentBone.fParent > -1 then
					local gandPaBone = skel:Bone(parentBone.fParent)
					parentAngleZero = self:GetGlobalBoneAngle(moho, gandPaBone.fRestMatrix, parentAngleZero, 0, false)
					parentAngle = self:GetGlobalBoneAngle(moho, gandPaBone.fMovedMatrix, parentAngle, frame, false)
				end
				parentAngleOffset = parentAngleZero - parentAngle
			end
		end	
		oldBoneAngle = oldBoneAngle - self.boneAngleOffsetList[boneNum]
		
		if oldParentBoneId > -1 then
			local oldParentBone = skel:Bone(oldParentBoneId)
			if self.considerOldParentRotation then
				local oldParentAngleZero = oldParentBone.fAnimAngle:GetValue(0)
				local oldParentAngle = oldParentBone.fAnimAngle:GetValue(frame)
				if	oldParentBone.fParent > -1 then
					local oldGandPaBone = skel:Bone(oldParentBone.fParent)
					oldParentAngleZero = self:GetGlobalBoneAngle(moho, oldGandPaBone.fRestMatrix, oldParentAngleZero, 0, false)
					oldParentAngle = self:GetGlobalBoneAngle(moho, oldGandPaBone.fMovedMatrix, oldParentAngle, frame, false)
				end
				oldParentAngleOffset = oldParentAngleZero - oldParentAngle
			end
		end
		boneAngleDifFrameZero = oldBoneAngle - newBoneAngleFrameZero
		newBoneAngle = (newBoneAngle - boneAngleDifFrameZero) + parentAngleOffset - oldParentAngleOffset
		self:CollectLog(myBone:Name(), actionName)
	elseif newParentBoneId > -1 and oldParentBoneId < 0 then
		local parentBone = skel:Bone(newParentBoneId)
		local newMatrix = LM.Matrix:new_local()
		newMatrix = parentBone.fMovedMatrix

		if self.considerNewParentRotation then
			local parentAngleZero = parentBone.fAnimAngle:GetValue(0)
			local parentAngle = parentBone.fAnimAngle:GetValue(frame)
			if	parentBone.fParent > -1 then
			local gandPaBone = skel:Bone(parentBone.fParent)
				parentAngleZero = self:GetGlobalBoneAngle(moho, gandPaBone.fRestMatrix, parentAngleZero, 0, false)
				parentAngle = self:GetGlobalBoneAngle(moho, gandPaBone.fMovedMatrix, parentAngle, frame, false)
			end	
			parentAngleOffset = parentAngleZero - parentAngle
		end	
		newBoneAngleFrameZero = newBoneAngleFrameZero + self.boneAngleOffsetList[boneNum]
		boneAngleDifFrameZero = oldBoneAngle - newBoneAngleFrameZero
		newBoneAngle = (newBoneAngle - boneAngleDifFrameZero) + parentAngleOffset 
		self:CollectLog(myBone:Name(), actionName)
	end
	myBone.fAnimAngle:SetValue(frame,newBoneAngle)
end

-- **************************************************
-- Localization
-- **************************************************

function MR_SmartBoneFixer:Localize(text)
	local phrase = {}

	phrase['Description'] = 'Fix smart bone actions after editing bones (For rigging only)'
	phrase['UILabel'] = 'Smart Bone Fixer'

	phrase['Get original transformation'] = 'Get original transformation'
	phrase['Fix Actions'] = 'Fix actions'
	phrase['Clear'] = 'Clear'
	phrase['Consider new parent rotation'] = 'Consider new parent rotation'
    phrase['Consider old parent rotation'] = 'Consider old parent rotation'
    phrase['Close'] = 'Close'
    phrase['Settings'] = 'Settings'
	phrase['Status:'] = 'Status: '.. self.status

	local fileWord = MOHO.Localize("/Menus/File/File=File")
	if fileWord == "Файл" then
		phrase['Description'] = 'Fix smart bone actions after editing bones (For rigging only)'
		phrase['UILabel'] = 'Smart Bone Fixer'

		phrase['Get original transformation'] = 'Get original transformation'
		phrase['Fix Actions'] = 'Fix actions'
		phrase['Clear'] = 'Clear'
		phrase['Consider new parent rotation'] = 'Consider new parent rotation'
		phrase['Consider old parent rotation'] = 'Consider old parent rotation'
		phrase['Close'] = 'Close'
		phrase['Settings'] = 'Settings'
		phrase['Status:'] = 'Status: ' .. self.status
	end

	return phrase[text]
end

Icon
Smartbone Fixer
Listed

Script type: Tool

Uploaded: Dec 30 2020, 14:11

Last modified: Dec 30 2020, 14:13

This script allows you to correct actions after insignificant bone transformations.
1. Select bones that you want to transform, activate the Smart Bone Fixer tool and click the “Get original transformation” button.
2. Transform the bones. You can change the position, rotation, scale, and also change the bone’s parent.
3. Activate the Smart Bone Fixer again and click the “Fix actions” button.

Currently the script only compatible with Moho 12 due to some undocumented breaking API changes in version 13.


This script, and all other scripts on this site are distributed as free software under the GNU General Public License 3.0 or later.
Downloads count: 308