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

# # Job Search Model
# 
# 
# **Randall Romero Aguilar, PhD**
# 
# This demo is based on the original Matlab demo accompanying the  <a href="https://mitpress.mit.edu/books/applied-computational-economics-and-finance">Computational Economics and Finance</a> 2001 textbook by Mario Miranda and Paul Fackler.
# 
# Original (Matlab) CompEcon file: **demdp04.m**
# 
# Running this file requires the Python version of CompEcon. This can be installed with pip by running
# 
#     !pip install compecon --upgrade
# 
# <i>Last updated: 2022-Oct-22</i>
# <hr>

# ## About
# 
# Infinitely-lived worker must decide whether to quit, if employed, or search for a job, if unemployed, given prevailing market wages.
# 
# ### States
# 
# -      w       prevailing wage
# -     i       unemployed (0) or employed (1) at beginning of period
# 
# ### Actions
# 
# -     j       idle (0) or active (i.e., work or search) (1) this period
# 
# ### Parameters
# 
# | Parameter | Meaning    |
# |-----------|-------------------------|
# | $v$       | benefit of pure leisure |
# | $\bar{w}$  | long-run mean wage |
# | $\gamma$  | wage reversion rate |
# | $p_0$      | probability of finding job |
# | $p_1$      | probability of keeping job |
# | $\sigma$   | standard deviation of wage shock |
# | $\delta$   | discount factor |
# 

# # Preliminary tasks

# In[1]:


import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from compecon import BasisSpline, DPmodel, qnwnorm


# ## FORMULATION
#   
# ### Worker's reward
# 
# The worker's reward is:
# 
# - $w$ (the prevailing wage rate), if he's employed and active (working)
# - $u=90$, if he's unemployed but active (searching)
# - $v=95$, if he's idle (quit if employed, not searching if unemployed)

# In[2]:


u     =  90
v     =  95

def reward(w, x, employed, active):
    if active:
        return w.copy() if employed else np.full_like(w, u)  # the copy is critical!!! otherwise it passes a pointer to w!!
    else:
        return np.full_like(w, v)


# ### Model dynamics
# 
# #### Stochastic Discrete State Transition Probabilities
# 
# An unemployed worker who is searching for a job has a probability $p_0=0.2$ of finding it, while an employed worker who doesn't want to quit his job has a probability $p_1 = 0.9$ of keeping it. An idle worker (someone who quits or doesn't search for a job) will definitely be unemployed next period. Thus, the transition probabilities are
# \begin{align}
#     q = \begin{bmatrix}1-p_0 &p_0\\1-p_1&p_1\end{bmatrix},&\qquad\text{if active} \\
#       = \begin{bmatrix}1 & 0\\1 &0 \end{bmatrix},&\qquad\text{if iddle}
# \end{align}

# In[3]:


p0    = 0.20
p1    = 0.90

q = np.zeros((2, 2, 2))
q[1, 0, 1] = p0
q[1, 1, 1] = p1
q[:, :, 0] = 1 - q[:, :, 1]


# #### Stochastic Continuous State Transition
# Assuming that the wage rate $w$ follows an exogenous Markov process 
# \begin{equation}
#     w_{t+1} = \bar{w} + \gamma(w_t − \bar{w}) + \epsilon_{t+1}
# \end{equation}
# 
# where $\bar{w}=100$ and $\gamma=0.4$. 

# In[4]:


wbar  = 100
ɣ = 0.40
def transition(w, x, i, j, in_, e):
    return wbar + ɣ * (w - wbar) + e


# Here, $\epsilon$ is normal $(0,\sigma^2)$ wage shock, where $\sigma=5$. We discretize this distribution with the function ```qnwnorm```.

# In[5]:


σ = 5
m = 15
e, w = qnwnorm(m, 0, σ**2)


# ### Approximation Structure
# 
# To discretize the continuous state variable, we use a cubic spline basis with $n=150$ nodes between $w_\text{min}=0$ and $w_\text{max}=200$.

# In[6]:


n = 150
wmin = 0
wmax = 200
basis = BasisSpline(n, wmin, wmax, labels=['wage'])


# ## SOLUTION
# 
# To represent the model, we create an instance of ```DPmodel```. Here, we assume a discout factor of $\delta=0.95$.

# In[7]:


model = DPmodel(basis, reward, transition,
                i =['unemployed', 'employed'],
                j = ['idle', 'active'],
                discount=0.95, e=e, w=w, q=q)


# Then, we call the method ```solve``` to solve the Bellman equation

# In[8]:


S = model.solve(show=True)
S.head()


# ### Compute and Print Critical Action Wages

# In[9]:


def critical(db):
    wcrit = np.interp(0, db['value[active]'] - db['value[idle]'], db['wage'])
    vcrit = np.interp(wcrit, db['wage'], db['value[idle]'])
    return wcrit, vcrit

wcrit0, vcrit0 = critical(S.loc['unemployed'])
print(f'Critical Search Wage = {wcrit0:5.1f}')

wcrit1, vcrit1 = critical(S.loc['employed'])
print(f'Critical Quit Wage   = {wcrit1:5.1f}')


# ### Plot Action-Contingent Value Function

# In[10]:


vv = ['value[idle]','value[active]']
fig1, (ax0, ax1) = plt.subplots(1,2,figsize=[16,4])

# UNEMPLOYED
ax0.set(title='Action-Contingent Value, Unemployed', xlabel='Wage', ylabel='Value')
ax0.plot(S.loc['unemployed',vv])
ax0.annotate(f'$w^*_0 = {wcrit0:.1f}$', [wcrit0, vcrit0], fontsize=12, va='top')
ax0.legend(['Do Not Search', 'Search'], loc='upper left', fontsize="x-small")


# EMPLOYED
ax1.set(title='Action-Contingent Value, Employed', xlabel='Wage', ylabel='Value')
ax1.plot(S.loc['employed',vv])
ax1.annotate(f'$w^*_0 = {wcrit1:.1f}$', [wcrit1, vcrit1], fontsize=12, va='top')
ax1.legend(['Quit', 'Work'], loc='upper left', fontsize="x-small")


# ### Plot Residual

# In[11]:


S['resid2'] = 100 * (S['resid'] / S['value'])

fig2, ax = plt.subplots()

ax.set(title='Bellman Equation Residual',
       xlabel='Wage',
       ylabel='Percent Residual')
ax.plot(S.loc['unemployed','resid2'])
ax.plot(S.loc['employed','resid2'])
ax.legend(model.labels.i)


# ## SIMULATION

# ### Simulate Model
# 
# We simulate the model 10000 times for a time horizon $T=40$, starting with an unemployed worker ($i=0$) at the long-term wage rate mean $\bar{w}$. To be able to reproduce these results, we set the random seed at an arbitrary value of 945.

# In[12]:


T = 40
nrep = 10000
sinit = np.full((1, nrep), wbar)
iinit = 0
data = model.simulate(T, sinit, iinit, seed=945)


# In[13]:


data.head()


# ### Print Ergodic Moments

# In[14]:


ff = '\t{:12s} = {:5.2f}'

print('\nErgodic Means')
print(ff.format('Wage', data['wage'].mean()))
print(ff.format('Employment', (data['i'] == 'employed').mean()))
print('\nErgodic Standard Deviations')
print(ff.format('Wage',data['wage'].std()))
print(ff.format('Employment', (data['i'] == 'employed').std()))


# In[15]:


ergodic = pd.DataFrame({
    'Ergodic Means' : [data['wage'].mean(), (data['i'] == 'employed').mean()],
    'Ergodic Standard Deviations': [data['wage'].std(), (data['i'] == 'employed').std()]},
    index=['Wage', 'Employment'])

ergodic.round(2)


# ### Plot Expected Discrete State Path

# In[16]:


data.head()


# In[17]:


data['ii'] = data['i'] == 'employed'
probability_of_employment = data[['ii','time']].groupby('time').mean()

fig3, ax = plt.subplots()
ax.set(title='Probability of Employment', xlabel='Period', ylabel='Probability')
ax.plot(probability_of_employment)


# ### Plot Simulated and Expected Continuous State Path

# In[31]:


subdata = data.query('_rep < 3').pivot("time", "_rep" ,'wage')

fig4, ax = plt.subplots()

ax.set(title='Simulated and Expected Wage', xlabel='Period', ylabel='Wage')
subdata.plot(ax=ax)
ax.plot(data[['time','wage']].groupby('time').mean(),'k--',label='mean')
ax.legend(fontsize='xx-small')