text:
#########################################################################
#                                                                       #
# This file provides functions for interfacing with Matpower data files #
#                                                                       #
#########################################################################

"Parses the matpwer data from either a filename or an IO object"
function parse_matpower(io::IO; validate = true)::Dict
    mp_data = _parse_matpower_string(read(io, String))
    pm_data = _matpower_to_powermodels!(mp_data)
    if validate
        correct_network_data!(pm_data)
    end
    return pm_data
end

function parse_matpower(file::String; kwargs...)::Dict
    mp_data = open(file) do io
        parse_matpower(io; kwargs...)
    end
    return mp_data
end

### Data and functions specific to Matpower format ###

const _mp_data_names = [
    "mpc.version",
    "mpc.baseMVA",
    "mpc.bus",
    "mpc.gen",
    "mpc.branch",
    "mpc.dcline",
    "mpc.gencost",
    "mpc.dclinecost",
    "mpc.bus_name",
    "mpc.storage",
    "mpc.switch",
]

const _mp_bus_columns = [
    ("bus_i", Int),
    ("bus_type", Int),
    ("pd", Float64),
    ("qd", Float64),
    ("gs", Float64),
    ("bs", Float64),
    ("area", Int),
    ("vm", Float64),
    ("va", Float64),
    ("base_kv", Float64),
    ("zone", Int),
    ("vmax", Float64),
    ("vmin", Float64),
    ("lam_p", Float64),
    ("lam_q", Float64),
    ("mu_vmax", Float64),
    ("mu_vmin", Float64),
]

const _mp_bus_name_columns = [("name", Union{String, SubString{String}})]

const _mp_gen_columns = [
    ("gen_bus", Int),
    ("pg", Float64),
    ("qg", Float64),
    ("qmax", Float64),
    ("qmin", Float64),
    ("vg", Float64),
    ("mbase", Float64),
    ("gen_status", Int),
    ("pmax", Float64),
    ("pmin", Float64),
    ("pc1", Float64),
    ("pc2", Float64),
    ("qc1min", Float64),
    ("qc1max", Float64),
    ("qc2min", Float64),
    ("qc2max", Float64),
    ("ramp_agc", Float64),
    ("ramp_10", Float64),
    ("ramp_30", Float64),
    ("ramp_q", Float64),
    ("apf", Float64),
    ("mu_pmax", Float64),
    ("mu_pmin", Float64),
    ("mu_qmax", Float64),
    ("mu_qmin", Float64),
]

const _mp_branch_columns = [
    ("f_bus", Int),
    ("t_bus", Int),
    ("br_r", Float64),
    ("br_x", Float64),
    ("br_b", Float64),
    ("rate_a", Float64),
    ("rate_b", Float64),
    ("rate_c", Float64),
    ("tap", Float64),
    ("shift", Float64),
    ("br_status", Int),
    ("angmin", Float64),
    ("angmax", Float64),
    ("pf", Float64),
    ("qf", Float64),
    ("pt", Float64),
    ("qt", Float64),
    ("mu_sf", Float64),
    ("mu_st", Float64),
    ("mu_angmin", Float64),
    ("mu_angmax", Float64),
]

const _mp_dcline_columns = [
    ("f_bus", Int),
    ("t_bus", Int),
    ("br_status", Int),
    ("pf", Float64),
    ("pt", Float64),
    ("qf", Float64),
    ("qt", Float64),
    ("vf", Float64),
    ("vt", Float64),
    ("pmin", Float64),
    ("pmax", Float64),
    ("qminf", Float64),
    ("qmaxf", Float64),
    ("qmint", Float64),
    ("qmaxt", Float64),
    ("loss0", Float64),
    ("loss1", Float64),
    ("mu_pmin", Float64),
    ("mu_pmax", Float64),
    ("mu_qminf", Float64),
    ("mu_qmaxf", Float64),
    ("mu_qmint", Float64),
    ("mu_qmaxt", Float64),
]

const _mp_storage_columns = [
    ("storage_bus", Int),
    ("ps", Float64),
    ("qs", Float64),
    ("energy", Float64),
    ("energy_rating", Float64),
    ("charge_rating", Float64),
    ("discharge_rating", Float64),
    ("charge_efficiency", Float64),
    ("discharge_efficiency", Float64),
    ("thermal_rating", Float64),
    ("qmin", Float64),
    ("qmax", Float64),
    ("r", Float64),
    ("x", Float64),
    ("standby_loss", Float64),
    ("status", Int),
]

const _mp_switch_columns = [
    ("f_bus", Int),
    ("t_bus", Int),
    ("psw", Float64),
    ("qsw", Float64),
    ("state", Int),
    ("thermal_rating", Float64),
    ("status", Int),
]

""
function _parse_matpower_string(data_string::String)
    matlab_data, func_name, colnames = parse_matlab_string(data_string, extended = true)

    case = Dict{String, Any}()

    if func_name !== nothing
        case["name"] = func_name
    else
        @info(
            string(
                "no case name found in matpower file.  The file seems to be missing \"function mpc = ...\"",
            )
        )
        case["name"] = "no_name_found"
    end

    case["source_type"] = "matpower"
    if haskey(matlab_data, "mpc.version")
        case["source_version"] = matlab_data["mpc.version"]
    else
        @info(
            string(
                "no case version found in matpower file.  The file seems to be missing \"mpc.version = ...\"",
            )
        )
        case["source_version"] = "0.0.0+"
    end

    if haskey(matlab_data, "mpc.baseMVA")
        case["baseMVA"] = matlab_data["mpc.baseMVA"]
    else
        @info(
            string(
                "no baseMVA found in matpower file.  The file seems to be missing \"mpc.baseMVA = ...\"",
            )
        )
        case["baseMVA"] = 1.0
    end

    if haskey(matlab_data, "mpc.bus")
        buses = []
        for bus_row in matlab_data["mpc.bus"]
            bus_data = row_to_typed_dict(bus_row, _mp_bus_columns)
            bus_data["index"] = check_type(Int, bus_row[1])
            bus_data["source_id"] = ["bus", bus_data["index"]]
            push!(buses, bus_data)
        end
        case["bus"] = buses
    else
        error(
            string(
                "no bus table found in matpower file.  The file seems to be missing \"mpc.bus = [...];\"",
            ),
        )
    end

    if haskey(matlab_data, "mpc.gen")
        gens = []
        for (i, gen_row) in enumerate(matlab_data["mpc.gen"])
            gen_data = row_to_typed_dict(gen_row, _mp_gen_columns)
            gen_data["index"] = i
            gen_data["source_id"] = ["gen", i]
            push!(gens, gen_data)
        end
        case["gen"] = gens
    else
        error(
            string(
                "no gen table found in matpower file.  The file seems to be missing \"mpc.gen = [...];\"",
            ),
        )
    end

    if haskey(matlab_data, "mpc.branch")
        branches = []
        for (i, branch_row) in enumerate(matlab_data["mpc.branch"])
            branch_data = row_to_typed_dict(branch_row, _mp_branch_columns)
            branch_data["index"] = i
            branch_data["source_id"] = ["branch", i]
            push!(branches, branch_data)
        end
        case["branch"] = branches
    else
        error(
            string(
                "no branch table found in matpower file.  The file seems to be missing \"mpc.branch = [...];\"",
            ),
        )
    end

    if haskey(matlab_data, "mpc.dcline")
        dclines = []
        for (i, dcline_row) in enumerate(matlab_data["mpc.dcline"])
            dcline_data = row_to_typed_dict(dcline_row, _mp_dcline_columns)
            dcline_data["index"] = i
            dcline_data["source_id"] = ["dcline", i]
            push!(dclines, dcline_data)
        end
        case["dcline"] = dclines
    end

    if haskey(matlab_data, "mpc.storage")
        storage = []
        for (i, storage_row) in enumerate(matlab_data["mpc.storage"])
            storage_data = row_to_typed_dict(storage_row, _mp_storage_columns)
            storage_data["index"] = i
            storage_data["source_id"] = ["storage", i]
            push!(storage, storage_data)
        end
        case["storage"] = storage
    end

    if haskey(matlab_data, "mpc.switch")
        switch = []
        for (i, switch_row) in enumerate(matlab_data["mpc.switch"])
            switch_data = row_to_typed_dict(switch_row, _mp_switch_columns)
            switch_data["index"] = i
            switch_data["source_id"] = ["switch", i]
            push!(switch, switch_data)
        end
        case["switch"] = switch
    end

    if haskey(matlab_data, "mpc.bus_name")
        bus_names = []
        for (i, bus_name_row) in enumerate(matlab_data["mpc.bus_name"])
            bus_name_data = row_to_typed_dict(bus_name_row, _mp_bus_name_columns)
            bus_name_data["index"] = i
            bus_name_data["source_id"] = ["bus_name", i]
            push!(bus_names, bus_name_data)
        end
        case["bus_name"] = bus_names

        if length(case["bus_name"]) != length(case["bus"])
            error(
                "incorrect Matpower file, the number of bus names ($(length(case["bus_name"]))) is inconsistent with the number of buses ($(length(case["bus"]))).\n",
            )
        end
    end

    if haskey(matlab_data, "mpc.gencost")
        gencost = []
        for (i, gencost_row) in enumerate(matlab_data["mpc.gencost"])
            gencost_data = _mp_cost_data(gencost_row)
            gencost_data["index"] = i
            gencost_data["source_id"] = ["gencost", i]
            push!(gencost, gencost_data)
        end
        case["gencost"] = gencost

        if length(case["gencost"]) != length(case["gen"]) &&
           length(case["gencost"]) != 2 * length(case["gen"])
            error(
                "incorrect Matpower file, the number of generator cost functions ($(length(case["gencost"]))) is inconsistent with the number of generators ($(length(case["gen"]))).\n",
            )
        end
    end

    if haskey(matlab_data, "mpc.dclinecost")
        dclinecosts = []
        for (i, dclinecost_row) in enumerate(matlab_data["mpc.dclinecost"])
            dclinecost_data = _mp_cost_data(dclinecost_row)
            dclinecost_data["index"] = i
            dclinecost_data["source_id"] = ["dclinecost", i]
            push!(dclinecosts, dclinecost_data)
        end
        case["dclinecost"] = dclinecosts

        if length(case["dclinecost"]) != length(case["dcline"])
            error(
                "incorrect Matpower file, the number of dcline cost functions ($(length(case["dclinecost"]))) is inconsistent with the number of dclines ($(length(case["dcline"]))).\n",
            )
        end
    end

    for k in keys(matlab_data)
        if !in(k, _mp_data_names) && startswith(k, "mpc.")
            case_name = k[5:length(k)]
            value = matlab_data[k]
            if isa(value, Array)
                column_names = []
                if haskey(colnames, k)
                    column_names = colnames[k]
                end
                tbl = []
                for (i, row) in enumerate(matlab_data[k])
                    row_data = row_to_dict(row, column_names)
                    row_data["index"] = i
                    row_data["source_id"] = [case_name, i]
                    push!(tbl, row_data)
                end
                case[case_name] = tbl
                @info(
                    "extending matpower format with data: $(case_name) $(length(tbl))x$(length(tbl[1])-1)"
                )
            else
                case[case_name] = value
                @info("extending matpower format with constant data: $(case_name)")
            end
        end
    end

    return case
end

