# Female labor force participation model
# Routines for solving, simulating, and estimating a finite-horizon model of the
# participation decisions of married women, in which work experience raises wages.
# State in period t: experience h ∈ {0,…,t-1}, husband's income (index iy), young child k ∈ {0,1}
# Value and choice probability arrays are indexed as [h+1, iy, k+1, t]

using Distributions, LinearAlgebra, Random

const γₑ = Base.MathConstants.eulergamma #<- Euler's constant

# Discretize x′ = ρx + η, η ~ N(0,σ²) on K points.
# Returns the grid and a row-stochastic matrix with F[i,i′] = P(x′ = grid[i′] | x = grid[i])
function tauchen(ρ,σ,K; m = 1.5)
    sd = σ / sqrt(1-ρ^2)
    grid = collect(range(-m*sd,m*sd,length=K))
    Δ = grid[2] - grid[1]
    F = zeros(K,K)
    for i in 1:K, i′ in 1:K
        μ = ρ*grid[i]
        lo = i′==1 ? -Inf : (grid[i′] - Δ/2 - μ)/σ
        hi = i′==K ? Inf : (grid[i′] + Δ/2 - μ)/σ
        F[i,i′] = cdf(Normal(),hi) - cdf(Normal(),lo)
    end
    return grid,F
end

# The environment: everything except preferences and the discount factor
function flfp_environment(; T = 30, Ky = 5, μy = log(5.), ρy = 0.9, σy = 0.15,
                            γ = (;γ0 = log(1.5), γ1 = 0.08, γ2 = 0.0015),
                            birth = t -> max(0.25 - 0.012*(t-1), 0.), stay = 0.8,
                            subsidy = t -> 0.)
    grid,Fy = tauchen(ρy,σy,Ky)
    ygrid = exp.(μy .+ grid)
    return (;T, Ky, ygrid, Fy, γ, birth, stay, subsidy)
end

wage(h,γ) = exp(γ.γ0 + γ.γ1*h - γ.γ2*h^2)

# transition matrix for the young child indicator between t and t+1: Fk[k+1,k′+1]
function child_transition(t,env)
    (;birth,stay) = env
    b = birth(t)
    return [1-b b; 1-stay stay]
end

# flow payoffs (before the shock) for each choice. A wage subsidy s(t) adds s(t)×wage to income.
u1(h,y,k,t,θ,env) = log(y + (1 + env.subsidy(t))*wage(h,env.γ))
u0(h,y,k,t,θ,env) = log(y) + θ.α0 + θ.α1*k

# Solve by backward induction. θ = (;α0,α1,σ). Returns log choice probabilities of working
# and the integrated value function. V[:,:,:,T+1] = 0.
function solve_flfp(θ,env,β)
    (;T,Ky,ygrid,Fy) = env
    (;σ) = θ
    R = promote_type(typeof(θ.α0),typeof(θ.α1),typeof(σ),typeof(β))
    V = zeros(R,T+1,Ky,2,T+1)
    logp1 = zeros(R,T,Ky,2,T)
    logp0 = zeros(R,T,Ky,2,T)
    for t in T:-1:1
        Fk = child_transition(t,env)
        for k in 0:1, iy in 1:Ky, h in 0:t-1
            # expected continuation values given next period's experience
            EV0 = zero(R); EV1 = zero(R)
            for k′ in 0:1, iy′ in 1:Ky
                f = Fy[iy,iy′] * Fk[k+1,k′+1]
                EV0 += f * V[h+1,iy′,k′+1,t+1]
                EV1 += f * V[h+2,iy′,k′+1,t+1]
            end
            y = ygrid[iy]
            v0 = u0(h,y,k,t,θ,env) + β*EV0
            v1 = u1(h,y,k,t,θ,env) + β*EV1
            vmax = max(v0,v1)
            logsum = vmax + σ*log(exp((v0-vmax)/σ) + exp((v1-vmax)/σ))
            V[h+1,iy,k+1,t] = σ*γₑ + logsum
            logp1[h+1,iy,k+1,t] = (v1 - logsum)/σ
            logp0[h+1,iy,k+1,t] = (v0 - logsum)/σ
        end
    end
    return (;logp1,logp0,V)
