import numpy as np
import matplotlib.pyplot as plt

# From https://matplotlib.org/stable/gallery/event_handling/resample.html
# A class that will downsample the data and recompute when zoomed.
class DataDisplayDownsampler:
    def __init__(self, xdata, ydata):
        self.origYData = ydata
        self.origXData = xdata
        self.max_points = 50
        self.delta = xdata[-1] - xdata[0]

    def downsample(self, xstart, xend):
        # get the points in the view range
        mask = (self.origXData > xstart) & (self.origXData < xend)
        # dilate the mask by one to catch the points just outside
        # of the view range to not truncate the line
        mask = np.convolve([1, 1, 1], mask, mode='same').astype(bool)
        # sort out how many points to drop
        ratio = max(np.sum(mask) // self.max_points, 1)

        # mask data
        xdata = self.origXData[mask]
        ydata = self.origYData[mask]

        # downsample data
        xdata = xdata[::ratio]
        ydata = ydata[::ratio]

        print("using {} of {} visible points".format(len(ydata), np.sum(mask)))

        return xdata, ydata

    def update(self, ax):
        # Update the line
        lims = ax.viewLim
        if abs(lims.width - self.delta) > 1e-8:
            self.delta = lims.width
            xstart, xend = lims.intervalx
            self.line.set_data(*self.downsample(xstart, xend))
            ax.figure.canvas.draw_idle()


# Create a signal
xdata = np.linspace(16, 365, (365-16)*4)
ydata = np.sin(2*np.pi*xdata/153) + np.cos(2*np.pi*xdata/127)


# --- This does not work: ax1 drawn after ax0 kills callbacks
d = DataDisplayDownsampler(xdata, ydata)
fig, (ax0, ax1) = plt.subplots(nrows=2, sharex=True)

# Hook up the line
d.line, = ax0.plot(xdata, ydata, 'o-')
ax0.set_autoscale_on(False)  # Otherwise, infinite loop

# Connect for changing the view limits
ax0.callbacks.connect('xlim_changed', d.update)
ax0.set_xlim(16, 365)

ax1.plot(xdata, -ydata)
plt.show()


# --- This does work: ax0 drawn after ax1
# --- Note: only works if axis limits are controlled via ax0, not ax1
# Create a signal
xdata = np.linspace(16, 365, (365-16)*4)
ydata = np.sin(2*np.pi*xdata/153) + np.cos(2*np.pi*xdata/127)

d = DataDisplayDownsampler(xdata, ydata)

fig, (ax0, ax1) = plt.subplots(nrows=2, sharex=True)

ax1.plot(xdata, -ydata)

# Hook up the line
d.line, = ax0.plot(xdata, ydata, 'o-')
ax0.set_autoscale_on(False)  # Otherwise, infinite loop

# Connect for changing the view limits
ax0.callbacks.connect('xlim_changed', d.update)
ax0.set_xlim(16, 365)


plt.show()

from typing import Any

import matplotlib.pyplot as plt
from matplotlib.axes import Axes


def add_identity(ax: Axes = None, **line_kwargs: Any) -> None:
    """Add a parity line (y = x) to the provided axis."""
    if ax is None:
        ax = plt.gca()

    # zorder=0 ensures other plotted data displays on top of line
    default_kwargs = dict(alpha=0.5, zorder=0, linestyle="dashed", color="black")
    (identity,) = ax.plot([], [], **default_kwargs, **line_kwargs)

    def callback(axes: Axes) -> None:
        x_min, x_max = axes.get_xlim()
        y_min, y_max = axes.get_ylim()
        low = max(x_min, y_min)
        high = min(x_max, y_max)
        identity.set_data([low, high], [low, high])

    callback(ax)
    # Register callbacks to update identity line when moving plots in interactive
    # mode to ensure line always extend to plot edges.
    ax.callbacks.connect("xlim_changed", callback)
    ax.callbacks.connect("ylim_changed", callback)


fig, (ax1, ax2) = plt.subplots(1, 2, sharex=True, sharey=True)

ax1.plot([0, 1], [1, 0])
add_identity(ax1)

ax2.plot([0, 2], [2, 0])
add_identity(ax2)

plt.savefig('tmp.png')