text:
using OpenMDAO: AbstractExplicitComp, VarData, PartialsData, get_rows_cols
import OpenMDAO: setup, compute!, compute_partials!
using Interpolations
using ReverseDiff
using LinearAlgebra
using ConcreteStructs
using PCHIPInterpolation
using Dates
using FLoops
using BenchmarkTools

include("normalization_engine_variables.jl")
include("shielding.jl")
include("geometry.jl")
include("source.jl")
include("fan.jl")
include("core.jl")
include("jet.jl")
include("airframe.jl")
include("propagation.jl")
include("split_subbands.jl")
include("lateral_attenuation.jl")
include("ground_reflections.jl")
include("spl.jl")
include("oaspl.jl")
include("pnlt.jl")
include("epnl.jl")
include("ipnlt.jl")
include("ioaspl.jl")
include("aspl.jl")
include("sel.jl")
include("smooth_max.jl")

# Define propagation struct
@concrete struct NoiseModel <: AbstractExplicitComp
    settings :: PyObject
    data :: PyObject
    ac :: PyObject
    n_t :: Int64
    idx :: Dict{Any, Any}
    objective :: String
    noise_model
    X :: Array{Float64, 1}
    J :: Array{Float64, 2}
    noise_model_fwd
end

function NoiseModel(settings::PyObject, data::PyObject, ac::PyObject, n_t::Int64, idx::Dict{Any, Any}, objective::String)

    function compute_noise_observer(settings::PyObject, data::PyObject, ac::PyObject, n_t::Int64, idx::Dict{Any, Any}, x_obs::Array{Float64, 1}, shield::Array{Float64, 2}, idx_supp::Array{CartesianIndex{2}, 1}, objective::String, input_v::Union{Array{Float64, 1}, ReverseDiff.TrackedArray{Float64,Float64,1,Array{Float64,1},Array{Float64,1}}})

        # Initialize
        spl = zeros(eltype(input_v), (n_t, settings.N_f))

        # Compute geometry
        r, beta, theta, phi, c_bar, t_o = geometry(settings, x_obs, input_v[idx["x"][1]:idx["x"][2]], input_v[idx["y"][1]:idx["y"][2]], input_v[idx["z"][1]:idx["z"][2]], input_v[idx["alpha"][1]:idx["alpha"][2]], input_v[idx["gamma"][1]:idx["gamma"][2]], input_v[idx["t_s"][1]:idx["t_s"][2]], input_v[idx["c_0"][1]:idx["c_0"][2]], input_v[idx["T_0"][1]:idx["T_0"][2]])

        # Normalize engine inputs
        if settings.jet_mixing == true && settings.jet_shock == false
            V_j_star, rho_j_star, A_j_star, Tt_j_star = normalization_jet_mixing(settings, input_v[idx["V_j"][1]:idx["V_j"][2]], input_v[idx["rho_j"][1]:idx["rho_j"][2]], input_v[idx["A_j"][1]:idx["A_j"][2]], input_v[idx["Tt_j"][1]:idx["Tt_j"][2]], input_v[idx["c_0"][1]:idx["c_0"][2]], input_v[idx["rho_0"][1]:idx["rho_0"][2]], input_v[idx["T_0"][1]:idx["T_0"][2]])

        elseif settings.jet_shock == true && settings.jet_mixing == false
            V_j_star, A_j_star, Tt_j_star = normalization_jet_shock(settings, input_v[idx["V_j"][1]:idx["V_j"][2]], input_v[idx["A_j"][1]:idx["A_j"][2]], input_v[idx["Tt_j"][1]:idx["Tt_j"][2]], input_v[idx["c_0"][1]:idx["c_0"][2]], input_v[idx["T_0"][1]:idx["T_0"][2]])

        elseif settings.jet_shock ==true && settings.jet_mixing == true
            V_j_star, rho_j_star, A_j_star, Tt_j_star = normalization_jet(settings, input_v[idx["V_j"][1]:idx["V_j"][2]], input_v[idx["rho_j"][1]:idx["rho_j"][2]], input_v[idx["A_j"][1]:idx["A_j"][2]], input_v[idx["Tt_j"][1]:idx["Tt_j"][2]], input_v[idx["c_0"][1]:idx["c_0"][2]], input_v[idx["rho_0"][1]:idx["rho_0"][2]], input_v[idx["T_0"][1]:idx["T_0"][2]])
        end
        if settings.core
            if settings.method_core_turb == "GE"
                mdoti_c_star, Tti_c_star, Ttj_c_star, Pti_c_star, DTt_des_c_star = normalization_core_ge(settings, input_v[idx["mdoti_c"][1]:idx["mdoti_c"][2]], input_v[idx["Tti_c"][1]:idx["Tti_c"][2]], input_v[idx["Ttj_c"][1]:idx["Ttj_c"][2]], input_v[idx["Pti_c"][1]:idx["Pti_c"][2]], input_v[idx["DTt_des_c"][1]:idx["DTt_des_c"][2]], input_v[idx["c_0"][1]:idx["c_0"][2]], input_v[idx["rho_0"][1]:idx["rho_0"][2]], input_v[idx["T_0"][1]:idx["T_0"][2]], input_v[idx["p_0"][1]:idx["p_0"][2]])

            elseif settings.method_core_turb == "PW"
                mdoti_c_star, Tti_c_star, Ttj_c_star, Pti_c_star, rho_te_c_star, c_te_c_star, rho_ti_c_star, c_ti_c_star = normalization_core_pw(settings, input_v[idx["mdoti_c"][1]:idx["mdoti_c"][2]], input_v[idx["Tti_c"][1]:idx["Tti_c"][2]], input_v[idx["Ttj_c"][1]:idx["Ttj_c"][2]], input_v[idx["Pti_c"][1]:idx["Pti_c"][2]], input_v[idx["rho_te_c"][1]:idx["rho_te_c"][2]], input_v[idx["c_te_c"][1]:idx["c_te_c"][2]], input_v[idx["rho_ti_c"][1]:idx["rho_ti_c"][2]], input_v[idx["c_ti_c"][1]:idx["c_ti_c"][2]], input_v[idx["c_0"][1]:idx["c_0"][2]], input_v[idx["rho_0"][1]:idx["rho_0"][2]], input_v[idx["T_0"][1]:idx["T_0"][2]], input_v[idx["p_0"][1]:idx["p_0"][2]])
            end
        end
        if settings.fan_inlet==true || settings.fan_discharge==true
            DTt_f_star, mdot_f_star, N_f_star, A_f_star, d_f_star = normalization_fan(settings, input_v[idx["DTt_f"][1]:idx["DTt_f"][2]], input_v[idx["mdot_f"][1]:idx["mdot_f"][2]], input_v[idx["N_f"][1]:idx["N_f"][2]], input_v[idx["A_f"][1]:idx["A_f"][2]], input_v[idx["d_f"][1]:idx["d_f"][2]], input_v[idx["c_0"][1]:idx["c_0"][2]], input_v[idx["rho_0"][1]:idx["rho_0"][2]], input_v[idx["T_0"][1]:idx["T_0"][2]])
        end

        # Compute source
        if settings.fan_inlet
            spl .+= fan(settings, data, ac, n_t, shield, input_v[idx["M_0"][1]:idx["M_0"][2]], input_v[idx["c_0"][1]:idx["c_0"][2]], input_v[idx["T_0"][1]:idx["T_0"][2]], input_v[idx["rho_0"][1]:idx["rho_0"][2]], theta, DTt_f_star, mdot_f_star, N_f_star, A_f_star, d_f_star, "fan_inlet")
        end
        if settings.fan_discharge
            spl .+= fan(settings, data, ac, n_t, shield, input_v[idx["M_0"][1]:idx["M_0"][2]], input_v[idx["c_0"][1]:idx["c_0"][2]], input_v[idx["T_0"][1]:idx["T_0"][2]], input_v[idx["rho_0"][1]:idx["rho_0"][2]], theta, DTt_f_star, mdot_f_star, N_f_star, A_f_star, d_f_star, "fan_discharge")
        end

        if settings.core
            if settings.method_core_turb == "GE"
                msap_core = core_ge(settings, data, ac, n_t, input_v[idx["M_0"][1]:idx["M_0"][2]], theta, mdoti_c_star, Tti_c_star, Ttj_c_star, Pti_c_star, DTt_des_c_star)
            elseif settings.method.method_core_turb == "PW"
                msap_core = core_pw(settings, data, ac, n_t, input_v[idx["M_0"][1]:idx["M_0"][2]], theta, mdoti_c_star, Tti_c_star, Ttj_c_star, Pti_c_star, rho_te_c_star, c_te_c_star, rho_ti_c_star, c_ti_c_star)
            end

            if settings.suppression && settings.case_name in ["nasa_stca_standard", "stca_enginedesign_standard"]
                msap_core[idx_supp] = (10. ^(-2.3 / 10.) * msap_core)[idx_supp]
            end
            spl .+= msap_core
        end

        if settings.jet_mixing
            msap_jet_mixing = jet_mixing(settings, data, ac, n_t, input_v[idx["M_0"][1]:idx["M_0"][2]], input_v[idx["c_0"][1]:idx["c_0"][2]], theta, V_j_star, rho_j_star, A_j_star, Tt_j_star)
            if settings.suppression && settings.case_name in ["nasa_stca_standard", "stca_enginedesign_standard"]
                msap_jet_mixing[idx_supp] = (10. ^(-2.3 / 10.) * msap_jet_mixing)[idx_supp]
            end
            spl .+= msap_jet_mixing
        end
        if settings.jet_shock
            msap_jet_shock = jet_shock(settings, data, ac, n_t, input_v[idx["M_0"][1]:idx["M_0"][2]], input_v[idx["c_0"][1]:idx["c_0"][2]], theta, V_j_star, input_v[idx["M_j"][1]:idx["M_j"][2]], A_j_star, Tt_j_star)
            if settings.suppression && settings.case_name in ["nasa_stca_standard", "stca_enginedesign_standard"]
                msap_jet_shock[idx_supp] = (10. ^(-2.3 / 10.) * msap_jet_shock)[idx_supp]
            end
            spl .+= msap_jet_shock
        end

        if settings.airframe
            spl .+= airframe(settings, data, ac, n_t, input_v[idx["M_0"][1]:idx["M_0"][2]], input_v[idx["mu_0"][1]:idx["mu_0"][2]], input_v[idx["c_0"][1]:idx["c_0"][2]], input_v[idx["rho_0"][1]:idx["rho_0"][2]], theta, phi, input_v[idx["theta_flaps"][1]:idx["theta_flaps"][2]], input_v[idx["I_landing_gear"][1]:idx["I_landing_gear"][2]])
        end

        # Compute propagation
        spl = propagation(settings, data, x_obs, spl, r, input_v[idx["x"][1]:idx["x"][2]], input_v[idx["z"][1]:idx["z"][2]], c_bar, input_v[idx["rho_0"][1]:idx["rho_0"][2]], input_v[idx["I_0"][1]:idx["I_0"][2]], beta)

        # Compute noise levels
        spl = f_spl(spl, input_v[idx["rho_0"][1]:idx["rho_0"][2]], input_v[idx["c_0"][1]:idx["c_0"][2]])

        oaspl = f_oaspl(spl)
        pnlt, C = f_pnlt(settings, data, spl)
        aspl = f_aspl(data, spl)

        # Compute integrated levels
        if settings.levels_int_metric == "ioaspl"
            level_int = f_ioaspl(t_o, oaspl)
        elseif settings.levels_int_metric == "ipnlt"
            level_int = f_ipnlt(t_o, pnlt)
        elseif settings.levels_int_metric == "epnl"
            level_int = f_epnl(t_o, pnlt)
        elseif settings.levels_int_metric == "sel"
            level_int = f_sel(t_o, aspl)
        end

        # Write output
        if objective == "noise"
            return level_int
        else
            return t_o, spl, aspl, oaspl, pnlt, C, level_int
        end

    end

    function noise_model(settings::PyObject, data::PyObject, ac::PyObject, n_t::Int64, idx::Dict{Any, Any}, objective::String, input_v::Union{Array{Float64, 1}, ReverseDiff.TrackedArray{Float64,Float64,1,Array{Float64,1},Array{Float64,1}}})

        # Get type of input vector
        T = eltype(input_v)

        # Number of observers
        n_obs = size(settings.x_observer_array)[1]

        # Compute noise for each observer
        t_o = zeros(T, (n_obs, n_t))
        spl = zeros(T, (n_obs, n_t, settings.N_f))
        aspl = zeros(T, (n_obs, n_t))
        oaspl = zeros(T, (n_obs, n_t))
        pnlt = zeros(T, (n_obs, n_t))
        C = zeros(T, (n_obs, n_t, settings.N_f))
        level_int = zeros(T, (n_obs, ))

        # Compute airframe shielding coefficients
        shield = shielding(settings, n_t)

        # Suppression  matrix
        idx_supp = findall(input_v[idx["TS"][1]:idx["TS"][2]].*ones(1, settings.N_f).>0.8)

        # Iterate over observers
        ncores=6
        # @floop ThreadedEx(basesize=n_obs÷ncores)
        for i in range(1, n_obs, step=1)

            println("Computing noise at observer: ", i)

            if objective == "noise"
                level_int[i] = compute_noise_observer(settings, data, ac, n_t, idx, settings.x_observer_array[i,:], shield[i,:,:], idx_supp, objective, input_v)
            else
                t_o[i,:], spl[i,:,:], aspl[i,:], oaspl[i,:], pnlt[i,:], C[i,:,:], level_int[i] = compute_noise_observer(settings, data, ac, n_t, idx, settings.x_observer_array[i,:], shield[i,:,:], idx_supp, objective, input_v)
            end

        end

        # Write output
        if objective == "noise"
            k_smooth = 50.
            level_lateral = smooth_max(k_smooth, level_int[1:end-1])

            return [level_lateral, level_int[end]]
        else
            return t_o, spl, aspl, oaspl, pnlt, C, level_int
        end
    end

    # Default values for input vector
    X = range(1, 10000, length=n_t)             # x
    X = vcat(X, zeros(Float64, n_t))            # y
    X = vcat(X, range(1, 1000, length=n_t))     # z
    X = vcat(X, 10. * ones(Float64, n_t))       # alpha
    X = vcat(X, 10. * ones(Float64, n_t))       # gamma
    X = vcat(X, range(0, 100, length=n_t))      # t_s
    X = vcat(X, 1.225 * ones(Float64, n_t))     # rho_0
    X = vcat(X, 1.789e-5 * ones(Float64, n_t))  # mu_0
    X = vcat(X, 340.294 * ones(Float64, n_t))   # c_0
    X = vcat(X, 288.15 * ones(Float64, n_t))    # T_0
    X = vcat(X, 101325. * ones(Float64, n_t))   # p_0
    X = vcat(X, 0.3 * ones(Float64, n_t))       # M_0
    X = vcat(X, 400. * ones(Float64, n_t))      # I_0
    X = vcat(X, 1. * ones(Float64, n_t))        # TS
    X = vcat(X, 10. * ones(Float64, n_t))       # theta_flaps
    if settings.jet_mixing == true && settings.jet_shock == false
        X = vcat(X, 400. * ones(Float64, n_t))  # V_j
        X = vcat(X, 0.8  * ones(Float64, n_t))  # rho_j
        X = vcat(X, 0.5  * ones(Float64, n_t))  # A_j
        X = vcat(X, 500. * ones(Float64, n_t))  # Tt_j
    elseif settings.jet_shock == true && settings.jet_mixing == false
        X = vcat(X, 400. * ones(Float64, n_t))  # V_j
        X = vcat(X, 0.5  * ones(Float64, n_t))  # A_j
        X = vcat(X, 500. * ones(Float64, n_t))  # Tt_j
        X = vcat(X, 1.   * ones(Float64, n_t))  # M_j
    elseif settings.jet_shock ==true && settings.jet_mixing == true
        X = vcat(X, 400. * ones(Float64, n_t))  # V_j
        X = vcat(X, 0.8  * ones(Float64, n_t))  # rho_j
        X = vcat(X, 0.5  * ones(Float64, n_t))  # A_j
        X = vcat(X, 500. * ones(Float64, n_t))  # Tt_j
        X = vcat(X, 1.   * ones(Float64, n_t))  # M_j
    end
    if settings.core
        if settings.method_core_turb == "GE"
            X = vcat(X, 30.   * ones(Float64, n_t))  # mdot_c
            X = vcat(X, 20.e6 * ones(Float64, n_t))  # Tti_c
            X = vcat(X, 800.  *  ones(Float64, n_t)) # Ttj_c
            X = vcat(X, 1600. * ones(Float64, n_t))  # Pti_c
            X = vcat(X, 800.  * ones(Float64, n_t))  # DTt_des_c
        elseif settings.method_core_turb == "PW"
            X = vcat(X, 30.   * ones(Float64, n_t))  # mdot_c
            X = vcat(X, 20e6  * ones(Float64, n_t))  # Tti_c
            X = vcat(X, 800.  * ones(Float64, n_t))  # Ttj_c
            X = vcat(X, 1600. * ones(Float64, n_t))  # Pti_c
            X = vcat(X, 0.5   * ones(Float64, n_t))  # rho_te_c
            X = vcat(X, 600.  * ones(Float64, n_t))  # c_te_c
            X = vcat(X, 3.5   * ones(Float64, n_t))  # rho_ti_c
            X = vcat(X, 800.  * ones(Float64, n_t))  # c_ti_c
        end
    end
    if settings.airframe
        X = vcat(X, ones(Float64, n_t))       # I_landing_gear
    end
    if settings.fan_inlet==true || settings.fan_discharge==true
        X = vcat(X, 70.   * ones(Float64, n_t))
        X = vcat(X, 200.  * ones(Float64, n_t))
        X = vcat(X, 8000. * ones(Float64, n_t))
        X = vcat(X, 0.9   * ones(Float64, n_t))
        X = vcat(X, 1.7   * ones(Float64, n_t))
    end

    # Default values for output value and jacobian
    Y = [100., 100.]
    J = Y.*X'
    #'

    # Define noise function
    noise_model_fwd = (x)->noise_model(settings, data, ac, n_t, idx, objective, x)

    return NoiseModel(settings, data, ac, n_t, idx, objective, noise_model, X, J, noise_model_fwd)
