<filename>src/pltviz/pie.py
"""
Pie Plot
--------

Contents:
    pie
"""

import matplotlib.pyplot as plt
import seaborn as sns
from colormath.color_objects import sRGBColor
from poli_sci_kit.appointment.methods import highest_averages

from pltviz import utils

default_sat = 0.95


def pie(
    counts,
    labels=None,
    faction_labels=None,
    colors=None,
    radius=1,
    outer_ring_density=100,
    donut_ratio=1,
    display_labels=False,
    display_counts=False,
    label_font_size=20,
    dsat=default_sat,
    axis=None,
):
    """
    Produces a donut plot of group (and faction) shares or allocations.

    Parameters
    ----------
        counts : list or list of lists (contains ints or floats)
            The data to be plotted.

            Note: a list of lists produces a two layer plot where sublists define factions.

        labels : list : optional (default=None; contains strs)
            The labels of the groups.

        faction_labels : list : optional (default=None; contains strs)
            The labels of potential factions.

        colors : list or list of lists : optional (default=None)
            The colors of the groups as hex keys.

        outer_ring_density : int (default=500)
            The density of the faction ring gradients.

        radius : float : optional (default=1)
            The size of the plot.

        donut_ratio : float (default=1, a full circle)
            The ratio of the center radius of a donut to the whole.

        display_labels : bool : optional (default=False)
            Whether to display the labels of the groups (or factions if they're included).

                Note: labels can only be displayed for groups or factions, faction labels are if they're included.

                Note: if factions are included, then a legend should be used for the inner group labels (see package examples).

        display_counts : bool : optional (default=False)
            Whether to display the counts of the groups or factions.

        label_font_size : int (default=20)
            The size of the text in the labels.

        dsat : float : optional (default=default_sat)
            The degree of desaturation to be applied to the colors.

        axis : str : optional (default=None)
            Adds an axis to plots so they can be combined.

    Returns
    -------
        ax : matplotlib.pyplot.subplot
            A donut plot that depicts shares or allocations (potentially including factions).
    """
    if faction_labels:
        assert (
            list(set([type(count) for count in counts]))[0] == list
            and len(set([type(count) for count in counts])) == 1
        ), "If plotting groups and their factions, then the 'counts' argument must be a list of lists, where sublists are group counts in the given faction."

    if (
        list(set([type(count) for count in counts]))[0] == list
        and len(set([type(count) for count in counts])) == 1
    ):
        assert (
            faction_labels
        ), "A list of lists has been provided for 'counts', implying that factions should also be represented, but no labels for the factions have been provided."

    if list in [type(item) for item in counts]:
        total_groups = len([item for sublist in counts for item in sublist])
    else:
        total_groups = len(counts)

    if colors:
        assert (
            len(colors) == total_groups
        ), "The number of colors provided doesn't match the number of counts to be displayed."

    elif colors == None:
        sns.set_palette("deep")  # default sns palette
        colors = [
            utils.rgb_to_hex(c)
            for c in sns.color_palette(n_colors=total_groups, desat=1)
        ]

    colors = [
        utils.scale_saturation(rgb_trip=utils.hex_to_rgb(c), sat=dsat) for c in colors
    ]
    colors = [utils.rgb_to_hex(c) for c in colors]

    if axis:
        ax = axis  # to mirror seaborn axis plotting
    else:
        ax = plt.subplots(1, 1)[1]

    if faction_labels:
        faction_counts = [sum(sub_list) for sub_list in counts]

        # Indexes for later colors.
        group_indexes = list(range(total_groups))
        factioned_indexes = utils.gen_list_of_lists(
            original_list=group_indexes,
            new_structure=[len(sublist) for sublist in counts],
        )

        # Outer sections to be colored and determined by outer_ring_density.
        outer_ring_sections = [1 for i in range(outer_ring_density)]

        # Convert colors to rgb and classify them into factions.
        rgb_colors = [utils.hex_to_rgb(c) for c in colors]
        faction_colors = [
            [rgb_colors[i] for i in sublist] for sublist in factioned_indexes
        ]

        # Use the Jefferson highest_averages method divide the outer ring
        # based on the proportions of the factions.
        faction_sections = highest_averages(
            shares=faction_counts, total_alloc=len(outer_ring_sections)
        )

        outer_ring_colors = []
        for faction_index in range(len(faction_labels)):
            # Use the Jefferson highest_averages method again to allocate
            # the faction's outer ring sections to colors.
            # This would contain allocations for len(counts[faction_index]) colors,
            # but there are len(counts[faction_index])-1 gradients.
            # Thus average over Jefferson highest_averages allocations when
            # each element is removed for appropriately weighted gradients.
            if len(counts[faction_index]) == 1:
                averaged_allocations = [faction_sections[faction_index]]
            else:
                one_removed_allocations = [
                    highest_averages(
                        shares=counts[faction_index][:i]
                        + counts[faction_index][i + 1 :],
                        total_alloc=faction_sections[faction_index],
                    )
                    for i in range(len(counts[faction_index]))
                ]
                averaged_allocations = [
                    sum(
                        [
                            sub_allocation[i]
                            for sub_allocation in one_removed_allocations
                        ]
                    )
                    / len(one_removed_allocations)
                    for i in range(len(one_removed_allocations[0]))
                ]

            averaged_allocations = [round(i) for i in averaged_allocations]

            # Correct in case of rounding errors.
            if sum(averaged_allocations) < faction_sections[faction_index]:
                averaged_allocations[0] += faction_sections[faction_index] - sum(
                    averaged_allocations
                )
            elif sum(averaged_allocations) > faction_sections[faction_index]:
                averaged_allocations[0] -= faction_sections[faction_index] - sum(
                    averaged_allocations
                )

            if len(faction_colors[faction_index]) == 1:
                # Assign sole group's color via a monochrome gradient.
                outer_ring_colors.append(
                    utils.create_color_palette(
                        start_rgb=faction_colors[faction_index][0],
                        end_rgb=faction_colors[faction_index][0],
                        num_colors=averaged_allocations[0],
                        colorspace=sRGBColor,
                    )
                )
            else:
                # Create a gradient mix of the faction colors for the section
                # of the outer ring.
                for color_index in range(len(faction_colors[faction_index]))[:-1]:
                    outer_ring_colors.append(
                        utils.create_color_palette(
                            start_rgb=faction_colors[faction_index][color_index],
                            end_rgb=faction_colors[faction_index][color_index + 1],
                            num_colors=averaged_allocations[color_index],
                            colorspace=sRGBColor,
                        )
                    )

        outer_ring_colors = [item for sublist in outer_ring_colors for item in sublist]

        # Flatten counts for the inner ring and labelling.
        counts = [item for sublist in counts for item in sublist]

        if display_labels:
            outer_ring_labels = []
            labels = [""] * len(counts)
            # Place labels in the middle of the faction's arc, and make the others blank.
            factions_index = 0
            label_index = faction_counts[factions_index] / 2  # in the middle
            for i in range(outer_ring_density):
                if i == round(label_index / sum(faction_counts) * outer_ring_density):
                    if display_counts:
                        outer_ring_labels.append(
                            f"{faction_labels[factions_index]}: {faction_counts[factions_index]}"
                        )
                    else:
                        outer_ring_labels.append(faction_labels[factions_index])

                    factions_index += 1
                    if factions_index < len(faction_counts):
                        label_index += faction_counts[factions_index - 1] / 2
                        label_index += faction_counts[factions_index] / 2
                else:
                    outer_ring_labels.append("")

        else:
            label_font_size = 0
            outer_ring_labels = [""] * outer_ring_density
            labels = [""] * len(counts)

        outer_ring, _ = ax.pie(
            x=outer_ring_sections,
            radius=radius + (0.2 * radius),
            labels=outer_ring_labels,
            colors=outer_ring_colors,
            textprops={"fontsize": label_font_size},
        )
        plt.setp(obj=outer_ring, width=0.3 * radius, linewidth=0)

    if labels == None:
        labels = [f"group_{i}" for i in range(len(counts))]

    else:
        if display_counts:
            labels = [f"{lbl}: {counts[i]}" for i, lbl in enumerate(labels)]
        else:
            # Remove labels for those that have 0 counts to avoid confusion.
            labels = [lbl if counts[i] > 0 else "" for i, lbl in enumerate(labels)]

        if not display_labels:
            labels = [""] * len(counts)

    inner_ring, _ = ax.pie(
        x=counts,
        radius=radius,
        labels=labels,
        colors=colors,
        textprops={"fontsize": label_font_size},
    )
    plt.setp(obj=inner_ring, width=radius * donut_ratio, edgecolor="white")

    return ax