""
function _mp_cost_data(cost_row)
    ncost = cost_row[4]
    model = cost_row[1]
    if model == 1
        nr_parameters = ncost * 2
    elseif model == 2
        nr_parameters = ncost
    end
    cost_data = Dict{String, Any}(
        "model" => check_type(Int, cost_row[1]),
        "startup" => check_type(Float64, cost_row[2]),
        "shutdown" => check_type(Float64, cost_row[3]),
        "ncost" => check_type(Int, cost_row[4]),
        "cost" => [check_type(Float64, x) for x in cost_row[5:(5 + nr_parameters - 1)]],
    )

    #=
    # skip this literal interpretation, as its hard to invert
    cost_values = [check_type(Float64, x) for x in cost_row[5:length(cost_row)]]
    if cost_data["model"] == 1:
        if length(cost_values)%2 != 0
            error("incorrect matpower file, odd number of pwl cost function values")
        end
        for i in 0:(length(cost_values)/2-1)
            p_idx = 1+2*i
            f_idx = 2+2*i
            cost_data["p_$(i)"] = cost_values[p_idx]
            cost_data["f_$(i)"] = cost_values[f_idx]
        end
    else:
        for (i,v) in enumerate(cost_values)
            cost_data["c_$(length(cost_values)+1-i)"] = v
        end
    =#
    return cost_data
end

### Data and functions specific to PowerModels format ###

"""
Converts a Matpower dict into a PowerModels dict
"""
function _matpower_to_powermodels!(mp_data::Dict{String, <:Any})
    pm_data = mp_data

    # required default values
    if !haskey(pm_data, "dcline")
        pm_data["dcline"] = []
    end
    if !haskey(pm_data, "gencost")
        pm_data["gencost"] = []
    end
    if !haskey(pm_data, "dclinecost")
        pm_data["dclinecost"] = []
    end
    if !haskey(pm_data, "storage")
        pm_data["storage"] = []
    end
    if !haskey(pm_data, "switch")
        pm_data["switch"] = []
    end

    # translate component models
    _mp2pm_branch!(pm_data)
    _mp2pm_dcline!(pm_data)

    # translate cost models
    _add_dcline_costs!(pm_data)

    # merge data tables
    _merge_bus_name_data!(pm_data)
    _merge_generator_cost_data!(pm_data)
    _merge_generic_data!(pm_data)

    # split loads and shunts from buses
    _split_loads_shunts!(pm_data)

    # use once available
    arrays_to_dicts!(pm_data)

    for optional in ["dcline", "load", "shunt", "storage", "switch"]
        if length(pm_data[optional]) == 0
            pm_data[optional] = Dict{String, Any}()
        end
    end

    return pm_data
end

"""
    _split_loads_shunts!(data)

Seperates Loads and Shunts in `data` under separate "load" and "shunt" keys in the
PowerModels data format. Includes references to originating bus via "load_bus"
and "shunt_bus" keys, respectively.
"""
function _split_loads_shunts!(data::Dict{String, Any})
    data["load"] = []
    data["shunt"] = []

    load_num = 1
    shunt_num = 1
    for (i, bus) in enumerate(data["bus"])
        if bus["pd"] != 0.0 || bus["qd"] != 0.0
            append!(
                data["load"],
                [
                    Dict{String, Any}(
                        "pd" => bus["pd"],
                        "qd" => bus["qd"],
                        "load_bus" => bus["bus_i"],
                        "status" => convert(Int8, bus["bus_type"] != 4),
                        "index" => load_num,
                        "source_id" => ["bus", bus["bus_i"]],
                    ),
                ],
            )
            load_num += 1
        end

        if bus["gs"] != 0.0 || bus["bs"] != 0.0
            append!(
                data["shunt"],
                [
                    Dict{String, Any}(
                        "gs" => bus["gs"],
                        "bs" => bus["bs"],
                        "shunt_bus" => bus["bus_i"],
                        "status" => convert(Int8, bus["bus_type"] != 4),
                        "index" => shunt_num,
                        "source_id" => ["bus", bus["bus_i"]],
                    ),
                ],
            )
            shunt_num += 1
        end

        for key in ["pd", "qd", "gs", "bs"]
            delete!(bus, key)
        end
    end
end

"sets all branch transformer taps to 1.0, to simplify branch models"
function _mp2pm_branch!(data::Dict{String, Any})
    branches = [branch for branch in data["branch"]]
    if haskey(data, "ne_branch")
        append!(branches, data["ne_branch"])
    end
    for branch in branches
        if branch["tap"] == 0.0
            branch["transformer"] = false
            branch["tap"] = 1.0
        else
            branch["transformer"] = true
        end

        branch["g_fr"] = 0.0
        branch["g_to"] = 0.0

        branch["b_fr"] = branch["br_b"] / 2.0
        branch["b_to"] = branch["br_b"] / 2.0

        delete!(branch, "br_b")

        if branch["rate_a"] == 0.0
            delete!(branch, "rate_a")
        end
        if branch["rate_b"] == 0.0
            delete!(branch, "rate_b")
        end
        if branch["rate_c"] == 0.0
            delete!(branch, "rate_c")
        end
    end
end

"adds pmin and pmax values at to and from buses"
function _mp2pm_dcline!(data::Dict{String, Any})
    for dcline in data["dcline"]
        pmin = dcline["pmin"]
        pmax = dcline["pmax"]
        loss0 = dcline["loss0"]
        loss1 = dcline["loss1"]

        delete!(dcline, "pmin")
        delete!(dcline, "pmax")

        if pmin >= 0 && pmax >= 0
            pminf = pmin
            pmaxf = pmax
            pmint = loss0 - pmaxf * (1 - loss1)
            pmaxt = loss0 - pminf * (1 - loss1)
        end
        if pmin >= 0 && pmax < 0
            pminf = pmin
            pmint = pmax
            pmaxf = (-pmint + loss0) / (1 - loss1)
            pmaxt = loss0 - pminf * (1 - loss1)
        end
        if pmin < 0 && pmax >= 0
            pmaxt = -pmin
            pmaxf = pmax
            pminf = (-pmaxt + loss0) / (1 - loss1)
            pmint = loss0 - pmaxf * (1 - loss1)
        end
        if pmin < 0 && pmax < 0
            pmaxt = -pmin
            pmint = pmax
            pmaxf = (-pmint + loss0) / (1 - loss1)
            pminf = (-pmaxt + loss0) / (1 - loss1)
        end

        dcline["pmaxt"] = pmaxt
        dcline["pmint"] = pmint
        dcline["pmaxf"] = pmaxf
        dcline["pminf"] = pminf

        # preserve the old pmin and pmax values
        dcline["mp_pmin"] = pmin
        dcline["mp_pmax"] = pmax

        dcline["pt"] = -dcline["pt"] # matpower has opposite convention
        dcline["qf"] = -dcline["qf"] # matpower has opposite convention
        dcline["qt"] = -dcline["qt"] # matpower has opposite convention
    end
end

"adds dcline costs, if gen costs exist"
function _add_dcline_costs!(data::Dict{String, Any})
    if length(data["gencost"]) > 0 &&
       length(data["dclinecost"]) <= 0 &&
       length(data["dcline"]) > 0
        @info("added zero cost function data for dclines")
        model = data["gencost"][1]["model"]
        if model == 1
            for (i, dcline) in enumerate(data["dcline"])
                dclinecost = Dict(
                    "index" => i,
                    "model" => 1,
                    "startup" => 0.0,
                    "shutdown" => 0.0,
                    "ncost" => 2,
                    "cost" => [dcline["pminf"], 0.0, dcline["pmaxf"], 0.0],
                )
                push!(data["dclinecost"], dclinecost)
            end
        else
            for (i, dcline) in enumerate(data["dcline"])
                dclinecost = Dict(
                    "index" => i,
                    "model" => 2,
                    "startup" => 0.0,
                    "shutdown" => 0.0,
                    "ncost" => 3,
                    "cost" => [0.0, 0.0, 0.0],
                )
                push!(data["dclinecost"], dclinecost)
            end
        end
    end
end

"merges generator cost functions into generator data, if costs exist"
function _merge_generator_cost_data!(data::Dict{String, Any})
    if haskey(data, "gencost")
        for (i, gencost) in enumerate(data["gencost"])
            gen = data["gen"][i]
            @assert(gen["index"] == gencost["index"])
            delete!(gencost, "index")
            delete!(gencost, "source_id")

            _check_keys(gen, keys(gencost))
            merge!(gen, gencost)
        end
        delete!(data, "gencost")
    end

    if haskey(data, "dclinecost")
        for (i, dclinecost) in enumerate(data["dclinecost"])
            dcline = data["dcline"][i]
            @assert(dcline["index"] == dclinecost["index"])
            delete!(dclinecost, "index")
            delete!(dclinecost, "source_id")

            _check_keys(dcline, keys(dclinecost))
            merge!(dcline, dclinecost)
        end
        delete!(data, "dclinecost")
    end
end

"merges bus name data into buses, if names exist"
function _merge_bus_name_data!(data::Dict{String, Any})
    if haskey(data, "bus_name")
        # can assume same length is same as bus
        # this is validated during matpower parsing
        for (i, bus_name) in enumerate(data["bus_name"])
            bus = data["bus"][i]
            delete!(bus_name, "index")
            delete!(bus_name, "source_id")

            _check_keys(bus, keys(bus_name))
            merge!(bus, bus_name)
        end
        delete!(data, "bus_name")
    end
end

"merges Matpower tables based on the table extension syntax"
function _merge_generic_data!(data::Dict{String, Any})
    mp_matrix_names = [name[5:length(name)] for name in _mp_data_names]

    key_to_delete = []
    for (k, v) in data
        if isa(v, Array)
            for mp_name in mp_matrix_names
                if startswith(k, "$(mp_name)_")
                    mp_matrix = data[mp_name]
                    push!(key_to_delete, k)

                    if length(mp_matrix) != length(v)
                        error(
                            "failed to extend the matpower matrix \"$(mp_name)\" with the matrix \"$(k)\" because they do not have the same number of rows, $(length(mp_matrix)) and $(length(v)) respectively.",
                        )
                    end

                    @info(
                        "extending matpower format by appending matrix \"$(k)\" in to \"$(mp_name)\""
                    )

                    for (i, row) in enumerate(mp_matrix)
                        merge_row = v[i]
                        #@assert(row["index"] == merge_row["index"]) # note this does not hold for the bus table
                        delete!(merge_row, "index")
                        delete!(merge_row, "source_id")
                        for key in keys(merge_row)
                            if haskey(row, key)
                                error(
                                    "failed to extend the matpower matrix \"$(mp_name)\" with the matrix \"$(k)\" because they both share \"$(key)\" as a column name.",
                                )
                            end
                            row[key] = merge_row[key]
                        end
                    end

                    break # out of mp_matrix_names loop
                end
            end
        end
    end

    for key in key_to_delete
        delete!(data, key)
    end
end

"Export power network data in the matpower format"
function export_matpower(data::Dict{String, Any})
    return sprint(export_matpower, data)
end

" Get a default value for dict entry "
function _get_default(dict, key, default = 0.0)
    if haskey(dict, key) && dict[key] != NaN
        return dict[key]
    end
    return default
end

""
function _check_keys(data, keys)
    for key in keys
        if haskey(data, key)
            error("attempting to overwrite value of $(key) in PowerModels data,\n$(data)")
        end
    end
end

