# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/02_learner/03_learner.ipynb.

# %% auto 0
__all__ = ['get_inp_out_size', 'RNNLearner', 'TCNLearner', 'CRNNLearner', 'AR_TCNLearner', 'AR_RNNLearner']

# %% ../../nbs/02_learner/03_learner.ipynb 2
from ..data import *
from ..models import *
from .callbacks import *
from .losses import *

from fastai.basics import *
from fastai.callback.tracker import *

# %% ../../nbs/02_learner/03_learner.ipynb 6
def get_inp_out_size(dls):
    '''returns input and output size of a timeseries databunch'''
    tup = dls.one_batch()
    inp = tup[0].shape[-1]
    out = tup[1].shape[-1]
    return inp,out

# %% ../../nbs/02_learner/03_learner.ipynb 9
@delegates(SimpleRNN, keep=True)
def RNNLearner(dls,loss_func=nn.L1Loss(),metrics=[fun_rmse],n_skip=0,num_layers=1,hidden_size=100,stateful=False,opt_func=Adam,cbs=None,**kwargs):
    if cbs is None: cbs = []

    inp,out = get_inp_out_size(dls)
    model = SimpleRNN(inp,out,num_layers,hidden_size,stateful=stateful,**kwargs)
  
    skip = partial(SkipNLoss,n_skip=n_skip)
        
    metrics= [skip(f) for f in metrics]

    if stateful: 
        cbs.append(TbpttResetCB())
        # if stateful apply n_skip with a callback for the first minibatch of a tbptt sequence
        cbs.append(SkipFirstNCallback(n_skip))
    else:
        loss_func = skip(loss_func)
        
    lrn = Learner(dls,model,loss_func=loss_func,opt_func=opt_func,metrics=metrics,cbs=cbs,lr=3e-3)
    return lrn

# %% ../../nbs/02_learner/03_learner.ipynb 14
@delegates(TCN, keep=True)
def TCNLearner(dls,num_layers=3,hidden_size=100,loss_func=nn.L1Loss(),metrics=[fun_rmse],n_skip=None,opt_func=Adam,cbs=None,**kwargs):
    inp,out = get_inp_out_size(dls)
    n_skip = 2**num_layers if n_skip is None else n_skip
    model = TCN(inp,out,num_layers,hidden_size,**kwargs)
  
    skip = partial(SkipNLoss,n_skip=n_skip)
        
    metrics= [skip(f) for f in metrics]
    loss_func = skip(loss_func)
        
    lrn = Learner(dls,model,loss_func=loss_func,opt_func=opt_func,metrics=metrics,cbs=cbs,lr=3e-3)
    return lrn

# %% ../../nbs/02_learner/03_learner.ipynb 17
@delegates(CRNN, keep=True)
def CRNNLearner(dls,loss_func=nn.L1Loss(),metrics=[fun_rmse],n_skip=0,opt_func=Adam,cbs=None,**kwargs):
    inp,out = get_inp_out_size(dls)
    model = CRNN(inp,out,**kwargs)
  
    skip = partial(SkipNLoss,n_skip=n_skip)
        
    metrics= [skip(f) for f in metrics]
    loss_func = skip(loss_func)
        
    lrn = Learner(dls,model,loss_func=loss_func,opt_func=opt_func,metrics=metrics,cbs=cbs,lr=3e-3)
    return lrn

# %% ../../nbs/02_learner/03_learner.ipynb 20
@delegates(TCN, keep=True)
def AR_TCNLearner(dls,hl_depth=3,alpha=1,beta=1,early_stop=0,metrics=None,n_skip=None,opt_func=Adam,**kwargs):
    n_skip = 2**hl_depth if n_skip is None else n_skip
    skip = partial(SkipNLoss,n_skip=n_skip)
    
    inp,out = get_inp_out_size(dls)
    model = AR_Model(TCN(inp+out,out,hl_depth,**kwargs),ar=False)
    model.init_normalize(dls.one_batch())
    
    cbs=[ARInitCB(),TimeSeriesRegularizer(alpha=alpha,beta=beta,modules=[model.model.conv_layers[-1]])]#SaveModelCallback()
    if early_stop > 0:
        cbs += [EarlyStoppingCallback(patience=early_stop)]
        
    if metrics is None: metrics=SkipNLoss(fun_rmse,n_skip)
        
    lrn = Learner(dls,model,loss_func=nn.L1Loss(),opt_func=opt_func,metrics=metrics,cbs=cbs,lr=3e-3)
    return lrn

# %% ../../nbs/02_learner/03_learner.ipynb 23
@delegates(SimpleRNN, keep=True)
def AR_RNNLearner(dls,alpha=0,beta=0,early_stop=0,metrics=None,n_skip=0,opt_func=Adam,**kwargs):
    skip = partial(SkipNLoss,n_skip=n_skip)
    
    inp,out = get_inp_out_size(dls)
    model = AR_Model(SimpleRNN(inp+out,out,**kwargs),ar=False)
    model.init_normalize(dls.one_batch())
    
    cbs=[ARInitCB(),TimeSeriesRegularizer(alpha=alpha,beta=beta,modules=[model.model.rnn])]#SaveModelCallback()
    if early_stop > 0:
        cbs += [EarlyStoppingCallback(patience=early_stop)]
        
    if metrics is None: metrics=SkipNLoss(fun_rmse,n_skip)
        
    lrn = Learner(dls,model,loss_func=nn.L1Loss(),opt_func=opt_func,metrics=metrics,cbs=cbs,lr=3e-3)
    return lrn
