import numpy as np
import tensorflow as tf


def _stack_past(x, steps):
    """Stack the past data for each step.
    Ex. x = [0, ..., 60]. steps = [10, 20]
    Result [x, x[:-10], x[:-20]] normalized to the same shape
    """
    # Sort the steps in ascending order [10, 20]
    sorted_steps = steps.copy()
    sorted_steps.sort()

    largest_step = sorted_steps[-1]

    # Include the original data
    stacks = [x[largest_step:]]

    for step in sorted_steps:
        # Normalize the shapes by skipping the difference from the largest step
        # Which are rows without enough past data
        skip = largest_step - step
        stacks.append(x[skip:-step])

    return stacks


def np_stack_past(x, steps):
    stacks = _stack_past(x, steps)
    return np.stack(stacks, axis=-1)


def tf_stack_past(tensor, steps):
    stacks = _stack_past(tensor, steps)
    return tf.stack(stacks, axis=-1)