import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.tri

# Example for TriContour based from:
# https://matplotlib.org/2.0.2/examples/pylab_examples/tricontour_smooth_delaunay.html


def experiment_res(x, y):
    """ An analytic function representing experiment results """
    x = 2.*x
    r1 = np.sqrt((0.5 - x)**2 + (0.5 - y)**2)
    theta1 = np.arctan2(0.5 - x, 0.5 - y)
    r2 = np.sqrt((-x - 0.2)**2 + (-y - 0.2)**2)
    theta2 = np.arctan2(-x - 0.2, -y - 0.2)
    z = (4*(np.exp((r1/10)**2) - 1)*30. * np.cos(3*theta1) +
         (np.exp((r2/10)**2) - 1)*30. * np.cos(5*theta2) +
         2*(x**2 + y**2))
    return (np.max(z) - z)/(np.max(z) - np.min(z))

n = 200
random_gen = np.random.mtrand.RandomState(seed=127260)
x = random_gen.uniform(-1., 1., size=n)
y = random_gen.uniform(-1., 1., size=n)
v = experiment_res(x, y)

## triangulate
tri = mpl.tri.Triangulation(x, y)
mask = mpl.tri.TriAnalyzer(tri).get_flat_tri_mask(min_circle_ratio=.01)
tri.set_mask(mask)

## refine triangular mesh
refiner = mpl.tri.UniformTriRefiner(tri)
triFiner, vFiner = refiner.refine_field(v, subdiv=2)  # uses CubicTriInterpolator

## plot the mesh
plt.close('all')
plt.triplot(tri, color='k', lw=.8)
plt.triplot(triFiner, color='k', alpha=.3, lw=.5)
plt.show()

## plot and create TriCounterSet object
levels = [.7, .9, 1]
cmap = mpl.cm.plasma
plt.close('all')
cl = plt.tricontour(triFiner, vFiner, levels=levels, colors='k', linewidths=2, linestyles='--')
cf = plt.tricontourf(triFiner, vFiner, levels=levels, cmap=cmap, extend='both', alpha=.8)
_ = [collection.set_edgecolor("face") for collection in cf.collections]  # https://stackoverflow.com/a/73805556/921580
plt.triplot(tri, color='k', lw=1.5, alpha=.1)  # refined mesh
plt.triplot(triFiner, color='k', lw=.5, alpha=.1)  # original mesh
cb = plt.colorbar(cf)
plt.show()

## retrieve contours from TriContourSet (paths) and convert to polygons
for i in range(len(cf.collections)):
    plt.gca().set_xlim(-1.1, 1.1)
    plt.gca().set_ylim(-1.1, 1.1)
    [plt.gca().add_patch(
            mpl.patches.PathPatch(
                cf.collections[i].get_paths()[j],
                facecolor=cf.collections[i].get_facecolor()
            )
        ) for j in range(len(cf.collections[i].get_paths()))]
    [plt.plot(*list(zip(*p)), c='red', ls='--') 
            for path in cf.collections[i].get_paths()
            for p in path.to_polygons()
        ]
    if i == (len(cf.collections) - 2):
        plt.savefig('bug_actual_outcome.png')
    plt.show()

## compare path vertices with output from to_polygons method
path = cf.collections[-2].get_paths()[0]
print("\npath vertices and codes:\n", path)
print("\npath.to_polygons() = \n", path.to_polygons())

def path2polygons(path):
    assert isinstance(path, mpl.path.Path), f"Type of input is {type(path)} instead of Path."
    codesOld = path.codes[path.codes != mpl.path.Path.CLOSEPOLY]
    vertsOld = path.vertices[path.codes != mpl.path.Path.CLOSEPOLY]
    jj = np.append(np.flatnonzero(codesOld == path.MOVETO), [len(codesOld)])
    verts, codes = [], []
    for ji, je in zip(jj[:-1], jj[1:]):
        codes.extend(   [mpl.path.Path.MOVETO]
                      + [mpl.path.Path.LINETO]*(je - ji - 1)
                      + [mpl.path.Path.CLOSEPOLY]
                    )
        verts.extend( list(vertsOld[ji:je]) + [vertsOld[ji]] )
    return mpl.path.Path(verts, codes).to_polygons()