end

function setup(self::NoiseModel)
    # Load options
    settings = self.settings
    n_t = self.n_t

    # Number of observers
    n_obs = size(settings.x_observer_array)[1]

    # Define inputs --------------------------------------------------------------------------------
    inputs = Vector{VarData}()
    push!(inputs, VarData("x", shape=(n_t, ), val=ones(n_t), units="m"))
    push!(inputs, VarData("y", shape=(n_t, ), val=ones(n_t), units="m"))
    push!(inputs, VarData("z", shape=(n_t, ), val=ones(n_t), units="m"))
    push!(inputs, VarData("alpha", shape=(n_t, ), val=ones(n_t), units="deg"))
    push!(inputs, VarData("gamma", shape=(n_t, ), val=ones(n_t), units="deg"))
    push!(inputs, VarData("t_s", shape=(n_t, ), val=ones(n_t), units="s"))
    push!(inputs, VarData("rho_0", shape=(n_t,), val=ones(n_t), units="kg/m**3"))
    push!(inputs, VarData("mu_0", shape=(n_t,), val=ones(n_t), units="kg/m/s"))
    push!(inputs, VarData("c_0", shape=(n_t, ), val=ones(n_t), units="m/s"))
    push!(inputs, VarData("T_0", shape=(n_t, ), val=ones(n_t), units="K"))
    push!(inputs, VarData("p_0", shape=(n_t, ), val=ones(n_t), units="Pa"))
    push!(inputs, VarData("M_0", shape=(n_t,), val=ones(n_t)))
    push!(inputs, VarData("I_0", shape=(n_t,), val=ones(n_t), units="kg/m**2/s"))
    push!(inputs, VarData("TS", shape=(n_t,), val=ones(n_t)))
    push!(inputs, VarData("theta_flaps", shape=(n_t,), val=ones(n_t), units="deg"))

    if settings.jet_mixing == true && settings.jet_shock == false
        push!(inputs, VarData("V_j", shape=(n_t,), val=ones(n_t), units="m/s"))
        push!(inputs, VarData("rho_j", shape=(n_t,), val=ones(n_t), units="kg/m**3"))
        push!(inputs, VarData("A_j", shape=(n_t,), val=ones(n_t), units="m**2"))
        push!(inputs, VarData("Tt_j", shape=(n_t,), val=ones(n_t), units="K"))
    elseif settings.jet_shock == true && settings.jet_mixing == false
        push!(inputs, VarData("V_j", shape=(n_t,), val=ones(n_t), units="m/s"))
        push!(inputs, VarData("M_j", shape=(n_t,), val=ones(n_t)))
        push!(inputs, VarData("A_j", shape=(n_t,), val=ones(n_t), units="m**2"))
        push!(inputs, VarData("Tt_j", shape=(n_t,), val=ones(n_t), units="K"))
    elseif settings.jet_shock ==true && settings.jet_mixing == true
        push!(inputs, VarData("V_j", shape=(n_t,), val=ones(n_t), units="m/s"))
        push!(inputs, VarData("rho_j", shape=(n_t,), val=ones(n_t), units="kg/m**3"))
        push!(inputs, VarData("A_j", shape=(n_t,), val=ones(n_t), units="m**2"))
        push!(inputs, VarData("Tt_j", shape=(n_t,), val=ones(n_t), units="K"))
        push!(inputs, VarData("M_j", shape=(n_t,), val=ones(n_t)))
    end
    if settings.core
        if settings.method_core_turb == "GE"
            push!(inputs, VarData("mdoti_c", shape=(n_t,), val=ones(n_t), units="kg/s"))
            push!(inputs, VarData("Tti_c", shape=(n_t,), val=ones(n_t), units="K"))
            push!(inputs, VarData("Ttj_c", shape=(n_t,), val=ones(n_t), units="K"))
            push!(inputs, VarData("Pti_c", shape=(n_t,), val=ones(n_t), units="Pa"))
            push!(inputs, VarData("DTt_des_c", shape=(n_t,), val=ones(n_t), units="K"))
        elseif settings.method_core_turb == "PW"
            push!(inputs, VarData("mdoti_c", shape=(n_t,), val=ones(n_t), units="kg/s"))
            push!(inputs, VarData("Tti_c", shape=(n_t,), val=ones(n_t), units="K"))
            push!(inputs, VarData("Ttj_c", shape=(n_t,), val=ones(n_t), units="K"))
            push!(inputs, VarData("Pti_c", shape=(n_t,), val=ones(n_t), units="Pa"))
            push!(inputs, VarData("rho_te_c", shape=(n_t,), val=ones(n_t), units="kg/m**3"))
            push!(inputs, VarData("c_te_c", shape=(n_t,), val=ones(n_t), units="m/s"))
            push!(inputs, VarData("rho_ti_c", shape=(n_t,), val=ones(n_t), units="kg/m**3"))
            push!(inputs, VarData("c_ti_c", shape=(n_t,), val=ones(n_t), units="m/s"))
        end
    end
    if settings.airframe
        push!(inputs, VarData("I_landing_gear", shape=(n_t,), val=ones(n_t)))
    end
    if settings.fan_inlet==true || settings.fan_discharge==true
        push!(inputs, VarData("DTt_f", shape=(n_t,), val=ones(n_t), units="K"))
        push!(inputs, VarData("mdot_f", shape=(n_t,), val=ones(n_t), units="kg/s"))
        push!(inputs, VarData("N_f", shape=(n_t,), val=ones(n_t), units="rpm"))
        push!(inputs, VarData("A_f", shape=(n_t,), val=ones(n_t), units="m**2"))
        push!(inputs, VarData("d_f", shape=(n_t,), val=ones(n_t), units="m"))
    end

    # Define outputs --------------------------------------------------------------------------------
    outputs = Vector{VarData}()
    if self.objective == "noise"
        push!(outputs, VarData("lateral", shape=(1, ), val=0.))
        push!(outputs, VarData("flyover", shape=(1, ), val=0.))
    else
        push!(outputs, VarData("t_o", shape=(n_obs, n_t), val=ones(n_obs, n_t)))
        push!(outputs, VarData("spl", shape=(n_obs, n_t, settings.N_f), val=ones(n_obs, n_t, settings.N_f)))
        push!(outputs, VarData("aspl", shape=(n_obs, n_t), val=ones(n_obs, n_t)))
        push!(outputs, VarData("oaspl", shape=(n_obs, n_t), val=ones(n_obs, n_t)))
        push!(outputs, VarData("pnlt", shape=(n_obs, n_t), val=ones(n_obs, n_t)))
        push!(outputs, VarData("C", shape=(n_obs, n_t, settings.N_f), val=ones(n_obs, n_t, settings.N_f)))
        push!(outputs, VarData(settings.levels_int_metric, shape=(n_obs, ), val=ones(n_obs, )))
    end

    ## Define partials --------------------------------------------------------------------------------
    partials = Vector{PartialsData}()

    if self.objective == "noise"
        for mic in ["lateral", "flyover"]
            push!(partials, PartialsData(mic, "x"))
            push!(partials, PartialsData(mic, "y"))
            push!(partials, PartialsData(mic, "z"))
            push!(partials, PartialsData(mic, "alpha"))
            push!(partials, PartialsData(mic, "gamma"))
            push!(partials, PartialsData(mic, "t_s"))
            push!(partials, PartialsData(mic, "rho_0"))
            push!(partials, PartialsData(mic, "mu_0"))
            push!(partials, PartialsData(mic, "c_0"))
            push!(partials, PartialsData(mic, "T_0"))
            push!(partials, PartialsData(mic, "p_0"))
            push!(partials, PartialsData(mic, "M_0"))
            push!(partials, PartialsData(mic, "I_0"))
            push!(partials, PartialsData(mic, "TS"))
            push!(partials, PartialsData(mic, "theta_flaps"))
            if settings.jet_mixing == true && settings.jet_shock == false
                push!(partials, PartialsData(mic, "V_j"))
                push!(partials, PartialsData(mic, "rho_j"))
                push!(partials, PartialsData(mic, "A_j"))
                push!(partials, PartialsData(mic, "Tt_j"))
            elseif settings.jet_shock == true && settings.jet_mixing == false
                push!(partials, PartialsData(mic, "V_j"))
                push!(partials, PartialsData(mic, "M_j"))
                push!(partials, PartialsData(mic, "A_j"))
                push!(partials, PartialsData(mic, "Tt_j"))
            elseif settings.jet_shock ==true && settings.jet_mixing == true
                push!(partials, PartialsData(mic, "V_j"))
                push!(partials, PartialsData(mic, "rho_j"))
                push!(partials, PartialsData(mic, "A_j"))
                push!(partials, PartialsData(mic, "Tt_j"))
                push!(partials, PartialsData(mic, "M_j"))
            end
            if settings.core
                if settings.method_core_turb == "GE"
                    push!(partials, PartialsData(mic, "mdoti_c"))
                    push!(partials, PartialsData(mic, "Tti_c"))
                    push!(partials, PartialsData(mic, "Ttj_c"))
                    push!(partials, PartialsData(mic, "Pti_c"))
                    push!(partials, PartialsData(mic, "DTt_des_c"))
                elseif settings.method_core_turb == "PW"
                    push!(partials, PartialsData(mic, "mdoti_c"))
                    push!(partials, PartialsData(mic, "Tti_c"))
                    push!(partials, PartialsData(mic, "Ttj_c"))
                    push!(partials, PartialsData(mic, "Pti_c"))
                    push!(partials, PartialsData(mic, "rho_te_c"))
                    push!(partials, PartialsData(mic, "c_te_c"))
                    push!(partials, PartialsData(mic, "rho_ti_c"))
                    push!(partials, PartialsData(mic, "c_ti_c"))
                end
            end
            if settings.airframe
                push!(partials, PartialsData(mic, "I_landing_gear"))
            end
            if settings.fan_inlet==true || settings.fan_discharge==true
                push!(partials, PartialsData(mic, "DTt_f"))
                push!(partials, PartialsData(mic, "mdot_f"))
                push!(partials, PartialsData(mic, "N_f"))
                push!(partials, PartialsData(mic, "A_f"))
                push!(partials, PartialsData(mic, "d_f"))
            end
        end
    end

    return inputs, outputs, partials
end