<filename>src/pltviz/pie.py
"""
Pie Plot
--------

Contents:
    pie
"""

import matplotlib.pyplot as plt
import seaborn as sns
from colormath.color_objects import sRGBColor
from poli_sci_kit.appointment.methods import highest_averages

from pltviz import utils

default_sat = 0.95


def pie(
    counts,
    labels=None,
    faction_labels=None,
    colors=None,
    radius=1,
    outer_ring_density=100,
    donut_ratio=1,
    display_labels=False,
    display_counts=False,
    label_font_size=20,
    dsat=default_sat,
    axis=None,
):
    """
    Produces a donut plot of group (and faction) shares or allocations.

    Parameters
    ----------
        counts : list or list of lists (contains ints or floats)
            The data to be plotted.

            Note: a list of lists produces a two layer plot where sublists define factions.

        labels : list : optional (default=None; contains strs)
            The labels of the groups.

        faction_labels : list : optional (default=None; contains strs)
            The labels of potential factions.

        colors : list or list of lists : optional (default=None)
            The colors of the groups as hex keys.

        outer_ring_density : int (default=500)
            The density of the faction ring gradients.

        radius : float : optional (default=1)
            The size of the plot.

        donut_ratio : float (default=1, a full circle)
            The ratio of the center radius of a donut to the whole.

        display_labels : bool : optional (default=False)
            Whether to display the labels of the groups (or factions if they're included).

                Note: labels can only be displayed for groups or factions, faction labels are if they're included.

                Note: if factions are included, then a legend should be used for the inner group labels (see package examples).

        display_counts : bool : optional (default=False)
            Whether to display the counts of the groups or factions.

        label_font_size : int (default=20)
            The size of the text in the labels.

        dsat : float : optional (default=default_sat)
            The degree of desaturation to be applied to the colors.

        axis : str : optional (default=None)
            Adds an axis to plots so they can be combined.

    Returns
    -------
        ax : matplotlib.pyplot.subplot
            A donut plot that depicts shares or allocations (potentially including factions).
    """
    if faction_labels:
        assert (
            list(set([type(count) for count in counts]))[0] == list
            and len(set([type(count) for count in counts])) == 1
        ), "If plotting groups and their factions, then the 'counts' argument must be a list of lists, where sublists are group counts in the given faction."

    if (
        list(set([type(count) for count in counts]))[0] == list
        and len(set([type(count) for count in counts])) == 1
    ):
        assert (
            faction_labels
        ), "A list of lists has been provided for 'counts', implying that factions should also be represented, but no labels for the factions have been provided."

    if list in [type(item) for item in counts]:
        total_groups = len([item for sublist in counts for item in sublist])
    else:
        total_groups = len(counts)

    if colors:
        assert (
            len(colors) == total_groups
        ), "The number of colors provided doesn't match the number of counts to be displayed."

    elif colors == None:
        sns.set_palette("deep")  # default sns palette
        colors = [
            utils.rgb_to_hex(c)
            for c in sns.color_palette(n_colors=total_groups, desat=1)
        ]

    colors = [
        utils.scale_saturation(rgb_trip=utils.hex_to_rgb(c), sat=dsat) for c in colors
    ]
    colors = [utils.rgb_to_hex(c) for c in colors]

    if axis:
        ax = axis  # to mirror seaborn axis plotting
    else:
        ax = plt.subplots(1, 1)[1]

    if faction_labels:
        faction_counts = [sum(sub_list) for sub_list in counts]

        # Indexes for later colors.
        group_indexes = list(range(total_groups))
        factioned_indexes = utils.gen_list_of_lists(
            original_list=group_indexes,
            new_structure=[len(sublist) for sublist in counts],
        )

        # Outer sections to be colored and determined by outer_ring_density.
        outer_ring_sections = [1 for i in range(outer_ring_density)]

        # Convert colors to rgb and classify them into factions.
        rgb_colors = [utils.hex_to_rgb(c) for c in colors]
        faction_colors = [
            [rgb_colors[i] for i in sublist] for sublist in factioned_indexes
        ]

        # Use the Jefferson highest_averages method divide the outer ring
        # based on the proportions of the factions.
        faction_sections = highest_averages(
            shares=faction_counts, total_alloc=len(outer_ring_sections)
        )

        outer_ring_colors = []
        for faction_index in range(len(faction_labels)):
            # Use the Jefferson highest_averages method again to allocate
            # the faction's outer ring sections to colors.
            # This would contain allocations for len(counts[faction_index]) colors,
            # but there are len(counts[faction_index])-1 gradients.
            # Thus average over Jefferson highest_averages allocations when
            # each element is removed for appropriately weighted gradients.
            if len(counts[faction_index]) == 1:
                averaged_allocations = [faction_sections[faction_index]]
            else:
                one_removed_allocations = [
                    highest_averages(
                        shares=counts[faction_index][:i]
                        + counts[faction_index][i + 1 :],
                        total_alloc=faction_sections[faction_index],
                    )
                    for i in range(len(counts[faction_index]))
                ]
                averaged_allocations = [
                    sum(
                        [
                            sub_allocation[i]
                            for sub_allocation in one_removed_allocations
                        ]
                    )
                    / len(one_removed_allocations)
                    for i in range(len(one_removed_allocations[0]))
                ]

            averaged_allocations = [round(i) for i in averaged_allocations]

            # Correct in case of rounding errors.
            if sum(averaged_allocations) < faction_sections[faction_index]:
                averaged_allocations[0] += faction_sections[faction_index] - sum(
                    averaged_allocations
                )
            elif sum(averaged_allocations) > faction_sections[faction_index]:
                averaged_allocations[0] -= faction_sections[faction_index] - sum(
                    averaged_allocations
                )

            if len(faction_colors[faction_index]) == 1:
                # Assign sole group's color via a monochrome gradient.
                outer_ring_colors.append(
                    utils.create_color_palette(
                        start_rgb=faction_colors[faction_index][0],
                        end_rgb=faction_colors[faction_index][0],
                        num_colors=averaged_allocations[0],
                        colorspace=sRGBColor,
                    )
                )
            else:
                # Create a gradient mix of the faction colors for the section
                # of the outer ring.
                for color_index in range(len(faction_colors[faction_index]))[:-1]:
                    outer_ring_colors.append(
                        utils.create_color_palette(
                            start_rgb=faction_colors[faction_index][color_index],
                            end_rgb=faction_colors[faction_index][color_index + 1],
                            num_colors=averaged_allocations[color_index],
                            colorspace=sRGBColor,
                        )
                    )

        outer_ring_colors = [item for sublist in outer_ring_colors for item in sublist]

        # Flatten counts for the inner ring and labelling.
        counts = [item for sublist in counts for item in sublist]

        if display_labels:
            outer_ring_labels = []
            labels = [""] * len(counts)
            # Place labels in the middle of the faction's arc, and make the others blank.
            factions_index = 0
            label_index = faction_counts[factions_index] / 2  # in the middle
            for i in range(outer_ring_density):
                if i == round(label_index / sum(faction_counts) * outer_ring_density):
                    if display_counts:
                        outer_ring_labels.append(
                            f"{faction_labels[factions_index]}: {faction_counts[factions_index]}"
                        )
                    else:
                        outer_ring_labels.append(faction_labels[factions_index])

                    factions_index += 1
                    if factions_index < len(faction_counts):
                        label_index += faction_counts[factions_index - 1] / 2
                        label_index += faction_counts[factions_index] / 2
                else:
                    outer_ring_labels.append("")

        else:
            label_font_size = 0
            outer_ring_labels = [""] * outer_ring_density
            labels = [""] * len(counts)

        outer_ring, _ = ax.pie(
            x=outer_ring_sections,
            radius=radius + (0.2 * radius),
            labels=outer_ring_labels,
            colors=outer_ring_colors,
            textprops={"fontsize": label_font_size},
        )
        plt.setp(obj=outer_ring, width=0.3 * radius, linewidth=0)

    if labels == None:
        labels = [f"group_{i}" for i in range(len(counts))]

    else:
        if display_counts:
            labels = [f"{lbl}: {counts[i]}" for i, lbl in enumerate(labels)]
        else:
            # Remove labels for those that have 0 counts to avoid confusion.
            labels = [lbl if counts[i] > 0 else "" for i, lbl in enumerate(labels)]

        if not display_labels:
            labels = [""] * len(counts)

    inner_ring, _ = ax.pie(
        x=counts,
        radius=radius,
        labels=labels,
        colors=colors,
        textprops={"fontsize": label_font_size},
    )
    plt.setp(obj=inner_ring, width=radius * donut_ratio, edgecolor="white")

    return ax

