import numpy as np
import matplotlib.pyplot as plt
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib.widgets import LassoSelector
from matplotlib.path import Path
import matplotlib

print(matplotlib.get_backend())


class SelectFromSharedCollection:
    mask = None

    def __init__(self, ax, collection, alpha_other=0.3):
        self.ax, self.canvas = ax, ax.figure.canvas
        self.collection = collection
        self.alpha_other = alpha_other

        self.xys = collection.get_offsets()
        self.Npts = len(self.xys)

        # Ensure that we have separate colors for each object
        self.fc = collection.get_facecolors()
        if len(self.fc) == 0:
            raise ValueError("Collection must have a facecolor")
        elif len(self.fc) == 1:
            self.fc = np.tile(self.fc, (self.Npts, 1))

        self.lasso = LassoSelector(ax, onselect=self.onselect)
        self.ind = []

    def onselect(self, verts):
        path = Path(verts)
        SelectFromSharedCollection.last_path = path
        SelectFromSharedCollection.mask = path.contains_points(self.xys)
        self.ind = np.nonzero(SelectFromSharedCollection.mask)[0]
        self.fc[:, -1] = self.alpha_other
        self.fc[self.ind, -1] = 1
        self.collection.set_facecolors(self.fc)
        self.canvas.draw_idle()

    def disconnect(self):
        self.lasso.disconnect_events()
        self.fc[:, -1] = 1
        self.collection.set_facecolors(self.fc)
        self.canvas.draw_idle()

    def accept(self, event):
        if event.key == "enter":
            print(SelectFromSharedCollection.mask)
            self.fc[~SelectFromSharedCollection.mask] = [0.0, 0.5, 0, 0.1]
            self.fc[SelectFromSharedCollection.mask] = [0.5, 0, 0, 1]
            self.canvas.draw_idle()


fig, axes = plt.subplots(3, 3)
axes = axes.flatten()

figures = {}


def accept(event):
    match event.key:
        case "enter":
            for perplexity, figure in figures.items():
                print(f"Perplexity: {perplexity}")
                figure["selector"].accept(event)
            plt.pause(0.5)
            fig.canvas.draw()
            fig.canvas.flush_events()
        case _:
            pass


for ax, perplexity in zip(axes, range(20, 200, 20)):
    X_embedded = np.random.rand(50, 2)
    pts = ax.scatter(X_embedded[:, 0], X_embedded[:, 1], s=1)
    figures[perplexity] = {
        "ax": ax,
        "selector": SelectFromSharedCollection(ax, pts),
    }

fig.canvas.mpl_connect("key_press_event", accept)
plt.show()