text:
function launch()

    @static if Sys.isapple()
        # OpenGL 3.2 + GLSL 150
        glsl_version = 150
        GLFW.WindowHint(GLFW.CONTEXT_VERSION_MAJOR, 3)
        GLFW.WindowHint(GLFW.CONTEXT_VERSION_MINOR, 2)
        GLFW.WindowHint(GLFW.OPENGL_PROFILE, GLFW.OPENGL_CORE_PROFILE) # 3.2+ only
        GLFW.WindowHint(GLFW.OPENGL_FORWARD_COMPAT, GL_TRUE) # required on Mac
    else
        # OpenGL 3.0 + GLSL 130
        glsl_version = 130
        GLFW.WindowHint(GLFW.CONTEXT_VERSION_MAJOR, 3)
        GLFW.WindowHint(GLFW.CONTEXT_VERSION_MINOR, 0)
        # GLFW.WindowHint(GLFW.OPENGL_PROFILE, GLFW.OPENGL_CORE_PROFILE) # 3.2+ only
        # GLFW.WindowHint(GLFW.OPENGL_FORWARD_COMPAT, GL_TRUE) # 3.0+ only
    end

    # setup GLFW error callback
    error_callback(err::GLFW.GLFWError) = @error "GLFW ERROR: code $(err.code) msg: $(err.description)"
    GLFW.SetErrorCallback(error_callback)

    # create window
    window = GLFW.CreateWindow(1280, 720, "Electrodermal Activity Analysis")
    @assert window != C_NULL
    GLFW.MakeContextCurrent(window)
    GLFW.SwapInterval(1)  # enable vsync

    # setup Dear ImGui context
    ctx = CImGui.CreateContext()

    # setup Dear ImGui style
    CImGui.StyleColorsLight()

    # setup Platform/Renderer bindings
    ImGui_ImplGlfw_InitForOpenGL(window, true)
    ImGui_ImplOpenGL3_Init(glsl_version)

    # Instantiate variables that are used to control input and output
    # of various widges.
    clear_color = Cfloat[0.45, 0.55, 0.60, 1.00]

    gui_events = Dict{String, Bool}()

    eda_data = nothing
    tags_data = nothing
    electrodermal_data = nothing

    # Initialize importers
    import_labels_key = "Import Interval Labels"
    import_labelled_intervals = create_import_context(GenericVendor(), GenericProduct(), IntervalLabels(), import_labels_key)

    import_eda_key = "Empatica E4 EDA.csv"
    import_eda = create_import_context(Empatica(), E4(), SkinConductance(), import_eda_key)

    import_tags_key = "Empatica E4 tags.csv"
    import_tags = create_import_context(Empatica(), E4(), Tags(), import_tags_key)

    select_timezone_key = "Select Time Zone"
    # Set the default time zome in accordance with the PC.
    select_timezone = create_timezone_context(localzone().name, timezone_names())

    # Initialize exporters
    labels_key = "Export Interval Labels"
    export_labelled_intervals = create_export_context(GenericVendor(), GenericProduct(), IntervalLabels(), labels_key)

    export_electrodermal_activity_key = "Export Electrodermal Activity"
    export_electrodermal_activity = create_export_context(Empatica(), E4(), SkinConductance(), export_electrodermal_activity_key)

    outline_layout = RectangularLayout(pos = ImVec2(0,0), width = Cfloat(600), height = Cfloat(80))
    data_layout = RectangularLayout(pos = ImVec2(0,0), width = Cfloat(600), height = Cfloat(400))

    # Configure overview plot context
    plotmodel = PlotlinesModel(rand(Float32,100))
    plotproperties₁ = PlotlinesDisplayProperties(id = "###overview"; caption = x -> "", createwindow = false, layout = outline_layout, xaxis = Axis(false, Cfloat(1)), yaxis = Axis(false, Cfloat(1)), padding = (0,0,0,0))
    plot_data_context = PlotContext(PlotlinesControl(true), plotmodel, plotproperties₁)

    # Configure interval selection context
    ni_control = NestedIntervalControl(true)
    ni_model = NestedInterval(start = 20.0, stop = 80.0, interval = 1:100)
    ni_properties = NestedIntervalDisplayProperties(id="###hidden control", caption = "Nested Interval", col = ImVec4(0.0, 0.0, 0.99, 0.2), layout = outline_layout, plotcontext = plot_data_context, padding = (0,0,0,0))
    ni_context = NestedIntervalContext(ni_control, ni_model, ni_properties)

    # Configure zoomed plot context.
    plotproperties₂ = PlotlinesDisplayProperties(id = "###eda"; caption = x -> "EDA", layout = data_layout, padding = (0,0,0,0))
    plotcontrol = PlotlinesControl(true)
    plot_outline_context = PlotContext(plotcontrol, plotmodel, plotproperties₂)

    plot_data_context = nothing
    ni_context = nothing
    plot_roi_context = nothing
    label_context = nothing
    tags_context = nothing
    enable_filter = false

    while !GLFW.WindowShouldClose(window)
        GLFW.PollEvents()
        # start the Dear ImGui frame
        ImGui_ImplOpenGL3_NewFrame()
        ImGui_ImplGlfw_NewFrame()
        CImGui.NewFrame()

        menubar!(gui_events)
        for (key, event) in pairs(gui_events)
            if key == import_eda_key
                if event == true
                    enable!(import_eda.control)
                    gui_events[key] = false
                end
                electrodermal_data₀ = import_eda()
                if !isrunning(import_eda) && !isnothing(electrodermal_data₀)
                    electrodermal_data = electrodermal_data₀
                    plot_data_context = construct_plot_data_context(electrodermal_data)
                    # The interval selection model will be shared with the PlotlinesDisplayProperties so that we can determine the mapping from the interval to
                    # an appropriate timestamp on the x-axis on the zoomed EDA plot.
                    data = get_eda(electrodermal_data)
                    ni_model = NestedInterval(start = Float64(1.0), stop = Float64(length(data)), interval = 1:length(data))
                    plot_outline_context =  construct_plot_outline_context(ni_model, plot_data_context.model, electrodermal_data)
                    ni_context = construct_nested_interval_context(ni_model, plot_data_context)
                    # Configure interactive truncated/zoomed plot context.
                    plot_roi_context = TruncatedPlotContext(PlotlinesControl(true),  plot_outline_context, ni_context)
                    #
                    dict = Dict{String, LabelledInterval}()
                    labelled_intervals = LabelledIntervals("Conditions", dict)
                    label_context = construct_labelled_interval_context(labelled_intervals, plot_data_context.model, ni_context.model, get_layout(plot_data_context.display_properties))
                    delete!(gui_events, key)
                end
            end
            if key == "Empatica E4 tags.csv"
                if isnothing(label_context) || isnothing(electrodermal_data)
                    CImGui.OpenPopup("Have you loaded electrodermal activity data?")
                    delete!(gui_events, key)
                else
                    if event == true
                        enable!(import_tags.control)
                        gui_events[key] = false
                    end
                    tags_data₀ = import_tags()
                    if !isrunning(import_tags) && !isnothing(tags_data₀)
                        tags_data = tags_data₀
                        tags_context  = construct_tags_context(tags_data, label_context.display_properties, electrodermal_data)
                        delete!(gui_events, key)
                    end
                end
            end
            if key == "Export Interval Labels"
                if isnothing(electrodermal_data)
                    CImGui.OpenPopup("Have you loaded electrodermal activity data?")
                    delete!(gui_events, key)
                else
                    if event == true
                        enable!(export_labelled_intervals)
                        gui_events[key] = false
                    end
                    export_labelled_intervals(label_context.model)
                    if !isrunning(export_labelled_intervals)
                        delete!(gui_events, key)
                    end
                end
            end
            if key == "Export Electrodermal Activity"
                if isnothing(electrodermal_data)
                    CImGui.OpenPopup("Have you loaded electrodermal activity data?")
                    delete!(gui_events, key)
                else
                    if event == true
                        enable!(export_electrodermal_activity)
                        gui_events[key] = false
                    end
                    export_electrodermal_activity(electrodermal_data, label_context.model)
                    if !isrunning(export_electrodermal_activity)
                        delete!(gui_events, key)
                    end
                end
            end
            if key == "Import Interval Labels"
                if isnothing(plot_data_context)
                    CImGui.OpenPopup("Have you loaded electrodermal activity data?")
                    delete!(gui_events, key)
                else
                    if event == true
                        enable!(import_labelled_intervals)
                        gui_events[key] = false
                    end
                    data = import_labelled_intervals()
                    if !isrunning(import_labelled_intervals) && !isnothing(data)
                        labelled_intervals = data
                        label_context = construct_labelled_interval_context(labelled_intervals, plot_data_context.model, ni_context.model, get_layout(plot_data_context.display_properties))
                        delete!(gui_events, key)
                    end
                end
            end
             if key == "Filter Lowpass 1Hz"
                 enable_filter = !enable_filter
                 if !isnothing(electrodermal_data)
                     apply_filter!(enable_filter, electrodermal_data)
                     # Update all contexts so that they reflect the new data
                     plot_data_context = construct_plot_data_context(electrodermal_data)
                     data = get_eda(electrodermal_data)
                     ni_model = NestedInterval(start = get_start(ni_context.model), stop = get_stop(ni_context.model), interval = 1:length(data))
                     plot_outline_context =  construct_plot_outline_context(ni_model, plot_data_context.model, electrodermal_data)
                     ni_context = construct_nested_interval_context(ni_model, plot_data_context)
                     # Configure interactive truncated/zoomed plot context.
                     plot_roi_context = TruncatedPlotContext(PlotlinesControl(true),  plot_outline_context, ni_context)
                     label_context = construct_labelled_interval_context(label_context.model, plot_data_context.model, ni_context.model, get_layout(plot_data_context.display_properties))
                     delete!(gui_events, key)
                end
             end
             if key == "Select Time Zone"
                 if isnothing(plot_data_context)
                     CImGui.OpenPopup("Have you loaded electrodermal activity data?")
                     delete!(gui_events, key)
                 else
                     if event == true
                         enable!(select_timezone)
                         gui_events[key] = false
                     end
                     data = select_timezone()
                     if !isrunning(select_timezone) && !isnothing(data)
                         tz = data
                         #label_context = construct_labelled_interval_context(labelled_intervals, plot_data_context.model, ni_context.model, get_layout(plot_data_context.display_properties))
                         plot_outline_context = plot_roi_context.plot_context
                         properties = plot_outline_context.display_properties
                         xtick = get_xtick(properties)
                         interpreter = get_interpreter(xtick)
                         timestamp = get_timestamp(interpreter)
                         intersample_duration = get_intersample_duration(interpreter)
                         nested_interval = get_nested_interval(interpreter)
                         xtick_new = Tickmark(; interpret = TimestampResolver(tz, timestamp, intersample_duration, nested_interval), spacing = Cfloat(15))
                         set_xtick!(properties, xtick_new )
                         delete!(gui_events, key)
                     end
                 end
             end
        end

        handle_import_error_messages()

        # Display the electrodermal activity data and concomitant interval selector.
        if !isnothing(plot_roi_context)
            CImGui.Begin("EDA (microsiemens)",C_NULL, CImGui.ImGuiWindowFlags_NoBringToFrontOnFocus)
                plot_roi_context()
            CImGui.End()
        end

        # Display the labelled conditions and markers.
        if !isnothing(label_context) || !isnothing(tags_context)
            CImGui.Begin("Conditions###Conditions and Tags",C_NULL, CImGui.ImGuiWindowFlags_NoBringToFrontOnFocus)
                pos = CImGui.GetCursorScreenPos()
                !isnothing(label_context) ? label_context() : nothing
                CImGui.SetCursorScreenPos(pos.x, pos.y + 3)
                !isnothing(tags_context) ? tags_context() : nothing
            CImGui.End()
        end

        # rendering
        CImGui.Render()
        GLFW.MakeContextCurrent(window)
        display_w, display_h = GLFW.GetFramebufferSize(window)
        glViewport(0, 0, display_w, display_h)
        glClearColor(clear_color...)
        glClear(GL_COLOR_BUFFER_BIT)
        ImGui_ImplOpenGL3_RenderDrawData(CImGui.GetDrawData())

        GLFW.MakeContextCurrent(window)
        GLFW.SwapBuffers(window)
    end

    # cleanup
    ImGui_ImplOpenGL3_Shutdown()
    ImGui_ImplGlfw_Shutdown()
    CImGui.DestroyContext(ctx)

    GLFW.DestroyWindow(window)

