import numpy as np
from scipy import optimize
import matplotlib.pyplot as plt
from types import SimpleNamespace


# Function to find steady state (ss) value for capital (k) of the basic Solow model using scipy.optimize
def ss_bas(n, s, B, alpha, delta):
    """
    Args:
    alpha  (float): Return on physical capital
    delta  (float): Depreciation rate on capital
    s      (float): Savings/investments
    n      (float): Growth rate of population/labour force
    B      (int)  : Total factor productivity (TFP) 
    
    Returns:
    Steady state value for capital, k*
    """
    res = lambda k: s*B*k**alpha - (n+delta)*k # Objective function we wish to find the root of
    
    return optimize.bisect(res, 1,100, full_output=False) # We use the bisect method optimizer from the scipy package to find the root of res


# Function to plot the basic solow model
def ss_bas_plot(n,s,B,alpha,delta,k_max):
    """
    Args:
    alpha  (float): Return on physical capital
    delta  (float): Depreciation rate on capital
    s      (float): Savings/investments
    n      (float): Growth rate of population/labour force
    B      (int)  : Total factor productivity (TFP)
    k_max  (int)  : Used to set the maximum value of x-axis
    
    Returns:
    The solow diagram with (n+delta)*k_t line and the growth paths of physical capital per capita and output per capita"""
    
    # Make an empty list to contain the k values
    k_vec = [0]
    
    # Make an empty list to contain the y values
    y_vec = [0]
    
    # An empty list is made to contain the (n+delta)*k_t line
    line = [0]
    
    # (n+delta)*k_t line
    for k in range(1,k_max):
        l = (n + delta)*k
        line.append(l)
        
    # We generate the growth in capital per capita
    for k in range(1,k_max):
        k_t = s*B*k**alpha
        k_vec.append(k_t)
        
    # We generate the growth in output per capita
    for k in range(1,k_max):
        y_t = B*k**alpha
        y_vec.append(y_t)
    
    # Create the plot
    plt.figure(figsize=(14,7))
    plt.plot(k_vec[:k_max], label=r'$sBk_t^{\alpha}$', color = 'blue')
    plt.plot(line[:k_max], label = r'$(n+\delta)k_t$', color = 'green')
    plt.plot(y_vec[:k_max], label=r'$y_t$', color = 'orange')
    plt.xlim(0,k_max-1)
    plt.ylim(0,y_vec[-1])
    plt.xlabel('$k_t$')
    
    plt.grid(True)
    
    plt.title('Figure 1. Solow diagram - Basic model')
    
    res = lambda k: s*B*k**alpha - (n+delta)*k
    
    k_star = optimize.bisect(res, 1e-8,100000, full_output=False)
    y_star = B*k_star**alpha
    c_star = (1-s)*y_star
    
    plt.scatter(k_star,y_star-c_star,color='black',s=80,zorder=3)
    plt.scatter(k_star,y_star,color='black',s=80,zorder=3)
    plt.axvline(k_star, ymin=(y_star-c_star)/y_vec[-1] ,ymax=y_star/y_vec[-1],color='red',linestyle='-', label='c*')
    plt.axvline(k_star, ymax=(y_star-c_star)/y_vec[-1] ,color='gray',linestyle='--')
    plt.text(k_star + 0.5,0.02 , 'k*', fontsize=12)
    plt.legend()
    
    print(f'For the chosen values of parameters the steady state for k is: k* = {k_star:.3f}')
    print(f'For the chosen values of parameters the steady state for y is: y* = {y_star:.3f}')
    print(f'For the chosen values of parameters the steady state for c is: c* = {c_star:.3f}')
    
    
    
    return plt.show()


# Function to find transition equation for human capital per capita
def h_func(k,s_H,s_K,g,n,alpha,phi,delta,h):
    """args:
    k     (float): Physical capital
    s_H   (float): Investments in human capital
    s_K   (float): Investments in physical capital
    g     (float): Growth in technology
    n     (float): Growth rate of population/labour force
    delta (float): Depreciation rate
    alpha (float): Return on physical capital
    phi   (float): Return on human capital
    h     (float): Human capital
    
    Returns: 
    The transition equation for human capital per capita
        """
    return 1/((1+n)*(1+g))*(s_H*k**(alpha)*h**(phi)+(1-delta)*h)-h

# Function to find transition equation for physical capital per capita
def k_func(h,s_H,s_K,g,n,alpha,phi,delta,k):
    """args:
    h     (float): Human capital
    s_h   (float): Investments in human capital
    s_k   (float): Investments in physical capital
    g     (float): Growth in technology
    n     (float): Growth rate of population/labour force
    delta (float): Depreciation rate
    alpha (float): Return on physical capital
    phi   (float): Return on human capital
    k     (float): Physical capital
    
    Returns: 
    The transition equation for physical capital per capita
    """
    return 1/((1+n)*(1+g))*(s_K*k**(alpha)*h**(phi)+(1-delta)*k)-k


#We define a function expressing the nullclines
def solve_ss(s_H,s_K,g,n,alpha,phi,delta):
    """args:
    sh    (float): Investments in human capital
    sk     float): Investments in physical capital
    g     (float): Growth in technology
    n     (float): Growth in labour force
    alpha (float): Return on physical capital
    phi   (float): Return on human capital
    delta (float): Depriciation rate
    
    Returns:
    Nullclines for physical and human capital
    """
    
    
    # a. grids
    
    k_vec = np.linspace(1e-4,100,500)
    h_vec_DeltaK0 = np.empty(500)
    h_vec_DeltaH0 = np.empty(500)
    
    # b. solve for each k in grid
    for i,k in enumerate(k_vec):
        #Solve for delta_h=0
        
        # i. objective
        obj = lambda h: -h_func(k,s_H,s_K,g,n,alpha,phi,delta,h)

        # ii. optimizer
        result = optimize.root_scalar(obj,method='brentq',bracket=[1e-20,1000])
        
        h_vec_DeltaH0[i] = result.root

       
        #Solve for delta_k=0
        # i. objective
        obj = lambda h: -k_func(h,s_H,s_K,g,n,alpha,phi,delta,k)

        # ii. optimizer
        result = optimize.root_scalar(obj,method='brentq',bracket=[1e-20,1000])

        h_vec_DeltaK0[i] = result.root

    return k_vec, h_vec_DeltaK0, h_vec_DeltaH0
