import torch


def get_loss_and_grads(model, train_x, train_y, flat=True, weights=None, item_loss=True,
                       create_graph=False, retain_graph=False):
    """Computes loss and gradients
    
    Apply model to data (train_x, train_y), compute the loss
    and obtain the gradients.
    
    Parameters
    ----------
    model : nn.Module
        Neural network. We assume the model has a criterion attribute 
    train_x : torch.Tensor
        Training inputs
    train_y : torch.Tensor
        Training targets
    
    Returns
    ----------
    loss
        torch.Tensor of size (#params in model) containing the loss value 
        if flat=True, else a single float
    gradients
        torch.Tensor of size (#params in model) containing gradients w.r.t.
        all model parameters if flat=True, else a structured list
    """
    
    model.zero_grad()
    if weights is None:
        weights = model.parameters()
        out = model(train_x)
    else:
        out = model.forward_weights(train_x, weights)
    
    loss = model.criterion(out, train_y)
    grads = torch.autograd.grad(loss, weights, create_graph=create_graph, 
                                retain_graph=retain_graph)
    
    if flat:
        gradients = torch.cat([p.reshape(-1) for p in grads])
        loss = torch.zeros(gradients.size()).to(train_x.device) + loss.item()
    else:
        gradients = list(grads)
        if item_loss:
            loss = loss.item()
    return loss, gradients