def hi():
    
    print('Hello world!')
    
hi()

"""
Explanation: Class 10: User-defined functions and a Solow growth model example
User-defined functions
Create a new function by using the def keyword followed by the designated name of the new function. In the definition, the function name has to be followed by a set of parentheses and a colon. If the function has arguments, put them inside the parentheses separated by commas. The code to be performed when the function is run follows the definition line and is indented.
Example: A function that prints Hello world!
End of explanation
"""


def cobbDouglas(A,alpha,k):
    ''' Computes output per worker y given A, alpha, and a value of capital per worker k
    
    Args:
        
        A (float):                TFP
        alpha (float):            Cobb-Douglas parameter
        k (float or numpy array): capital per worker
        
    Returns
        
        float or numpy array'''
    
    return A*k**alpha

"""
Explanation: Notice that in the previous example, the function takes no arguments and returns nothing. It just does the task that it's supposed to.
Example:
A function that returns the computes the fllowing production function:
\begin{align}
y & = A k^{\alpha}
\end{align}
End of explanation
"""


# Use cobbDouglas() to plot the production function for a bunch of values of alpha between 0 and 1.


"""
Explanation: Note the cobbDouglas() has a docstring. The docstring is optional, but it tells users about the function. The contents of the docstring can be accessed with the help() function. It's good practice to make use of doc strings.
End of explanation
"""


def solow_example(A,alpha,delta,s,n,K0,L0,T):
    '''Returns DataFrame with simulated values for a Solow model with labor growth and constant TFP
    
        Args:
            A (float):     TFP
            alpha (float): Cobb-Douglas production function parameter
            delta (float): capital deprection rate
            s (float):     saving rate
            n (float):     labor force growth rate
            K0 (float):    initial capital stock
            L0 (float):    initial labor force
            T (int):       number of periods to simulate
            
        Returns:
            pandas DataFrame with columns: 
                
                'capital', 'labor', 'output', 'consumption', 'investment',
                'capital_pw','output_pw', 'consumption_pw', 'investment_pw'
            '''
    
    # Initialize a variable called capital as a (T+1)x1 array of zeros and set first value to K0
    capital = np.zeros(T+1)
    capital[0] = K0
    
    # Initialize a variable called labor as a (T+1)x1 array of zeros and set first value to L0
    labor = np.zeros(T+1)
    labor[0] = L0


    # Compute all capital and labor values by iterating over t from 0 through T
    for t in np.arange(T):
        labor[t+1] = (1+n)*labor[t]
        capital[t+1] = s*A*capital[t]**alpha*labor[t]**(1-alpha) + (1-delta)*capital[t]
    
    # Store the simulated capital df in a pandas DataFrame called data
    df = pd.DataFrame({'capital':capital,'labor':labor})
    
    # Create columns in the DataFrame to store computed values of the other endogenous variables
    df['output'] = df['capital']**alpha*df['labor']**(1-alpha)
    df['consumption'] = (1-s)*df['output']
    df['investment'] = df['output'] - df['consumption']
    
    # Create columns in the DataFrame to store capital per worker, output per worker, consumption per worker, and investment per worker
    df['capital_pw'] = df['capital']/df['labor']
    df['output_pw'] = df['output']/df['labor']
    df['consumption_pw'] = df['consumption']/df['labor']
    df['investment_pw'] = df['investment']/df['labor']
    
    return df

"""
Explanation: The Solow model with exogenous population growth
Recall the Solow growth model with exogenous labor growth:
\begin{align}
Y_t & = AK_t^{\alpha} L_t^{1-\alpha}\tag{1}
\end{align}
The supply of labor grows at an exogenously determined rate $n$ and so it's value is determined recursively by a first-order difference equation:
\begin{align}
L_{t+1} & = (1+n) L_t \tag{2}
\end{align}
The rest of the economy is characterized by the same equations as before:
\begin{align}
C_t & = (1-s)Y_t \tag{3}\
Y_t & = C_t + I_t \tag{4}\
K_{t+1} & = I_t + ( 1- \delta)K_t \tag{5}\
\end{align}
Combine Equations (1), (3), (4), and (5) to eliminate $C_t$, $I_t$, and $Y_t$ and obtain a recurrence relation specifying $K_{t+1}$ as a funtion of $K_t$ and $L_t$:
\begin{align}
K_{t+1} & =  sAK_t^{\alpha}L_t^{1-\alpha} + ( 1- \delta)K_t \tag{6}
\end{align}
Given an initial values for capital and labor, Equations (2) and (6) can be iterated on to compute the values of the capital stock and labor supply at some future date $T$. Furthermore, the values of consumption, output, and investment at date $T$ can also be computed using Equations (1), (3), (4), and (5).
Use a function to simulate the model
Suppose that we wanted to simulate the Solow model with different parameter values so that we could compare the simulations. Since we'd be doing the same basic steps multiple times using different numbers, it would make sense to define a function so that we could avoid repetition.
The code below defines a function called solow_example() that simulates the Solow model with exogenous labor growth. solow_example() takes as arguments the parameters of the Solow model $A$, $\alpha$, $\delta$, $s$, and $n$; the initial values $K_0$ and $L_0$; and the number of simulation periods $T$. solow_example() returns a Pandas DataFrame with computed values for aggregate and per worker quantities.
End of explanation
"""


# Create the DataFrame with simulated values
df = solow_example(A=10,alpha=0.35,delta=0.1,s=0.15,n=0.01,K0=20,L0=1,T=100)

# Create a 2x2 grid of plots of the capital per worker, output per worker, consumption per worker, and investment per worker
fig = plt.figure(figsize=(12,8))
ax = fig.add_subplot(2,2,1)
ax.plot(df['capital_pw'],lw=3)
ax.grid()
ax.set_title('Capital per worker')

ax = fig.add_subplot(2,2,2)
ax.plot(df['output_pw'],lw=3)
ax.grid()
ax.set_title('Output per worker')

ax = fig.add_subplot(2,2,3)
ax.plot(df['consumption_pw'],lw=3)
ax.grid()
ax.set_title('Consumption per worker')

ax = fig.add_subplot(2,2,4)
ax.plot(df['investment_pw'],lw=3)
ax.grid()
ax.set_title('Investment per worker')

"""
Explanation: Example: A single simulation
Use the function solow_example() to simulate the Solow growth model with exogenous labor growth for $t=0\ldots 100$. For the simulation, assume the following values of the parameters:
\begin{align}
A & = 10\
\alpha & = 0.35\
s & = 0.15\
\delta & = 0.1\
n & = 0.01
\end{align}
Furthermore, suppose that the initial values of capital and labor are:
\begin{align}
K_0 & = 20\
L_0 & = 1
\end{align}
End of explanation
"""


# Construct two dataframes df1 and df2 with the simulated values.


# Create a 2x2 grid of plots of the capital per worker, output per worker, consumption per worker, and investment per worker


"""
Explanation: Example: Two simulations with different initial capital stocks
Repeat the previous exercise for two simulations of the Solow model having two different initial values of capital: $K_0 = 20$ and $K_0'=10$.
End of explanation
"""


# Construct three dataframes df1, df2, df3 with the simulated values.


# Create a 2x2 grid of plots of the capital per worker, output per worker, consumption per worker, and investment per worker


"""
Explanation: Example: Three simulations with different TFPs
Repeat the previous exercise for three simulations of the Solow model having three different values of TFP: $A = 5$,  $A=10$, and $A = 15$.
End of explanation
"""