function get_noise_input_vector!(X::Array{Float64, 1}, settings::PyObject, inputs::PyDict{String, PyArray, true}, idx::Dict{Any, Any})

    # Extract inputs
    X[idx["x"][1]:idx["x"][2]] = inputs["x"]
    X[idx["y"][1]:idx["y"][2]] = inputs["y"]
    X[idx["z"][1]:idx["z"][2]] = inputs["z"]
    X[idx["alpha"][1]:idx["alpha"][2]] = inputs["alpha"]
    X[idx["gamma"][1]:idx["gamma"][2]] = inputs["gamma"]
    X[idx["t_s"][1]:idx["t_s"][2]] = inputs["t_s"]
    X[idx["rho_0"][1]:idx["rho_0"][2]] = inputs["rho_0"]
    X[idx["mu_0"][1]:idx["mu_0"][2]] = inputs["mu_0"]
    X[idx["c_0"][1]:idx["c_0"][2]] = inputs["c_0"]
    X[idx["T_0"][1]:idx["T_0"][2]] = inputs["T_0"]
    X[idx["p_0"][1]:idx["p_0"][2]] = inputs["p_0"]
    X[idx["M_0"][1]:idx["M_0"][2]] = inputs["M_0"]
    X[idx["I_0"][1]:idx["I_0"][2]] = inputs["I_0"]
    X[idx["TS"][1]:idx["TS"][2]] = inputs["TS"]
    X[idx["theta_flaps"][1]:idx["theta_flaps"][2]] = inputs["theta_flaps"]

    if settings.jet_mixing == true && settings.jet_shock == false
        X[idx["V_j"][1]:idx["V_j"][2]] = inputs["V_j"]
        X[idx["rho_j"][1]:idx["rho_j"][2]] = inputs["rho_j"]
        X[idx["A_j"][1]:idx["A_j"][2]] = inputs["A_j"]
        X[idx["Tt_j"][1]:idx["Tt_j"][2]] = inputs["Tt_j"]
    elseif settings.jet_shock == true && settings.jet_mixing == false
        X[idx["V_j"][1]:idx["V_j"][2]] = inputs["V_j"]
        X[idx["M_j"][1]:idx["M_j"][2]] = inputs["M_j"]
        X[idx["A_j"][1]:idx["A_j"][2]] = inputs["A_j"]
        X[idx["Tt_j"][1]:idx["Tt_j"][2]] = inputs["Tt_j"]
    elseif settings.jet_shock ==true && settings.jet_mixing == true
        X[idx["V_j"][1]:idx["V_j"][2]] = inputs["V_j"]
        X[idx["rho_j"][1]:idx["rho_j"][2]] = inputs["rho_j"]
        X[idx["A_j"][1]:idx["A_j"][2]] = inputs["A_j"]
        X[idx["Tt_j"][1]:idx["Tt_j"][2]] = inputs["Tt_j"]
        X[idx["M_j"][1]:idx["M_j"][2]] = inputs["M_j"]
    end
    if settings.core
        if settings.method_core_turb == "GE"
            X[idx["mdoti_c"][1]:idx["mdoti_c"][2]] = inputs["mdoti_c"]
            X[idx["Tti_c"][1]:idx["Tti_c"][2]] = inputs["Tti_c"]
            X[idx["Ttj_c"][1]:idx["Ttj_c"][2]] = inputs["Ttj_c"]
            X[idx["Pti_c"][1]:idx["Pti_c"][2]] = inputs["Pti_c"]
            X[idx["DTt_des_c"][1]:idx["DTt_des_c"][2]] = inputs["DTt_des_c"]
        elseif settings.method_core_turb == "PW"
            X[idx["mdoti_c"][1]:idx["mdoti_c"][2]] = inputs["mdoti_c"]
            X[idx["Tti_c"][1]:idx["Tti_c"][2]] = inputs["Tti_c"]
            X[idx["Ttj_c"][1]:idx["Ttj_c"][2]] = inputs["Ttj_c"]
            X[idx["Pti_c"][1]:idx["Pti_c"][2]] = inputs["Pti_c"]
            X[idx["rho_te_c"][1]:idx["rho_te_c"][2]] = inputs["rho_te_c"]
            X[idx["c_te_c"][1]:idx["c_te_c"][2]] = inputs["c_te_c"]
            X[idx["rho_ti_c"][1]:idx["rho_ti_c"][2]] = inputs["rho_ti_c"]
            X[idx["c_ti_c"][1]:idx["c_ti_c"][2]] = inputs["c_ti_c"]
        end
    end
    if settings.airframe
        X[idx["I_landing_gear"][1]:idx["I_landing_gear"][2]] = inputs["I_landing_gear"]
    end
    if settings.fan_inlet==true || settings.fan_discharge==true
        X[idx["DTt_f"][1]:idx["DTt_f"][2]] = inputs["DTt_f"]
        X[idx["mdot_f"][1]:idx["mdot_f"][2]] = inputs["mdot_f"]
        X[idx["N_f"][1]:idx["N_f"][2]] = inputs["N_f"]
        X[idx["A_f"][1]:idx["A_f"][2]] = inputs["A_f"]
        X[idx["d_f"][1]:idx["d_f"][2]] = inputs["d_f"]
    end

end

function compute!(self::NoiseModel, inputs, outputs)
    # Load options
    settings = self.settings
    data = self.data
    ac = self.ac
    n_t = self.n_t
    idx = self.idx
    X = self.X

    # Get input vector
    get_noise_input_vector!(X, settings, inputs, idx)

    if self.objective == "noise"
        levels_int = self.noise_model(settings, data, ac, n_t, idx, self.objective, X)
        @. outputs["lateral"] = levels_int[1]
        @. outputs["flyover"] = levels_int[2]

        # Print inputs to file
        open(settings.pyNA_directory * "/cases/" * settings.case_name * "/output/" * settings.output_directory_name * "/" * "inputs_TS.txt","a") do io
            println(io, inputs["TS"])
        end

        open(settings.pyNA_directory * "/cases/" * settings.case_name * "/output/" * settings.output_directory_name * "/" * "inputs_alpha.txt","a") do io
            println(io, inputs["alpha"])
        end

        open(settings.pyNA_directory * "/cases/" * settings.case_name * "/output/" * settings.output_directory_name * "/" * "inputs_theta_flaps.txt","a") do io
            println(io, inputs["theta_flaps"])
        end

        # Print outputs to file
        open(settings.pyNA_directory * "/cases/" * settings.case_name * "/output/" * settings.output_directory_name * "/" * "outputs_levels_int.txt", "a") do io
            println(io, levels_int)
        end

    else
        t_o, spl, aspl, oaspl, pnlt, C, levels_int = self.noise_model(settings, data, ac, n_t, idx, self.objective, X)
        @. outputs["t_o"] = t_o
        @. outputs["spl"] = spl
        @. outputs["aspl"] = aspl
        @. outputs["oaspl"] = oaspl
        @. outputs["pnlt"] = pnlt
        @. outputs["C"] = C
        @. outputs[settings.levels_int_metric] = levels_int
    end

end

function compute_partials!(self::NoiseModel, inputs, partials)
    # Load options
    settings = self.settings
    n_t = self.n_t
    idx = self.idx
    X = self.X
    J = self.J

    # Print start statement
    println("Computing partials noise")

    # Get input vector
    get_noise_input_vector!(X, settings, inputs, idx)

    # compute Jacobian
    ReverseDiff.jacobian!(J, self.noise_model_fwd, X)

    for (i, mic) in enumerate(["lateral", "flyover"])
        dnoise_dx = reshape(J[i, idx["x"][1]:idx["x"][2]], (1, n_t))
        @. partials[mic, "x"] = dnoise_dx
        dnoise_dy = reshape(J[i, idx["y"][1]:idx["y"][2]], (1, n_t))
        @. partials[mic, "y"] = dnoise_dy
        dnoise_dz = reshape(J[i, idx["z"][1]:idx["z"][2]], (1, n_t))
        @. partials[mic, "z"] = dnoise_dz
        dnoise_dalpha = reshape(J[i, idx["alpha"][1]:idx["alpha"][2]], (1, n_t))
        @. partials[mic, "alpha"] = dnoise_dalpha
        dnoise_dgamma = reshape(J[i, idx["gamma"][1]:idx["gamma"][2]], (1, n_t))
        @. partials[mic, "gamma"] = dnoise_dgamma
        dnoise_dts = reshape(J[i, idx["t_s"][1]:idx["t_s"][2]], (1, n_t))
        @. partials[mic, "t_s"] = dnoise_dts
        dnoise_drho_0 = reshape(J[i, idx["rho_0"][1]:idx["rho_0"][2]], (1, n_t))
        @. partials[mic, "rho_0"] = dnoise_drho_0
        dnoise_dmu_0 = reshape(J[i, idx["mu_0"][1]:idx["mu_0"][2]], (1, n_t))
        @. partials[mic, "mu_0"] = dnoise_dmu_0
        dnoise_dc_0 = reshape(J[i, idx["c_0"][1]:idx["c_0"][2]], (1, n_t))
        @. partials[mic, "c_0"] = dnoise_dc_0
        dnoise_dT_0 = reshape(J[i, idx["T_0"][1]:idx["T_0"][2]], (1, n_t))
        @. partials[mic, "T_0"] = dnoise_dT_0
        dnoise_dp_0 = reshape(J[i, idx["p_0"][1]:idx["p_0"][2]], (1, n_t))
        @. partials[mic, "p_0"] = dnoise_dp_0
        dnoise_dM_0 = reshape(J[i, idx["M_0"][1]:idx["M_0"][2]], (1, n_t))
        @. partials[mic, "M_0"] = dnoise_dM_0
        dnoise_dI_0 = reshape(J[i, idx["I_0"][1]:idx["I_0"][2]], (1, n_t))
        @. partials[mic, "I_0"] = dnoise_dI_0
        dnoise_dTS = reshape(J[i, idx["TS"][1]:idx["TS"][2]], (1, n_t))
        @. partials[mic, "TS"] = dnoise_dTS
        dnoise_dtheta_flaps = reshape(J[i, idx["theta_flaps"][1]:idx["theta_flaps"][2]], (1, n_t))
        @. partials[mic, "theta_flaps"] = dnoise_dtheta_flaps
        if settings.jet_mixing == true && settings.jet_shock == false
            dnoise_dvj = reshape(J[i, idx["V_j"][1]:idx["V_j"][2]], (1, n_t))
            @. partials[mic, "V_j"] = dnoise_dvj
            dnoise_drhoj = reshape(J[i, idx["rho_j"][1]:idx["rho_j"][2]], (1, n_t))
            @. partials[mic, "rho_j"] = dnoise_drhoj
            dnoise_dAj = reshape(J[i, idx["A_j"][1]:idx["A_j"][2]], (1, n_t))
            @. partials[mic, "A_j"] = dnoise_dAj
            dnoise_dTtj = reshape(J[i, idx["Tt_j"][1]:idx["Tt_j"][2]], (1, n_t))
            @. partials[mic, "Tt_j"] = dnoise_dTtj
        elseif settings.jet_shock == true && settings.jet_mixing == false
            dnoise_dvj = reshape(J[i, idx["V_j"][1]:idx["V_j"][2]], (1, n_t))
            @. partials[mic, "V_j"] = dnoise_dvj
            dnoise_dMj = reshape(J[i, idx["M_j"][1]:idx["M_j"][2]], (1, n_t))
            @. partials[mic, "M_j"] = dnoise_dMj
            dnoise_dAj = reshape(J[i, idx["A_j"][1]:idx["A_j"][2]], (1, n_t))
            @. partials[mic, "A_j"] = dnoise_dAj
            dnoise_dTtj = reshape(J[i, idx["Tt_j"][1]:idx["Tt_j"][2]], (1, n_t))
            @. partials[mic, "Tt_j"] = dnoise_dTtj
        elseif settings.jet_shock ==true && settings.jet_mixing == true
            dnoise_dvj = reshape(J[i, idx["V_j"][1]:idx["V_j"][2]], (1, n_t))
            @. partials[mic, "V_j"] = dnoise_dvj
            dnoise_drhoj = reshape(J[i, idx["rho_j"][1]:idx["rho_j"][2]], (1, n_t))
            @. partials[mic, "rho_j"] = dnoise_drhoj
            dnoise_dAj = reshape(J[i, idx["A_j"][1]:idx["A_j"][2]], (1, n_t))
            @. partials[mic, "A_j"] = dnoise_dAj
            dnoise_dTtj = reshape(J[i, idx["Tt_j"][1]:idx["Tt_j"][2]], (1, n_t))
            @. partials[mic, "Tt_j"] = dnoise_dTtj
            dnoise_dMj = reshape(J[i, idx["M_j"][1]:idx["M_j"][2]], (1, n_t))
            @. partials[mic, "M_j"] = dnoise_dMj
        end
        if settings.core
            if settings.method_core_turb == "GE"
                dnoise_dmdoti_c = reshape(J[i, idx["mdoti_c"][1]:idx["mdoti_c"][2]], (1, n_t))
                @. partials[mic, "mdoti_c"] = dnoise_dmdoti_c
                dnoise_dTti_c = reshape(J[i, idx["Tti_c"][1]:idx["Tti_c"][2]], (1, n_t))
                @. partials[mic, "Tti_c"] = dnoise_dTti_c
                dnoise_dTtj_c = reshape(J[i, idx["Ttj_c"][1]:idx["Ttj_c"][2]], (1, n_t))
                @. partials[mic, "Ttj_c"] = dnoise_dTtj_c
                dnoise_dPti_c = reshape(J[i, idx["Pti_c"][1]:idx["Pti_c"][2]], (1, n_t))
                @. partials[mic, "Pti_c"] = dnoise_dPti_c
                dnoise_dDTt_des_c = reshape(J[i, idx["DTt_des_c"][1]:idx["DTt_des_c"][2]], (1, n_t))
                @. partials[mic, "DTt_des_c"] = dnoise_dDTt_des_c
            elseif settings.method_core_turb == "PW"
                dnoise_dmdoti_c = reshape(J[i, idx["mdoti_c"][1]:idx["mdoti_c"][2]], (1, n_t))
                @. partials[mic, "mdoti_c"] = dnoise_dmdoti_c
                dnoise_dTti_c = reshape(J[i, idx["Tti_c"][1]:idx["Tti_c"][2]], (1, n_t))
                @. partials[mic, "Tti_c"] = dnoise_dTti_c
                dnoise_dTtj_c = reshape(J[i, idx["Ttj_c"][1]:idx["Ttj_c"][2]], (1, n_t))
                @. partials[mic, "Ttj_c"] = dnoise_dTtj_c
                dnoise_dPti_c = reshape(J[i, idx["Pti_c"][1]:idx["Pti_c"][2]], (1, n_t))
                @. partials[mic, "Pti_c"] = dnoise_dPti_c
                dnoise_drho_te_c = reshape(J[i, idx["rho_te_c"][1]:idx["rho_te_c"][2]], (1, n_t))
                @. partials[mic, "rho_te_c"] = dnoise_drho_te_c
                dnoise_dc_te_c = reshape(J[i, idx["c_te_c"][1]:idx["c_te_c"][2]], (1, n_t))
                @. partials[mic, "c_te_c"] = dnoise_dc_te_c
                dnoise_drho_ti_c = reshape(J[i, idx["rho_ti_c"][1]:idx["rho_ti_c"][2]], (1, n_t))
                @. partials[mic, "rho_ti_c"] = dnoise_drho_ti_c
                dnoise_dc_ti_c = reshape(J[i, idx["c_ti_c"][1]:idx["c_ti_c"][2]], (1, n_t))
                @. partials[mic, "c_ti_c"] = dnoise_dc_ti_c
            end
        end
        if settings.airframe
            dnoise_dI_landing_gear = reshape(J[i, idx["I_landing_gear"][1]:idx["I_landing_gear"][2]], (1, n_t))
            @. partials[mic, "I_landing_gear"] = dnoise_dI_landing_gear
        end
        if settings.fan_inlet==true || settings.fan_discharge==true
            dnoise_dDTt_f = reshape(J[i, idx["DTt_f"][1]:idx["DTt_f"][2]], (1, n_t))
            @. partials[mic, "DTt_f"] = dnoise_dDTt_f
            dnoise_dmdot_f = reshape(J[i, idx["mdot_f"][1]:idx["mdot_f"][2]], (1, n_t))
            @. partials[mic, "mdot_f"] = dnoise_dmdot_f
            dnoise_dN_f = reshape(J[i, idx["N_f"][1]:idx["N_f"][2]], (1, n_t))
            @. partials[mic, "N_f"] = dnoise_dN_f
            dnoise_dA_f = reshape(J[i, idx["A_f"][1]:idx["A_f"][2]], (1, n_t))
            @. partials[mic, "A_f"] = dnoise_dA_f
            dnoise_dd_f = reshape(J[i, idx["d_f"][1]:idx["d_f"][2]], (1, n_t))
            @. partials[mic, "d_f"] = dnoise_dd_f
        end
    end

    println("Done computing partials noise.")

