#!/usr/bin/env python
# coding: utf-8

# ## JupyterLab-compatible for SO https://stackoverflow.com/q/75017358/8508004
# 
# This builds on the Jupyter notebook with related code that can be viewed [here](https://nbviewer.org/gist/fomightez/f539a3770f2f3287bf12de2d2a549e3a).  
# [Click here](https://mybinder.org/v2/gh/fomightez/animated_matplotlib-binder/master?urlpath=git-pull%3Frepo%3Dhttps%253A%252F%252Fgist.github.com%252Ffomightez%252Ff539a3770f2f3287bf12de2d2a549e3a%26urlpath%3Dtree%252Ff539a3770f2f3287bf12de2d2a549e3a%252Fsegments_for_training_updated_animation.ipynb) to launch that notebook in an active, temporary Jupyter session served via MyBinder.org with the environment already having Python and ALMOST everything needed to make the animation with a widget work in the notebook installed and working. Here, I'm going to include the `%pip install torchflow` line as a cell so that the notebook can just be run. (If you use [that link above]((https://mybinder.org/v2/gh/fomightez/animated_matplotlib-binder/master?urlpath=git-pull%3Frepo%3Dhttps%253A%252F%252Fgist.github.com%252Ffomightez%252Ff539a3770f2f3287bf12de2d2a549e3a%26urlpath%3Dtree%252Ff539a3770f2f3287bf12de2d2a549e3a%252Fsegments_for_training_updated_animation.ipynb), you'll need to run the same process first.)
# 
# Go ahead and run that below to get the final preparation of the environment out of the way.

# In[1]:


get_ipython().run_line_magic('pip', 'install torch')


# Restart the kernel after running that cell and then try running the one below. **With that preparation complete**, continue on to examine and run this method...

# Right now this method (based on based on [here](https://stackoverflow.com/a/52672859/8508004)) is more universal than the one illustrated in the first notebook. This method works in both the traditional Jupyter notebook and JupyterLab as of now (January 2023).
# 
# In the previous notebook with the other method ([here](https://nbviewer.org/gist/fomightez/f539a3770f2f3287bf12de2d2a549e3a)), the data was collected first and then the collected data used to plot using `fig.canvas.draw()` to reset the image each round. A pause is built in to make sure the animation steps through adding showing each next segment with time in between instead of just quickly blinking through doing that too fast to really see and just displaying the last step.  
# 
# Here the data is plotted in the course of each iteration, which corresponds to the original loop in the OP's code. The time it takes to make each round of data actually build in more pause than the other method where the data had already been collected. I've added in additional cushion here though to make it even smoother. This additional cushio can be removed or adjusted as needed here. It would be important to have this pause to control running through displaying the data if the data was already existing or not computational time intensive to produce.
# 
# In addition to the extra delay, the `live_plot()` funnction that handles updating the plot with each round uses IPython's display control to clear the output with each frame that will be shown in the animation. You can see a somewhat simpler implementation in [the StackOverflow answer that this approach is based on](https://stackoverflow.com/a/52672859/8508004
# ).
# 
# Run this code below to see the animation. Note that it doesn't need setting any `%matplotlib` magic explicitly. This is as opposed to the use of `%matplotlib notebook` being necessaryfor related code that can be viewed [here](https://nbviewer.org/gist/fomightez/f539a3770f2f3287bf12de2d2a549e3a) to work in the tradtional Jupyter notbeook interface.

# In[2]:


# based on https://stackoverflow.com/a/52672859/8508004 to help with https://stackoverflow.com/q/75017358/8508004
from IPython.display import clear_output
from matplotlib import pyplot as plt
import numpy as np
import collections
import time

def live_plot(data_dict, figsize=(12,5), title=''):
    clear_output(wait=True)
    plt.figure(figsize=figsize)
    #plt.plot(data_dict["steps"],data_dict["r"] , 'r-', label = "real")
    #plt.plot(data_dict["steps"],data_dict["b"] , 'b-', label = "prediction")
    for i,_ in enumerate(data_dict["steps"]):
        plt.plot(data_dict["steps"][i], list(data_dict["r"][i]) , 'r-', )
        plt.plot(data_dict["steps"][i], list(data_dict["b"][i]) , 'b-', )
    plt.title(title)
    plt.grid(True)
    #plt.legend(loc='center left') # the plot evolves to the right
    plt.show()
    time.sleep(0.2) # extend delay between adding next frame in animation

data = collections.defaultdict(list)
    
import torch
from torch import nn

# torch.manual_seed(1)    # reproducible

# Hyper Parameters
TIME_STEP = 10      # rnn time step
INPUT_SIZE = 1      # rnn input size
LR = 0.02           # learning rate

# data
steps = np.linspace(0, np.pi*2, 100, dtype=np.float32)  # float32 for converting torch FloatTensor
x_np = np.sin(steps)
y_np = np.cos(steps)

