import torch


def format_abstract_tensile(cycles, predictions):
    """
    Format abstracted tensile test data, where only the yield strength
    and ultimate tensile strength are available.

    This method relies on the input cycles being:
    * 0: in the elastic regime
    * 1: for the first 1/2 of the remaining time points outside of the elastic regime
    * 2: for the second 1/2 of the remaining time points.

    The inputs are the full stress history.  This function:

    * Cycle 0: overwrite with zeros
    * Cycle 1: overwrite with the first value in cycle 1
    * Cycle 2: overwrite with the maximum stress

    Args:
      cycles (torch.tensor):        cycle count
      predictions (torch.tensor):   input to format

    Returns:
      torch.tensor:                 tensor of processed results
    """
    result = torch.zeros_like(predictions)

    # cycles == 0 -> 0
    result += torch.where(cycles == 0, 0.0, 0.0)

    # cycles == 1 -> first value in that region
    _, index = torch.max(cycles == 1, 0)
    result += torch.where(cycles == 1, predictions.gather(0, index.view(1, -1))[0], 0.0)

    # cycles == 2 -> max value overall
    mvalues, _ = torch.max(predictions, 0)
    result += torch.where(cycles == 2, mvalues, 0.0)

    return result