Cell (WoW addon) LibHealComm heal prediction arbitration snippet - CHAP/CHAD

LibHealCommHealPrediction.lua

-- Cell Heal Arbitration Prediction (CHAP) with Cell Heal Arbitration
-- Diagnostics (CHAD)
--
-- Improves Cell's incoming-heal prediction on raid/party frames by
-- combining Blizzard's native prediction with LibHealComm-4.0's, since
-- native alone misses HoTs/AoE healing and any healer whose own client
-- isn't broadcasting via LibHealComm.
--
-- Requires some addon on this client to provide LibHealComm-4.0, either
-- as its own standalone library addon or bundled inside another addon
-- (e.g. YAHP). If none is found, this prints a message and does nothing
-- else - click Run again in the Code Snippets editor once one is
-- installed if it doesn't pick it up automatically.
--
-- Commands (/chap also works as /cellhealpred, /chad also works as
-- /cellhealdiag; /chap with no args lists all of these in-game):
-- /chap - show status and this command list.
-- /chap debug - toggle a raw trace of widget discovery/hooking.
-- /chap diag - toggle CHAD (the /chad diagnostics below) on/off.
-- /chad - show how many heals were won by native prediction vs
--   LibHealComm.
-- /chad reset - clear the diagnostic counters.
-- /chad debug - toggle a raw trace of each diagnostic decision.
-- /chad test - show collected native/LibHealComm/actual-landed data per
--   target+spell (only while CHAD is on) - useful for checking whether
--   native tracks spell rank/gear the same way LibHealComm does.
-- /chad test reset - clear collected test data.

if CELL_CHAD_ENABLED == nil then
    CELL_CHAD_ENABLED = false
end

-- How native and LibHealComm get combined into the displayed value,
-- whenever both have something to say:
--   "max" (default) - always the larger of the two.
--   "min" - always the smaller - more conservative, if you'd rather
--           under-predict than over-predict incoming healing.
--   "avg" - split the difference.
-- Falls back to whichever single source is nonzero if the other is 0
-- (e.g. HoTs, which native never predicts).
local CHAP_AGGREGATE = "max"