"Export power network data in the matpower format"
function export_matpower(io::IO, data::Dict{String, Any})
    data = deepcopy(data)

    #convert data to mixed unit
    if data["per_unit"]
        make_mixed_units!(data)
    end

    # make all costs have the name number of items (to prepare for table export)
    standardize_cost_terms!(data)

    # create some useful maps and data structures
    buses = Dict{Int, Dict}()
    for (idx, bus) in data["bus"]
        buses[bus["index"]] = bus
    end
    generators = Dict{Int, Dict}()
    for (idx, gen) in data["gen"]
        generators[gen["index"]] = gen
    end
    storage = Dict{Int, Dict}()
    for (idx, strg) in data["storage"]
        storage[strg["index"]] = strg
    end
    branches = Dict{Int, Dict}()
    for (idx, branch) in data["branch"]
        branches[branch["index"]] = branch
    end
    dclines = Dict{Int, Dict}()
    for (idx, dcline) in data["dcline"]
        dclines[dcline["index"]] = dcline
    end
    ne_branches = Dict{Int, Dict}()
    if haskey(data, "ne_branch")
        for (idx, branch) in data["ne_branch"]
            ne_branches[branch["index"]] = branch
        end
    end

    pd = Dict{Int, Float64}()
    qd = Dict{Int, Float64}()
    gs = Dict{Int, Float64}()
    bs = Dict{Int, Float64}()

    # collect all the loads
    for (idx, bus) in sort(data["bus"])
        pd[bus["index"]] = 0
        qd[bus["index"]] = 0
    end
    for (idx, load) in sort(data["load"])
        bus = buses[load["load_bus"]]
        pd[bus["index"]] = pd[bus["index"]] + load["pd"]
        qd[bus["index"]] = qd[bus["index"]] + load["qd"]
    end

    # collect all the shunts
    for (idx, bus) in sort(data["bus"])
        gs[bus["index"]] = 0
        bs[bus["index"]] = 0
    end
    for (idx, shunt) in data["shunt"]
        bus = buses[shunt["shunt_bus"]]
        bs[bus["index"]] = bs[bus["index"]] + shunt["bs"]
        gs[bus["index"]] = gs[bus["index"]] + shunt["gs"]
    end

    mvaBase = data["baseMVA"]

    # Print the header information
    println(io, "%% MATPOWER Case Format : Version 2")
    println(io, "function mpc = ", data["name"])
    println(io, "mpc.version = '2';")
    println(io)
    println(io, "%%-----  Power Flow Data  -----%%")
    println(io, "%% system MVA base")
    println(io, "mpc.baseMVA = ", mvaBase, ";")

    # Print the bus data
    println(io, "%% bus data")
    println(
        io,
        "%    bus_i    type    Pd    Qd    Gs    Bs    area    Vm    Va    baseKV    zone    Vmax    Vmin",
    )
    println(io, "mpc.bus = [")
    for (idx, bus) in sort(buses)
        println(
            io,
            "\t",
            bus["index"],
            "\t",
            bus["bus_type"],
            "\t",
            _get_default(pd, bus["index"]),
            "\t",
            _get_default(qd, bus["index"]),
            "\t",
            _get_default(gs, bus["index"]),
            "\t",
            _get_default(bs, bus["index"]),
            "\t",
            _get_default(bus, "area"),
            "\t",
            _get_default(bus, "vm"),
            "\t",
            _get_default(bus, "va"),
            "\t",
            _get_default(bus, "base_kv"),
            "\t",
            _get_default(bus, "zone"),
            "\t",
            _get_default(bus, "vmax"),
            "\t",
            _get_default(bus, "vmin"),
            "\t",
            _get_default(bus, "lam_p", ""),
            "\t",
            _get_default(bus, "lam_q", ""),
            "\t",
            _get_default(bus, "mu_vmax", ""),
            "\t",
            _get_default(bus, "mu_vmin", ""),
        )
    end
    println(io, "];")
    println(io)

    # Print the bus names
    if haskey(collect(values(buses))[1], "bus_name")
        println(io, "%% bus names")
        println(io, "mpc.bus_name = {")
        for (idx, bus) in sort(buses)
            println(io, "\t'", bus["bus_name"], "'")
        end
        println(io, "};")
        println(io)
    end

    # Print the generator data
    println(io, "%% generator data")
    println(
        io,
        "%    bus    Pg    Qg    Qmax    Qmin    Vg    mBase    status    Pmax    Pmin    Pc1    Pc2    Qc1min    Qc1max    Qc2min    Qc2max    ramp_agc    ramp_10    ramp_30    ramp_q    apf",
    )
    println(io, "mpc.gen = [")
    i = 1
    for (idx, gen) in sort(generators)
        if idx != gen["index"]
            @info(
                "The index of the generator does not match the matpower assigned index. Any data that uses generator indexes for reference is corrupted."
            )
        end
        println(
            io,
            "\t",
            gen["gen_bus"],
            "\t",
            _get_default(gen, "pg"),
            "\t",
            _get_default(gen, "qg"),
            "\t",
            _get_default(gen, "qmax"),
            "\t",
            _get_default(gen, "qmin"),
            "\t",
            _get_default(gen, "vg"),
            "\t",
            _get_default(gen, "mbase"),
            "\t",
            _get_default(gen, "gen_status"),
            "\t",
            _get_default(gen, "pmax"),
            "\t",
            _get_default(gen, "pmin"),
            "\t",
            _get_default(gen, "pc1", ""),
            "\t",
            _get_default(gen, "pc2", ""),
            "\t",
            _get_default(gen, "qc1min", ""),
            "\t",
            _get_default(gen, "qc1max", ""),
            "\t",
            _get_default(gen, "qc2min", ""),
            "\t",
            _get_default(gen, "qc2max", ""),
            "\t",
            _get_default(gen, "ramp_agc", ""),
            "\t",
            (haskey(gen, "ramp_10") ? gen["ramp_10"] : haskey(gen, "ramp_30") ? 0 : ""),
            "\t",
            _get_default(gen, "ramp_30", ""),
            "\t",
            _get_default(gen, "ramp_q", ""),
            "\t",
            _get_default(gen, "apf", ""),
            "\t",
            _get_default(gen, "mu_pmax", ""),
            "\t",
            _get_default(gen, "mu_pmin", ""),
            "\t",
            _get_default(gen, "mu_qmax", ""),
            "\t",
            _get_default(gen, "mu_qmin", ""),
        )
        i = i + 1
    end
    println(io, "];")
    println(io)

    if length(storage) > 0
        # Print the storage data
        println(io, "%% storage data")
        println(
            io,
            "%    storage_bus    ps    qs    energy    energy_rating    charge_rating    discharge_rating    charge_efficiency    discharge_efficiency    thermal_rating    qmin    qmax    r    x    standby_loss    status",
        )
        println(io, "mpc.storage = [")
        i = 1
        for (idx, strg) in sort(storage)
            if idx != strg["index"]
                @info(
                    "The index of the storage does not match the matpower assigned index. Any data that uses storage indexes for reference is corrupted."
                )
            end
            println(
                io,
                "\t",
                strg["storage_bus"],
                "\t",
                _get_default(strg, "ps"),
                "\t",
                _get_default(strg, "qs"),
                "\t",
                _get_default(strg, "energy"),
                "\t",
                _get_default(strg, "energy_rating"),
                "\t",
                _get_default(strg, "charge_rating"),
                "\t",
                _get_default(strg, "discharge_rating"),
                "\t",
                _get_default(strg, "charge_efficiency"),
                "\t",
                _get_default(strg, "discharge_efficiency"),
                "\t",
                _get_default(strg, "thermal_rating"),
                "\t",
                _get_default(strg, "qmin"),
                "\t",
                _get_default(strg, "qmax"),
                "\t",
                _get_default(strg, "r"),
                "\t",
                _get_default(strg, "x"),
                "\t",
                _get_default(strg, "standby_loss"),
                "\t",
                _get_default(strg, "status"),
            )
            i = i + 1
        end
        println(io, "];")
        println(io)
    end

    # Print the branch data
    println(io, "%% branch data")
    println(
        io,
        "%    fbus    tbus    r    x    b    rateA    rateB    rateC    ratio    angle    status    angmin    angmax",
    )
    println(io, "mpc.branch = [")
    i = 1
    for (idx, branch) in sort(branches)
        if idx != branch["index"]
            @info(
                "The index of the branch does not match the matpower assigned index. Any data that uses branch indexes for reference is corrupted."
            )
        end
        println(
            io,
            "\t",
            _get_default(branch, "f_bus"),
            "\t",
            _get_default(branch, "t_bus"),
            "\t",
            _get_default(branch, "br_r"),
            "\t",
            _get_default(branch, "br_x"),
            "\t",
            (branch["b_to"] + branch["b_fr"]) != NaN ? (branch["b_to"] + branch["b_fr"]) :
            0,
            "\t",
            _get_default(branch, "rate_a"),
            "\t",
            _get_default(branch, "rate_b"),
            "\t",
            _get_default(branch, "rate_c"),
            "\t",
            (branch["transformer"] ? branch["tap"] : 0),
            "\t",
            (branch["transformer"] ? branch["shift"] : 0),
            "\t",
            _get_default(branch, "br_status"),
            "\t",
            _get_default(branch, "angmin"),
            "\t",
            _get_default(branch, "angmax"),
            "\t",
            _get_default(branch, "pf", ""),
            "\t",
            _get_default(branch, "qf", ""),
            "\t",
            _get_default(branch, "pt", ""),
            "\t",
            _get_default(branch, "qt", ""),
            "\t",
            _get_default(branch, "mu_sf", ""),
            "\t",
            _get_default(branch, "mu_st", ""),
            "\t",
            _get_default(branch, "mu_angmin", ""),
            "\t",
            _get_default(branch, "mu_angmax", ""),
        )
        i = i + 1
    end
    println(io, "];")
    println(io)

    if length(dclines) > 0
        # print the dcline data
        println(io, "%% dcline data")
        println(
            io,
            "%	fbus	tbus	status	Pf	Pt	Qf	Qt	Vf	Vt	Pmin	Pmax	QminF	QmaxF	QminT	QmaxT	loss0	loss1",
        )
        println(io, "mpc.dcline = [")
        for (idx, dcline) in sort(dclines)
            println(
                io,
                "\t",
                _get_default(dcline, "f_bus"),
                "\t",
                _get_default(dcline, "t_bus"),
                "\t",
                _get_default(dcline, "br_status"),
                "\t",
                _get_default(dcline, "pf"),
                "\t",
                -_get_default(dcline, "pt"),  # opposite convention
                "\t",
                -_get_default(dcline, "qf"),  # opposite convention
                "\t",
                -_get_default(dcline, "qt"),  # opposite convention
                "\t",
                _get_default(dcline, "vf"),
                "\t",
                _get_default(dcline, "vt"),
                "\t",
                (
                    haskey(dcline, "mp_pmin") ? dcline["mp_pmin"] :
                    min(dcline["pmaxt"], dcline["pmint"], dcline["pmaxf"], dcline["pminf"])
                ),
                "\t",
                (
                    haskey(dcline, "mp_pmax") ? dcline["mp_pmax"] :
                    max(dcline["pmaxt"], dcline["pmint"], dcline["pmaxf"], dcline["pminf"])
                ),
                "\t",
                _get_default(dcline, "qminf"),
                "\t",
                _get_default(dcline, "qmaxf"),
                "\t",
                _get_default(dcline, "qmint"),
                "\t",
                _get_default(dcline, "qmaxt"),
                "\t",
                _get_default(dcline, "loss0"),
                "\t",
                _get_default(dcline, "loss1"),
                "\t",
                _get_default(dcline, "mu_pmin", ""),
                "\t",
                _get_default(dcline, "mu_pmax", ""),
                "\t",
                _get_default(dcline, "mu_qminf", ""),
                "\t",
                _get_default(dcline, "mu_qmaxf", ""),
                "\t",
                _get_default(dcline, "mu_qmint", ""),
                "\t",
                _get_default(dcline, "mu_qmaxt", ""),
            )
        end
        println(io, "];")
        println(io)
    end

    # Print the gen cost data
    _export_cost_data(io, generators, "mpc.gencost")

    # Print the dcline cost data
    _export_cost_data(io, dclines, "mpc.dclinecost")

    # ne branch is not part of the matpower specs. However, it is treated as a special case by the matpower parser
    # for example, br_b is converted into b_to and b_fr
    if haskey(data, "ne_branch")
        println(
            io,
            "%column_names%	f_bus	t_bus	br_r	br_x	br_b	rate_a	rate_b	rate_c	tap	shift	br_status	angmin	angmax	construction_cost",
        )
        println(io, "mpc.ne_branch = [")
        i = 1
        for (idx, branch) in sort(ne_branches)
            if idx != branch["index"]
                @info(
                    "The index of the ne_branch does not match the matpower assigned index. Any data that uses branch indexes for reference is corrupted."
                )
            end
            println(
                io,
                "\t",
                branch["f_bus"],
                "\t",
                branch["t_bus"],
                "\t",
                _get_default(branch, "br_r"),
                "\t",
                _get_default(branch, "br_x"),
                "\t",
                (haskey(branch, "b_to") ? branch["b_to"] + branch["b_fr"] : 0),
                "\t",
                _get_default(branch, "rate_a"),
                "\t",
                _get_default(branch, "rate_b"),
                "\t",
                _get_default(branch, "rate_c"),
                "\t",
                (branch["transformer"] ? branch["tap"] : 0),
                "\t",
                (branch["transformer"] ? branch["shift"] : 0),
                "\t",
                _get_default(branch, "br_status"),
                "\t",
                _get_default(branch, "angmin"),
                "\t",
                _get_default(branch, "angmax"),
                "\t",
                _get_default(branch, "construction_cost"),
            )
            i = i + 1
        end
        println(io, "];")
        println(io)
    end

    if all(haskey(bus, "name") for (i, bus) in buses)
        # Print the bus name data
        println(io, "%% bus names")
        println(io, "mpc.bus_name = {")
        for (idx, bus) in sort(buses)
            println(io, "\t", "'", bus["name"], "'")
        end
        println(io, "};")
        println(io)
    end

    # Print the extra bus data
    _export_extra_data(
        io,
        data,
        "bus",
        Set([
            "index",
            "source_id",
            "gs",
            "bs",
            "zone",
            "bus_i",
            "bus_type",
            "qd",
            "vmax",
            "area",
            "vmin",
            "va",
            "vm",
            "base_kv",
            "pd",
            "name",
            "lam_p",
            "lam_q",
            "mu_vmax",
            "mu_vmin",
        ]);
        postfix = "_data",
    )

    # Print the extra generator data
    _export_extra_data(
        io,
        data,
        "gen",
        Set([
            "index",
            "source_id",
            "gen_bus",
            "pg",
            "qg",
            "qmax",
            "qmin",
            "vg",
            "mbase",
            "gen_status",
            "pmax",
            "pmin",
            "pc1",
            "pc2",
            "qc1min",
            "qc1max",
            "qc2min",
            "qc2max",
            "ramp_agc",
            "ramp_10",
            "ramp_30",
            "ramp_q",
            "apf",
            "ncost",
            "model",
            "shutdown",
            "startup",
            "cost",
            "mu_pmax",
            "mu_pmin",
            "mu_qmax",
            "mu_qmin",
        ]);
        postfix = "_data",
    )

    # Print the extra storage data
    _export_extra_data(
        io,
        data,
        "storage",
        Set([
            "index",
            "source_id",
            "storage_bus",
            "ps",
            "qs",
            "energy",
            "energy_rating",
            "charge_rating",
            "discharge_rating",
            "charge_efficiency",
            "discharge_efficiency",
            "thermal_rating",
            "qmin",
            "qmax",
            "r",
            "x",
            "standby_loss",
            "status",
        ]);
        postfix = "_data",
    )

    # Print the extra branch data
    _export_extra_data(
        io,
        data,
        "branch",
        Set([
            "index",
            "source_id",
            "f_bus",
            "t_bus",
            "br_r",
            "br_x",
            "br_b",
            "b_to",
            "b_fr",
            "rate_a",
            "rate_b",
            "rate_c",
            "tap",
            "shift",
            "br_status",
            "angmin",
            "angmax",
            "transformer",
            "g_to",
            "g_fr",
            "pf",
            "qf",
            "pt",
            "qt",
            "mu_sf",
            "mu_st",
            "mu_angmin",
            "mu_angmax",
        ]);
        postfix = "_data",
    )

    # Print the extra dcline data
    _export_extra_data(
        io,
        data,
        "dcline",
        Set([
            "index",
            "source_id",
            "mu_qmaxt",
            "mu_qmint",
            "mu_qmaxf",
            "mu_qminf",
            "mu_pmax",
            "mu_pmin",
            "loss0",
            "loss1",
            "qmint",
            "qmaxt",
            "pmin",
            "pmax",
            "qminf",
            "qmaxf",
            "f_bus",
            "t_bus",
            "br_status",
            "pf",
            "pt",
            "qf",
            "qt",
            "vf",
            "vt",
            "ncost",
            "model",
            "shutdown",
            "pmaxt",
            "startup",
            "pmint",
            "cost",
            "pminf",
            "pmaxf",
            "mp_pmax",
            "mp_pmin",
        ]);
        postfix = "_data",
    )

    # Print the extra ne_branch data
    if haskey(data, "ne_branch")
        _export_extra_data(
            io,
            data,
            "ne_branch",
            Set([
                "index",
                "source_id",
                "f_bus",
                "t_bus",
                "br_r",
                "br_x",
                "br_b",
                "b_to",
                "b_fr",
                "rate_a",
                "rate_b",
                "rate_c",
                "tap",
                "shift",
                "br_status",
                "angmin",
                "angmax",
                "transformer",
                "construction_cost",
                "g_to",
                "g_fr",
            ]);
            postfix = "_data",
        )
    end

    # print the extra load data
    _export_extra_data(
        io,
        data,
        "load",
        Set(["index", "source_id", "load_bus", "status", "qd", "pd"]);
        postfix = "_data",
    )

    # print the extra shunt data
    _export_extra_data(
        io,
        data,
        "shunt",
        Set(["index", "source_id", "shunt_bus", "status", "gs", "bs"]);
        postfix = "_data",
    )

    # print the extra component data
    for (key, value) in data
        if key != "bus" &&
           key != "gen" &&
           key != "branch" &&
           key != "load" &&
           key != "shunt" &&
           key != "storage" &&
           key != "dcline" &&
           key != "ne_branch" &&
           key != "version" &&
           key != "baseMVA" &&
           key != "per_unit" &&
           key != "name" &&
           key != "source_type" &&
           key != "source_version"
            _export_extra_data(io, data, key)
        end
    end