end

# stationary distribution of a row-stochastic matrix
function stationary(F)
    π∞ = real(eigvecs(Matrix(F'))[:,end])
    return π∞ ./ sum(π∞)
end

# Simulate N women over the full life-cycle. All begin with zero experience.
# Observed log wages include measurement error ξ ~ N(0,σξ²) and are NaN when not working.
function simulate_flfp(logp1,env,N; σξ = 0.3, seed = nothing)
    if !isnothing(seed)
        Random.seed!(seed)
    end
    (;T,Ky,Fy,γ) = env
    π∞ = stationary(Fy)
    H = zeros(Int,T,N)
    Y = zeros(Int,T,N)
    K = zeros(Int,T,N)
    D = zeros(Int,T,N)
    logW = fill(NaN,T,N)
    for n in 1:N
        h = 0
        iy = rand(Categorical(π∞))
        k = rand() < 0.3
        for t in 1:T
            d = rand() < exp(logp1[h+1,iy,k+1,t])
            H[t,n] = h; Y[t,n] = iy; K[t,n] = k; D[t,n] = d
            if d==1
                logW[t,n] = log(wage(h,γ)) + σξ*randn()
            end
            h += d
            iy = rand(Categorical(Fy[iy,:]))
            k = rand() < child_transition(t,env)[k+1,2]
        end
    end
    return (;H,Y,K,D,logW)
end

# OLS estimates of the wage equation using the sample of workers
function estimate_wages(data)
    (;H,D,logW) = data
    idx = findall(D.==1)
    X = [ones(length(idx)) H[idx] -H[idx].^2]
    b = X \ logW[idx]
    return (;γ0 = b[1], γ1 = b[2], γ2 = b[3])
end

# log-likelihood of choices given an array of log choice probabilities
function choice_loglikelihood(logp1,logp0,data)
    (;H,Y,K,D) = data
    T = size(D,1)
    ll = zero(eltype(logp1))
    for n in axes(D,2), t in 1:T
        i = (H[t,n]+1,Y[t,n],K[t,n]+1,t)
        ll += D[t,n]==1 ? logp1[i...] : logp0[i...]
    end
    return ll
end

# partial log-likelihood with x = (α0,α1,log σ)
function log_likelihood(x,data,env,β)
    θ = (;α0 = x[1], α1 = x[2], σ = exp(x[3]))
    (;logp1,logp0) = solve_flfp(θ,env,β)
    return choice_loglikelihood(logp1,logp0,data)
end

# Exact distribution over states implied by the model. Returns participation and mean experience
# in each period, starting from h=0, the stationary distribution of y, and P(k=1) = 0.3.
function life_cycle_profile(logp1,env)
    (;T,Ky,Fy) = env
    μ = zeros(T+1,Ky,2) #<- μ[h+1,iy,k+1]: distribution over the state in period t
    μ[1,:,1] .= 0.7 .* stationary(Fy)
    μ[1,:,2] .= 0.3 .* stationary(Fy)
    part = zeros(T)
    exper = zeros(T)
    for t in 1:T
        Fk = child_transition(t,env)
        μ′ = zeros(T+1,Ky,2)
        for k in 0:1, iy in 1:Ky, h in 0:t-1
            m = μ[h+1,iy,k+1]
            p = exp(logp1[h+1,iy,k+1,t])
            part[t] += m*p
            exper[t] += m*h
            for k′ in 0:1, iy′ in 1:Ky
                f = Fy[iy,iy′]*Fk[k+1,k′+1]
                μ′[h+1,iy′,k′+1] += m*(1-p)*f
                μ′[h+2,iy′,k′+1] += m*p*f
            end
        end
        μ = μ′
    end
    return (;part,exper)
end