# RNN for classification

import torch

import numpy as np

import torch.nn as nn

import torch.utils.data as Data

import matplotlib.pyplot as plt

import torchvision

#hyper parameters

TIME_STEP=10     #run time step

INPUT_SIZE=1

LR=0.02   #learning rate

# t=np.linspace(0,np.pi*2,100,dtype=float)  #from zero to pi*2, and one hundred point there

# x=np.sin(t)

# y=np.cos(t)

# plt.plot(t,x,'r-',label='input (sin)')

# plt.plot(t,y,'b-',label='target (cos)')

# plt.legend(loc='best')

# plt.show()

class RNN_Net(nn.Module):

    def __init__(self):

        super(RNN_Net,self).__init__()

        self.rnn=nn.RNN(

            input_size=INPUT_SIZE,

            hidden_size=32,

            num_layers=1,

            batch_first=True,

        )

        self.out=nn.Linear(32,1)

    def forward(self,x,h_state):

        r_out,h_state=self.rnn(x,h_state)

        outs=[]

        for time_step in range(r_out.size(1)):

            outs.append(self.out(r_out[:,time_step,:]))

        return torch.stack(outs,dim=1),h_state  # the type of return data is torch, and the return data also include h_state

rnn=RNN_Net()

# print(rnn)

optimizer=torch.optim.Adam(rnn.parameters(),lr=LR)

loss_func=nn.MSELoss()

plt.ion()

h_state=None

for step in range (60):

    start,end=step*np.pi,(step+1)*np.pi

    #using sin predicts cos

    steps=np.linspace(start,end,TIME_STEP,dtype=np.float32)

    x_np=np.sin(steps)

    y_np=np.cos(steps)

    x=torch.from_numpy(x_np[np.newaxis,:,np.newaxis]) # np.newaxis means increase a dim

    y=torch.from_numpy(y_np[np.newaxis,:,np.newaxis])

    predition,h_state=rnn(x,h_state) #the first h_state is None

    h_state=h_state.data  #?????

    loss=loss_func(predition,y)

    optimizer.zero_grad()

    loss.backward()

    optimizer.step()

    plt.plot(steps,y_np,'r-')

    plt.plot(steps,predition.detach().numpy().flatten(),'b-')  #flatten() 展平维度

    plt.draw()

    plt.pause(0.05)

plt.ioff()

plt.show()