end

max_stars_repo_name: MIT-LAE/pyNA

max_issues_repo_name: MIT-LAE/pyNA

max_forks_repo_name: MIT-LAE/pyNA


text:
function subtract(d1::Dict, d2::Dict)
    res = Dict()
    for key in keys(d1)
        if key ∉ keys(d2)
            res[key] = nothing
        end
    end
    res
end

# using Sets instead
"""
    Subtract two sets.
    Result is a Set.
"""
function subtract(s1::Set, s2::Set)
    setdiff(s1, s2)
end

#   ---------- Duplicates ----------

# ... using dictionaries

function hasduplicates(t::Dict)
    d = Dict()
    for x in t
        if x ∈ d
            return true
        end
        d[x] = nothing
    end
    false
end

function usesonly(word, available)
    for letter in word
        if word ∉ available
            return false
        end
    end
    true
end

# ... using sets
function hasduplicates(t)
    length(Set(t)) < length(t)
end

function usesonly_set(word, available)
    Set(word) ⊆ Set(available)
end

# ---------- Math -----------

x = 0:0.1:2π

cos.(x) == 0.5*(ℯ.^(im*x) + ℯ.^(-im*x))

# ----------- Strings --------

function usesonly_regex(word::String, available::String)
    r = Regex("[^$(available)]")
    @show r;
    !occursin(r,word)
end

usesonly_regex("banana", "abn")
usesonly_regex("bananas", "abn")

match(r"[^abn]", "banana")

m = match(r"[^abn]", "bananas")

# ---------------- Matricies  -----------------

z = zeros(Float64, 2,3)

z[1,2] = 2

z[2,3] = 1
z

size(z)

# 1x3 row matrix
s = ones(String, 1, 3)
size(s)

# 3 x _ column matrix
s_col = ["", "", ""]
size(s_col)
s == s_col

# also a column matrix
s_col2 = ["";"";""]
size(s_col2)

s_col == s_col2

s == s_col2

a = [1 2 3; 4 5 6]

@show a
print(a)
size(a)

#slicing

u = z[:,2:end]

#broadcasting
@show im*u
size(im*u)
g = ℯ.^(im*u)
@show g
typeof(g)

typeof(u)

# ----------- Interfaces: p243 -----------------

#parametric type with no fields
struct Fibonacci{T<:Real} end

#outer constructor
Fibonacci(d::DataType) = d<:Real ? Fibonacci{d}() : error("Not Real type!")

# Iterator interface implementation: 2 parts:

# Part 1: called to initialize the iterator
Base.iterate(::Fibonacci{T}) where {T<:Real} = (zero(T), (one(T), one(T)))

# Part 2:
Base.iterate(::Fibonacci{T}, state::Tuple{T,T}) where {T<:Real} = (state[1], (state[2], state[1] + state[2]))

for e in Fibonacci(Int64)
    e > 100 && break
    print(e, " ")
end

#Equivalently:
f = Fibonacci(Int64)

#initialize state sequence
(a, (b,c)) = iterate(f)

#iterate ...
(a, (b,c)) = iterate(f, (b,c))

(a, (b,c)) = iterate(f, (b,c))

(a, (b,c)) = iterate(f, (b,c))

(a, (b,c)) = iterate(f, (b,c))

(a, (b,c)) = iterate(f, (b,c))

(a, (b,c)) = iterate(f, (b,c))

(a, (b,c)) = iterate(f, (b,c))

(a, (b,c)) = iterate(f, (b,c))

(a, (b,c)) = iterate(f, (b,c))

(a, (b,c)) = iterate(f, (b,c))

#= a for loop in julia:

for i in iter
    body
end

    equivalant to:

next = iterate(iter)
while next !== nothing
    (i, state) = next
    body
    next = iterate(iter, state)
end

=#

ext: jl


text:
using Test
using EstimateSDE
using Distributions
using LinearAlgebra
using Random
using Zygote

@testset "EstimateSDE" begin
ty =[
  0.69   4.6   5.72   8.09   8.19   8.21  10.11  10.37  11.88  12.46
 -0.31  -0.02  0.12  -0.23  -0.24  -0.24   0.15   0.24   0.55   0.47
]
data = map(x->Obs(x...),eachcol(ty))
a = [0.371, 0.0693, 1.26, 0.124]
eq = Sde(a,0.05)
fixed_state = State([data[1].y,0.2,-0.05,2.1])

@testset "Kalman equations" begin
  μ = [5.99,0.45,3.49,0.23]
  Σ = [
     5.71  -1.11  -3.71  -0.49
    -1.11   2.58   0.26  -1.94
    -3.71   0.26   2.84   0.36
    -0.49  -1.94   0.36  11.48
  ]
  proc = SdeProcess(eq,State(μ,Σ))
  Δt = 2.3
  for engine in EstimateSDE.ENGINES
    proc′ = transform_ss(sdetype(engine),proc)
    eq′ = proc′.eq
    state′ = proc′.state
    e = engine(eq)

    μy,σ2y = EstimateSDE.obsdist(e,state′)
    @test μy≈5.99
    @test σ2y≈5.71

    state_m′ = EstimateSDE.measure(e,state′,-4.3,μy,σ2y)
    state_m = transform_ss(Sde,SdeProcess(eq′,state_m′)).state
    @test state_m.μ≈[-4.299999999999999, 2.450332749562172, 10.175796847635727, 1.1130297723292468]
    @test state_m.Σ≈[-8.881784197001252e-16 0.0 0.0 0.0; 0.0 2.3642206654991242 -0.4612084063047286 -2.0352539404553416; 0.0 -0.4612084063047286 0.4294746059544656 0.04162872154115588; 0.0 -2.0352539404553416 0.04162872154115588 11.437950963222418]

    state_p′ = EstimateSDE.predict(e,state′,Δt)
    state_p = transform_ss(Sde,SdeProcess(eq′,state_p′)).state
    @test state_p.μ≈[ 10.30350225480549,-1.0284495977990458,-6.621680617195209,-4.072984824995861]
    @test state_p.Σ≈[21.053696357059337 19.2706890985522 0.7602965515146454 -15.198802594403409; 19.2706890985522 21.706321456728876 2.464682525360381 -17.17083607634992; 0.7602965515146454 2.464682525360381 2.365832518117748 -2.1570291778727433; -15.198802594403409 -17.17083607634992 -2.1570291778727433 13.789300684647449]
  end
end

@testset "logpdf" begin
  logpdf(SdeProcess(eq),data)≈-25538.711601624374
  logpdf(SdeProcess(eq,fixed_state),data)≈-29370.38457270086
  lps_ref = logpdf(SdeProcess(eq),data)
  lpf_ref = logpdf(SdeProcess(eq,fixed_state),data)
  proc = SdeProcess(eq,fixed_state)
  for E in EstimateSDE.ENGINES
    e = E(eq)
    lps,_ = logpdf_state(e,statstate(e),data)
    @test lps≈lps_ref atol=0.1
    state_fixed′ = transform_ss(sdetype(E),proc).state
    lpf,_ = logpdf_state(e,state_fixed′,data)
    @test lpf≈lpf_ref atol=0.1
  end
end

@testset "SDE roots" begin
  eqr = testprocess(:C).eq
  eq = Sde(eqr)
  r = eqr.r
  p = order(eqr)
  U = [rc^(j-1) for j in 1:p, rc in r]
  g = zeros(p); g[end] = 1
  v = U\g
  Q′ = eqr.σϵ^2*v*v'
  Σs′ = [-Q′[i,j]/(r[i]+conj(r[j])) for i in 1:p,j in 1:p]
  A′ = Diagonal(r)
  # Continuous Lyapunov equation
  @test A′*Σs′+Σs′*A′'+Q′≈zeros(p,p) atol=1e-14
  # Compare to diagonalized vector SDE(1) process
  veq = diag_vectorsde(convert(VectorSde,eq))
  @test A′≈veq.A
  @test Q′≈veq.Q
  @test Σs′≈statcovz(veq)
end

@testset "logpdf_ml_σϵ" begin
  ty =[
    0.69   4.6   5.72   8.09   8.19   8.21  10.11  10.37  11.88  12.46
  -0.31  -0.02  0.12  -0.23  -0.24  -0.24   0.15   0.24   0.55   0.47
  ]
  data = map(x->Obs(x...),eachcol(ty))
  a = [0.371, 0.0693, 1.26, 0.124]
  eq = Sde(a,0.05)
  fixed_state = State([data[1].y,0.2,-0.05,2.1])
  proc = SdeProcess(eq,fixed_state)
  eq1 = Sde(eq.a,1)
  for E in EstimateSDE.ENGINES
    e = E(eq1)
    for stationary in [true,false]
      state′ = if stationary
        statstate(eq1)
      else
        transform_ss(sdetype(E),proc).state
      end
      lp,σϵ_ml = logpdf_ml_σϵ(e,state′,data)
      σϵ_ml_test = stationary ? 3.575206949538863 : 4.041520808012871
      @test σϵ_ml≈σϵ_ml_test atol=0.05
    end
  end
end

@testset "transform variable" begin
  using TransformVariables
  eq = Sde([0.371, 0.0693, 1.26, 0.124],0.05)
  z = [0.3,1.1,-0.7,0.2]
  proc = SdeProcess(eq,State(z))
  data = rand(MersenneTwister(12202),proc,200,1.0)

  problem = ProblemLike(data;p=order(proc),engine=ElimQ)
  𝓣 = transformation(problem,proc)
  𝓣⁻¹ = inverse(𝓣)
  θ = paramstuple(problem,proc)
  θ_rec = (𝓣∘𝓣⁻¹)(θ)
  eq_rec = Sde(θ_rec,eq.σϵ)
  @test eq_rec.a≈eq.a atol=1e-10

  problem = ProblemLike(data;p=order(proc),engine=SdeRootsEngine)
  proc′ = transform_ss(problem,proc)
  𝓣 = transformation(problem,proc′)
  𝓣⁻¹ = inverse(𝓣)
  θ = paramstuple(problem,proc′)
  θ_rec = (𝓣∘𝓣⁻¹)(θ)
  eq′_rec = SdeRoots(θ_rec,eq.σϵ)
  @test eq′_rec.r≈proc′.eq.r atol=1e-8
end

@testset "initial estimates" begin
  proc = testprocess(:C)
  data = rand(MersenneTwister(12202),proc,200,1.0)
  problem = ProblemLike(data;p=order(proc)+2)
  Lwn,σ = logpdf_ml_white_noise(data,false)
  est_wn = WhiteNoise(σ)
  init_tors = [AddRoot();AddComplexRoot()]
  proc_prev = proc
  for m in init_tors
    esti = initial_estimate(m,problem,[proc_prev])
    L = logpdf(esti,problem.data)
    # NOTE: L>Lwn is not necessarily met for all signals and initial estimators
    # However, this condition IS met for the current test data and estimators.
    @test L>Lwn
  end
end

@testset "autodiff" begin
  # The numerical gradient computed in `ngradient` can be very inaccurate because it is
  # affected by the rouding errors in the Kalman likelihood computation.
  function ngradient(f, xs::AbstractArray...)
    grads = zero.(xs)
    for (x, Δ) in zip(xs, grads), i in 1:length(x)
      δ = 1e-5 # sqrt(eps())=1e-8 is too small; results in large rounding errors
      tmp = x[i]
      x[i] = tmp - δ/2
      y1 = f(xs...)
      x[i] = tmp + δ/2
      y2 = f(xs...)
      x[i] = tmp
      Δ[i] = (y2-y1)/δ
    end
    return grads
  end
  proc = SdeProcess(testprocess(:C).eq)
  data = rand(MersenneTwister(12202),proc,200,1.0)

  for E in [ElimQ,MatFrac,SdeRootsEngine]
    problem = ProblemLike(data;p=order(proc),engine=E)
    proc′ = transform_ss(problem,proc)
    θ = paramstuple(problem,proc′)
    𝓣 = transformation(problem,proc′)
    η = inverse(𝓣,θ)
    Zygote.refresh()
    yf,B = Zygote.pullback(problem∘𝓣,η)
    dη = B(1)[1]
    ng = ngradient(problem∘𝓣,η)[1]
    @test dη≈ng atol=0.5
  end
