# Heckman's sample selection estimated using maximum likelihood

using Distributions
using LinearAlgebra
using GLM
using StatFiles
using DataFrames
using Optim
using NLSolversBase
using ForwardDiff
using BenchmarkTools
using Distributions

# read the data
women = DataFrame(load("/Users/maciej/git/dydaktyka/cda-2019/data-raw/womenwk.dta"))
size(women)

# build models
women.selection = ismissing.(women.wage) .== false
women_sel =  women[women.selection,:]
women_sel.wage_new = Float64.(women_sel.wage)
m1 = glm(@formula(selection ~ married + children + education +age),women, Bernoulli(), ProbitLink());
m2 = lm(@formula(wage_new ~ education + age), women_sel);

θ = vcat(coef(m1), coef(m2)[:], 0.5, std(m2.mf.df.wage_new))'
xs = m1.mm.m
xo  = m2.mm.m
sel = women.selection
y = m2.mf.df.wage_new


function heckman(θ, xs, xo, sel, y)
      ncol_xs = size(xs,2)
      ncol_xo =  size(xo, 2)

      βs  = θ[1:ncol_xs]
      βo = θ[(ncol_xs+1):(ncol_xs+ncol_xo)]
      ρ = θ[end-1]
      σ = θ[end]

      xsβs_nosel = xs[ .!sel,:]*βs
      xsβs_sel = xs[sel,:]*βs
      xoβo = xo * βo


      global ll_s = 0
      for i in 1:length(xsβs_nosel)
            global ll_s += log(cdf(Normal(), -xsβs_nosel[i]))
      end

      global ll_o = 0
      for i in 1:length(xoβo)
            global ll_o += log(cdf(Normal(), (xsβs_sel[i]+ρ/σ*(y[i]-xoβo[i])) / √(1-ρ^2)))  - 0.5( (y[i] - xoβo[i])/σ )^2 - log(√(2π)σ)
      end
      l = ll_s+ll_o
      -l
end

heckman(θ, xs, xo, sel, y)


res = optimize(θ -> heckman(θ, xs, xo, sel, y), θ;  autodiff = :forward);

function heckit(θ, xs, xo, sel, y)
      res = optimize(θ -> heckman(θ, xs, xo, sel, y), θ; autodiff = :forward)
      θ̂ = Optim.minimizer(res)
      ∇h(θ) = ForwardDiff.gradient(θ -> heckman(θ, xs, xo, sel, y), θ̂);
      ∇h²(θ) = ForwardDiff.hessian(θ -> heckman(θ, xs, xo, sel, y), θ̂);
      se = sqrt.(diag(inv(∇h²(θ))))
      return (θ̂[:], se)
end


par, se = heckit(θ, xs, xo, sel, y);
tab_est = DataFrame(est = par, se_est = se)
tab_est.t = tab_est.est./tab_est.se_est

## fastforward
od = OnceDifferentiable(x -> heckman(θ, xs, xo, sel, y), θ; autodiff = :forward)

### comparison
@btime optimize(od, θ, NelderMead())
@btime optimize(od, θ, BFGS());
@btime optimize(θ -> heckman(θ, xs, xo, sel, y), θ; autodiff = :forward);