end

"Export fields of a component type"
function _export_extra_data(
    io::IO,
    data::Dict{String, <:Any},
    component,
    excluded_fields = Set(["index", "source_id"]);
    postfix = "",
)
    if isa(data[component], Int) ||
       isa(data[component], Int) ||
       isa(data[component], Float64)
        println(io, "mpc.", component, " = ", data[component], ";")
        println(io)
        return
    end

    if isa(data[component], String) || isa(data[component], SubString{String})
        println(io, "mpc.", component, " = '", data[component], "';")
        println(io)
        return
    end

    if !isa(data[component], Dict)
        return
    end

    if length(data[component]) == 0
        return
    end

    # Gather the fields
    included_fields = []
    c = collect(values(data[component]))[1]
    for (key, value) in c
        if !in(key, excluded_fields)
            push!(included_fields, key)
        end
    end

    if length(included_fields) == 0
        return
    end

    # Print the header
    print(io, "%column_names% ")
    for field in included_fields
        print(io, field)
        print(io, " ")
    end
    println(io)
    println(io, "mpc.", component, postfix, " = {")

    # sort the data
    components = Dict{Int, Dict}()
    for (idx, c) in data[component]
        components[c["index"]] = c
    end

    # print the data
    i = 1
    for (idx, c) in sort(components)
        if idx != c["index"]
            @info(
                "The index of a component does not match the matpower assigned index. Any data that uses component indexes for reference is corrupted."
            )
        end

        for field in included_fields
            print(io, "\t")
            if isa(c[field], Union{String, SubString})
                print(io, "'")
            end
            print(io, c[field])
            if isa(c[field], Union{String, SubString})
                print(io, "'")
            end
        end
        println(io)
        i = i + 1
    end
    println(io, "};")
    println(io)
end

"Export cost data"
function _export_cost_data(io::IO, components::Dict{Int, Dict}, prefix::String)
    if length(components) <= 0
        return
    end

    a_comp = collect(values(components))[1]
    if haskey(a_comp, "cost")
        ncost = length(a_comp["cost"])
        model = a_comp["model"]

        for (i, comp) in components
            if length(comp["cost"]) != ncost || comp["model"] != model
                @info("heterogeneous cost functions will be ommited in Matpower data")
                return
            end
        end

        println(io, "%%-----  OPF Data  -----%%")
        println(io, "%% cost data")
        println(io, "%    1    startup    shutdown    n    x1    y1    ...    xn    yn")
        println(io, "%    2    startup    shutdown    n    c(n-1)    ...    c0")
        println(io, prefix, " = [")

        for (idx, gen) in (sort(components))
            if gen["model"] == 1
                print(
                    io,
                    "\t1\t",
                    gen["startup"],
                    "\t",
                    gen["shutdown"],
                    "\t",
                    (length(gen["cost"]) / 2),
                ),
                for l in 1:length(gen["cost"])
                    print(io, "\t", gen["cost"][l])
                end
            else
                print(
                    io,
                    "\t2\t",
                    gen["startup"],
                    "\t",
                    gen["shutdown"],
                    "\t",
                    length(gen["cost"]),
                ),
                for l in 1:length(gen["cost"])
                    print(io, "\t", gen["cost"][l])
                end
            end
            println(io)
        end
        println(io, "];")
        println(io)
    end
end


text:
using Arpack, SparseArrays, LinearAlgebra

# define functions used here
include("ades.jl")
include("acre.jl")
include("pascaltriangle.jl")
include("in2b.jl")
include("b2in.jl")
include("epsilon.jl")
include("vijkl.jl")

function Hintfunc(Msize0::Int64, Np::Int64)

    Msize = Msize0*2

    matp = zeros(Int,Msize+1,Np+1);
    pascaltriangle!(Msize,Np,matp) # the size is Msize+1 times Np+1
    maxmatp = matp[Msize+1,Np+1] # the indices are m+1 and n+1 for N^m_ns

    # defines vectors and matrices
    vecmbnn = spzeros(Int64,Msize+1);
    vecmbnnj = spzeros(Int64,Msize+1);
    vecmbnnij = spzeros(Int64,Msize+1);
    # Hsoc = spzeros(ComplexF64,maxmatp,maxmatp);
    Hintdown = spzeros(Float64,maxmatp,maxmatp);
    Hintup = spzeros(Float64,maxmatp,maxmatp);
    Hintdu = spzeros(Float64,maxmatp,maxmatp);

    # define a matrix for the Hamiltonian
    for nn = 1:maxmatp

        vecmbnn = in2b(nn,Msize,Np)

        # Interactions
        for ll = 1:Msize

            vecmbnnl = ades(ll,vecmbnn)
            if vecmbnnl[Msize+1] == 0
               continue
            end

            for kk = 1:Msize

                vecmbnnkl = ades(kk,vecmbnnl)
                if vecmbnnkl[Msize+1] == 0
                   continue
                end

                for jj = 1:Msize

                    vecmbnnjkl = acre(jj,vecmbnnkl)

                    for ii = 1:Msize

                        vecmbnnijkl = acre(ii,vecmbnnjkl)

                        # if isodd(ii+jj+kk+ll)
                        #    continue
                        # end

                        for mm = 1:nn

                            vecmbmm = in2b(mm,Msize,Np)

                            if vecmbnnijkl[1:Msize] == vecmbmm[1:Msize]

                               n1 = ceil(Int64,ii/2)-1
                               n2 = ceil(Int64,jj/2)-1
                               n3 = ceil(Int64,kk/2)-1
                               n4 = ceil(Int64,ll/2)-1

                               if isodd(ii) && isodd(jj) && isodd(kk) && isodd(ll)
                                  Hintdown[mm,nn] = Hintdown[mm,nn] + Vijkl(n1,n2,n3,n4)*sqrt(vecmbnnijkl[Msize+1])
                               elseif iseven(ii) && iseven(jj) && iseven(kk) && iseven(ll)
                                  Hintup[mm,nn] = Hintup[mm,nn] + Vijkl(n1,n2,n3,n4)*sqrt(vecmbnnijkl[Msize+1])
                               elseif isodd(ii) && iseven(jj) && isodd(kk) && iseven(ll)
                                  Hintdu[mm,nn] = Hintdu[mm,nn] + Vijkl(n1,n2,n3,n4)*sqrt(vecmbnnijkl[Msize+1])
                               elseif iseven(ii) && isodd(jj) && iseven(kk) && isodd(ll)
                                  Hintdu[mm,nn] = Hintdu[mm,nn] + Vijkl(n1,n2,n3,n4)*sqrt(vecmbnnijkl[Msize+1])
                               end

                            end

                        end

                    end

                end

            end

        end
    end

    Hintdown = Hintdown + Hintdown' - spdiagm(diag(Hintdown))
    Hintup = Hintup + Hintup' - spdiagm(diag(Hintup))
    Hintdu = Hintdu + Hintdu' - spdiagm(diag(Hintdu))

    return Hintdown, Hintup, Hintdu

