<commit_before>import pyx
from pyx.metapost.path import beginknot, endknot, smoothknot, tensioncurve
import networkx as nx
import numpy as np
from collections import OrderedDict
from shapely import geometry as geom
from geom_utils import linestring_to_coords, get_projecting_line, get_projection, place_point_on_line
import copy
from collections import Counter
from colour_utils import convert_colour, categories_to_float
import random


class HivePlot():
    def __init__(self, network, node_class_attribute="type", node_class_values=None, **kwargs):
        """
        :param network: NetworkX representation of a network to be plotted
        :param node_class_attribute: The attribute of nodes on which to split them onto 1-3 axes (default: "type")
        :param node_class_values: Values of node_class_attribute which will be included in the plot (in clockwise order)
        :param kwargs: Dictionary of attributes which can override default behaviour of the plot
        :return:
        """
        self.network = network
        self.node_class_attribute = node_class_attribute
        self.node_classes = self._split_nodes(node_class_values)

        # define default behaviour, which can be overridden with kwargs

        # background parameters
        self.background_colour = "Black"

        # axis parameters
        self.axis_length = 10
        self.proportional_offset_from_intersection = 0.2  # as a proportion of the longest axis
        self.split_axes = []
        if self.split_axes:
            raise NotImplementedError("Axis splitting is not yet implemented.")
        self.normalise_axis_length = False
        self.axis_colour = "Gray"
        self.axis_thickness = 0.15

        # edge parameters
        self.edge_thickness_range = (0.005, 0.16)
        self.edge_colour_attribute = "weight"
        self.edge_colour_gradient = "Jet"
        self.edge_category_colours = None
        self.edge_curvature = 1.7
        self.normalise_edge_colours = False

        # axis label parameters
        self.label_colour = "White"
        self.label_size = 15

        # node parameters
        self.normalise_node_distribution = False
        self.node_size = 0.08
        self.node_colour_gradient = "GreenRed"

        # fields to be filled by the object
        self.canvas = None

        self.__dict__.update(kwargs)

    def _split_nodes(self, node_class_names):
        """
        Split nodes based on the attribute specified in self.node_class_attribute
        :param node_class_names: Values of node_class_attribute which will be included in the plot (in clockwise order)
        :return: A dictionary whose keys are the node class attribute values, and values are lists of nodes belonging to that class
        """
        node_attribute_dict = nx.get_node_attributes(self.network, self.node_class_attribute)

        if node_class_names is None:
            node_class_names = list(node_attribute_dict.values())
            if len(node_class_names) > 3:
                raise ValueError("Nodes should be in 3 or fewer classes based on their {} attribute.".format(
                    self.node_class_attribute))
            node_class_names = sorted(node_class_names)
        else:
            for_deletion = []
            for node in node_attribute_dict:
                if node_attribute_dict[node] not in node_class_names:
                    for_deletion.append(node)
            for fd in for_deletion:
                node_attribute_dict.pop(fd)

        split_nodes = OrderedDict([(node_class, []) for node_class in node_class_names])
        node_list = list(node_attribute_dict)
        for i, node_class in enumerate(node_attribute_dict.values()):
            split_nodes[node_class].append(node_list[i])

        return split_nodes

    def _get_working_nodes(self):
        """
        :return: The set of all nodes to be included in the plot
        """
        working_nodes = set()

        for node_class_values in self.node_classes.values():
            working_nodes |= set(node_class_values)

        return working_nodes

    def _draw_background(self):
        max_length = (self.axis_length * (1 + self.proportional_offset_from_intersection)) * self.edge_curvature
        colour = convert_colour(self.background_colour)
        self.canvas.fill(pyx.path.rect(-max_length, -max_length, 2 * max_length, 2 * max_length), [colour])

    def _draw_axes(self, axes):
        axis_colour = convert_colour(self.axis_colour)
        for axis in axes.values():
            self.canvas.stroke(pyx.path.line(*linestring_to_coords(axis)),
                               [pyx.style.linewidth(self.axis_thickness), axis_colour])

    def _draw_labels(self, axes):
        text_alignment_list = [
            [pyx.text.halign.boxcenter, pyx.text.valign.bottom],
            [pyx.text.halign.boxleft, pyx.text.valign.top],
            [pyx.text.halign.boxright, pyx.text.valign.top]
        ]
        text_alignment = dict(zip(list(axes), text_alignment_list))

        for axis_name in axes:
            line_end = axes[axis_name].coords[1]
            # txt_str = axis_name
            txt_str = self._colour_text(self._size_text(axis_name, self.label_size))
            self.canvas.text(line_end[0], line_end[1],
                             txt_str,
                             text_alignment[axis_name])

    def _size_text(self, text, size):
        return r"{\fontsize{" + str(size) + r"}{" + str(round(size*1.2)) + r"}\selectfont " + text + r"}"

    def _colour_text(self, text):
        return r"\textcolor{COL}{" + text + "}"

    def _draw_edges(self, edge_lines):

        for edge_dict in edge_lines:
            start = edge_dict["start"]
            end = edge_dict["end"]
            colour = edge_dict["colour"]

            if "crossing_point" in edge_dict:
                midpoint = edge_dict["crossing_point"]

                edgepath = pyx.metapost.path.path([
                    beginknot(*start), tensioncurve(),
                    smoothknot(*midpoint), tensioncurve(),
                    endknot(*end)
                ])
            else:
                edgepath = pyx.path.line(pyx.path.line(start[0], start[1], end[0], end[1]))

            self.canvas.stroke(edgepath, [
                pyx.style.linewidth(edge_dict["thickness"]),
                colour
            ])

    def _draw_nodes(self, node_positions):
        gradient = eval("pyx.color.gradient." + self.node_colour_gradient)
        node_position_weights = self._weight_by_colocation(node_positions)
        for coords, weight in node_position_weights.items():
            self.canvas.stroke(pyx.path.circle(coords[0], coords[1], self.node_size))
            self.canvas.fill(pyx.path.circle(coords[0], coords[1], self.node_size),
                             [gradient.getcolor(weight)])

    def draw(self, save_path=None):
        """
        Draw the graph using the current settings
        :param save_path: path of PDF file to save graph into
        :return: True for successful completion
        """

        self._setup_latex()

        self.canvas = pyx.canvas.canvas()
        axes = self._create_axes()
        node_positions = self._place_nodes(axes)

        if len(axes) == 3:
            mid_ax_lines = self._create_mid_ax_lines(axes)

            # edge_lines = self._create_straight_edge_info(node_positions, curved=False)
            edge_lines = self._create_edge_info(node_positions, curved=True, mid_ax_lines=mid_ax_lines)
        else:
            raise NotImplementedError("2-axis plots not yet implemented")

        self._draw_background()
        self._draw_axes(axes)
        self._draw_labels(axes)
        self._draw_edges(edge_lines)
        self._draw_nodes(node_positions)

        if not save_path is None:
            self.save_canvas(save_path)

        return True

    def _create_mid_ax_lines(self, axes):<commit_msg>Improved background size, start adding types to docstring<commit_after>import pyx
