import os

import torch
import torch.nn as nn
import torch.optim as optim
import tqdm


class CharRNNTrainer:

    def __init__(self, config):
        self.config = config

    def fit(self, model, data):
        def get_params():
            return (p for p in model.parameters() if p.requires_grad)

        if isinstance(data, tuple):
            train_dataloader = data[0]
            val_dataloader = data[1]
        else:
            train_dataloader = data
            val_dataloader = None

        num_epochs = self.config.num_epochs

        criterion = nn.CrossEntropyLoss()
        optimizer = optim.Adam(get_params(), lr=self.config.lr)

        best_val_loss = float('inf')

        for epoch in range(num_epochs):
            model.train()
            train_dataloader = tqdm.tqdm(train_dataloader)
            train_dataloader.set_description('Train (epoch #{})'.format(epoch))

            self._pass_data(model, train_dataloader, criterion, optimizer)

            if val_dataloader is not None:
                val_dataloader = tqdm.tqdm(val_dataloader)
                val_dataloader.set_description('Validation (epoch #{})'.format(epoch))

                val_loss = self._pass_data(model, val_dataloader, criterion)

                if best_val_loss > val_loss:
                    best_val_loss = val_loss
                    torch.save(model.state_dict(),
                               os.path.join(self.config.model_save, "model-iter" + str(epoch) + ".pt"))  # TODO
            else:
                torch.save(model.state_dict(),
                           os.path.join(self.config.model_save, "model-iter" + str(epoch) + ".pt"))  # TODO

    def _pass_data(self, model, dataloader, criterion, optimizer=None):
        if optimizer is None:
            model.eval()
        else:
            model.train()

        running_loss = 0

        for i, (prevs, nexts, lens) in enumerate(dataloader):

            outputs, _, _ = model(prevs, lens)

            loss = criterion(outputs.view(-1, outputs.shape[-1]), nexts.view(-1))

            postfix = {'loss': loss.item()}
            dataloader.set_postfix(postfix)

            running_loss += loss * len(prevs)

            if optimizer is not None:
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

        return running_loss.item() / len(dataloader)

@@ -648,46 +648,8 @@
 r)%0A%0A
-        best_val_loss = float('inf')%0A%0A
     
@@ -1198,435 +1198,8 @@
 n)%0A%0A
-                if best_val_loss %3E val_loss:%0A                    best_val_loss = val_loss%0A                    torch.save(model.state_dict(),%0A                               os.path.join(self.config.model_save, %22model-iter%22 + str(epoch) + %22.pt%22))  # TODO%0A            else:%0A                torch.save(model.state_dict(),%0A                           os.path.join(self.config.model_save, %22model-iter%22 + str(epoch) + %22.pt%22))  # TODO%0A%0A
     

fix char rnn trainer