end

ext: jl

max_stars_repo_name: leios/SOCED

max_issues_repo_name: leios/SOCED

max_forks_repo_name: leios/SOCED


text:
"""
    MTSampler{ST} <: Sampler where ST<:Sampler

A multithreaded wrapper for a sampler. Can be created by calling
`multithread(<:Sampler)`. This structure has two fields, a sampler, and a list
of thread-local samplers. It can be accessed as a standard sampler. Whenever it
is accessed as that, it will recompute the parameters of every thread-local sampler.
"""
mutable struct MTSampler{ST} <: Sampler
    sampler::ST
    samplers::Vector{ST}
end

"""
    multithread(sampler)

Creates a multithreaded version of the sampler, which will use `Threads.nthreads()`
independent processes when sampling.
"""
function multithread(sampler::Sampler)
    mts = MTSampler(sampler, [deepcopy(sampler) for i=1:Threads.nthreads()])
    recompute_params!(mts)
    mts
end
# dont multithread multithreaded code
multithread(sampl::MTSampler) = sampl

# Every time we get a property from the MTSampler, we forward the call to the
# original sampler.
@inline function Base.getproperty(s::MTSampler, val::Symbol)
    return getfield(getfield(s,:sampler), val)
end

function Base.setproperty!(s::MTSampler, sim::Symbol, val)
    Base.setfield!(getfield(s, :sampler), sim, val)
    _mt_recompute_sampler_params!(getfield(s, :samplers), getfield(s, :sampler))
end

Base.propertynames(x::MTSampler, private=false) =
    Base.propertynames(get_sampler(x), private)

recompute_params!(s::MTSampler) = _mt_recompute_sampler_params!(getfield(s, :samplers), getfield(s, :sampler))

get_sampler(s::MTSampler) = getfield(s, :sampler)
sampler_list(s::MTSampler) = getfield(s, :samplers)

function Base.show(io::IO, s::MTSampler)
    print(io, "MTSampler ($(length(sampler_list(s))) threads) of $(get_sampler(s))")
end

##
"""
    MTSamplerCache

A multithreading wrapper for SamplerCache, holding `num_threads` independent
copies of a `SamplerCache` and of the network, states.
"""
mutable struct MTSamplerCache{S,SC<:SamplerCache,V,N} <: SamplerCache{S}
    caches::Vector{SC}
    nets::Vector{N}
    σs::Vector{V}
end

MTSamplerCache(::S, caches, net::N, σ::V, op=copy) where {N,V,S<:Sampler} =
    MTSamplerCache{S,eltype(caches),V,N}(caches,
                                        [op(net) for i=1:length(caches)],
                                        [deepcopy(σ) for i=1:length(caches)])

# When initializing a MTCache the chain of function calls is the following:
# instead of calling directly _sampler_cache, _mt_sampler_cache is called, which
# is responsible for creating the multithreaded cache. You can override this, or
cache(s::MTSampler, v::State, net) =
    _mt_sampler_cache(s, v, net, ParallelThreaded())

# Mt sampler cache normally calls _sampler_cache for every thread. with the thread-local
# sampler. If this is fine enough, good, overwise you can override both _sampler_cache(.... ::ParallelThreaded)
# or the whole function.
function _mt_sampler_cache(s::MTSampler, v, net, ::ParallelThreaded)
    T = typeof(_sampler_cache(sampler_list(s)[1], v, net, ParallelThreaded()))
    scs = Vector{T}(undef, Threads.nthreads())

    Threads.@threads for i=1:Threads.nthreads()
        scs[i] = _sampler_cache(sampler_list(s)[i], v, net, ParallelThreaded(), i)
    end

    return MTSamplerCache(s, scs, net, v)
end

# Basic fallback: some samplers dont need any extra computation in case they are
# parallel, and therefore the cache creation has a signature with 4 arguments, not
# five. To fallback to use standard sampler cache if possible, we add this method.
_sampler_cache(s::Sampler, v, net, t::ParallelThreaded, thread_i) =
    _sampler_cache(s, v, net, t)

# The call tree for init_sampler is the following: if called with multithreading
# structures, first it updates the weight in the cached networks, then it calls
# `mt_init_sampler`, which can be overridden by the user. By default
# `mt_init_sampler` calls `mt_init_sampler_perthread` on every thread, which in
# turns default to calling `init_sampler!`.
# If a sampler must run some specific thread-local operations, it must override
# `*_perthread`.  If it must run more complicated logic involving syncronization
# among threads, then `mt_init_sampler` must be overridden. See respectively
# the examples of `FullSumSampler` and `ExactSampler`.
function init_sampler!(s::MTSampler, net, σ, c::MTSamplerCache)
    for cn=c.nets
        update!(Optimisers.Identity(), cn, net, nothing)
    end
    mt_init_sampler(s, net, σ, c)
end

"""
    mt_init_sampler(s::MTSampler, net, σ, c::MTSamplerCache)

Initializes a MTSamplerCache object.
"""
function mt_init_sampler(s::MTSampler, net, σ, c::MTSamplerCache)
    Threads.@threads for i=1:Threads.nthreads()
        mt_init_sampler_perthread(sampler_list(s)[i],
                                  c.nets[i],
                                  c.σs[i],
                                  c.caches[i])
    end
    return c
end

mt_init_sampler_perthread(args...) = init_sampler!(args...)


text:
module Abaqus_import_export_examples
using FinEtools
using FinEtools.MeshExportModule
using FinEtools.MeshImportModule
using FinEtools.MeshImportModule: import_ABAQUS

function Abaqus_import_export()
    ## Solid cylinder/taper/sphere—-temperature loading; quadratic brick mesh

    # The mesh  will be created in a very coarse representation from the
    # key points in the drawing. The first coordinate is radial, the second coordinate is axial.
    rz=[1.     0.;#A
    1.4    0.;#B
    0.995184726672197   0.098017140329561;
    1.393258617341076 0.137223996461385;
    0.980785  0.195090;#
    1.37309939 0.27312645;
    0.956940335732209   0.290284677254462
    1.339716470025092 0.406398548156247
    0.9238795  0.38268;#C
    1.2124  0.7;#D
    0.7071  0.7071;#E
    1.1062  1.045;#F
    0.7071  (0.7071+1.79)/2;#(E+H)/2
    1.      1.39;#G
    0.7071  1.79;#H
    1.      1.79;#I
    ]*phun("M")
    tolerance =1.e-6*phun("M")

    # This is the quadrilateral mesh of the cross-section.   It will be modified and
    # refined as  we go.
    fens = FENodeSet(rz);
    fes = FESetQ4([1 2 4 3; 3 4 6 5; 5 6 8 7; 7 8 10 9; 9 10 12 11; 11 12 14 13; 13 14 16 15]);

    nref = 1;
    for ref = 1:nref
        fens,fes = Q4refine(fens,fes);
        list = selectnode(fens, distance=1.0+0.1/2^nref, from=[0. 0.], inflate=tolerance);
        fens.xyz[list,:] = FinEtools.MeshUtilModule.ontosphere(fens.xyz[list,:],1.0);
    end

    ##
    # The mesh is extruded by sweeping around the axis of symmetry.
    # Only a single layer of elements is generated of internal angle
    # |angslice|.
    nLayers = 7;
    angslice  = 5*pi/16;

    ##
    # First the mesh is extruded to a block whose third dimension
    # represents the angular coordinate.
    fens,fes = H8extrudeQ4(fens, fes, nLayers, (rz,k)->[rz[1],rz[2],0.0]-(k)/nLayers*[0.,0.,angslice]);
    # The block is now converted  to the axially symmetric geometry by using the
    # third (angular) coordinate  to sweep out  an axially symmetric domain. The
    # ccoordinates of the nodes at this point are |rza|,  radial distance,
    # Z-coordinate, angle.
    sweep(rza) = [-rza[1]*sin(rza[3]+angslice/2.0), rza[1]*cos(rza[3]+angslice/2.0), rza[2]]
    for j=1:size(fens.xyz,1)
        fens.xyz[j,:] = sweep(fens.xyz[j,:])
    end

    AE = AbaqusExporter("LE11NAFEMS_H8");
    HEADING(AE, "LE11NAFEMS: Linear bricks.");
    PART(AE, "part1");
    END_PART(AE);
    ASSEMBLY(AE, "ASSEM1");
    INSTANCE(AE, "INSTNC1", "PART1");
    NODE(AE, fens.xyz);
    ELEMENT(AE, "c3d8rh", "AllElements", 1, connasarray(fes))
    END_INSTANCE(AE);
    END_ASSEMBLY(AE);
    close(AE)

    output = MeshImportModule.import_ABAQUS("./LE11NAFEMS_H8.inp")
    fens, fes = output["fens"], output["fesets"][1]

    File = "LE11NAFEMS_H8.vtk"
    MeshExportModule.vtkexportmesh(File, fens, fes)
    @async run(`"paraview.exe" $File`)
end # Abaqus_import_export

function Abaqus_import_export_LE10_T10()
    output = import_ABAQUS("$(@__DIR__)/" * "NLE10-coarse-T10.inp")
    fens, fes = output["fens"], output["fesets"][1]

    File = "LE10NAFEMS_T10.vtk"
    MeshExportModule.vtkexportmesh(File, fens, fes)
    @async run(`"paraview.exe" $File`)
end # Abaqus_import_export_LE10_T10

function allrun()
    println("#####################################################")
    println("# Abaqus_import_export ")
    Abaqus_import_export()
    println("#####################################################")
    println("# Abaqus_import_export_LE10_T10 ")
    Abaqus_import_export_LE10_T10()
    return true
end # function allrun

end # module Abaqus_import_export_examples


text:
# Code to be factored into a plotting package

using Plots, Test
using PlotsLH

## ------------  Setup

function make_matrix_xyz(nx, ny, nz)
    m = zeros(nx, ny, nz);
    for ix = 1 : nx
        m[ix,:,:] .+= 0.1 * ix;
    end
    for iy = 1 : ny
        m[:,iy,:] .+= 0.2 * iy;
    end
    for iz = 1 : nz
        m[:,:,iz] .+= 0.3 * iz;
    end
    return m
end

function make_test_labels(nx, ny, nz)
    xyzLabelV = [make_test_labels("x", nx), make_test_labels("y", ny), make_test_labels("z", nz)];
    return xyzLabelV
end

make_test_labels(lStr :: AbstractString, n :: Integer) =
    ["$lStr$j" for j = 1 : n];

## ---------------  Tests

function grouped_bar_xyz_test()
    @testset "Grouped bar xyz" begin
        nx = 5; ny = 4; nz = 3;
        m = make_matrix_xyz(nx, ny, nz);
        xyzLabelV = make_test_labels(nx, ny, nz);
        yStr = "y values";

        for subPlotIdx = 1 : 3
            for groupIdx = 1 : 3
                if groupIdx != subPlotIdx
                    p = grouped_bar_xyz(m, xyzLabelV;
                        yLabel = yStr, legendPos = :outerright,
                        groupIdx = groupIdx, subPlotIdx = subPlotIdx);
                    fName = "plot_by_xyz_$groupIdx$subPlotIdx";
                    delete_test_figure(fName);
                    save_test_figure(p,  fName);
                    @test isfile(test_fig_path(fName))
                end
            end
        end
    end