end

end # @testset EstimateSDE begin


text:
mutable struct BifurcationCurve{DType,SPType}
  points::DType
  stab::Vector{String}
  special_points::SPType
  changes::Vector{Int}
end

@recipe function f(bif::BifurcationCurve,coords)
  x = bif.points[coords[1]]
  y = bif.points[coords[2]]
  g = bif.stab
  style_order = unique(bif.stab)

  linestyle --> reshape([style == "S" ? :solid : :dash for style in style_order],1,length(style_order))
  color --> reshape([style == "S" ? :blue : :red for style in style_order],1,length(style_order))
  legend --> false
  xlabel --> coords[1]
  ylabel --> coords[2]

  for k in keys(bif.special_points)
    @series begin
      seriestype --> :scatter
      markersize --> 5
      color := :red
      [bif.special_points[k][coords[1]]],[bif.special_points[k][coords[2]]]
    end
  end
  CompositeLine((x, y, g))
end

mutable struct CompositeLine
  args
end

@recipe function f(h::CompositeLine)
    x, y, linegroups = h.args

    seriesx = Dict{eltype(linegroups), Vector{Float64}}()
    seriesy = Dict{eltype(linegroups), Vector{Float64}}()

    for i in 1:length(linegroups)
        if ! haskey(seriesx, linegroups[i])
            seriesx[linegroups[i]] = Float64[]
            seriesy[linegroups[i]] = Float64[]
        end

        push!(seriesx[linegroups[i]], x[i])
        push!(seriesy[linegroups[i]], y[i])

        if i > 1 && linegroups[i] != linegroups[i-1]
            # End the previous series with the same point to avoid gaps
            push!(seriesx[linegroups[i-1]], x[i])
            push!(seriesy[linegroups[i-1]], y[i])

            # Add an NaN to stop the group
            push!(seriesx[linegroups[i-1]], NaN)
            push!(seriesy[linegroups[i-1]], NaN)

        end
    end

    seriestype := :path

    for key in keys(seriesx)
        @series begin
            seriesx[key], seriesy[key]
        end
    end

end

function bifurcation_curve(PC,bif_type,freepars;max_num_points=450,
                          max_stepsize=2,min_stepsize=1e-5,
                          stepsize=2e-2,loc_bif_points="all",
                          save_eigen=true,name="DefaultName",
                          print_info=true,calc_stab=true,
                          var_tol = 1e-6, func_tol = 1e-6,
                          test_tol = 1e-4,
                          initpoint=nothing,solver_sequence=[:forward])

  curve_point_type = bif_type[1:end-2]

  if !(typeof(freepars)<:AbstractArray)
    freepars = [freepars]
  end

  # Setup Parameters
  PCargs = ds[:args](name=name)
  PCargs[:type]         = bif_type
  PCargs[:freepars]     = freepars
  PCargs[:MaxNumPoints] = max_num_points
  PCargs[:MaxStepSize]  = max_stepsize
  PCargs[:MinStepSize]  = min_stepsize
  PCargs[:StepSize]     = stepsize
  PCargs[:LocBifPoints] = loc_bif_points
  PCargs[:SaveEigen]    = save_eigen
  PCargs[:VarTol]       = var_tol
  PCargs[:FuncTol]      = func_tol
  PCargs[:TestTol]      = test_tol
  if initpoint != nothing
    PCargs[:initpoint] = initpoint
  end

  # Run Solver
  PC[:newCurve](PCargs)
  for step in solver_sequence
    PC[:curves][name][step]()
  end

  # Print Info
  if print_info
    PC[:curves][name][:info]()
  end

  # Get the curve
  pts = PyDict(PC[:curves][name][:_curveToPointset]())
  points = OrderedDict{Symbol,Vector{Float64}}()
  for k in keys(pts)
    points[Symbol(k)] = pts[k]
  end
  len = length(points[first(keys(points))])

  # Get the stability
  # S => Stable
  # U => Unstable
  # N => Neutral

  # Get this from the information at:
  # https://github.com/robclewley/pydstool/blob/master/PyDSTool/PyCont/ContClass.py#L218
  curve = PC[:curves][name]
  if calc_stab
    stab = [curve[:CurveInfo][i][curve_point_type]["stab"] for i in 1:len]
  else
    stab = []
  end

  # Get information for special points, ex limit points
  special_points = Dict{String,Any}()
  for k in keys(curve[:BifPoints])
    for i in 1:length(curve[:BifPoints][k][:found])
      tmp_dict = PyDict(curve[:BifPoints][k][:found][i]["X"])
      dd = Dict{Symbol,Float64}()
      for k2 in keys(tmp_dict)
        dd[Symbol(k2)]=tmp_dict[k2]
      end
      special_points[k*string(i)] = dd
    end
  end
  #=
  # Start and endpoints
  for i in 1:len
    if "P" in keys(curve[i])
      tmp_dict = PyDict(curve[i]["P"]["data"]["V"])
      dd = Dict{Symbol,Float64}()
      for k in keys(tmp_dict)
        dd[Symbol(k)]=tmp_dict[k]
      end
      special_points[curve[i]["P"]["name"]] = dd
    end
  end
  =#

  changes = find_changes(stab)

  BifurcationCurve(points,stab,special_points,changes)
end

function find_changes(stab)
  changes = Int[]
  for i in 2:length(stab)
    if stab[i]!= stab[i-1]
      push!(changes,i)
    end
  end
  changes
end


text:
@testset "Load data misspecification" begin
    device_model = DeviceModel(InterruptibleLoad, DispatchablePowerLoad)
    c_sys5 = PSB.build_system(PSITestSystems, "c_sys5")
    info_message = "The data doesn't include devices of type InterruptibleLoad, consider changing the device models"
    model = DecisionModel(MockOperationProblem, DCPPowerModel, c_sys5)
    @test_logs (:info, info_message) match_mode = :any mock_construct_device!(
        model,
        device_model,
    )
    device_model = DeviceModel(PowerLoad, DispatchablePowerLoad)
    warn_message = "The Formulation DispatchablePowerLoad only applies to FormulationControllable Loads, \n Consider Changing the Device Formulation to StaticPowerLoad"
    model = DecisionModel(MockOperationProblem, DCPPowerModel, c_sys5)
    @test_logs (:warn, warn_message) match_mode = :any mock_construct_device!(
        model,
        device_model,
    )
end

@testset "StaticPowerLoad" begin
    models = [StaticPowerLoad, DispatchablePowerLoad, InterruptiblePowerLoad]
    c_sys5_il = PSB.build_system(PSITestSystems, "c_sys5_il")
    networks = [DCPPowerModel, ACPPowerModel]
    for m in models, n in networks
        device_model = DeviceModel(PowerLoad, m)
        model = DecisionModel(MockOperationProblem, n, c_sys5_il)
        mock_construct_device!(model, device_model)
        moi_tests(model, false, 0, 0, 0, 0, 0, false)
        psi_checkobjfun_test(model, GAEVF)
    end
end

@testset "DispatchablePowerLoad DC- PF" begin
    models = [DispatchablePowerLoad]
    c_sys5_il = PSB.build_system(PSITestSystems, "c_sys5_il")
    networks = [DCPPowerModel]
    for m in models, n in networks
        device_model = DeviceModel(InterruptibleLoad, m)
        model = DecisionModel(MockOperationProblem, n, c_sys5_il)
        mock_construct_device!(model, device_model)
        moi_tests(model, false, 24, 0, 24, 0, 0, false)
        psi_checkobjfun_test(model, GAEVF)
    end
end

@testset "DispatchablePowerLoad AC- PF" begin
    models = [DispatchablePowerLoad]
    c_sys5_il = PSB.build_system(PSITestSystems, "c_sys5_il")
    networks = [ACPPowerModel]
    for m in models, n in networks
        device_model = DeviceModel(InterruptibleLoad, m)
        model = DecisionModel(MockOperationProblem, n, c_sys5_il)
        mock_construct_device!(model, device_model)
        moi_tests(model, false, 48, 0, 24, 0, 24, false)
        psi_checkobjfun_test(model, GAEVF)
    end
end

@testset "InterruptiblePowerLoad DC- PF" begin
    models = [InterruptiblePowerLoad]
    c_sys5_il = PSB.build_system(PSITestSystems, "c_sys5_il")
    networks = [DCPPowerModel]
    for m in models, n in networks
        device_model = DeviceModel(InterruptibleLoad, m)
        model = DecisionModel(MockOperationProblem, n, c_sys5_il)
        mock_construct_device!(model, device_model)
        moi_tests(model, false, 48, 0, 24, 0, 0, true)
        psi_checkobjfun_test(model, GAEVF)
    end
end

@testset "InterruptiblePowerLoad AC- PF" begin
    models = [InterruptiblePowerLoad]
    c_sys5_il = PSB.build_system(PSITestSystems, "c_sys5_il")
    networks = [ACPPowerModel]
    for m in models, n in networks
        device_model = DeviceModel(InterruptibleLoad, m)
        model = DecisionModel(MockOperationProblem, n, c_sys5_il)
        mock_construct_device!(model, device_model)
        moi_tests(model, false, 72, 0, 24, 0, 24, true)
        psi_checkobjfun_test(model, GAEVF)
    end
end


text:
# Parameters:
#	FFT1::Array{Array{Float64, 2}, 1}
#	- the first Fourier transform
#	FFT2::Array{Array{Float64, 2}, 1}
#	- the second Fourier transform
# Return Values:
#	Convolution::Array{Array{Float64, 2}, 1}
#	- the convolution of FFT1 and FFT2
# Notes:
#	- FFT1 and FFT2 have to be Fourier transforms of functions over the same Sn
#	- However, the don't have to have the same degree of bandlimitedness
function sn_convolution(FFT1::Array{Array{Float64, 2}, 1}, FFT2::Array{Array{Float64, 2}, 1})
	l1 = length(FFT1)
	l2 = length(FFT2)
	L = min(l1, l2)
	Convolution = Array(Array{Float64, 2}, L)
	while L != 0
		Convolution[L] = FFT1[l1] * FFT2[l2]
		l1 -= 1
		l2 -= 1
		L -= 1
	end
	return Convolution
end


text:
using BandedMatrices, LinearAlgebra, Test, Random
import BandedMatrices: _BandedMatrix

# set prng to some value that avoids test failure
Random.seed!(0)
struct _foo <: Number end

@testset "Banded A\\b" begin
    @testset "banded" begin
        A = brand(5, 1, 1)
        b = Float64[1, 2, 3, 4, 5]

        # dense storage
        Af = Matrix(A)
        bf = copy(b)

        @test Af\bf ≈ A\b

        # check A\b has not overwritten b, but has made a copy
        @test b == [1, 2, 3, 4, 5]
    end

    @testset "advanced interface" begin
        # banded
        A = brand(5, 1, 1)
        b = rand(5)

        # dense storage
        Af = Matrix(A)
        bf = copy(b)
        L,U,p = lu(A)
        Lf,Uf,pf = lu(Af)
        @test L ≈ Lf # storage format is different
        @test U ≈ Uf
        @test p ≈ pf
        lua = lu(A)
        @test lu(lua) === lua
        @test Factorization{Float64}(lua) === lua
        @test Matrix(Factorization{Float32}(lua)) ≈ convert(Matrix{Float32}, Af)
        @test Matrix(copy(lua)) ≈ Af
        @test issuccess(lua)

        # note lu makes copies; these need revision
        # once the lapack storage is built in to a BandedMatrix
        @test Af\bf ≈ lu(A)\copy(b)
        @test Af\bf ≈ ldiv!(A, copy(b))
        @test Af\bf ≈ ldiv!(lu(A), copy(b))
        @test transpose(Af)\bf ≈ transpose(lu(A))\copy(b)
        @test adjoint(Af)\bf ≈ adjoint(lu(A))\copy(b)
    end

    @testset "conversion of inputs if needed" begin
        # factorisation performs conversion
        Ai = _BandedMatrix(rand(1:10, 3, 5), 5, 1, 1)
        @test eltype(lu(Ai)) == Float64

        # no op
        Af = _BandedMatrix(rand(Float32, 3, 5), 5, 1, 1)
        @test eltype(lu(Af)) == Float32

        # linear systems of integer data imply promotion
        Ai = _BandedMatrix(rand(1:10, 3, 5), 5, 1, 1)
        bi = collect(1:5)
        @test eltype(Ai\bi) == Float64
        # this code                     ≈ julia base
        @test Ai\bi                     ≈ Matrix(Ai)\copy(bi)
        @test lu(Ai)\bi             ≈ Matrix(Ai)\copy(bi)

        # check A\b makes a copy of b
        Ai = _BandedMatrix(rand(1:10, 3, 5), 5, 1, 1)
        bi = collect(1:5)
        Ai\bi
        @test bi == [1, 2, 3, 4, 5]
    end

    @testset "matrix must be square error" begin
        for (n, m) in zip([7, 5], [5, 7])
            A = brand(m, n, 1, 1)
            b = rand(m)
            @test_throws DimensionMismatch (A\b)
            @test_throws DimensionMismatch lu(A)\b
            @test_throws DimensionMismatch ldiv!(A, b)
            @test_throws DimensionMismatch ldiv!(lu(A), b)
        end
    end

    @testset "transposed algorithm" begin
        # banded
        A = brand(5, 1, 1)
        b = rand(5)

        # dense storage
        Af = Matrix(A)
        bf = copy(b)

        # note lu makes copies; these need revision
        # once the lapack storage is built in to a BandedMatrix
        @test Af'\bf ≈ ldiv!(transpose(lu(A)),  copy(b))
        @test Af'\bf ≈ ldiv!(lu(A'), copy(b))
    end

    @testset "complex input" begin
        # banded
        A = brand(5, 1, 1) + brand(5, 1, 1)*im
        b = rand(5) + rand(5)*im

        # dense storage
        Af = Matrix(A)
        bf = copy(b)

        # note lu makes copies; these need revision
        # once the lapack storage is built in to a BandedMatrix
        @test Af\bf  ≈ A\copy(b)
        @test Af\bf  ≈ lu(A)\copy(b)
        @test Af'\bf ≈ A'\copy(b)
        @test Af'\bf ≈ lu(A')\copy(b)
        @test Af\bf  ≈ ldiv!(lu(A),  copy(b))
        @test Af'\bf ≈ ldiv!(lu(A'), copy(b))
        @test transpose(Af)\bf ≈ ldiv!(transpose(lu(A)),  copy(b))
        @test adjoint(Af)\bf ≈ ldiv!(lu(A'), copy(b))
        @test adjoint(Af)\bf ≈ ldiv!(lu(A)', copy(b))
    end

    @testset "multiple rhs" begin
        # banded
        A = brand(5, 1, 1)
        b = rand(5, 10)

        # dense storage
        Af = Matrix(A)
        bf = copy(b)

        # note lu makes copies; these need revision
        # once the lapack storage is built in to a BandedMatrix
        @test Af\bf  ≈ A\copy(b)
        @test Af\bf  ≈ lu(A)\copy(b)
        @test Af'\bf ≈ A'\copy(b)
        @test Af'\bf ≈ lu(A')\copy(b)
        @test Af\bf  ≈ ldiv!(lu(A),  copy(b))
        @test Af'\bf ≈ ldiv!(lu(A'), copy(b))
        @test Af'\bf ≈ ldiv!(transpose(lu(A)),  copy(b))
        @test Af'\bf ≈ ldiv!(lu(A'), copy(b))
    end

    @testset "properties of factorisation" begin
        BLU = lu(brand(5, 4, 1, 1))
        @test size(BLU) == (5, 4)
        @test size(BLU, 1) == 5
        @test size(BLU, 2) == 4
        @test size(BLU, 3) == 1
        @test_throws BoundsError size(BLU, -1)
        @test_throws BoundsError size(BLU,  0)
    end
