#
# Plots a function defined on a two-dimensional space.
#
# This file is part of PINTS (https://github.com/pints-team/pints/) which is
# released under the BSD 3-clause license. See accompanying LICENSE.md for
# copyright notice and full license details.
#
# This ``surface`` method was adapted from Myokit (see http://myokit.org) by
# its original author.
#
# The code to plot voronoi regions was based on an example shown here:
# http://stackoverflow.com/a/20678647/423420
#
import numpy as np

import pints


def surface(points, values, boundaries=None, markers='+', figsize=None):
    """
    Takes irregularly spaced points and function evaluations in a
    two-dimensional parameter space and creates a coloured surface plot using a
    voronoi diagram.

    Returns a ``matplotlib`` figure object and axes handle.

    Parameters
    ----------
    points
        A list of (two-dimensional) points in parameter space.
    values
        The values (e.g. error measure evaluations) corresponding to these
        points.
    boundaries
        An optional :class:`pints.RectangularBoundaries` object to restrict the
        area shown. If set to ``None`` boundaries will be determined from the
        given ``points``.
    markers
        An optional string indicating the matplotlib markers to use to plot
        the ``points``. Set to ``None`` to hide.
    figsize
        An optional tuple ``(width, height)`` that will be passed to
        matplotlib when creating the figure. If set to ``None`` matplotlib will
        use its default figure size.

    """
    import matplotlib.pyplot as plt
    from matplotlib.collections import PolyCollection

    # Check points and values
    points = pints.matrix2d(points)
    n, d = points.shape
    if d != 2:
        raise ValueError('Only two-dimensional parameters are supported.')
    values = pints.vector(values)
    if len(values) != n:
        raise ValueError(
            'The number of values must match the number of points.')

    # Extract x and y points
    x = points[:, 0]
    y = points[:, 1]
    del(points)

    # Check boundaries
    if boundaries is None:
        xmin, xmax = np.min(x), np.max(x)
        r = 0.1 * (xmax - xmin)
        xmin -= r
        xmax += r
        ymin, ymax = np.min(y), np.max(y)
        r = 0.1 * (ymax - ymin)
        ymin -= r
        ymax += r
    else:
        if boundaries.n_parameters() != 2:
            raise ValueError('The boundaries must be two-dimensional.')

        xmin, ymin = boundaries.lower()
        xmax, ymax = boundaries.upper()

    # Get voronoi regions (and filter points and evaluations)
    xlim = xmin, xmax
    ylim = ymin, ymax
    x, y, values, regions = _voronoi_regions(x, y, values, xlim, ylim)

    # Create figure and axes
    figure, axes = plt.subplots(figsize=figsize)
    axes.set_xlim(xmin, xmax)
    axes.set_ylim(ymin, ymax)

    # Add coloured voronoi regions
    c = PolyCollection(
        regions, array=values, edgecolors='none', cmap='viridis_r')
    axes.add_collection(c)

    # Add markers
    if markers:
        axes.plot(x, y, markers, color='#ffffff', alpha=0.75)

    # Add colorbar
    figure.colorbar(c, ax=axes)

    return figure, axes