class RNN(nn.Module):
    def __init__(self):
        super(RNN, self).__init__()

        self.rnn = nn.RNN(
            input_size=INPUT_SIZE,
            hidden_size=32,     # rnn hidden unit
            num_layers=1,       # number of rnn layer
            batch_first=True,   # input & output will has batch size as 1s dimension. e.g. (batch, time_step, input_size)
        )
        self.out = nn.Linear(32, 1)

    def forward(self, x, h_state):
        # x (batch, time_step, input_size)
        # h_state (n_layers, batch, hidden_size)
        # r_out (batch, time_step, hidden_size)
        r_out, h_state = self.rnn(x, h_state)

        outs = []    # save all predictions
        for time_step in range(r_out.size(1)):    # calculate output for each time step
            outs.append(self.out(r_out[:, time_step, :]))
        return torch.stack(outs, dim=1), h_state

        # instead, for simplicity, you can replace above codes by follows
        # r_out = r_out.view(-1, 32)
        # outs = self.out(r_out)
        # outs = outs.view(-1, TIME_STEP, 1)
        # return outs, h_state
        
        # or even simpler, since nn.Linear can accept inputs of any dimension 
        # and returns outputs with same dimension except for the last
        # outs = self.out(r_out)
        # return outs

rnn = RNN()
print(rnn)

optimizer = torch.optim.Adam(rnn.parameters(), lr=LR)   # optimize all cnn parameters
loss_func = nn.MSELoss()

h_state = None      # for initial hidden state

for step in range(100):
    start, end = step * np.pi, (step+1)*np.pi   # time range
    # use sin predicts cos
    steps = np.linspace(start, end, TIME_STEP, dtype=np.float32, endpoint=False)  # float32 for converting torch FloatTensor
    x_np = np.sin(steps)
    y_np = np.cos(steps)

    x = torch.from_numpy(x_np[np.newaxis, :, np.newaxis])    # shape (batch, time_step, input_size)
    y = torch.from_numpy(y_np[np.newaxis, :, np.newaxis])

    prediction, h_state = rnn(x, h_state)   # rnn output
    # !! next step is important !!
    h_state = h_state.data        # repack the hidden state, break the connection from last iteration

    loss = loss_func(prediction, y)         # calculate loss
    optimizer.zero_grad()                   # clear gradients for this training step
    loss.backward()                         # backpropagation, compute gradients
    optimizer.step()                        # apply gradients

    # plotting
    data['steps'].append(list(steps))
    data['r'].append(y_np.flatten())
    data['b'].append(prediction.data.numpy().flatten())
    live_plot(data);


# #### Will the method with `FuncAnimation()` with associated widget controller also work well in JupyterLab?

# Now that I pulled [the original code](https://stackoverflow.com/q/75017358/8508004) apart enough to realize each is a segment and implemented a couple approaches, I also wondered if could use the method with `FuncAnimation()` with associated widget controller that is illustrated at the bottom of [here](https://nbviewer.org/github/fomightez/animated_matplotlib-binder/blob/master/index.ipynb) (and that I had used to answer [here](https://stackoverflow.com/a/75009196/8508004), recently) to make something that also would work in JupyterLab. (**Note it is not all animations methods involving `FuncAnimation()` that work in JupyterLab**, see [here](https://stackoverflow.com/a/73451172/8508004) for one that only works in classic notebook at this time. The widget is key.)

# In[3]:


# If upon first running, it shows a non-interactive, single static shot of the plot below the interactive one with the widget controller,
# JUST RE-RUN TWICE. Re-run usually fixes that display quirk.
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
from IPython.display import HTML, display
plt.rcParams["animation.html"] = "jshtml"
#plt.ioff() #needed so the second time you run it you get only single plot
import collections
import time


fig = plt.figure(figsize=(12,5))
ax = plt.axes(xlim=(0, 320), ylim=(-1.75, 1.75))
lineplot, = ax.plot([], [], "r-")
lineplot2, = ax.plot([], [], "b-")


import torch
from torch import nn

# torch.manual_seed(1)    # reproducible

# Hyper Parameters
TIME_STEP = 10      # rnn time step
INPUT_SIZE = 1      # rnn input size
LR = 0.02           # learning rate

# data
steps = np.linspace(0, np.pi*2, 100, dtype=np.float32)  # float32 for converting torch FloatTensor
x_np = np.sin(steps)
y_np = np.cos(steps)

