import torch


def format_relaxation(cycles, predictions):
    """
    Format stress relaxation test data to our "post-processed" form for
    comparison. This works for both creep and stress relaxation tests.

    Input data are stresses for stress relaxation and strains for creep

    Cycle 0 indicates the loading part of the test, cycle 1 is the hold

    Zero out the loading results, replace the relaxation results with the
    normalized (subtract t=0) curve

    Args:
      cycles (torch.tensor):        cycle count/listing
      predictions (torch.tensor):   input data

    Returns:
      torch.tensor:                 processed results
    """
    result = torch.zeros_like(predictions)
    rcurve = cycles[:, 0] == 1  # This is right, but dangerous in the future

    curve = predictions[rcurve]
    curve = curve - curve[0]

    result[rcurve] = curve

    return result