text:
""" Generic functions on spectra that are to be overloaded for each instrument.  """

# Declare functions that should be specialized for each instrument here, so they can be imported into their instrument's module.
"""Returns minimum allowable order for the specified instrument"""
function min_order end

"""Returns maximum allowable order for the specified instrument"""
function max_order end

"""Returns minimum allowable pixel index for the specified instrument and order index"""
function min_pixel_in_order end

"""Returns maximum allowable pixel index for the specified instrument and order index"""
function max_pixel_in_order end

"""Returns minimum allowable pixel index for the specified instrument and any order"""
function min_pixel end

"""Returns maximum allowable pixel index for the specified instrument and any order"""
function max_pixel end

"""Returns range of order indices to be used by default for the specified instrument"""
function orders_to_use_default(inst::AbstractInstrument)
    get_inst_module(inst).orders_to_use_default(inst)
end

"""Returns minimum pixel index to be used by default for the specified instrument"""
function min_col_default end

"""Returns maximum pixel index to be read from header into metadata by default for the specified instrument"""
function max_col_default end

"""Returns ranges of bad columns"""
function bad_col_ranges end

"""Returns integer range of min and maximum pixels to use for specified instrument and order."""
function get_pixel_range end

"""Returns array of symbols to be stored from metadata in header by default for the specified instrument"""
function metadata_symbols_default end

"""Returns array of strings to be read from header into metadata by default for the specified instrument"""
function metadata_strings_default end

"""Returns the maximum range of wavelengths (in solar system barycenter frame) to be used for given spectrum"""
function get_λ_range end

""" `filter_line_list`( df, inst )
Returns a dataframe filtered to only include lines appropriate for the specified instrument.
Inputs:
- df: Dataframe containing a column with key `:lambda` for wavelengths to be considered.
- inst: instrument whose properties will be used in choosing which lines to keep.
Optional Inputs:
- λmin: Overide the instrument's default minimum wavelength
- λmax: Overide the instrument's default maximum wavelength
"""
function filter_line_list end

"""Returns the module containing code specific to the specified instrument"""
function get_inst_module end

"""Returns the default velocity width (in m/s) of a tophat CCF mask for the specified instrument"""
function default_ccf_mask_v_width end

# Trait-based functions that provide defaults (can be overwritten by instrument-specific versions)
""" Returns range of all allowable orders for the specified instrument"""
function orders_all end

orders_all(inst::AbstractInstrument2D) = min_order(inst):max_order(inst)

""" Returns range of all allowable pixels for the specified instrument"""
function pixels_all end

pixels_all(inst::AbstractInstrument2D) = min_pixel(inst):max_pixel(inst)
pixels_all(inst::AbstractInstrument1D) = min_pixel(inst):max_pixel(inst)

""" Returns maximum number of pixels in a spectrum for the specified instrument"""
function max_pixels_in_spectra end

max_pixels_in_spectra(inst::AbstractInstrument1D) = length(pixels_all(inst))
max_pixels_in_spectra(inst::AbstractInstrument2D) = (max_order(inst)-min_order(inst)+1) * (max_pixel_in_order(inst)-min_pixel_in_order(inst)+1)

""" Returns maximum number of pixels in a chunk from the specified instrument"""
min_pixels_in_chunk(inst::AbstractInstrument) = 6

function make_clean_line_list_from_tellurics end
#export make_clean_line_list_from_tellurics

function choose_obs_idx_for_init_guess(df::DataFrame, inst::AbstractInstrument)
   @assert size(df,1) >= 1
   @warn "Using first observation as reference, since choose_obs_idx_for_init_guess not specialized for instrument."
   return 1
end

#export orders_to_use_default
#export filter_line_list, get_inst_module, default_ccf_mask_v_width
#export orders_all, pixels_all, max_pixels_in_spectra, min_pixels_in_chunk

max_stars_repo_name: RvSpectML/RvSpectMLTemplate

max_issues_repo_name: RvSpectML/RvSpectMLTemplate

max_forks_repo_name: RvSpectML/RvSpectMLTemplate


text:
# NOTE: Adapted from the implementation found in Julia package
# `DifferentialEquations.jl`

using Parameters

@with_kw struct Vern7Tableau
    c₂::Float64 = 0.005
    c₃::Float64 = 0.10888888888888888
    c₄::Float64 = 0.16333333333333333
    c₅::Float64 = 0.4555
    c₆::Float64 = 0.6095094489978381
    c₇::Float64 = 0.884
    c₈::Float64 = 0.925
    a₂₁::Float64 = 0.005
    a₃₁::Float64 = -1.07679012345679
    a₃₂::Float64 = 1.185679012345679
    a₄₁::Float64 = 0.04083333333333333
    a₄₃::Float64 = 0.1225
    a₅₁::Float64 = 0.6389139236255726
    a₅₃::Float64 = -2.455672638223657
    a₅₄::Float64 = 2.272258714598084
    a₆₁::Float64 = -2.6615773750187572
    a₆₃::Float64 = 10.804513886456137
    a₆₄::Float64 = -8.3539146573962
    a₆₅::Float64 = 0.820487594956657
    a₇₁::Float64 = 6.067741434696772
    a₇₃::Float64 = -24.711273635911088
    a₇₄::Float64 = 20.427517930788895
    a₇₅::Float64 = -1.9061579788166472
    a₇₆::Float64 = 1.006172249242068
    a₈₁::Float64 = 12.054670076253203
    a₈₃::Float64 = -49.75478495046899
    a₈₄::Float64 = 41.142888638604674
    a₈₅::Float64 = -4.461760149974004
    a₈₆::Float64 = 2.042334822239175
    a₈₇::Float64 = -0.09834843665406107
    a₉₁::Float64 = 10.138146522881808
    a₉₃::Float64 = -42.6411360317175
    a₉₄::Float64 = 35.76384003992257
    a₉₅::Float64 = -4.3480228403929075
    a₉₆::Float64 = 2.0098622683770357
    a₉₇::Float64 = 0.3487490460338272
    a₉₈::Float64 = -0.27143900510483127
    a₁₀₁::Float64 = -45.030072034298676
    a₁₀₃::Float64 = 187.3272437654589
    a₁₀₄::Float64 = -154.02882369350186
    a₁₀₅::Float64 = 18.56465306347536
    a₁₀₆::Float64 = -7.141809679295079
    a₁₀₇::Float64 = 1.3088085781613787
    b₁::Float64 = 0.04715561848627222
    b₄::Float64 = 0.25750564298434153
    b₅::Float64 = 0.26216653977412624
    b₆::Float64 = 0.15216092656738558
    b₇::Float64 = 0.4939969170032485
    b₈::Float64 = -0.29430311714032503
    b₉::Float64 = 0.08131747232495111
    b̃₁::Float64 = 0.002547011879931045
    b̃₄::Float64 = -0.00965839487279575
    b̃₅::Float64 = 0.04206470975639691
    b̃₆::Float64 = -0.0666822437469301
    b̃₇::Float64 = 0.2650097464621281
    b̃₈::Float64 = -0.29430311714032503
    b̃₉::Float64 = 0.08131747232495111
    b̃₁₀::Float64 = -0.02029518466335628
end

function vern7(f, t, y, dt, P, tableau)
    (@unpack c₂,c₃,c₄,c₅,c₆,c₇,c₈,a₂₁,a₃₁,a₃₂,a₄₁,a₄₃,a₅₁,a₅₃,a₅₄,a₆₁,a₆₃,a₆₄,
             a₆₅,a₇₁,a₇₃,a₇₄,a₇₅,a₇₆,a₈₁,a₈₃,a₈₄,a₈₅,a₈₆,a₈₇,a₉₁,a₉₃,a₉₄,a₉₅,
             a₉₆,a₉₇,a₉₈,b₁,b₄,b₅,b₆,b₇,b₈,b₉ = tableau)
    k1 = f(t, y, P)
    k2 = f(t + c₂*dt, y + dt*a₂₁*k1, P)
    k3 = f(t + c₃*dt, y + dt*(a₃₁*k1 + a₃₂*k2), P)
    k4 = f(t + c₄*dt, y + dt*(a₄₁*k1 +        + a₄₃*k3), P)
    k5 = f(t + c₅*dt, y + dt*(a₅₁*k1 +        + a₅₃*k3 + a₅₄*k4), P)
    k6 = f(t + c₆*dt, y + dt*(a₆₁*k1 +        + a₆₃*k3 + a₆₄*k4 + a₆₅*k5), P)
    k7 = f(t + c₇*dt, y + dt*(a₇₁*k1 +        + a₇₃*k3 + a₇₄*k4 + a₇₅*k5
                              + a₇₆*k6), P)
    k8 = f(t + c₈*dt, y + dt*(a₈₁*k1 +        + a₈₃*k3 + a₈₄*k4 + a₈₅*k5
                              + a₈₆*k6 + a₈₇*k7), P)
    k9 = f(t + dt, y + dt*(a₉₁*k1 +        + a₉₃*k3 + a₉₄*k4 + a₉₅*k5
                           + a₉₆*k6 + a₉₇*k7 + a₉₈*k8), P)
    y + dt*(b₁*k1 + b₄*k4 + b₅*k5 + b₆*k6 + b₇*k7 + b₈*k8 + b₉*k9)
end

struct Vern7 <: ODESolverType end

#=
        FOUR components
=#
function update(::Vern7, fs, t, A, B, C, D, dt, P, tableau)
    (@unpack c₂,c₃,c₄,c₅,c₆,c₇,c₈,a₂₁,a₃₁,a₃₂,a₄₁,a₄₃,a₅₁,a₅₃,a₅₄,a₆₁,a₆₃,a₆₄,
             a₆₅,a₇₁,a₇₃,a₇₄,a₇₅,a₇₆,a₈₁,a₈₃,a₈₄,a₈₅,a₈₆,a₈₇,a₉₁,a₉₃,a₉₄,a₉₅,
             a₉₆,a₉₇,a₉₈,b₁,b₄,b₅,b₆,b₇,b₈,b₉ = tableau)
    kA1 = update(fs[1], t, A, B, C, D, P)
    kB1 = update(fs[2], t, A, B, C, D, P)
    kC1 = update(fs[3], t, A, B, C, D, P)
    kD1 = update(fs[4], t, A, B, C, D, P)

    Ai = A + dt*a₂₁*kA1
    Bi = B + dt*a₂₁*kB1
    Ci = C + dt*a₂₁*kC1
    Di = D + dt*a₂₁*kD1

    kA2 = update(fs[1], t + c₂*dt, Ai, Bi, Ci, Di, P)
    kB2 = update(fs[2], t + c₂*dt, Ai, Bi, Ci, Di, P)
    kC2 = update(fs[3], t + c₂*dt, Ai, Bi, Ci, Di, P)
    kD2 = update(fs[4], t + c₂*dt, Ai, Bi, Ci, Di, P)

    Ai = A + dt*(a₃₁*kA1 + a₃₂*kA2)
    Bi = B + dt*(a₃₁*kB1 + a₃₂*kB2)
    Ci = C + dt*(a₃₁*kC1 + a₃₂*kC2)
    Di = D + dt*(a₃₁*kD1 + a₃₂*kD2)

    kA3 = update(fs[1], t + c₃*dt, Ai, Bi, Ci, Di, P)
    kB3 = update(fs[2], t + c₃*dt, Ai, Bi, Ci, Di, P)
    kC3 = update(fs[3], t + c₃*dt, Ai, Bi, Ci, Di, P)
    kD3 = update(fs[4], t + c₃*dt, Ai, Bi, Ci, Di, P)

    Ai = A + dt*(a₄₁*kA1 +         + a₄₃*kA3)
    Bi = B + dt*(a₄₁*kB1 +         + a₄₃*kB3)
    Ci = C + dt*(a₄₁*kC1 +         + a₄₃*kC3)
    Di = D + dt*(a₄₁*kD1 +         + a₄₃*kD3)

    kA4 = update(fs[1], t + c₄*dt, Ai, Bi, Ci, Di, P)
    kB4 = update(fs[2], t + c₄*dt, Ai, Bi, Ci, Di, P)
    kC4 = update(fs[3], t + c₄*dt, Ai, Bi, Ci, Di, P)
    kD4 = update(fs[4], t + c₄*dt, Ai, Bi, Ci, Di, P)

    Ai = A + dt*(a₅₁*kA1 +         + a₅₃*kA3 + a₅₄*kA4)
    Bi = B + dt*(a₅₁*kB1 +         + a₅₃*kB3 + a₅₄*kB4)
    Ci = C + dt*(a₅₁*kC1 +         + a₅₃*kC3 + a₅₄*kC4)
    Di = D + dt*(a₅₁*kD1 +         + a₅₃*kD3 + a₅₄*kD4)

    kA5 = update(fs[1], t + c₅*dt, Ai, Bi, Ci, Di, P)
    kB5 = update(fs[2], t + c₅*dt, Ai, Bi, Ci, Di, P)
    kC5 = update(fs[3], t + c₅*dt, Ai, Bi, Ci, Di, P)
    kD5 = update(fs[4], t + c₅*dt, Ai, Bi, Ci, Di, P)

    Ai = A + dt*(a₆₁*kA1 +         + a₆₃*kA3 + a₆₄*kA4 + a₆₅*kA5)
    Bi = B + dt*(a₆₁*kB1 +         + a₆₃*kB3 + a₆₄*kB4 + a₆₅*kB5)
    Ci = C + dt*(a₆₁*kC1 +         + a₆₃*kC3 + a₆₄*kC4 + a₆₅*kC5)
    Di = D + dt*(a₆₁*kD1 +         + a₆₃*kD3 + a₆₄*kD4 + a₆₅*kD5)

    kA6 = update(fs[1], t + c₆*dt, Ai, Bi, Ci, Di, P)
    kB6 = update(fs[2], t + c₆*dt, Ai, Bi, Ci, Di, P)
    kC6 = update(fs[3], t + c₆*dt, Ai, Bi, Ci, Di, P)
    kD6 = update(fs[4], t + c₆*dt, Ai, Bi, Ci, Di, P)

    Ai = A + dt*(a₇₁*kA1 +         + a₇₃*kA3 + a₇₄*kA4 + a₇₅*kA5 + a₇₆*kA6)
    Bi = B + dt*(a₇₁*kB1 +         + a₇₃*kB3 + a₇₄*kB4 + a₇₅*kB5 + a₇₆*kB6)
    Ci = C + dt*(a₇₁*kC1 +         + a₇₃*kC3 + a₇₄*kC4 + a₇₅*kC5 + a₇₆*kC6)
    Di = D + dt*(a₇₁*kD1 +         + a₇₃*kD3 + a₇₄*kD4 + a₇₅*kD5 + a₇₆*kD6)

    kA7 = update(fs[1], t + c₇*dt, Ai, Bi, Ci, Di, P)
    kB7 = update(fs[2], t + c₇*dt, Ai, Bi, Ci, Di, P)
    kC7 = update(fs[3], t + c₇*dt, Ai, Bi, Ci, Di, P)
    kD7 = update(fs[4], t + c₇*dt, Ai, Bi, Ci, Di, P)

    Ai = A + dt*(a₈₁*kA1 +         + a₈₃*kA3 + a₈₄*kA4 + a₈₅*kA5 + a₈₆*kA6
                 + a₈₇*kA7)
    Bi = B + dt*(a₈₁*kB1 +         + a₈₃*kB3 + a₈₄*kB4 + a₈₅*kB5 + a₈₆*kB6
                 + a₈₇*kB7)
    Ci = C + dt*(a₈₁*kC1 +         + a₈₃*kC3 + a₈₄*kC4 + a₈₅*kC5 + a₈₆*kC6
                 + a₈₇*kC7)
    Di = D + dt*(a₈₁*kD1 +         + a₈₃*kD3 + a₈₄*kD4 + a₈₅*kD5 + a₈₆*kD6
                 + a₈₇*kD7)

    kA8 = update(fs[1], t + c₈*dt, Ai, Bi, Ci, Di, P)
    kB8 = update(fs[2], t + c₈*dt, Ai, Bi, Ci, Di, P)
    kC8 = update(fs[3], t + c₈*dt, Ai, Bi, Ci, Di, P)
    kD8 = update(fs[4], t + c₈*dt, Ai, Bi, Ci, Di, P)

    Ai = A + dt*(a₉₁*kA1 +         + a₉₃*kA3 + a₉₄*kA4 + a₉₅*kA5 + a₉₆*kA6
                 + a₉₇*kA7 + a₉₈*kA8)
    Bi = B + dt*(a₉₁*kB1 +         + a₉₃*kB3 + a₉₄*kB4 + a₉₅*kB5 + a₉₆*kB6
                 + a₉₇*kB7 + a₉₈*kB8)
    Ci = C + dt*(a₉₁*kC1 +         + a₉₃*kC3 + a₉₄*kC4 + a₉₅*kC5 + a₉₆*kC6
                 + a₉₇*kC7 + a₉₈*kC8)
    Di = D + dt*(a₉₁*kD1 +         + a₉₃*kD3 + a₉₄*kD4 + a₉₅*kD5 + a₉₆*kD6
                 + a₉₇*kD7 + a₉₈*kD8)

    kA9 = update(fs[1], t + dt, Ai, Bi, Ci, Di, P)
    kB9 = update(fs[2], t + dt, Ai, Bi, Ci, Di, P)
    kC9 = update(fs[3], t + dt, Ai, Bi, Ci, Di, P)
    kD9 = update(fs[4], t + dt, Ai, Bi, Ci, Di, P)

    (A + dt*(b₁*kA1 + b₄*kA4 + b₅*kA5 + b₆*kA6 + b₇*kA7 + b₈*kA8 + b₉*kA9),
     B + dt*(b₁*kB1 + b₄*kB4 + b₅*kB5 + b₆*kB6 + b₇*kB7 + b₈*kB8 + b₉*kB9),
     C + dt*(b₁*kC1 + b₄*kC4 + b₅*kC5 + b₆*kC6 + b₇*kC7 + b₈*kC8 + b₉*kC9),
     D + dt*(b₁*kD1 + b₄*kD4 + b₅*kD5 + b₆*kD6 + b₇*kD7 + b₈*kD8 + b₉*kD9))