end

function labels_to_dataframe(labelled_intervals)
    df = DataFrame(label = String[], start = Int64[], stop = Int64[], first = Float64[], step = Float64[], last  = Float64[])
    for (key, labelled_interval) in pairs(labelled_intervals)
        nested_interval = labelled_interval.nested_interval
        start = get_start(nested_interval)
        stop = get_stop(nested_interval)
        interval = get_interval(nested_interval)
        push!(df, [key start stop first(interval) step(interval) last(interval)])
    end
    return df
end

function construct_tags_context(tags_data::AbstractVector, label_display_properties::LabelIntervalDisplayProperties, electrodermal_data::ElectrodermalData)
    outline_layout = get_layout(label_display_properties)
    padding = get_padding(label_display_properties)
    start_time = get_timestamp(electrodermal_data)
    Δt = (1 / get_sampling_frequency(electrodermal_data)) * 1000
    time₀ = Δt
    time₁ = Δt * length(get_unprocessed_eda(electrodermal_data))
    tags_model = TagsModel(start_time, time₀, time₁,  tags_data)
    tags_display_properties = TagsDisplayProperties(id = "###tags"; caption = x -> "", createwindow = false, layout = outline_layout, padding = padding)
    tags_context = TagsContext(TagsControl(true),  tags_model, tags_display_properties)
end

function construct_plot_data_context(data::ElectrodermalData)
    outline_layout = RectangularLayout(pos = ImVec2(0,0), width = Cfloat(1820), height = Cfloat(80))
    eda = Float32.(get_eda(data))
    # Configure overview plot context
    plotmodel = PlotlinesModel(eda)
    plotproperties₁ = PlotlinesDisplayProperties(id = "###overview"; caption = x -> "", createwindow = false, layout = outline_layout, xaxis = Axis(false, Cfloat(1)), yaxis = Axis(false, Cfloat(1)), padding = (0,80,0,0))
    plot_data_context = PlotContext(PlotlinesControl(true), plotmodel, plotproperties₁)
    return plot_data_context
end

function construct_plot_outline_context(ni_model, plotmodel,  data::ElectrodermalData)
    Δt = (1 / get_sampling_frequency(data)) * 1000
    timestamp₀ = get_timestamp(data)
    data_layout = RectangularLayout(pos = ImVec2(0,0), width = Cfloat(1820), height = Cfloat(400))
    # Configure zoomed plot context.
    plotproperties₂ = PlotlinesDisplayProperties(id = "###eda"; caption = DurationResolver(Δt, ni_model), layout = data_layout, padding = (0,80,0,0), xtick = Tickmark(; interpret = TimestampResolver(localzone().name,timestamp₀, Δt, ni_model), spacing = Cfloat(15)))
    plotcontrol = PlotlinesControl(true)
    plot_outline_context = PlotContext(plotcontrol, plotmodel, plotproperties₂)
    return plot_outline_context
end

function construct_labelled_interval_context(labelled_intervals::LabelledIntervals, plotmodel::AbstractModel, ni_model::NestedInterval, outline_layout)
    label_control = LabelIntervalControl(true, ni_model, "\0"^64*"\0")
    label_display_properties = LabelIntervalDisplayProperties(id = "###event markers" ; caption = "Conditions", createwindow =  false, layout = outline_layout)
    plotproperties = PlotlinesDisplayProperties(id = "###events overview"; caption = x->"", createwindow = false, layout = outline_layout, xaxis = Axis(false, Cfloat(1)), yaxis = Axis(false, Cfloat(1)), padding = (0,0,0,0))
    plot_outline_context = PlotContext(PlotlinesControl(true), plotmodel, plotproperties)
    label_context = LabelIntervalContext(label_control, labelled_intervals, label_display_properties, plot_outline_context)
    return label_context
end

function construct_nested_interval_context(ni_model, plot_data_context)
    # Configure interval selection context
    outline_layout = get_layout(plot_data_context.display_properties)
    ni_control = NestedIntervalControl(true)
    ni_properties = NestedIntervalDisplayProperties(id="###hidden control", caption = "Nested Interval", col = ImVec4(0.0, 0.0, 0.99, 0.2), layout = outline_layout, plotcontext = plot_data_context, padding = (0,80,0,0))
    ni_context = NestedIntervalContext(ni_control, ni_model, ni_properties)
    return ni_context
end

function apply_filter!(enabled::Bool, electrodermal_data::ElectrodermalData)
    eda₀ = get_unprocessed_eda(electrodermal_data)
    if enabled
        response_type = Lowpass(1; fs = get_sampling_frequency(electrodermal_data))
        design_method = Butterworth(2)
        eda₁ = filt(digitalfilter(response_type, design_method), eda₀)
        set_eda!(electrodermal_data, eda₁)
    else
        set_eda!(electrodermal_data, eda₀)
    end
end

function create_timezone_context(zone::String, possible_zones::Vector{String})
    TimeZoneContext(TimeZoneControl(false), TimeZoneModel(zone,  possible_zones), TimeZoneDisplayProperties(), SelectTimeZone(true))
end

function handle_import_error_messages()
    if CImGui.BeginPopupModal("Have you loaded electrodermal activity data?", C_NULL, CImGui.ImGuiWindowFlags_AlwaysAutoResize)
        CImGui.Text("Please first import the accompanying electrodermal activity data.\n\n")
        CImGui.Separator()
        CImGui.Button("OK", (120, 0)) && CImGui.CloseCurrentPopup()
        CImGui.SetItemDefaultFocus()
        CImGui.EndPopup()
    end

    # if CImGui.BeginPopupModal("Do you have permission to read the file?", C_NULL, CImGui.ImGuiWindowFlags_AlwaysAutoResize)
    #     CImGui.Text("Unable to access the specified file.\nPlease verify that: \n   (1) the file exists; \n   (2) you have permission to read the file.\n\n")
    #     CImGui.Separator()
    #     CImGui.Button("OK", (120, 0)) && CImGui.CloseCurrentPopup()
    #     CImGui.SetItemDefaultFocus()
    #     CImGui.EndPopup()
    # end
end


text:
# Autogenerated wrapper script for timg_jll for x86_64-linux-gnu
export timg

