wow-addons/z-BigBrother/bb/graph.py

import base64
import io
from collections import defaultdict, Counter
from contextlib import suppress
from datetime import datetime, timedelta
from typing import List

import matplotlib.pyplot as plt

from . import config


def graph(raid: List[dict]) -> None:
    # TODO: This entire function is shit.
    figs = []
    buff_ids = {
        buff_id: buff_name
        for buff_name, buff_data in config.buffs.items()
        for buff_id in buff_data["ids"]
    }
    classes = defaultdict(lambda: defaultdict(lambda: defaultdict(list)))
    player_buff_counts = defaultdict(Counter)
    xticks = []
    first_encounter_date = None
    last_encounter_date = None
    previous_encounter_date = None
    for encounter in raid:
        encounter_date = datetime.strptime(encounter["date"], "%Y-%m-%d %H:%M:%S")
        if first_encounter_date is None:
            first_encounter_date = encounter_date
            previous_encounter_date = encounter_date - timedelta(minutes=5)
        last_encounter_date = encounter_date
        xticks.append((encounter_date, encounter["encounterName"]))
        for class_name, players in encounter["buffs"].items():
            for player_name, buffs in sorted(players.items()):
                real_class_name = config.class_overrides.get(player_name, str.title(class_name))
                classes[real_class_name][player_name]  # ensure player is added, irregardless if they have any buffs
                player_buff_names = set()
                for buff_id in buffs:
                    with suppress(KeyError):  # KeyError on buff_id => buff not tracked
                        buff_name = buff_ids[buff_id]
                        player_buff_names.add(buff_name)
                        classes[real_class_name][player_name][buff_name].append(
                            (previous_encounter_date, (encounter_date - previous_encounter_date))
                        )
                player_buff_counts[player_name].update(player_buff_names)
        previous_encounter_date = encounter_date

    for class_name, players in sorted(classes.items()):
        fig, ax = plt.subplots()
        class_buffs = {
            buff_name: buff_data
            for buff_name, buff_data in config.buffs.items()
            if any(x in buff_data["classes"] for x in ("ALL", class_name))
        }
        required_class_buffs = {
            buff_name: buff_data
            for buff_name, buff_data in class_buffs.items()
            if buff_data["required"]
        }
        #legend_handles = []
        bar_width = 1 / (len(class_buffs) + 1)
        for b, (buff_name, buff) in enumerate(class_buffs.items()):
            bar_color = "#{color}".format(**buff)
            #legend_handles.append(Patch(label=buff_name, color=bar_color))
            for y0, (player_name, player_buffs) in enumerate(players.items()):
                ax.broken_barh(
                    player_buffs[buff_name],
                    (bar_width/2 + y0 + b * bar_width, bar_width),
                    color=bar_color
                )

                annotation_color = "black"
                try:
                    if player_buffs[buff_name][1][0] == first_encounter_date:
                        annotation_color = contrast_color(*tuple(int(buff["color"][i:i+2], 16) for i in (0, 2, 4)))
                except IndexError:
                    pass
                ax.annotate(
                    f" {buff_name}",
                    xy=(first_encounter_date, bar_width + y0 + b * bar_width),
                    ha="left",
                    va="center",
                    color=annotation_color
                )

        for y0 in range(len(players) + 1):
            ax.axhline(y=y0, linewidth=1.0, color="black")

        plt.xticks(*zip(*xticks), rotation=30, ha="right")
        required_buff_fraction = {
            p: sum(
                c
                for b, c in player_buff_counts[p].items()
                if b in required_class_buffs
            ) / (len(raid) * len(required_class_buffs))
            for p in players
        }
        plt.yticks(
            list(range(len(players))),
            [f"{p}\n{required_buff_fraction[p]:.0%}" for p in players],
            verticalalignment="top"
        )
        ax.invert_yaxis()
        ax.grid(axis="x", linestyle="--")
        ax.set_xlim(first_encounter_date - timedelta(minutes=10), last_encounter_date + timedelta(minutes=1))

        plt.title(class_name, fontweight="bold", fontsize=16)
        #plt.legend(handles=legend_handles, bbox_to_anchor=(1, 1))
        fig.set_size_inches(15, 2 + len(players) * 2)
        plt.tight_layout()
        #plt.show()

        buffer = io.BytesIO()
        fig.savefig(buffer, format="svg")
        figs.append(base64.b64encode(buffer.getbuffer()).decode("ascii"))

    with open("{date:10.10}-{zone}.html".format(**raid[0]), "w") as f:
        f.write("<html>")
        f.writelines(
            f"<img src='data:image/svg+xml;base64,{fig}'/>"
            for fig in figs
        )
        f.write("</html>")


def contrast_color(r, g, b):
    luma = ((0.299 * r) + (0.587 * g) + (0.114 * b)) / 255
    return "black" if luma > 0.5 else "white"