end


text:
# This file was generated, do not modify it.

using Pkg # hideall
Pkg.activate("_literate/A-model-choice/Project.toml")
Pkg.update()

using RDatasets
using MLJ
MLJ.color_off() # hide
iris = dataset("datasets", "iris")

first(iris, 3) |> pretty

iris2 = coerce(iris, :PetalWidth => OrderedFactor)
first(iris2[:, [:PetalLength, :PetalWidth]], 1) |> pretty

y, X = unpack(iris, ==(:Species))
first(X, 1) |> pretty

y, X = unpack(iris, ==(:Species), !=(:PetalLength))
first(X, 1) |> pretty

X, y = @load_iris;

for m in models(matching(X, y))
    if m.prediction_type == :probabilistic
        println(rpad(m.name, 30), "($(m.package_name))")
    end
end

knc = @load KNeighborsClassifier

linreg = @load LinearRegressor pkg=GLM

max_issues_repo_name: alan-turing-institute/MLJTutorials

max_forks_repo_name: alan-turing-institute/MLJTutorials


text:
module TestImages
using Images

export testimage

const imagedir = joinpath(dirname(@__FILE__), "..", "images")

REPO_URL = "https://github.com/timholy/TestImages.jl/blob/gh-pages/images/"

remotefiles = [
    "autumn_leaves.png" ,
    "blobs.gif" ,
    "cameraman.tif" ,
    "earth_apollo17.jpg" ,
    "fabio_color_256.png" ,
    "fabio_color_512.png" ,
    "fabio_gray_256.png" ,
    "fabio_gray_512.png" ,
    "hela-cells.tif" ,
    "house.tif" ,
    "jetplane.tif" ,
    "lake.tif" ,
    "lena_color_256.tif" ,
    "lena_color_512.tif" ,
    "lena_gray_256.tif" ,
    "lena_gray_512.tif" ,
    "lena_gray_16bit.png" ,
    "livingroom.tif" ,
    "lighthouse.png" ,
    "mandril_color.tif" ,
    "mandril_gray.tif" ,
    "mandrill.tiff" ,
    "m51.tif" ,
    "moonsurface.tiff" ,
    "mountainstream.png" ,
    "mri-stack.tif" ,
    "multi-channel-time-series.ome.tif" ,
    "peppers_color.tif" ,
    "peppers_gray.tif" ,
    "pirate.tif" ,
    "toucan.png" ,
    "walkbridge.tif" ,
    "woman_blonde.tif" ,
    "woman_darkhair.tif" ,
]

function testimage(filename, ops...)
    imagefile = joinpath(imagedir, filename)
    if !isfile(imagefile)
        fls = readdir(imagedir)
        havefile = false
        for f in fls
            if startswith(f, filename)
                imagefile = joinpath(imagedir, f)
                havefile = true
                break
            end
        end

        if !havefile
            info("Could not find "*filename*" in directory. Checking if it exists in the online repository.")
            for f in remotefiles
                if startswith(f, filename)
                    info("Found "*filename*" in the online repository. Downloading to the images directory.")
                    download(REPO_URL*f*"?raw=true", joinpath(imagedir, f))
                    havefile = true
                    imagefile = joinpath(imagedir, f)
                    break
                end
            end
        end
        havefile || error("$filename not found in the directory or the online repository. Here are the contents of the images/ directory:\n$(join(fls, '\n'))")
    end
    load(imagefile, ops...)
end
end


text:
# This file is a part of Julia. License is MIT: https://julialang.org/license

module PkgTests

include("utils.jl")
include("pkg.jl")
include("binaryplatforms.jl")
include("platformengines.jl")
include("sandbox.jl")
include("resolve.jl")

# clean up locally cached registry
rm(joinpath(@__DIR__, "registries"); force = true, recursive = true)

end # module


text:
Base.convert(::Type{T}, x::AbstractString) where {T<:AbstractICD10} = T(x)

Base.print(io::IO, icd::T, punct = ICDOPTS[:punct]) where {T<:AbstractICD10} =
  if punct
    if icd.level == 3
      write(io, Ref(icd.data[1:3]), '.')
    elseif icd.level == 4
      write(io, Ref(icd.data[1:3]), '.', Ref(icd.data[4]))
    elseif icd.level == 5
      write(io, Ref(icd.data[1:3]), '.', Ref(icd.data[4:5]))
    elseif icd.level == 6
      write(io, Ref(icd.data[1:3]), '.', Ref(icd.data[4:6]))
    elseif icd.level == 7
      write(io, Ref(icd.data[1:3]), '.', Ref(icd.data[4:7]))
    end
  else
    write(io, Ref(icd.data[1:icd.level]))
  end

Base.show(
  io::IO,
  ::MIME"text/plain",
  icd::T,
  punct = ICDOPTS[:punct],
) where {T<:AbstractICD10} = print(io, icd)

Base.isless(icd1::AbstractICD10, icd2::AbstractICD10) = icd1.data < icd2.data

Base.length(icd::T) where {T<:AbstractICD10} = icd.level

Base.:(==)(icd1::AbstractICD10, icd2::AbstractICD10) =
  icd1.level == icd2.level && icd1.data == icd2.data

## ICD10 functions #####

"""
    isvalidcode(icd, validcodes)

Takes an AbstractICD10 and compares it with a list of valid codes.
"""
isvalidcode(icd::T, validcodes) where {T<:AbstractICD10} = icd in validcodes

"""
    icd3(icdcode)

Returns the initial 3 digits of an ICD10 code. Use for
comparing codes just on first 3 digits.
"""
icd3(icdcode::T) where {T<:AbstractICD10} = icdcode.data[1:3]

"""
    icd4(icdcode)

Returns the initial 4 digits of an ICD10 code. Use for
comparing codes just on first 4 digits.
"""
icd4(icdcode::T) where {T<:AbstractICD10} = icdcode.data[1:4]


text:
# automatically generated by the FlatBuffers compiler, do not modify

MyGame.Example.eval(quote

FlatBuffers.@STRUCT struct Ability
    id::UInt32
    distance::UInt32
end
FlatBuffers.@ALIGN(Ability, 4)

Ability(buf::AbstractVector{UInt8}) = FlatBuffers.read(Ability, buf)
Ability(io::IO) = FlatBuffers.deserialize(io, Ability)

end)

max_stars_repo_name: rjkat/flatbuffers-julia

max_issues_repo_name: rjkat/flatbuffers-julia

max_forks_repo_name: rjkat/flatbuffers-julia


text:
import Base.copy
using Parameters

export
    solver_name

abstract type AbstractSolver{T} end
abstract type AbstractSolverOptions{T<:Real} end

include("solvers/ilqr/ilqr_solver.jl")
include("solvers/ilqr/ilqr_methods.jl")
include("solvers/ilqr/backward_pass.jl")
include("solvers/ilqr/forward_pass.jl")

include("solvers/augmented_lagrangian/augmented_lagrangian_solver.jl")
include("solvers/augmented_lagrangian/augmented_lagrangian_methods.jl")

include("solvers/direct/direct_solvers.jl")
include("solvers/direct/sequential_newton.jl")
include("solvers/direct/dircol.jl")
include("solvers/direct/dircol_ipopt.jl")
include("solvers/direct/dircol_snopt.jl")
include("solvers/direct/moi.jl")
include("solvers/direct/sequential_newton_solve.jl")
include("solvers/direct/projected_newton.jl")

include("solvers/altro/altro_solver.jl")
include("solvers/altro/altro_methods.jl")
include("solvers/altro/infeasible.jl")
include("solvers/altro/minimum_time.jl")

include("solvers/direct/primals_mintime.jl")
include("solvers/direct/direct_solvers_mintime.jl")
include("solvers/direct/dircol_mintime.jl")
include("solvers/direct/moi_mintime.jl")

# Get name of solver as a string
solver_name(::iLQRSolverOptions) = "iLQR"
solver_name(::ALTROSolverOptions) = "ALTRO"
solver_name(opts::DIRCOLSolverOptions) = string(optimizer_name(opts.nlp))
solver_name(opts::AugmentedLagrangianSolverOptions) = "AL-" * solver_name(opts.opts_uncon)
solver_name(solver::AbstractSolver) = solver_name(solver.opts)

# Solver interface
"""$(SIGNATURES)
Solve trajectory optimization problem `prob` using `solver`.
    The problem will be modified in place, with the solution stored in `prob.X` and `prob.U`.
    The solver will also be modified, and may either return the same solver or a new one that may or may not be the same type as the one given.
"""
solve!(prob::Problem{T,D}, solver::AbstractSolver{T}) where {T<:AbstractFloat,D<:DynamicsType} =
    error("Cannot solve with `AbstractSolver`")

