from collections import deque

try:
    import matplotlib
    matplotlib.use('Agg')  # non-interactive safe default
    import matplotlib.pyplot as plt
    MATPLOTLIB_AVAILABLE = True
except ImportError:
    MATPLOTLIB_AVAILABLE = False


class Plotter:
    """Tracks gradient norms per layer and plots them every N steps."""

    def __init__(self, window=100, log_every=10):
        self.window = window
        self.log_every = log_every
        self._history = {}  # layer_name -> deque of norm values
        self._fig = None
        self._ax = None

    def update(self, stats, step):
        """Record norms and optionally refresh plot."""
        for name, s in stats.items():
            if name not in self._history:
                self._history[name] = deque(maxlen=self.window)
            self._history[name].append(s.get('norm', 0.0))

        if step % self.log_every == 0 and MATPLOTLIB_AVAILABLE:
            self._draw(step)

    def _draw(self, step):
        if not self._history:
            return
        if self._fig is None:
            self._fig, self._ax = plt.subplots(figsize=(10, 4))
            plt.ion()

        self._ax.clear()
        for name, vals in list(self._history.items())[:10]:  # max 10 layers
            self._ax.plot(list(vals), label=name, linewidth=1)

        self._ax.set_title(f'Gradient norms — step {step}')
        self._ax.set_xlabel('Steps (last 100)')
        self._ax.set_ylabel('Gradient norm')
        self._ax.legend(fontsize=7, loc='upper right')
        self._ax.set_yscale('log')
        try:
            self._fig.tight_layout()
            plt.pause(0.001)
        except Exception:
            pass

    def save(self, path='gradient_norms.png'):
        if self._fig is not None:
            self._fig.savefig(path, dpi=100, bbox_inches='tight')
            print(f'[torch-surgeon] Plot saved to {path}')