def _voronoi_regions(x, y, f, xlim, ylim):
    """
    Takes a set of ``(x, y, f)`` points and returns the edgepoints of the
    voronoi region around each point within the boundaries specified by
    ``xlim`` and ``ylim``.

    The actual voronoi diagram is constructed by ``scipy.spatial.Voronoi``.
    This method builds on the regions generated by scipy and truncates them to
    a rectangular area. Voronoi regions entirely outside of the bounds are
    removed, and a version of ``x``, ``y``, and ``f`` with these points
    filtered out is also created and returned.

    Parameters
    ----------
    x
        A list of x-coorindates.
    y
        A list of y-coordinates.
    f
        The score function at the given x and y coordinates.
    xlim
        Lower and upper bound for the x coordinates.
    ylim
        Lower and upper bound for the y coordinates.

    Returns
    -------
    A tuple ``(x, y, f, regions)`` where ``x``, ``y`` and ``f`` are the
    coordinates of the accepted points and each ``regions[i]`` is a list of the
    vertices making up the voronoi region for point ``(x[i], y[i])`` with score
    ``f[i]``.
    """
    from scipy.spatial import Voronoi

    # Don't check x, y, f: handled by calling method
    n = len(x)

    # Check limits
    xmin, xmax = [float(a) for a in sorted(xlim)]
    ymin, ymax = [float(a) for a in sorted(ylim)]

    # Drop any points outside the bounds
    # within_bounds = (x >= xmin) & (x <= xmax) & (y >= ymin) & (y <= ymax)
    # x = x[within_bounds]
    # y = y[within_bounds]
    # f = f[within_bounds]

    # Create voronoi diagram
    vor = Voronoi(np.array([x, y]).transpose())

    # The voronoi diagram consists of a number of ridges drawn between a set
    # of points.
    #
    #   points          Are the points the diagram is based on.
    #   vertices        The coordinates of the vertices connecting the ridges
    #   ridge_points    Is a list of tuples (p1, p2) defining the points each
    #                   ridge belongs to. Points are given as their index in
    #                   the list of points.
    #   ridge_vertices  Is a list of vertices (v1, v2) defining the vertices
    #                   between which each ridge is drawn. Vertices are given
    #                   as their index in the list of vertice coordinates. For
    #                   ridges extending to infinity, one of the vertices will
    #                   be given as -1.
    #
    # Get the center of the voronoi diagram's points and define a radius that
    # will bring us outside the visible area for any starting point / direction
    #
    center = vor.points.mean(axis=0)
    radius2 = 2 * np.sqrt((xmax - xmin)**2 + (ymax - ymin)**2)

    # Create a list containing the set of vertices defining each region
    regions = [set() for i in range(n)]
    for (p1, p2), (v1, v2) in zip(vor.ridge_points, vor.ridge_vertices):
        # Order the edges: if one of the edges extends to infinity, the value
        # v1/v2 will be -1. By ordering here we ensure that only v1 can ever be
        # negative, so we don't need to check for v2 < 0 in the code below.
        v1, v2 = (v2, v1) if v1 > v2 else (v1, v2)

        # Get vertice coordinates
        x2 = vor.vertices[v2]  # Only v1 can be -1
        if v1 >= 0:
            # Finite vertex: use as is
            x1 = vor.vertices[v1]
        else:
            # Replacement vertex needed
            # Tangent line to points involved
            y1, y2 = vor.points[p1], vor.points[p2]
            t = y2 - y1
            t /= np.linalg.norm(t)

            # Normal line
            q = np.array([-t[1], t[0]])

            # Midpoint between involved points
            midpoint = np.mean([y1, y2], axis=0)

            # Point beyond the outer boundary
            x1 = x2 + np.sign(np.dot(midpoint - center, q)) * q * radius2

        # Add vertice coordinates to both region coordinate lists
        x1, x2 = tuple(x1), tuple(x2)  # arrays and lists aren't hashable
        regions[p1].update((x1, x2))
        regions[p2].update((x1, x2))

    # Order vertices in regions counter clockwise, truncate the regions at the
    # border, and remove regions outside of the bounds.
    selection = []
    for k, region in enumerate(regions):
        # Regions can be empty if points appear more than once
        if len(region) == 0:
            continue

        # Convert set of tuples to 2d array
        region = np.asarray([np.asarray(v) for v in region])

        # Filter out any regions lying entirely outside the bounds
        xmn = region[:, 0] < xmin
        xmx = region[:, 0] > xmax
        ymn = region[:, 1] < ymin
        ymx = region[:, 1] > ymax
        if np.all(xmn | xmx) and np.all(ymn | ymx):
            continue

        # Sort vertices counter clockwise
        p = vor.points[k]
        angles = np.arctan2(region[:, 1] - p[1], region[:, 0] - p[0])
        regions[k] = region[np.argsort(angles)]

        # Region fully contained? Then keep in selection and continue
        if not (np.any(xmn) or np.any(xmx) or np.any(ymn) or np.any(ymx)):
            selection.append(k)
            continue

        # Region needs truncating

        # Drop points outside of boundary and replace by 0, 1 or 2 new points
        # on the actual boundaries.
        # Run twice: once for x violations, once for y violations (two
        # successive corrections may be needed, to solve corner issues).
        new_region = []
        for j, p in enumerate(region):
            if p[0] < xmin:
                q = region[j - 1] if j > 0 else region[-1]
                r = region[j + 1] if j < len(region) - 1 else region[0]
                if q[0] < xmin and r[0] < xmin:
                    # Point connecting two outsiders: drop
                    continue
                if q[0] >= xmin:
                    # Add point on line p-q
                    s = p[1] + (xmin - p[0]) * (q[1] - p[1]) / (q[0] - p[0])
                    new_region.append(np.array([xmin, s]))
                if r[0] >= xmin:
                    # Add point on line p-r
                    s = p[1] + (xmin - p[0]) * (r[1] - p[1]) / (r[0] - p[0])
                    new_region.append(np.array([xmin, s]))
            elif p[0] > xmax:
                q = region[j - 1] if j > 0 else region[-1]
                r = region[j + 1] if j < len(region) - 1 else region[0]
                if q[0] > xmax and r[0] > xmax:
                    # Point connecting two outsiders: drop
                    continue
                if q[0] <= xmax:
                    # Add point on line p-q
                    s = p[1] + (xmax - p[0]) * (q[1] - p[1]) / (q[0] - p[0])
                    new_region.append(np.array([xmax, s]))
                if r[0] <= xmax:
                    # Add point on line p-r
                    s = p[1] + (xmax - p[0]) * (r[1] - p[1]) / (r[0] - p[0])
                    new_region.append(np.array([xmax, s]))
            else:
                # Point is fine, just add
                new_region.append(p)
        region = new_region

        # Run again for y-violations
        new_region = []
        for j, p in enumerate(region):
            if p[1] < ymin:
                q = region[j - 1] if j > 0 else region[-1]
                r = region[j + 1] if j < len(region) - 1 else region[0]
                if q[1] < ymin and r[1] < ymin:
                    # Point connecting two outsiders: drop
                    continue
                if q[1] >= ymin:
                    # Add point on line p-q
                    s = p[0] + (ymin - p[1]) * (q[0] - p[0]) / (q[1] - p[1])
                    new_region.append(np.array([s, ymin]))
                if r[1] >= ymin:
                    # Add point on line p-r
                    s = p[0] + (ymin - p[1]) * (r[0] - p[0]) / (r[1] - p[1])
                    new_region.append(np.array([s, ymin]))
            elif p[1] > ymax:
                q = region[j - 1] if j > 0 else region[-1]
                r = region[j + 1] if j < len(region) - 1 else region[0]
                if q[1] > ymax and r[1] > ymax:
                    # Point connecting two outsiders: drop
                    continue
                if q[1] <= ymax:
                    # Add point on line p-q
                    s = p[0] + (ymax - p[1]) * (q[0] - p[0]) / (q[1] - p[1])
                    new_region.append(np.array([s, ymax]))
                if r[1] <= ymax:
                    # Add point on line p-r
                    s = p[0] + (ymax - p[1]) * (r[0] - p[0]) / (r[1] - p[1])
                    new_region.append(np.array([s, ymax]))
            else:
                # Point is fine, just add
                new_region.append(p)
        region = new_region

        # Store regions that are still OK
        if len(region) > 2:
            selection.append(k)

    # Filter out bad regions
    regions = np.array(regions, dtype=object)
    regions = regions[selection]
    x = x[selection]
    y = y[selection]
    f = f[selection]

    # Return output
    # Note: Regions is transformed back to a list, which fixes an issue with
    # matplotlib 3.3.0 (which does not expect "ragged" ndarrays made of
    # objects).
    return x, y, f, regions.tolist()

