import numpy as np
import matplotlib.pyplot as plt


def u(n):
    return np.heaviside(n, 1)


def main():
    n = np.arange(101)
    x = u(n-1)/n

    # Avoid divide by zero
    x[0] = 0
    # Calculate the energy as a cumulative sum from n=0 to each sample
    energy = np.cumsum(np.abs(x)**2)

    plt.stem(n, energy)
    plt.hlines(np.pi**2 / 6, n[0], n[-1], color='r')
    plt.xlabel("Time index")
    plt.ylabel("Energy")
    plt.show()


if __name__ == "__main__":
    main()