#!/usr/bin/env python

from pylab import array, plot, zeros, ones, arange, figure, NaN

tmax=50

alpha1=0.6; alpha2=0.25 # consumption parameters
theta =0.2 # tax rate

# define model SIM

H_time = zeros(tmax); H_time[0]=25
C_time = zeros(tmax); C_time[0]=NaN
G_time = ones(tmax) * 20
T_time = zeros(tmax); T_time[0]=NaN
Y_time = zeros(tmax); Y_time[0]=NaN
Yd_time = zeros(tmax); Yd_time[0]=NaN

for tstep in range(1,tmax):
    C_time[tstep] = (alpha2*H_time[tstep-1]+alpha1*G_time[tstep]*(1-theta))/(1.-alpha1*(1-theta))
    H_time[tstep] = H_time[tstep-1]+(1-theta)*G_time[tstep]-theta*C_time[tstep]
    Y_time[tstep] = C_time[tstep] + G_time[tstep]
    T_time[tstep] = theta * Y_time[tstep]
    Yd_time[tstep] = Y_time[tstep]-T_time[tstep]

# plot time evolution of stocks
fig1 = figure()
ax1 = fig1.add_subplot(111)
ax1.set_xlabel(r"$t$",fontsize=14)

ax1.plot(arange(tmax),H_time,'ok',label=r"$H_h = H_s$")
ax1.plot(arange(tmax),Y_time,'^m',label=r"$Y$")
ax1.plot(arange(tmax),Yd_time,'vw',label=r"$Y_d$")
ax1.plot(arange(tmax),C_time,'sb',label=r"$C$")
ax1.plot(arange(tmax),G_time,'<r',label=r"$G$")
ax1.plot(arange(tmax),T_time,'>y',label=r"$T$")

ax1.legend(loc="upper left", frameon=False).get_frame().set_facecolor('none')

fig1.savefig("godley-lavoie-chapter-3-sfc-model-plot.svg", bbox_inches = 'tight')