import numpy as np
from scipy.integrate import odeint, quad
from scipy.optimize import brentq
import matplotlib.pyplot as plt
from matplotlib import animation, rc
import seaborn as sbs
#rc('font', **{'family': 'serif', 'serif': ['Computer Modern'], 'size': 20})
#rc('text', usetex=True)
#rc('animation', html='jshtml')
#plt.rcParams["animation.html"] = "jshtml"


def tsanim(tss):
    "Plot an array of time series data t, x = [ti], [xi]."
    nrows = len(tss)
    fig, axes = plt.subplots(nrows=nrows, ncols=1)
    tsmax = lambda ts: max(len(ts[0]), len(ts[1]))
    ts_maxlen = max(map(tsmax, tss))
    tinit = int(0.1*ts_maxlen)
    print("Animate time series from t=%d with %d frames" % (tinit, ts_maxlen))

    def tsplot(ax, ts):
        t, x = ts
        ax.set_ylabel(r'$x$')
        ax.set_ylim(np.min(x), np.max(x))
        line, = ax.plot(t[:tinit], x[:tinit])
        return line
    
    lines = list(map(lambda axts: tsplot(*axts), zip(axes, tss)))
    axes[-1].set_xlabel(r'$t$')

    def animate(i):
        """Update the image for iteration i of the Matplotlib animation."""
        for ax, line, ts in zip(axes, lines, tss):
            t, x = ts
            line.set_data(t[:i+1], x[:i+1])
            ax.set_xlim(0, t[i])
        return

    plt.tight_layout()
    anim = animation.FuncAnimation(fig, animate, frames=ts_maxlen, interval=10)
    return anim