using Zlib_jll
using FFMPEG_jll
using JpegTurbo_jll
using libexif_jll
using OpenSlide_jll
using libgraphicsmagic_jll
JLLWrappers.@generate_wrapper_header("timg")
JLLWrappers.@declare_executable_product(timg)
function __init__()
    JLLWrappers.@generate_init_header(Zlib_jll, FFMPEG_jll, JpegTurbo_jll, libexif_jll, OpenSlide_jll, libgraphicsmagic_jll)
    JLLWrappers.@init_executable_product(
        timg,
        "bin/timg",
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()

ext: jl


text:
# This file contains the necessary ingredients to create a PackageManager for BinDeps
using BinDeps

struct EnvManager{T} <: BinDeps.PackageManager
    packages::Vector{String}
end

"Manager for root environment"
const Manager = EnvManager{Symbol(PREFIX)}

function Base.show(io::IO, manager::EnvManager)
    write(io, "Conda packages: ", join(manager.packages, ", "))
end

BinDeps.can_use(::Type{EnvManager}) = true

function BinDeps.package_available(manager::EnvManager{T}) where {T}
    pkgs = manager.packages
    # For each package, see if we can get info about it. If not, fail out
    for pkg in pkgs
        if !exists(pkg, T)
            return false
        end
    end
    return true
end

BinDeps.libdir(m::EnvManager{T}, ::Any) where {T} = lib_dir(T)
BinDeps.bindir(m::EnvManager{T}, ::Any) where {T} = bin_dir(T)

BinDeps.provider(::Type{EnvManager{T}}, packages::AbstractVector{<:AbstractString}; opts...) where {T} = EnvManager{T}(packages)
BinDeps.provider(::Type{EnvManager{T}}, packages::AbstractString; opts...) where {T} = EnvManager{T}([packages])

function BinDeps.generate_steps(dep::BinDeps.LibraryDependency, manager::EnvManager, opts)
    pkgs = manager.packages
    ()->install(pkgs, manager)
end

function install(pkgs, manager::EnvManager{T}) where {T}
    for pkg in pkgs
        add(pkg, T)
    end
end


text:
# functions to add a hybridization
# originally in functions.jl
# Claudia March 2015

# ------------------------------------------------ add new hybridization------------------------------------

# function to add hybridization event
# input: edge1, edge2 are the edges to remove
#        edge3, edge4, edge5 are the new tree edges to add
#        net is the network
#        gamma is the gamma for the hybridization
# warning: assumes that edge1, edge2 are tree edges with inCycle=-1
#          assumes the hybrid edge goes from edge1 to edge2
#          sets minor hybrid edge length to zero
# this function create the hybrid node/edges and connect everything
# and deletes edge1,2 from the nodes, and removes the nodes from edge1,2
# returns the hybrid node to start future updates there
function createHybrid!(edge1::Edge, edge2::Edge, edge3::Edge, edge4::Edge, net::HybridNetwork, gamma::Float64)
    0 < gamma < 1 || error("gamma must be between 0 and 1: $(gamma)")
    (edge1.hybrid || edge2.hybrid) ? error("edges to delete must be tree edges") : nothing
    (edge3.hybrid || edge4.hybrid) ? error("edges to add must be tree edges") : nothing
    pushEdge!(net,edge3);
    pushEdge!(net,edge4);
    # create hybridization
    max_node = maximum([e.number for e in net.node]);
    max_edge = maximum([e.number for e in net.edge]);
    gamma < 0.5 || @warn "adding a major hybrid edge with gamma $(gamma), this can cause problems when updating incycle"
    hybrid_edge = Edge(max_edge+1,0.0,true,gamma,gamma>=0.5);
    pushEdge!(net,hybrid_edge);
    hybrid_node = Node(max_node+1,false,true,[edge2,hybrid_edge,edge4]);
    tree_node = Node(max_node+2,false,false,[edge1,edge3,hybrid_edge]);
    setNode!(hybrid_edge,[tree_node,hybrid_node]);
    setNode!(edge3,[tree_node,edge1.node[2]]);
    setNode!(edge4,[hybrid_node,edge2.node[2]]);
    setEdge!(edge1.node[2],edge3);
    setEdge!(edge2.node[2],edge4);
    removeEdge!(edge2.node[2],edge2);
    removeEdge!(edge1.node[2],edge1);
    removeNode!(edge1.node[2],edge1);
    setNode!(edge1,tree_node);
    removeNode!(edge2.node[2],edge2);
    #[n.number for n in edge2.node]
    setNode!(edge2,hybrid_node)
    pushNode!(net,hybrid_node);
    pushNode!(net,tree_node);
    #pushHybrid!(net,hybrid_node);
    return hybrid_node
end

# aux function for chooseEdgesGamma to identify
# if two edges are sisters and if they are cherry
# (have leaves)
# returns true/false for sisters, true/false for cherry
#         true/false for nonidentifiable (two leaves, k=1 node crossed by hybridization)
function sisterOrCherry(edge1::Edge,edge2::Edge)
    sisters = false
    cherry = false
    nonidentifiable = false
    node = nothing;
    if isEqual(edge1.node[1],edge2.node[1]) || isEqual(edge1.node[1],edge2.node[2])
        node = edge1.node[1];
    elseif isEqual(edge1.node[2],edge2.node[1]) || isEqual(edge1.node[2],edge2.node[2])
        node = edge1.node[2];
    end
    if node !== nothing
        size(node.edge,1) == 3 || error("node found $(node.number) that does not have exactly 3 edges, it has $(size(node.edge,1)) edges instead.")
        sisters = true
        if getOtherNode(edge1,node).leaf && getOtherNode(edge2,node).leaf
            cherry = true
        elseif getOtherNode(edge1,node).leaf || getOtherNode(edge2,node).leaf
            edge = nothing
            for e in node.edge
                if(!isEqual(e,edge1) && !isEqual(e,edge2))
                    edge = e
                end
            end
            if getOtherNode(edge,node).leaf
                nonidentifiable = true
            end
        end
    end
    return sisters, cherry, nonidentifiable
end

# aux function to addHybridization
# it chooses the edges in the network and the gamma value
# warning: chooses edge1, edge2, gamma randomly, but
#          we could do better later
# check: gamma is uniform(0,1/2) to avoid big gammas
# fixit: add different function to choose gamma
# fixit: how to stop from infinite loop if there are no options
# blacklist used for afterOptBLAll
# input: edges, list of edges from which to choose, default is net.edge
# warning: if edges is not net.edge, it still need to contain Edge objects from net (not deepcopies)
function chooseEdgesGamma(net::HybridNetwork, blacklist::Bool, edges::Vector{Edge})
    index1 = 1;
    index2 = 1;
    inlimits = false
    inblack = true
    cherry = false
    nonidentifiable = false
    while !inlimits || edges[index1].inCycle != -1 || edges[index2].inCycle != -1 || inblack || cherry || nonidentifiable
        index1 = round(Integer,rand()*size(edges,1));
        index2 = round(Integer,rand()*size(edges,1));
        if index1 != index2 && index1 != 0 && index2 != 0 && index1 <= size(edges,1) && index2 <= size(edges,1)
            inlimits = true
            sisters, cherry, nonidentifiable = sisterOrCherry(edges[index1],edges[index2]);
        else
            inlimits = false
        end
        if blacklist && !isempty(net.blacklist)
            length(net.blacklist) % 2 == 0 || error("net.blacklist should have even number of entries, not length: $(length(net.blacklist))")
            i = 1
            while i < length(net.blacklist)
                if edges[index1].number == net.blacklist[i]
                    if edges[index2].number == net.blacklist[i+1]
                        inblack = true
                    else
                        inblack = false
                    end
                elseif edges[index2].number == net.blacklist[i]
                    if edges[index1].number == net.blacklist[i+1]
                        inblack = true
                    else
                        inblack = false
                    end
                end
                i += 2
            end
        else
            inblack = false
        end
    end
    gamma = rand()*0.5;
    @debug "choose edges and gamma: from $(edges[index1].number) to $(edges[index2].number), $(gamma)"
    return edges[index1],edges[index2],gamma
end

chooseEdgesGamma(net::HybridNetwork) = chooseEdgesGamma(net, false, net.edge)
chooseEdgesGamma(net::HybridNetwork, blacklist::Bool) = chooseEdgesGamma(net, blacklist, net.edge)

# aux function for addHybridization
# that takes the output edge1, edge2, gamma from
# chooseEdgesGamma and created necessary edges
# returns edge3, edge4, and adjusts edge1, edge2 to shorter length
function parameters4createHybrid!(edge1::Edge, edge2::Edge,net::HybridNetwork)
    max_edge = maximum([e.number for e in net.edge]);
    t1 = rand()*edge1.length;
    t3 = edge1.length - t1;
    edge3 = Edge(max_edge+1,t3);
    edge1.length = t1;
    t1 = rand()*edge2.length;
    t3 = edge2.length - t1;
    edge4 = Edge(max_edge+2,t3);
    edge2.length = t1;
    edge3.containRoot = edge1.containRoot
    edge4.containRoot = edge2.containRoot
    return edge3, edge4
end

# aux function to add the hybridization
# without checking all the updates
# returns the hybrid node of the new hybridization
# calls chooseEdgesGamma, parameter4createHybrid and createHybrid
# blacklist used in afterOptBLAll
# usePartition=true if we use the information on net.partition, default true
function addHybridization!(net::HybridNetwork, blacklist::Bool, usePartition::Bool)
    if(net.numHybrids > 0 && usePartition)
        !isempty(net.partition) || error("net has $(net.numHybrids) but net.partition is empty")
        index = choosePartition(net)
        if(index == 0) #no place for new hybrid
            @debug "no partition suitable to place new hybridization"
            return nothing
        end
        partition = splice!(net.partition,index) #type partition
        @debug "add hybrid with partition $([n.number for n in partition.edges])"
        edge1, edge2, gamma = chooseEdgesGamma(net, blacklist,partition.edges);
    else
        edge1, edge2, gamma = chooseEdgesGamma(net, blacklist);
    end
    @debug "add hybridization between edge1, $(edge1.number) and edge2 $(edge2.number) with gamma $(gamma)"
    edge3, edge4 = parameters4createHybrid!(edge1,edge2,net);
    hybrid = createHybrid!(edge1, edge2, edge3, edge4, net, gamma);
    return hybrid
end

addHybridization!(net::HybridNetwork) = addHybridization!(net, false, true)

# function to update who is the major hybrid
# after a new hybridization is added and
# inCycle is updated
# warning: needs updateInCycle! run first
# can return the updated edge for when undoing network moves, not needed now
function updateMajorHybrid!(net::HybridNetwork, node::Node)
    node.hybrid || error("node $(node.number) is not hybrid, cannot update major hybrid after updateInCycle")
    length(node.edge) == 3 || error("hybrid node $(node.number) has $(length(node.edge)) edges, should have 3")
    hybedge = nothing
    edgecycle = nothing
    for e in node.edge
        if(e.hybrid)
            hybedge = e
        elseif(e.inCycle != -1 && !e.hybrid)
            edgecycle = e
        end
    end
    !isa(hybedge,Nothing) || error("hybrid node $(node.number) does not have hybrid edge")
    !isa(edgecycle,Nothing) || error("hybrid node $(node.number) does not have tree edge in cycle to update to hybrid edge after updateInCycle")
    #println("updating hybrid status to edgeincycle $(edgecycle.number) for hybedge $(hybedge.number)")
    makeEdgeHybrid!(edgecycle,node,1-hybedge.gamma)
end

# function to update everything of a new hybridization
# it follows the flow diagram in ipad
# input: new added hybrid, network,
#        updatemajor (bool) to decide if we need to update major edge
#        only need to update if new hybrid added, if read from file not needed
#        allow=true allows extreme/very bad triangles, needed when reading
#        updatePart = true will update PArtition at this moment, it makes sense with a newly added hybrid
#        but not if net just read (because in this case it needs all inCycle updated before)
# returns: success bool, hybrid, flag, nocycle, flag2, flag3
function updateAllNewHybrid!(hybrid::Node,net::HybridNetwork, updatemajor::Bool, allow::Bool, updatePart::Bool)
    flag, nocycle, edgesInCycle, nodesInCycle = updateInCycle!(net,hybrid);
    if(nocycle)
        return false, hybrid, flag, nocycle, false, false
    else
        if(flag)
            if(updatemajor)
                updateMajorHybrid!(net,hybrid);
            end
            flag2, edgesGammaz = updateGammaz!(net,hybrid,allow);
            if(flag2)
                flag3, edgesRoot = updateContainRoot!(net,hybrid);
                if(updatePart)
                    updatePartition!(net,nodesInCycle)
                end
                if(flag3)
                    parameters!(net)
                    return true, hybrid, flag, nocycle, flag2, flag3
                else
                    #undoContainRoot!(edgesRoot);
                    #undoistIdentifiable!(edgesGammaz);
                    #undoGammaz!(hybrid,net);
                    #undoInCycle!(edgesInCycle, nodesInCycle);
                    return false, hybrid, flag, nocycle, flag2, flag3
                end
            else
                if(updatePart)
                    updatePartition!(net,nodesInCycle)
                end
                flag3, edgesRoot = updateContainRoot!(net,hybrid); #update contain root even if it is bad triangle to writeTopologyLevel1 correctly
                #undoistIdentifiable!(edgesGammaz);
                #undoGammaz!(hybrid,net);
                #undoInCycle!(edgesInCycle, nodesInCycle);
                return false, hybrid, flag, nocycle, flag2, flag3
            end
        else
            # no need to do updatePartition in this case, because we only call deleteHybrid after
            undoInCycle!(edgesInCycle, nodesInCycle);
            return false, hybrid, flag, nocycle, true, true
        end
    end
end

updateAllNewHybrid!(hybrid::Node,net::HybridNetwork, updatemajor::Bool) = updateAllNewHybrid!(hybrid,net, updatemajor, false, true)

# function to add a new hybridization event
# it calls chooseEdgesGamma and createHybrid!
# input: network
# check: assumes that one of the two possibilities for
#        major hybrid edge gives you a cycle, true?
# warning: "while" removed, it does not attempt to add until
#          success, it attempts to add once
# returns: success (bool), hybrid, flag, nocycle, flag2, flag3
# blacklist used in afterOptBLAll
function addHybridizationUpdate!(net::HybridNetwork, blacklist::Bool, usePartition::Bool)
    hybrid = addHybridization!(net,blacklist, usePartition);
    isa(hybrid,Nothing) && return false,nothing,false,false,false,false
    updateAllNewHybrid!(hybrid,net,true)
end

addHybridizationUpdate!(net::HybridNetwork) = addHybridizationUpdate!(net, false, true)
addHybridizationUpdate!(net::HybridNetwork, blacklist::Bool) = addHybridizationUpdate!(net, blacklist::Bool, true)

# function that will add a hybridization with addHybridizationUpdate,
# if success=false, it will try to move the hybridization before
# declaring failure
# blacklist used in afterOptBLAll
function addHybridizationUpdateSmart!(net::HybridNetwork, blacklist::Bool, N::Integer)
    global CHECKNET
    @debug "MOVE: addHybridizationUpdateSmart"
    success, hybrid, flag, nocycle, flag2, flag3 = addHybridizationUpdate!(net, blacklist)
    @debug begin
        printEverything(net)
        "success $(success), flag $(flag), flag2 $(flag2), flag3 $(flag3)"
    end
    i = 0
    if !success
        if isa(hybrid,Nothing)
            @debug "MOVE: could not add hybrid by any means"
        else
            while((nocycle || !flag) && i < N) #incycle failed
                @debug "MOVE: added hybrid causes conflict with previous cycle, need to delete and add another"
                deleteHybrid!(hybrid,net,true)
                success, hybrid, flag, nocycle, flag2, flag3 = addHybridizationUpdate!(net, blacklist)
            end
            if(nocycle || !flag)
                @debug "MOVE: added hybridization $(i) times trying to avoid incycle conflicts, but failed"
            else
                if(!flag3 && flag2) #containRoot failed
                    @debug "MOVE: added hybrid causes problems with containRoot, will change the direction to fix it"
                    success = changeDirectionUpdate!(net,hybrid) #change dir of minor
                elseif(!flag2 && flag3) #gammaz failed
                    @debug "MOVE: added hybrid has problem with gammaz (not identifiable bad triangle)"
                    if(flag3)
                        @debug "MOVE: we will move origin to fix the gammaz situation"
                        success = moveOriginUpdateRepeat!(net,hybrid,true)
                    else
                        @debug "MOVE: we will move target to fix the gammaz situation"
                        success = moveTargetUpdateRepeat!(net,hybrid,true)
                    end
                elseif(!flag2 && !flag3) #containRoot AND gammaz failed
                    @debug "MOVE: containRoot and gammaz both fail"
                end
            end
            if !success
                @debug "MOVE: could not fix the added hybrid by any means, we will delete it now"
                CHECKNET && checkNet(net)
                @debug begin printEverything(net); "printed everything" end
                deleteHybridizationUpdate!(net,hybrid)
                @debug begin printEverything(net); "printed everything" end
                CHECKNET && checkNet(net)
            end
        end
    end
    success && @debug "MOVE: added hybridization SUCCESSFUL: new hybrid $(hybrid.number)"
    return success
end

addHybridizationUpdateSmart!(net::HybridNetwork, N::Integer) = addHybridizationUpdateSmart!(net, false,N)

# ----------------------------------- add alternative hybridizations found in bootstrap ------------------------------------
## function to add the top hybridizations (not in best network)
## found in bootstrap networks
## it modifies the network and BSe to add the new edge number,
## it also adds a new column 0/1, with 1 if the edge is not present in the original estimated network
## cutoff=10, only show hybridizations with bootstrap support>10%
## top=3, only show the top three alternative hybridizations (sorted by bootstrap support)
## can we put in a different color in the plot?
## Cecile: you could add a new option to our RCall-based plot function to take a vector of colors.
"""
    addAlternativeHybridizations!(net::HybridNetwork, BSe::DataFrame;
                                  cutoff=10::Number, top=3::Int)

Modify the network `net` (the best network estimated with snaq) by adding other hybridizations
that are present in the bootstrap networks. By default, it will only consider hybrid edges with
more than 10% bootstrap support (`cutoff`) and it will only include the three top hybridizations
(`top`) sorted by bootstrap support.
The function also modifies the dataframe `BSe` obtained with `hybridBootstrapSupport`. In the original
`BSe` dataframe, hybrid edges that do not appear in the best network have a missing number.
After the hybrid edges are added with `addAlternativeHybridizations`, `BSe` is modified to include the
edge numbers of the newly added hybrid edges. Note that the function only adds the hybrid edges as minor to
keep the underlying tree topology.

# example

```julia
bootnet = readMultiTopology("bootstrap-networks.txt")
bestnet = readTopology("best.tre")
BSn, BSe, BSc, BSgam, BSedgenum = hybridBootstrapSupport(bootnet, bestnet);
addAlternativeHybridizations!(bestnet,BSe)
using PhyloPlots
plot(bestnet, edgeLabel=BSe[[:edge,:BS_hybrid_edge]])
```
"""
function addAlternativeHybridizations!(net::HybridNetwork,BSe::DataFrame; cutoff=10::Number,top=3::Int)
    top > 0 || error("top must be greater than 0")
    BSe[:alternative] = falses(length(BSe[:hybrid]))
    newBSe = BSe[BSe[:BS_hybrid_edge] .> cutoff,:]
    newBSe = newBSe[.!ismissing.(newBSe[:hybrid]) & .!ismissing.(newBSe[:sister]),:]
    newBSe = newBSe[ismissing.(newBSe[:edge]),:]
    newHyb = newBSe[1:top,:]

    if(size(newHyb,1) == 0)
        @warn "Did not find any alternative hybridizations with bootstrap support greater than the cutoff, so nothign added"
        return
    end

    for i in 1:top
        hybnum = newHyb[:hybrid][i]
        sisnum = newHyb[:sister][i]
        edgenum = addHybridBetweenClades!(hybnum,sisnum,net)
        ind1 = findall(x->!ismissing(x) && x==hybnum,BSe[:hybrid])
        ind2 = findall(x->!ismissing(x) && x==sisnum,BSe[:sister])
        ind = intersect(ind1,ind2)
        BSe[ind,:edge] = edgenum
        BSe[ind,:alternative] = true
    end
end

## function to a hybrid edge between hybrid clade (hybnum = node number) and sister clade (sisnum = node number)
## the function finds the parent edges to these nodes, and puts a hybrid edge between them
## the function modifies net, and it returns the number of the minor hybrid edge added
function addHybridBetweenClades!(hybnum::Number,sisnum::Number,net::HybridNetwork)
    hybind = getIndexNode(hybnum,net)
    sisind = getIndexNode(sisnum,net)

    ## hybridization ed1->ed2
    edge1 = getMajorParentEdge(net.node[sisind])
    edge2 = getMajorParentEdge(net.node[hybind])

    edge3,edge4 = parameters4createHybrid!(edge1, edge2,net)
    hybridnode = createHybrid!(edge1, edge2, edge3, edge4, net, 0.1) ## gamma=0.1, fixed later
    if(edge1.length < 0)
        setBranchLength!(edge1,-1.0)
        setBranchLength!(edge3,-1.0)
    end
    if(edge2.length < 0)
        setBranchLength!(edge2,-1.0)
        setBranchLength!(edge4,-1.0)
    end

    if(edge2.isChild1)
        edge4.hybrid = true
        setGamma!(edge4,0.9)
        edge4.isChild1 = true
    else
        edge2.hybrid = true
        setGamma!(edge2,0.9)
        edge2.isChild1 = false
    end
    ## used gamma=0.1 to make the new edge a minor edge, but we really do not have gamma value:
    emaj = getMajorParentEdge(hybridnode)
    emaj.gamma = -1
    e = getMinorParentEdge(hybridnode)
    e.gamma = -1
    return e.number
end


text:
#Assistant functions for calculating the number of particles hitting the wall per frame

function countcollisions(A)         #This function counts collisions, uses same principle as the particle-wall
                                    #function.
    count=0 #Counter.
    for x in A
        if x==0
            return 0
        end

        if x%22==0
            count+=1
        end
        if x%39==0
            count+=1
        end
        if x%85==0
            count+=1
        end
        if x%133==0
            count+=1
        end

    end
    return count
end
function arraygenerator(n,dim)     #This function generates a Grid with a fixed number of randomly placed particles.
    A=Array{Int64,2}(ones(dim))
    lengthA=length(A)
    for k=1:n
        randint=rand(1:lengthA)

        while A[randint]!=1
            randint=rand(1:lengthA)
        end
        A[randint]=rand([2 3 5 7])

    end
    return A
end
function average_of_array(A)      #Average function.
    return (sum(A))/length(A)
end


text:
using ModelConstructors, HDF5

regenerate_data = false

###################################################################
# Set Up Linear Model
###################################################################

function setup_linear_model(; regime_switching::Bool = false)
    m = GenericModel()

    # Set up linear parameters
    prior_para = regime_switching ? 1e1 : 1e3
    m <= parameter(:α1, 0., (-1e5, 1e5), (-1e5, 1e5), Untransformed(), Normal(0, prior_para),
                  fixed = false)
    m <= parameter(:β1, 0., (-1e5, 1e5), (-1e5, 1e5), Untransformed(), Normal(0, prior_para),
                  fixed = false)
    m <= parameter(:σ1, 1., (1e-5, 1e5), (1e-5, 1e5), SquareRoot(), Uniform(0, prior_para),
                  fixed = false)
    m <= parameter(:α2, 0., (-1e5, 1e5), (-1e5, 1e5), Untransformed(), Normal(0, prior_para),
                  fixed = false)
    m <= parameter(:β2, 0., (-1e5, 1e5), (-1e5, 1e5), Untransformed(), Normal(0, prior_para),
                  fixed = false)
    m <= parameter(:σ2, 1., (1e-5, 1e5), (1e-5, 1e5), SquareRoot(), Uniform(0, prior_para),
                  fixed = false)
    m <= parameter(:α3, 0., (-1e5, 1e5), (-1e5, 1e5), Untransformed(), Normal(0, prior_para),
                  fixed = false)
    m <= parameter(:β3, 0., (-1e5, 1e5), (-1e5, 1e5), Untransformed(), Normal(0, prior_para),
                  fixed = false)
    m <= parameter(:σ3, 1., (1e-5, 1e5), (1e-5, 1e5), SquareRoot(), Uniform(0, prior_para),
                  fixed = false)
    m <= Setting(:n_particles, 400) #-> will get commented eventually
    m <= Setting(:n_Φ, 100)
    m <= Setting(:λ, 2.0)
    m <= Setting(:n_smc_blocks, 1)
    m <= Setting(:use_parallel_workers, false)
    m <= Setting(:step_size_smc, 0.5)
    m <= Setting(:n_mh_steps_smc, 1)
    m <= Setting(:resampler_smc, :polyalgo)
    m <= Setting(:target_accept, 0.25)

    m <= Setting(:mixture_proportion, .9)
    m <= Setting(:tempering_target, 0.95)
    m <= Setting(:resampling_threshold, .5)
    m <= Setting(:use_fixed_schedule, true)

    if regime_switching
        for i in 1:3
            ModelConstructors.set_regime_val!(m[Symbol("α$(i)")], 1, (i == 3) ? 3. : -.1 * i)
            ModelConstructors.set_regime_val!(m[Symbol("α$(i)")], 2, (i == 3) ? 3. : .1 * i)
            ModelConstructors.set_regime_val!(m[Symbol("α$(i)")], 3, 3.)
            ModelConstructors.set_regime_fixed!(m[Symbol("α$(i)")], 1, (i == 3) ? true : false)
            ModelConstructors.set_regime_fixed!(m[Symbol("α$(i)")], 2, (i == 3) ? true : false) # Do not estimate α3,
            ModelConstructors.set_regime_fixed!(m[Symbol("α$(i)")], 3, (i == 3) ? true : false) # just to check functionality
        end
        for i in 1:3
            ModelConstructors.set_regime_val!(m[Symbol("β$(i)")], 1, .2 * i)
            ModelConstructors.set_regime_prior!(m[Symbol("β$(i)")], 1, Normal(0, prior_para)) # regime-switching prior,
            ModelConstructors.set_regime_val!(m[Symbol("β$(i)")], 2, -.1 * i)                 # just to check functionality
            ModelConstructors.set_regime_prior!(m[Symbol("β$(i)")], 2, Normal(0, prior_para * 1.2))
            ModelConstructors.set_regime_val!(m[Symbol("β$(i)")], 3, .1 * i)
            ModelConstructors.set_regime_prior!(m[Symbol("β$(i)")], 3, Normal(0, prior_para * 1.5))
        end
    end

    return m
end

# Generate Data
if regenerate_data
    Random.seed!(1793)

    X = randn(Float64, (3, 100))
    err = randn(Float64, (3, 100))
    β = collect(1:3)
    α = collect(1:3)

    data = β .* X .+ α .+ err

    # Regime-switching version
    Xrs = randn(Float64, (3, 300))
    err = randn(Float64, (3, 300))
    β₁ = collect(1:3)
    β₂ = collect(2:4)
    β₃ = collect(3:5)

    α₁ = collect(1:3)
    α₂ = collect(1:3)

    rsdata = similar(err)
    reg1   = 1:100
    reg2   = 101:200
    reg3   = 201:300
    rsdata[:, reg1] = β₁ .* Xrs[:, reg1] .+ α₁ .+ err[:, reg1]
    rsdata[:, reg2] = β₂ .* Xrs[:, reg2] .+ α₂ .+ err[:, reg2]
    rsdata[:, reg3] = β₃ .* Xrs[:, reg3] .+ α₂ .+ err[:, reg3]

    # Save Data

    h5open("reference/test_data.h5", "w") do file
        write(file, "data", data)
        write(file, "rsdata", rsdata)
        write(file, "X", X)
        write(file, "Xrs", Xrs)
    end
else # Need to define reg1, reg2, and reg3 for regime-switching log likelihood function
    reg1 = 1:100
    reg2 = 101:200
    reg3 = 201:300
end

# Read Predictors from data
X = h5read("reference/test_data.h5", "X")
Xrs = h5read("reference/test_data.h5", "Xrs")

# Log Likelihood Function
N = 3
function loglik_fn(p, d)
    # we assume the ordering of (α_i, β_i, σ_i)
    Σ = zeros(N,N)
    α = Vector{Float64}(undef,N)
    β = Vector{Float64}(undef,N)
    for i in 1:N
        α[i]   = p[i * 3 - 2] # note p is a ParameterVector
        β[i]   = p[i * 3 - 1] # but since LHS are Vector{Float64},
        Σ[i,i] = p[i * 3]^2   # the value of p is autoamtically stored
    end
    det_Σ = det(Σ)
    inv_Σ = inv(Σ)
    term1 = -N / 2 * log(2 * π) - 1 /2 * log(det_Σ) # N should be number of equations
    logprob = 0.
    errors = d .- α .- β .* X[:, 1:size(d, 2)] # bridging test uses first 1/2 of sample so need to be robust there
    for t in 1:size(d, 2) # Number of time periods
        logprob += term1 - 1/2 * dot(errors[:, t], inv_Σ * errors[:, t]) # Should be error in each time period
    end
    return logprob
end

function rs_loglik_fn(p, d)
    # we assume the ordering of (α_i, β_i, σ_i)
    Σ = zeros(N,N)
    α = Vector{Float64}(undef, 3 * N)
    β = Vector{Float64}(undef, 3 * N)
    for i in 1:N
        α[i]   = regime_val(p[N * (i - 1) + 1], 1)
        β[i]   = regime_val(p[N * (i - 1) + 2], 1)
        Σ[i,i] = p[N * (i - 1) + 3]
    end
    for i in 1:N # Looping over each entry of the constant vector and betas in a given regime
        α[N + i] = regime_val(p[N * (i - 1) + 1], 2) # α1, α2, α3 each has 3 regimes. This line is for regime 2.
        α[(2 * N) + i] = regime_val(p[N * (i - 1) + 1], 3) # This line is for regime 3.
        β[N + i] = regime_val(p[N * (i - 1) + 2], 2) # β1, β2, β3 each has 3 regimes.
        β[(2 * N) + i] = regime_val(p[N * (i - 1) + 2], 3)
    end

    det_Σ = det(Σ)
    inv_Σ = inv(Σ)
    term1 = -N / 2 * log(2 * π) - 1 /2 * log(det_Σ)
    logprob = 0.
    errors = similar(d)
    errors[:, reg1] = d[:, reg1] .- α[1:N] .- β[1:N] .* Xrs[:, reg1]
    errors[:, reg2] = d[:, reg2] .- α[(N + 1):(2 * N)] .- β[(N + 1):(2 * N)] .* Xrs[:, reg2]
    errors[:, reg3] = d[:, reg3] .- α[(2 * N + 1):(3 * N)] .- β[(2 * N + 1):(3 * N)] .* Xrs[:, reg3]
    for t in 1:size(d, 2) # see above
        logprob += term1 - 1/2 * dot(errors[:, t], inv_Σ * errors[:, t])
    end
    return logprob
end


text:
module DDCModelsExample

# Types and functions package exports
export DDCbasic
export gen_data, estimateΠ
export estimate_model, calc_stderr

# Load additional packages
using Parameters
using LinearAlgebra, Distributions, Random
using GalacticOptim, Optim, ForwardDiff

# Include  the source files
include("types.jl")
include("support.jl")
include("models/DDCbasic.jl")
end

max_stars_repo_name: rgreminger/DDCModelsExample

max_issues_repo_name: rgreminger/DDCModelsExample

max_forks_repo_name: rgreminger/DDCModelsExample


text:
const MACHEP = 1.2e-16
const DWARF = 1.0e-38

function enorm(n::Int, x::Vector{Float64})
  const rdwarf = 3.834e-20
  const rgiant = 1.304e19

  s1 = s2 = s3 = x1max = x3max = 0.0

  floatn = float(n)
  agiant = rgiant / floatn

  for i = 1:n
    @inbounds xabs = abs(x[i])
    if xabs > rdwarf && xabs < agiant
      # sum for intermediate components
      s2 += xabs * xabs
      continue
    end

    if xabs > rdwarf
      # sum for larger components
      if xabs > x1max
        temp = x1max / xabs
        s1 = 1.0 + s1 * temp * temp
        x1max = xabs
      else
        temp = xabs / x1max
        s1 += temp * temp
      end

      continue
    end

    # sum for smaller components
    if xabs > x3max
      temp = x3max / xabs
      s3 = 1.0 + s3 * temp * temp
      x3max = xabs
    else
      if xabs != 0.0
        temp = xabs / x3max
        s3 += temp * temp
      end
    end
  end

  # Calculation of norm
  if s1 != 0.0
    temp = s1 + (s2 / x1max) / x1max
    return x1max * sqrt(temp)
  end

  if s2 != 0.0
    if s2 >= x3max
      temp = s2 * (1.0 + (x3max / s2) * (x3max * s3))
    else
      temp = x3max * ((s2 / x3max) + (x3max * s3))
    end
    ans = sqrt(temp)
  else
    ans = x3max * sqrt(s3)
  end

  return ans
end

# Forward Difference Approximation #
function fdjac2!(m::Int, n::Int, x::Vector{Float64}, fvec::Vector{Float64}, fjac::Matrix{Float64}, ::Int, iflag::Int, epsfcn::Float64, wa::Vector{Float64}, fcn!::Function)
  # returns:
  # fjac :  is an output m by n array which contains the
  #         approximation to the jacobian matrix evaluated at x.
  temp  = max(epsfcn, MACHEP)
  eps_ = sqrt(temp)
  ij = 1
  @simd for j = 1:n
    @inbounds temp = x[j]
    h = eps_ * abs(temp)
    if h == 0.0
      h = eps_
    end

    @inbounds x[j] = temp + h
    fcn!(m, n, x, wa)
    if iflag < 0
      return fjac
    end
    @inbounds x[j] = temp
    for i = 1:m
      @inbounds fjac[ij] = (wa[i] - fvec[i]) / h
      ij += 1
    end
  end

  return fjac
end

function qrfac!(m::Int, n::Int, a::Matrix{Float64}, ::Int, pivot::Int, ipvt::Vector{Int}, ::Int, rdiag::Vector{Float64}, acnorm::Vector{Float64}, wa::Vector{Float64})
  # returns:
  # a :     is an m by n array. on input a contains the matrix for
  #         which the qr factorization is to be computed. on output
  #         the strict upper trapezoidal part of a contains the strict
  #         upper trapezoidal part of r, and the lower trapezoidal
  #         part of a contains a factored form of q (the non-trivial
  #         elements of the u vectors described above).
  # ipvt:   is an integer output array of length lipvt. ipvt
  #         defines the permutation matrix p such that a*p = q*r.
  #         column j of p is column ipvt(j) of the identity matrix.
  #         if pivot is false, ipvt is not referenced.
  # rdiag:  is an output array of length n which contains the
  #         diagonal elements of r.
  # acnorm: is an output array of length n which contains the
  #         norms of the corresponding columns of the input matrix a.
  #         if this information is not needed, then acnorm can coincide
  #         with rdiag.
  # Compute the initial column norms and initialize several arrays
  ij = 1
  @simd for j = 1:n
    @inbounds acnorm[j] = enorm(m, a[:, ij])
    @inbounds rdiag[j] = acnorm[j]
    @inbounds wa[j] = rdiag[j]
    if pivot != 0
      @inbounds ipvt[j] = j
    end
    ij += 1
  end

  # Reduce a to r with householder transformations
  minmn = min(m, n)
  for j = 1:minmn
    if pivot == 0
      @goto L40
    end

    # bring the column of the largest norm into the pivot position
    kmax = j
    @simd for k = j:n
      @inbounds if rdiag[k] > rdiag[kmax]
        kmax = k
      end
    end

    if kmax == j
      @goto L40
    end
    ij = (m * (j - 1)) + 1
    jj = (m * (kmax - 1)) + 1

    @simd for i = 1:m
      @inbounds temp = a[ij]
      @inbounds a[ij] = a[jj]
      @inbounds a[jj] = temp
      ij += 1
      jj += 1
    end

    @inbounds rdiag[kmax] = rdiag[j]
    @inbounds wa[kmax] = wa[j]
    @inbounds k = ipvt[j]
    @inbounds ipvt[j] = ipvt[kmax]
    @inbounds ipvt[kmax] = k

    @label L40
    # Compute the householder transformation to reduce the j-th column of a to
    # a multiple of the j-th unit vector
    jj = j + m * (j - 1)
    # jj_end = jj + (m - j)
    @inbounds ajnorm = enorm(m - (j - 1), a[jj:end])
    if ajnorm == 0.0
      @goto L100
    end
    @inbounds if a[jj] < 0.0
      ajnorm = -ajnorm
    end
    ij = jj

    @simd for i = j:m
      @inbounds a[ij] /= ajnorm
      ij += 1

    end

    @inbounds a[jj] += 1.0

    # Apply the transformation to the remaining columns
    # and update the norms
    jp1 = j + 1
    if jp1 <= n
      @simd for k = jp1:n
        sum_ = 0.0
        ij = j + m * (k - 1)
        jj = j + m * (j - 1)

        for i = j:m
          @inbounds sum_ += a[jj] * a[ij]
          ij += 1
          jj += 1
        end

        @inbounds temp = sum_ / a[j + m * (j - 1)]
        ij = j + m * (k - 1)
        jj = j + m * (j - 1)
        for i = j:m
          @inbounds a[ij] -= temp * a[jj]
          ij += 1
          jj += 1
        end

        if pivot != 0 && rdiag[k] != 0.0
          @inbounds temp = a[j + m * (k - 1)]  / rdiag[k]
          temp = max(0.0, 1.0 - temp * temp)
          @inbounds rdiag[k] *= sqrt(temp)
          @inbounds temp = rdiag[k]/wa[k]

          if 0.05 * temp * temp <= MACHEP
            a_start = jp1 + m * (k - 1)
            @inbounds rdiag[k] = enorm(m - j, a[a_start:end])
            @inbounds wa[k] = rdiag[k]
          end
        end
      end
    end

    @label L100

    @inbounds rdiag[j] = -ajnorm
  end

  return a, ipvt, rdiag, acnorm
end

function qrsolv!(n::Int, r::Matrix{Float64}, ldr::Int, ipvt::Vector{Int}, diag_::Vector{Float64}, qtb::Vector{Float64}, x::Vector{Float64},
                sdiag::Vector{Float64}, wa::Vector{Float64})
  # copy r and (q transpose) * b to preserve input and initialize s, in particular
  # save the diagonal elements of r in x
  kk = 1
  @simd for j = 1:n
    ij = kk
    ik = kk
    for i = j:n
      @inbounds r[ij] = r[ik]
      ij += 1
      ik += ldr
    end
    @inbounds x[j] = r[kk]
    @inbounds wa[j] = qtb[j]
    kk += ldr + 1
  end

  # eliminate the diagonal matrix d using a given rotation
  for j = 1:n
    @inbounds l = ipvt[j]
    @inbounds if diag_[l] == 0.0
      @goto L90
    end
    for k = j:n
      @inbounds sdiag[k] = 0.0
    end
    @inbounds sdiag[j] = diag_[l]

    # transformations to eliminate the row of d, modify only a single element
    # of (q transpose) * b beyond the first n, which is initially zero
    qtbpj = 0.0
    for k = j:n
      # determine a given rotation which elminates the appropriate element
      # in the current row of d
      @inbounds if sdiag[k] == 0.0
        continue
      end
      kk = k + ldr * (k - 1)
      @inbounds if abs(r[kk]) < abs(sdiag[k])
        @inbounds cotan_ = r[kk] / sdiag[k]
        sin_ = 0.5 / sqrt(0.25 + 0.25 * cotan_ * cotan_)
        cos_ = sin_ * cotan_
      else
        @inbounds tan_ = sdiag[k] / r[kk]
        cos_ = 0.5 / sqrt(0.25 + 0.25 * tan_ * tan_)
        sin_ = cos_ * tan_
      end

      # compute the modified diagonal element of r and the modified element of
      # (q transpose) * b, 0
      @inbounds r[kk] = cos_ * r[kk] + sin_* sdiag[k]
      @inbounds temp = cos_ * wa[k] + sin_ * qtbpj
      @inbounds qtbpj = -sin_ * wa[k] + cos_ * qtbpj
      @inbounds wa[k] = temp

      # accumulate the transformation of the row s
      kp1 = k + 1
      if n >= kp1
        ik = kk + 1
        for i = kp1:n
          @inbounds temp = cos_ * r[ik] + sin_ * sdiag[i]
          @inbounds sdiag[i] = -sin_ * r[ik] + cos_ * sdiag[i]
          @inbounds r[ik] = temp
          ik += 1
        end
      end
    end

    @label L90

    # store the diagonal element of s and restore the corresponding diagonal
    # element of r
    kk = j + ldr * (j - 1)
    @inbounds sdiag[j] = r[kk]
    @inbounds r[kk] = x[j]
  end

  # Solve the triangular system for z, if the system is singular, then obtain
  # a least squares solution
  nsing = n
  @simd for j = 1:n
    @inbounds if sdiag[j] == 0.0 && nsing == n
      nsing = j - 1
    end

    if nsing < n
      @inbounds wa[j] = 0.0
    end
  end

  if nsing < 1
    @goto L150
  end

  @simd for k = 1:nsing
    j = nsing - k + 1
    sum_ = 0.0
    jp1 = j + 1
    if nsing >= jp1
      ij = jp1 + ldr * (j - 1)
      for i = jp1:nsing
        @inbounds sum_ += r[ij] * wa[i]
        ij += 1
      end
    end
    @inbounds wa[j] = (wa[j] - sum_) / sdiag[j]
  end

  @label L150
  # permute the components of z back to components of x
  @simd for j = 1:n
    # l = ipvt[j]
    @inbounds x[ipvt[j]] = wa[j]
  end

  return r, x, sdiag
end

function lmpar!(n::Int, r::Matrix{Float64}, ldr::Int, ipvt::Vector{Int}, diag_::Vector{Float64}, qtb::Vector{Float64}, delta::Float64, par::Float64,
              x::Vector{Float64}, sdiag::Vector{Float64}, wa1::Vector{Float64}, wa2::Vector{Float64})

  # Compute and store in x the gauss-newton direction.  If the jacobin is rank-deficient
  # obtain a least-squares solution

  nsing = n
  jj = 1
  @simd for j=1:n
    @inbounds wa1[j] = qtb[j]
    @inbounds if r[jj] == 0.0 && nsing == n
      nsing = j - 1
    end
    if nsing < n
      @inbounds wa1[j] = 0.0
    end

    jj += ldr + 1
  end

  if nsing >= 1
    @simd for k = 1:nsing
      j = nsing - k + 1
      @inbounds wa1[j] = wa1[j] / r[j + ldr * (j - 1)]
      @inbounds temp = wa1[j]
      jm1 = j - 1
      if jm1 > 0
        ij = ldr * (j - 1) + 1
        for i = 1:jm1 + 1
          @inbounds wa1[i] -= r[ij] * temp
          ij += 1
        end
      end
    end
  end

  @simd for j = 1:n
    @inbounds l = ipvt[j]
    @inbounds x[l] = wa1[j]
  end

  # initialize the iteration counter, evaluate the function at the origin
  # and test for acceptance of the gauss-newton direction
  iter = 0
  @simd for j = 1:n
    @inbounds wa2[j] = diag_[j] * x[j]
  end

  dxnorm = enorm(n, wa2)
  fp = dxnorm - delta

  fp <= 0.1 * delta && @goto L220

  # if the jacobin is not rank deficient, the newton step provides a lower bound
  # parl, for the zero of the function.  Otherwise set this bound to zero
  parl = 0.0
  if nsing >= n
    @simd for j = 1:n
      @inbounds l = ipvt[j]
      @inbounds wa1[l] = diag_[l] * (wa2[l] / dxnorm)
    end
    jj = 1
    @simd for j = 1:n
      sum_ = 0.0
      jm1 = j - 1
      if jm1 >= 1
        ij = jj
        for i = 1:jm1 + 1
          @inbounds sum_ += r[ij] * wa1[i]
          ij += 1
        end
      end
      @inbounds wa1[j] = (wa1[j] - sum_) / r[j + ldr * (j - 1)]
      jj += ldr
    end
    temp = enorm(n, wa1)
    parl = ((fp / delta) / temp) / temp
  end

  jj = 1
  @simd for j = 1:n
    sum_ = 0.0
    ij = jj
    for i = 1:j
      @inbounds sum_ += r[ij] * qtb[i]
      ij += 1
    end
    @inbounds l = ipvt[j]
    @inbounds wa1[j] = sum_ / diag_[l]
    jj += ldr
  end

  gnorm = enorm(n, wa1)
  paru = gnorm / delta
  if paru == 0.0
    paru = DWARF / min(delta, 0.1)
  end

  # if the input par lies outside of the interval (parl, paru), set the par to
  # the closer endpoint
  par = max(par, parl)
  par = min(par, paru)

  if par == 0.0
    par = gnorm / dxnorm
  end

  @label L150
  iter += 1

  # eval the function at the current value of par
  if par == 0.0
    par = max(DWARF, 0.001 * paru)
  end

  temp = sqrt(par)
  # for j = 1:n
  #   wa1[j] = temp * diag_[j]
  # end
  wa1 = temp * diag_

  qrsolv!(n, r, ldr, ipvt, wa1, qtb, x, sdiag, wa2)

  @simd for j = 1:n
    @inbounds wa2[j] = diag_[j] * x[j]
  end

  dxnorm = enorm(n, wa2)
  temp = fp
  fp = dxnorm - delta

  # if the function is small enough, accept the current value of par.  also test
  # for the exceptional cases where parl is zero or the number of iterations
  # has reached 10
  (abs(fp) <= 0.1 * delta || (parl == 0.0 && fp <= temp && temp < 0.0) || iter == 10) && @goto L220

  # compute the newton correction
  @simd for j = 1:n
    @inbounds l = ipvt[j]
    @inbounds wa1[j] = diag_[l] * (wa2[l] / dxnorm)
  end

  jj = 1
  @simd for j = 1:n
    @inbounds wa1[j] = wa1[j] / sdiag[j]
    @inbounds temp = wa1[j]
    jp1 = j + 1
    if jp1 <= n
      ij = jp1 + jj
      for i = jp1:n
        @inbounds wa1[i] -= r[ij] * temp
        ij += 1
      end
    end
    jj += ldr
  end
  temp = enorm(n, wa1)
  parc = ((fp / delta) / temp) / temp

  # depending on the sign of the function, update parl or paru
  if fp > 0.0
    parl = max(parl, par)
  end

  if fp < 0.0
    paru = min(paru, par)
  end

  # compute an improved estimate for par
  par = max(parl, par + parc)

  # end of iteration
  @goto L150

  @label L220
  # termination
  if iter == 0
    par = 0.0
  end

  return par
end

function lmdif!(m::Int, n::Int, x::Vector{Float64}, fvec::Vector{Float64}, ftol::Float64, xtol::Float64, gtol::Float64, maxFev::Int, epsfcn::Float64,
                diag_::Vector{Float64}, mode::Int, factor_::Float64, nprint::Int, info_::Int, nfev::Int, fjac::Matrix{Float64}, ldfjac::Int, ipvt::Vector{Int}, qtf::Vector{Float64},
                wa1::Vector{Float64}, wa2::Vector{Float64}, wa3::Vector{Float64}, wa4::Vector{Float64}, fcn!::Function)

  delta = 0.0
  xnorm = 0.0
  info_ = 0
  iflag = 0
  nfev = 0

  # checking for errors
  (n <= 0 || m < n || ldfjac < m || ftol < 0.0 || xtol < 0.0 || gtol < 0.0 || maxFev <= 0 || factor_ <= 0.0) && @goto L300

  if mode == 2
    for j = 1:n
      @inbounds if diag_[j] <= 0.0
        @goto L300
      end
    end
  end

  # Evaluate the function at its starting point and calculate its norm
  iflag = 1
  fcn!(m, n, x, fvec)
  nfev = 1
  if iflag < 0
    @goto L300
  end
  fnorm = enorm(m, fvec)

  # Initialize the levenberg-marquardt param and iteration counter
  par = 0.0
  iter = 1

  @label L30
  # Calcualte the Jacobian matrix
  iflag = 2
  # if useJac
  #   jacFcn!(m, n, x, fjac, iflag)
  # else
  fdjac2!(m, n, x, fvec, fjac, ldfjac, iflag, epsfcn, wa4, fcn!)
  # end
  # fjac = jacFcn!(x)

  nfev += n

  if iflag < 0
    @goto L300
  end

  # if requested, call fcn to enable printing of iterates
  if nprint > 0
    iflag = 0
    if mod(iter - 1, nprint) == 0
      fcn!(m, n, x, fvec)
      if iflag < 0
        @goto L300
      end
    end
  end

  # compute the QR factorization of the Jacobian
  qrfac!(m, n, fjac, ldfjac, 1, ipvt, n, wa1, wa2, wa3)

  # on the first iteration and if mode is 1, scale according to the norms
  # of the columns in the initial jacobin
  if iter == 1
    if mode != 2
      @simd for j = 1:n
        @inbounds diag_[j] = wa2[j]
        if wa2[j] == 0.0
          @inbounds diag_[j] = 1.0
        end
      end
    end

    # on the first iteration, calculate the norm of the scaled x and initialize
    # the step bound delta
    @simd for j = 1:n
      @inbounds wa3[j] = diag_[j] * x[j]
    end

    xnorm = enorm(n, wa3)
    delta = factor_ * xnorm
    if delta == 0.0
      delta = factor_
    end

  end

  # form (q transpose) * fvec and store the first n components in qtf
  @simd for i = 1:m
    @inbounds wa4[i] = fvec[i]
  end

  jj = 1
  @simd for j = 1:n
    @inbounds temp3 = fjac[jj]
    if temp3 != 0.0
      sum_ = 0.0
      ij = jj
      for i = j:m
        @inbounds sum_ += fjac[ij] * wa4[i]
        ij += 1
      end
      temp = -sum_ / temp3
      ij = jj
      for i = j:m
        @inbounds wa4[i] += fjac[ij] * temp
        ij += 1
      end
    end
    @inbounds fjac[jj] = wa1[j]
    jj += m + 1
    @inbounds qtf[j] = wa4[j]
  end

  # println("wa1: ", wa1)
  # println("wa2: ", wa2)
  # println("qtf: ", qtf)
  # println("wa4: ", wa4)
  # println("fjac: ", fjac)
  # error("break")

  # compute the norm of the scaled gradient
  gnorm = 0.0
  if fnorm != 0.0
    jj = 1
    @simd for j = 1:n
      @inbounds l = ipvt[j]
      @inbounds if wa2[l] != 0.0
        sum_ = 0.0
        ij = jj
        for i = 1:j
          @inbounds sum_ += fjac[ij] * (qtf[i] / fnorm)
          ij += 1
        end
        @inbounds gnorm = max(gnorm, abs(sum_ / wa2[l]))
      end
      jj += m
    end
  end

  # test for convergence of the gradient norm
  if gnorm <= gtol
    info_ = 4
  end

  if info_ != 0
    @goto L300
  end

  # rescale if necessary
  if mode != 2
    @simd for j = 1:n
      @inbounds diag_[j] = max(diag_[j], wa2[j])
    end
  end

  # Beginning of inner loop
  @label L200

  # determine the levenberg-marquardt param
  par = lmpar!(n, fjac, ldfjac, ipvt, diag_, qtf, delta, par, wa1, wa2, wa3, wa4)

  # println("AFTER: ")
  # println("fvec: ", fvec)
  # println("fjac: ", fjac)
  # println("wa1: ", wa1)
  # println("wa2: ", wa2)
  # println("wa3: ", wa3)
  # println("wa4: ", wa4)
  # println("ipvt: ", ipvt)
  # println("diag: ", diag_)
  # println("qtf: ", qtf)
  # println("delta: ", delta)
  #
  # error("STOP")

  # store the direction of p and x + p, calculate the norm of p
  @simd for j = 1:n
    @inbounds wa1[j] = -wa1[j]
    @inbounds wa2[j] = x[j] + wa1[j]
    @inbounds wa3[j] = diag_[j] * wa1[j]
  end

  pnorm = enorm(n, wa3)

  # on the initial iteration, adjust the initial step bound
  if iter == 1
    delta = min(delta, pnorm)
  end

  iflag = 1
  fcn!(m, n, wa2, wa4)
  nfev += 1

  if iflag < 0
    @goto L300
  end

  fnorm1 = enorm(m, wa4)
  # println(wa4)

  # Compute the scaled actual reduction
  actred = -1.0

  # println("fnorm: ", fnorm)
  # println("fnorm1: ", fnorm1)

  if (0.1 * fnorm1) < fnorm
    temp = fnorm1 / fnorm
    actred = 1.0 - temp * temp
  end

  # Compute the scaled predicted reduction and the scaled directional derivative
  jj = 1
  @simd for j = 1:n
    @inbounds wa3[j] = 0.0
    @inbounds l = ipvt[j]
    @inbounds temp = wa1[l]
    ij = jj
    for i = 1:j
      @inbounds wa3[i] += fjac[ij] * temp
      ij += 1
    end

    jj += m
  end

  temp1 = enorm(n, wa3) / fnorm
  temp2 = (sqrt(par) * pnorm) / fnorm
  prered = temp1 * temp1 + (temp2 * temp2) / 0.5
  dirder = -(temp1 * temp1 + temp2 * temp2)

  # println("temp1: ", temp1)
  # println("temp2: ", temp2)
  # println("prered: ", prered)
  # println("dirder: ", dirder)

  # Compute the ratio of the actual to the predicted reduction
  ratio = 0.0
  if prered != 0.0
    ratio = actred / prered
  end

  if ratio <= 0.25
    if actred >= 0.0
      temp = 0.5
    else
      temp = 0.5 * dirder / (dirder + 0.5 * actred)
    end

    if ((0.1 * fnorm1) >= fnorm) || temp < 0.1
      temp = 0.1
    end

    delta = temp * min(delta, pnorm/0.1)
    par = par / temp
  else
    if par == 0.0  || ratio >= 0.75
      delta = pnorm / 0.5
      par = 0.5 * par
    end
  end
  # test for successful iteration
  if ratio >= 0.0001
    @simd for j = 1:n
      @inbounds x[j] = wa2[j]
      @inbounds wa2[j] = diag_[j] * x[j]
    end
    @simd for i = 1:m
      @inbounds fvec[i] = wa4[i]
    end

    xnorm = enorm(n, wa2)
    fnorm = fnorm1
    iter += 1
  end

  # Tests for convergence
  if abs(actred) <= ftol && prered <= ftol && 0.5 * ratio <= 1.0
    info_ = 1
  end

  if delta <= xtol * xnorm
    info_ = 2
  end

  if abs(actred) <= ftol && prered <= ftol && 0.5 * ratio <= 1.0 && info_ == 2
    info_ = 3
  end

  info_ != 0 && @goto L300

  # Tests for termination and stringent tolerances
  if nfev >= maxFev
    info_ = 5
  end

  if abs(actred) <= MACHEP && prered <= MACHEP && 0.5 * ratio <= 1.0
    info_ = 6
  end

  if delta <= MACHEP * xnorm
    info_ = 7
  end

  if gnorm <= MACHEP
    info_ = 8
  end

  info_ != 0 && @goto L300

  # end of inner loop, repeat if iteration unsuccessful
  if ratio < 0.0001
    @goto L200
  end

  # end of outer loop
  @goto L30

  @label L300
  # Termination either normal or user imposed
  if iflag < 0
    info_ = iflag
  end
  iflag = 0
  if nprint > 0
    fcn!(m, n, x, fvec)
  end

  return x, fvec, info_, nfev, fjac, ipvt, qtf
end


text:
#! This is the package in which we put all our dependencies, functions, types etc.
#! In desparate need of cleaning up!
module InnovationPrivacy

using Parameters, Setfield, Optim, NLsolve, DataFrames, Interpolations

using QuantEcon: qnwnorm, gridmake
#using StatsFuns: normcdf
using Roots: find_zero
using Distributions: cdf, Pareto
using StatsBase: ecdf
using Statistics: quantile

using Plots
#pgfplotsx()

include("Types.jl")
include("Functions_olg.jl")
include("Functions_v3.jl")
include("FOCS_v3.jl")
#include("Retired.jl")
#include("PlotFuncs.jl")
include("Auxiliary_Funcs.jl")
include("Wages.jl")
include("Residuals.jl")

#include("Functions.jl")
#include("FunctionsDeprInfo.jl")
#include("Functions_Replication.jl")
#include("TestFuncs.jl")
#include("Functions_simple.jl")
#include("Functions_simple_solving.jl")
#include("Functions_simple_model.jl")
#include("Functions_bundle.jl")

export Pars, ParsGrowth, Pars_Bundle, BPars, BParsGrowth, Pars_Simple, Pars_OLG, Pars_v3, residuals, solve_firm_problem, FOC_data_sharing, bundle, dD_dDI, dD_dDE, FOC_data_gen, FOC_data_buying, FOC_int_labor, data_gen, get_solution, weight_of_type, sold_data, firm_output, aggregate_output, residuals_SS, find_steady_state, get_ss_solution, residuals_SS_single, find_p_D, find_steady_state_iterative, comp_statics, residuals_p_D, firm_revenue, sol_f_problem, find_steady_state_PE, firm_profits, solve_firm_problem_unc, solve_firm_problem_con, solve_firm_problem_unc2, solve_firm_problem_nl, sol_f_problem_nl, find_value_func, Firm_sol, Save_Firm_Solution, solve_firm_problem_unc_nl2, solve_firm_problem_unc_nl, get_solution_nl2, A_G_of_type,
@unpack_Pars, @unpack_ParsGrowth, @unpack_BPars, @unpack_BParsGrowth, @unpack_Pars_OLG, @unpack_Pars_Simple, @unpack_Pars_Bundle, @unpack_Pars_v3, @unpack_Firm_sol,
firm_profit,
capital_demand,
test_func,
Eπ, Er, EA,
extractK,
signal,
Tau_of_K_old,
solve_model,
K_solve,
Transition_Probability,
find_limit,
find_V_n_max,
find_V_n,
find_all_n,
find_all_V_n,
find_V_zero,
find_u_star,
u_to_z,
z_to_u,
Firm_Distribution,
iterate_population,
Eπ_Distribution,
Er_Distribution,
labor_costs,
μ_n,
Transition_Matrix2,
Equilibrium_Distribution,
find_V_0,
mass_and_entry,
mean_preserving_noise,
CalcAllUncertainty,
BPars_dep,
Eπ_dep,
Transition_Probability_dep,
ParsV2,
find_price_info,
solve_agg_C,
find_agg_C,
wage_rate,
find_wage_rate,
capital_prod,
solve_agg_C_iterative,
capital_of_tau,
profit,
k_I_prod_A,
k_I_prod_B,
solve_firm_problems,
k_I_prod_A,
k_I_prod_B,
loss_A,
loss_B,
loss_A_con,
solve_agg_Y_iterative_unc,
profit_unc,
find_mass,
solve_U_problem
end # module


text:
# ------------------------------------------------------------------
# Licensed under the ISC License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    WeightedHoldOut(weighter, fraction; loss=Dict())

Weighted loss validation with train / hold-out split.
The source data of the problem is split with a `fraction`
into train and hold-out sets, and the samples are weighted
with a `weighter` method. Optionally, specify `loss` function
from `LossFunctions.jl` for some of the variables.
"""
struct WeightedHoldOut{W<:AbstractWeighter,T<:Real} <: AbstractErrorEstimator
  weighter::W
  fraction::T
  loss::Dict{Symbol,SupervisedLoss}
end

WeightedHoldOut(weighter::AbstractWeighter, fraction::Real; loss=Dict()) =
  WeightedHoldOut{typeof(weighter),typeof(fraction)}(weighter, fraction, loss)

function Base.error(solver::AbstractLearningSolver,
                    problem::LearningProblem,
                    eestimator::WeightedHoldOut)
  # retrieve problem info
  sdata = sourcedata(problem)
  ovars = outputvars(task(problem))
  loss  = eestimator.loss
  for var in ovars
    if var ∉ keys(loss)
      loss[var] = defaultloss(sdata[1,var])
    end
  end

  # train and hold-out sets
  partitioner = FractionPartitioner(eestimator.fraction)
  train, hold = partition(sdata, partitioner)

  # solve subproblem
  subproblem = LearningProblem(train, hold, task(problem))
  solution   = solve(subproblem, solver)

  # weight hold-out set
  w = weight(hold, eestimator.weighter)

  result = pmap(ovars) do var
    y = hold[var]
    ŷ = solution[var]
    var => value(loss[var], y, ŷ, AggMode.WeightedSum(w)) /length(y)
  end

  Dict(result)
end