end

@testset "Plot (x,y,z)" begin
    grouped_bar_xyz_test()
end

# ------------


text:
@testset "Modularity" begin

    # 1. undirected test cases
    n = 10
    m = n * (n - 1) / 2
    c = ones(Int, n)
    gint = CompleteGraph(n)
    for g in testgraphs(gint)
      @test @inferred(modularity(g, c)) == 0
    end

    gint = SimpleGraph(n)
    for g in testgraphs(gint)
      @test @inferred(modularity(g, c)) == 0
    end

    barbell = blockdiag(CompleteGraph(3), CompleteGraph(3))
    add_edge!(barbell, 1, 4)
    c = [1, 1, 1, 2, 2, 2]

    for g in testgraphs(barbell)
        Q1 = @inferred(modularity(g, c))
        @test isapprox(Q1, 0.35714285714285715, atol=1e-3)
        Q2 = @inferred(modularity(g, c, γ=0.5))
        @test isapprox(Q2, 0.6071428571428571, atol=1e-3)
    end

    # 2. directed test cases
    triangle = SimpleDiGraph(3)
    add_edge!(triangle, 1, 2)
    add_edge!(triangle, 2, 3)
    add_edge!(triangle, 3, 1)

    barbell = blockdiag(triangle, triangle)
    add_edge!(barbell, 1, 4)
    c = [1, 1, 1, 2, 2, 2]

    for g in testdigraphs(barbell)
        Q1 = @inferred(modularity(g, c))
        @test isapprox(Q1, 0.3673469387755103, atol=1e-3)

        Q2 = @inferred(modularity(barbell, c, γ=0.5))
        @test isapprox(Q2, 0.6122448979591837, atol=1e-3)
    end

    add_edge!(barbell, 4, 1)
    for g in testdigraphs(barbell)
        @test @inferred(modularity(g, c)) == 0.25
    end

    # 3. weighted test cases
    # 3.1. undirected and weighted test cases
    triangle = SimpleGraph(3)
    add_edge!(triangle, 1, 2)
    add_edge!(triangle, 2, 3)
    add_edge!(triangle, 3, 1)

    barbell = blockdiag(triangle, triangle)
    add_edge!(barbell, 1, 4) # this edge has a weight of 5
    c = [1, 1, 1, 2, 2, 2]
    d = [[0  1  1  5  0  0]
         [1  0  1  0  0  0]
         [1  1  0  0  0  0]
         [5  0  0  0  1  1]
         [0  0  0  1  0  1]
         [0  0  0  1  1  0]]
    for g in testgraphs(barbell)
        Q = @inferred(modularity(g, c, distmx=d))
        @test isapprox(Q, 0.045454545454545456, atol=1e-3)
    end

    # 3.2. directed and weighted test cases
    triangle = SimpleDiGraph(3)
    add_edge!(triangle, 1, 2)
    add_edge!(triangle, 2, 3)
    add_edge!(triangle, 3, 1)

    barbell = blockdiag(triangle, triangle)
    add_edge!(barbell, 1, 4) # this edge has a weight of 5
    c = [1, 1, 1, 2, 2, 2]
    d = [[0  1  0  5  0  0]
         [0  0  1  0  0  0]
         [1  0  0  0  0  0]
         [0  0  0  0  1  0]
         [0  0  0  0  0  1]
         [0  0  0  1  0  0]]
    for g in testdigraphs(barbell)
        Q = @inferred(modularity(g, c, distmx=d))
        @test isapprox(Q, 0.1487603305785124, atol=1e-3)
    end
end


text:
function part1(input)
    state = parse_input(input, 3)
    for i = 1:6
        state = simulate_cycle(state)
    end
    return length(state)
end

function part2(input)
    state = parse_input(input, 4)
    for i = 1:6
        state = simulate_cycle(state)
    end
    return length(state)
end

function parse_input(input, dims)
    state = Set{CartesianIndex{dims}}()
    for (i, line) in enumerate(readlines(input))
        for (j, c) in enumerate(split(line, ""))
            if c == "#"
                push!(state, CartesianIndex(i, j, fill(0, dims - 2)...))
            end
        end
    end
    return state
end

function simulate_cycle(state::Set{CartesianIndex{N}}) where N
    neighbors = Dict{CartesianIndex{N}, Int}()
    neighborhood = CartesianIndices(ntuple(i -> -1:1, N))
    for position in state
        for Δ in neighborhood
            if Δ == zero(CartesianIndex{N})
                neighbors[position + Δ] = get(neighbors, position + Δ, 0) + 1
            else
                neighbors[position + Δ] = get(neighbors, position + Δ, 0) + 2
            end
        end
    end
    return Set(position
               for (position, num_neighbors) in neighbors
               if 5 <= num_neighbors <= 7)
end


text:
# Kept only for backwards compatibility
macro fix(ex) esc(ex) end


text:
__precompile__(false)

module GDAXClient

using Requests
using Nettle
using DandelionWebSockets
using JSON
using FIX
using MbedTLS
using Reexport

import Base: connect

abstract type AbstractGDAXMessageHandler <: FIX.AbstractMessageHandler end

export getProducts, getServerTime, getProductOrderBook, FIXQuote
export getProductTicker, getProductTrades, get24HourStats, getHistoricRates
export getCurrencies,  getPosition, getAccounts
export getAccountHistory, getAccountHolds
export placeOrder, cancelOrder, cancelAll, listOrders

export GDAXUser, GDAXWebSocketClient, GDAXOrderBooks
export fixconnect, wsconnect, subscribe, unsubscribe, onWSMessage, onFIXMessage
export AbstractGDAXMessageHandler

struct GDAXUser
    rest_url::String
    ws_feed::String
    key::String
    secret::String
    passphrase::String
end

isregistered(this::GDAXUser) = (!isempty(this.key) && !isempty(this.secret) && !isempty(this.passphrase))

include("sloworderbook.jl")
@reexport using .SlowOrderBook

include("rest.jl")
include("websocket.jl")
include("fix.jl")
include("historical_data.jl")

end

max_stars_repo_name: DatName/gdax-julia

max_issues_repo_name: DatName/gdax-julia

max_forks_repo_name: DatName/gdax-julia


text:
# #
# A proof of concept:
# 1) PlotSolitaryWaveWhithamBoussinesq() numerically computes the solitary wave of a Whitham-Boussinesq equation with prescribed velocity
# 2) IntegrateSolitaryWaveWhithamBoussinesq() integrates the equation in time
# #
export PlotSolitaryWaveWhithamBoussinesq,IntegrateSolitaryWaveWhithamBoussinesq
using WaterWaves1D,FFTW,Plots;

#----
"""
	PlotSolitaryWaveWhithamBoussinesq(;kwargs...)

Construct and plot a solitary wave of the Whitham-Boussinesq equation.

All keyword arguments are optional.
- `c` the velocity of the wave,
- `α` defines the model used,
- `L` the half-length of the mesh,
- `N` the number of collocation points,
- `μ` the shallowness parameter,
- `ϵ` the nonlinearity parameter.

"""
function PlotSolitaryWaveWhithamBoussinesq(;c=1.05,α=1,L=20,N=2^9,μ=0.1,ϵ=0.1)
	param = ( μ  = μ, ϵ = ϵ,
			N = N, L = L,
			c = c, α = α)

	(η,v,mesh) = SolitaryWaveWhithamBoussinesq( param;
										α = 1, model  = α, iterative = false, max_iter=30)

	ηSGN = (c^2-1)/ϵ*sech.(sqrt(3*(c^2-1)/(c^2)/μ)/2*(mesh.x)).^2
	uSGN = c*η./(1 .+ ϵ*η)
	k = mesh.k
	F₀ = sqrt(param.μ)*1im * k
	DxF(v) = real.(ifft(F₀ .* fft(v)))
	h = 1 .+ param.ϵ*ηSGN
	vSGN = uSGN - 1/3 ./h .* (DxF(h.^3 .*DxF(uSGN)))

	plt = plot(layout=(1,2))
	plot!(plt[1,1], mesh.x, [η v];
     title=string("initial data"),
     label=["η" "v"])
	plot!(plt[1,1], mesh.x, [ηSGN vSGN];
      title=string("initial data"),
      label=["ηSGN" "vSGN"])

    plot!(plt[1,2], fftshift(mesh.k),
     [abs.(fftshift(fft(η))).+eps() abs.(fftshift(fft(v))).+eps()];
     yscale=:log10,
	 title="frequency",
	 label=["|η̂|" "|v̂|"])
	plot!(plt[1,2], fftshift(mesh.k),
      [abs.(fftshift(fft(ηSGN))).+eps() abs.(fftshift(fft(vSGN))).+eps()];
	  yscale=:log10,
	  title="frequency",
 	 label=["|η̂SGN|" "|v̂SGN|"])

	 display(plt)
end

"""
	IntegrateSolitaryWaveWhithamBoussinesq()

Integrate in time a Whitham-Boussinesq equation with a solitary wave initial data
"""
function IntegrateSolitaryWaveWhithamBoussinesq()
	param = ( μ  = 1,
			ϵ  = 1,
			c = 1.05,
        	N  = 2^9,
            L  = 20,
			T  = 40/1.05,
            dt = 0.001/1.05,
			ns = 200
			)
	α = 1/2 # determines the model

	(η,v,mesh) = SolitaryWaveWhithamBoussinesq( param;
										α = 1, model  = α, iterative = false, max_iter=30)

	init=Init(mesh,η,v)

	plt = plot(layout=(1,2))
	plot!(plt[1,1], mesh.x, [η v];
     title=string("initial data"),
     label=["η" "u"])

    plot!(plt[1,2], fftshift(mesh.k),
     [log10.(abs.(fftshift(fft(η)))) log10.(abs.(fftshift(fft(v))))];
     title="frequency",
	 label=["η" "u"])
	 display(plt)

	model  = WhithamBoussinesq(param;α=α)
	problem = Problem(model, init, param);

	solve!( problem )
	ηf = model.mapfro(problem.data.U[end])[1]
	p = plot(layout=(1,2))
	plot!(p[1,1],mesh.x,[η ηf];
	title="surface deformation",
	label=["initial" "final"])
	plot!(p[1,2],mesh.x,η-ηf;
	title="difference",
	label="")
	display(p)

	anim = create_animation( problem )
	gif(anim, fps=15)
end
nothing


text:
## Generic ##

Tracker.dual(x::Bool, p) = x
Tracker.dual(x::Int, p) = x

Base.prevfloat(r::TrackedReal) = track(prevfloat, r)
@grad function prevfloat(r::Real)
    prevfloat(data(r)), Δ -> Δ
end
Base.nextfloat(r::TrackedReal) = track(nextfloat, r)
@grad function nextfloat(r::Real)
    nextfloat(data(r)), Δ -> Δ
end

for f = [:hcat, :vcat]
    for c = [
        [:TrackedReal],
        [:AbstractVecOrMat, :TrackedReal],
        [:TrackedVecOrMat, :TrackedReal],
    ]
        cnames = map(_ -> gensym(), c)
        @eval begin
            function Base.$f(
                $([:($x::$c) for (x, c) in zip(cnames, c)]...),
                x::Union{TrackedArray,TrackedReal},
                xs::Union{AbstractArray,Number}...,
            )
                return track($f, $(cnames...), x, xs...)
            end
        end
    end
    @eval begin
        @grad function $f(x::Real)
            $f(data(x)), (Δ) -> (Δ[1],)
        end
        @grad function $f(x1::Real, x2::Real)
            $f(data(x1), data(x2)), (Δ) -> (Δ[1], Δ[2])
        end
        @grad function $f(x1::AbstractVector, x2::Real)
            $f(data(x1), data(x2)), (Δ) -> (Δ[1:length(x1)], Δ[length(x1)+1])
        end
    end