local function ChapInit()
    local HealComm = LibStub("LibHealComm-4.0", true)

    if not HealComm then
        print("|cff0099ffCHAP|r: LibHealComm-4.0 isn't available on this client - install it as a standalone library addon, or install any addon that bundles it (e.g. YAHP), for heal prediction arbitration to work. If it IS installed and this just didn't pick it up automatically, click Run on this snippet in the Code Snippets editor to activate it now.")
        return
    end

    if not Cell or not Cell.funcs or not Cell.funcs.HandleUnitButton then
        print("|cff0099ffCHAP|r: Cell.funcs.HandleUnitButton isn't available - Cell may not be fully loaded yet, or its internal structure changed. Click Run again once Cell is fully loaded, or report this if it keeps happening.")
        return
    end

    local F = Cell.funcs

    local chapDebug = false
    local chapErrorCount = 0
    local chapHookedCount = 0

    local chadDebug = false
    local healPredictionDiagLastValues = {}
    local healPredictionDiagLastWinner = {}
    local HEAL_PREDICTION_DIAG_EPSILON = 1

    local chapSpellData = {}
    local chapCurrentSpellByGuid = {}

    local function GetSpellBucket(destGUID, spellID, spellName)
        local targetData = chapSpellData[destGUID]
        if not targetData then
            targetData = {spells = {}}
            chapSpellData[destGUID] = targetData
        end

        local bucket = targetData.spells[spellID]
        if not bucket then
            bucket = {spellName = spellName, blizzardMax = 0, libMax = 0, actualHeals = {}, actualCritHeals = {}}
            targetData.spells[spellID] = bucket
        elseif spellName and not bucket.spellName then
            bucket.spellName = spellName
        end

        return bucket
    end

    local function HealPredictionDiag_Classify(blizzardValue, libValue)
        if libValue > blizzardValue then return "lib"
        elseif blizzardValue > libValue then return "blizzard"
        else return "tie" end
    end

    local function HealPredictionDiag_Record(guid, blizzardValue, libValue)
        if not guid then return end

        local last = healPredictionDiagLastValues[guid]
        healPredictionDiagLastValues[guid] = {blizzard = blizzardValue, lib = libValue}

        if blizzardValue == 0 and libValue == 0 then
            healPredictionDiagLastWinner[guid] = nil
            return
        end

        local grew = (not last)
            or (blizzardValue - last.blizzard > HEAL_PREDICTION_DIAG_EPSILON)
            or (libValue - last.lib > HEAL_PREDICTION_DIAG_EPSILON)
        local winner = HealPredictionDiag_Classify(blizzardValue, libValue)
        local winnerChanged = healPredictionDiagLastWinner[guid] ~= winner

        if chadDebug and grew then
            print(string.format("|cff00FF00CHAD|r %s: blizzard=%.0f lib=%.0f last=%s winner=%s(%s)",
                guid:sub(-6), blizzardValue, libValue,
                last and string.format("(%.0f/%.0f)", last.blizzard, last.lib) or "none",
                winner, winnerChanged and "new" or "same"))
        end

        if not grew or not winnerChanged then return end
        healPredictionDiagLastWinner[guid] = winner

        CellDB.healPredictionDiag = CellDB.healPredictionDiag or {total = 0, blizzardWins = 0, libHealCommWins = 0, tie = 0}
        local diag = CellDB.healPredictionDiag

        diag.total = diag.total + 1
        if winner == "lib" then
            diag.libHealCommWins = diag.libHealCommWins + 1
        elseif winner == "blizzard" then
            diag.blizzardWins = diag.blizzardWins + 1
        else
            diag.tie = diag.tie + 1
        end
    end

    local function AggregateValue(blizzardValue, libValue)
        if blizzardValue == 0 or libValue == 0 then
            return max(blizzardValue, libValue)
        end

        if CHAP_AGGREGATE == "min" then
            return math.min(blizzardValue, libValue)
        elseif CHAP_AGGREGATE == "avg" then
            return (blizzardValue + libValue) / 2
        else
            return max(blizzardValue, libValue)
        end
    end

    local function ComputeArbitratedValue(button)
        local unit = button.states and button.states.displayedUnit
        if not unit then return nil end

        local guid = UnitGUID(unit)
        if not guid then return nil end

        local blizzardValue = UnitGetIncomingHeals(unit) or 0

        local modifier = HealComm:GetHealModifier(guid) or 1
        local libValue = (HealComm:GetHealAmount(guid, HealComm.CASTED_HEALS) or 0) * modifier
        local hot = (HealComm:GetHealAmount(guid, HealComm.OVERTIME_AND_BOMB_HEALS, GetTime() + 3) or 0) * modifier
        libValue = libValue + hot

        local arbitratedValue = AggregateValue(blizzardValue, libValue)

        if CELL_CHAD_ENABLED then
            HealPredictionDiag_Record(guid, blizzardValue, libValue)

            local spellID = chapCurrentSpellByGuid[guid]
            if spellID then
                local spellName = GetSpellInfo and GetSpellInfo(spellID)
                local bucket = GetSpellBucket(guid, spellID, spellName)
                bucket.blizzardMax = max(bucket.blizzardMax, blizzardValue)
                bucket.libMax = max(bucket.libMax, libValue)
            end

            if chapSpellData[guid] and not chapSpellData[guid].name then
                chapSpellData[guid].name = UnitName(unit)
            end
        end

        return arbitratedValue
    end

    local hookedWidgets = setmetatable({}, {__mode = "k"})
    local reentryGuard = setmetatable({}, {__mode = "k"})

    local function ApplyArbitratedValue(widget, button)
        if reentryGuard[widget] then return end

        local ok, arbitratedValue = pcall(ComputeArbitratedValue, button)
        if not ok then
            chapErrorCount = chapErrorCount + 1
            print(string.format("|cff0099ffCHAP ERROR|r (leaving native value as-is, %d total): %s", chapErrorCount, tostring(arbitratedValue)))
            return
        end
        if not arbitratedValue or arbitratedValue == 0 then return end

        local healthMax = button.states and button.states.healthMax
        if not healthMax or healthMax == 0 then return end

        local newFraction = arbitratedValue / healthMax

        if chapDebug then
            print(string.format("|cff0099ffCHAP APPLY|r unit=%s arbitrated=%.0f healthMax=%.0f fraction=%.3f",
                tostring(button.states and button.states.displayedUnit), arbitratedValue, healthMax, newFraction))
        end

        reentryGuard[widget] = true
        local setOk, setErr = pcall(function()
            widget:Show()

            local getOk, currentFraction = pcall(widget.GetValue, widget)
            if not (getOk and currentFraction and math.abs(newFraction - currentFraction) < 0.0001) then
                widget:SetValue(newFraction)
            end
        end)
        reentryGuard[widget] = false

        if not setOk then
            chapErrorCount = chapErrorCount + 1
            print(string.format("|cff0099ffCHAP ERROR|r (Show/SetValue failed, %d total): %s", chapErrorCount, tostring(setErr)))
        end
    end

    local function HookButtonWidget(button)
        local widget = button and button.widgets and button.widgets.incomingHeal
        if not widget or hookedWidgets[widget] then return end
        hookedWidgets[widget] = true
        chapHookedCount = chapHookedCount + 1

        if chapDebug then
            print(string.format("|cff0099ffCHAP|r hooked incomingHeal widget #%d for unit %s", chapHookedCount, tostring(button.states and button.states.displayedUnit)))
        end

        hooksecurefunc(widget, "SetValue", function(w)
            ApplyArbitratedValue(w, button)
        end)

        hooksecurefunc(widget, "Hide", function(w)
            ApplyArbitratedValue(w, button)
        end)
    end

    local function HookAndRefresh(button)
        HookButtonWidget(button)

        local widget = button and button.widgets and button.widgets.incomingHeal
        if widget then
            ApplyArbitratedValue(widget, button)
        end
    end

    local chapHealComm = {}

    local function HealComm_Discover(_, event, casterGUID, spellID, healType, endTime, ...)
        for i = 1, select("#", ...) do
            local destGUID = select(i, ...)
            if CELL_CHAD_ENABLED and spellID then
                chapCurrentSpellByGuid[destGUID] = spellID
            end
            F.HandleUnitButton("guid", destGUID, HookAndRefresh)
        end
    end
    chapHealComm.HealComm_Discover = HealComm_Discover

    local function HealComm_DiscoverSingleGUID(_, event, guid)
        F.HandleUnitButton("guid", guid, HookAndRefresh)
    end
    chapHealComm.HealComm_DiscoverSingleGUID = HealComm_DiscoverSingleGUID

    HealComm.RegisterCallback(chapHealComm, "HealComm_HealStarted", "HealComm_Discover")
    HealComm.RegisterCallback(chapHealComm, "HealComm_HealUpdated", "HealComm_Discover")
    HealComm.RegisterCallback(chapHealComm, "HealComm_HealStopped", "HealComm_Discover")
    HealComm.RegisterCallback(chapHealComm, "HealComm_HealDelayed", "HealComm_Discover")
    HealComm.RegisterCallback(chapHealComm, "HealComm_ModifierChanged", "HealComm_DiscoverSingleGUID")
    HealComm.RegisterCallback(chapHealComm, "HealComm_GUIDDisappeared", "HealComm_DiscoverSingleGUID")

    local chapNativeDiscoveryFrame = CreateFrame("Frame")
    chapNativeDiscoveryFrame:RegisterEvent("UNIT_HEAL_PREDICTION")
    chapNativeDiscoveryFrame:SetScript("OnEvent", function(_, event, unit)
        local guid = unit and UnitGUID(unit)
        if guid then
            F.HandleUnitButton("guid", guid, HookAndRefresh)
        end
    end)

    local chapCombatLogFrame = CreateFrame("Frame")
    chapCombatLogFrame:RegisterEvent("COMBAT_LOG_EVENT_UNFILTERED")
    chapCombatLogFrame:SetScript("OnEvent", function()
        if not CELL_CHAD_ENABLED then return end

        local _, subevent, _, sourceGUID, sourceName, _, _, destGUID, destName, _, _,
            spellID, spellName, spellSchool, amount, overhealing, absorbed, critical = CombatLogGetCurrentEventInfo()

        if subevent ~= "SPELL_HEAL" and subevent ~= "SPELL_PERIODIC_HEAL" then return end
        if not spellID or not destGUID then return end

        if chapDebug then
            print(string.format("|cff0099ffCHAP CLEU|r %s spellID=%s school=%s amount=%s overhealing=%s absorbed=%s critical=%s",
                subevent, tostring(spellID), tostring(spellSchool), tostring(amount), tostring(overhealing), tostring(absorbed), tostring(critical)))
        end

        local bucket = GetSpellBucket(destGUID, spellID, spellName)
        if destName and not (chapSpellData[destGUID] and chapSpellData[destGUID].name) then
            chapSpellData[destGUID].name = destName
        end

        table.insert(critical and bucket.actualCritHeals or bucket.actualHeals, amount or 0)
    end)

    local CHAP_TAG = "|cff0099ff[CHAP]|r "
    local CHAD_TAG = "|cff00FF00[CHAD]|r "
    local CMD_CHAP = "|cff33ccff"
    local CMD_CHAD = "|cff99ff99"

    local function StateTag(isOn)
        return isOn and "[|cff00ff00ON|r]" or "[|cff888888OFF|r]"
    end

    local function SpellLink(spellID, name)
        return "|cff71d5ff|Hspell:" .. spellID .. ":0|h[" .. (name or "?") .. "]|h|r(" .. spellID .. ")"
    end

    SLASH_CELLHEALARBITRATION1 = "/chap"
    SLASH_CELLHEALARBITRATION2 = "/cellhealpred"
    SlashCmdList["CELLHEALARBITRATION"] = function(msg)
        if msg == "debug" then
            chapDebug = not chapDebug
            print(CHAP_TAG .. "Widget-hook/discovery trace: " .. (chapDebug and "ON" or "OFF"))
            return
        end

        if msg == "diag" then
            CELL_CHAD_ENABLED = not CELL_CHAD_ENABLED
            print(CHAP_TAG .. "CHAD diagnostics: " .. (CELL_CHAD_ENABLED and ("ON, use " .. CMD_CHAD .. "/chad|r for stats.") or "OFF"))
            return
        end

        print(CHAP_TAG .. "Active (LibHealComm-4.0 found). "
            .. chapHookedCount .. " widgets hooked. Errors caught: " .. chapErrorCount .. ".\n" ..
            CMD_CHAP .. "/chap|r (or " .. CMD_CHAP .. "/cellhealpred|r): show this command list and whether LibHealComm is active.\n" ..
            CMD_CHAP .. "/chap debug|r: toggle widget-hook/discovery trace. " .. StateTag(chapDebug) .. "\n" ..
            CMD_CHAP .. "/chap diag|r: toggle CHAD (" .. CMD_CHAD .. "/chad|r" .. ") diagnostics on/off. " .. StateTag(CELL_CHAD_ENABLED) .. "\n" ..
            CMD_CHAD .. "/chad|r (or " .. CMD_CHAD .. "/cellhealdiag|r): show heal-arbitration diagnostic stats.\n" ..
            CMD_CHAD .. "/chad reset|r: clear diagnostic counters.\n" ..
            CMD_CHAD .. "/chad debug|r: toggle per-arbitration trace. " .. StateTag(chadDebug) .. "\n" ..
            CMD_CHAD .. "/chad test|r: show collected blizzard/libhealcomm/actual data per target+spell.\n" ..
            CMD_CHAD .. "/chad test reset|r: clear collected test data.")
    end

    SLASH_CELLHEALDIAG1 = "/cellhealdiag"
    SLASH_CELLHEALDIAG2 = "/chad"
    SlashCmdList["CELLHEALDIAG"] = function(msg)
        if msg == "reset" then
            CellDB.healPredictionDiag = nil
            print(CHAD_TAG .. "Diagnostic reset.")
            return
        end

        if msg == "debug" then
            chadDebug = not chadDebug
            local note = (chadDebug and not CELL_CHAD_ENABLED) and (" (won't print anything until CHAD is also on - use " .. CMD_CHAP .. "/chap diag|r)") or ""
            print(CHAD_TAG .. "Per-arbitration trace: " .. (chadDebug and "ON" or "OFF") .. note)
            return
        end

        if msg == "test reset" then
            wipe(chapSpellData)
            wipe(chapCurrentSpellByGuid)
            print(CHAD_TAG .. "Test data cleared.")
            return
        end

        if msg == "test" then
            if not CELL_CHAD_ENABLED then
                print(CHAD_TAG .. "Disabled - use " .. CMD_CHAP .. "/chap diag|r to enable it (data is only collected while CHAD is on).")
                return
            end

            local function StatsText(label, values)
                local n = #values
                if n == 0 then return nil end

                local vMin, vMax, vSum = math.huge, 0, 0
                for _, v in ipairs(values) do
                    vMin, vMax, vSum = math.min(vMin, v), math.max(vMax, v), vSum + v
                end
                return string.format("%s: n=%d min=%.0f max=%.0f avg=%.0f", label, n, vMin, vMax, vSum / n)
            end

            local any = false
            for destGUID, targetData in pairs(chapSpellData) do
                for spellID, bucket in pairs(targetData.spells) do
                    any = true

                    local actualText = StatsText("actual", bucket.actualHeals) or "actual: no non-crit landed data yet"
                    local critText = StatsText("crits", bucket.actualCritHeals)
                    local arbitratedFromPeaks = AggregateValue(bucket.blizzardMax, bucket.libMax)

                    print(CHAD_TAG .. string.format("%s - %s predictions: blizzard=%.0f, libhealcomm=%.0f, arbitrated (%s)=%.0f | %s%s",
                        targetData.name or destGUID:sub(-6), SpellLink(spellID, bucket.spellName),
                        bucket.blizzardMax, bucket.libMax, CHAP_AGGREGATE, arbitratedFromPeaks, actualText,
                        critText and (" | " .. critText) or ""))
                end
            end

            if not any then
                print(CHAD_TAG .. "No data yet - cast heals and let them land.")
                return
            end

            print(CHAD_TAG .. "Use " .. CMD_CHAD .. "/chad test reset|r before switching test conditions.")
            return
        end

        if not CELL_CHAD_ENABLED then
            print(CHAD_TAG .. "Disabled - use " .. CMD_CHAP .. "/chap diag|r to enable it.")
            return
        end

        local diag = CellDB.healPredictionDiag
        if not diag or diag.total == 0 then
            print(CHAD_TAG .. "No data recorded yet.")
            return
        end

        print(CHAD_TAG .. string.format("%d incoming heals recorded:", diag.total))
        print(string.format("  Blizzard native: %d (%.1f%%)", diag.blizzardWins, diag.blizzardWins / diag.total * 100))
        print(string.format("  LibHealComm: %d (%.1f%%)", diag.libHealCommWins, diag.libHealCommWins / diag.total * 100))
        print(string.format("  Tied: %d (%.1f%%)", diag.tie, diag.tie / diag.total * 100))
        print(CHAD_TAG .. "Use " .. CMD_CHAD .. "/chad reset|r to clear.")
    end

    print(CHAP_TAG .. "LibHealComm-4.0 found - heal prediction arbitration active. Type " .. CMD_CHAP .. "/chap|r (or " .. CMD_CHAP .. "/cellhealpred|r) for commands.")
end

if IsLoggedIn() then
    ChapInit()
else
    local chapLoginFrame = CreateFrame("Frame")
    chapLoginFrame:RegisterEvent("PLAYER_LOGIN")
    chapLoginFrame:SetScript("OnEvent", function(self)
        self:UnregisterEvent("PLAYER_LOGIN")
        ChapInit()
    end)
end