"""
This module defines some symmetry metrics.

Methods
-------

"""
from __future__ import annotations

import logging
import math
from enum import Enum
from typing import Union, Tuple

import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
from scipy import ndimage
from scipy.spatial import ConvexHull
from skimage import io

from metricX import crossings, common, boundary, logConfig
from metricX.common import numeric, Vector

logger = logConfig.get_logger(__name__)


def _flip_point_around_axis(p: np.array, a: np.array, b: np.array) -> np.array:
    m = (a + b) / 2
    v_ab = b - a
    v_mp = p - m

    # Obtain the perpendicular bisector
    v_ab_rot = np.asarray((v_ab[1], v_ab[0] * -1))
    v_ab_rot_norm = v_ab_rot / np.linalg.norm(v_ab_rot)

    projection = m + v_mp.dot(
        v_ab_rot_norm) * v_ab_rot_norm

    v_pl = projection - p
    return p + 2 * v_pl


def reflective_symmetry(g: nx.Graph, pos: Union[str, dict, None] = None, threshold: int = 10,
                        tolerance: float = 1e-2, fraction: float = 1) -> float:
    r"""
    Computes a metric for axial symmetry between 0 and 1 as defined by H. Purchase [1].

    The metric sums up the symmetry for all axes defined by pairs of points. For each such axis, a symmetry value
    is calculated for all subgraphs with sufficiently many reflected nodes. The symmetry values for each subgraph
    are weighted by the area of the subgraph and summed up.

    Note that, with a worst-case runtime of $O(n^7)$ and a best-case runtime of $O(n^5)$,
    the metric is computationally very expensive.

    [1] H. Purchase, "Metrics for Graph Drawing Aesthetics",
    Journal of Visual Languages & Computing 13.5 (2002): 501-516.
    :param g: A networkX graph
    :type g: nx.Graph
    :param pos: Optional node position dictionary. If not supplied, node positions are read from the graph directly.
            If given as a string, the property under the given name in the networkX graph is used.
    :type pos: Union[str, dic, None]
    :param threshold: The minimum number of edges a subgraph needs in order to be considered to be symmetric
    :type threshold: int
    :param tolerance: The tolerance for when two points should be considered to be at the same position.
    :type tolerance: float
    :param fraction: A weighing of how much crossings and endpoints should be distinguished between 0 and 1. 1 means
            that we do not care about whether a point is a crossing or an endpoint regarding detecting symmetry and the
            two are treated equally.
    :type fraction:
    :return: Axial symmetry estimate between 0 and 1
    :rtype: float
    """

    def _flip_edge(e, a, b):
        edge_a = np.asarray(pos[e[0]])
        edge_b = np.asarray(pos[e[1]])

        edge_a_flipped = _flip_point_around_axis(edge_a, a, b)
        edge_b_flipped = _flip_point_around_axis(edge_b, a, b)

        return edge_a_flipped, edge_b_flipped

    def _points_within_tolerance(point_a, point_b) -> bool:
        return abs(point_a[0] - point_b[0]) <= tolerance and abs(point_a[1] - point_b[1]) <= tolerance

    def _endpoints_overlap(edge_to_be_flipped, edge_other, a, b):
        edge_flipped = _flip_edge(edge_to_be_flipped, a, b)
        other_edge_pos = (pos[edge_other[0]], pos[edge_other[1]])

        if (_points_within_tolerance(edge_flipped[0], other_edge_pos[0])
                and _points_within_tolerance(edge_flipped[1], other_edge_pos[1])):
            return (edge_to_be_flipped[0], edge_other[0]), (edge_to_be_flipped[1], edge_other[1])
        elif (_points_within_tolerance(edge_flipped[0], other_edge_pos[1])
              and _points_within_tolerance(edge_flipped[1], other_edge_pos[0])):
            return (edge_to_be_flipped[0], edge_other[1]), (edge_to_be_flipped[1], edge_other[0])

        return None

    def _subgraph_area(nodes) -> float:
        ch = ConvexHull([pos[node] for node in nodes])
        return ch.volume

    def _is_crossing_node(node):
        node = g.nodes[node]
        return 'is_elevated_crossing' in node and node['is_elevated_crossing']

    def _subgraph_symmetry(edge_pairs) -> float:

        if fraction == 1:
            return 1
        else:
            total = 0

            for pair_1, pair_2 in edge_pairs:
                factor_1 = fraction if _is_crossing_node(pair_1[0]) != _is_crossing_node(pair_1[1]) else 1
                factor_2 = fraction if _is_crossing_node(pair_2[0]) != _is_crossing_node(pair_2[1]) else 1

                total += factor_1 * factor_2

        return total / len(edge_pairs)

    if g.order() <= 1:
        return 1

    convex_hull_area = boundary.area_tight(g, pos)
    if convex_hull_area <= 0:
        return 1

    # Get node positions
    g = nx.edge_subgraph(g, g.edges()).copy()

    pos = common.get_node_positions(g, pos)
    nx.set_node_attributes(g, pos, 'pos')

    # Planarize by replacing crossings with nodes
    crossings.planarize(g)
    pos = common.get_node_positions(g)

    m = len(g.edges())
    n = len(g.nodes())

    total_area = 0.0
    total_symmetry = 0.0

    edge_list = list(g.edges())
    node_list = list(g.nodes())

    for i_a, node_a in enumerate(node_list):
        for i_b in range(i_a, n):
            node_b = node_list[i_b]

            if node_a == node_b or pos[node_a] == pos[node_b]:
                continue

            # Iterate over all node pairs

            pos_a = np.asarray(pos[node_a])
            pos_b = np.asarray(pos[node_b])

            subgraph_edge_pairs = []
            subgraph_size = 0

            # All edges that are symmetrical along the flip axis get added to the new subgraph
            for i, edge in enumerate(edge_list):
                for j in range(i, m):

                    other_edge = edge_list[j]

                    edge_tuple = _endpoints_overlap(edge, other_edge, pos_a, pos_b)

                    if edge_tuple is not None:
                        if edge == other_edge:
                            subgraph_size += 1
                        else:
                            subgraph_size += 2

                        subgraph_edge_pairs.append(edge_tuple)

            if subgraph_size >= threshold:
                sub_sym = _subgraph_symmetry(subgraph_edge_pairs)
                contained_nodes = set(
                    [node for edge_pair in subgraph_edge_pairs for edge in edge_pair for node in edge])
                sub_area = _subgraph_area(contained_nodes)

                total_symmetry += sub_sym * sub_area
                total_area += sub_area

    return total_symmetry / max(total_area, convex_hull_area)


