# -*- coding: utf-8 -*-

"""
Created on Mon Feb 14 10:20:38 2022

@author: gbongio
"""

import numpy as np
from matplotlib import pyplot as plt
from functools import partial

from copy import copy, deepcopy

import time

class A(object):
def __init__(self):
print ('init')
self.v = 10
self.z = [2,3,4]

def __copy__(self):
cls = self.__class__
result = cls.__new__(cls)
result.__dict__.update(self.__dict__)
return result

def __deepcopy__(self, memo):

cls = self.__class__
result = cls.__new__(cls)
memo[id(self)] = result
for k, v in self.__dict__.items():
setattr(result, k, deepcopy(v, memo))
return result

nkk = 400
nzz = 12

alpha = 0.36
beta = 0.96
delta = 0.025
rho = 0.95
var_e = 0.007
theta = 0.64

# a.Create a tensor which stores the optimal amount of hours worked conditional
# for each possible combination tensor = 12 matrics 400x400
tensor_l = np.empty([nzz, nkk, nkk])
tensor_u = np.empty([nzz, nkk, nkk])

# b.
# solve for ss w/o uncertainty optimal labor
l_ss = ((1-theta)*(1-alpha)) / ((1-theta)*(1-alpha) - ((alpha*beta*theta*delta)/(1-
beta*(1-delta))) + theta)

# solve for ss w/o uncertainty capital where l=l', k=k' and z=0
k_ss = l_ss * ((1-beta*(1-delta))/(alpha*beta))**(1/(alpha-1))
# grid
k_grid = np.linspace(0.5*k_ss,1.5*k_ss, num=400)

# c.

import numpy as np
from scipy.stats import norm

def approx_markov(rho, sigma_u, m=2, n=12):

"""
Computes the Markov matrix associated with a discretized version of
the linear Gaussian AR(1) process

y_{t+1} = rho * y_t + u_{t+1}

according to Tauchen's method. Here {u_t} is an iid Gaussian

process with zero mean.

Parameters
----------
rho : scalar(float)
The autocorrelation coefficient
sigma_u : scalar(float)
The standard deviation of the random process
m : scalar(int), optional(default=3)
The number of standard deviations to approximate out to
n : scalar(int), optional(default=7)
The number of states to use in the approximation

Returns
-------

x : array_like(float, ndim=1)
The state space of the discretized process
P : array_like(float, ndim=2)
The Markov transition matrix where P[i, j] is the probability
of transitioning from x[i] to x[j]

"""
F = norm(loc=0, scale=sigma_u).cdf

# standard deviation of y_t

std_y = np.sqrt(sigma_u**2 / (1-rho**2))

# top of discrete state space

x_max = m * std_y

# bottom of discrete state space

x_min = - x_max

# discretized state space

x = np.linspace(x_min, x_max, n)

step = (x_max - x_min) / (n - 1)

half_step = 0.5 * step

# initialize an empty transition matrix

 P = np.empty((n, n))

for i in range(n):
# for every step compute the probability mass for the extreme values of the
grid
P[i, 0] = F(x[0]-rho * x[i] + half_step) #lowest value
P[i, n-1] = 1 - F(x[n-1] - rho * x[i] - half_step) #highest value

for j in range(1, n-1):

# compute the probabilities of the elements not at the extremes of
# the grid i.e. get the
z = x[j] - rho * x[i]
P[i, j] = F(z + half_step) - F(z - half_step)

return x, P

var_z = (var_e**2)/(1-rho**2)

z_grid, transition = approx_markov(rho, var_e)

# d.Fill in a nkk # nkk # nzz tensor containing the utility level associated with
any feasible combination
# of (k; z; k') using the optimal labor choice u(k; z; k')

# conceptually what we are going to do now is:

# for given triplet (k, k1, z) compute the optimal labor supply which we do
with a bisection on the labor supply FOC
# follow by computing consumption

def c(k, k1, z, l, alpha, delta):

consumption = np.exp(z)*(k**alpha)*(l**(1-alpha))+(1-delta)*k-k1
return consumption

def lmin(k, k1, z, alpha, delta):

if k1-(1-delta)*k > 0:
return (k1-(1-delta)*k)/(np.exp(z)*k**alpha)**(1/(1-alpha))
else:
return 0

def labor_foc(k, k1, z, theta, alpha, delta, l):

consumption = c(k, k1, z, l, alpha, delta)
return (((1-theta)*(1-alpha)*(1-l)*(np.exp(z))*(k**alpha)) / consumption ) -
(l**(alpha))*theta

#np.exp(z)*(1-theta)*(1-alpha)*(1-l) - np.exp(z)*theta*l - (1-delta) * k**(1-alpha)

* l**alpha * theta + (k1/k)*theta

def my_bisection(f, a, b, tol):

# approximates a root, R, of f bounded
# by a and b to within tolerance
# | f(m) | < tol with m the midpoint
# between a and b Recursive implementation
 # check if a and b bound a root
if np.sign(f(a)) == np.sign(f(b)) & f(a) < 0:
return min(1,a)
elif a > 1:
return 1

# get midpoint
m = (a + b)/2

if np.abs(f(m)) < tol:

# stopping condition, report m as root
return m
elif np.sign(f(a)) == np.sign(f(m)):
# case where m is an improvement on a.
# Make recursive call with a = m
return my_bisection(f, m, b, tol)
elif np.sign(f(b)) == np.sign(f(m)):
# case where m is an improvement on b.
# Make recursive call with b = m
return my_bisection(f, a, m, tol)

