import matplotlib.pyplot as plt
import numpy as np


class Monitor:
    """
    Monitor v and i data in a matplotlib window
    """
    def __init__(self, xlabel: str, datas: list[dict], max_points_shown=None, use_print=False):
        self.max_points_shown = max_points_shown
        self.use_print = use_print
        self.index = []
        self.xdata = []
        self.options = [{"ax": 0, "ylabel": None, "lim": None, "color": "orange", "grid": True, "label": None} | datas[i] for i in range(len(datas))]
        self.ydatas = [[] for _ in range(len(datas))]
        self.lines = []
        n_ax = max(self.options, key=lambda v: v["ax"])["ax"] + 1
        plt.ion()
        self.fig1, self.ax = plt.subplots(n_ax, 1, figsize=(7, 8), sharex=True)
       
        
        for i in range(len(self.ydatas)):
            opt = self.options[i]
            ax = self.ax[opt["ax"]]
            self.lines.append(ax.plot(self.xdata, self.ydatas[i], color=opt["color"], label=opt["label"])[0])
            if opt["label"]:
                ax.legend()
            if opt["grid"]:
                ax.grid(True)
            if opt["lim"]:
                ax.set_ylim(opt["lim"][0], opt["lim"][1])
            if opt["ylabel"]:
                ax.set_ylabel(opt["ylabel"])

    def override(self, xval, *yvals):
        idx = self.xdata.index(xval)
        for i, y in enumerate(yvals):
            if y is None: continue
            self.ydatas[i][idx] = y
            self.lines[i].set_ydata(self.ydatas[i])

        self._post_update()

    def update(self, xval, *yvals):
        n = len(self.xdata)
        if self.use_print:
            s = f"n = {n}"
            for i, y in enumerate(yvals):
                s += f", {self.options[i]['label']} = {y}"
                print(s + " "*10, end='\r')
        self.xdata.append(xval)
        for i, y in enumerate(yvals):
            if y is None: y = np.nan
            self.ydatas[i].append(y)
            self.lines[i].set_xdata(self.xdata)
            self.lines[i].set_ydata(self.ydatas[i])
        self._post_update()

    def update_array(self, xdata, *ydatas):
        self.xdata.extend(xdata)
        for i, y in enumerate(ydatas):
            self.ydatas[i].extend(y)
            self.lines[i].set_xdata(self.xdata)
            self.lines[i].set_ydata(self.ydatas[i])
        self._post_update()

    def _post_update(self):
        n = len(self.xdata)
        if self.max_points_shown and n > self.max_points_shown:
            # recalculate limits and set them for the view
            for ax in self.ax:
                self.ax[-1].set_xlim(self.xdata[n - self.max_points_shown], self.xdata[n])
        for ax in self.ax:
            ax.relim()  # required for autoscale_view to work
            ax.autoscale_view()

        self.ax[-1].autoscale_view()
        # update plot
        self.fig1.canvas.draw()
        self.fig1.canvas.flush_events()

    def __del__(self):
        plt.close(self.fig1)