class SymmetryType(Enum):
    """ Type of symmetry used for :func:`edge_based_symmetry` """
    REFLECTIVE = 0,
    ROTATIONAL = 1,
    TRANSLATIONAL = 2


# TODO use precision for comparisons


class _SIFTFeature:

    def __init__(self, edge: Tuple[numeric, numeric], edge_pos, score):

        point_a = common.Vector.from_point(edge_pos[0])
        point_b = common.Vector.from_point(edge_pos[1])

        self.mid = (point_a + point_b) / 2
        self.length = common.euclidean_distance(edge_pos[0], edge_pos[1])
        self.angle = ((common.Vector.from_point(edge_pos[1]) - (common.Vector.from_point(edge_pos[0])))
                      .upwardAndForwardDirection().angleToX())
        self.edge = edge
        self.score = score

    def compare_scale(self, other: _SIFTFeature, sigma_scale: float) -> float:
        """
            Calculates the scale facture between two SIFT features in the range of [0,1]
                :param other: The second sift feature to compare to
        :type other: _SIFTFeature
        :param sigma_scale: TODO
        :type sigma_scale: TODO
        :return: A facture between 0 and 1
        :rtype: float
        """
        top = -abs(self.length - other.length)
        bottom = sigma_scale * (self.length + other.length)

        if bottom == 0:
            return

        return math.exp(top / bottom) ** 2

    def compare_distance(self, other: _SIFTFeature, sigma_distance: float, is_distance_bound: bool):
        if not is_distance_bound:
            return 1
        dist = self.mid.euclidean_distance(other.mid)
        return math.exp(- (dist ** 2) / (2 * sigma_distance * sigma_distance))

    def compare_rotation(self, other: _SIFTFeature, theta: float = None):
        if theta is None:
            return abs(math.cos(self.angle - other.angle))
        else:
            return abs(math.cos(self.angle + other.angle - 2 * theta))


class _Axis:

    def __init__(self, pos, edge_a, edge_b, score):
        self.pos = pos
        self.edge_a = edge_a
        self.edge_b = edge_b
        self.score = score

    def __str__(self):
        return f"_Vote({self.pos}, {self.edge_a}, {self.edge_b}, {self.score})"