end

function Base.copy(
    A::TrackedArray{T, 2, <:Adjoint{T, <:AbstractTriangular{T, <:AbstractMatrix{T}}}},
) where {T <: Real}
    return track(copy, A)
end
@grad function Base.copy(
    A::TrackedArray{T, 2, <:Adjoint{T, <:AbstractTriangular{T, <:AbstractMatrix{T}}}},
) where {T <: Real}
    return copy(data(A)), ∇ -> (copy(∇),)
end

Base.:*(A::TrackedMatrix, B::AbstractTriangular) = track(*, A, B)
Base.:*(A::AbstractTriangular{T}, B::TrackedVector) where {T} = track(*, A, B)
Base.:*(A::AbstractTriangular{T}, B::TrackedMatrix) where {T} = track(*, A, B)
Base.:*(A::Adjoint{T, <:AbstractTriangular{T}}, B::TrackedMatrix) where {T} = track(*, A, B)
Base.:*(A::Adjoint{T, <:AbstractTriangular{T}}, B::TrackedVector) where {T} = track(*, A, B)

function Base.fill(
    value::TrackedReal,
    dims::Vararg{Union{Integer, AbstractUnitRange}},
)
    return track(fill, value, dims...)
end
@grad function Base.fill(value::Real, dims...)
    return fill(data(value), dims...), function(Δ)
        size(Δ) ≢  dims && error("Dimension mismatch")
        return (sum(Δ), map(_->nothing, dims)...)
    end
end

## StatsFuns ##

logsumexp(x::TrackedArray; dims=:) = _logsumexp(x, dims)
_logsumexp(x::TrackedArray, dims=:) = track(_logsumexp, x, dims)
@grad function _logsumexp(x::TrackedArray, dims)
    lse = logsumexp(data(x), dims = dims)
    return lse, Δ -> (Δ .* exp.(x .- lse), nothing)
end

## Linear algebra ##

