#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@author: Zach Swaziek
Value Function Iteration in Baby Growth Model
Econ 8105
"""

#Import packages
import numpy as np
from matplotlib import pyplot as plt

#Define parameters
#You need to plug in the correct values for PS4
delta = 1
beta = 0.6
theta = 5
alpha = 0.5

#Define capital grid space
#Plug in correct values for grid
k_min = 1
k_max = 5
k_grid_pts = 10
k_grid = np.linspace(k_min, k_max, k_grid_pts)
#Print k_grid to check that it is correct
print(k_grid)

#Define utility function
def logu(c):
    return np.log(c)

#Define function for value function iteration
def VFI(utility, v0, delta, beta, theta, alpha, k_grid, max_iterations, tolerance):
    #This is also the initial guess v_0
    v_n = v0
    #This will be updated value function
    v_n_plus_1 = np.zeros(len(k_grid))
    #Need a vector for the policy function
    policy = np.zeros(len(k_grid))
    
    #Begin value function iteration
    #I like to set a maximum number of iterations just to make sure my code doesn't get
    #stuck in an infinite loop
    iterations = 1
    for i in range(max_iterations):
        #Loop over each value of k in capital grid
        #k_loc means the location of k on the k_grid
        for k_loc in range(len(k_grid)):
            #Initialize v_compare_at_k
            #This is the value of the value function if the optimal choice was k' = k_min; 
            v_compare_at_k = utility(theta*(k_grid[k_loc]**alpha) + (1-delta)*k_grid[k_loc] - k_grid[0]) + beta*v_n[0]
            #Initialize v_max_at_k
            v_max_at_k = 0
            #Initialize the location (within the kgrid vector) of the optimal policy at k 
            opt_policy_loc_at_k = 0
            #Loop over k' values and find the k' that maximizes the value function
            for k_prime_loc in range(len(k_grid)):
                #Only check for optimal k' when consumption is positive
                #But we also know that by Inada, optimal consumption will be non-zero
                c = theta*(k_grid[k_loc]**alpha) + (1-delta)*k_grid[k_loc] - k_grid[k_prime_loc]
                if c > 0:
                    #Find value of u(c) + BV(k') at k_prime
                    v_cur = utility(c) + beta*v_n[k_prime_loc]
                    #Check if v_cur is greater than or equal to v_compare_at_k
                    if v_cur >= v_compare_at_k:
                        #If greater than previous value at k, this is the new candidate for v_max_at_k
                        v_max_at_k = v_cur
                        #Our new v_compare is also v_cur - when we check the VF at the next candidate for k'
                        #we want to compare it the value function at the previous k'
                        v_compare_at_k = v_cur
                        #Store location of this candidate for optimal k'
                        opt_policy_loc_at_k = k_prime_loc

                    #Otherwise, try a different value of k'
                    else:
                        pass

                #If consumption is negative, break this loop and start at the next k' candidate
                else:
                    break

            #For this k, store the value of the value function at the maximum
            v_n_plus_1[k_loc] = v_max_at_k
            #And the value of the policy function at k
            policy[k_loc] = k_grid[opt_policy_loc_at_k]

        #Once have foundnd value function and policy function, calculate maximum |v_n(k) and v_n+1(k)|
        max_abs_dif_vn_and_vn_plus_1 = np.max(np.abs(v_n-v_n_plus_1))
        
        #If maximum difference is less than our tolerance, break out of loop
        if max_abs_dif_vn_and_vn_plus_1 < tolerance:
            print("Converged in", iterations, "iterations")
            break
        else:
            #Otherwise, start another iteration with v_n+1 as the value function
            v_n = v_n_plus_1
            
            #Reset v_n+1
            v_n_plus_1 = np.zeros(len(k_grid))
            
        #Update number of iterations
        iterations = iterations + 1
    
    #Return the value function and policy function
    return v_n_plus_1, policy
                
################################################################################
#Calculate Value Function and Policy Function
################################################################################

#Initial guess is v0(k) = 0
v0 = np.zeros(len(k_grid))
max_iterations = 1000
tolerance =1e-8
valf, pol = VFI(logu, v0, delta, beta, theta, alpha, k_grid, max_iterations, tolerance)

#Plot the policy function
plt.figure(1)
plt.plot(k_grid, pol, label = "Numerical")
#Plot the analytical solution from question 2 when delta = 1
#plt.plot(k_grid, , label="Analytical")
plt.title("Policy Function")
plt.legend(loc="upper left")


#Plot for value function
plt.figure(2)
plt.plot(k_grid, valf, label = "Numerical")
#Plot the analytical solution from question 2 when delta = 1
#plt.plot(k_grid,  label="Analytical")
plt.title("Value Function")
plt.legend(loc="upper left")

plt.show