# Labor supply and program participation with a lifetime time limit
# Routines for solving, simulating, and estimating the model in the chapter
# "Labor Supply and Welfare Participation with Time Limits".
#
# Choices j ∈ {1,2,3,4} correspond to (H,P) ∈ {(0,0),(1,0),(0,1),(1,1)}
#   H: work, P: program participation
# State in period t: (ω,κ) where
#   ω ∈ {0,…,Ω}: cumulative periods of prior program use (Ω is the time limit)
#   κ ∈ {0,…,T-1}: accumulated work experience
# Arrays indexed as [j, ω+1, κ+1, t]

using Distributions, Random

const H_j = (0,1,0,1)
const P_j = (0,0,1,1)
choice_index(H,P) = 2P + H + 1

# default parameters: a period is a year, income is measured in dollars per week
default_pars() = (;
    T = 20,          #<- number of periods a family is potentially eligible
    Ω = 5,           #<- lifetime limit on program use
    b = 200.,        #<- maximum weekly benefit
    τ = 0.5,         #<- benefit reduction rate on earnings
    ϕ = (5.4,0.05),  #<- log wage equation: log(W) = ϕ[1] + ϕ[2]κ
    γ = 50.,         #<- weekly income from other sources
    α = 0.9,         #<- disutility of program participation (stigma, hassle)
    φ = 1.2,         #<- disutility of work
    σ = 0.5,         #<- scale of the taste shocks
    β = 0.95,        #<- discount factor
    σ_ξ = 0.2)       #<- std dev of measurement error in log wages

wage(κ,ϕ) = exp(ϕ[1] + ϕ[2]*κ)
benefit(W,H,b,τ) = max(0., b - τ*H*W)
net_income(W,H,P,b,τ) = H*W + P*benefit(W,H,b,τ)
flow_utility(W,H,P,pars) = log(pars.γ + net_income(W,H,P,pars.b,pars.τ)) - pars.α*P - pars.φ*H
available(P,ω,Ω) = P==0 || ω<Ω #<- program participation requires remaining eligibility

# fill in log choice probabilities and the value at state (ω,κ) in period t
function iterate!(logP,V,v,pars,ω,κ,t)
    (;β,ϕ,Ω,σ,T) = pars
    W = wage(κ,ϕ)
    nJ = ω<Ω ? 4 : 2 #<- choices 3 and 4 (participation) unavailable once the limit is reached
    for j in 1:nJ
        H,P = H_j[j],P_j[j]
        κ′ = min(κ + H, T-1)
        ω′ = min(ω + P, Ω)
        v[j] = flow_utility(W,H,P,pars) + β*V[ω′+1,κ′+1,t+1]
    end
    vmax = maximum(v[1:nJ])
    logsum = vmax + σ*log(sum(exp((v[j] - vmax)/σ) for j in 1:nJ))
    for j in 1:4
        logP[j,ω+1,κ+1,t] = j<=nJ ? (v[j] - logsum)/σ : -Inf
    end
    V[ω+1,κ+1,t] = logsum #<- drops σ times Euler's constant, which shifts all values equally
end

# solve by backward induction. Returns log choice probabilities and values
function solve(pars)
    (;T,Ω,α,φ,σ,γ) = pars
    R = promote_type(typeof(α),typeof(φ),typeof(σ),typeof(γ))
    V = zeros(R,Ω+1,T,T+1) #<- values in period T+1 are zero
    logP = zeros(R,4,Ω+1,T,T)
    v = zeros(R,4)
    for t in reverse(1:T), κ in 0:t-1, ω in 0:min(t-1,Ω) #<- only reachable states
        iterate!(logP,V,v,pars,ω,κ,t)
    end
    return logP,V
end

# Simulate N families from t=1 with ω=κ=0, using log choice probabilities logP.
# Participation is only possible while ω<Ω (from pars): if logP allows participation
# at an ineligible state, those choices are removed and probabilities renormalized.
# This lets us simulate the "mechanical" effect of a time limit. The same seed delivers
# the same draws, so different policies can be compared using common random numbers.
function simulate(logP,pars,N; seed = 1234)
    (;T,Ω,ϕ,σ_ξ) = pars
    rng = MersenneTwister(seed)
    u = rand(rng,N,T) #<- draws for choices
    ξ = σ_ξ .* randn(rng,N,T) #<- measurement error in wages
    J = zeros(Int,N,T)
    ω = zeros(Int,N,T)
    κ = zeros(Int,N,T)
    logWo = fill(NaN,N,T) #<- wages only observed for workers
    p = zeros(4)
    for n in 1:N
        for t in 1:T
            for j in 1:4
                p[j] = available(P_j[j],ω[n,t],Ω) ? exp(logP[j,ω[n,t]+1,κ[n,t]+1,t]) : 0.
            end
            p ./= sum(p)
            j = min(findfirst(cumsum(p) .> u[n,t]),4)
            J[n,t] = j
            H,P = H_j[j],P_j[j]
            if H==1
                logWo[n,t] = log(wage(κ[n,t],ϕ)) + ξ[n,t]
            end
            if t<T
                κ[n,t+1] = κ[n,t] + H
                ω[n,t+1] = ω[n,t] + P
            end
        end
    end
    return (;J,ω,κ,logWo)
end

# OLS estimates of the wage equation using observed wages of workers
function estimate_wages(data)
    (;κ,logWo) = data
    work = .!isnan.(logWo)
    X = [ones(sum(work)) κ[work]]
    return X \ logWo[work]
end

function log_likelihood(logP,data)
    (;J,ω,κ) = data
    ll = zero(eltype(logP))
    N,T = size(J)
    for n in 1:N, t in 1:T
        ll += logP[J[n,t],ω[n,t]+1,κ[n,t]+1,t]
    end
    return ll
end

# preference parameters x = (α, φ, log σ, log γ), other parameters fixed at pars
function update_pars(x,pars)
    return (;pars..., α = x[1], φ = x[2], σ = exp(x[3]), γ = exp(x[4]))
end

function log_likelihood(x::AbstractVector,data,pars)
    logP,_ = solve(update_pars(x,pars))
    return log_likelihood(logP,data)
end

# summary statistics by period: shares working, participating, and at the limit
function profiles(data,pars)
    (;J,ω) = data
    work = vec(sum([H_j[j] for j in J]; dims=1)) ./ size(J,1)
    part = vec(sum([P_j[j] for j in J]; dims=1)) ./ size(J,1)
    at_limit = vec(sum(ω .>= pars.Ω; dims=1)) ./ size(J,1)
    return (;work,part,at_limit)
end