import matplotlib.pyplot as plt
import numpy as np
%matplotlib inline 

G     = 100 # government spending
theta = 0.2 # tax rate
alpha = 0.9 # propensity to consume

n_rounds = 30 # number of rounds we'll consider

# create an array of numbers from 0-30, one for each spending round
r = np.arange(0,n_rounds) 

# solve equation 1 for each individual round
y_n = G*(alpha*(1-theta))**r

# solve equation 2 for each individual round
sum_y = G*(1-(alpha*(1-theta))**(r+1))/(1-alpha*(1-theta))

# plot
plt.bar(r,sum_y, color='r',label='cumulative income')
plt.bar(r,y_n, label='spending round income')
plt.grid()
plt.legend(loc='center right')
plt.xlabel('Spending round, n')
plt.ylabel('Income')

plt.tight_layout()
plt.show()

Y = G/(1-alpha*(1-theta))
print(Y)

print(1/(1-alpha*(1-theta)))

T = theta * Y
print(T)

print(T - G)

S = (1-alpha) * (1-theta) * Y
print(S)