# AUTOGENERATED! DO NOT EDIT! File to edit: Model.ipynb (unless otherwise specified).

__all__ = ['Sequential']

# Cell
import torch
import torch.nn as nn

from functools import partial
from fastai import *
from fastai.basic_train import Learner

from .data import create_db
from .layers import *
from .losses import *

from IPython.display import clear_output

# Cell
class Sequential:
    def __init__(self, model=None):
        self.layers = []
        self.last_dim = None
        self.model = model
        self.device = torch.device('cpu')
        if torch.cuda.is_available():
            self.device = torch.device('cuda')

    def add(self, layer):
        layer = layer.get_layer(self.last_dim)
        self.last_dim = layer['output_dim']
        self.layers.extend(layer['layers'])

    def compile(self, loss, optimizer=None):
        if len(self.layers) > 0:
            self.model = nn.Sequential(*self.layers)
        self.loss = loss

    def fit(self, x, y, bs, epochs, lr=1e-3, one_cycle=True, get_lr=True):
        db = create_db(x, y, bs=bs)
        self.learn = Learner(db, self.model, loss_func=self.loss)
        if one_cycle:
            self.learn.fit_one_cycle(epochs, lr)
        else:
            self.learn.fit(epochs, lr)

    def lr_find(self, x, y, bs):
        db = create_db(x, y, bs=bs)
        learn = Learner(db, self.model, loss_func=self.loss)
        learn.lr_find()
        clear_output()
        learn.recorder.plot(suggestion=True)

    def predict(self, x):
        self.learn.model.eval()
        with torch.no_grad():
            y_preds = self.learn.model(torch.Tensor(x).to(device))
        return y_preds.cpu().numpy()