class _AxesSystem:

    def __init__(self, pos: dict, sigma_scale: float, distance_bound: bool, sigma_distance: float, num_features: int,
                 angle_merge: float, pixel_merge: float, x_min: int, y_min: int):
        self.pos = pos
        self.y_min = y_min
        self.x_min = x_min
        self.y_merge = pixel_merge
        self.x_merge = angle_merge
        self.num_features = num_features
        self.sigma_distance = sigma_distance
        self.is_distance_bound = distance_bound
        self.sigma_scale = sigma_scale
        self.votes = []

    def add_rotational_vote(self, feature):
        self.votes.append(_Axis(feature.mid, feature.edge, feature.edge, 1))
        pass

    def add_rotational_vote_from_two(self, feature_a, feature_b):
        sca = feature_a.compare_scale(feature_b, self.sigma_scale)
        dis = feature_a.compare_distance(feature_b, self.sigma_distance, self.is_distance_bound)
        score = sca * dis

        mid = feature_a.mid.mid(feature_b.mid)
        mid_diff = feature_b.mid - feature_a.mid
        theta_ij = math.atan2(mid_diff.y, mid_diff.x) + math.pi / 2

        delta_dist = feature_a.mid.euclidean_distance(feature_b.mid) / 2
        delta_angl = abs(feature_b.angle - feature_a.angle) / 2

        if math.tan(delta_angl) == 0:
            # The two edges are parallel -> rotation equals 0° -> skip
            return

        radius = delta_dist / math.tan(delta_angl)

        if abs(radius) < 0.001 or abs(delta_angl) < 0.5 or abs(delta_angl - math.pi) < 0.05:
            # The rotational symmetry is centered at the mid itself
            vote = _Axis(mid, feature_a.edge, feature_b.edge, score)
            self.votes.append(vote)
        else:
            # Otherwise, find rotation points on boundary of radius (two opposite points)

            direction = common.Vector(1, 0).rotate(common.Angle(theta_ij))

            scaled_direction = direction * radius

            center_a = mid + scaled_direction
            center_b = mid - scaled_direction

            theta_diff = feature_b.angle - feature_a.angle

            test_a = (feature_a.mid - center_a).rotate(theta_diff) + center_a
            test_b = (feature_a.mid - center_b).rotate(theta_diff) + center_b

            if (test_a.euclidean_distance(feature_b.mid) < 0.0001):
                center = center_a
                factor = -1
            elif (test_b.euclidean_distance(feature_b.mid) < 0.0001):
                center = center_b
                factor = 1
            else:
                # Should never by the case
                return

            vote = _Axis(center, feature_a.edge, feature_b.edge, score)
            self.votes.append(vote)

            # Opposite point
            if radius > 0:
                delta_angl_flipped = (math.pi - delta_angl * 2) / 2
                radius_2 = delta_dist / math.tan(delta_angl_flipped)

                scaled_direction = direction * radius_2 * factor
                center = mid + scaled_direction

                vote = _Axis(center, feature_a.edge, feature_b.edge, score)
                self.votes.append(vote)

    @staticmethod
    def _parse_reflective_vote(angle: common.Angle, radius: float, score: float, edge_a, edge_b) -> _Axis:
        while angle.rad < 0:
            angle += math.pi

        if radius < 0:
            radius = abs(radius)
            angle += math.pi
            while angle > math.pi * 2:
                angle -= math.pi

        return _Axis(common.Vector(angle.deg(), radius), edge_a, edge_b, score)

    def add_reflective_vote(self, feature):

        mid = feature.mid

        edge_pos_a = common.Vector.from_point(self.pos[feature.edge[0]])
        edge_pos_b = common.Vector.from_point(self.pos[feature.edge[1]])

        theta_ij = common.Angle((edge_pos_a - edge_pos_b).angleToX() % math.pi)
        theta_ij_bisector = common.Angle(theta_ij + math.pi / 2)

        # Symmetry axis along perpendicular bisector
        r_ij = mid.x * math.cos(theta_ij_bisector.rad) + mid.y * math.sin(theta_ij_bisector.rad)
        self.votes.append(self._parse_reflective_vote(theta_ij_bisector, r_ij, 1, feature.edge, feature.edge))

        # Symmetry axis along edge itself
        r_ij = mid.x * math.cos(theta_ij.rad) + mid.y * math.sin(theta_ij.rad)
        self.votes.append(self._parse_reflective_vote(theta_ij, r_ij, 1, feature.edge, feature.edge))

    def add_reflective_vote_from_two(self, feature_a, feature_b):
        theta_ij = (feature_a.mid - feature_b.mid).angleToX() % math.pi
        mid = feature_a.mid.mid(feature_b.mid)
        r_ij = mid.x * math.cos(theta_ij) + mid.y * math.sin(theta_ij)

        rot = feature_a.compare_rotation(feature_b, theta_ij)
        sca = feature_a.compare_scale(feature_b, self.sigma_scale)
        dis = feature_a.compare_distance(feature_b, self.sigma_distance, self.is_distance_bound)
        score = rot * sca * dis

        self.votes.append(
            self._parse_reflective_vote(theta_ij, r_ij, score, feature_a.edge, feature_b.edge))

    def add_translative_vote(self, feature_a, feature_b):
        rot = feature_a.compare_rotation(feature_b)
        sca = feature_a.compare_scale(feature_b, self.sigma_scale)
        dis = feature_a.compare_distance(feature_b, self.sigma_distance, self.is_distance_bound)

        delta = feature_a.mid - feature_b.mid

        if delta.y < 0:
            delta *= -1

        score = rot * sca * dis
        vote = _Axis(delta, feature_a.edge, feature_b.edge, score)
        logger.debug(vote)
        self.votes.append(vote)

    def _points_too_close(self, pos_a, pos_b):
        return abs(pos_a.x - pos_b.x) < self.x_min and abs(pos_a.y - pos_b.y) < self.y_min

    def _find_maxima(self):
        def _merge_close_point(value: float, merge_span: float):
            return (int(value / merge_span)) * merge_span

        voting_dic = {}

        # Group together close points
        # TODO this can be done when adding already, might save us some time/space
        for vote in self.votes:
            x_merged = _merge_close_point(vote.pos.x, self.x_merge)
            y_merged = _merge_close_point(vote.pos.y, self.y_merge)

            key = Vector(x_merged, y_merged)

            if vote.score <= 0:
                continue

            if key in voting_dic:
                voting_dic[key] += vote.score
            else:
                voting_dic[key] = vote.score

        # Find max-scoring points
        sorted_axes = list(sorted(voting_dic.keys(), key=lambda x: voting_dic[x], reverse=True))

        # Remove double-entries that are too close
        for i in range(min(self.num_features, len(sorted_axes))):
            j = i + 1
            while j < len(sorted_axes):
                if self._points_too_close(sorted_axes[i], sorted_axes[j]):
                    sorted_axes.pop(j)
                else:
                    j += 1

        return sorted_axes[:min(self.num_features, len(sorted_axes))]

    def conclude_voting(self, g: nx.Graph) -> float:
        top_features = self._find_maxima()

        if len(top_features) == 0:
            return 0

        matching_edge_dic = set()

        for vote in self.votes:
            matched = False

            for feature in top_features:
                if self._points_too_close(vote.pos, feature):
                    matched = True
                    break

            if matched:
                matching_edge_dic.add(vote.edge_a)
                matching_edge_dic.add(vote.edge_b)

        return len(matching_edge_dic) / float(len(top_features) * len(g.edges()))