end

#=
        THREE components
=#
function update(::Vern7, fs, t, A, B, C, dt, P, tableau)
    (@unpack c₂,c₃,c₄,c₅,c₆,c₇,c₈,a₂₁,a₃₁,a₃₂,a₄₁,a₄₃,a₅₁,a₅₃,a₅₄,a₆₁,a₆₃,a₆₄,
             a₆₅,a₇₁,a₇₃,a₇₄,a₇₅,a₇₆,a₈₁,a₈₃,a₈₄,a₈₅,a₈₆,a₈₇,a₉₁,a₉₃,a₉₄,a₉₅,
             a₉₆,a₉₇,a₉₈,b₁,b₄,b₅,b₆,b₇,b₈,b₉ = tableau)
    kA1 = update(fs[1], t, A, B, C, P)
    kB1 = update(fs[2], t, A, B, C, P)
    kC1 = update(fs[3], t, A, B, C, P)

    Ai = A + dt*a₂₁*kA1
    Bi = B + dt*a₂₁*kB1
    Ci = C + dt*a₂₁*kC1

    kA2 = update(fs[1], t + c₂*dt, Ai, Bi, Ci, P)
    kB2 = update(fs[2], t + c₂*dt, Ai, Bi, Ci, P)
    kC2 = update(fs[3], t + c₂*dt, Ai, Bi, Ci, P)

    Ai = A + dt*(a₃₁*kA1 + a₃₂*kA2)
    Bi = B + dt*(a₃₁*kB1 + a₃₂*kB2)
    Ci = C + dt*(a₃₁*kC1 + a₃₂*kC2)

    kA3 = update(fs[1], t + c₃*dt, Ai, Bi, Ci, P)
    kB3 = update(fs[2], t + c₃*dt, Ai, Bi, Ci, P)
    kC3 = update(fs[3], t + c₃*dt, Ai, Bi, Ci, P)

    Ai = A + dt*(a₄₁*kA1 +         + a₄₃*kA3)
    Bi = B + dt*(a₄₁*kB1 +         + a₄₃*kB3)
    Ci = C + dt*(a₄₁*kC1 +         + a₄₃*kC3)

    kA4 = update(fs[1], t + c₄*dt, Ai, Bi, Ci, P)
    kB4 = update(fs[2], t + c₄*dt, Ai, Bi, Ci, P)
    kC4 = update(fs[3], t + c₄*dt, Ai, Bi, Ci, P)

    Ai = A + dt*(a₅₁*kA1 +         + a₅₃*kA3 + a₅₄*kA4)
    Bi = B + dt*(a₅₁*kB1 +         + a₅₃*kB3 + a₅₄*kB4)
    Ci = C + dt*(a₅₁*kC1 +         + a₅₃*kC3 + a₅₄*kC4)

    kA5 = update(fs[1], t + c₅*dt, Ai, Bi, Ci, P)
    kB5 = update(fs[2], t + c₅*dt, Ai, Bi, Ci, P)
    kC5 = update(fs[3], t + c₅*dt, Ai, Bi, Ci, P)

    Ai = A + dt*(a₆₁*kA1 +         + a₆₃*kA3 + a₆₄*kA4 + a₆₅*kA5)
    Bi = B + dt*(a₆₁*kB1 +         + a₆₃*kB3 + a₆₄*kB4 + a₆₅*kB5)
    Ci = C + dt*(a₆₁*kC1 +         + a₆₃*kC3 + a₆₄*kC4 + a₆₅*kC5)

    kA6 = update(fs[1], t + c₆*dt, Ai, Bi, Ci, P)
    kB6 = update(fs[2], t + c₆*dt, Ai, Bi, Ci, P)
    kC6 = update(fs[3], t + c₆*dt, Ai, Bi, Ci, P)

    Ai = A + dt*(a₇₁*kA1 +         + a₇₃*kA3 + a₇₄*kA4 + a₇₅*kA5 + a₇₆*kA6)
    Bi = B + dt*(a₇₁*kB1 +         + a₇₃*kB3 + a₇₄*kB4 + a₇₅*kB5 + a₇₆*kB6)
    Ci = C + dt*(a₇₁*kC1 +         + a₇₃*kC3 + a₇₄*kC4 + a₇₅*kC5 + a₇₆*kC6)

    kA7 = update(fs[1], t + c₇*dt, Ai, Bi, Ci, P)
    kB7 = update(fs[2], t + c₇*dt, Ai, Bi, Ci, P)
    kC7 = update(fs[3], t + c₇*dt, Ai, Bi, Ci, P)

    Ai = A + dt*(a₈₁*kA1 +         + a₈₃*kA3 + a₈₄*kA4 + a₈₅*kA5 + a₈₆*kA6
                 + a₈₇*kA7)
    Bi = B + dt*(a₈₁*kB1 +         + a₈₃*kB3 + a₈₄*kB4 + a₈₅*kB5 + a₈₆*kB6
                 + a₈₇*kB7)
    Ci = C + dt*(a₈₁*kC1 +         + a₈₃*kC3 + a₈₄*kC4 + a₈₅*kC5 + a₈₆*kC6
                 + a₈₇*kC7)

    kA8 = update(fs[1], t + c₈*dt, Ai, Bi, Ci, P)
    kB8 = update(fs[2], t + c₈*dt, Ai, Bi, Ci, P)
    kC8 = update(fs[3], t + c₈*dt, Ai, Bi, Ci, P)

    Ai = A + dt*(a₉₁*kA1 +         + a₉₃*kA3 + a₉₄*kA4 + a₉₅*kA5 + a₉₆*kA6
                 + a₉₇*kA7 + a₉₈*kA8)
    Bi = B + dt*(a₉₁*kB1 +         + a₉₃*kB3 + a₉₄*kB4 + a₉₅*kB5 + a₉₆*kB6
                 + a₉₇*kB7 + a₉₈*kB8)
    Ci = C + dt*(a₉₁*kC1 +         + a₉₃*kC3 + a₉₄*kC4 + a₉₅*kC5 + a₉₆*kC6
                 + a₉₇*kC7 + a₉₈*kC8)

    kA9 = update(fs[1], t + dt, Ai, Bi, Ci, P)
    kB9 = update(fs[2], t + dt, Ai, Bi, Ci, P)
    kC9 = update(fs[3], t + dt, Ai, Bi, Ci, P)

    (A + dt*(b₁*kA1 + b₄*kA4 + b₅*kA5 + b₆*kA6 + b₇*kA7 + b₈*kA8 + b₉*kA9),
     B + dt*(b₁*kB1 + b₄*kB4 + b₅*kB5 + b₆*kB6 + b₇*kB7 + b₈*kB8 + b₉*kB9),
     C + dt*(b₁*kC1 + b₄*kC4 + b₅*kC5 + b₆*kC6 + b₇*kC7 + b₈*kC8 + b₉*kC9))
end

#=
        TWO components
=#
function update(::Vern7, fs, t, A, B, dt, P, tableau)
    (@unpack c₂,c₃,c₄,c₅,c₆,c₇,c₈,a₂₁,a₃₁,a₃₂,a₄₁,a₄₃,a₅₁,a₅₃,a₅₄,a₆₁,a₆₃,a₆₄,
             a₆₅,a₇₁,a₇₃,a₇₄,a₇₅,a₇₆,a₈₁,a₈₃,a₈₄,a₈₅,a₈₆,a₈₇,a₉₁,a₉₃,a₉₄,a₉₅,
             a₉₆,a₉₇,a₉₈,b₁,b₄,b₅,b₆,b₇,b₈,b₉ = tableau)
    kA1 = update(fs[1], t, A, B, P)
    kB1 = update(fs[2], t, A, B, P)

    Ai = A + dt*a₂₁*kA1
    Bi = B + dt*a₂₁*kB1

    kA2 = update(fs[1], t + c₂*dt, Ai, Bi, P)
    kB2 = update(fs[2], t + c₂*dt, Ai, Bi, P)

    Ai = A + dt*(a₃₁*kA1 + a₃₂*kA2)
    Bi = B + dt*(a₃₁*kB1 + a₃₂*kB2)

    kA3 = update(fs[1], t + c₃*dt, Ai, Bi, P)
    kB3 = update(fs[2], t + c₃*dt, Ai, Bi, P)

    Ai = A + dt*(a₄₁*kA1 +         + a₄₃*kA3)
    Bi = B + dt*(a₄₁*kB1 +         + a₄₃*kB3)

    kA4 = update(fs[1], t + c₄*dt, Ai, Bi, P)
    kB4 = update(fs[2], t + c₄*dt, Ai, Bi, P)

    Ai = A + dt*(a₅₁*kA1 +         + a₅₃*kA3 + a₅₄*kA4)
    Bi = B + dt*(a₅₁*kB1 +         + a₅₃*kB3 + a₅₄*kB4)

    kA5 = update(fs[1], t + c₅*dt, Ai, Bi, P)
    kB5 = update(fs[2], t + c₅*dt, Ai, Bi, P)

    Ai = A + dt*(a₆₁*kA1 +         + a₆₃*kA3 + a₆₄*kA4 + a₆₅*kA5)
    Bi = B + dt*(a₆₁*kB1 +         + a₆₃*kB3 + a₆₄*kB4 + a₆₅*kB5)

    kA6 = update(fs[1], t + c₆*dt, Ai, Bi, P)
    kB6 = update(fs[2], t + c₆*dt, Ai, Bi, P)

    Ai = A + dt*(a₇₁*kA1 +         + a₇₃*kA3 + a₇₄*kA4 + a₇₅*kA5 + a₇₆*kA6)
    Bi = B + dt*(a₇₁*kB1 +         + a₇₃*kB3 + a₇₄*kB4 + a₇₅*kB5 + a₇₆*kB6)

    kA7 = update(fs[1], t + c₇*dt, Ai, Bi, P)
    kB7 = update(fs[2], t + c₇*dt, Ai, Bi, P)

    Ai = A + dt*(a₈₁*kA1 +         + a₈₃*kA3 + a₈₄*kA4 + a₈₅*kA5 + a₈₆*kA6
                 + a₈₇*kA7)
    Bi = B + dt*(a₈₁*kB1 +         + a₈₃*kB3 + a₈₄*kB4 + a₈₅*kB5 + a₈₆*kB6
                 + a₈₇*kB7)

    kA8 = update(fs[1], t + c₈*dt, Ai, Bi, P)
    kB8 = update(fs[2], t + c₈*dt, Ai, Bi, P)

    Ai = A + dt*(a₉₁*kA1 +         + a₉₃*kA3 + a₉₄*kA4 + a₉₅*kA5 + a₉₆*kA6
                 + a₉₇*kA7 + a₉₈*kA8)
    Bi = B + dt*(a₉₁*kB1 +         + a₉₃*kB3 + a₉₄*kB4 + a₉₅*kB5 + a₉₆*kB6
                 + a₉₇*kB7 + a₉₈*kB8)

    kA9 = update(fs[1], t + dt, Ai, Bi, P)
    kB9 = update(fs[2], t + dt, Ai, Bi, P)

    (A + dt*(b₁*kA1 + b₄*kA4 + b₅*kA5 + b₆*kA6 + b₇*kA7 + b₈*kA8 + b₉*kA9),
     B + dt*(b₁*kB1 + b₄*kB4 + b₅*kB5 + b₆*kB6 + b₇*kB7 + b₈*kB8 + b₉*kB9))
end

"""
    createTableau(::Vern7)

Tableau of coefficients for the `Vern7` ODE solver
"""
createTableau(::Vern7) = Vern7Tableau()

ext: jl


text:
using DataFrames, FileIO, Dates, CSV

root_dir = dirname(pwd())
unzipped_files_path = *(root_dir, "/Datasets/Dataset_Bovespa/Unzipped/")
csv_file_name = (unzipped_files_path * "COTAHIST_PIVOTED" * ".csv")

#csv = open(csv_file_name,"r")

df = CSV.read(csv_file_name)
println(size(names(df),1)-1) #remove data
#header = true

liquidez = 0.8
lqdf = DataFrame(df[2:size(names(df),1)-1])
println(lqdf)


text:
using BinaryBuilder

name = "HelloCExecutable"
version = v"0.1.0"

# No sources, we're just building the testsuite
sources = [
    DirectorySource("./src"),
]

# Bash recipe for building across all platforms
script = raw"""
mkdir -p ${prefix}/bin
cc -o ${prefix}/bin/hello_world${exeext} -g -O2 hello_world.c
install_license /usr/share/licenses/MIT
"""

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
platforms = [
    Linux(:armv7l; libc=:glibc, compiler_abi=CompilerABI(cxxstring_abi=:cxx11)),
    Windows(:x86_64; compiler_abi=CompilerABI(cxxstring_abi=:cxx11)),
    Linux(:x86_64; libc=:glibc, compiler_abi=CompilerABI(cxxstring_abi=:cxx11)),
    MacOS(:x86_64; compiler_abi=CompilerABI(cxxstring_abi=:cxx11)),
]

