import torch


def collect_results(tgt_ts_output, tgt_cat_output, patent_tgt):
    """Prepare results and ground truth for evaluation.

    Parameters
    ----------
    tgt_ts_output : :class:`torch.Tensor`
        Prediction for arrival time, a tensor of shape (seq_len, batch, 1).
    tgt_cat_output: :class:`torch.Tensor`
        Prediction for category, a tensor of shape (seq_len, batch,
        num_categories).
    patent_tgt : dict
        Ground truth data for patent prediction including real timestamp,
        category and mask.
    Returns
    -------
    mae : :class:`torch.Tensor`
        (loss, # of points)
    acc : :class:`torch.Tensor`
        (# of correct predictions, # of points)

    """

    # get predictions and ground ready. y'all agree on dimensions first.
    tgt_pts = patent_tgt['pts'].unsqueeze(-1)
    tgt_pcat = patent_tgt['pcat'].unsqueeze(-1)
    mask = patent_tgt['mask'].unsqueeze(-1)
    pred_cat = tgt_cat_output.topk(1, dim=2)[1]
    assert(tgt_pts.size() == tgt_pcat.size() == mask.size() == pred_cat.size()
           == tgt_ts_output.size())
    # timestamp predictions and ground truth
    pred_ts = tgt_ts_output.masked_select(mask)
    tgt_ts = tgt_pts.masked_select(mask)
    # category predictions and ground truth
    pred_cat = pred_cat.masked_select(mask)
    tgt_cat = tgt_pcat.masked_select(mask)
    # calculate mae and accuracy
    mae = torch.tensor([torch.abs(pred_ts - tgt_ts).sum().item(),
                        pred_ts.size(0)])
    acc = torch.tensor([torch.sum((pred_cat == tgt_cat)).item(),
                        pred_cat.size(0)], dtype=torch.float)
    return mae, acc