def u(c, l, theta, alpha, delta):

# if consumption is negative return nan
if k1 > np.exp(z) * (k**alpha) * (l**(1-alpha)) + (1-delta)*k:
return -np.inf
else:
return (1-theta)*np.log(c) + theta*np.log(1-l)

# in this loop we are going to compute the optimal labor supply given a state
# (k,k1,z) and and the associated utilities in period t = 0.
# This information will come handy when performing VF iteration
count=0
for k in range(nkk):
count+=1
for k1 in range(nkk):
for z in range(nzz):

# define partial function where the only input becomes the labor
# this means that for each triplet (k,k1,z) we fix such values in the
# formula for the labor foc and then solve with the bisection for the
root
# of the derivative.
labor_partial = partial(labor_foc, k_grid[k], k_grid[k1], z_grid[z],
theta, alpha, delta)

# set the threshold delivering positive consumption, this comes at hand

# when applying the bisection.
l_min = lmin(k_grid[k], k_grid[k1], z_grid[z], alpha, delta)

#fill the respective spot in the tensor of optimal labor supply

responses
tol = 10**-4

tensor_l[z, k, k1] = my_bisection(labor_partial, l_min, 1, tol)

# compute consumption associated with the optimal labor supply given

# the triplet (k,k1,z)
cons = c(k_grid[k], k_grid[k1], z_grid[z], tensor_l[z, k, k1], alpha,
delta)

if cons < 0:
tensor_u[z, k, k1] = -np.inf
else:
tensor_u[z, k, k1] = u(cons, tensor_l[z, k, k1], theta, alpha,
delta)
print(count)

# e.
# given a state we will look for the combination of 12 k1s that achieves the
# highest value.
V0 = np.zeros([nzz, nkk])
V1 = np.zeros([nzz, nkk])

pf_k_idx = np.zeros([nzz, nkk])

error = 1
count = 0
# we economize on one loop my running vector operations for all k1 (k primes)
start = time.time()
while error > tol:
count += 1
for z in range(nzz):
for k in range(nkk):
#compute consumption not useful probably
#consump = c(k_grid[k], k_grid[k1], z_grid[z], tensor_l[k, k1, z])

# here we fix the state of the world by fixing a pair (z,k) and then
# we proceed computing the expected future value of each choice of
capital
# by taking the inner product of the probabilities' vector with the
# guesses we have made.
# first summand is the vector of today's utilities associated with
# all possible values of k1.
# the second summand is the vector of expected continuation values
# given different choices of k1.
# We then sum the two summands element-wise.
temp = tensor_u[z, k, :] + beta*(np.inner(V0.T, transition[z]))

#then we take the index of the k1 that maximizes the Value given the
# state (z,k).
pf_k_idx[z, k] = np.argmax(temp)
V1[z, k] = temp[np.argmax(temp)]

#save the index and the value.

error = np.linalg.norm(V1-V0)/np.linalg.norm(V0)
V0 = copy(V1)
end = time.time()
seconds = end-start
print("The iteration converged in:",seconds," seconds")

plt.title("VF")
plt.plot(k_grid, V1[0], label = "V1")
plt.plot(k_grid, V1[5], label = "V6")
plt.plot(k_grid, V1[11], label = "V12")
plt.legend()

# now we want to retrieve the values of capitals from the indices

# we want a grid of 400 policy functions for each of the 12 states
pf_k = np.zeros([nzz, nkk])

for k in range(nkk):
for z in range(nzz):
index = pf_k_idx[z, k]
pf_k[z, k] = k_grid[int(index)]

plt.title("PF capital (k)")

plt.plot(k_grid, pf_k[0], label = "V1")
plt.plot(k_grid, pf_k[5], label = "V6")
plt.plot(k_grid, pf_k[11], label = "V12")
plt.plot(k_grid, k_grid, color="black", label = "45 degree line",
linestyle="dashed")
plt.legend()

pf_l = np.zeros([nzz, nkk])

for k in range(nkk):
for z in range(nzz):
pf_l[z, k] = tensor_l[z, k, int(pf_k_idx[z, k])]

plt.title("PF labor (l)")

plt.plot(k_grid, pf_l[0], label = "z=1")
plt.plot(k_grid, pf_l[5], label = "z=6")
plt.plot(k_grid, pf_l[11], label = "z=12")
plt.legend()

# simulation part
T = 50000
ts_z_idx = np.zeros(T)
ts_z_idx[0] = 1
ts_k_idx = np.zeros(T)
ts_k_idx[0] = 1

for t in range(1,T):
prev_idx = ts_z_idx[t-1]
ts_z_idx[t] = np.random.choice(np.array(range(12)),
p=list(transition[int(prev_idx)]))

# we now simulate the path of capital for T per

for t in range(T-1):
ts_k_idx[t+1] = pf_k_idx[int(ts_z_idx[t]), int(ts_k_idx[t])]

ts_k = np.zeros(T)

for i in range(T):
index = ts_k_idx[i]
ts_k[i] = k_grid[int(index)]

plt.title("Capital path simulation")

plt.xlabel("Time")
plt.ylabel("Capital")
plt.plot(ts_k[20:], color="Blue", linewidth=0.5, label = "simulated path")
plt.plot(np.linspace(0,50000), np.linspace(k_ss,k_ss), color="Red", label="steady
state capital")
plt.legend()

plt.title("Stationary distribution of capital")

plt.xlabel("Capital")
plt.ylabel("Density")
plt.hist(ts_k[20:], bins=100, color="Blue", density=True)