# The products that we will ensure are always built
products = [
    ExecutableProduct("hello_world", :hello_world),
]

# Dependencies that must be installed before this package can be built
dependencies = Dependency[
]

# Build the tarballs, and possibly a `build.jl` as well.
build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies; preferred_gcc_version = v"7")

max_issues_repo_name: terasakisatoshi/HelloCExecutable

max_forks_repo_name: terasakisatoshi/HelloCExecutable


text:
#
# OptimPackNextGen.jl --
#
# Package for numerical optimization in Julia.
#
# -----------------------------------------------------------------------------
#
# This file is part of OptimPackNextGen.jl which is licensed under the MIT
# "Expat" License.
#
# Copyright (C) 2015-2018, Éric Thiébaut,
# <https://github.com/emmt/OptimPackNextGen.jl>.
#

module OptimPackNextGen

export
    conjgrad!,
    conjgrad,
    fmin,
    fzero,
    getreason,
    gqtpar!,
    gqtpar,
    nllsq!,
    nllsq,
    spg!,
    spg,
    vmlmb!,
    vmlmb

using LazyAlgebra
import LazyAlgebra: conjgrad, conjgrad!

"""
`Float` is the type of all floating point scalars, it is currently an alias to
`Cdouble` which is itself an alias to `Float64`.
"""
const Float = Cdouble

include("bounds.jl")

include("gqtpar.jl")
import .MoreSorensen: gqtpar, gqtpar!

include("linesearches.jl")
import .LineSearches: getreason

include("quasinewton.jl")
import .QuasiNewton: vmlmb!, vmlmb, EMULATE_BLMVM

include("brent.jl")
import .Brent: fmin, fzero

include("powell.jl")

include("nllsq.jl")
import .NonLinearLeastSquares: nllsq, nllsq!

include("bradi.jl")

include("step.jl")

include("spg.jl")
import .SPG: spg, spg!

end # module


text:
#=~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~=#
# Problem set 4 solutions
# Written by Tyler Ransom
# Commented by Giuseppe Grasso
#=~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~=#

using Random
using LinearAlgebra
using Statistics
using Optim
using DataFrames
using CSV
using HTTP
using GLM
using FreqTables
using ForwardDiff
using LineSearches
using Distributions
cd("/Users/peppegrass/Documents/GitHub/fall-2020/ProblemSets/PS4-mixture/") # GG: sets working directory
pwd() ## GG: prints working directory
readdir() # GG: equivalent to -ls- to see elements of working directory

include("lgwt.jl")

