import torch


def train(model, iterator, optimizer, criterion, clip):

    """
    TRAINING: At each iteration:

        get the source and target sentences from the batch, $X$ and $Y$
        zero the gradients calculated from the last batch
        feed the source and target into the model to get the output, $\hat{Y}$
        as the loss function only works on 2d inputs with 1d targets we need to flatten each of them with .view
            we slice off the first column of the output and target tensors as mentioned above
        calculate the gradients with loss.backward()
        clip the gradients to prevent them from exploding (a common issue in RNNs)
        update the parameters of our model by doing an optimizer step
        sum the loss value to a running total

    Finally, we return the loss that is averaged over all batches.
    """
    model.train()
    epoch_loss = 0
    for i, batch in enumerate(iterator):  # trg is labels, src is input
        src = batch.src
        trg = batch.trg

        optimizer.zero_grad()

        output = model(src, trg)

        # trg = [trg len, batch size]
        # output = [trg len, batch size, output dim]

        output_dim = output.shape[-1]  # ?

        output = output[1:].view(-1, output_dim)  # prediction
        trg = trg[1:].view(-1)  # labels

        loss = criterion(
            output, trg
        )  # calculates both the log softmax as well as the negative log-likelihood of our predictions.

        loss.backward()

        torch.nn.utils.clip_grad_norm_(model.parameters(), clip)  # prevent explosion

        optimizer.step()  # gradient descent / adam

        epoch_loss += loss.item()  # sum all the loss
    return epoch_loss / len(iterator)