from pyx.metapost.path import beginknot, endknot, smoothknot, tensioncurve
import networkx as nx
import numpy as np
from collections import OrderedDict
from shapely import geometry as geom
from geom_utils import linestring_to_coords, get_projecting_line, get_projection, place_point_on_line
import copy
from collections import Counter
from colour_utils import convert_colour, categories_to_float
import random
import warnings


class HivePlot():
    def __init__(self, network, node_class_attribute="type", node_class_values=None, **kwargs):
        """
        :param network: network to be plotted
        :type network: nx.Graph
        :param node_class_attribute: The attribute of nodes on which to split them onto 1-3 axes (default: "type")
        :type node_class_attribute: str
        :param node_class_values: Values of node_class_attribute which will be included in the plot (in clockwise order)
        :type node_class_values: list
        :param kwargs: Dictionary of attributes which can override default behaviour of the plot
        :type kwargs: dict
        """
        self.network = network
        self.node_class_attribute = node_class_attribute
        self.node_classes = self._split_nodes(node_class_values)

        # define default behaviour, which can be overridden with kwargs

        # background parameters
        self.background_proportion = 1.2
        self.background_colour = "Black"

        # axis parameters
        self.axis_length = 10
        self.proportional_offset_from_intersection = 0.2  # as a proportion of the longest axis
        self.split_axes = []
        if self.split_axes:
            raise NotImplementedError("Axis splitting is not yet implemented.")
        self.normalise_axis_length = False
        self.axis_colour = "Gray"
        self.axis_thickness = 0.15

        # edge parameters
        self.edge_thickness_range = (0.005, 0.16)
        self.edge_colour_attribute = "weight"
        self.edge_colour_gradient = "Jet"
        self.edge_category_colours = None
        self.edge_curvature = 1.7
        self.normalise_edge_colours = False

        # axis label parameters
        self.label_colour = "White"
        self.label_size = 15

        # node parameters
        self.normalise_node_distribution = False
        self.node_size = 0.08
        self.node_colour_gradient = "GreenRed"

        # fields to be filled by the object
        self.canvas = None
        self._background_layer = None
        self._foreground_layer = None

        self.__dict__.update(kwargs)

    def _split_nodes(self, node_class_names):
        """
        Split nodes based on the attribute specified in self.node_class_attribute
        :param node_class_names: Values of node_class_attribute which will be included in the plot (in clockwise order)
        :type node_class_names: list
        :return: A dictionary whose keys are the node class attribute values, and values are lists of nodes belonging to that class
        :rtype: dict
        """
        node_attribute_dict = nx.get_node_attributes(self.network, self.node_class_attribute)

        if node_class_names is None:
            node_class_names = list(node_attribute_dict.values())
            if len(node_class_names) > 3:
                raise ValueError("Nodes should be in 3 or fewer classes based on their {} attribute.".format(
                    self.node_class_attribute))
            node_class_names = sorted(node_class_names)
        else:
            for_deletion = []
            for node in node_attribute_dict:
                if node_attribute_dict[node] not in node_class_names:
                    for_deletion.append(node)
            for fd in for_deletion:
                node_attribute_dict.pop(fd)

        split_nodes = OrderedDict([(node_class, []) for node_class in node_class_names])
        node_list = list(node_attribute_dict)
        for i, node_class in enumerate(node_attribute_dict.values()):
            split_nodes[node_class].append(node_list[i])

        return split_nodes

    def _get_working_nodes(self):
        """
        :return: The set of all nodes to be included in the plot
        :rtype: set
        """
        working_nodes = set()

        for node_class_values in self.node_classes.values():
            working_nodes |= set(node_class_values)

        return working_nodes

    def _draw_background(self):
        # max_length = self.axis_length * (self.background_proportion + self.proportional_offset_from_intersection)
        min_x, min_y, max_x, max_y = self._get_bbox()
        colour = convert_colour(self.background_colour)
        self._background_layer.fill(pyx.path.rect(min_x, min_y, max_x-min_x, max_y-min_y), [colour])

    def _get_bbox(self, extend=1.1):
        """
        :param extend: the proportion of the extent which should be added to the size of the bounding box
        :type extend: float or int
        :return: extents of bounding box
        :rtype: tuple
        """
        bbox = self.canvas.bbox()
        min_x = pyx.unit.tocm(bbox.left())
        min_y = pyx.unit.tocm(bbox.bottom())
        max_x = pyx.unit.tocm(bbox.right())
        max_y = pyx.unit.tocm(bbox.top())

        return min_x*extend, min_y*extend, max_x*extend, max_y*extend

    def _draw_axes(self, axes):
        axis_colour = convert_colour(self.axis_colour)
        for axis in axes.values():
            self._foreground_layer.stroke(pyx.path.line(*linestring_to_coords(axis)),
                               [pyx.style.linewidth(self.axis_thickness), axis_colour])

    def _draw_labels(self, axes):
        text_alignment_list = [
            [pyx.text.halign.boxcenter, pyx.text.valign.bottom],
            [pyx.text.halign.boxleft, pyx.text.valign.top],
            [pyx.text.halign.boxright, pyx.text.valign.top]
        ]
        text_alignment = dict(zip(list(axes), text_alignment_list))

        for axis_name in axes:
            line_end = axes[axis_name].coords[1]
            # txt_str = axis_name
            txt_str = self._colour_text(self._size_text(axis_name, self.label_size))
            with warnings.catch_warnings():
                warnings.simplefilter("ignore")
                self._foreground_layer.text(line_end[0], line_end[1],
                                 txt_str,
                                 text_alignment[axis_name])

    def _size_text(self, text, size):
        return r"{\fontsize{" + str(size) + r"}{" + str(round(size*1.2)) + r"}\selectfont " + text + r"}"

    def _colour_text(self, text):
        return r"\textcolor{COL}{" + text + "}"

    def _draw_edges(self, edge_lines):

        for edge_dict in edge_lines:
            start = edge_dict["start"]
            end = edge_dict["end"]
            colour = edge_dict["colour"]

            if "crossing_point" in edge_dict:
                midpoint = edge_dict["crossing_point"]

                edgepath = pyx.metapost.path.path([
                    beginknot(*start), tensioncurve(),
                    smoothknot(*midpoint), tensioncurve(),
                    endknot(*end)
                ])
            else:
                edgepath = pyx.path.line(pyx.path.line(start[0], start[1], end[0], end[1]))

            self._foreground_layer.stroke(edgepath, [
                pyx.style.linewidth(edge_dict["thickness"]),
                colour
            ])

    def _draw_nodes(self, node_positions):
        gradient = eval("pyx.color.gradient." + self.node_colour_gradient)
        node_position_weights = self._weight_by_colocation(node_positions)
        for coords, weight in node_position_weights.items():
            self._foreground_layer.stroke(pyx.path.circle(coords[0], coords[1], self.node_size))
            self._foreground_layer.fill(pyx.path.circle(coords[0], coords[1], self.node_size),
                             [gradient.getcolor(weight)])

    def draw(self, save_path=None):
        """
        Draw the graph using the current settings
        :param save_path: path of PDF file to save graph into
        :return: True for successful completion
        """

        self._setup_latex()

        self.canvas = pyx.canvas.canvas()
        self._background_layer = self.canvas.layer("background")
        self._foreground_layer = self.canvas.layer("foreground", above="background")
        axes = self._create_axes()
        node_positions = self._place_nodes(axes)

        if len(axes) == 3:
            mid_ax_lines = self._create_mid_ax_lines(axes)

            # edge_lines = self._create_straight_edge_info(node_positions, curved=False)
            edge_lines = self._create_edge_info(node_positions, curved=True, mid_ax_lines=mid_ax_lines)
        else:
            raise NotImplementedError("2-axis plots not yet implemented")

        self._draw_axes(axes)
        self._draw_labels(axes)
        self._draw_edges(edge_lines)
        self._draw_nodes(node_positions)
        self._draw_background()

        if not save_path is None:
            self.save_canvas(save_path)

        return True

    def _create_mid_ax_lines(self, axes):
{'fraction_non_alphanumeric': 0.08704572098475967, 'fraction_numerical': 0.005099648300117233, 'mean_word_length': 2.439717741935484, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2042623', 'n_tokens_mistral': 5125, 'n_tokens_neox': 4838, 'n_words': 1140}