# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_animate.ipynb.

# %% auto 0
__all__ = ['default_fig_setup', 'animate_lines']

# %% ../nbs/00_animate.ipynb 2
import numpy as np
import matplotlib.pyplot as plt
plt.style.use('fivethirtyeight')

from matplotlib import animation
from IPython.display import HTML

# %% ../nbs/00_animate.ipynb 4
def default_fig_setup():
    fig, ax = plt.subplots()
    return fig, ax
    

# %% ../nbs/00_animate.ipynb 5
def animate_lines(X, 
                     Y, 
                     xlabel = 'x-position', 
                     ylabel = 'y-position',
                     setup_fig_function = default_fig_setup,
                     xlims = [], ylims = [], 
                    labels = [],
                     linewidth = 5,
                     legend = False,
                     interval = 200,
                     **plotargs
                    ):
    '''
    `animate_lines(X, Y)` is a function that takes x2 3D arrays, `X` and `Y` 
    and returns an animation of the changing lines
    
    The `X` and `Y` arrays need to be arranged in 3D arrays such that, 
    
    `X[line_N, point_i, time_j], Y[line_N, point_i, time_j]`
    
    where each line is arranged in columns and the next frame is aranged in the 
    third dimension of the array
    
    If you are plotting a single line, you can use the columns as the timestep e.g. 
    
    `X[point_i, time_j], Y[point_i, time_j]`
    
     the function will add an extra axis to the beginning of the arrays as such
     
     `X[:, np.new_axis :], Y[:, np.new_axis, :]`
    
    Parameters:
    ----------
    `X`: The x-axis data for the animated lines, its shape is such that each column 
         is a set of x-values for a given line and each frame is organized along the 
         third dimension
    `Y`: The y-axis data for the animated lines, its shape is such that each column 
         is a set of y-values for a given line and each frame is organized along the 
         third dimension
    xlabel:  plot x-axis label, default is 'x-position', 
    ylabel: plot y-axis label, default is 'y-position', 
    setup_fig_function: a function that returns fig and ax that can be used to plot static lines 
                        before animating, default is an empty plot
    xlims: Manually set the x-axis limits. If its not specified, the animation uses 1.1*(max&min)
    ylims: [], 
    labels: [],
    linewidth: 5
    legend: False,
    `interval` : int, default: 200
                 Delay between frames in milliseconds.
    **plotargs: used as **kwargs for matplotlib's plot command
    
    
    Returns:
    --------
    `anim`: an `animation.FuncAnimation` result that can be viewed with Jupyter's HTML
          e.g. `HTML(anim.to_html5_video())` or `HTML(anim.to_jshtml())`
    '''
    if len(X.shape) < 3:
        X = X[:, np.newaxis, :]
    if len(Y.shape) < 3:
        Y = Y[:, np.newaxis,:]
    
    fig, ax = setup_fig_function()

    if len(xlims)<2: 
        ax.set_xlim((np.min(X)*1.1, np.max(X)*1.1));
    else: ax.set_xlim(xlims)
    if len(ylims)<2:
        ax.set_ylim((np.min(Y)*1.1, np.max(Y)*1.1));
    else: ax.set_ylim(ylims)
    
    ax.set_xlabel(xlabel);
    ax.set_ylabel(ylabel);

    i = 0
    lines = ax.plot(X[:, :, i], Y[:, :, i], lw=linewidth, **plotargs);
    if legend:
        plt.legend(lines, labels)

    def init():
        for j in range(len(lines)):
            lines[j].set_data([], [])
        return lines

    def animate(i):
        for j in range(len(lines)):
            lines[j].set_data(X[:, j, i], Y[:, j, i])
        return lines

    anim = animation.FuncAnimation(fig, animate, init_func=init,
                                   frames=range(0,Y.shape[-1]), 
                                   blit=True, interval = interval)
    plt.close(fig)
    HTML(anim.to_html5_video())
    return anim
