import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation

class Plot:
	def __init__(self, count):
		self.fig, self.plots = plt.subplots(nrows=count, ncols=1)
		self.fig.tight_layout()
		self.fig.canvas.set_window_title('FFT Cube')
		
		self.count = count
		self.entities = []
		
	def addLabel(self, ax, title="", xlabel="", ylabel=""):
		self.plots[ax].set_title(title)
		self.plots[ax].set_xlabel(xlabel)
		self.plots[ax].set_ylabel(ylabel)
		
	def addLine(self, ax, x, yMin, yMax, color = 'r-', xscale = 'linear'):
		self.plots[ax].set_xlim(min(x), max(x))
		self.plots[ax].set_ylim(yMin, yMax)
		self.plots[ax].set_xscale(xscale)
		line, = self.plots[ax].plot(x, [0]*len(x), color)
		self.entities.append(line)
		return line
	
	def addBars(self, ax, labels, yMin, yMax, width, offset, color = 'r'):
		self.plots[ax].set_ylim(yMin, yMax)
		self.plots[ax].set_xticklabels(labels)
		x = range(len(labels))
		y = [0] * (len(labels))
		bars = self.plots[ax].bar([i+offset for i in x], y, width, color=color)
		for bar in bars:
			self.entities.append(bar)
		return bars

	def update(self, frame = None):
		self._cb()
		return self.entities
		
	def show(self, cb, framerate = 60):
		self._cb = cb
		self.animation = animation.FuncAnimation(self.fig, self.update, interval=1000/framerate, blit=True)
		plt.show()