"```
AbstractSolver(prob::Problem, opts::AbstractSolverOptions)::AbsractSolver
```
Create a solver, with the type specified by the type of the solver options `opts` "
AbstractSolver(::Problem{T,D}, ::AbstractSolverOptions{T}) where {T<:AbstractFloat,D<:DynamicsType} =
    error("Can't create an Abstract Solver without knowing the type of the Solver Options")

"```
reset!(solver::AbstractSolver)
```
Reset the solver, including initial values for the fields and solve statistics"
reset!(::AbstractSolver) = nothing

"```
copy(solver::AbstractSolver)::AbstractSolver
```
Create a copy of the solver with zero associated memory between the two solvers"
copy(::AbstractSolver)::AbstractSolver = error("Cannot copy `AbstractSolver`")

"```
size(solver::AbstractSolver)::NTuple{3,Int}
```
 Return the number of controls (n), number of states (m), and the number of knot points (N) as a tuple, i.e. (n,m,N)"
size(::AbstractSolver)::NTuple{3,Int} = error("`AbstractSolver` has no size")

# Generic methods for calling solve
"""```
solve!(prob, opts)::AbstractSolver
```
Solve the trajectory optimization problem `prob` using the solver specified by solver options `opts`.
    The problem will be modified in place, storing the solution in `prob.X` and `prob.U`.
"""
function solve!(prob::Problem{T,D}, opts::AbstractSolverOptions{T}) where {T<:AbstractFloat, D<:DynamicsType}
    solver = AbstractSolver(prob, opts)
    solve!(prob, solver)
end

"""```
solve(prob, opts)::Tuple{Problem,AbstractSolver}
```
Solve the trajectory optimization problem `prob` using the solver specified by solver options `opts`.
    The problem will not be modified in place,
    and will instead return a new problem with the solution in `prob.X` and `prob.U`,
    along with the solver.
"""
function solve(prob::Problem{T,D}, opts::AbstractSolverOptions{T}) where {T<:AbstractFloat, D<:DynamicsType}
    prob0 = copy(prob)
    solver = solve!(prob0, opts)
    return prob0, solver
end

""" ```
solve(prob, solver)::Tuple{Problem,AbstractSolver}
```
Solve the trajectory optimization problem `prob` using `solver`.
    The problem will not be modified in place,
    and will instead return a new problem with the solution in `prob.X` and `prob.U`,
    along with the solver. The solver will be modified in place,
        and may or may not be the same solver returned.
"""
function solve(prob::Problem{T,D}, solver::AbstractSolver{T}) where {T<:AbstractFloat, D<:DynamicsType}
    prob0 = copy(prob)
    solver = solve!(prob0, solver)
    return prob0, solver
end

jacobian!(prob::Problem{T,Continuous}, solver::AbstractSolver) where T = jacobian!(solver.∇F, prob.model, prob.X, prob.U)
jacobian!(prob::Problem{T,Discrete},   solver::AbstractSolver) where T = jacobian!(solver.∇F, prob.model, prob.X, prob.U, get_dt_traj(prob))

function check_convergence_criteria(opts_uncon::AbstractSolverOptions{T},cost_tolerance::T,gradient_norm_tolerance::T) where T
    if opts_uncon.cost_tolerance != cost_tolerance
        @warn "Augmented Lagrangian cost tolerance overriding unconstrained solver option\n >>cost tolerance=$cost_tolerance"
    end

    if opts_uncon.gradient_norm_tolerance != gradient_norm_tolerance
        @warn "Augmented Lagrangian gradient norm tolerance overriding unconstrained solver option\n >>gradient norm tolerance=$gradient_norm_tolerance"
    end
    return nothing
end

max_issues_repo_name: MehrdadZareian/GMARA

max_forks_repo_name: MehrdadZareian/GMARA


text:
const GenericTapirEx = Union{TapirEx,StaticTapirEx}

function Base.NamedTuple(ex::GenericTapirEx)
    nts = ntuple(Val(nfields(ex))) do i
        (; fieldname(typeof(ex), i) => getfield(ex, i))
    end
    return foldl(merge, nts)
end

function Base.show(io::IO, ex::GenericTapirEx)
    @nospecialize ex
    print(io, nameof(typeof(ex)))
    print(io, NamedTuple(ex))
end

# executor_from(::Type{Ex}, exc::PreferParallel) API
(::Type{Ex})(kwargs::NamedTuple) where {Ex <: GenericTapirEx} = Ex(; kwargs...)

# FIXME: Compatibility for Folds.Testing:
# <Workaround>
function Base.getproperty(ex::GenericTapirEx, field::Symbol)
    if field === :kwargs
        return NamedTuple(ex)
    else
        return getfield(ex, field)
    end
end

ConstructionBase.getproperties(ex::GenericTapirEx) = (kwargs = NamedTuple(ex),)

ConstructionBase.setproperties(ex::GenericTapirEx, patch::NamedTuple{(:kwargs,)}) =
    ConstructionBase.constructorof(typeof(ex))(; patch.kwargs...)
# </Workaround>


text:
let
    f = Hop.Floquet.FloquetHamiltonian([ones(2, 2), 2*ones(2, 2)*im], 1, 2)
    @test f.Hf[1:2, 1:2] ≈ ones(2, 2)+2*I
    @test f.Hf[5:6, 3:4] ≈ 2*ones(2, 2)*im
    @test f.Hf[3:4, 5:6] ≈ -2*ones(2, 2)*im
end

let
    graphene = getgraphene()
    A = [1.0, 0.0, 0.0]
    Ω = 1.0
    Γ = 1.0
    μ = 0.0
    f = Hop.Floquet.get_illuminated_hamiltonian(graphene, [0.2, 0.0, 0.0], A=A, Ω=Ω, harmonics_cutoff=5)
    @test Hop.Floquet.get_floquet_occupation(f; μ=μ, Γ=Γ) ≈ [0.23, 0.77] atol=1.0e-2
end

let
    graphene = getgraphene()
    kp = KPath(graphene.rlat, [1 0; 0 1; 0 0], 10)
    bands = Hop.Floquet.get_illuminated_band(graphene, A=[1.0, 0.0, 0.0],
        Ω=1.0, kp=kp, harmonics_cutoff=5)
    @test bands[21, 4] ≈ 5.589253 atol=1.0e-6
end


text:
module Grid
import Logging
import FFTW
import Printf: @sprintf
import Luna: PhysData, Maths

abstract type AbstractGrid end

abstract type TimeGrid <: AbstractGrid end
abstract type SpaceGrid <: AbstractGrid end

struct RealGrid <: TimeGrid
    zmax::Float64
    referenceλ::Float64
    t::Array{Float64, 1}
    ω::Array{Float64, 1}
    to::Array{Float64, 1}
    ωo::Array{Float64, 1}
    sidx::BitArray{1}
    ωwin::Array{Float64, 1}
    twin::Array{Float64, 1}
    towin::Array{Float64, 1}
end

"""
    RealGrid(zmax, referenceλ, λ_lims, trange; δt=1)

Time grid for simulations with real-valued (field-resolved) fields

# Arguments
- `zmax::Real` : Total distance to propagate
- `referenceλ::Real` : Reference wavelength (e.g. centre wavelength of input pulse)
- `λ_lims::Tuple{Real, Real}` : Wavelength limits of the frequency window
- `trange::Real` : Total extent of the time window required
- `δt::Real` : Sample spacing in time. The value actually used is either δt or the value
    required to satisfy `trange` and `λ_lims`, whichever is smaller.
"""
function RealGrid(zmax, referenceλ, λ_lims, trange, δt=1)
    f_lims = PhysData.c./λ_lims
    Logging.@info @sprintf("Freq limits %.2f - %.2f PHz", f_lims[2]*1e-15, f_lims[1]*1e-15)
    δto = min(1/(6*maximum(f_lims)), δt) # 6x maximum freq, or user-defined if finer
    samples = 2^(ceil(Int, log2(trange/δto))) # samples for fine grid (power of 2)
    trange_even = δto*samples # keep frequency window fixed, expand time window as necessary
    Logging.@info @sprintf("Samples needed: %.2f, samples: %d, δt = %.2f as",
                            trange/δto, samples, δto*1e18)
    Logging.@info @sprintf("Requested time window: %.1f fs, actual time window: %.1f fs", trange*1e15, trange_even*1e15)
    δωo = 2π/trange_even # frequency spacing for fine grid
    # Make fine grid
    Nto = collect(range(0, length=samples))
    to = @. (Nto-samples/2)*δto # centre on 0
    Nωo = collect(range(0, length=Int(samples/2 +1)))
    ωo = Nωo*δωo

    ωmin = 2π*minimum(f_lims)
    ωmax = 2π*maximum(f_lims)
    ωmax_win = 1.1*ωmax

    cropidx = findfirst(x -> x>ωmax_win, ωo)
    cropidx = 2^(ceil(Int, log2(cropidx))) + 1 # make coarse grid power of 2 as well
    ω = ωo[1:cropidx]
    δt = π/maximum(ω)
    tsamples = (cropidx-1)*2
    Nt = collect(range(0, length=tsamples))
    t = @. (Nt-tsamples/2)*δt

    # Make apodisation windows
    ωwindow = Maths.planck_taper(ω, ωmin/2, ωmin, ωmax, ωmax_win)

    twindow = Maths.planck_taper(t, minimum(t), -trange/2, trange/2, maximum(t))
    towindow = Maths.planck_taper(to, minimum(to), -trange/2, trange/2, maximum(to))

    # Indices to select real frequencies (for dispersion relation)
    sidx = (ω .> ωmin/2) .& (ω .< ωmax_win)

    @assert δt/δto ≈ length(to)/length(t)
    @assert δt/δto ≈ maximum(ωo)/maximum(ω)

    Logging.@info @sprintf("Grid: samples %d / %d, ωmax %.2e / %.2e",
                           length(t), length(to), maximum(ω), maximum(ωo))
    return RealGrid(float(zmax), referenceλ, t, ω, to, ωo, sidx, ωwindow, twindow, towindow)
end

function RealGrid(;zmax, referenceλ, t, ω, to, ωo, sidx, ωwin, twin, towin)
    RealGrid(zmax, referenceλ, t, ω, to, ωo, sidx, ωwin, twin, towin)
end

struct EnvGrid{T} <: TimeGrid
    zmax::Float64
    referenceλ::Float64
    ω0::Float64
    t::Array{Float64, 1}
    ω::Array{Float64, 1}
    to::Array{Float64, 1}
    ωo::Array{Float64, 1}
    sidx::T
    ωwin::Array{Float64, 1}
    twin::Array{Float64, 1}
    towin::Array{Float64, 1}
end

"""
    EnvGrid(zmax, referenceλ, λ_lims, trange; δt=1, thg=false)

Time grid for simulations with envelope (a.k.a. analytic) fields

# Arguments
- `zmax::Real` : Total distance to propagate
- `referenceλ::Real` : Reference wavelength (e.g. centre wavelength of input pulse)
- `λ_lims::Tuple{Real, Real}` : Wavelength limits of the frequency window
- `trange::Real` : Total extent of the time window required
- `δt::Real` : Sample spacing in time. The value actually used is either δt or the value
    required to satisfy `trange` and `λ_lims`, whichever is smaller.
- `thg::Bool` : Whether the grid should include space for the third hamonic (default: false)
"""
function EnvGrid(zmax, referenceλ, λ_lims, trange; δt=1, thg=false)
    fmin = PhysData.c/maximum(λ_lims)
    fmax = PhysData.c/minimum(λ_lims)
    fmax_win = 1.1*fmax # extended frequency window to accommodate apodisation
    ω0 = 2π*PhysData.c/referenceλ # central frequency
    Logging.@info @sprintf("Freq limits %.2f - %.2f PHz", fmin*1e-15, fmax*1e-15)
    if thg
        ffac = 6 # need to sample at 6x maximum desired frequency
        f_lims = (fmin, fmax)
    else
        ffac = 2 # need to sample at Nyquist limit for desired frequency only
        # Frequency window is not extended without THG, so need to extend additionally
        # to accommodate apodisation
        f_lims = (fmin, fmax_win)
    end
    rf_lims = f_lims .- PhysData.c/referenceλ # Relative frequency limits
    δt_f = 1/(ffac*maximum(rf_lims)) # time spacing as required by frequency window
    if δt_f <= δt
        # Demands of frequency window are more stringent
        δto = δt_f
        oversampling = thg # if no thg, then no oversampling
    else
        # User-defined time spacing is more stringent
        δto = δt
        oversampling = true
    end
    samples = 2^(ceil(Int, log2(trange/δto))) # samples for grid (power of 2)
    trange_even = δto*samples # keep frequency window fixed, expand time window as necessary
    Logging.@info @sprintf("Samples needed: %.2f, samples: %d, δt = %.2f as",
    trange/δto, samples, δto*1e18)
    δωo = 2π/trange_even # frequency spacing for grid
    # Make fine grid
    No = collect(range(0, length=samples))
    to = @. (No-samples/2)*δto # time grid, centre on 0
    vo = @. (No-samples/2)*δωo # freq grid relative to ω0
    vo = FFTW.fftshift(vo)
    ωo = vo .+ ω0

    ωmin = 2π*fmin
    ωmax = 2π*fmax
    ωmax_win = 2π*fmax_win

    # Find cropping area for coarse grid (contains frequencies of interest + apodisation)
    if oversampling
        cropidx = findfirst(x -> x>=(ωmax_win-δωo), ωo)
        cropidx = 2^(ceil(Int, log2(cropidx))) # make coarse grid power of 2 as well
        v = vcat(vo[1:cropidx], vo[end-cropidx+1:end])
    else
        v = vo
    end

    δt = -π/minimum(v)
    tsamples = length(v)
    Nt = collect(range(0, length=tsamples))
    t = @. (Nt - tsamples/2)*δt

    ω = v .+ ω0 # True frequency grid
    # Indices to select real frequencies (for dispersion relation)
    sidx = (ω .> ωmin/2) .& (ω .< ωmax_win)

    # Make apodisation windows
    ωwindow = Maths.planck_taper(ω, ωmin/2, ωmin, ωmax, ωmax_win)
    twindow = Maths.planck_taper(t, minimum(t), -trange/2, trange/2, maximum(t))
    towindow = Maths.planck_taper(to, minimum(to), -trange/2, trange/2, maximum(to))

    # Check that grids are correct
    @assert δt/δto ≈ length(to)/length(t)
    @assert δt/δto ≈ minimum(vo)/minimum(v) # FFT grid -> sample at -fs/2 but not +fs/2
    factor = Int(length(to)/length(t))
    zeroidx = findfirst(x -> x==0, to)
    # The time samples should be exactly the same (except fewer in t)
    @assert all(to[zeroidx:factor:end] .≈ t[t .>= 0])
    @assert all(to[zeroidx:-factor:1] .≈ t[t .<= 0][end:-1:1])

    return EnvGrid(float(zmax), referenceλ, ω0, t, ω, to, ωo, sidx, ωwindow, twindow, towindow)
end

function EnvGrid(;zmax, referenceλ, ω0, t, ω, to, ωo, sidx, ωwin, twin, towin)
    EnvGrid(zmax, referenceλ, ω0, t, ω, to, ωo, sidx, ωwin, twin, towin)
end

struct FreeGrid <: SpaceGrid
    x::Vector{Float64}
    y::Vector{Float64}
    kx::Vector{Float64}
    ky::Vector{Float64}
    r::Array{Float64, 3}
    xywin::Array{Float64, 3}
end

"""
    FreeGrid(Rx, Nx, Ry, Ny; window_factor=0.1)

Spatial grid for full 3D freespace propagation with `x`/`y` half-width `Rx`/`Ry` and
`Nx`/`Ny` samples. `window_factor` determines by how much the grid size is extended to fit
a filtering window.
"""
function FreeGrid(Rx, Nx, Ry, Ny; window_factor=0.1)
    Rxw = Rx * (1 + window_factor)
    Ryw = Ry * (1 + window_factor)

    δx = 2Rxw/Nx
    nx = collect(range(0, length=Nx))
    x = @. (nx-Nx/2) * δx
    kx = 2π*FFTW.fftfreq(Nx, 1/δx)

    δy = 2Ryw/Ny
    ny = collect(range(0, length=Ny))
    y = @. (ny-Ny/2) * δy
    ky = 2π*FFTW.fftfreq(Ny, 1/δy)

    r = sqrt.(reshape(y, (1, Ny)).^2 .+ reshape(x, (1, 1, Nx)).^2)

    xwin = Maths.planck_taper(x, -Rxw, -Rx, Rx, Rxw)
    ywin = Maths.planck_taper(y, -Ryw, -Ry, Ry, Ryw)
    xywin = reshape(xwin, (1, length(xwin))) .* reshape(xwin, (1, 1, length(xwin)))

    FreeGrid(x, y, kx, ky, r, xywin)
end

FreeGrid(R, N) = FreeGrid(R, N, R, N)

function to_dict(g::GT) where GT <: AbstractGrid
    d = Dict{String, Any}()
    for field in fieldnames(GT)
        d[string(field)] = getfield(g, field)
    end
    d
end

function from_dict(gridtype, d)
    kwargs = (Symbol(k) => v for (k, v) in pairs(d))
    grid = gridtype(;kwargs...)

    # Make sure the grid is valid
    validate(grid)
    return grid
end

RealGrid(d::AbstractDict) = from_dict(RealGrid, d)
EnvGrid(d::AbstractDict) = from_dict(EnvGrid, d)

function validate(grid::TimeGrid)
    δt = grid.t[2] - grid.t[1]
    δto = grid.to[2] - grid.to[1]
    @assert δt/δto ≈ length(grid.to)/length(grid.t)
    @assert length(grid.towin) == length(grid.to)
    @assert length(grid.ωwin) == length(grid.ω)
    @assert length(grid.sidx) == length(grid.ω)
    if grid isa EnvGrid
        δω = grid.ω[2] - grid.ω[1]
        Δω = length(grid.ω)*δω
        @assert δt ≈ 2π/Δω
        @assert length(grid.t) == length(grid.ω)
        @assert length(grid.to) == length(grid.ωo)
    else
        Δω = maximum(grid.ω)
        @assert δt ≈ π/Δω
        @assert length(grid.t) == 2*(length(grid.ω)-1)
        @assert length(grid.to) == 2*(length(grid.ωo)-1)
    end
end

end


text:
let
    θ = 1.0*π
    @test θ == wrap_to_pi(θ)
    @test θ == wrap_to_pi(-θ)

    θ = 20.0
    @test -π < wrap_to_pi(θ) <= π
end

let
    v1 = [0,0]
    v2 = [0,1]
    pts = GraphUtils.nearest_points_between_circles(v1,v2,0.0,0.0)
    @test all(isapprox.(pts[1],v1))
    @test all(isapprox.(pts[2],v2))
    pts = GraphUtils.nearest_points_between_circles(v1,v2,0.5,0.5)
    @test all(isapprox.(pts[1],(v1+v2)/2))
    @test all(isapprox.(pts[2],(v1+v2)/2))

    pts = GraphUtils.nearest_points_between_circles(v1,v2,1,4)
    @test all(isapprox.(pts[1],[0.0,-1.0]))
    @test all(isapprox.(pts[2],[0.0,-3.0]))
    pts = GraphUtils.nearest_points_between_circles(v1,v2,4,1)
    @test all(isapprox.(pts[1],[0.0,4.0]))
    @test all(isapprox.(pts[2],[0.0,2.0]))

    pts = GraphUtils.nearest_points_between_circles(v1,v2,sqrt(2)/2,sqrt(2)/2)
    @test isapprox(abs(pts[1][1]),0.5)
    @test isapprox(abs(pts[2][1]),0.5)

end


text:
using Convex
using ECOS
# Formulate and solve a quadratic problem:
#     min ||g||^2 st. Afriat inequalities for random d

# Initialize W^c=W (consumption measurement error).
W=ones(M,T,N)
# Initialize simulated q^*=q.
qsim=zeros(M,T,N)
# Initialize simulated Afriat numbers δ=δsim.
δsim=zeros(M,T,N)
# Initialize optimal value to a high number
optimval=ones(n,ndelta+1)*10000
# Initialize matrix to verify Afriat inequalities of the simulated data.
aiverify2=zeros(n,T,T)
# Convex model, variable definition.
v=Variable(T, Positive())
c=Variable(T,K,Positive())
# Identity Matrix in Convex.
P=I+zeros(1,1)

# Loop to check all entries of deltavec per individual and save in Delta the
# value corresponding to the lowest objective value
for dt=2:ndelta+1
    for id=1:n
        Deltatemp[id]=deltavec[dt-1]
        # Convex objective value, see Appendix C.
        modvex=minimize(quadform(rho[id,1,:]'*(c[1,:]'-cve[id,1,:]),P)
            +quadform(rho[id,2,:]'*(c[2,:]'-cve[id,2,:]),P)
            +quadform(rho[id,3,:]'*(c[3,:]'-cve[id,3,:]),P)
            +quadform(rho[id,4,:]'*(c[4,:]'-cve[id,4,:]),P))
        # Define the Afriat's constraints.
        for t=1:T
            for s=1:T
                modvex.constraints+=v[t]-v[s]-Deltatemp[id]^(-(t-1))*rho[id,t,:]'*(c[t,:]'-c[s,:]')>=0
            end
        end
        # Solve the model using ECOS.
        solve!(modvex,() -> ECOS.Optimizer(verbose=false))
        # Save objective value.
        optimval[id,dt]=modvex.optval
        # Re-initialize the verification matrix.
        aiverify=zeros(n,T,T)

        # Keep the lowest objective value, associated Delta
        if (optimval[id,dt]<optimval[id,dt-1])
            Delta[id]=Deltatemp[id]
            for i=1:T
                vsim[id,i]=v.value[i]
                for j=1:K
                    cvesim[id,i,j]=c.value[i,j]
                end
            end
            # Verify Afriat's inequalities with the solution.
            for t=1:T
                for s=1:T
                    aiverify2[id,t,s]=vsim[id,t]-vsim[id,s]-Delta[id]^(-(t-1))*rho[id,t,:]'*(cvesim[id,t,:]-cvesim[id,s,:])
                end
            end
        end
    end
    #Memory cleaning.
    modvex=nothing
    GC.gc()
end

minimum(aiverify2)

max_stars_repo_name: hans-mtz/SCR

max_issues_repo_name: hans-mtz/SCR

max_forks_repo_name: hans-mtz/SCR


text:
using Convex, ECOS

"""
	contact constraints
		-differentiate through second-order cone program to get friction
"""
function friction_socp(model, x⁺, u, h) # version for particle
	q3 = view(x⁺, model.nq .+ (1:model.nq))
	q2 = view(x⁺, 1:model.nq)
	λ = view(u, model.idx_λ)

	v = (q3 - q2)[1:2] ./ h #
	y = model.μ * λ

	"Convex.jl"
	b = Variable(2)
	prob = minimize(v' * b)
	prob.constraints += norm(b) <= y
	solve!(prob, ECOS.Optimizer(verbose = false));

	# @show prob.status
	# @show b.value
	# @show prob.constraints[1].dual
	# prob.optval

	# # norms
	# _norm(x) = sqrt(x' * x)
	function vec_norm(x)
		if norm(x) == 0.0
			return ones(length(x)) ./ norm(ones(length(x)))
		else
			x ./ norm(x)
		end
	end
	# function d_vec_norm(x)
	# 	if norm(x) == 0.0
	# 		y = 1.0 * ones(length(x))#./norm(ones(length(x)))
	# 		return (I - y * y' / (y' * y)) /_norm(y)
	# 	else
	# 		(I - x * x' / (x' * x)) /_norm(x)
	# 	end
	# end
	function r(z,θ)
		b = z[1:2]
		ψ = z[3]
		y = θ[1]
		v = θ[2:3]

		return [norm(b) * v + ψ * b;
			    ψ * (y - norm(b))]
	end

	function drz(z,θ)
		b = z[1:2]
		ψ = z[3]
		y = θ[1]
		v = θ[2:3]

		return [v * vec_norm(b)' + ψ * I b;
			    -ψ * vec_norm(b)' (y - norm(b))]
	end

	function drθ(z,θ)
		b = z[1:2]
		ψ = z[3]
		y = θ[1]
		v = θ[2:3]
		return Array([zeros(2) Diagonal(norm(b) * ones(2));
					  ψ zeros(1, 2)])
	end

	# if y == 0.0
	# 	z = [zeros(2); ψ]
	# else
	# 	z = [b.value; ψ]
	# end
	z = [b.value; prob.constraints[1].dual]
	θ = [y; v]
	rz(x) = r(x, θ)
	rθ(x) = r(z, x)

	drdz = ForwardDiff.jacobian(rz, z)
	drdθ = ForwardDiff.jacobian(rθ, θ)
	# drz(z, θ)
	# drθ(z, θ)
	# norm(drdz - drz(z, θ))
	# norm(drdθ - drθ(z, θ))
	# rank(drdz)
	# eigen(drdz).values
	#
	# rank(drz(z, θ))
	# eigen(drz(z, θ)).values

	# x1 = (-drdz \ drdθ)[1:2,:]
	# x3 = (-(drdz' * drdz) \ (drdz' * drdθ))[1:2,:]

	ρ = 1.0e-5
	sol = (-(drdz' * drdz + ρ * I) \ (drdz' * drdθ))[1:2,:]
	# x4 = (-drdz' * ((drdz * drdz' + ρ * I) \ drdθ))[1:2,:]
	dλ = sol[:, 1] .* model.μ
	dq2 = -1.0 * sol[:, 2:3] ./ h
	dq3 = sol[:, 2:3] ./ h
	# print(b)
	return b, dλ, dq2, dq3
end

struct ContactDiffConstraints <: Constraints
	n
	ineq
end

function contact_diff_constraints(model, T)
	n = model.nc * (T + 1) + model.nb * (T - 1) + 1 * (T - 1)
	ineq = con_ineq_contact(model, T)

	return ContactDiffConstraints(n, ineq)
end

function constraints!(c, Z, con::ContactDiffConstraints, model, idx, h, T)
	shift = 0

	# signed-distance function
	for t = 1:T
		x = view(Z, idx.x[t])

		if t == 1
			q = view(x, 1:model.nq)
			c[1:model.nc] = ϕ_func(model, q)
		end

		q⁺ = view(x, model.nq .+ (1:model.nq))
		c[model.nc + (t-1) * model.nc .+ (1:model.nc)] = ϕ_func(model, q⁺)
	end

	shift += model.nc * (T + 1)

	# maximum dissipation
	for t = 1:T-1
		x⁺ = view(Z, idx.x[t+1])
		u = view(Z, idx.u[t])
		b = view(u, model.idx_b)

		b_sol, dλ, dq2, dq3 = friction_socp(model, x⁺, u, h)
		c[shift + (t - 1) * model.nb .+ (1:model.nb)] = b_sol - b
	end

	shift += model.nb * (T - 1)

	# impact complementarity
	for t = 1:T-1
		x⁺ = view(Z, idx.x[t + 1])
		u = view(Z, idx.u[t])

		q = view(x⁺, model.nq .+ (1:model.nq))
		s = view(u, model.idx_s)
		λ = view(u, model.idx_λ)

		c[shift + (t-1) * 1 + 1] = s[1] - (λ' * ϕ_func(model, q))[1]
	end

	shift += (T - 1)

    return nothing
end

function constraints_jacobian!(∇c, Z, con::ContactDiffConstraints, model, idx, h, T)
	shift = 0
	c_shift = 0

	# signed-distance function
	ϕ(y) = ϕ_func(model, y)

	for t = 1:T
		x = view(Z, idx.x[t])

		if t == 1
			q = view(x, 1:model.nq)
			r_idx = 1:model.nc
			c_idx = idx.x[t][1:model.nq]
			len = length(r_idx) * length(c_idx)
			∇c[shift .+ (1:len)] = vec(ForwardDiff.jacobian(ϕ, q))
			shift += len
		end

		q⁺ = view(x, model.nq .+ (1:model.nq))
		r_idx = model.nc + (t-1) * model.nc .+ (1:model.nc)
		c_idx = idx.x[t][model.nq .+ (1:model.nq)]
		len = length(r_idx) * length(c_idx)
		∇c[shift .+ (1:len)] = vec(ForwardDiff.jacobian(ϕ, q⁺))
		shift += len
	end

	c_shift += model.nc * (T + 1)

	# friction socp
	for t = 1:T-1
		x⁺ = view(Z, idx.x[t+1])
		u = view(Z, idx.u[t])

		b_sol, dλ, dq2, dq3 = friction_socp(model, x⁺, u, h)

		r_idx = c_shift + (t - 1) * model.nb .+ (1:model.nb)
		c_idx = idx.x[t + 1][collect([(1:2)..., (4:5)...])]
		len = length(r_idx) * length(c_idx)
		∇c[shift .+ (1:len)] = vec([dq2; dq3])
		shift += len

		c_idx = idx.u[t][model.idx_λ]
		len = length(r_idx) * length(c_idx)
		∇c[shift .+ (1:len)] = vec(dλ)
		shift += len

		c_idx = idx.u[t][model.idx_b]
		len = length(r_idx) * length(c_idx)
		∇c[shift .+ (1:len)] = vec(Diagonal(-1.0 * ones(model.nb)))
		shift += len
	end

	c_shift += model.nb * (T - 1)

	# impact complementarity
	for t = 1:T-1
		x⁺ = view(Z, idx.x[t + 1])
		u = view(Z, idx.u[t])

		q = view(x⁺, model.nq .+ (1:model.nq))
		s = view(u, model.idx_s)
		λ = view(u, model.idx_λ)

		cq(y) = s - [λ' * ϕ_func(model, y)]
		cu(y) = y[model.idx_s] - [y[model.idx_λ]' * ϕ_func(model, q)]

		r_idx = c_shift + (t-1) * 1 + 1
		c_idx = idx.x[t + 1][model.nq .+ (1:model.nq)]
		len = length(r_idx) * length(c_idx)
		∇c[shift .+ (1:len)] = vec(ForwardDiff.jacobian(cq, q))
		shift += len

		c_idx = idx.u[t]
		len = length(r_idx) * length(c_idx)
		∇c[shift .+ (1:len)] = vec(ForwardDiff.jacobian(cu, u))
		shift += len
	end

	c_shift += (T - 1)

	return nothing
end

function constraints_sparsity(con::ContactDiffConstraints, model, idx, T;
	shift_row = 0, shift_col = 0)

	row = []
    col = []

	con_shift = 0

	# signed-distance function
	for t = 1:T

		if t == 1
			r_idx = shift_row .+ (1:model.nc)
			c_idx = shift_col .+ idx.x[t][1:model.nq]
			row_col!(row, col, r_idx, c_idx)
		end

		r_idx = shift_row + model.nc + (t-1) * model.nc .+ (1:model.nc)
		c_idx = shift_col .+ idx.x[t][model.nq .+ (1:model.nq)]
		row_col!(row, col, r_idx, c_idx)
	end

	con_shift += model.nc * (T + 1)

	# friction socp
	for t = 1:T-1
		r_idx = shift_row + con_shift + (t - 1) * model.nb .+ (1:model.nb)
		c_idx = shift_col .+ idx.x[t + 1][collect([(1:2)..., (4:5)...])]
		row_col!(row, col, r_idx, c_idx)

		c_idx = shift_col .+ idx.u[t][model.idx_λ]
		row_col!(row, col, r_idx, c_idx)

		c_idx = idx.u[t][model.idx_b]
		row_col!(row, col, r_idx, c_idx)
	end

	con_shift += model.nb * (T - 1)

	# impact complementarity
	for t = 1:T-1
		r_idx = shift_row + con_shift + (t-1) * 1 + 1
		c_idx = shift_col .+ idx.x[t + 1][model.nq .+ (1:model.nq)]
		row_col!(row, col, r_idx, c_idx)

		c_idx = shift_col .+ idx.u[t]
		row_col!(row, col, r_idx, c_idx)
	end

	con_shift += (T - 1)

    return collect(zip(row, col))
end

function con_ineq_contact(model, T)

	con_ineq = []
	shift = 0

	# signed-distance function
	push!(con_ineq, [i for i = 1:model.nc * (T + 1)])
	shift += model.nc * (T + 1)

	# maximum dissipation
	# push!(con_ineq, [i for i = shift .+ (1:model.nb * (T - 1))])
	shift += model.nb * (T - 1)

	# impact complementarity
	push!(con_ineq, [i for i = shift .+ (1:(T - 1))])
	shift += (T - 1)

    return vcat(con_ineq...)
end