function allwrap()
    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # Question 1
    # Estimate a multinomial logit (with alternative-speciﬁc covariates Z) on the following data set, which is a panel
    # form of the same data as Problem Set 3. You should be able to simply re-use your code from Problem Set 3.
    # However, I would ask that you use AUTOMATIC DIFFERENTIATION to SPEED UP your estimation, and to obtain the
    # STANDARD ERRORS of your estimates.
    # Note: this took my machine about 30 minutes to estimate using random starting values.
    # You might consider using the estimated values from Question 1 of PS3.
    #:::::::::::::::::::::::::::::::::::::::::::::::::::

    url = "https://raw.githubusercontent.com/OU-PhD-Econometrics/fall-2020/master/ProblemSets/PS4-mixture/nlsw88t.csv"
    df = CSV.read(HTTP.get(url).body)
    X = [df.age df.white df.collgrad]
    Z = hcat(df.elnwage1, df.elnwage2, df.elnwage3, df.elnwage4,
             df.elnwage5, df.elnwage6, df.elnwage7, df.elnwage8)
    y = df.occ_code

    #=
    GG:
    The choice set is identical to that of Problem Set 3 and represents possible occupations and is structured as follows.

    1 Professional/Technical
    2 Managers/Administrators
    3 Sales
    4 Clerical/Unskilled
    5 Craftsmen
    6 Operatives
    7 Transport
    8 Other

    =#

    # GG: most of it is copy-pasted
    function mlogit_with_Z(theta, X, Z, y)

        alpha = theta[1:end-1]
        gamma = theta[end]
        K = size(X,2)
        J = length(unique(y))
        N = length(y)
        bigY = zeros(N,J)
        for j=1:J
            bigY[:,j] = y.==j
        end
        bigAlpha = [reshape(alpha,K,J-1) zeros(K)]

        T = promote_type(eltype(X),eltype(theta)) # GG: getting conformability right so that X and theta talk to each other
        num   = zeros(T,N,J)
        dem   = zeros(T,N)
        for j=1:J
            num[:,j] = exp.(X*bigAlpha[:,j] .+ (Z[:,j] .- Z[:,J])*gamma)
            dem .+= num[:,j]
        end

        P = num./repeat(dem,1,J)

        loglike = -sum( bigY.*log.(P) )

        return loglike
    end
    startvals = [2*rand(7*size(X,2)).-1; .1] # GG: kinda random
    # GG: solutions to previous PS
    startvals = [ .0403744; .2439942; -1.57132; .0433254; .1468556; -2.959103; .1020574; .7473086; -4.12005; .0375628; .6884899; -3.65577; .0204543; -.3584007; -4.376929; .1074636; -.5263738; -6.199197; .1168824; -.2870554; -5.322248; 1.307477]
    td = TwiceDifferentiable(theta -> mlogit_with_Z(theta, X, Z, y), startvals; autodiff = :forward) # GG: declare twice-differential object
    # run the optimizer
    theta_hat_optim_ad = optimize(td, startvals, LBFGS(), Optim.Options(g_tol = 1e-5, iterations=100_000, show_trace=true, show_every=50)) # GG: run optimizer
    theta_hat_mle_ad = theta_hat_optim_ad.minimizer
    # evaluate the Hessian at the estimates
    H  = Optim.hessian!(td, theta_hat_mle_ad) # GG: get hessian
    theta_hat_mle_ad_se = sqrt.(diag(inv(H))) # GG: get SEs
    println([theta_hat_mle_ad theta_hat_mle_ad_se]) # these standard errors match Stata # GG: see do-file and log-file in folder

    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # Question 2
    # Does the estimated coefﬁcient γˆ make more sense now than in Problem Set 3?
    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # The coefficient makes much more sense than it did before. It's large, positive and significant

    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # Question 3
    # Now we will estimate the mixed logit version of the model in Question 1, where γ is
    # distributed N(˜γ,σ_γ^2). Following the notes, the formula for the choice probabilities will be [see assignment text]
    # and the lig likelihood function will be [see assignment text]
    #
    # While this looks daunting, we can slightly modify the objective function from Question 1.
    # The ﬁrst step is to recognize that we will need to approximate the integral in the log likelihood function in (1).
    # There are many ways of doing this, but we will use something called Gauss-Legendre QUADRATURE (Another popular method of approximating the integral is by simulation)
    # We can rewrite the integral in (1) as a (weighted) discrete summation [see assignment text]:
    # where ω_r are the quadrature weights and ξ_r are the quadrature points. ˜γ and σ γ are parameters of the distribution function f(·).
    #:::::::::::::::::::::::::::::::::::::::::::::::::::

    # GG: QUADRATURE: a fancy polynomial that can trace out any sort of flexible function (TR layman definition)
    # GG: called quadrature because it uses a quadratic approximation. we can boil down our integral into some small numner of points and then treat is as discrete.
    # GG: variable of integration is ξ_r, so its values are real values on a grid point at which we evaluate the quadrature function (QUADRATURE POINTS/NODES)
    # GG: ω_r are weights, i.e. the density points (QUADRATURE WEIGHTS)
    # GG: f(.) is the normal pdf

    # GG: so in the log likelihood we change the integral with a summation and add the quadrature weights (substituting the "function" of the dγ)

    #---------------------------------------------------
    # a. Before we dive in, let’s learn how to use quadrature. In the folder where this problem set is posted on GitHub, there is a
    # ﬁle called lgwt.jl. This is a function that returns the ω’s and ξ’s for a given choice of K (number of quadrature
    # points) and bounds of integration [a,b]. Let’s practice doing quadrature using the density of the Normal distribution.
    #---------------------------------------------------

    # GG: lgwt() return the ω and the ξ, based on the number of quadrature points and the bounds of integration that you give it

    # define distribution
    d = Normal(0,1) # mean=0, standard deviation=1

    #=
    GG:
    Once we have the distribution deﬁned, we can do things like evaluate its density or probability, take draws from it, etc.
    We want to verify that ∫φ(x)dx is equal to 1 (i.e. integrating over the density of the support equals 1).
    We also want to verify that ∫xφ(x)dx is equal to µ (which for the distribution above is 0).
    When using quadrature, we should try to pick a number of points and bounds that will minimize computational resources,
    but still give us a good approximation to the integral. For Normal distributions, ±4σ will get us there.
    =#

    # GG: name of the game is take bounds of integration which will give us the lowest number of quadrature points but the most coverage of the support of the distribution so that we can have the best of both worlds

    # get quadrature nodes and weights for 7 grid points # GG: typically you want to pick an odd number of grid points; don't know exactly why, maybe it is about symmetry
    nodes, weights = lgwt(7,-4,4) # GG: Input: #points; bounds of integrations
    # GG: output is nodes (ξ) and weights (ω)

    # now compute the integral over the density and verify it's 1
    sum(weights.*pdf.(d,nodes)) # GG: evaluate pdf at nodes and compute weighted sum

    # now compute the expectation and verify it's 0
    sum(weights.*nodes.*pdf.(d,nodes)) # GG: in this case nodes is basically x

    #---------------------------------------------------
    # b. To get some more practice, I’d like you to use quadrature to compute the following integrals:
    #   1. ∫_{−5σ}^{5σ} x^2 f(x)dx where f(·) is the pdf of a N (0,2) distribution using 7 quadrature points # GG: that's the variance of the distribution; should be around 4
    #   2. The same as above, but with 10 quadrature points
    #   The above integrals are the variance of the distribution f. Comment on how well the quadrature approximates the true value.
    #---------------------------------------------------
    d = Normal(0,2)

    nodes, weights = lgwt(7,-10,10)
    sum(weights .* nodes.^2 .* pdf.(d,nodes)) # GG: 1. does not approximate really well

    nodes, weights = lgwt(10,-10,10)
    sum(weights .* nodes.^2 .* pdf.(d,nodes)) # GG: 2. much better approximation

    # it's more accurate with more quadrature points

    # GG: TRADE-OFF: approximation accuracy increases with the number of grid points, but the downside is that it becomes cumbersome
    # GG: you wanna be as parsimonious as possible when you nest it inside an optimization routine

    #---------------------------------------------------
    # c. An alternative to quadrature is MONTE CARLO INTEGRATION. Under this approach, we approximate the integral
    # of f by averaging over a function of many random numbers. Formally, we have that [see assignment text]
    # where D is the number of random draws and where each X i is drawn from a U[a,b] interval:
    #   • With D = 1,000,000, use the formula in (3) to approximate ∫ −5σ x 2 f (x)dx where f (·) is the pdf of a N (0,2) and verify that it is (very) close to 4
    #   • Do the same for ∫ −5σ x f (x)dx where f(·) and verify that it is very close to 0
    #   • Do the same for ∫ −5σ f (x)dx and verify that it is very close to 1
    #   • Comment on how well the simulated integral approximates the true value when D = 1,000 compared to when D = 1,000,000.
    #---------------------------------------------------

    #=
    GG: An alternative to quadrature https://en.wikipedia.org/wiki/Monte_Carlo_integration
    In this case, we approximate the integral by just taking a bunch of random numbers from the
    distribution of the integrand, and then average over those.
    At each draw, we sample points 𝐗_i uniformly on the bounds of integration [a,b]:
    Then we evaluate f(·) at those points and compute an average (corrected by the width of the bounds / # draws)
    [see assignment text] to undertand better
    This can be generalized to higher-order integrals, in which case draws are taken from a joint uniform distribution
    D needs to be large to get a nice approximation
    =#

    b = 10
    a = -10
    c = Uniform(a,b)
    draws = rand(c,1_000_000) # GG: taking a million draws from a uniform distribution with bounds -10 and 10

    # same distribution as in part (b) above "d"

    # estimate of variance # GG: |> is the equivalent of piping in Julia
    (b-a)*mean(draws.^2 .* pdf.(d,draws)) |> println # GG: returns 4

    # estimate of mean
    (b-a)*mean(draws .* pdf.(d,draws)) |> println # GG: returns 0

    # estimate of integral over support
    (b-a)*mean(pdf.(d,draws)) |> println # GG: this is the Monte Carlo integration of equation (3) in the PS

    # now repeat with only 1_000 draws
    draw2 = rand(c,1_000) # GG: this time just a thousand draws
    (b-a)*mean(draw2.^2 .* pdf.(d,draw2)) |> println
    (b-a)*mean(draw2 .* pdf.(d,draw2)) |> println
    (b-a)*mean(pdf.(d,draw2)) |> println

    # GG: definitely not as good with only a thousand draws (he says really really bad)

    #---------------------------------------------------
    # d. Note the similarity between quadrature and Monte Carlo. With quadrature, we approximate the integral with [see assignment text]
    # where ω_i is the quadrature weight and ξ_i the quadrature node. With Monte Carlo, we approximate the integral with [see assignment text]
    # So the “quadrature weight” in Monte Carlo integration is the same (b−a)/D at each node, and the “quadrature node” is a U[a,b] random number
    #---------------------------------------------------
    # Similarity noted
    #=
    GG: Go see text of problem set
    WEIGHTS: Quadrature -> ω_i [changes at each node/point] | Monte Carlo integration -> (b-a)/D [same at each node]
    POINTS: Quadrature -> ξ_i [chosen deterministically by quadrature function] | Monte Carlo integration -> 𝑋_i [drawn randomly]
    The reason you need to take so many draws with MCI is that you are giving all of them the same weight

    GG: quadrature smarter [but outcome of a chosen deterministic (quadratic-like) function], MCI brute force.
    Sometimes Quadrature can't handle integrals, especially crazy integrals (e.g. 50 dimensions)
    There is also "cubiture" that uses a cubic-like function
    =#

    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # Question 4
    # Try to modify (but not run!) your code from Question 1 to OPTIMIZE the likelihood function in (2)
    # (This took about 4 hours to estimate on my machine, given the multinomial logit starting values.
    # “Don’t try this at home,” as they say.)
    #:::::::::::::::::::::::::::::::::::::::::::::::::::

    # GG: Now using QUADRATURE to approximate the integral in the log likelihood function

    function mixlogit_quad_with_Z(theta, X, Z, y, R) # GG: here R is the number of quadrature points; so that it can be fine-tuned

        alpha = theta[1:end-2]
        gamma = theta[end-1]
        # GG: one additional parameter here σ; (γ is the μ of the gamma distribution; σ is the SD of the gamma distribution)
        sigma = exp(theta[end]) # makes sure sigma is never 0 or negative (which would make the Normal density undefined)
        K = size(X,2)
        J = length(unique(y))
        N = length(y)
        bigY = zeros(N,J)
        for j=1:J
            bigY[:,j] = y.==j
        end
        bigAlpha = [reshape(alpha,K,J-1) zeros(K)]

        # now write P as a function of variable of integration
        T = promote_type(eltype(X),eltype(theta))
        # function like_int() integrates the likelihood for the given value of parameter vector theta
        function like_int(T,N,J,X,Z,bigAlpha,gamma,sigma,R) # GG: typo: the bounds are actually also function of μ (i.e. γ) so γ should also be an argument
            nodes, weights = lgwt(R,-4*sigma,4*sigma) # GG: as sigma is part of the optimization, nodes and weights will change at every step!
            out = zeros(T,N)
            for r=1:R # GG: for every grid point of our quadrature...
                num   = zeros(T,N,J)
                dem   = zeros(T,N)
                P     = zeros(T,N,J)
                for j=1:J
                    num[:,j] = exp.(X*bigAlpha[:,j] .+ (Z[:,j] .- Z[:,J])*nodes[r])
                    dem .+= num[:,j]
                end
                P = num./repeat(dem,1,J)
                # GG: ... here's where we add the result of the quadrature [see formula of the original integral]
                out .+= vec( weights[r]*( prod(P.^bigY; dims=2) )*pdf(Normal(gamma,sigma),nodes[r]) ) # why do we have to use vec here? because vectors in Julia are "flat" and the vec operator flattens a one-dimensional array.
            end
            return out
        end

        intlike = like_int(T,N,J,X,Z,bigAlpha,gamma,sigma,R) # GG: getting the approximation through the quadrature
        loglike = -sum(log.(intlike)) # GG: evaluating the loglikelihood # GG: it is interesting to notice that sum() will sum over both dimensions N and T because we are using a stacker person-time vector
        return loglike
    end
    startvals = [ .0403744; .2439942; -1.57132; .0433254; .1468556; -2.959103; .1020574; .7473086; -4.12005; .0375628; .6884899; -3.65577; .0204543; -.3584007; -4.376929; .1074636; -.5263738; -6.199197; .1168824; -.2870554; -5.322248; 1.307477; 2]
    startvals = [-.0493358;-.7359427;1.161792;-.0127601;.1499798;-2.438863;.0979789;.6508201;-4.049595;.055819;.2773671;-4.34042;.0203963;-.6880248;-4.794188;.1288023;-.9974547;-7.914539;.118583;-.7813687;-5.609388;4.50982;2.2925]
    td = TwiceDifferentiable(theta -> mixlogit_quad_with_Z(theta, X, Z, y, 7), startvals; autodiff = :forward) # GG: same old stuff; here we specify 7 quadrature points
    # run the optimizer
    theta_hat_mix_optim_ad = optimize(td, startvals, LBFGS(linesearch = BackTracking()), Optim.Options(g_tol = 1e-5, iterations=2, show_trace=true, show_every=1))
    # theta_hat_mix_optim_ad = optimize(td, startvals, LBFGS(), Optim.Options(g_tol = 1e-5, iterations=100_000, show_trace=true, show_every=50))
    theta_hat_mix_mle_ad = theta_hat_mix_optim_ad.minimizer
    println(theta_hat_mix_mle_ad)
    # # evaluate the Hessian at the estimates
    # H  = Optim.hessian!(td, theta_hat_mix_optim_ad) # it bugs out at this point, and I don't now why # GG: he had some problem here
    # # theta_hat_mix_mle_ad_se = sqrt.(diag(inv(H)))
    # # println([theta_hat_mix_mle_ad theta_hat_mix_mle_ad_se]) # these standard errors match Stata

    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # Question 5
    # Try to modify (but not run!) your code from Question 1 to optimize the likelihood function in (1),
    # where the integral is approximated by Monte Carlo. Your program will take basically the same form as under
    # Quadrature, but the weights will be slightly different.
    # (This takes even longer because, instead of 7 quadrature points, we need to use many, many simulation draws.)
    #:::::::::::::::::::::::::::::::::::::::::::::::::::

    # GG: Now using MONTE CARLO INTEGRATION to approximate the integral in the log likelihood function
    # GG: We can use most of the previous exercise code; we're just gonna change the weights and the dimensionality of the number of points (we need many many more draws with MCI)

    function mixlogit_MC_with_Z(theta, X, Z, y, R)

        alpha = theta[1:end-2]
        gamma = theta[end-1]
        sigma = exp(theta[end]) # makes sure sigma is never 0 or negative (which would make the Normal density undefined)
        K = size(X,2)
        J = length(unique(y))
        N = length(y)
        bigY = zeros(N,J)
        for j=1:J
            bigY[:,j] = y.==j
        end
        bigAlpha = [reshape(alpha,K,J-1) zeros(K)]

        # now write P as a function of variable of integration
        T = promote_type(eltype(X),eltype(theta))
        function like_int(T,N,J,X,Z,bigAlpha,gamma,sigma,R)
            draws = 8*sigma.*Base.rand(R).-4*sigma # GG: same problem as above -> should also be function of gamma
            out = zeros(T,N)
            for r=1:R
                num   = zeros(T,N,J)
                dem   = zeros(T,N)
                P     = zeros(T,N,J)
                for j=1:J
                    num[:,j] = exp.(X*bigAlpha[:,j] .+ (Z[:,j] .- Z[:,J])*draws[r])
                    dem .+= num[:,j]
                end
                P = num./repeat(dem,1,J)
                out .+= vec( (8*sigma/R)*( prod(P.^bigY; dims=2) )*pdf(Normal(gamma,sigma),draws[r]) ) # why do we have to use vec here? because vectors in Julia are "flat" and the vec operator flattens a one-dimensional array.
            end
            return out
        end

        intlike = like_int(T,N,J,X,Z,bigAlpha,gamma,sigma,R)
        loglike = -sum(log.(intlike))
        return loglike
    end
    startvals = [ .0403744; .2439942; -1.57132; .0433254; .1468556; -2.959103; .1020574; .7473086; -4.12005; .0375628; .6884899; -3.65577; .0204543; -.3584007; -4.376929; .1074636; -.5263738; -6.199197; .1168824; -.2870554; -5.322248; 1.307477; 2]
    startvals = [-.0493358;-.7359427;1.161792;-.0127601;.1499798;-2.438863;.0979789;.6508201;-4.049595;.055819;.2773671;-4.34042;.0203963;-.6880248;-4.794188;.1288023;-.9974547;-7.914539;.118583;-.7813687;-5.609388;4.50982;2.2925]

    td = TwiceDifferentiable(theta -> mixlogit_MC_with_Z(theta, X, Z, y, 100), startvals; autodiff = :forward)
    # run the optimizer
    theta_hat_mix_optim_ad = optimize(td, startvals, LBFGS(), Optim.Options(g_tol = 1e-5, iterations=2, show_trace=true, show_every=1))
    # theta_hat_mix_optim_ad = optimize(td, startvals, LBFGS(), Optim.Options(g_tol = 1e-5, iterations=100_000, show_trace=true, show_every=50))
    theta_hat_mix_mle_ad = theta_hat_mix_optim_ad.minimizer
    println(theta_hat_mix_mle_ad)
    # evaluate the Hessian at the estimates
    # H  = Optim.hessian!(td, theta_hat_mix_optim_ad) # it bugs out at this point, and I don't now why (it's possibly due to me making sigma piecewise above)
    # theta_hat_mix_mle_ad_se = sqrt.(diag(inv(H)))
    # println([theta_hat_mix_mle_ad theta_hat_mix_mle_ad_se]) # these standard errors match Stata

    return nothing
end

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# Question 6
# Wrap all of your code above into a function and then call that function at the very bottom of your script.
# Make sure you add println() statements after obtaining each set of estimates so that you can read them.
#:::::::::::::::::::::::::::::::::::::::::::::::::::
allwrap()


text:
function fibonacci_search(f, a, b, n; ϵ=0.01)
    ϕ = (1+√5)/2;
    s = (1-√5) / (1+√5);
    ρ = 1 / (ϕ*(1-s^(n+1))/(1-s^n));
    d = ρ*b + (1-ρ) * a;
    yd = f(d);
    for i in 1: n-1
        if i == n-1
            c = ϵ*a + (1-ϵ)*d;
        else
            c = ρ*a + (1-ρ)*b;
        end

        yc = f(c);
        if yc < yd
            b, d, yd = b, c, yc;
        else
            a, b = b, c;
        end
        ρ = 1 / (ϕ*(1-s^(n-i+1))/(1-s^(b-i)));;
    end
    return a < b ? (a, b) : (b, a)
end


text:
using TimeZones.TZData: ARCHIVE_DIR, TZDATA_VERSION_REGEX, TZDATA_NEWS_REGEX
using TimeZones.TZData: read_news, extract, tzdata_version_dir, tzdata_version_archive
using TimeZones.TZData: active_version, active_archive

for year = ("12", "1234"), letter = ("", "z")
    version = year * letter
    @test match(TZDATA_VERSION_REGEX, version).match == version
    @test match(TZDATA_VERSION_REGEX, "tzdata$version.tar.gz").match == version
end

for year = ("1", "123", "12345"), letter = ("", "a")
    version = year * letter
    @test match(TZDATA_VERSION_REGEX, version) === nothing
    @test match(TZDATA_VERSION_REGEX, "tzdata$version.tar.gz") === nothing
end

@test match(TZDATA_NEWS_REGEX, "Release 92")[:version] == "92"
@test match(TZDATA_NEWS_REGEX, "Release 92c")[:version] == "92c"
@test match(TZDATA_NEWS_REGEX, "Release data94a")[:version] == "94a"
@test match(TZDATA_NEWS_REGEX, "Release code94c") === nothing
@test match(TZDATA_NEWS_REGEX, "Release code95f-data95j")[:version] == "95j"
@test match(TZDATA_NEWS_REGEX, "Release data1996m")[:version] == "1996m"
@test match(TZDATA_NEWS_REGEX, "Release code1996n") === nothing
@test match(TZDATA_NEWS_REGEX, "Release code1996o-data1996n")[:version] == "1996n"
@test match(TZDATA_NEWS_REGEX, "Release 1997a")[:version] == "1997a"

@test match(TZDATA_NEWS_REGEX, "Release 1") === nothing
@test match(TZDATA_NEWS_REGEX, "Release 199") === nothing
@test match(TZDATA_NEWS_REGEX, "Release 19999") === nothing

archive = joinpath(ARCHIVE_DIR, "tzdata$TZDATA_VERSION.tar.gz")

mktempdir() do temp_dir
    # Read the first tzdata version
    extract(archive, temp_dir, "NEWS")
    versions = read_news(joinpath(temp_dir, "NEWS"), 1)
    @test versions == [TZDATA_VERSION]

    # Read all tzdata versions
    # Note: Avoid testing for a complete set the NEWS file may be truncated in the future.
    year, letter = TZDATA_VERSION[1:end - 1], TZDATA_VERSION[end]
    latest_versions = map(c -> "$year$c", letter:-1:'a')

    versions = read_news(joinpath(temp_dir, "NEWS"))
    @test length(versions) > 1
    @test versions[1:length(latest_versions)] == latest_versions

    # Determine tzdata version from a directory
    @test tzdata_version_dir(temp_dir) == TZDATA_VERSION
    @test_throws ErrorException tzdata_version_dir(dirname(@__FILE__))
end

@test tzdata_version_archive(archive) == TZDATA_VERSION
@test_throws ProcessFailedException tzdata_version_archive(@__FILE__) == TZDATA_VERSION

# Active/built tzdata version
version = active_version()
@test version != "latest"  # Could happen if the logic to resolve the version fails
@test match(TZDATA_VERSION_REGEX, version) !== nothing

archive = active_archive()
@test isfile(archive)
@test dirname(archive) == ARCHIVE_DIR
@test basename(archive) == "tzdata$version.tar.gz"


text:
export Sigmoid

"""
`Sigmoid([domainType=Float64::Type,] dim_in::Tuple, γ = 1.)`

Creates the sigmoid non-linear operator with input dimensions `dim_in`.
```math
\\sigma(\\mathbf{x}) = \\frac{1}{1+e^{-\\gamma \\mathbf{x} } }
```

"""
struct Sigmoid{T,N,G<:Real} <: NonLinearOperator
	dim::NTuple{N,Int}
	gamma::G
end

function Sigmoid(DomainType::Type, DomainDim::NTuple{N,Int}, gamma::G=1.) where {N, G <: Real}
	Sigmoid{DomainType,N,G}(DomainDim,gamma)
end

Sigmoid(DomainDim::NTuple{N,Int}, gamma::G=1.) where {N,G} = Sigmoid{Float64,N,G}(DomainDim,gamma)

function mul!(y::AbstractArray{T,N}, L::Sigmoid{T,N,G}, x::AbstractArray{T,N}) where {T,N,G}
	y .= (1 .+exp.(-L.gamma.*x)).^(-1)
end

function mul!(y::AbstractArray,
              J::AdjointOperator{Jacobian{A,TT}},
              b::AbstractArray) where {T,N,G, A<: Sigmoid{T,N,G}, TT <: AbstractArray{T,N}}
    L = J.A
	y .= exp.(-L.A.gamma.*L.x)
	y ./= (1 .+y).^2
	y .= conj.(L.A.gamma.*y)
	y .*= b
end

fun_name(L::Sigmoid) = "σ"

size(L::Sigmoid) = (L.dim, L.dim)

domainType(L::Sigmoid{T,N,D}) where {T,N,D} = T
codomainType(L::Sigmoid{T,N,D}) where {T,N,D} = T


text:
# Vector fields (section 4.15 in the mgl manual)

function traj(gr::Graph, x::Array{cmgl.Float},  y::Array{cmgl.Float},
                                 ax::Array{cmgl.Float}, ay::Array{cmgl.Float};
                                 stl::String="",  opt::String="")
    cmgl.traj_xy(gr.ptr, Data(x).ptr, Data(y).ptr, Data(ax).ptr, Data(ay).ptr, stl, opt)
end

function traj(gr::Graph, x::Array{cmgl.Float},  y::Array{cmgl.Float},  z::Array{cmgl.Float},
                                 ax::Array{cmgl.Float}, ay::Array{cmgl.Float}, az::Array{cmgl.Float};
                                 stl::String="",  opt::String="")
    cmgl.traj_xyz(gr.ptr, Data(x).ptr,  Data(y).ptr,  Data(z).ptr,
                     Data(ax).ptr, Data(ay).ptr, Data(az).ptr, stl, opt)
end

function vect(gr::Graph, x::Array{cmgl.Float},  y::Array{cmgl.Float},
                                 ax::Array{cmgl.Float}, ay::Array{cmgl.Float};
                                 stl::String="",  opt::String="")
    cmgl.vect_xy(gr.ptr, Data(x).ptr, Data(y).ptr, Data(ax).ptr, Data(ay).ptr, stl, opt)
end

function vect(gr::Graph, x::Array{cmgl.Float},  y::Array{cmgl.Float},  z::Array{cmgl.Float},
                                 ax::Array{cmgl.Float}, ay::Array{cmgl.Float}, az::Array{cmgl.Float};
                                 stl::String="",  opt::String="")
    cmgl.vect_xyz(gr.ptr, Data(x).ptr,  Data(y).ptr,  Data(z).ptr,
                     Data(ax).ptr, Data(ay).ptr, Data(az).ptr, stl, opt)
end

function vect(gr::Graph, ax::Array{cmgl.Float}, ay::Array{cmgl.Float};
                                 stl::String="",  opt::String="")
    cmgl.vect_2d(gr.ptr, Data(ax).ptr, Data(ay).ptr, stl, opt)
end

function vect(gr::Graph, ax::Array{cmgl.Float}, ay::Array{cmgl.Float}, az::Array{cmgl.Float};
                                 stl::String="",  opt::String="")
    cmgl.vect_3d(gr.ptr, Data(ax).ptr, Data(ay).ptr, Data(az).ptr, stl, opt)
end

function vect3(gr::Graph, x::Array{cmgl.Float},  y::Array{cmgl.Float},  z::Array{cmgl.Float},
                                 ax::Array{cmgl.Float}, ay::Array{cmgl.Float}, az::Array{cmgl.Float},
                                 sval::cmgl.Float=-1.; stl::String="", opt::String="")
    cmgl.vect3_xyz(gr.ptr, Data(x).ptr,  Data(y).ptr,  Data(z).ptr,
                      Data(ax).ptr, Data(ay).ptr, Data(az).ptr, stl, sval, opt)
end

function vect3(gr::Graph, ax::Array{cmgl.Float}, ay::Array{cmgl.Float}, az::Array{cmgl.Float},
                                 sval::cmgl.Float=-1; stl::String="",  opt::String="")
    cmgl.vect3(gr.ptr, Data(ax).ptr, Data(ay).ptr, Data(az).ptr, stl, sval, opt)
end

function dew(gr::Graph, x::Array{cmgl.Float},  y::Array{cmgl.Float},
                                ax::Array{cmgl.Float}, ay::Array{cmgl.Float};
                                stl::String="",  opt::String="")
    cmgl.dew_xy(gr.ptr, Data(x).ptr, Data(y).ptr, Data(ax).ptr, Data(ay).ptr, stl, opt)
end

function dew(gr::Graph, ax::Array{cmgl.Float}, ay::Array{cmgl.Float};
                                stl::String="",  opt::String="")
    cmgl.dew(gr.ptr, Data(x).ptr, Data(y).ptr, Data(ax).ptr, Data(ay).ptr, stl, opt)
end

function flow(gr::Graph, x::Array{cmgl.Float},  y::Array{cmgl.Float},
                                 ax::Array{cmgl.Float}, ay::Array{cmgl.Float};
                                 stl::String="",  opt::String="")
    cmgl.flow_xy(gr.ptr, Data(x).ptr, Data(y).ptr, Data(ax).ptr, Data(ay).ptr, stl, opt)
end

function flow(gr::Graph, x::Array{cmgl.Float},  y::Array{cmgl.Float},  z::Array{cmgl.Float},
                                 ax::Array{cmgl.Float}, ay::Array{cmgl.Float}, az::Array{cmgl.Float};
                                 stl::String="",  opt::String="")
    cmgl.flow_xyz(gr.ptr, Data(x).ptr,  Data(y).ptr,  Data(z).ptr,
                     Data(ax).ptr, Data(ay).ptr, Data(az).ptr, stl, opt)
end

function flow(gr::Graph, ax::Array{cmgl.Float}, ay::Array{cmgl.Float};
                                 stl::String="",  opt::String="")
    cmgl.flow_2d(gr.ptr, Data(ax).ptr, Data(ay).ptr, stl, opt)
end

function flow(gr::Graph, ax::Array{cmgl.Float}, ay::Array{cmgl.Float}, az::Array{cmgl.Float};
                                 stl::String="",  opt::String="")
    cmgl.flow_3d(gr.ptr, Data(ax).ptr, Data(ay).ptr, Data(az).ptr, stl, opt)
end

function flow(gr::Graph, x0::cmgl.Float,        y0::cmgl.Float,
                                 x::Array{cmgl.Float},  y::Array{cmgl.Float},
                                 ax::Array{cmgl.Float}, ay::Array{cmgl.Float};
                                 stl::String="",  opt::String="")
    cmgl.flowp_xy(gr.ptr, x0, y0, Data(x).ptr, Data(y).ptr, Data(ax).ptr, Data(ay).ptr, stl, opt)
end

function flow(gr::Graph, x0::cmgl.Float,        y0::cmgl.Float,        z0::cmgl.Float,
                                 x::Array{cmgl.Float},  y::Array{cmgl.Float},  z::Array{cmgl.Float},
                                 ax::Array{cmgl.Float}, ay::Array{cmgl.Float}, az::Array{cmgl.Float};
                                 stl::String="",  opt::String="")
    cmgl.flowp_xyz(gr.ptr, x0, y0, z0, Data(x).ptr,  Data(y).ptr,  Data(z).ptr,
                      Data(ax).ptr, Data(ay).ptr, Data(az).ptr, stl, opt)
end

function flow(gr::Graph, x0::cmgl.Float,  y0::cmgl.Float,
                                 ax::Array{cmgl.Float}, ay::Array{cmgl.Float};
                                 stl::String="",  opt::String="")
    cmgl.flowp_2d(gr.ptr, x0, y0, Data(ax).ptr, Data(ay).ptr, stl, opt)
end

function flow(gr::Graph, x0::cmgl.Float,  y0::cmgl.Float,  z0::cmgl.Float,
                                 ax::Array{cmgl.Float}, ay::Array{cmgl.Float}, az::Array{cmgl.Float};
                                 stl::String="",  opt::String="")
    cmgl.flowp_3d(gr.ptr, x0, y0, z0, Data(ax).ptr, Data(ay).ptr, Data(az).ptr, stl, opt)
end

function grad(gr::Graph, phi::Array{cmgl.Float}; stl::String="", opt::String="")
    cmgl.grad(gr.ptr, Data(phi).ptr, stl, opt)
end

function grad(gr::Graph, x::Array{cmgl.Float}, y::Array{cmgl.Float}, phi::Array{cmgl.Float};
                                 stl::String="", opt::String="")
    cmgl.grad_xy(gr.ptr, Data(x).ptr, Data(y).ptr, Data(phi).ptr, stl, opt)
end

function grad(gr::Graph, x::Array{cmgl.Float}, y::Array{cmgl.Float}, z::Array{cmgl.Float}, phi::Array{cmgl.Float};
                                 stl::String="", opt::String="")
    cmgl.grad_xyz(gr.ptr, Data(x).ptr, Data(y).ptr, Data(z).ptr, Data(phi).ptr, stl, opt)
end

function pipe(gr::Graph, x::Array{cmgl.Float}, y::Array{cmgl.Float},
                                 ax::Array{cmgl.Float}, ay::Array{cmgl.Float}, r0::cmgl.Float=0.05;
                                 stl::String="", opt::String="")
    cmgl.pipe_xy(gr.ptr, Data(x).ptr, Data(y).ptr, Data(ax).ptr, Data(ay).ptr, stl, r0, opt)
end

function pipe(gr::Graph, x::Array{cmgl.Float}, y::Array{cmgl.Float}, z::Array{cmgl.Float},
                                 ax::Array{cmgl.Float}, ay::Array{cmgl.Float}, az::Array{cmgl.Float}, r0::cmgl.Float=0.05;
                                 stl::String="", opt::String="")
    cmgl.pipe_xyz(gr.ptr, Data(x).ptr,  Data(y).ptr,  Data(z).ptr,
                         Data(ax).ptr, Data(ay).ptr, Data(az).ptr, stl, r0, opt)
end

function pipe(gr::Graph, ax::Array{cmgl.Float}, ay::Array{cmgl.Float}, r0::cmgl.Float=0.05;
                                 stl::String="", opt::String="")
    cmgl.pipe_2d(gr.ptr, Data(ax).ptr, Data(ay).ptr, stl, r0, opt)
end

function pipe(gr::Graph, ax::Array{cmgl.Float}, ay::Array{cmgl.Float}, az::Array{cmgl.Float}, r0::cmgl.Float=0.05;
                                 stl::String="", opt::String="")
    cmgl.pipe_3d(gr.ptr, Data(ax).ptr, Data(ay).ptr, Data(az).ptr, stl, r0, opt)
end


text:
abstract type MaternLikelihood{T<:Real} <: RegressionLikelihood{T} end

"""
**Matern 3/2 likelihood**

Matern 3/2 likelihood for regression: ````
see [wiki page](https://en.wikipedia.org/wiki/Mat%C3%A9rn_covariance_function)

```julia
Matern3_2Likelihood(ρ::T) #ρ is the lengthscale
```

---

For the analytical solution, it is augmented via:
```math
p(y|f,\\omega) = \\mathcal{N}(y|f,\\sigma^2\\omega)
```
Where ``\\omega \\sim \\mathcal{IG}(\\frac{\\nu}{2},\\frac{\\nu}{2})`` where ``\\mathcal{IG}`` is the inverse gamma distribution
See paper [Robust Gaussian Process Regression with a Student-t Likelihood](http://www.jmlr.org/papers/volume12/jylanki11a/jylanki11a.pdf)
"""
struct Matern3_2Likelihood{T<:Real} <: MaternLikelihood{T}
    ρ::T
    c²::LatentArray{Vector{T}}
    θ::LatentArray{Vector{T}}
    function Matern3_2Likelihood{T}(ρ::T) where {T<:Real}
        new{T}(ρ)
    end
    function Matern3_2Likelihood{T}(ρ::T,c²::AbstractVector{<:AbstractVector{T}},θ::AbstractVector{<:AbstractVector{T}}) where {T<:Real}
        new{T}(ρ,c²,θ)
    end
end

function Matern3_2Likelihood(ρ::T=1.0) where {T<:Real}
    Matern3_2Likelihood{T}(ρ)
end

function init_likelihood(likelihood::Matern3_2Likelihood{T},inference::Inference{T},nLatent::Int,nSamplesUsed::Int,nFeatures::Int) where T
    if inference isa AnalyticVI || inference isa GibbsSampling
        Matern3_2Likelihood{T}(
        likelihood.ρ,
        [abs2.(T.(rand(T,nSamplesUsed))) for _ in 1:nLatent],
        [zeros(T,nSamplesUsed) for _ in 1:nLatent])
    else
        Matern3_2Likelihood{T}(likelihood.ρ)
    end
end

function pdf(l::Matern3_2Likelihood{T},y::Real,f::Real) where {T}
    u = sqrt(3)*abs(y-f)/l.ρ
    4*l.ρ/sqrt(3)*(one(T)+u)*exp(-u)
end

function logpdf(l::Matern3_2Likelihood{T},y::Real,f::Real) where {T}
    u = sqrt(3)*abs(y-f)/l.ρ
    log(4*l.ρ/sqrt(3)) + log(one(T)+u) - u
end

function Base.show(io::IO,model::Matern3_2Likelihood{T}) where T
    print(io,"Matern 3/2 likelihood")
end

function compute_proba(l::Matern3_2Likelihood{T},μ::AbstractVector{T},σ²::AbstractVector{T}) where {T<:Real}
    return μ,max.(σ²,0.0).+ 4*l.ρ^2/3
end

## Local Updates ##
function local_updates!(model::VGP{T,<:Matern3_2Likelihood,<:AnalyticVI}) where {T}
    model.likelihood.c² .= broadcast((Σ,μ,y)->Σ+abs2.(μ-y),diag.(model.Σ),model.μ,model.y)
    model.likelihood.θ .= broadcast(c²->3.0./(2.0.*sqrt.(3*c²)*model.likelihood.ρ.+2*model.likelihood.ρ^2),model.likelihood.c²)
end

function local_updates!(model::SVGP{T,<:Matern3_2Likelihood,<:AnalyticVI}) where {T}
    model.likelihood.c² .= broadcast((Σ,μ,y)->K̃ + opt_diag(κ*Σ,κ) + abs2.(κ*μ-y),diag.(model.Σ),model.K̃,model.κ,model.Σ,model.μ,model.inference.y)
    model.likelihood.θ .= broadcast(c²->3.0./(2.0.*sqrt.(3*c²)*model.likelihood.ρ.+2*model.likelihood.ρ^2),model.likelihood.c²)
end

function sample_local!(model::VGP{T,<:Matern3_2Likelihood,<:GibbsSampling}) where {T}
    model.likelihood.c² .= broadcast((μ,y,ρ)->rand.(GeneralizedInverseGaussian.(3/(2*ρ^2),2.0.*abs2.(μ-y),1.5)),model.μ,model.inference.y,model.likelihood.ρ)
    model.likelihood.θ .= broadcast(c²->c²,model.likelihood.c²)
    return nothing
end

## Global Gradients ##

@inline ∇E_μ(model::AbstractGP{T,<:Matern3_2Likelihood,<:GibbsorVI}) where {T} = 2.0.*hadamard.(model.likelihood.θ,model.inference.y)
@inline ∇E_μ(model::AbstractGP{T,<:Matern3_2Likelihood,<:GibbsorVI},i::Int) where {T} = 2.0.*hadamard(model.likelihood.θ[i],model.inference.y[i])

@inline ∇E_Σ(model::AbstractGP{T,<:Matern3_2Likelihood,<:GibbsorVI}) where {T} = model.likelihood.θ
@inline ∇E_Σ(model::AbstractGP{T,<:Matern3_2Likelihood,<:GibbsorVI},i::Int) where {T} = model.likelihood.θ[i]

## ELBO Section ##

function ELBO(model::AbstractGP{T,<:Matern3_2Likelihood,<:AnalyticVI}) where {T}
    return NaN;
    #expecLogLikelihood(model) - InverseGammaKL(model) - GaussianKL(model)
end

function expecLogLikelihood(model::VGP{T,<:Matern3_2Likelihood,<:AnalyticVI}) where {T}
    # model.likelihood.β .= broadcast((Σ,μ,y)->0.5*(Σ+abs2.(μ-y).+model.likelihood.σ^2*model.likelihood.ν),diag.(model.Σ),model.μ,model.y)
    # model.likelihood.θ .= broadcast(β->model.likelihood.α./β,model.likelihood.β)
    tot = -0.5*model.nLatent*model.nSample*(log(twoπ*model.likelihood.σ^2))
    tot += -sum(broadcast((α,β)->sum(log.(β).-digamma(α)), model.likelihood.α,model.likelihood.β))
    tot += -0.5.*sum(broadcast((θ,Σ,μ,y)->dot(θ,Σ)+dot(θ,abs2.(μ))-2.0*dot(θ,μ.*y)+dot(θ,abs2.(y)),model.likelihood.θ,diag.(model.Σ),model.μ,model.y))
    return tot
end

function expecLogLikelihood(model::SVGP{T,<:Matern3_2Likelihood,<:AnalyticVI}) where {T}
    tot = -0.5*model.nLatent*model.inference.nSamplesUsed*(log(twoπ*model.likelihood.σ^2))
    tot += -sum(broadcast(β->sum(digamma(model.likelihood.α).-log.(β)),model.likelihood.β))
    tot += -0.5.*sum(broadcast((θ,K̃,κ,Σ,κμ,y)->dot(θ,(K̃+opt_diag(κ*Σ,κ)+abs2.(κμ)-2.0*(κμ).*y+abs2.(y))),model.likelihood.θ,model.K̃,model.κ,model.Σ,model.κ.*model.μ,model.inference.y))
    return model.inference.ρ*tot
end

function InverseGammaKL(model::AbstractGP{T,<:Matern3_2Likelihood}) where {T}
    α_p = model.likelihood.ν/2; β_p= α_p*model.likelihood.σ^2
    model.inference.ρ*sum(broadcast(InverseGammaKL,model.likelihood.α,model.likelihood.β,α_p,β_p))
end

## PDF and Log PDF Gradients ## (verified gradients)

function grad_log_pdf(l::Matern3_2Likelihood{T},y::Real,f::Real) where {T<:Real}
    3.0 * (y-f) / (l.ρ*(abs(f-y)*sqrt(3)+l.ρ))
end

function hessian_log_pdf(l::Matern3_2Likelihood{T},y::Real,f::Real) where {T<:Real}
    3.0 / (l.ρ+sqrt(3)*abs(f-y))^2
end


text:
using SparseArrays
using SphericalOperators

import SphericalOperators: ord

function hamiltonian(basis::FEDVR.Basis, L::AbstractSphericalBasis,
                     ::Type{O}=SphericalOperators.LexicalOrdering;
                     v::Function=coulomb(1.0)) where {O<:SphericalOperators.Ordering}
    nᵣ = basecount(basis.grid)
    @assert nᵣ == size(L,2)
    M = prod(size(L))
    H₀ = spzeros(M,M)

    T = sparse(kinop(basis)) # One-body operator, identical for all partial waves
    rsel = 1:nᵣ

    for ℓ in eachℓ(L)
        Vℓ = potop(basis, v(ℓ)).lmap
        Hℓ = T + Vℓ
        for m in eachm(L,ℓ)
            sel = ord(L,O,ℓ,m,rsel)
            H₀[sel,sel] += Hℓ
        end
    end

    H₀
end
hamiltonian(basis::FEDVR.Basis, ℓ::Integer; kwargs...) =
    hamiltonian(basis, SphericalBasis2d(ℓ,basecount(basis.grid),ℓₘᵢₙ=ℓ); kwargs...)

function interaction_common(fun::Function, basis::FEDVR.Basis, L::AbstractSphericalBasis, component=:z,
                            ::Type{O}=SphericalOperators.LexicalOrdering) where {O<:SphericalOperators.Ordering}
    m = basecount(basis.grid)
    @assert m == size(L,2)
    M = prod(size(L))
    Hᵢ = spzeros(M,M)

    𝔞,𝔟 = fun()

    op = Dict(:z => SphericalOperators.ζ,
              :x => SphericalOperators.ξ)[component]

    materialize!(Hᵢ, op, L, 𝔞, 𝔟, O)
end

function hamiltonian_E_R(basis::FEDVR.Basis, L::AbstractSphericalBasis, component=:z,
                         ::Type{O}=SphericalOperators.LexicalOrdering) where {O<:SphericalOperators.Ordering}
    """Dipole interaction Hamiltonian in the length gauge
    Ĥᵢ(t) = 𝓔(t)⋅r, where r = [x,y,z]."""
    interaction_common(basis, L, component, O) do
        R = potop(basis, r -> r).lmap
        rℓ = ℓ -> R
        rℓ,rℓ
    end
end

function APℓ(basis::FEDVR.Basis, L::AbstractSphericalBasis, component=:z,
             ::Type{O}=SphericalOperators.LexicalOrdering) where {O<:SphericalOperators.Ordering}
    """Dipole interaction Hamiltonian (centrifugal part) in the
    velocity gauge Ĥᵢ(t) = 𝓐(t)⋅p, where p = -im*[∂x,∂y,∂z]."""
    interaction_common(basis, L, component, O) do
        R⁻¹ = potop(basis, r -> 1/r).lmap
        ℓ -> (ℓ+1)*R⁻¹,ℓ -> -ℓ*R⁻¹
    end
end

function ∂ᵣ(basis::FEDVR.Basis, L::AbstractSphericalBasis, component=:z,
            ::Type{O}=SphericalOperators.LexicalOrdering) where {O<:SphericalOperators.Ordering}
    """Dipole interaction Hamiltonian (differential part) in the
    velocity gauge Ĥᵢ(t) = 𝓐(t)⋅p, where p = -im*[∂x,∂y,∂z]."""
    interaction_common(basis, L, component, O) do
        ∂ᵣop = sparse(derop(basis, 1))
        𝔞𝔟 = ℓ -> ∂ᵣop
        𝔞𝔟,𝔞𝔟
    end
end

hamiltonian_A_P(basis::FEDVR.Basis, L::AbstractSphericalBasis,
                component=:z,
                ::Type{O}=SphericalOperators.LexicalOrdering) where {O<:SphericalOperators.Ordering} =
                    APℓ(basis, L, component, O) + ∂ᵣ(basis, L, component, O)

export hamiltonian, hamiltonian_E_R, hamiltonian_A_P


text:
"""
    EqOddsWrapper

It is a postprocessing algorithm which uses Linear Programming to optimise the constraints for Equalized Odds.
"""
mutable struct EqOddsWrapper <: DeterministicComposite
	grp::Symbol
	classifier::MLJBase.Model
end

"""
    EqOddsWrapper(classifier; grp=:class)

Instantiates EqOddsWrapper which wraps the classifier
"""
function EqOddsWrapper(classifier::MLJBase.Model; grp::Symbol=:class)
	return EqOddsWrapper(grp, classifier)
end

# Corresponds to eq_odds_optimal_mix_rates function, mix_rates are returned as fitresult
function MMI.fit(model::EqOddsWrapper, verbosity::Int,
	X, y)
	grps = X[:, model.grp]
	n = length(levels(grps)) # Number of different values for sensitive attribute

	# As equalized odds is a postprocessing algorithm, the model needs to be fitted a_
	mch = machine(model.classifier, X, y)
	fit!(mch)
	ŷ = MMI.predict(mch, X)

	if typeof(ŷ[1]) <: MLJBase.UnivariateFinite
		ŷ = MLJBase.mode.(ŷ)
	end

	ŷ = convert(Array, ŷ) # Incase ŷ is categorical array, convert to normal array to support various operations
	y = convert(Array, y)

	# Finding the probabilities of changing predictions is a Linear Programming Problem
	# JuMP and Cbc are used to for this Linear Programming Problem
	m = JuMP.Model(GLPK.Optimizer)

	@variable(m, 0<= p2p[1:n] <=1)
	@variable(m, 0<= p2n[1:n] <=1)
	@variable(m, 0<= n2p[1:n] <=1)
	@variable(m, 0<= n2n[1:n] <=1)

	@constraint(m, cons1[i=1:n], p2p[i] == 1 - p2n[i])
	@constraint(m, cons2[i=1:n], n2p[i] == 1 - n2n[i])

	a_Class = levels(grps)[1]
	a_Grp = grps .== a_Class

	a_flip = 1 .- ŷ[a_Grp]
	a_const = ŷ[a_Grp]

	a_BaseRate = mean(y[a_Grp]) # Base rate for a_ class

	ft = fair_tensor(categorical(ŷ), categorical(y), categorical(grps))

	error = 0
	for i in 1:n
		i_fpr = fpr(ft; grp=levels(grps)[i]) * p2p[i] + tnr(ft; grp=levels(grps)[i]) * n2p[i]
		i_fnr = fnr(ft; grp=levels(grps)[i]) * n2n[i] + tpr(ft; grp=levels(grps)[i]) * p2n[i]
		error += (i_fpr + i_fnr)
	end
	@objective(m, Min, error)

	a_tn = ŷ[a_Grp].==0 .& y[a_Grp].==0
	a_fn = ŷ[a_Grp].==0 .& y[a_Grp].==1
	a_tp = ŷ[a_Grp].==1 .& y[a_Grp].==1
	a_fp = ŷ[a_Grp].==1 .& y[a_Grp].==0

	# Following variables names have been changed from the implementation by Equalized Odds postprocessing algorithm
	# These variables better explain the corresponding quantity.
	# For eg. a_pp_given_p Corresponds to Predicted Positive given Negative for class named a

	a_pp_given_p = ((n2p[1] * mean(a_flip .& a_fn) + n2n[1] * mean(a_const .& a_fn)) / a_BaseRate +
				  (p2p[1] * mean(a_const .& a_tp) + p2n[1] * mean(a_flip .& a_tp)) / a_BaseRate)

	a_pn_given_n = ((p2n[1] * mean(a_flip .& a_fp) + p2p[1] * mean(a_const .& a_fp)) / (1 - a_BaseRate) +
				  (n2p[1] * mean(a_flip .& a_tn) + n2n[1] * mean(a_const .& a_tn)) / (1 - a_BaseRate))

	for i in 2:n
		i_Class = levels(grps)[i]

		i_Grp = grps .== i_Class

		i_flip = 1 .- ŷ[i_Grp]
		i_const = ŷ[i_Grp]

		i_BaseRate = mean(y[i_Grp]) # Base rate for a_ class

		i_tn = ŷ[i_Grp].==0 .& y[i_Grp].==0
		i_fn = ŷ[i_Grp].==0 .& y[i_Grp].==1
		i_tp = ŷ[i_Grp].==1 .& y[i_Grp].==1
		i_fp = ŷ[i_Grp].==1 .& y[i_Grp].==0

		i_pp_given_p = ((n2p[i] * mean(i_flip .& i_fn) + n2n[i] * mean(i_const .& i_fn)) / i_BaseRate +
					  (p2p[i] * mean(i_const .& i_tp) + p2n[i] * mean(i_flip .& i_tp)) / i_BaseRate)

		i_pn_given_n = ((p2n[i] * mean(i_flip .& i_fp) + p2p[i] * mean(i_const .& i_fp)) / (1 - i_BaseRate) +
					  (n2p[i] * mean(i_flip .& i_tn) + n2n[i] * mean(i_const .& i_tn)) / (1 - i_BaseRate))

		@constraint(m, a_pp_given_p==i_pp_given_p)
		@constraint(m, a_pn_given_n==i_pn_given_n)
	end

	optimize!(m)

	fitresult = [[JuMP.value.(p2n), JuMP.value.(n2p)], mch.fitresult]

	return fitresult, nothing, nothing
end

# Corresponds to eq_odds function which uses mix_rates to modify results
function MMI.predict(model::EqOddsWrapper, fitresult, Xnew)

	(p2n, n2p), classifier_fitresult = fitresult

	ŷ = MMI.predict(model.classifier, classifier_fitresult, Xnew)

	if typeof(ŷ[1]) <: MLJBase.UnivariateFinite
		ŷ = MLJBase.mode.(ŷ)
	end

	ŷ = convert(Array, ŷ) # Need to convert to normal array as categorical array doesn't support sub
	grps = Xnew[:, model.grp]

	n = length(levels(grps)) # Number of different values for sensitive attribute

	for i in 1:n
		Class = levels(grps)[i]
		Grp = grps .== Class

		pp_indices = shuffle(findall((grps.==Grp) .& (ŷ.==1))) # predicted positive for iᵗʰ class
		pn_indices = shuffle(findall((grps.==Class) .& (ŷ.==0))) # predicted negative for iᵗʰ class

		# Note : arrays in julia start from 1
		p2n_indices = pp_indices[1:convert(Int, floor(length(pp_indices)*p2n[i]))]
		n2p_indices = pn_indices[1:convert(Int, floor(length(pn_indices)*n2p[i]))]

		ŷ[p2n_indices] = 1 .- ŷ[p2n_indices]
		ŷ[n2p_indices] = 1 .- ŷ[n2p_indices]
	end
	return ŷ
end


text:
include("31-OnlineApp.jl")
while true
    try
        julia_main([""])
    catch(e)
        println( e)
    end
    sleep(30)
end

max_stars_repo_name: wangzhongtian/fiberAlarmAnalyseTool

max_issues_repo_name: wangzhongtian/fiberAlarmAnalyseTool

max_forks_repo_name: wangzhongtian/fiberAlarmAnalyseTool


text:
# MIT License
# Copyright (c) 2017: Xavier Gandibleux, Anthony Przybylski, Gauthier Soleilhac, and contributors.
struct _2UMFLP
  m::Int #nbCustomers
  n::Int #nbFacilities
  A1::Matrix{Float64} #Assignment costs of clients to facilities
  A2::Matrix{Float64} #Assignment costs of clients to facilities
  R1::Vector{Float64} #Running cost of facilities
  R2::Vector{Float64} #Running cost of facilities
end

Base.show(io::IO, id::_2UMFLP) = println("Bi-Objective Facility Location Problem with $(id.m) customers and $(id.n) facilities.")

set2UMFLP(m,n,A1,A2,R1,R2) = begin
    @assert size(A1,2) == length(R1)
    @assert n == size(A1, 2) == size(A2, 2) == length(R1) == length(R2)
    @assert m == size(A1, 1) == size(A2, 1)
    return _2UMFLP(m,n,A1,A2,R1,R2)
end
set2UMFLP(A1,A2,R1,R2) = set2UMFLP(size(A1,1), size(A1,2), A1, A2, R1, R2)

struct UMFLPsolver
    solve::Function
end

function vSolve(id::_2UMFLP, solver::UMFLPsolver = UMFLP_Delmee2017())
    solver.solve(id)
end

function load2UMFLP(fname)
   f = open(fname)
   nbC = parse(Int, readline(f))
   nbF = parse(Int, readline(f))

   d = readdlm(f)
   c_alloc1 = convert(Matrix{Float64}, d[1:nbC, 1:nbF])
   c_alloc2 = convert(Matrix{Float64}, d[nbC+1:2nbC, 1:nbF])
   c_loc1 = convert(Vector{Float64}, d[2nbC+1, 1:nbF])
   c_loc2 = convert(Vector{Float64}, d[2nbC+2, 1:nbF])
   return _2UMFLP(nbC, nbF, c_alloc1, c_alloc2, c_loc1, c_loc2)
end

function UMFLP_Delmee2017(; modeVerbose = false, modeUpperBound = true, modeLowerBound = true, modeImprovedLB = false, modeParam = true)::UMFLPsolver
    f = (id::_2UMFLP) -> begin
        n, m = id.n, id.m  #n facilities, m customers
        A1, A2 = vec(permutedims(id.A1, [2,1])), vec(permutedims(id.A2, [2,1]))
        R1, R2 = id.R1, id.R2
        p_z1,p_z2 = Ref{Ptr{Cdouble}}() ,Ref{Ptr{Cdouble}}()
        p_facility, p_isEdge, p_isExtremityDominated = Ref{Ptr{Cuchar}}(), Ref{Ptr{Cuchar}}(), Ref{Ptr{Cuchar}}()
        p_nbAlloc, p_customerAlloc, p_correspondingFac = Ref{Ptr{Cint}}(), Ref{Ptr{Cint}}(), Ref{Ptr{Cint}}()
        p_percentageAlloc = Ref{Ptr{Cdouble}}()
        nbSol = ccall((:solve, libUMFLPpath),
                       Int,
                       (Cint, Cint, Ref{Cdouble}, Ref{Cdouble}, Ref{Cdouble}, Ref{Cdouble}, #Inputs
                           Cuchar, Cuchar, Cuchar, Cuchar, Cuchar, Cuchar, Cuchar, Cuchar, #Parameters
                           # Outputs :
                           Ref{Ptr{Cdouble}}, Ref{Ptr{Cdouble}}, #z1, z2
                           Ref{Ptr{Cuchar}}, Ref{Ptr{Cuchar}}, Ref{Ptr{Cuchar}},#facility, isEdge, isExtremityDominated,
                           Ref{Ptr{Cint}}, Ref{Ptr{Cint}}, Ref{Ptr{Cint}}, Ref{Ptr{Cdouble}}), #nb Alloc, SparseMatrix : i,j,k

                       m, n, A1, A2, R1, R2,
                       modeVerbose, modeLowerBound, false, false, modeImprovedLB, modeUpperBound, modeParam, false,
                       p_z1, p_z2,
                       p_facility, p_isEdge, p_isExtremityDominated,
                       p_nbAlloc, p_customerAlloc, p_correspondingFac, p_percentageAlloc)
        z1 = _unsafe_wrap(Array, p_z1.x, nbSol, own=true)
        z2 = _unsafe_wrap(Array, p_z2.x, nbSol, own=true)

        facility = convert(BitArray, _unsafe_wrap(Array, p_facility.x, nbSol*n, own=true))
        facility_res = [facility[i:i+n-1] for i = 1:n:nbSol*n]
        isEdge = convert(BitArray, _unsafe_wrap(Array, p_isEdge.x, nbSol, own=true))
        isExtremityDominated = convert(BitArray, _unsafe_wrap(Array, p_isExtremityDominated.x, nbSol, own=true))
        nbAlloc = _unsafe_wrap(Array, p_nbAlloc.x, nbSol, own=true)

        totalAlloc = sum(nbAlloc)

        customerAlloc = 1 .+ _unsafe_wrap(Array, p_customerAlloc.x, totalAlloc, own=true)
        correspondingFac = 1 .+ _unsafe_wrap(Array,  p_correspondingFac.x, totalAlloc, own=true)
        percentageAlloc = _unsafe_wrap(Array,  p_percentageAlloc.x, totalAlloc, own=true)

        ind = 1

        X = SparseMatrixCSC{Float64, Int}[]

        for i= 1:nbSol
            sm = sparse(customerAlloc[ind:ind+nbAlloc[i]-1], correspondingFac[ind:ind+nbAlloc[i]-1], percentageAlloc[ind:ind+nbAlloc[i]-1])
            ind += nbAlloc[i]
            push!(X, sm)
        end
        return z1, z2, facility_res, X, isEdge, isExtremityDominated
    end

    return UMFLPsolver(f)
end

max_stars_repo_name: vOptSolver/vOptSpecific

max_issues_repo_name: vOptSolver/vOptSpecific

max_forks_repo_name: vOptSolver/vOptSpecific


text:
# Use baremodule to shave off a few KB from the serialized `.ji` file
baremodule finufft_jll
using Base
using Base: UUID
import JLLWrappers

JLLWrappers.@generate_main_file_header("finufft")
JLLWrappers.@generate_main_file("finufft", UUID("c41cd5a2-72a3-5203-9076-a500b088fc82"))
end  # module finufft_jll


text:
#==========================================================
This routine initializes an optimization problem to recover
a CNOT gate on a single qudit with 4 energy levels (and 2
guard states) and showcases the use of uncoupled controls.
The drift Hamiltonian in the rotating frame is
        H_0 = - 0.5*ξ_a(a^†a^†aa),
where a is the annihilation operator for the qudit. Here
the control Hamiltonian includes the usual symmetric and
anti-symmetric terms
        H_s = p(t)(a + a^†),    H_a = q(t)(a - a^†),
which come from the rotating frame approximation and hence
we refer to these as "coupled" controls. In
addition, we consider a flux-charge term
                    H_f = f(t) a^†a,
which we refer to as an "uncoupled" control as it is
invariant to the rotating frame approximation. The problem
parameters for this example are from Jonathan and Pranav at
UChicago:
                    ω_a =  2π × 5.0 Grad/s,
                    ξ_a =  2π × 0.2 Grad/s.
We use Bsplines with carrier waves with frequencies
0, ξ_a Grad/s.
==========================================================#
using LinearAlgebra
#using Plots
#pyplot()
#using FFTW
#using DelimitedFiles
using Printf
#using Ipopt
using Random
using SparseArrays

Base.show(io::IO, f::Float64) = @printf(io, "%20.13e", f)

import Juqbox

verbose = false
N = 4
Nguard = 2
Ntot = N + Nguard

samplerate = 64 # default number of time steps per unit time (plotting only)
casename = "cnot-lab" # base file name (used in optimize-once.jl)

# Set to false for dense matrix operations
use_sparse = true

# frequencies (in GHz, will be multiplied by 2*pi to get angular frequencies in the Hamiltonian matrix)
fa = 5.0 # 4.1
xa = 0.2

# duration
T = 12.0 # Tperiod/4

Ident = Matrix{Float64}(I, Ntot, Ntot)
utarget = Matrix{ComplexF64}(I, Ntot, N)

# CNOT target
utarget[:,4] = Ident[:,3]
utarget[:,3] = Ident[:,4]

Nosc  = 1
Nfreq = 3 # number of carrier frequencies

Random.seed!(2456)

# setup drift Hamiltonian
number = Diagonal(collect(0:Ntot-1))

H0 = 2*pi*fa*number-0.5*(2*pi)*xa* (number*number - number)

# lowering matrix
amat = Bidiagonal(zeros(Ntot),sqrt.(collect(1:Ntot-1)),:U) # standard lowering operator matrix
# raising matrix
adag = transpose(amat) # raising operator matrix

Hunc_ops=[Array(amat + adag)]
H0 = Array(H0)

Ncoupled = 0
Nunc = length(Hunc_ops)

# setup carrier frequencies
om = zeros(1,Nfreq)
use_bcarrier = true

if use_bcarrier
    @assert(Nfreq==1 || Nfreq==2 || Nfreq==3)
    if Nfreq == 2
        om[1:1,2] .= -2.0*pi*fa
    elseif Nfreq == 3
        om[:,2] .= -2.0*pi*fa
        om[:,3] .= 2.0*pi*fa
    end
end

# Note: same frequencies for each p(t) (x-drive) and q(t) (y-drive)
#println("Carrier frequencies [GHz]: ", om[:,:]./(2*pi))

# max parameter amplitude
max_unc = 2*pi*5.0

# Initial conditions
Ident = Matrix{Float64}(I, Ntot, Ntot)
U0 = Ident[1:Ntot,1:N]

# setup the initial parameter vector, either randomized or from file
startFromScratch = true # true
startFile = "flux-pcof-opt-alpha-0.5.dat"
useBarrier = true

if startFromScratch
    # D1 smaller than 3 does not work
    D1 = 30 # Number of B-spline coefficients per frequency, sin/cos and real/imag
    nCoeff = (2*Ncoupled + Nunc)*Nfreq*D1
    pcof0  = zeros(nCoeff)
    pcof0 = (rand(nCoeff) .- 0.5).*max_unc*0.1
else
    # the data on the startfile must be consistent with the setup!
    # use if you want to have initial coefficients read from file
    pcof0 = vec(readdlm(startFile))
    nCoeff = length(pcof0)
    D1 = div(nCoeff, (2*Ncoupled + Nunc)*Nfreq) # factor '2' is for sin/cos

    nCoeff = (2*Ncoupled + Nunc)*Nfreq*D1 # just to be safe if the file doesn't contain the right number of elements

#    println("*** Starting from B-spline coefficients in file: ", startFile)
end

# Estimate time step for simulation
nsteps = Juqbox.calculate_timestep(T, H0, Hunc_ops, [max_unc])

# setup the simulation parameters
params = Juqbox.objparams([N], [Nguard], T, nsteps, Uinit=U0, Utarget=utarget, Cfreq=om, Rfreq=[fa], Hconst=H0, Hunc_ops=Hunc_ops)
params.saveConvHist = true
params.nsteps *= 5

# Quiet mode for testing
params.quiet = true

#Tikhonov regularization coefficients
params.tik0 = 1e-3

# Set bounds on coefficients
minCoeff, maxCoeff = Juqbox.assign_thresholds(params,D1,[0.0],[max_unc])

# For ipopt
maxIter = 50 # optional argument
lbfgsMax = 250 # optional argument

if verbose
    println("*** Settings ***")
    println("System Hamiltonian coefficients: (fa, xa) =  ", fa, xa)
    println("Total number of states, Ntot = ", Ntot, " Total number of guard states, Nguard = ", Nguard)
    if use_bcarrier
        println("Using B-spline basis functions with carrier wave, # freq = ", Nfreq)
    else
        println("Using regular B-spline basis functions")
    end
    println("Number of coefficients per spline = ", D1, " Total number of parameters = ", nCoeff)
    println("Max parameter amplitudes: max_unc = ", max_unc)
    println("Tikhonov coefficient: tik0 = ", params.tik0)
end

# Estimate number of terms in Neumann series for time stepping (Default 3)
tol = eps(1.0); # machine precision
Juqbox.estimate_Neumann!(tol, params, Float64[], [max_unc])

wa = Juqbox.Working_Arrays(params, nCoeff)
prob = Juqbox.setup_ipopt_problem(params, wa, nCoeff, minCoeff, maxCoeff, maxIter, lbfgsMax)

# uncomment to run the gradient checker for the initial pcof
addOption( prob, "derivative_test", "first-order"); # for testing the gradient

#println("Initial coefficient vector stored in 'pcof0'")


text:
struct V2016R <: Version end

Base.show(io::IO, v::V2016R) = print(io, "v2016R beta")

"""
    v2016R()

Identifier for the 2016R beta version of the model.
Note that this is a beta version following `DICE2016R-091916ap.gms`.

# Examples
```jldoctest
julia> v2016R()
v2016R beta
```
"""
function v2016R()
    V2016R()
end

export v2016R

function options(version::V2016R;
    N::Int = 100, #Number of years to calculate (from 2015 onwards)
    tstep::Int = 5, #Years per Period
    α::Float64 = 1.45, #Elasticity of marginal utility of consumption
    ρ::Float64 = 0.015, #Initial rate of social time preference per year ρ
    γₑ::Float64 = 0.3, #Capital elasticity in production function
    pop₀::Int = 7403, #Initial world population 2015 (millions)
    popadj::Float64 = 0.134, #Growth rate to calibrate to 2050 pop projection
    popasym::Int = 11500, #Asymptotic population (millions)
    δk::Float64 = 0.1, #Depreciation rate on capital (per year)
    q₀::Float64 = 105.5, #Initial world gross output 2015 (trill 2010 USD)
    k₀::Float64 = 223.0, #Initial capital value 2015 (trill 2010 USD)
    a₀::Float64 = 5.115, #Initial level of total factor productivity
    ga₀::Float64 = 0.076, #Initial growth rate for TFP per 5 years
    δₐ::Float64 = 0.005, #Decline rate of TFP per 5 years
    gσ₁::Float64 = -0.0152, #Initial growth of sigma (continuous per year)
    δσ::Float64 = -0.001, #Decline rate of decarbonization per period
    eland₀::Float64 = 2.6, #Carbon emissions from land 2015 (GtCO2 per year)
    deland::Float64 = 0.115, #Decline rate of land emissions (per period)
    e₀::Float64 = 35.85, #Industrial emissions 2015 (GtCO2 per year)
    μ₀::Float64 = 0.03, #Initial emissions control rate for base case 2015
    mat₀::Float64 = 851.0, #Initial Concentration in atmosphere 2015 (GtC)
    mu₀::Float64 = 460.0, #Initial Concentration in upper strata 2015 (GtC)
    ml₀::Float64 = 1740.0, #Initial Concentration in lower strata 2015 (GtC)
    mateq::Float64 = 588.0, #Equilibrium concentration atmosphere  (GtC)
    mueq::Float64 = 360.0, #Equilibrium concentration in upper strata (GtC)
    mleq::Float64 = 1720.0, #Equilibrium concentration in lower strata (GtC)
    ϕ₁₂::Float64 = 0.12, #Carbon cycle transition matrix coefficient
    ϕ₂₃::Float64 = 0.007, #Carbon cycle transition matrix coefficient
    t2xco2::Float64 = 3.1, #Equilibrium temp impact (oC per doubling CO2)
    fₑₓ0::Float64 = 0.5, #2015 forcings of non-CO2 GHG (Wm-2)
    fₑₓ1::Float64 = 1.0, #2100 forcings of non-CO2 GHG (Wm-2)
    tocean₀::Float64 = 0.0068, #Initial lower stratum temp change (C from 1900)
    tatm₀::Float64 = 0.85, #Initial atmospheric temp change (C from 1900)
    ξ₁::Float64 = 0.1005, #Climate equation coefficient for upper level
    ξ₃::Float64 = 0.088, #Transfer coefficient upper to lower stratum
    ξ₄::Float64 = 0.025, #Transfer coefficient for lower level
    η::Float64 = 3.6813, #Forcings of equilibrium CO2 doubling (Wm-2)
    ψ₁₀::Float64 = 0.0, #Initial damage intercept
    ψ₁::Float64 = 0.0, #Damage intercept
    ψ₂::Float64 = 0.00236, #Damage quadratic term
    ψ₃::Float64 = 2.0, #Damage exponent
    θ₂::Float64 = 2.6, #Exponent of control cost function
    pback::Float64 = 550.0, #Cost of backstop 2010$ per tCO2 2015
    gback::Float64 = 0.025, #Initial cost decline backstop cost per period
    limμ::Float64 = 1.2, #Upper limit on control rate after 2150
    tnopol::Float64 = 45.0, #Period before which no emissions controls base
    cprice₀::Float64 = 2.0, #Initial base carbon price (2010$ per tCO2)
    gcprice::Float64 = 0.02, #Growth rate of base carbon price per year
    fosslim::Float64 = 6000.0, #Maximum cumulative extraction fossil fuels (GtC)
    scale1::Float64 = 0.0302455265681763, #Multiplicative scaling coefficient
    scale2::Float64 = -10993.704) #Additive scaling coefficient
    OptionsV2016R(N,tstep,α,ρ,γₑ,pop₀,popadj,popasym,δk,q₀,k₀,a₀,ga₀,δₐ,gσ₁,δσ,eland₀,deland,mat₀,mu₀,ml₀,mateq,mueq,mleq,ϕ₁₂,ϕ₂₃,t2xco2,fₑₓ0,fₑₓ1,tocean₀,tatm₀,ξ₁,ξ₃,ξ₄,η,ψ₁,ψ₂,ψ₃,θ₂,pback,gback,limμ,fosslim,scale1,scale2,e₀,μ₀,tnopol,cprice₀,gcprice,ψ₁₀)
end

function Base.show(io::IO, ::MIME"text/plain", opt::OptionsV2016R)
    println(io, "Options for 2016R beta version");
    println(io, "Time step");
    println(io, "N: $(opt.N), tstep: $(opt.tstep)");
    println(io, "Preferences");
    println(io, "α: $(opt.α), ρ: $(opt.ρ)");
    println(io, "Population and Technology");
    println(io, "γₑ: $(opt.γₑ), pop₀: $(opt.pop₀), popadj: $(opt.popadj), popasym: $(opt.popasym), δk: $(opt.δk)");
    println(io, "q₀: $(opt.q₀), k₀: $(opt.k₀), a₀: $(opt.a₀), ga₀: $(opt.ga₀), δₐ: $(opt.δₐ)");
    println(io, "Emissions Parameters");
    println(io, "gσ₁: $(opt.gσ₁), δσ: $(opt.δσ), eland₀: $(opt.eland₀)");
    println(io, "deland: $(opt.deland), e₀: $(opt.e₀), μ₀: $(opt.μ₀)");
    println(io, "Carbon Cycle");
    println(io, "mat₀: $(opt.mat₀), mu₀: $(opt.mu₀), ml₀: $(opt.ml₀)");
    println(io, "mateq: $(opt.mateq), mueq: $(opt.mueq), mleq: $(opt.mleq)");
    println(io, "Flow Parameters");
    println(io, "ϕ₁₂: $(opt.ϕ₁₂), ϕ₂₃: $(opt.ϕ₂₃)");
    println(io, "Climate Model Parameters");
    println(io, "t2xco2: $(opt.t2xco2), fₑₓ0: $(opt.fₑₓ0), fₑₓ1: $(opt.fₑₓ1)");
    println(io, "tocean₀: $(opt.tocean₀), tatm₀: $(opt.tatm₀), ξ₁: $(opt.ξ₁)");
    println(io, "ξ₃: $(opt.ξ₃), ξ₄: $(opt.ξ₄), η: $(opt.η)");
    println(io, "Climate Damage Parameters");
    println(io, "ψ₁₀: $(opt.ψ₁₀), ψ₁: $(opt.ψ₁), ψ₂: $(opt.ψ₂), ψ₃: $(opt.ψ₃)");
    println(io, "Abatement Cost");
    println(io, "θ₂: $(opt.θ₂), pback: $(opt.pback), gback: $(opt.gback), limμ: $(opt.limμ)");
    println(io, "tnopol: $(opt.tnopol), cprice₀: $(opt.cprice₀), gcprice: $(opt.gcprice)");
    println(io, "Fossil Fuel Availability");
    println(io, "fosslim: $(opt.fosslim)");
    println(io, "Scaling Parameters");
    print(io, "scale1: $(opt.scale1), scale2: $(opt.scale2)");
end

function generate_parameters(c::OptionsV2016R, model::JuMP.Model)
    optlrsav::Float64 = (c.δk + .004)/(c.δk + .004c.α + c.ρ)*c.γₑ; # Optimal savings rate
    ϕ₁₁::Float64 = 1 - c.ϕ₁₂; # Carbon cycle transition matrix coefficient
    ϕ₂₁::Float64 = c.ϕ₁₂*c.mateq/c.mueq; # Carbon cycle transition matrix coefficient
    ϕ₂₂::Float64 = 1 - ϕ₂₁ - c.ϕ₂₃; # Carbon cycle transition matrix coefficient
    ϕ₃₂::Float64 = c.ϕ₂₃*c.mueq/c.mleq; # Carbon cycle transition matrix coefficient
    ϕ₃₃::Float64 = 1 - ϕ₃₂; # Carbon cycle transition matrix coefficient
    σ₀::Float64 = c.e₀/(c.q₀*(1-c.μ₀)); # Carbon intensity 2010 (kgCO2 per output 2010 USD 2010)
    λ::Float64 = c.η/c.t2xco2; # Climate model parameter

    @NLparameter(model, ψ₂ == c.ψ₂);

    # Backstop price
    pbacktime = Array{Float64}(undef, c.N);
    # Growth rate of productivity from 0 to N
    gₐ = Array{Float64}(undef, c.N);
    # Emissions from deforestation
    Etree = Array{Float64}(undef, c.N);
    # Average utility social discount rate
    rr = Array{Float64}(undef, c.N);
    # Carbon price in base case
    cpricebase = Array{Float64}(undef, c.N);

    for i in 1:c.N
        pbacktime[i] = c.pback*(1-c.gback)^(i-1);
        gₐ[i] = c.ga₀*exp.(-c.δₐ*5*(i-1));
        Etree[i] = c.eland₀*(1-c.deland)^(i-1);
        rr[i] = 1/((1+c.ρ)^(c.tstep*(i-1)));
        cpricebase[i] = c.cprice₀*(1+c.gcprice)^(5*(i-1));
    end

    # Initial conditions and offset required
    # Level of population and labor
    L = Array{Float64}(undef, c.N);
    L[1] = c.pop₀;
    # Level of total factor productivity
    A = Array{Float64}(undef, c.N);
    A[1] = c.a₀;
    # Change in sigma (cumulative improvement of energy efficiency)
    gσ = Array{Float64}(undef, c.N);
    gσ[1] = c.gσ₁;
    # CO2-equivalent-emissions output ratio
    σ = Array{Float64}(undef, c.N);
    σ[1] = σ₀;

    # Cumulative from land
    cumtree = Array{Float64}(undef, c.N);
    cumtree[1] = 100.0;

    for i in 1:c.N-1
        L[i+1] = L[i]*(c.popasym/L[i])^c.popadj;
        A[i+1] = A[i]/(1-gₐ[i]);
        gσ[i+1] = gσ[i]*((1+c.δσ)^c.tstep);
        σ[i+1] = σ[i]*exp.(gσ[i]*c.tstep);
        cumtree[i+1] = cumtree[i]+Etree[i]*(5/3.666);
    end

    # Adjusted cost for backstop
    θ₁ = Array{Float64}(undef, c.N);
    # Exogenous forcing for other greenhouse gases
    fₑₓ = Array{Float64}(undef, c.N);

    for i in 1:c.N
        θ₁[i] = pbacktime[i]*σ[i]/c.θ₂/1000.0;
        fₑₓ[i] = if i <= 18
                        c.fₑₓ0+(1/17)*(c.fₑₓ1-c.fₑₓ0)*(i-1)
                    else
                        c.fₑₓ1
                    end;
    end
    ParametersV2016(ϕ₁₁,ϕ₂₁,ϕ₂₂,ϕ₃₂,ϕ₃₃,σ₀,λ,gₐ,Etree,L,A,gσ,σ,θ₁,fₑₓ,pbacktime,cpricebase,rr,optlrsav,ψ₂,cumtree)
end

#NOTE: CCATOT and the Tₐₜ upper bound is the only difference tho the 2013R models.
function model_vars(version::V2016R, model::Model, N::Int64, cca_ubound::Float64, μ_ubound::Array{Float64,1}, cprice_ubound::Array{Float64,1})
    # Variables #
    @variable(model, 0.0 <= μ[i=1:N] <= μ_ubound[i]); # Emission control rate GHGs
    @variable(model, FORC[1:N]); # Increase in radiative forcing (watts per m2 from 1900)
    @variable(model, 0.0 <= Tₐₜ[1:N] <= 12.0); # Increase temperature of atmosphere (degrees C from 1900)
    @variable(model, -1.0 <= Tₗₒ[1:N] <= 20.0); # Increase temperatureof lower oceans (degrees C from 1900)
    @variable(model, Mₐₜ[1:N] >= 10.0); # Carbon concentration increase in atmosphere (GtC from 1750)
    @variable(model, Mᵤₚ[1:N] >= 100.0); # Carbon concentration increase in shallow oceans (GtC from 1750)
    @variable(model, Mₗₒ[1:N] >= 1000.0); # Carbon concentration increase in lower oceans (GtC from 1750)
    @variable(model, E[1:N]); # Total CO2 emissions (GtCO2 per year)
    @variable(model, Eind[1:N]); # Industrial emissions (GtCO2 per year)
    @variable(model, C[1:N] >= 2.0); # Consumption (trillions 2010 US dollars per year)
    @variable(model, K[1:N] >= 1.0); # Capital stock (trillions 2010 US dollars)
    @variable(model, CPC[1:N] >= 0.01); #  Per capita consumption (thousands 2010 USD per year)
    @variable(model, I[1:N] >= 0.0); # Investment (trillions 2010 USD per year)
    @variable(model, S[1:N]); # Gross savings rate as fraction of gross world product
    @variable(model, RI[1:N]); # Real interest rate (per annum)
    @variable(model, Y[1:N] >= 0.0); # Gross world product net of abatement and damages (trillions 2010 USD per year)
    @variable(model, YGROSS[1:N] >= 0.0); # Gross world product GROSS of abatement and damages (trillions 2010 USD per year)
    @variable(model, YNET[1:N]); # Output net of damages equation (trillions 2010 USD per year)
    @variable(model, DAMAGES[1:N]); # Damages (trillions 2010 USD per year)
    @variable(model, Ω[1:N]); # Damages as fraction of gross output
    @variable(model, Λ[1:N]); # Cost of emissions reductions  (trillions 2010 USD per year)
    @variable(model, MCABATE[1:N]); # Marginal cost of abatement (2010$ per ton CO2)
    @variable(model, CCA[1:N] <= cca_ubound); # Cumulative industrial carbon emissions (GTC)
    @variable(model, CCATOT[1:N]); # Total carbon emissions (GTC)
    @variable(model, PERIODU[1:N]); # One period utility function
    @variable(model, CPRICE[i=1:N]); # Carbon price (2010$ per ton of CO2)
    @variable(model, CEMUTOTPER[1:N]); # Period utility
    @variable(model, UTILITY); # Welfare function
    VariablesV2016(μ,FORC,Tₐₜ,Tₗₒ,Mₐₜ,Mᵤₚ,Mₗₒ,E,C,K,CPC,I,S,RI,Y,YGROSS,YNET,DAMAGES,MCABATE,CCA,PERIODU,UTILITY,Eind,Ω,Λ,CPRICE,CEMUTOTPER,CCATOT)
end

#NOTE: MCABATE and CPRICE are the same in the original, can one of these be removed?...
function model_eqs(scenario::Scenario, model::Model, config::OptionsV2016R, params::ParametersV2016, vars::VariablesV2016, isMumps::Bool)
    N = config.N;
    scale = 5/3.666;

    # Equations #
    # Emissions Equation
    eeq = @constraint(model, [i=1:N], vars.E[i] == vars.Eind[i] + params.Etree[i]);
    # Industrial Emissions
    @constraint(model, [i=1:N], vars.Eind[i] == params.σ[i] * vars.YGROSS[i] * (1-vars.μ[i]));
    # Cumulative total carbon emissions
    @constraint(model, [i=1:N], vars.CCATOT[i] == vars.CCA[i] + params.cumtree[i]);
    # Radiative forcing equation
    @NLconstraint(model, [i=1:N], vars.FORC[i] == config.η * (log(vars.Mₐₜ[i]/588.0)/log(2)) + params.fₑₓ[i]);
    # Equation for damage fraction
    @NLconstraint(model, [i=1:N], vars.Ω[i] == config.ψ₁*vars.Tₐₜ[i]+params.ψ₂*vars.Tₐₜ[i]^config.ψ₃);
    # Damage equation
    @constraint(model, [i=1:N], vars.DAMAGES[i] == vars.YGROSS[i]*vars.Ω[i]);
    # Cost of exissions reductions equation
    @NLconstraint(model, [i=1:N], vars.Λ[i] == vars.YGROSS[i] * params.θ₁[i] * vars.μ[i]^config.θ₂);
    # Equation for MC abatement
    @NLconstraint(model, [i=1:N], vars.MCABATE[i] == params.pbacktime[i] * vars.μ[i]^(config.θ₂-1));
    # Carbon price equation from abatement
    @NLconstraint(model, [i=1:N], vars.CPRICE[i] == params.pbacktime[i] * vars.μ[i]^(config.θ₂-1));
    # Output gross equation
    @NLconstraint(model, [i=1:N], vars.YGROSS[i] == params.A[i]*(params.L[i]/1000.0)^(1-config.γₑ)*vars.K[i]^config.γₑ);
    # Output net of damages equation
    @constraint(model, [i=1:N], vars.YNET[i] == vars.YGROSS[i]*(1-vars.Ω[i]));
    # Output net equation
    @constraint(model, [i=1:N], vars.Y[i] == vars.YNET[i] - vars.Λ[i]);
    # Consumption equation
    cc = @constraint(model, [i=1:N], vars.C[i] == vars.Y[i] - vars.I[i]);
    # Per capita consumption definition
    @constraint(model, [i=1:N], vars.CPC[i] == 1000.0 * vars.C[i] / params.L[i]);
    # Savings rate equation
    @constraint(model, [i=1:N], vars.I[i] == vars.S[i] * vars.Y[i]);
    # Period utility
    @constraint(model, [i=1:N], vars.CEMUTOTPER[i] == vars.PERIODU[i] * params.L[i] * params.rr[i]);
    # Instantaneous utility function equation
    @NLconstraint(model, [i=1:N], vars.PERIODU[i] == ((vars.C[i]*1000.0/params.L[i])^(1-config.α)-1)/(1-config.α)-1);

    # Equations (offset) #
    # Cumulative carbon emissions
    @constraint(model, [i=1:N-1], vars.CCA[i+1] == vars.CCA[i] + vars.Eind[i]*scale);
    # Atmospheric concentration equation
    @constraint(model, [i=1:N-1], vars.Mₐₜ[i+1] == vars.Mₐₜ[i]*params.ϕ₁₁ + vars.Mᵤₚ[i]*params.ϕ₂₁ + vars.E[i]*scale);
    # Lower ocean concentration
    @constraint(model, [i=1:N-1], vars.Mₗₒ[i+1] == vars.Mₗₒ[i]*params.ϕ₃₃ + vars.Mᵤₚ[i]*config.ϕ₂₃);
    # Shallow ocean concentration
    @constraint(model, [i=1:N-1], vars.Mᵤₚ[i+1] == vars.Mₐₜ[i]*config.ϕ₁₂ + vars.Mᵤₚ[i]*params.ϕ₂₂ + vars.Mₗₒ[i]*params.ϕ₃₂);
    # Temperature-climate equation for atmosphere
    @constraint(model, [i=1:N-1], vars.Tₐₜ[i+1] == vars.Tₐₜ[i] + config.ξ₁*((vars.FORC[i+1]-params.λ*vars.Tₐₜ[i])-(config.ξ₃*(vars.Tₐₜ[i]-vars.Tₗₒ[i]))));
    # Temperature-climate equation for lower oceans
    @constraint(model, [i=1:N-1], vars.Tₗₒ[i+1] == vars.Tₗₒ[i] + config.ξ₄*(vars.Tₐₜ[i]-vars.Tₗₒ[i]));
    # Capital balance equation
    @NLconstraint(model, [i=1:N-1], vars.K[i+1] <= (1-config.δk)^config.tstep * vars.K[i] + config.tstep*vars.I[i]);
    # Interest rate equation
    @NLconstraint(model, [i=1:N-1], vars.RI[i] == (1+config.ρ)*(vars.CPC[i+1]/vars.CPC[i])^(config.α/config.tstep)-1);

    # Savings rate for asympotic equilibrium
    for i=N-9:N
        JuMP.fix(vars.S[i], params.optlrsav; force=true);
    end
    # Initial conditions
    JuMP.fix(vars.CCA[1], 400.0; force=true);
    JuMP.fix(vars.Mₐₜ[1], config.mat₀; force=true);
    JuMP.fix(vars.Mᵤₚ[1], config.mu₀; force=true);
    JuMP.fix(vars.Mₗₒ[1], config.ml₀; force=true);
    JuMP.fix(vars.Tₗₒ[1], config.tocean₀; force=true);
    if isMumps
        # We can't fix these, the solution becomes concave.
        # This is something buggy in JuMP I think. Haven't been able to pin it down.
        @NLconstraint(model, vars.K[1] == config.k₀);
        @constraint(model, vars.Tₐₜ[1] == config.tatm₀);
    else
        JuMP.fix(vars.Tₐₜ[1], config.tatm₀; force=true);
        JuMP.fix(vars.K[1], config.k₀; force=true);
    end

    @constraint(model, vars.UTILITY == config.tstep * config.scale1 * sum(vars.CEMUTOTPER[i] for i=1:N) + config.scale2);

    # Objective function
    @objective(model, Max, vars.UTILITY);

    assign_scenario(scenario, model, config, params, vars);

    EquationsV2016(eeq,cc)
end

function solve(scenario::Scenario, version::V2016R;
    config::OptionsV2016R = options(version),
    linear_solver::ipoptLinearSolver=ma97,
    optimizer::MOI.OptimizerWithAttributes = optimizer_with_attributes(Ipopt.Optimizer, "print_level" => 5, "max_iter" => 99900, "print_frequency_iter" => 250, "sb" => "yes", "linear_solver" => selectLinearSolver(linear_solver)))

    model = Model(optimizer);

    # Generate a solver test to implement DICE.jl#35 hacks.
    isMumps = JuMP.get_optimizer_attribute(model, "linear_solver") == "mumps";

    params = generate_parameters(config, model);

    # Rate limit
    μ_ubound = [if t < 30 1.0 else config.limμ end for t in 1:config.N];
    cprice_ubound = fill(Inf, config.N); #No initial price bound

    variables = model_vars(version, model, config.N, config.fosslim, μ_ubound, cprice_ubound);

    equations = model_eqs(scenario, model, config, params, variables, isMumps);

    optimize!(model);

    results = model_results(model, config, params, variables, equations);

    DICENarrative(config,params,model,scenario,version,variables,equations,results)
end