Pie Plot --------  Contents:     pie Produces a donut plot of group (and faction) shares or allocations.      Parameters     ----------         counts : list or list of lists (contains ints or floats)             The data to be plotted.              Note: a list of lists produces a two layer plot where sublists define factions.          labels : list : optional (default=None; contains strs)             The labels of the groups.          faction_labels : list : optional (default=None; contains strs)             The labels of potential factions.          colors : list or list of lists : optional (default=None)             The colors of the groups as hex keys.          outer_ring_density : int (default=500)             The density of the faction ring gradients.          radius : float : optional (default=1)             The size of the plot.          donut_ratio : float (default=1, a full circle)             The ratio of the center radius of a donut to the whole.          display_labels : bool : optional (default=False)             Whether to display the labels of the groups (or factions if they're included).                  Note: labels can only be displayed for groups or factions, faction labels are if they're included.                  Note: if factions are included, then a legend should be used for the inner group labels (see package examples).          display_counts : bool : optional (default=False)             Whether to display the counts of the groups or factions.          label_font_size : int (default=20)             The size of the text in the labels.          dsat : float : optional (default=default_sat)             The degree of desaturation to be applied to the colors.          axis : str : optional (default=None)             Adds an axis to plots so they can be combined.      Returns     -------         ax : matplotlib.pyplot.subplot             A donut plot that depicts shares or allocations (potentially including factions). # default sns palette # to mirror seaborn axis plotting # Indexes for later colors. # Outer sections to be colored and determined by outer_ring_density. # Convert colors to rgb and classify them into factions. # Use the Jefferson highest_averages method divide the outer ring # based on the proportions of the factions. # Use the Jefferson highest_averages method again to allocate # the faction's outer ring sections to colors. # This would contain allocations for len(counts[faction_index]) colors, # but there are len(counts[faction_index])-1 gradients. # Thus average over Jefferson highest_averages allocations when # each element is removed for appropriately weighted gradients. # Correct in case of rounding errors. # Assign sole group's color via a monochrome gradient. # Create a gradient mix of the faction colors for the section # of the outer ring. # Flatten counts for the inner ring and labelling. # Place labels in the middle of the faction's arc, and make the others blank. # in the middle # Remove labels for those that have 0 counts to avoid confusion.