def edge_based_symmetry(g: nx.Graph, symmetry_type: SymmetryType, pos: Union[str, dict, None] = None,
                        axes_count: int = 1,
                        sigma_scale: float = 0.1, sigma_distance: float = 2,
                        is_distance_bound: bool = False,
                        angle_merge: float = 5, pixel_merge: float = 5,
                        x_min=10, y_min=10) -> float:
    pos = common.get_node_positions(g, pos)

    if len(g.edges()) == 0:
        return 1

    votes = _AxesSystem(pos, sigma_scale, is_distance_bound, sigma_distance, axes_count, angle_merge, pixel_merge,
                        x_min,
                        y_min)
    sift_features = [_SIFTFeature(e, (pos[e[0]], pos[e[1]]), 1) for e in g.edges()]

    for a in range(len(sift_features)):
        feature_a = sift_features[a]

        for b in range(a + 1, len(sift_features)):
            feature_b = sift_features[b]

            if symmetry_type == SymmetryType.REFLECTIVE:
                votes.add_reflective_vote_from_two(feature_a, feature_b)
            elif symmetry_type == SymmetryType.ROTATIONAL:
                votes.add_rotational_vote_from_two(feature_a, feature_b)
            elif symmetry_type == SymmetryType.TRANSLATIONAL:
                votes.add_translative_vote(feature_a, feature_b)

        if symmetry_type == SymmetryType.REFLECTIVE:
            votes.add_reflective_vote(feature_a)
        elif symmetry_type == SymmetryType.ROTATIONAL:
            votes.add_rotational_vote(feature_a)

    return votes.conclude_voting(g)


