import torch


def sequence_loss(preds, flow_gt, gamma, valid=None, max_flow=400):
    """Compute sequence loss between prediction and ground truth.

    Args:
        preds (list(torch.Tensor)): List of flow prediction from
            flow_estimator.
        flow_gt (torch.Tensor): Ground truth flow map.
        gamma (float): Scale factor gamma in loss calculation.
        valid (torch.Tensor, optional): Tensor Used to exclude invalid pixels.
            Default: None.
        max_flow (int, optional): Used to exclude extremely large
            displacements. Default: 400.

    Returns:
        flow_loss (float): Total sequence loss.
    """
    n_preds = len(preds)
    flow_loss = 0.
    mag = torch.sum(flow_gt**2, dim=1).sqrt()
    if valid is None:
        valid = torch.ones(flow_gt[:, 0, :, :].shape).to(flow_gt)
    else:
        valid = ((valid >= 0.5) & (mag < max_flow)).to(flow_gt)

    for i, pred in enumerate(preds):
        i_weight = gamma**(n_preds - i - 1)
        i_loss = (pred - flow_gt).abs()
        flow_loss += i_weight * (valid[:, None] * i_loss).mean()
    return flow_loss