# Work around https://github.com/FluxML/Tracker.jl/pull/9#issuecomment-480051767
upper(A::AbstractMatrix) = UpperTriangular(A)
lower(A::AbstractMatrix) = LowerTriangular(A)
function upper(C::Cholesky)
    if C.uplo == 'U'
        return upper(C.factors)
    else
        return copy(lower(C.factors)')
    end
end
function lower(C::Cholesky)
    if C.uplo == 'U'
        return copy(upper(C.factors)')
    else
        return lower(C.factors)
    end
end

LinearAlgebra.LowerTriangular(A::TrackedMatrix) = lower(A)
lower(A::TrackedMatrix) = track(lower, A)
@grad lower(A) = lower(Tracker.data(A)), ∇ -> (lower(∇),)

LinearAlgebra.UpperTriangular(A::TrackedMatrix) = upper(A)
upper(A::TrackedMatrix) = track(upper, A)
@grad upper(A) = upper(Tracker.data(A)), ∇ -> (upper(∇),)

function LinearAlgebra.cholesky(A::TrackedMatrix; check=true)
    factors_info = turing_chol(A, check)
    factors = factors_info[1]
    info = data(factors_info[2])
    return Cholesky{eltype(factors), typeof(factors)}(factors, 'U', info)
end
function LinearAlgebra.cholesky(A::Symmetric{<:Any, <:TrackedMatrix}; check=true)
    uplo = A.uplo == 'U' ? (:U) : (:L)
    factors_info = symm_turing_chol(parent(A), check, uplo)
    factors = factors_info[1]
    info = data(factors_info[2])
    return Cholesky{eltype(factors), typeof(factors)}(factors, 'U', info)
end

turing_chol(A::TrackedMatrix, check) = track(turing_chol, A, check)
@grad function turing_chol(A::AbstractMatrix, check)
    Y, back = DistributionsAD.turing_chol_back(data(A),check)
    Y, Δ->back(data.(Δ))
end

symm_turing_chol(A::TrackedMatrix, check, uplo) = track(symm_turing_chol, A, check, uplo)
@grad function symm_turing_chol(A::AbstractMatrix, check, uplo)
    Y, back = DistributionsAD.symm_turing_chol_back(data(A),check,uplo)
    Y, Δ->back(data.(Δ))
end

# Specialised logdet for cholesky to target the triangle directly.
logdet_chol_tri(U::AbstractMatrix) = 2 * sum(log, U[diagind(U)])
logdet_chol_tri(U::TrackedMatrix) = track(logdet_chol_tri, U)
@grad function logdet_chol_tri(U::AbstractMatrix)
    U_data = data(U)
    return logdet_chol_tri(U_data), Δ->(Matrix(Diagonal(2 .* Δ ./ diag(U_data))),)
end

function LinearAlgebra.logdet(C::Cholesky{<:TrackedReal, <:TrackedMatrix})
    return logdet_chol_tri(C.U)
end

function zygote_ldiv(A::TrackedMatrix, B::TrackedVecOrMat)
    return track(zygote_ldiv, A, B)
end
function zygote_ldiv(A::TrackedMatrix, B::AbstractVecOrMat)
    return track(zygote_ldiv, A, B)
end
zygote_ldiv(A::AbstractMatrix, B::TrackedVecOrMat) =  track(zygote_ldiv, A, B)
@grad function zygote_ldiv(A, B)
    Y, dY_pullback = rrule(\, data(A), data(B))
    function back(Δ)
        _, dA, dB = dY_pullback(data(Δ))
        (unthunk(dA), unthunk(dB))
    end
    Y, back
end

function Base.:\(a::Cholesky{<:TrackedReal, <:TrackedArray}, b::AbstractVecOrMat)
    return (a.U \ (a.U' \ b))
end

## SpecialFunctions ##

SpecialFunctions.logabsgamma(x::TrackedReal) = track(logabsgamma, x)
@grad function SpecialFunctions.logabsgamma(x::Real)
    return logabsgamma(data(x)), Δ -> (digamma(data(x)) * Δ[1],)
end

# isprobvec
function Distributions.isprobvec(p::TrackedArray{<:Real})
    pdata = Tracker.data(p)
    all(x -> x ≥ zero(x), pdata) && isapprox(sum(pdata), one(eltype(pdata)), atol = 1e-6)
end

# Some array functions - workaround https://github.com/FluxML/Tracker.jl/issues/4
import Base: +, -, *, /, \
import LinearAlgebra: dot

for f in (:+, :-, :*, :/, :\, :dot), (T1, T2) in [
    (:TrackedArray, :AbstractArray),
    (:TrackedMatrix, :AbstractMatrix),
    (:TrackedMatrix, :AbstractVector),
    (:TrackedVector, :AbstractMatrix),
]
    @eval begin
        function $f(a::$T1{T}, b::$T2{<:TrackedReal}) where {T <: Real}
            return $f(convert(AbstractArray{TrackedReal{T}}, a), b)
        end
        function $f(a::$T2{<:TrackedReal}, b::$T1{T}) where {T <: Real}
            return $f(a, convert(AbstractArray{TrackedReal{T}}, b))
        end
    end
end

## Uniform ##

Distributions.Uniform(a::TrackedReal, b::Real) = TuringUniform{TrackedReal}(a, b)
Distributions.Uniform(a::Real, b::TrackedReal) = TuringUniform{TrackedReal}(a, b)
Distributions.Uniform(a::TrackedReal, b::TrackedReal) = TuringUniform{TrackedReal}(a, b)
Distributions.logpdf(d::Uniform, x::TrackedReal) = uniformlogpdf(d.a, d.b, x)
Distributions.logpdf(d::Uniform, x::TrackedArray) = uniformlogpdf.(d.a, d.b, x)

uniformlogpdf(a::Real, b::Real, x::TrackedReal) = track(uniformlogpdf, a, b, x)
uniformlogpdf(a::TrackedReal, b::TrackedReal, x::Real) = track(uniformlogpdf, a, b, x)
uniformlogpdf(a::TrackedReal, b::TrackedReal, x::TrackedReal) = track(uniformlogpdf, a, b, x)
@grad function uniformlogpdf(a, b, x)
    diff = data(b) - data(a)
    T = typeof(diff)
    if a <= data(x) <= b && a < b
        l = -log(diff)
        da = 1/diff^2
        return l, Δ -> (da * Δ, -da * Δ, zero(T) * Δ)
    else
        n = T(NaN)
        return n, Δ -> (n, n, n)
    end
end

## Binomial ##

binomlogpdf(n::Int, p::TrackedReal, x::Int) = track(binomlogpdf, n, p, x)
@grad function binomlogpdf(n::Int, p::TrackedReal, x::Int)
    return binomlogpdf(n, data(p), x),
        Δ->(nothing, Δ * (x / p - (n - x) / (1 - p)), nothing)
end

## Poisson ##

poislogpdf(v::TrackedReal, x::Int) = track(poislogpdf, v, x)
@grad function poislogpdf(v::TrackedReal, x::Int)
      return poislogpdf(data(v), x),
          Δ->(Δ * (x/v - 1), nothing)
end

## PoissonBinomial ##

PoissonBinomial(p::TrackedArray{<:Real}; check_args=true) =
    TuringPoissonBinomial(p; check_args = check_args)
poissonbinomial_pdf_fft(x::TrackedArray) = track(poissonbinomial_pdf_fft, x)
@grad function poissonbinomial_pdf_fft(x::TrackedArray)
    x_data = data(x)
    T = eltype(x_data)
    fft = poissonbinomial_pdf_fft(x_data)
    return  fft, Δ -> begin
        ((ForwardDiff.jacobian(x -> poissonbinomial_pdf_fft(x), x_data)::Matrix{T})' * Δ,)
    end
end

## Semicircle ##

function semicircle_dldr(r, x)
    diffsq = r^2 - x^2
    return -2 / r + r / diffsq
end
function semicircle_dldx(r, x)
    diffsq = r^2 - x^2
    return -x / diffsq
end

logpdf(d::Semicircle{<:Real}, x::TrackedReal) = semicirclelogpdf(d.r, x)
logpdf(d::Semicircle{<:TrackedReal}, x::Real) = semicirclelogpdf(d.r, x)
logpdf(d::Semicircle{<:TrackedReal}, x::TrackedReal) = semicirclelogpdf(d.r, x)

semicirclelogpdf(r, x) = logpdf(Semicircle(r), x)
M, f, arity = DiffRules.@define_diffrule DistributionsAD.semicirclelogpdf(r, x) =
    :(semicircle_dldr($r, $x)), :(semicircle_dldx($r, $x))
da, db = DiffRules.diffrule(M, f, :a, :b)
f = :($M.$f)
@eval begin
    @grad $f(a::TrackedReal, b::TrackedReal) = $f(data(a), data(b)), Δ -> (Δ * $da, Δ * $db)
    @grad $f(a::TrackedReal, b::Real) = $f(data(a), b), Δ -> (Δ * $da, Tracker._zero(b))
    @grad $f(a::Real, b::TrackedReal) = $f(a, data(b)), Δ -> (Tracker._zero(a), Δ * $db)
    $f(a::TrackedReal, b::TrackedReal)  = track($f, a, b)
    $f(a::TrackedReal, b::Real) = track($f, a, b)
    $f(a::Real, b::TrackedReal) = track($f, a, b)
end

## Negative binomial ##

# Note the definition of NegativeBinomial in Julia is not the same as Wikipedia's.
# Check the docstring of NegativeBinomial, r is the number of successes and
# k is the number of failures
_nbinomlogpdf_grad_1(r, p, k) = k == 0 ? log(p) : sum(1 / (k + r - i) for i in 1:k) + log(p)
_nbinomlogpdf_grad_2(r, p, k) = -k / (1 - p) + r / p

nbinomlogpdf(n::TrackedReal, p::TrackedReal, x::Int) = track(nbinomlogpdf, n, p, x)
nbinomlogpdf(n::Real, p::TrackedReal, x::Int) = track(nbinomlogpdf, n, p, x)
nbinomlogpdf(n::TrackedReal, p::Real, x::Int) = track(nbinomlogpdf, n, p, x)
@grad function nbinomlogpdf(r::TrackedReal, p::TrackedReal, k::Int)
    return nbinomlogpdf(data(r), data(p), k),
        Δ->(Δ * _nbinomlogpdf_grad_1(r, p, k), Δ * _nbinomlogpdf_grad_2(r, p, k), nothing)
end
@grad function nbinomlogpdf(r::Real, p::TrackedReal, k::Int)
    return nbinomlogpdf(data(r), data(p), k),
        Δ->(Tracker._zero(r), Δ * _nbinomlogpdf_grad_2(r, p, k), nothing)
end
@grad function nbinomlogpdf(r::TrackedReal, p::Real, k::Int)
    return nbinomlogpdf(data(r), data(p), k),
        Δ->(Δ * _nbinomlogpdf_grad_1(r, p, k), Tracker._zero(p), nothing)
end

## Categorical ##

function Distributions.DiscreteNonParametric{T,P,Ts,Ps}(
    vs::Ts,
    ps::Ps;
    check_args=true,
) where {T<:Real,P<:Real,Ts<:AbstractVector{T},Ps<:TrackedArray{P, 1, <:SubArray{P, 1}}}
    cps = ps[:]
    return DiscreteNonParametric{T,P,Ts,typeof(cps)}(vs, cps; check_args = check_args)
end

## Dirichlet ##

Distributions.Dirichlet(alpha::TrackedVector) = TuringDirichlet(alpha)
Distributions.Dirichlet(d::Integer, alpha::TrackedReal) = TuringDirichlet(d, alpha)

function Distributions.logpdf(d::Dirichlet{T}, x::TrackedVecOrMat) where {T}
    TV = typeof(d.alpha)
    logpdf(TuringDirichlet{T, TV}(d.alpha, d.alpha0, d.lmnB), x)
end

for T in (:TrackedVector, :TrackedMatrix)
    @eval begin
        function Distributions.logpdf(d::MvNormal{<:Any, <:PDMats.ScalMat}, x::$T)
            logpdf(TuringScalMvNormal(d.μ, sqrt(d.Σ.value)), x)
        end
        function Distributions.logpdf(d::MvNormal{<:Any, <:PDMats.PDiagMat}, x::$T)
            logpdf(TuringDiagMvNormal(d.μ, sqrt.(d.Σ.diag)), x)
        end
        function Distributions.logpdf(d::MvNormal{<:Any, <:PDMats.PDMat}, x::$T)
            logpdf(TuringDenseMvNormal(d.μ, d.Σ.chol), x)
        end

        function Distributions.logpdf(d::MvLogNormal{<:Any, <:PDMats.ScalMat}, x::$T)
            logpdf(TuringMvLogNormal(TuringScalMvNormal(d.normal.μ, sqrt(d.normal.Σ.value))), x)
        end
        function Distributions.logpdf(d::MvLogNormal{<:Any, <:PDMats.PDiagMat}, x::$T)
            logpdf(TuringMvLogNormal(TuringDiagMvNormal(d.normal.μ, sqrt.(d.normal.Σ.diag))), x)
        end
        function Distributions.logpdf(d::MvLogNormal{<:Any, <:PDMats.PDMat}, x::$T)
            logpdf(TuringMvLogNormal(TuringDenseMvNormal(d.normal.μ, d.normal.Σ.chol)), x)
        end
    end
end

# zero mean, dense covariance
MvNormal(A::TrackedMatrix) = TuringMvNormal(A)

# zero mean, diagonal covariance
MvNormal(σ::TrackedVector) = TuringMvNormal(σ)

# dense mean, dense covariance
MvNormal(m::TrackedVector{<:Real}, A::TrackedMatrix{<:Real}) = TuringMvNormal(m, A)
MvNormal(m::TrackedVector{<:Real}, A::Matrix{<:Real}) = TuringMvNormal(m, A)
MvNormal(m::AbstractVector{<:Real}, A::TrackedMatrix{<:Real}) = TuringMvNormal(m, A)

# dense mean, diagonal covariance
function MvNormal(
    m::TrackedVector{<:Real},
    D::Diagonal{T, <:TrackedVector{T}} where {T<:Real},
)
    return TuringMvNormal(m, D)
end
function MvNormal(
    m::AbstractVector{<:Real},
    D::Diagonal{T, <:TrackedVector{T}} where {T<:Real},
)
    return TuringMvNormal(m, D)
end
function MvNormal(
    m::TrackedVector{<:Real},
    D::Diagonal{T, <:AbstractVector{T}} where {T<:Real},
)
    return TuringMvNormal(m, D)
end

# dense mean, diagonal covariance
MvNormal(m::TrackedVector{<:Real}, σ::TrackedVector{<:Real}) = TuringMvNormal(m, σ)
MvNormal(m::TrackedVector{<:Real}, σ::AbstractVector{<:Real}) = TuringMvNormal(m, σ)
MvNormal(m::TrackedVector{<:Real}, σ::Vector{<:Real}) = TuringMvNormal(m, σ)
MvNormal(m::AbstractVector{<:Real}, σ::TrackedVector{<:Real}) = TuringMvNormal(m, σ)

# dense mean, constant variance
MvNormal(m::TrackedVector{<:Real}, σ::TrackedReal) = TuringMvNormal(m, σ)
MvNormal(m::TrackedVector{<:Real}, σ::Real) = TuringMvNormal(m, σ)
MvNormal(m::AbstractVector{<:Real}, σ::TrackedReal) = TuringMvNormal(m, σ)

# dense mean, constant variance
function MvNormal(m::TrackedVector{<:Real}, A::UniformScaling{<:TrackedReal})
    return TuringMvNormal(m, A)
end
function MvNormal(m::AbstractVector{<:Real}, A::UniformScaling{<:TrackedReal})
    return TuringMvNormal(m, A)
end
function MvNormal(m::TrackedVector{<:Real}, A::UniformScaling{<:Real})
    return TuringMvNormal(m, A)
end

# zero mean,, constant variance
MvNormal(d::Int, σ::TrackedReal{<:Real}) = TuringMvNormal(d, σ)

for T in (:TrackedVector, :TrackedMatrix)
    @eval Distributions.logpdf(d::TuringMvLogNormal, x::$T) = _logpdf(d, x)
end

# zero mean, dense covariance
MvLogNormal(A::TrackedMatrix) = TuringMvLogNormal(TuringMvNormal(A))

# zero mean, diagonal covariance
MvLogNormal(σ::TrackedVector) = TuringMvLogNormal(TuringMvNormal(σ))

# dense mean, dense covariance
MvLogNormal(m::TrackedVector{<:Real}, A::TrackedMatrix{<:Real}) = TuringMvLogNormal(TuringMvNormal(m, A))
MvLogNormal(m::TrackedVector{<:Real}, A::Matrix{<:Real}) = TuringMvLogNormal(TuringMvNormal(m, A))
MvLogNormal(m::AbstractVector{<:Real}, A::TrackedMatrix{<:Real}) = TuringMvLogNormal(TuringMvNormal(m, A))

## MvLogNormal ##

# dense mean, diagonal covariance
function MvLogNormal(
    m::TrackedVector{<:Real},
    D::Diagonal{T, <:TrackedVector{T}} where {T<:Real},
)
    return TuringMvLogNormal(TuringMvNormal(m, D))
end
function MvLogNormal(
    m::AbstractVector{<:Real},
    D::Diagonal{T, <:TrackedVector{T}} where {T<:Real},
)
    return TuringMvLogNormal(TuringMvNormal(m, D))
end
function MvLogNormal(
    m::TrackedVector{<:Real},
    D::Diagonal{T, <:AbstractVector{T}} where {T<:Real},
)
    return TuringMvLogNormal(TuringMvNormal(m, D))
end

# dense mean, diagonal covariance
MvLogNormal(m::TrackedVector{<:Real}, σ::TrackedVector{<:Real}) = TuringMvLogNormal(TuringMvNormal(m, σ))
MvLogNormal(m::TrackedVector{<:Real}, σ::AbstractVector{<:Real}) = TuringMvLogNormal(TuringMvNormal(m, σ))
MvLogNormal(m::TrackedVector{<:Real}, σ::Vector{<:Real}) = TuringMvLogNormal(TuringMvNormal(m, σ))
MvLogNormal(m::AbstractVector{<:Real}, σ::TrackedVector{<:Real}) = TuringMvLogNormal(TuringMvNormal(m, σ))

# dense mean, constant variance
function MvLogNormal(m::TrackedVector{<:Real}, σ::TrackedReal)
    return TuringMvLogNormal(TuringMvNormal(m, σ))
end
function MvLogNormal(m::TrackedVector{<:Real}, σ::Real)
    return TuringMvLogNormal(TuringMvNormal(m, σ))
end
function MvLogNormal(m::AbstractVector{<:Real}, σ::TrackedReal)
    return TuringMvLogNormal(TuringMvNormal(m, σ))
end

# dense mean, constant variance
function MvLogNormal(m::TrackedVector{<:Real}, A::UniformScaling{<:TrackedReal})
    return TuringMvLogNormal(TuringMvNormal(m, A))
end
function MvLogNormal(m::AbstractVector{<:Real}, A::UniformScaling{<:TrackedReal})
    return TuringMvLogNormal(TuringMvNormal(m, A))
end
function MvLogNormal(m::TrackedVector{<:Real}, A::UniformScaling{<:Real})
    return TuringMvLogNormal(TuringMvNormal(m, A))
end

# zero mean,, constant variance
MvLogNormal(d::Int, σ::TrackedReal{<:Real}) = TuringMvLogNormal(TuringMvNormal(d, σ))

## MvCategorical ##

_mv_categorical_logpdf(ps::Tracker.TrackedMatrix, x) = Tracker.track(_mv_categorical_logpdf, ps, x)
Tracker.@grad function _mv_categorical_logpdf(ps, x)
    ps_data = Tracker.data(ps)
    probs = view(ps_data, x, :)
    ps_grad = zero(ps_data)
    sum(log, probs), Δ -> begin
        ps_grad .= 0
        ps_grad[x,:] .= Δ ./ probs
        return (ps_grad, nothing)
    end
end

## MatrixBeta ##

function Distributions.logpdf(d::MatrixBeta, X::AbstractArray{<:TrackedMatrix{<:Real}})
    return map(x -> logpdf(d, x), X)
end

## Wishart ##

function Distributions.logpdf(d::Wishart, X::TrackedMatrix)
    return logpdf(TuringWishart(d), X)
end
function Distributions.logpdf(d::Wishart, X::AbstractArray{<:TrackedMatrix})
    return logpdf(TuringWishart(d), X)
end
Distributions.Wishart(df::TrackedReal, S::Matrix{<:Real}) = TuringWishart(df, S)
Distributions.Wishart(df::TrackedReal, S::AbstractMatrix{<:Real}) = TuringWishart(df, S)
Distributions.Wishart(df::Real, S::TrackedMatrix) = TuringWishart(df, S)
Distributions.Wishart(df::TrackedReal, S::TrackedMatrix) = TuringWishart(df, S)

## Inverse Wishart ##

function Distributions.logpdf(d::InverseWishart, X::TrackedMatrix)
    return logpdf(TuringInverseWishart(d), X)
end
function Distributions.logpdf(d::InverseWishart, X::AbstractArray{<:TrackedMatrix})
    return logpdf(TuringInverseWishart(d), X)
end

Distributions.InverseWishart(df::TrackedReal, S::Matrix{<:Real}) = TuringInverseWishart(df, S)
Distributions.InverseWishart(df::TrackedReal, S::AbstractMatrix{<:Real}) = TuringInverseWishart(df, S)
Distributions.InverseWishart(df::Real, S::TrackedMatrix) = TuringInverseWishart(df, S)
Distributions.InverseWishart(df::TrackedReal, S::TrackedMatrix) = TuringInverseWishart(df, S)