#!/usr/bin/env python
# coding: utf-8

# # Productividad Específica y Decisión de Aceptación Endógena

# El modelo esta caracterizado completamente por tres ecuaciones. La primera es la que caracteriza la productividad de reserva (PR):
# $$y_{R} =b+\frac{m(\theta )\beta }{r+\lambda }\int_{y_{R}}^{\overline{y}}\left( y-y_{R}\right) f(y)dy$$
# 
# La segunda es la condición de creación de empleo (CCE):
# $$c =\frac{m(\theta )}{\theta }\frac{1-\beta }{r+\lambda }\int_{y_{R}}^{\overline{y}}(y-y_{R})f(y)dy$$
# 
# La tercera es la curva de Beveridge:
# $$m\left( \theta \right) \left( 1-F\left( y_{R}\right) \right) u=\lambda\left( 1-u\right)$$
# 
# La parametrización que usamos es la siguiente: (1) $m(\theta)= A \theta ^\alpha$ y (2) $\log x \sim N(\mu, \sigma)$.
# 
# El algoritmo de solución es el siguiente:
# 
# 1. Conjeturar $\theta_0$.
# 2. Resolver por $y_R$, el punto fijo de PR, dado $\theta_0$.
# 3. Resolver $\theta_1$ usando CCE y $y_R$.
# 4. Si $\theta_1$ = $\theta_0$ tenemos la solución, caso contrario retornar a 1 con $\theta_1$ como nueva conjetura.

# In[1]:


# Importamos los módulos necesarios
import numpy as np
import scipy.stats as stats
from scipy import integrate
from scipy.optimize import root


# Definimos la ecuación de Bellman para $y_R$:

# In[2]:


def TyR(yr, θ, prms):
    b, λ, A, α, β, r, c, F = prms
    
    intf = lambda y: (y - yr)*F.pdf(y)
    
    yr1 = b + ((A*(θ**α)*β)/(λ+r))*integrate.quad(intf,yr,np.inf)[0]
    
    return yr1


# Definimos la ecuación de creación de empleo:

# In[3]:


def JCC(θ, yr, prms):
    b, λ, A, α, β, r, c, F = prms
    
    intf = lambda y: (y - yr)*F.pdf(y)
    
    out = c - (A*(θ**α)*(1 - β))/(θ*(r+λ))*integrate.quad(intf,yr,np.inf)[0]
    
    return out


# Definimos el procedimiento que resuelve el modelo:

# In[4]:


def SolveModel(prms, tol = 1e-6, step=0.5):
    b, λ, A, α, β, r, c, F = prms
    
    yr0, θ0 = 1, 1
    
    diftheta = 1
    
    while diftheta>tol:
           
        difyr = 1
        
        print('**** Solving for yr given theta')
    
        while difyr>tol:
            yr1 = TyR(yr0, θ0, prms)
            difyr = np.abs(yr1-yr0)
            print('**** convergence: ', difyr)
            yr0 = step*yr1 + (1-step)*yr0
        
        yrout = yr1
        
        print('Solving for theta')
        
        θ1 = root(lambda x: JCC(x, yrout, prms), θ0).x[0]
        
        diftheta = np.abs(θ1-θ0)
        print('convergence: ',diftheta)
        θ0 = step*θ1 + (1-step)*θ0
    
    θout = θ1
    
    uout = λ / ((A*(θout**α))*(1-F.cdf(yrout)) + λ)
    vout = θout*uout
    
    return yrout, θout, uout, vout


# Resolvemos el modelo:

# In[5]:


# Parametros
b = 1
λ = 0.05
A = 1
α = 0.5
β = 0.5
r = 0.1
μ = 0.8
σ = 0.5
c = 1.0
F = stats.lognorm(s = σ, loc = 0, scale = np.exp(μ))    

parameters = [b, λ, A, α, β, r, c, F]        

# Solución
yR, θ, u, v = SolveModel(parameters)


# In[6]:


print('Productividad de reserva = ', yR)
print('Tightness                = ', θ)
print('Tasa de desempleo        = ', u)
print('Tasa de vacancias        = ', v)