# -*- coding: utf-8 -*-
"""
Created on Sun May  5 11:08:05 2019

@author: jbertolai
"""

import numpy as np
################################################################
def construa_objetos_parametros(nk=5):
    global k_grid, n_k, alpha, beta, z_grid, n_z, delta, pi
#    k_grid = np.array( [0.04, 0.08, 0.12, .16, .2] )
    k_grid = np.linspace(0.04, 0.20, nk)
#    k_grid = np.linspace(0.04, 0.20, 5)
    n_k = len( k_grid )
    alpha, beta = 0.3, 0.6
    ###########################
    # Inovação aqui !!
    ###########################
    delta = 1
    z_grid = np.array( [-1/5, 0, 1/5] )
    n_z = len( z_grid )
    pi = np.array( [[3/5, 2/5, 0/5],
                    [1/5, 3/5, 1/5],
                    [0/5, 2/5, 3/5] ] )
    

def construa_objetivo(V):
    ## contruindo a função objetivo
    F_obj = np.zeros( (n_k,n_k,n_z) )
    for i_z, z in enumerate(z_grid):
#        for i in range(n_k):
        for i_k, k in enumerate(k_grid):
#            for j in range(n_k):
            for i_kk, kk in enumerate(k_grid):
                c = np.exp(z)*(k**alpha) +(1-delta)*k - kk
                if c>0:
                    Ev = np.dot(pi[i_z,:],V[i_kk,:])
                    F_obj[i_k,i_kk,i_z] = np.log(c) + beta*Ev
                else:
                    F_obj[i_k,i_kk,i_z] = - np.inf
    return F_obj

def aplique_Operador_Bellman(F_obj):        
    """aplicando o operador de Bellman e armazenando a nova $v$ em v_n"""
    Tv = np.zeros( (n_k, n_z) )
    Tg = np.zeros( (n_k, n_z) , dtype=np.int)
    for i_z in range(n_z):
        for i_k in range(n_k):
            Tv[i_k,i_z] = np.max(F_obj[i_k,:,i_z])
            Tg[i_k,i_z] = np.argmax(F_obj[i_k,:,i_z])
    return Tv, Tg


def iteracao_funcoes_valor(tol_n = 1e-2, tol_it=5):
    """ITERAÇÃo ATÉ A CONVERGÊNCIA"""
    global v, v_n, ig
    v = np.zeros( (n_k, n_z) )
    v_n = [ v[:] ]

    norma, it = 1.0, 0
    while norma>tol_n and it<tol_it:

        F_obj = construa_objetivo(v)

        T_v, T_ig = aplique_Operador_Bellman(F_obj)
        
        v_n.append(T_v)   

        norma = np.max(abs(v-T_v))
        v = np.copy(T_v)
        ig = T_ig
    ##    it = it + 1
        it += 1
    ##    print('iteração, norma = ',it,',',norma)
        print('A iteração {} terminou com norma igual a {:.5f}'.format(it,norma))
    return norma, it

def funcao_principal():
    construa_objetos_parametros(nk=200)
#    iteracao_funcoes_valor()
    iteracao_funcoes_valor(tol_n = 1e-5, tol_it=500)
#    visualizar_sequencia_funcoes()
    
    
    
    
##########################################################################
import matplotlib.pyplot as plt
funcao_principal()