###
# Copyright (c) 2016, 2017 Diamond Light Source Ltd.
#
# Contributors:
#    Tom Cobb - initial API and implementation and/or initial documentation
#    Gary Yendell - initial API and implementation and/or initial documentation
#    Charles Mita - initial API and implementation and/or initial documentation
#
###

from scanpointgenerator import CompoundGenerator, RectangularROI, CircularROI

MARKER_SIZE = 10


def plot_generator(gen, excluder=None, show_indexes=True):
    from matplotlib.patches import Rectangle, Circle
    import matplotlib.pyplot as plt
    import numpy as np
    from scipy import interpolate

    if excluder is not None:
        for roi in excluder.rois:
            overlay = plt.subplot(111, aspect='equal')
            if isinstance(roi, RectangularROI):
                overlay.add_patch(Rectangle(roi.start, roi.width, roi.height, fill=False))
            if isinstance(roi, CircularROI):
                overlay.add_patch(Circle(roi.centre, roi.radius, fill=False))

    if not isinstance(gen, CompoundGenerator):
        excluders = [] if excluder is None else [excluder]
        gen = CompoundGenerator([gen], excluders, [])
    gen.prepare()

    # points for spline generation
    x, y = [], []
    # capture points and indexes
    capx, capy, capi = [], [], []
    # segment start for colour changing
    starts = []
    for point in gen.iterator():
        # If lower is different from last then include it
        xlower = point.lower["x"]
        ylower = point.lower.get("y", 0)
        xpos = point.positions["x"]
        ypos = point.positions.get("y", 0)
        if len(x) == 0 or x[-1] != xlower or y[-1] != ylower:
            if len(x) != 0:
                # add in a tiny fractional distance to extend last point
                xdiff = (x[-1] - x[-2]) * 0.01
                ydiff = (y[-1] - y[-2]) * 0.01
                for i in range(3):
                    x.append(x[-1] + xdiff)
                    y.append(y[-1] + ydiff)
                # add in padding for the next point
                xdiff = (xpos - xlower) * 0.01
                ydiff = (ypos - ylower) * 0.01
                for i in reversed(range(3)):
                    x.append(xlower - xdiff * (i + 1))
                    y.append(ylower - ydiff * (i + 1))
            starts.append(len(x))
            x.append(xlower)
            y.append(ylower)
        # Add in capture points
        x.append(xpos)
        y.append(ypos)
        capx.append(xpos)
        capy.append(ypos)
        capi.append(point.indexes)
        # And upper point
        starts.append(len(x))
        x.append(point.upper["x"])
        y.append(point.upper.get("y", 0))

    # # Plot labels
    plt.xlabel("X (%s)" % gen.units["x"])
    if "y" in gen.units:
        plt.ylabel("Y (%s)" % gen.units["y"])
    else:
        plt.tick_params(left='off', labelleft='off')

    # Define curves parametrically
    x = np.array(x)
    y = np.array(y)
    t = np.zeros(len(x))
    t[1:] = np.sqrt((x[1:] - x[:-1])**2 + (y[1:] - y[:-1])**2)
    t = np.cumsum(t)
    t /= t[-1]
    tck, _ = interpolate.splprep([x, y], s=0)

    # Plot each line
    for i, start in enumerate(starts):
        if i + 1 < len(starts):
            end = starts[i+1]
        else:
            end = len(x) - 1
        tnew = np.linspace(t[start], t[end], num=1001, endpoint=True)
        sx, sy = interpolate.splev(tnew, tck)
        plt.plot(sx, sy, linewidth=2)

    # And the capture points
    plt.plot(capx, capy, linestyle="", marker="x", color="k",
             markersize=MARKER_SIZE)

    # And a start position
    plt.plot([x[0]], [y[0]], 'bo')
    plt.annotate("Start", (x[0], y[0]), xytext=(MARKER_SIZE/2, MARKER_SIZE/2),
                 textcoords='offset points')

    # And the indexes
    if show_indexes:
        for i, x, y in zip(capi, capx, capy):
            plt.annotate(i, (x, y), xytext=(MARKER_SIZE/2, MARKER_SIZE/2),
                         textcoords='offset points')
        #indexes = ["%s (size %d)" % z for z in zip(gen.index_names, gen.index_dims)]
        #plt.title("Dataset: [%s]" % (", ".join(indexes)))
    plt.show()