def visual_symmetry(g: nx.Graph, pos: Union[str, dict, None] = None, resolution: int = 100, sigma: float = 2) -> float:
    """
        Tries to estimate the perceived symmetry of the drawing by visually testing reflective, rotational and dihedral
        symmetry.
    :param g: graph
    :type g: nx.Graph
    :param pos: Dictionary of node positions. If not supplied, it is expected that the positions are included in g.
    :type pos: dic
    :param resolution: Width and height of the image the graph is drawn to. The higher the resolution, the more
    sensitive the measure is to local details.
    :type resolution: int
    :param sigma: Sigma of the Gaussian blur applied to the drawing of g. The higher the value, the more the symmetry
    reflects just the general shape of the graph instead of local details.
    :type sigma: float
    :return: A value in [0,1] estimating the perceived symmetry.
    :rtype: float
    """

    pos = common.get_node_positions(g, pos)

    if len(pos) <= 1:
        return 1

    pos = boundary.normalize_positions(g, pos)

    def _build_canvas():
        fig = plt.figure(1, figsize=(1, 1), dpi=resolution)
        ax = plt.axes([0., 0., 1., 1.], frameon=False, xticks=[], yticks=[])
        ax.set_xlim(-0.8, 0.8)
        ax.set_ylim(-0.8, 0.8)
        ax.set_facecolor((0, 0, 0))
        fig.patch.set_facecolor((0, 0, 0))
        canvas = fig.canvas
        return canvas, fig, ax

    def _draw_to_array(canvas):
        return np.frombuffer(canvas.tostring_rgb(), dtype='uint8').reshape(*reversed(canvas.get_width_height()), 3)

    def _get_edge_image():
        canvas, fig, ax = _build_canvas()
        nx.draw_networkx_edges(g, pos, edge_color='#f00', ax=ax)
        canvas.draw()
        image = _draw_to_array(canvas)
        fig.clf()
        return image

    def _get_node_image():
        canvas, fig, ax = _build_canvas()
        nx.draw_networkx_nodes(g, pos, node_size=.1, node_color='#0f0', ax=ax)
        canvas.draw()
        image = _draw_to_array(canvas)
        fig.clf()
        return image

    def _show_image(img):
        if logger.level == logging.DEBUG:
            plt.figure()
            plt.show()
            io.imshow(img)

    def _normalize_array(img):
        min_val = np.min(img)
        max_val = np.max(img)
        if min_val == max_val:
            return np.zeros_like(img)
        return (img - min_val) / (max_val - min_val)

    edge_image = _normalize_array(_get_edge_image()[:, :, 0])
    node_image = _normalize_array(_get_node_image()[:, :, 1])
    graph_image = edge_image + node_image
    _show_image(graph_image)

    # Blur
    graph_image = ndimage.gaussian_filter(graph_image, sigma=sigma)
    _show_image(graph_image)

    # Translate to centroid
    centroid = ndimage.center_of_mass(graph_image)
    shape = graph_image.shape
    graph_image = ndimage.shift(graph_image,
                                np.array([shape[0] / 2 - centroid[0],
                                          shape[1] / 2 - centroid[1]]))
    _show_image(graph_image)

    original_sum = np.sum(graph_image)
    rotational_estimate = 0
    reflective_estimate = 0
    dihedral_estimate = 0

    def _get_axis_weight(degree) -> float:
        if degree % 180 == 0:
            # Vertical
            return .5
        if degree % 90 == 0:
            # Horizontal
            return .3
        if degree % 45 == 0:
            # In between
            return .1
        return 0

    for i in range(0, 180, 45):
        # Rotate image
        rotated_image = ndimage.rotate(graph_image, i, reshape=False)
        flipped_and_rotated = np.flip(rotated_image, 0)

        # Rotational symmetry
        rotational_difference_image = rotated_image - graph_image
        rotational_difference = np.sum(np.abs(rotational_difference_image))
        rotational_estimate += _get_axis_weight(i) * rotational_difference

        # Reflective symmetry
        reflective_difference_image = flipped_and_rotated - rotated_image
        reflective_difference = np.sum(np.abs(reflective_difference_image))
        reflective_estimate += _get_axis_weight(i) * reflective_difference
        _show_image(reflective_difference_image)

        # Dihedral symmetry
        dihedral_difference_image = flipped_and_rotated - graph_image
        dihedral_difference = np.sum(np.abs(dihedral_difference_image))
        dihedral_estimate += _get_axis_weight(i) * dihedral_difference

    return 1 - math.pow(min(rotational_estimate, reflective_estimate, dihedral_estimate) / original_sum, 2)
