import torch


def calculate_cummulate_survive(max_len, gamma, surrogate_step_survival):
    """
    estimate a overall surrogate survival values
    :param input: the src tensor to be attacked. shape: [batch, timestep]
    :param gamma: used in reinforced rewards
    :param surrogate_survival: surrogate single step survival rewards
    :return: a list of cummulated survival for every step,
    with estimate_accumulate_survive[timestep]=accumualted survive of sen_len "timestep"
    """
    estimate_accumulate_survive = [surrogate_step_survival]
    for i in range(1,max_len):
        estimate_accumulate_survive.append(
            estimate_accumulate_survive[i-1]*gamma+surrogate_step_survival
        )
    return torch.tensor(estimate_accumulate_survive)