using Distributions
using LinearAlgebra
using Optim
using Statistics
using DataFrames
using CSVFiles
using FastGaussQuadrature
using NLSolversBase

df = load("fake_data.csv") |> DataFrame
T = convert(Int64,maximum(df.age[df[:caseid] .== 1, :]) - minimum(df.age[df[:caseid] .== 1, :]) + 1)
Ni = convert(Int64,maximum(df.caseid))

#create choice specific dataframes, with its own unique individual caseid from 1 to Ni0
df0 = df[df.school .== 0,:]
sort!(df0, (:t, :caseid))
df0.caseid = repeat(1:size(df0[df0.t .== 1,:].caseid)[1],T)
sort!(df0, (:caseid, :t))

df1 = df[df.school .== 1,:]
sort!(df1, (:t, :caseid))
df1.caseid = repeat(1:size(df1[df1.t .== 1,:].caseid)[1],T)
sort!(df1, (:caseid, :t))


#selection equation sum of X variables across t=1:T
Xs = by(df, :caseid) do x
    DataFrame(xsa = sum(x.xa), xsb = sum(x.xb), xsb2 = sum(x.xb2), xsbc = sum(x.xbc), xsbc2 = sum(x.xbc2), xsc = sum(x.xc), za = mean(x.za), zb = mean(x.zb), school = mean(x.school))
end

#Number of unique individuals in each choice dataframe
Ni0 = convert(Int64,maximum(df0.caseid))
Ni1 = convert(Int64,maximum(df1.caseid))

constants = [T, Ni, Ni0, Ni1]


include("MLE_function.jl")

#just guess actual parameters
β0 = [1., 2., -0.01, 0.5]
β1 = [.85, 3.4, -0.03, 1.]
σ0 = repeat([log(sqrt(0.1))], T)
σ1 = repeat([log(sqrt(0.1))],T)
σw = 1.
σt = log(sqrt(.4))
δz = [4., 1.1]
δt = 0.1
ρ = 0.9

theta = vcat(β0, β1, δz, δt, ρ, σ0, σ1, σw, σt)

#define nodes and weights of gauss hermite
nnodes = 20
nodes, weights = gausshermite(nnodes)
quad = [nnodes, nodes, weights]

function wrapmle1(theta)
    return mle1(theta, df0, df1, Xs, quad, constants)
end

#check that function works
@time wrapmle1(theta)

#fails
func = TwiceDifferentiable(wrapmle1, theta; autodiff = :forward)
@time mini1 = optimize(func, theta, BFGS(), Optim.Options(show_trace=true))

#goes through
func = TwiceDifferentiable(wrapmle1, theta; autodiff = :finite)
@time mini1 = optimize(func, theta, BFGS(), Optim.Options(show_trace=true))