class RNN(nn.Module):
    def __init__(self):
        super(RNN, self).__init__()

        self.rnn = nn.RNN(
            input_size=INPUT_SIZE,
            hidden_size=32,     # rnn hidden unit
            num_layers=1,       # number of rnn layer
            batch_first=True,   # input & output will has batch size as 1s dimension. e.g. (batch, time_step, input_size)
        )
        self.out = nn.Linear(32, 1)

    def forward(self, x, h_state):
        # x (batch, time_step, input_size)
        # h_state (n_layers, batch, hidden_size)
        # r_out (batch, time_step, hidden_size)
        r_out, h_state = self.rnn(x, h_state)

        outs = []    # save all predictions
        for time_step in range(r_out.size(1)):    # calculate output for each time step
            outs.append(self.out(r_out[:, time_step, :]))
        return torch.stack(outs, dim=1), h_state

        # instead, for simplicity, you can replace above codes by follows
        # r_out = r_out.view(-1, 32)
        # outs = self.out(r_out)
        # outs = outs.view(-1, TIME_STEP, 1)
        # return outs, h_state
        
        # or even simpler, since nn.Linear can accept inputs of any dimension 
        # and returns outputs with same dimension except for the last
        # outs = self.out(r_out)
        # return outs

rnn = RNN()
print(rnn)

optimizer = torch.optim.Adam(rnn.parameters(), lr=LR)   # optimize all cnn parameters
loss_func = nn.MSELoss()

h_state = None      # for initial hidden state

data = collections.defaultdict(list)
    
def init():
    global h_state, data 
    lineplot.set_data([], [])
    lineplot2.set_data([], [])
    data = collections.defaultdict(list)
    return lineplot, #return [lineplot] also works like in https://nbviewer.org/github/raphaelquast/jupyter_notebook_intro/blob/master/jupyter_nb_introduction.ipynb#pre-render-animations-and-export-to-HTML

def animate(i):
    global h_state, data 
    step = i
    start, end = step * np.pi, (step+1)*np.pi   # time range
    # use sin predicts cos
    steps = np.linspace(start, end, TIME_STEP, dtype=np.float32, endpoint=False)  # float32 for converting torch FloatTensor
    x_np = np.sin(steps)
    y_np = np.cos(steps)

    x = torch.from_numpy(x_np[np.newaxis, :, np.newaxis])    # shape (batch, time_step, input_size)
    y = torch.from_numpy(y_np[np.newaxis, :, np.newaxis])

    prediction, h_state = rnn(x, h_state)   # rnn output
    # !! next step is important !!
    h_state = h_state.data        # repack the hidden state, break the connection from last iteration

    loss = loss_func(prediction, y)         # calculate loss
    optimizer.zero_grad()                   # clear gradients for this training step
    loss.backward()                         # backpropagation, compute gradients
    optimizer.step()                        # apply gradients

    # plotting
    data['steps'].append(list(steps))
    data['r'].append(y_np.flatten())
    data['b'].append(prediction.data.numpy().flatten())
    #lineplot.set_data([x], [y])
    #lineplot2.set_data([x], [z])
    lineplot.set_data(data["steps"],data["r"])
    lineplot2.set_data(data["steps"],data["b"])
    '''
    for i,_ in enumerate(data_dict["steps"]):
        plt.plot(data_dict["steps"][i], list(data_dict["r"][i]) , 'r-', )
        plt.plot(data_dict["steps"][i], list(data_dict["b"][i]) , 'b-', )
    '''
    return [lineplot]

anim = animation.FuncAnimation(fig, animate, init_func=init,
                           frames=100, interval=20, blit=True)
anim


# Manually scrubbing back and forth with the slider allows you to choose a point in the building of the plot.
# 
# Note, that for some reason I've seen it break the widget normal looping ability at this time. I'm not sure what I did to break it. Manually scrubbing back and forth with the slider did still work even when that happend. I don't know what I broke to make the one above not loop? I tried adding to `init()` and that didn't seem to help. In fact when it broke it, **it would also break it for the one below that is simpler below**. Weird
# 
# I had seen the glitch I saw was not simply due to including multiple lines because this related, simple code one works to keep looping IN A SEPARATE, or new, NOTEBOOK: (It may or may not work here after running the one above.)

# In[4]:


# If upon first running, it shows a non-interactive, single static shot of the plot below the interactive one with the widget controller,
# JUST RE-RUN TWICE. Re-run usually fixes that display quirk.
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
from IPython.display import HTML, display
plt.rcParams["animation.html"] = "jshtml"
plt.ioff() #needed so the second time you run it you get only single plot

fig = plt.figure()
ax = plt.axes(xlim=(0, 4), ylim=(-2, 2))
lineplot, = ax.plot([], [], lw=3)
lineplot2, = ax.plot([], [], "r-")
    
def init():
    lineplot.set_data([], [])
    return lineplot, #return [lineplot] also works like in https://nbviewer.org/github/raphaelquast/jupyter_notebook_intro/blob/master/jupyter_nb_introduction.ipynb#pre-render-animations-and-export-to-HTML

def animate(i):
    x = np.linspace(0, 4, 1000)
    y = np.sin(2 * np.pi * (x - 0.01 * i))
    z = np.sin(2.2 * np.pi * (x - 0.31 * i))
    lineplot.set_data([x], [y])
    lineplot2.set_data([x], [z])
    return [lineplot,lineplot2]

anim = animation.FuncAnimation(fig, animate, init_func=init,
                           frames=200, interval=20, blit=True)
anim


# In[ ]: