text:
let
    S = CBSEnv.State
    A = CBSEnv.Action
    vtx_grid = initialize_regular_vtx_grid(;n_obstacles_x=1,n_obstacles_y=1)
    #  1   2   3   4   5   6
    #  7   8   9  10  11  12
    # 13  14   0   0  15  16
    # 17  18   0   0  19  20
    # 21  22  23  24  25  26
    # 27  28  29  30  31  32
    G = initialize_grid_graph_from_vtx_grid(vtx_grid)
    starts = [CBSEnv.State(1,0),CBSEnv.State(2,0)]
    goals = [CBSEnv.State(vtx=6),CBSEnv.State(vtx=5)]
    cost_model = SumOfTravelTime()
    env = CBSEnv.LowLevelEnv(graph=G,cost_model=cost_model)
    mapf = MAPF(env,starts,goals)

    get_infeasible_solution(mapf)

    solution = get_initial_solution(mapf)
    path = get_paths(solution)[1]
    for (t,path) in zip([4,5],get_paths(solution))
        extend_path!(env,path,t)
        @test length(path) == t
        @test get_cost(path) == length(path)
        trim_path!(env,path,t-1)
        @test length(path) == t-1
        @test get_cost(path) == length(path)
    end
    T = 2
    trim_solution!(env,solution,T)
    for path in get_paths(solution)
        @test length(path) == T
        @test get_cost(path) == T
    end
    @test get_cost(solution) == T*num_agents(mapf)

end
let
    mapf = init_mapf_1()
    env = mapf.env
    s0 = CBSEnv.State(1,0)
    path = Path{CBSEnv.State,CBSEnv.Action,Float64}(s0 = s0)
    for a in [
        CBSEnv.Action(Edge(1,2),1),
        CBSEnv.Action(Edge(2,3),1),
        CBSEnv.Action(Edge(3,4),1),
        ]
        s = get_final_state(path)
        sp = get_next_state(env,s,a)
        add_to_path!(path,env,s,a,sp)
    end
    @test get_initial_state(path) == s0
    p = deepcopy(path)
    trim_path!(env,path,1)
    @test get_initial_state(path) == s0
    @test get_final_state(path) == get_sp(get_path_node(p,1))
    extend_path!(env,path,3)
    @test get_vtx(get_final_state(path)) == get_vtx(get_sp(get_path_node(p,1)))
end


text:
## HC

function hc_score(pValues)
    validPvalues(pValues)
    n = length(pValues)
    f = competerank(pValues) / n ## ties = max
    var = f .* (1-f) / n
    var[var == 0] = min( var[var > 0] ) # just to make sure we have no zero variance
    hc = abs(f-pValues) / sqrt(var)
    hc
end

ext: jl


text:
include("map_transform.jl")

using NCDatasets
using .MapTransform

using ArgParse
using JSON

function parse_commandline()

    s = ArgParseSettings()
    @add_arg_table s begin

        "--input-file"
            help = "Input Qflux file."
            arg_type = String
            required = true

        "--output-file"
            help = "Output Qflux file."
            arg_type = String
            required = true

        "--domain-file"
            help = "Domain file."
            arg_type = String
            required = true

        "--varname-qflx"
            help = "Varname of Q-flux."
            arg_type = String
            required = true

    end

    return parse_args(ARGS, s)
end

println("Running ", @__FILE__)

parsed = parse_commandline()
print(json(parsed, 4))

Dataset(parsed["domain-file"], "r") do ds
    global area = replace(ds["area"][:], missing => NaN)
    global mask = replace(ds["mask"][:], missing => NaN)
    global lat  = replace(ds["yc"][:], missing => NaN)
end

lat_bnd = collect(Float64, -90:1:90)

println("Before: sum_area = ", sum(area))
area *= 4π * (6371e3)^2 / sum(area)
println("After sum_area = ", sum(area))

r = MapTransform.Relation(
    lat = lat,
    area = area,
    mask = mask,
    lat_bnd = lat_bnd,
)

_proxy = area * 0 .+ 1.0
sum_valid_area = MapTransform.∫∂a(r, _proxy)[end]

println("Sum of valid area: ", sum_valid_area, "; ratio: ", sum_valid_area / sum(area))

Dataset(parsed["input-file"], "r") do ids

    Nt = ids.dim["time"]
    global Q_ano  = zeros(Float64, length(lat_bnd)-1, Nt)
    global ∫Q_ano∂a = zeros(Float64, length(lat_bnd), Nt)
    global Q_avg = zeros(Float64, Nt)

    for t=1:Nt

        _data = nomissing(ids[parsed["varname-qflx"]][:, :, t], NaN)
        Q_avg[t] = MapTransform.∫∂a(r, _data)[end] / sum_valid_area

        _data .-= Q_avg[t]

        Q_ano[:, t] = MapTransform.transform(r, _data)
        ∫Q_ano∂a[:, t] = MapTransform.∫∂a(r, _data)

    end

end

Dataset(parsed["output-file"], "c") do ds

    defDim(ds, "lat_bnd", length(r.lat_bnd))
    defDim(ds, "lat",     length(r.lat_bnd)-1)
    defDim(ds, "time", Inf)

    for (varname, vardata, vardim, attrib) in [
        ("lat_bnd",   lat_bnd, ("lat_bnd",         ), Dict()),
        ("Q_avg",       Q_avg, ("time", ), Dict()),
        ("Q_ano",       Q_ano, ("lat",     "time", ), Dict()),
        ("intQ_ano", ∫Q_ano∂a, ("lat_bnd", "time", ), Dict()),
    ]
        println("Doing varname:", varname)
        var = defVar(ds, varname, Float64, vardim)
        var.attrib["_FillValue"] = 1e20

        for (k, v) in attrib
            var.attrib[k] = v
        end

        rng = []
        for i in 1:length(vardim)-1
            push!(rng, Colon())
        end
        push!(rng, 1:size(vardata)[end])
        var[rng...] = vardata

    end

end

max_stars_repo_name: meteorologytoday/CESM-diagnostic

max_issues_repo_name: meteorologytoday/CESM-diagnostic

max_forks_repo_name: meteorologytoday/CESM-diagnostic


text:
using StatsBase

function sortear(n,k)
  return sample(1:n, k, replace=false, ordered=true);
end

include("jogos/megasena.jl");
include("jogos/lotofacil.jl")
include("jogos/lotomania.jl")

max_stars_repo_name: fabiom/loteria

max_issues_repo_name: fabiom/loteria

max_forks_repo_name: fabiom/loteria


text:
export Finel

"""
Struct defining a finite element

* `pts`: corner points (vector of vectors)
* `domain`: `Domain` struct with the finite element's generic definition
* `abs_coord_regressors`: dictionary mapping `InterpolationFunction` structs to
    regressors
"""
mutable struct Finel
    pts::Vector{
        Vector{Float64}
    }
    domain::Domain
    abs_coord_regressors::Dict{
        InterpolationFunction, Regressor
    }
    _coord_linpars::Vector{Vector{Float64}}
end

"""
Constructor for finite element

* `pts`: corner points (vector of vectors)
* `domain`: `Domain` struct with the finite element's generic definition
"""
function Finel(pts::Vector{Vector{Float64}}, domain::Domain)
    pts_transpose=Vector{Vector{Float64}}(undef, length(pts[1]))
    for i=1:length(pts_transpose)
        pts_transpose[i]=[
            pts[j][i] for j=1:length(pts)
        ]
    end

    Finel(
        pts,
        domain,
        Dict{InterpolationFunction, Regressor}(),
        [
            regress(domain, coords) for coords in pts_transpose
        ]
    )
end

"""
Either create a new regressor for the given corner points or fetch a pre-existing one from the
struct's dictionary
"""
function finel_get_regr(fin::Finel, interp::InterpolationFunction)
    if haskey(fin.abs_coord_regressors, interp)
        return fin.abs_coord_regressors[interp]
    end

    regr=Regressor(
        interp,
        fin.pts
    )

    fin.abs_coord_regressors[interp]=regr

    regr
end

"""
Function to get absolute coordinates from finite element
"""
function get_abs_coords!(absc::Vector{Float64}, fin::Finel, xi::Vector{Float64})
    for (i, lps) in enumerate(fin._coord_linpars)
        absc[i]=fin.domain.interp._funcs[1](xi..., lps...)
    end
end


text:
using FileIO
using DelimitedFiles
using Random: seed!

# Experiment values
MinimumReductionsMSE = [1e-1, 1e-2, 1e-3, 1e-4, 1e-5, 1e-6, 1e-7];
MaxNumNodes = 5:5:200;
numFolds = 10;

function loadDataset(datasetName::String; datasetType::DataType=Float64)
    dataset = DelimitedFiles.readdlm(string("datasets/", datasetName, ".tsv"));
    @assert(dataset[1,end]=="target");
    @assert(all(dataset[1,1:end-1].!="target"));
    dataset = datasetType.(dataset[2:end,:]);
    inputs = dataset[:,1:end-1];
    targets = dataset[:, end];
    return (inputs, targets);
end;

using Random:shuffle!

function crossvalidationIndices(N, numFolds)
    # Create the same cross-validation indices used in the rest of experiments
    indicesKFold = repeat(1:numFolds, Int(trunc(N/numFolds)+1))[1:N];
    seed!(1);
    shuffle!(indicesKFold);
    return indicesKFold;
end;

function readFilesInFolder(folder::String; extension::String="")
    fileNames = String[];
    for fileName in readdir(folder)
        if isempty(extension) || ((length(fileName)>length(extension)+1) && (uppercase(fileName[end-length(extension):end]) == uppercase(string(".", extension))))
            push!(fileNames, fileName);
        end;
    end;
    return fileNames;
end;

function datasetNames()
    names = readFilesInFolder("datasets"; extension="tsv")
    return [name[1:end-4] for name in names];
end;

function datasetNameFromFile(name::String)
    i = findfirst("_", name); i = i[1];
    @assert(tryparse(Int, name[1:i-1])!=nothing);
    name = replace(name, "_"=>" ");
    name = name[i+1:end];
    return name;
end;

resultsFileName(datasetName, strategy) = string("results/results_DoME_", datasetName, "_strategy_", strategy, ".jld2");

function parseARGS()
    args = Array{Any,1}(undef, length(ARGS));
    args[:] .= ARGS[:];
    integers = tryparse.(Int,ARGS);
    args[integers.!=nothing] .= integers[integers.!=nothing];
    return args;
end;

function experimentsDoMEFinished(results::Array{<:Real,4})
    if size(results,1)<length(MinimumReductionsMSE) || size(results,2)<length(MaxNumNodes) || (size(results,3)<numFolds)
        return false
    else
        @assert(size(results,3)==numFolds);
        @assert(size(results,4)==5);
        return (!any(isnan.(results[:,:,:,[1,3,4,5]])) && !any(results[:,:,:,3:end].==0))
    end;
end

function experimentsDoMEFinished(datasetName, strategy)
    fileName = resultsFileName(datasetName, strategy);
    !isfile(fileName) && return (false, []);
    results = load(fileName, "results");
    results = results[1:min(size(results,1),length(MinimumReductionsMSE)), 1:min(size(results,2),length(MaxNumNodes)), :, :];
    return (experimentsDoMEFinished(results), results);
end

function experimentsDoMEFinished(datasetName)
    results = Array{Float64,5}(undef, 4, length(MinimumReductionsMSE), length(MaxNumNodes), numFolds, 5);
    for strategy in 1:4
        (finished, resultsThisStrategy) = experimentsDoMEFinished(datasetName, strategy);
        !finished && return (false, []);
        results[strategy, :, :, :, :] .= resultsThisStrategy;
    end;
    return (true, results);
end


text:
immutable WAVHeader
    RIFF::NTuple{4,UInt8}
    fsz1::UInt32
    WAVE::NTuple{4,UInt8}
    fmt::NTuple{4,UInt8}
    hlen::UInt32
    typ::UInt16
    channels::Uint16
    rate::UInt32
    chunk::UInt32
    ttyp::UInt16
    samplebits::UInt16
    data::NTuple{4,UInt8}
    dsiz::UInt32
end

function Base.show(io::IO,hdr::WAVHeader)
    println(io, "WAV file header")
    println(io, " size of header: ", Int(hdr.hlen), " bytes")
    println(io, " type: ", hdr.typ == 1 ? "PCM" : "unknown")
    println(io, " channels: ", Int(hdr.channels))
    println(io, " rate: ", Int(hdr.rate), "Hz.")
    println(io, " track type: ", ["8-bit mono","16-bit mono", "8-bit stereo", "16-bit stereo"][hdr.ttyp])
    println(io, " samplebits: ", Int(hdr.samplebits))
    println(io, " data size: ", Int(hdr.dsiz), " bytes")
end

type InvalidWAVHeader <: Exception
end

function WAVheader(bb::Vector{Uint8})
    length(bb) == 44 || error("WAV header must be 44 bytes")
    h = unsafe_load(convert(Ptr{WAVHeader},pointer(bb)))
    ok = h.RIFF == ('R','I','F','F')
    ok &= h.WAVE == ('W','A','V','E')
    ok &= h.fmt == ('f','m','t',' ')
    ok && h.data == ('d','a','t','a') || throw(InvalidWAVHeader())
    h
end

function WAVopen(fnm::AbstractString)
    ENDIAN_BOM == 0x04030201 || error("Requires a little endian host")
    ff = open(fnm,"r")
    h = WAVheader(readbytes(ff,44))
    mm = mmap_array(Int16,(h.dsiz >> 1,),ff)
    close(ff)
    h, mm
end


text:
using Documenter
using TranscodingStreams
using CodecZlib

makedocs(
    sitename="TranscodingStreams.jl",
    modules=[TranscodingStreams],
    pages=["index.md", "examples.md", "reference.md", "devnotes.md"],
    assets=["assets/custom.css"],
)

deploydocs(repo="github.com/bicycle1885/TranscodingStreams.jl.git")


text:
using ReactiveExtensions
using Base.Test

tests = ["subject",
         "chaining",
         "sequence",
         ]

println("Running tests:")

for t in tests
    println(" * $(t)")
    include("$(t).jl")
end


text:
@testset "Display" begin
    let
        foo() = foo()
        @test Atom.render(Atom.Inline(), Atom.@errs foo()) != nothing
    end

    # issue: https://github.com/JunoLab/Juno.jl/issues/376
    @test Atom.render(Atom.Inline(), :abc)[:contents] == [":abc"]
    @test Atom.render(Atom.Inline(), Symbol("A B"))[:contents] == ["Symbol(\"A B\")"]

    @test length(Atom.trim(rand(50), 20)) == 20

    @test Atom.isanon(sin) == false
    @test Atom.isanon(x -> x) == true # the `== true` might fix a weird travis error

    @test Atom.pluralize([1,2,3], "entry", "entries") == "3 entries"
    @test Atom.pluralize([1], "entry", "entries") == "1 entry"
end


text:
using NewtonsMethod, LinearAlgebra
using Test

@testset "NewtonsMethod.jl" begin

#Test with analytical derivates
f(x)=x^2
f′(x)=2x
@test newtonroot(f,f′,x₀=0.5).root ≈ 0.0 atol=0.0000001

f(x)=x^3-3x^2+2
f′(x)=3x^2-6x
@test newtonroot(f,f′,x₀=0.5).root ≈ 1.0 atol=0.0000001

f(x)=log(x)-1
f′(x)=1/x
@test newtonroot(f,f′,x₀=0.5).root ≈ 2.71828182845904523536

#Test with automatic derivatives
f(x)=x^2
@test newtonroot(f,x₀=0.5).root ≈ 0.0 atol=0.0000001

f(x)=x^3-3x^2+2
@test newtonroot(f,x₀=0.5).root ≈ 1.0 atol=0.0000001

f(x)=log(x)-1
@test newtonroot(f,x₀=0.5).root ≈ 2.71828182845904523536

#Test Big Float
f(x)=sin(x)-1
f′(x)=cos(x)
@test newtonroot(f, f′, x₀ = BigFloat(1.0), tolerance = 1E-32).root ≈ BigFloat(pi/2) atol = 1E-16
@test newtonroot(f, x₀ = BigFloat(1.0), tolerance = 1E-32).root ≈ BigFloat(pi/2) atol = 1E-16

#Test non-convergence
f(x)=x^2+2
@test newtonroot(f,x₀=0.5) == nothing

#Test maxiter
f(x)=log(x)-1
@test newtonroot(f,x₀=0.5,maxiter=5) == nothing

#Test tolerance
f(x)=log(x)-1
@test newtonroot(f,x₀=0.5).root ≈ 2.71828182845904523536
r1=newtonroot(f,x₀=1.0).root
r2=newtonroot(f,x₀=1.0, tolerance=0.00001).root
r3=newtonroot(f,x₀=1.0, tolerance=0.001).root
@test norm(f(r1)) <= norm(f(r2)) <= norm(f(r3)) # accuracy decreases as tolerance increases

end


text:
""
function run_ops(file, model_constructor, optimizer; kwargs...)
    return _PM.run_model(file, model_constructor, optimizer, build_ops;
        ref_extensions=[_PM.ref_add_on_off_va_bounds!], kwargs...)
end

function build_ops(pm::_PM.AbstractPowerModel)

    variable_bus_active_indicator(pm)
    _PMR.variable_bus_voltage_on_off(pm)

    _PM.variable_gen_indicator(pm)
    _PM.variable_gen_power_on_off(pm)

    _PM.variable_storage_indicator(pm)
    _PM.variable_storage_power_mi_on_off(pm)

    _PM.variable_branch_indicator(pm)
    _PM.variable_branch_power(pm)

    _PM.variable_dcline_power(pm)

    _PM.variable_load_power_factor(pm, relax=true)
    _PM.variable_shunt_admittance_factor(pm, relax=true)

    _PMR.constraint_model_voltage_damage(pm)
    for i in _PM.ids(pm, :ref_buses)
        _PM.constraint_theta_ref(pm, i)
    end

    for i in _PM.ids(pm, :gen)
        constraint_generation_active(pm, i)
        _PM.constraint_gen_power_on_off(pm, i)
    end

    for i in _PM.ids(pm, :bus)
        constraint_bus_active(pm, i)
        _PMR.constraint_power_balance_shed(pm, i)
    end

    for i in _PM.ids(pm, :storage)
        constraint_storage_active(pm, i)
        _PM.constraint_storage_state(pm, i)
        _PM.constraint_storage_complementarity_mi(pm, i)
        _PM.constraint_storage_on_off(pm,i)
        _PM.constraint_storage_loss(pm, i)
        _PM.constraint_storage_thermal_limit(pm, i)
    end

    for i in _PM.ids(pm, :branch)
        constraint_branch_active(pm, i)
        _PM.constraint_ohms_yt_from_on_off(pm, i)
        _PM.constraint_ohms_yt_to_on_off(pm, i)

        _PM.constraint_voltage_angle_difference_on_off(pm, i)

        _PM.constraint_thermal_limit_from_on_off(pm, i)
        _PM.constraint_thermal_limit_to_on_off(pm, i)
    end

    for i in _PM.ids(pm, :load)
        constraint_load_active(pm, i)
    end

    for i in _PM.ids(pm, :dcline)
        _PM.constraint_dcline_power_losses(pm, i) #not active decision variables
    end

    # Add Objective
    # ------------------------------------
    # Maximize power delivery while minimizing wildfire risk
    z_demand = _PM.var(pm, nw_id_default, :z_demand)
    # z_storage = _PM.var(pm, nw_id_default, :z_storage)
    z_gen = _PM.var(pm, nw_id_default, :z_gen)
    z_branch = _PM.var(pm, nw_id_default, :z_branch)
    z_bus = _PM.var(pm, nw_id_default, :z_bus)

    if haskey(_PM.ref(pm), :risk_weight)
        alpha = _PM.ref(pm, :risk_weight)
    else
        Memento.warn(_PM._LOGGER, "network data should specify risk_weight, using 0.5 as a default")
        alpha = 0.5
    end

    for comp_type in [:gen, :load, :bus, :branch]
        for (id,comp) in  _PM.ref(pm, comp_type)
            if ~haskey(comp, "power_risk")
                Memento.warn(_PM._LOGGER, "$(comp_type) $(id) does not have a power_risk value, using 0.1 as a default")
                comp["power_risk"] = 0.1
            end
            if ~haskey(comp, "base_risk")
                Memento.warn(_PM._LOGGER, "$(comp_type) $(id) does not have a base_risk value, using 0.1 as a default")
                comp["base_risk"] = 0.1
            end
        end
    end

    JuMP.@objective(pm.model, Max,
        (1-alpha)*(
                sum(z_demand[i]*load["pd"] for (i,load) in _PM.ref(pm,:load))
        )
        - alpha*(
            sum(z_gen[i]*gen["power_risk"]+gen["base_risk"] for (i,gen) in _PM.ref(pm, :gen))
            + sum(z_bus[i]*bus["power_risk"]+bus["base_risk"] for (i,bus) in _PM.ref(pm, :bus))
            + sum(z_branch[i]*branch["power_risk"]+branch["base_risk"] for (i,branch) in _PM.ref(pm, :branch))
            + sum(z_demand[i]*load["power_risk"]+load["base_risk"] for (i,load) in _PM.ref(pm,:load))
            # + sum(z_storage[i]*storage["power_risk"]+storage["base_risk"] for (i,storage) in _PM.ref(pm, :storage))
        )
    )

end


text:
using Documenter, UnivariateFunctions

makedocs(
    format = Documenter.HTML(),
    sitename = "UnivariateFunctions",
    modules = [UnivariateFunctions],
    pages = Any["Overview" => "index.md",
                "Examples" => "Examples.md",
                "API" => "api.md"]
)

deploydocs(
    repo   = "github.com/s-baumann/UnivariateFunctions.jl.git",
    target = "build",
    deps   = nothing,
    make   = nothing
)


text:
"""
This file constructs the functions to find lower and upper bound of optimal
set of sourcing countries according to Jia's algorithm.
"""

## Define module and things to be exported
module JiaAlgorithm

export lowerbound_setup, lowerbound, upperbound_setup, upperbound, optimalset

## Load packages
using LinearAlgebra, Random, Distributions, Statistics, DataFrames, StatsBase
using Combinatorics

## Initialize lower bound and source potential matrix
function lowerbound_setup(N, ξ, my_exp)
    J_lb = zeros(1, N)  # the lower bound is empty set
    source_start_lb = (J_lb * ξ).^my_exp
    temp = repeat(J_lb, N, 1) + I
    temp2 = ones(size(temp,1), size(temp,2))
    check_matrix_lb = min.(temp, temp2)
    source_check_lb = (check_matrix_lb * ξ).^my_exp

    return source_start_lb, source_check_lb
end

## Jia's lower bound algorithm
function lowerbound(source_start, source_check, ϕ_σ_B, fc, N, ξ, my_exp, firm)

    # Start iteration for Marginal Benefit (MB)
    k = 1
    Z_start  = zeros(1, N)

    while k <= N

    # compute MB
        if k > 1
            source_start = (Z_start * ξ).^my_exp
            source_check = (Z_start * ξ .+ ξ .* (1 .- Z_start')).^my_exp
        end

        source_potential_start = ϕ_σ_B[firm] .* source_start
        source_potential_new_vec = ϕ_σ_B[firm] .* source_check

        # generate matrix with 1 if MB positive and update set of sourcing countries
        MB = (source_potential_new_vec' - fc[firm, :]' .- source_potential_start .> 0)
        Z_new = min.(Z_start + MB, ones(size(MB,1), size(MB,2)))

        if Z_start == Z_new
            return Z_new
            @goto end_lb_algorithm
        end

        k += 1
        Z_start = Z_new
        return Z_new
    end

    @label end_lb_algorithm
    return Z_new, k
end

## Initialize upper bound and source potential matrix
function upperbound_setup(N, ξ, my_exp)
    J_ub = ones(1, N)  # the upper bound is full set
    source_start_ub = (J_ub * ξ).^my_exp
    temp = repeat(J_ub, N, 1) - I
    temp2 = zeros(size(temp,1), size(temp,2))
    check_matrix_ub = max.(temp, temp2)
    source_check_ub = (check_matrix_ub * ξ).^my_exp

    return source_start_ub, source_check_ub
end

## Jia's upper bound algorithm
function upperbound(source_start, source_check, ϕ_σ_B, fc, N, ξ, my_exp, firm)

    # Start iteration for Marginal Benefit (MB)
    k = 1
    Z_start  = ones(1, N)

    while k <= N

    # compute MB
        if k > 1
            source_start = (Z_start * ξ).^my_exp
            source_check = (Z_start * ξ .- ξ .* (Z_start')).^my_exp
        end

        source_potential_start = ϕ_σ_B[firm] .* source_start
        source_potential_new_vec = ϕ_σ_B[firm] .* source_check

        # generate matrix with 1 if MB positive and update set of sourcing countries
        MB = (source_potential_start .- source_potential_new_vec' - fc[firm, :]'  .< 0)
        Z_new = max.(Z_start - MB, zeros(size(MB,1), size(MB,2)))

        if Z_start == Z_new
            return Z_new
            @goto end_ub_algorithm
        end

        k += 1
        Z_start = Z_new
        return Z_new
    end

    @label end_ub_algorithm
    return Z_new, k
end

## Check if Jia's algorithm produce the same lower and upper bound
function optimalset(Z, gap_bounds, firm, Z_lb, Z_ub, S, N, num_rand_checks, rand_check_matrix, fc, ξ, my_exp, ϕ_σ_B)

    if Z_lb == Z_ub
        print("")
        Z[firm,:] = Z_lb

    # I don't know how to do lines 64-88 of gmm_objective.m: the case when lower and
    # upper bounds are not too different

    else
        #print("WARNING! The sourcing strategy may not be solved correctly")
        #print("for firm number")
        #println("$firm")
        Z_check = repeat(Z_lb, num_rand_checks, 1)
        ind_diffZ = zeros(1,N)
        for i in 1:N
            if Z_ub[i] == Z_lb[i]
                ind_diffZ[i] = 0.0
            else
                ind_diffZ[i] = 1.0
            end
        end
        gap_bounds[firm] = sum(ind_diffZ)
        K_top = Int(sum(ind_diffZ))
        K_diff = [i[2] for i in findall(x->x!=0.0, ind_diffZ)]
        for K in 1:K_top
            Z_check[:,K_diff[K]] = (rand_check_matrix[:,K_diff[K]] .> 0.5)
        end

        # Use the check and both bounds to find new set of sourcing countries
        Z_check = [Z_check; Z_lb; Z_ub]
        fc_payments = sum(Z_check .* fc[firm,:]', dims=2)
        source_potential_shock_mat_check = (Z_check * ξ).^my_exp
        source_potential_vec = ϕ_σ_B[firm] .* source_potential_shock_mat_check
        tot_profit = source_potential_vec - fc_payments
        loc_firm_best = findmax(tot_profit, dims = 1)[2]

        # replace Z as the set of locations from Z_check that maximize profits
        Z[firm,:] = Z_check[loc_firm_best[1][1],:]
    end

    return Z
end

end     # end module


text:
using Distributed
using ArDCA
using ArDCA: _outputarnet!
using JLD
using ExtractMacro: @extract
using LinearAlgebra: norm
using Base.Threads: @threads, nthreads, threadid
using ProgressMeter
using UnicodePlots
using PdbTool
using LoopVectorization
using Statistics
using ArgParse
using DCAUtils
using Printf
using HDF5
using Random

function make_chunks(M, chunk_size)
    return [i:min(M, i+chunk_size-1) for i in 1:chunk_size:M]
end

function make_pairs(N)
    npairs = binomial(N,2)
    l = 0
    pairs = []
    for i in 1:N
        for j in (i+1):N
            l+=1
            push!(pairs, (i,j,l))
        end
    end
end

function load_model(model_path)

    if !isfile(model_path)
        error("cannot load $model_path")
    end
    data = load(model_path)
    arnet = data["arnet"]
    episcore = data["episcore"]
    return arnet, episcore

end

function train_model(fasta_path, model_path, overwrite, randomize; lambdaJ=0.01, lambdaH=0.0001, theta=0.2, permorder=:NATURAL, max_gap_fraction=1)

    if isfile(model_path) && !overwrite
        error("Model $model_path exists")
    end

    @show fasta_path
    arnet, arvar = ardca(fasta_path, verbose=true, lambdaJ=lambdaJ, lambdaH=lambdaH, theta=theta, permorder=permorder, max_gap_fraction=max_gap_fraction)
    if randomize
        for _H in arnet.H
            _H[:] = randn(length(_H))
        end
        for _J in arnet.J
            _J[:] = randn(length(_J))
        end
    end
    episcore = epistatic_score(arnet, arvar, 1)
    save(model_path, "arnet", arnet, "episcore", episcore)

    return arnet, episcore

end

function getNq(arnet::ArNet)
    N = length(arnet.idxperm)
    q = length(arnet.H[1])
    return N, q
end

function get_plike(x, J, H, p0, N, q)
    plike = log(p0[x[1]])
    totH = zeros(q)
    for site in 1:N-1
        Js = J[site]
        h = H[site]
        copy!(totH,h)
        @avx for i in 1:site
            for a in 1:q
                totH[a] += Js[a,x[i],i]
            end
        end
        ArDCA.softmax!(totH)
        plike += log(totH[x[site+1]])
    end
    return plike
end

function extract_couplings_q0(arnet, gauge_aa; pc=0.01)

    idxperm = arnet.idxperm
    H = arnet.H
    J = arnet.J
    p0 = arnet.p0

    N, q = getNq(arnet)
    L = binomial(N,2)

    total = L * q * q + N * q + 1

    backorder = sortperm(arnet.idxperm)
    plike_mean = 0
    R = zeros(L, q, q)
    Q = zeros(N, q)

    progress = Progress(total)

    ppc = (1-pc) * arnet.p0 + pc * ones(q)/q

    # get gauge
    seq = fill(gauge_aa, N)
    gauge = get_plike(seq[idxperm], J, H, ppc, N, q)
    next!(progress)

    for i in 1:N
        for a in 1:q
            a == gauge_aa && continue
            seq = fill(gauge_aa, N)
            seq[i] = a
            Q[i,a] = get_plike(seq[idxperm], J, H, ppc, N, q) - gauge
            next!(progress)
        end
    end

    l = 0
    for i in 1:N
        for j in (i+1):N
            l += 1
            for a in 1:q
                for b in 1:q
                    b == gauge_aa && continue
                    seq = fill(gauge_aa, N)
                    seq[i] = a
                    seq[j] = b
                    R[l,a,b] = get_plike(seq[idxperm], J, H, ppc, N, q) - Q[i,a] - Q[j,b] - gauge
                    next!(progress)
                end
            end
        end
    end

    return Q, R

end

function extract_couplings_thread!(id, arnet, ppc, R, Q, R_count, Q_count, M_thread)

    @extract arnet : H J p0 idxperm

    N, q = getNq(arnet)
    Z = rand(1:q, N, M_thread)

    plikes = zeros(M_thread)
    for m in 1:M_thread
        plikes[m] = get_plike(Z[:, m], J, H, ppc, N, q)
    end

    ArDCA.permuterow!(Z, sortperm(idxperm))

    for i in 1:N
        @avx for m in 1:M_thread
            plike = plikes[m]
            a = Z[i,m]
            Q[i, a] += plike
            Q_count[i, a] += 1
        end
    end

    l = 0
    for i in 1:N
        for j in (i+1):N
            l += 1
            @avx for m in 1:M_thread
                plike = plikes[m]
                a = Z[i,m]
                b = Z[j,m]
                R[l,a,b] += plike
                R_count[l,a,b] += 1
            end
        end
    end

    return sum(plikes)

end

function extract_couplings(arnet, M; pc=0.01, chunks_per_thread=1000)

    idxperm = arnet.idxperm
    H = arnet.H
    J = arnet.J
    p0 = arnet.p0

    N, q = getNq(arnet)
    L = binomial(N,2)

    nt = nthreads()

    backorder = sortperm(arnet.idxperm)
    plike_mean = 0
    R = zeros(nt, L, q, q)
    Q = zeros(nt, N, q)
    R_count = zeros(nt, L, q, q)
    Q_count = zeros(nt, N, q)

    chunk_size_M = cld(M, nthreads()*chunks_per_thread)
    chunks_M = make_chunks(M, chunk_size_M)
    progress = Progress(length(chunks_M), 1)

    ppc = (1-pc) * arnet.p0 + pc * ones(q)/q

    plikes = zeros(nt)
    plikes_count = zeros(nt)

    @threads for chunk in chunks_M

       id = threadid()

       R_thread = @view R[id, :, :, :]
       Q_thread = @view Q[id, :, :]
       R_count_thread = @view R_count[id, :, :, :]
       Q_count_thread = @view Q_count[id, :, :]

       m = extract_couplings_thread!(id,
                                     arnet,
                                     ppc,
                                     R_thread,
                                     Q_thread,
                                     R_count_thread,
                                     Q_count_thread,
                                     length(chunk))

       plikes[id] += m
       plikes_count[id] += length(chunk)
       next!(progress)

    end

    plike_mean = sum(plikes)/sum(plikes_count)

    R = dropdims(sum(R, dims=1), dims=1)
    R_count = dropdims(sum(R_count, dims=1), dims=1)
    Q = dropdims(sum(Q, dims=1), dims=1)
    Q_count = dropdims(sum(Q_count, dims=1), dims=1)

    R[:] = R[:]./R_count[:]
    Q[:] = Q[:]./Q_count[:] .- plike_mean
    l = 0
    for i in 1:N
        for j in (i+1):N
            l+=1
            for a in 1:q
                for b in 1:q
                    R[l, a, b] -= (Q[i, a] + Q[j,b] + plike_mean)
                end
            end
        end
    end
    return Q, R
end

function extract_or_load_couplings(arnet, model_path, M)

    coupling_file = model_path*".$M.jld"
    if isfile(coupling_file)
        println("Loading couplings from $coupling_file")
        data = load(coupling_file)
        Q = data["Q"]
        R = data["R"]
        return Q, R
    else
        Q, R = extract_couplings(arnet; M=M)
        save(coupling_file, "Q", Q, "R", R)
        return Q, R
    end
end

function frobnorm(R, N)
    L = binomial(N,2)
    F = zeros(N, N)
    l = 0
    for i in 1:N
        for j in (i+1):N
            l += 1
            #f = norm(R[l, 1:end-1, 1:end-1])
            f = norm(R[l, :, :])
            F[i,j] = f
            F[j,i] = f
        end
    end
    return F
end

function get_num2letter(q)
    num2letter = fill('-', q)
    for c in 'A':'Z'
        num = DCAUtils.ReadFastaAlignment.letter2num(c)
        if num == q
            continue
        end
        num2letter[num] = c
    end
    return num2letter
end

function get_plike!(Z, q, arnet, pc)
    N, M = size(Z)
    ppc = (1-pc) * arnet.p0 + pc * ones(q)/q
    ArDCA.permuterow!(Z, arnet.idxperm)
    plikes = [get_plike(Z[:,m], arnet.J, arnet.H, ppc, N, q) for m in 1:M]
    return plikes
end

function get_plike(Z, q, arnet, pc)
    plikes = get_plike!(Z, q, arnet, pc)
    ArDCA.permuterow!(Z, sortperm(arnet.idxperm))
    return plikes
end

function sample_arnet(arnet, N, q, pc, nsamples, sample_batch_size, out_file, sample_dist="M")
    if nsamples % sample_batch_size != 0
        error("sample_batch_size has to divide nsamples")
    end
    num2letter = get_num2letter(q)
    samples = zeros(Int8, nsamples, N)
    logp = zeros(Float64, nsamples)
    nbatches = div(nsamples, sample_batch_size)
    progress = Progress(nbatches)
    sample_id = 1
    for batch_id in 1:nbatches
        if sample_dist=="M"
            samples_batch = sample(arnet, sample_batch_size)
        else
            samples_batch = rand(1:q, N, sample_batch_size)
        end
        logp_batch = get_plike(samples_batch, q, arnet, pc)
        for m_batch in 1:sample_batch_size
            logp[sample_id] = logp_batch[m_batch]
            for pos in 1:N
                samples[sample_id, pos] = samples_batch[pos, m_batch]
            end
            sample_id += 1
        end
        next!(progress)
    end
    @assert sample_id == nsamples+1
    # need to transpose and subtract 1
    h5write(out_file, "samples", convert.(Int8, transpose(samples) .- 1))
    h5write(out_file, "logp", convert(Array{Float32}, logp))
end

args = ArgParseSettings()
@add_arg_table! args begin
    "--gauge_aa"
    arg_type=Int
    default=21
    "--model_path"
    default=nothing
    "--fasta_path"
    default=nothing
    "--out_file"
    default="stdout"
    "--overwrite"
    action=:store_true
     "--extracted_couplings_path"
    default=nothing
    "--randomize"
    action=:store_true
    "--M"
    arg_type=Int
    default=nothing
    "--pc"
    arg_type=Float64
    default=0.1
    "--lambdaJ"
    arg_type=Float64
    default=0.01
    "--lambdaH"
    arg_type=Float64
    default=0.0001
    "--theta"
    arg_type=Float64
    default=0.2
    "mode"
    help="either train, extract, evaluate or randomize"
    required=true
    "--sample_batch_size"
    arg_type=Int
    default=1000
    "--nsamples"
    arg_type=Int
    default=10000
    "--sample_dist"
    arg_type=String
    "--chain"
    default="A"
    "--pdb_path"
    default=nothing
    "--seed"
    default=1
end

parsed_args = parse_args(ARGS, args)

# first of all we seed the RNG
Random.seed!(parsed_args["seed"])

# train
if parsed_args["mode"]=="train"
    if !("fasta_path" in keys(parsed_args))
        error("fasta_path not set")
    end
    if !("model_path" in keys(parsed_args))
        error("model_path not set")
    end
    train_model(parsed_args["fasta_path"],
                parsed_args["model_path"],
                parsed_args["overwrite"],
                parsed_args["randomize"];
                lambdaH = parsed_args["lambdaH"],
                lambdaJ = parsed_args["lambdaJ"])
end

# extract
if parsed_args["mode"]=="extract"
    if !("model_path" in keys(parsed_args))
        error("model_path not set")
    end
    if !("extracted_couplings_path" in keys(parsed_args))
        error("extracted_couplings_path not set")
    end
    if !("M" in keys(parsed_args))
        error("M not set")
    end

    arnet, episcore = load_model(parsed_args["model_path"])
    Q, R = extract_couplings(arnet, parsed_args["M"])
    h5open(parsed_args["extracted_couplings_path"], "w") do fid
        write(fid, "fields", Q)
        write(fid, "couplings", R)
    end
end

# extract q0
if parsed_args["mode"]=="extract_q0"
    if !("model_path" in keys(parsed_args))
        error("model_path not set")
    end
    if !("extracted_couplings_path" in keys(parsed_args))
        error("extracted_couplings_path not set")
    end
    if !("gauge_aa" in keys(parsed_args))
        error("gauge_aa not set")
    end

    arnet, episcore = load_model(parsed_args["model_path"])
    Q, R = extract_couplings_q0(arnet, parsed_args["gauge_aa"])
    h5open(parsed_args["extracted_couplings_path"], "w") do fid
        write(fid, "fields", Q)
        write(fid, "couplings", R)
    end
end

if parsed_args["mode"]=="evaluate"
    if !("model_path" in keys(parsed_args))
        error("model_path not set")
    end
    if !("fasta_path" in keys(parsed_args))
        error("fasta_path not set")
    end
    if !("out_file" in keys(parsed_args))
        error("out_file not set")
    end

    arnet, episcore = load_model(parsed_args["model_path"])
    N, q = getNq(arnet)
    pc = parsed_args["pc"]
    Z = DCAUtils.read_fasta_alignment(parsed_args["fasta_path"], 1.0)
    ArDCA.permuterow!(Z, arnet.idxperm)
    @assert N == size(Z,1)
    M = size(Z,2)

    ppc = (1-pc) * arnet.p0 + pc * ones(q)/q
    plikes = [get_plike(Z[:,m], arnet.J, arnet.H, ppc, N, q) for m in 1:M]
    out_file = parsed_args["out_file"]
    if parsed_args["out_file"]=="stdout"
        for m in 1:length(plikes)
            @printf("%f\n", plikes[m])
        end
    else
        # attention: need to revert permutation
        backorder = sortperm(arnet.idxperm)
        ArDCA.permuterow!(Z, backorder)
        h5write(out_file, "samples", convert.(Int8, Z .- 1))
        h5write(out_file, "logp", convert(Array{Float32}, plikes))
   end
end

function get_ppv(episcore, pdb_path, fasta_path, chain)
    hmm_path = tempname()
    run(pipeline(fasta_path, `sed '/^>/!s/[a-z,\.]//g'`, `hmmbuild -n hmm --symfrac=0.0 --informat=a2m $hmm_path - `))
    pdb = PdbTool.parsePdb(pdb_path)
    PdbTool.mapChainToHmm(pdb.chain[chain], hmm_path)
    roc = PdbTool.makeIntraRoc(episcore, pdb.chain[chain])
    println(lineplot([roc[i][3] for i in 1:100]))
end

if parsed_args["mode"]=="randomize"
    if !("model_path" in keys(parsed_args))
        error("model_path not set")
    end
    if !("out_file" in keys(parsed_args))
        error("out_file not set")
    end
    arnet, episcore = load_model(parsed_args["model_path"])
    N, q = getNq(arnet)
    for subJ in arnet.J
        subJ[:] = randn(size(subJ))[:]
    end
    for subH in arnet.H
        subH[:] = randn(size(subH))
    end
    arnet.idxperm[:] = randperm(N)[:]
    save(parsed_args["out_file"]*".jld", "arnet", arnet, "episcore", episcore)
end

if parsed_args["mode"]=="sample"
    if !("model_path" in keys(parsed_args))
        error("model_path not set")
    end
    if !("out_file" in keys(parsed_args))
        error("out_file not set")
    end
    if !("nsamples" in keys(parsed_args))
        error("nsamples not set")
    end
    arnet, episcore = load_model(parsed_args["model_path"])
    N, q = getNq(arnet)
    samples = sample_arnet(arnet, N, q, parsed_args["pc"], parsed_args["nsamples"], parsed_args["sample_batch_size"], parsed_args["out_file"], parsed_args["sample_dist"])
end

if parsed_args["mode"]=="ppv"
    if !("model_path" in keys(parsed_args))
        error("model_path not set")
    end
    if !("pdb_path" in keys(parsed_args))
        error("pdb_path not set")
    end
    if !("fasta_path" in keys(parsed_args))
        error("fasta_path not set")
    end
    arnet, episcore = load_model(parsed_args["model_path"])
    get_ppv(episcore, parsed_args["pdb_path"], parsed_args["fasta_path"], parsed_args["chain"])
end

max_stars_repo_name: christophfeinauer/PairwiseDistillations

max_issues_repo_name: christophfeinauer/PairwiseDistillations

max_forks_repo_name: christophfeinauer/PairwiseDistillations


text:
struct BalancedTernary <: Signed
    digits::Vector{Int8}
end
BalancedTernary() = zero(BalancedTernary)
BalancedTernary(n) = convert(BalancedTernary, n)

const sgn2chr = Dict{Int8,Char}(-1 => '-', 0 => '0', +1 => '+')
Base.show(io::IO, bt::BalancedTernary) = print(io, join(sgn2chr[x] for x in reverse(bt.digits)))
Base.copy(bt::BalancedTernary) = BalancedTernary(copy(bt.digits))
Base.zero(::Type{BalancedTernary}) = BalancedTernary(Int8[0])
Base.iszero(bt::BalancedTernary) = bt.digits == Int8[0]
Base.convert(::Type{T}, bt::BalancedTernary) where T<:Number = sum(3 ^ T(ex - 1) * s for (ex, s) in enumerate(bt.digits))
function Base.convert(::Type{BalancedTernary}, n::Signed)
    r = BalancedTernary(Int8[])
    if iszero(n) push!(r.digits, 0) end
    while n != 0
        if mod(n, 3) == 0
            push!(r.digits, 0)
            n = fld(n, 3)
        elseif mod(n, 3) == 1
            push!(r.digits, 1)
            n = fld(n, 3)
        else
            push!(r.digits, -1)
            n = fld(n + 1, 3)
        end
    end
    return r
end
const chr2sgn = Dict{Char,Int8}('-' => -1, '0' => 0, '+' => 1)
function Base.convert(::Type{BalancedTernary}, s::AbstractString)
    return BalancedTernary(getindex.(chr2sgn, collect(reverse(s))))
end

macro bt_str(s)
    convert(BalancedTernary, s)
end

const table = NTuple{2,Int8}[(0, -1), (1, -1), (-1, 0), (0, 0), (1, 0), (-1, 1), (0, 1)]
function _add(a::Vector{Int8}, b::Vector{Int8}, c::Int8=Int8(0))
    if isempty(a) || isempty(b)
        if c == 0 return isempty(a) ? b : a end
        return _add([c], isempty(a) ? b : a)
    else
        d, c = table[4 + (isempty(a) ? 0 : a[1]) + (isempty(b) ? 0 : b[1]) + c]
        r = _add(a[2:end], b[2:end], c)
        if !isempty(r) || d != 0
            return unshift!(r, d)
        else
            return r
        end
    end
end
function Base.:+(a::BalancedTernary, b::BalancedTernary)
    v = _add(a.digits, b.digits)
    return isempty(v) ? BalancedTernary(0) : BalancedTernary(v)
end
Base.:-(bt::BalancedTernary) = BalancedTernary(-bt.digits)
Base.:-(a::BalancedTernary, b::BalancedTernary) = a + (-b)
function _mul(a::Vector{Int8}, b::Vector{Int8})
    if isempty(a) || isempty(b)
        return Int8[]
    else
        if a[1] == -1 x = (-BalancedTernary(b)).digits
        elseif a[1] == 0 x = Int8[]
        elseif a[1] == 1 x = b end
        y = append!(Int8[0], _mul(a[2:end], b))
        return _add(x, y)
    end
end
function Base.:*(a::BalancedTernary, b::BalancedTernary)
    v = _mul(a.digits, b.digits)
    return isempty(v) ? BalancedTernary(0) : BalancedTernary(v)
end

a = bt"+-0++0+"
println("a: $(Int(a)), $a")
b = BalancedTernary(-436)
println("b: $(Int(b)), $b")
c = BalancedTernary("+-++-")
println("c: $(Int(c)), $c")
r = a * (b - c)
println("a * (b - c): $(Int(r)), $r")

@assert Int(r) == Int(a) * (Int(b) - Int(c))

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
import GLAbstraction: bind, unbind, draw

#TODO think about contexts
#TODO handle resizing properly
@component mutable struct Canvas
    name             :: Symbol
    area             :: Area{Int}
    native_window    :: GLFW.Window
    imgui_context    :: UInt
    background       :: Colorant{Float32,4}
    cursor_position  :: NTuple{2,Float64}
    scroll           :: NTuple{2,Float64}
    has_focus        :: Bool
    mouse_buttons    :: Tuple{GLFW.MouseButton,GLFW.Action,Int}
    keyboard_buttons :: Tuple{GLFW.Key,Int,GLFW.Action,Int}
    framebuffer_size :: NTuple{2,Int}

    function Canvas(name::Symbol, area, nw, background)
        ctx = convert(UInt, CImGui.GetCurrentContext())
        if ctx == 0
            ctx = convert(UInt, CImGui.CreateContext())
        else
            CImGui.DestroyContext(CImGui.GetCurrentContext())
            ctx = convert(UInt, CImGui.CreateContext())
        end
        ImGui_ImplGlfw_InitForOpenGL(nw, true)
        ImGui_ImplOpenGL3_Init(420)
        CImGui.Set_KeyMap(CImGui.GetIO(), CImGui.LibCImGui.ImGuiKey_Backspace,
                          GLFW.KEY_BACKSPACE)
        CImGui.Set_KeyMap(CImGui.GetIO(), CImGui.LibCImGui.ImGuiKey_Delete, GLFW.KEY_DELETE)

        obj = new(name, area, nw, ctx, background, (0.0, 0.0), (0.0, 0.0), true,
                  (GLFW.MOUSE_BUTTON_1, GLFW.RELEASE, 0),
                  (GLFW.KEY_UNKNOWN, 0, GLFW.RELEASE, 0), (0, 0))

        GLFW.SetCursorPosCallback(nw,
                                  (nw, x::Cdouble, y::Cdouble) -> begin
                                      obj.cursor_position = (x, y)
                                  end)

        GLFW.SetScrollCallback(nw,
                               (nw, xoffset::Cdouble, yoffset::Cdouble) -> begin
                                   obj.scroll = (obj.scroll[1] + xoffset,
                                                 obj.scroll[2] + yoffset)
                               end)

        GLFW.SetWindowFocusCallback(nw, (nw, focus::Bool) -> begin
                                        obj.has_focus = focus
                                    end)

        GLFW.SetMouseButtonCallback(nw,
                                    (nw, button::GLFW.MouseButton, action::GLFW.Action, mods::Cint) -> begin
                                        obj.mouse_buttons = (button, action, Int(mods))
                                    end)

        GLFW.SetKeyCallback(nw,
                            (nw, button::GLFW.Key, scancode::Cint, action::GLFW.Action, mods::Cint) -> begin
                                obj.keyboard_buttons = (button, Int(scancode), action,
                                                        Int(mods))
                            end)

        GLFW.SetFramebufferSizeCallback(nw,
                                        (nw, w::Cint, h::Cint) -> begin
                                            obj.framebuffer_size = (Int(w), Int(h))
                                        end)

        finalizer(free!, obj)
        return obj
    end
    # framebuffer::FrameBuffer # this will become postprocessing passes. Each pp has a
end

Base.size(area::Area) = (GeometryBasics.widths(area)...,)

function Canvas(name = :Glimpse; kwargs...)
    defaults = mergepop!(canvas_defaults(), kwargs)

    window_hints  = GLFW_DEFAULT_WINDOW_HINTS
    context_hints = GLFW.standard_context_hints(defaults[:major], defaults[:minor])

    area = defaults[:area]
    # if !isassigned(GLFW_context)
    # GLFW_context[] = GLFW.Window(name = string(name),
    window = GLFW.Window(; name = string(name),
                         resolution = (GeometryBasics.widths(area)...,),
                         debugging = defaults[:debugging], major = defaults[:major],
                         minor = defaults[:minor], windowhints = window_hints,
                         contexthints = context_hints, visible = false, focus = false,
                         fullscreen = defaults[:fullscreen], monitor = defaults[:monitor])

    GLFW.SwapInterval(0) # deactivating vsync seems to make everything quite a bit smoother

    background = defaults[:background]

    c = Canvas(name, area, window, background)
    return c
end

function make_current(c::Canvas)
    GLFW.MakeContextCurrent(c.native_window)
    return GLA.set_context!(c)
end

function expose(c::Canvas)
    make_current(c)
    GLFW.SetWindowShouldClose(c.native_window, false)
    return GLFW.ShowWindow(c.native_window)
end

function swapbuffers(c::Canvas)
    if c.native_window.handle == C_NULL
        warn("Native Window handle of canvas $(c.name) == C_NULL!")
        return
    end
    GLFW.SwapBuffers(c.native_window)
    return
end

function isopen(canvas::Canvas)
    canvas.native_window.handle == C_NULL && return false
    return GLFW.GetWindowAttrib(canvas.native_window, GLFW.VISIBLE) == 1
end

#Should this clear the context?
function close(c::Canvas)
    GLFW.HideWindow(c.native_window)
    return should_close!(c, false)
end

should_close!(c::Canvas, b) = GLFW.SetWindowShouldClose(c.native_window, b)

should_close(c::Canvas) = GLFW.WindowShouldClose(c.native_window)

function clear!(c::Canvas, color = c.background)
    glClearColor(color.r, color.g, color.b, color.alpha)
    # glClearColor(1,1,1,1)
    return glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
end

pollevents(c::Canvas) = GLFW.PollEvents()

waitevents(c::Canvas) = GLFW.WaitEvents()

function GLA.free!(c::Canvas)
    if GLA.is_current_context(c)
        GLFW.DestroyWindow(c.native_window)
        clear_context!()
    end
end

bind(c::Canvas, target = GL_FRAMEBUFFER) = glBindFramebuffer(target, 0)

draw(c::Canvas) = nothing

nativewindow(c::Canvas) = c.native_window

Base.size(canvas::Canvas) = size(canvas.area)
Base.size(canvas::Canvas, i::Int) = size(canvas.area)[i]

function Base.resize!(c::Canvas, wh::NTuple{2,Integer}, resize_window = false)
    resize!(GLA.context_framebuffer(), wh)
    return c.area = Area{Int}(0, 0, wh...)
end

callback_value(c::Canvas, cb::Symbol) = c.callbacks[cb][]

callback(c::Canvas, cb::Symbol) = c.callbacks[cb]

windowsize(canvas::Canvas) = GLFW.GetWindowSize(nativewindow(canvas))

set_background_color!(canvas::Canvas, color::Colorant) = canvas.background = convert(RGBA{Float32}, color)
set_background_color!(canvas::Canvas, color::NTuple)   = canvas.background = convert(RGBA{Float32}, color)

#---------------------DEFAULTS-------------------#

# canvas_defaults() = SymAnyDict(:area       => Area{Int}(0, 0, glfw_standard_screen_resolution()...),
function canvas_defaults()
    return SymAnyDict(:area       => Area{Int}(0, 0, 3840, 2160),
                      :background => RGBA(1.0f0), :depth      => GLA.Depth{Float32},
                      :callbacks  => standard_callbacks(), :debugging  => false,
                      :major      => 3, :minor      => 3, :clear      => true,
                      :hidden     => false, :visible    => true, :focus      => false,
                      :fullscreen => false, :monitor    => nothing)
end

@component struct FullscreenVao
    vao::VertexArray
end

function FullscreenVao()
    return FullscreenVao(VertexArray([BufferAttachmentInfo(:position, GLint(0),
                                                           Buffer(fullscreen_pos),
                                                           GEOMETRY_DIVISOR),
                                      BufferAttachmentInfo(:uv, GLint(1),
                                                           Buffer(fullscreen_uv),
                                                           GEOMETRY_DIVISOR)],
                                     GL_TRIANGLE_STRIP))
end

bind(v::FullscreenVao) = bind(v.vao)

draw(v::FullscreenVao) = draw(v.vao)

unbind(v::FullscreenVao) = unbind(v.vao)

# I'm not sure this is nice design idk
abstract type RenderTarget <: ComponentData end

macro render_target(name)
    return esc(quote
                   @component struct $name <: RenderTarget
                       target     :: Union{FrameBuffer,Canvas}
                       background :: RGBAf0
                   end
               end)
end

@render_target IOTarget

bind(r::RenderTarget, args...) = bind(r.target, args...)

draw(r::RenderTarget, args...) = draw(r.target, args...)

clear!(r::RenderTarget, c = r.background) = clear!(r.target, c)

Base.size(r::RenderTarget) = size(r.target)

GLA.depth_attachment(r::RenderTarget, args...) = GLA.depth_attachment(r.target, args...)

GLA.color_attachment(r::RenderTarget, args...) = GLA.color_attachment(r.target, args...)

GLA.free!(r::RenderTarget) = free!(r.target)

Base.resize!(r::RenderTarget, args...) = resize!(r.target, args...)

@component @with_kw mutable struct TimingData
    time::Float64          = time()
    dtime::Float64         = 0.0
    frames::Int            = 0
    preferred_fps::Float64 = 60
    reversed::Bool         = false
    timer::TimerOutput     = TimerOutput()
end

abstract type RenderProgram <: ComponentData end

macro render_program(name)
    return esc(quote
                   @component struct $name <: RenderProgram
                       program::GLA.Program
                   end
               end)
end

bind(p::RenderProgram) = bind(p.program)

unbind(p::RenderProgram) = unbind(p.program)

GLA.gluniform(p::RenderProgram, args...) = GLA.gluniform(p.program, args...)

function generate_buffers(p::RenderProgram, args...; kwargs...)
    return generate_buffers(p.program, args...; kwargs...)
end

@component @with_kw struct UpdatedComponents
    components::Vector{DataType} = DataType[]
end

Base.empty!(uc::UpdatedComponents) = empty!(uc.components)

Base.iterate(uc::UpdatedComponents, r...) = iterate(uc.components, r...)

Base.push!(uc::UpdatedComponents, t::T) where {T<:ComponentData} = push!(uc.components, T)
Base.push!(uc::UpdatedComponents, t::DataType)                   = push!(uc.components, t)

function update_component!(uc::UpdatedComponents, ::Type{T}) where {T<:ComponentData}
    if !in(T, uc.components)
        push!(uc, T)
    end
end


text:
"""

$(SIGNATURES)

Helper infrastructure to compile and sample models using `cmdstan`.

"""
module StanOptimize

using StanBase
using DocStringExtensions: FIELDS, SIGNATURES, TYPEDEF

import StanBase: stan_sample, get_cmdstan_home
import StanBase: cmdline, stan_summary, read_summary
import StanBase: RandomSeed, Init, Output, StanModelError

include("stanmodel/optimize_types.jl")
include("stanmodel/OptimizeModel.jl")
include("stanrun/cmdline.jl")
include("stansamples/read_optimize.jl")

stan_optimize = stan_sample

export
  OptimizeModel,
  StanModelError,
  stan_optimize,
  read_optimize,
  read_summary,
  stan_summary

end # module


text:
"""
requireprobs \\

For returning the probabilities required to calculate the entropies \\
Arguments:  x = random variable, usually input \\
            xn = number of variables in x \\
            xmax = the maximum number of possible values across the dimensions in x \\
            y = random variable, usually output/classification \\
            yn = number of variables in y \\
            ymax = the maximum number of possible values across the dimensions in y \\
            calc = The entropies to calculate ex: ["HX", "HXY"] \\
            sampling = method for sampling ex: "naive" \\
Returns:    probs, dictionary probability string => probabilities \\
"""

function requireprobs(x::Array{Int}, xn::Int, xmax::Int, y::Array{Int}, yn::Int, ymax::Int, calc::Array{String, 1}, sampling::String)
  probs = Dict{String, Array{Float64}}()
  dec_x = decimalise(x, xn, xmax)
  dec_y = decimalise(y, yn, ymax)
  x_dist_dim = xmax ^ xn # the dimension of the probability distribution required for x
  y_dist_dim = ymax ^ yn
  reduce(|, [in(ent, calc) for ent in ["HXY" "HY"]]) && (probs["PY"] = prob(dec_y, y_dist_dim, sampling))
  reduce(|, [in(ent, calc) for ent in ["HXY" "HX"]]) && (probs["PX"] = prob(dec_x, x_dist_dim, sampling))
  reduce(|, [in(ent, calc) for ent in ["HXY"]]) && (probs["PXY"] = condprob(dec_x, dec_y, x_dist_dim, y_dist_dim, sampling)[2])
  return probs
end


text:
"""
    sum(dv::AbstractDistributionVector, ms=PassMissing(); <keyword arguments>)
    sum(dv::AbstractDistributionVector, corr, ms=PassMissing(); <keyword arguments>)
    sum(dv::AbstractDistributionVector, acf, ms=PassMissing(); <keyword arguments>)

    mean(dv::AbstractDistributionVector, ms=PassMissing(); <keyword arguments>)
    mean(dv::AbstractDistributionVector, corr, ms=PassMissing(); <keyword arguments>)
    mean(dv::AbstractDistributionVector, acf, ms=PassMissing(); <keyword arguments>)

Compute the distribution of the sum or mean of correlated random variables.

# Arguments
- `dv`: The vector of distributions, see [`AbstractDistributionVector`](@ref)
- `ms`: `MissingStrategy`: set to `SkipMissing()` or `ExactMissing()`to
    consciously care for missing values in `dv`.

An optional second arguments supports correlation between random variables.
- `corr::Symmetric`: correlation matrix, or
- `acf::AutoCorrelationFunction`: coefficients of the
   [`AutoCorrelationFunction`](@ref)

Keyword arguments:
- `isgapfilled::AbstractVector{Bool}`: set to true for records that should
   contribute to the sum but not to the decrease of relative uncertainty
   with increasing number of records, e.g. for missing records that have
   been estimated (gapfilled).

The sums of correlated variables require extra allocation and
support an additional keyword parameter
- `storage`: a mutable `AbstractVector{eltype(D)}` of length of `dv`
  that provides storage space to avoid additional allocations.

When implementing sum and mean for another Distribution `MyDist`,
the `AutoCorrelationFunction`-based method falls back to the `Symmetric`
correlation based method. Hence one only needs to define methods
```julia
sum(dv::AbstractDistributionVector{<:MyDist},
    ms::MissingStrategy=PassMissing())
sum(dv::AbstractDistributionVector{<:MyDist}, corr::Symmetric,
    ms::MissingStrategy=PassMissing())
```
"""
function sum(dv::AbstractDistributionVector, ms::MissingStrategy=PassMissing(); kwargs...)
    error("sum not defined yet for " *
    "Distributionvector{$(nonmissingtype(eltype(dv)))}")
end,
function sum(dv::AbstractDistributionVector, acf::AutoCorrelationFunction, ms::MissingStrategy=PassMissing(); kwargs...)
    corrmat = cormatrix_for_acf(length(dv), acf)
    sum(dv, Symmetric(corrmat), ms)
end,
function sum(dv::AbstractDistributionVector, corr::Symmetric, ms::MissingStrategy=PassMissing(); kwargs...)
    error("sum with correlations not defined yet for " *
    "Distributionvector{$(nonmissingtype(eltype(dv)))}")
end,
function mean(dv::AbstractDistributionVector, ms::MissingStrategy=PassMissing(); kwargs...)
    error("mean not defined yet for " *
    "Distributionvector{$(nonmissingtype(eltype(dv)))}")
end,
function mean(dv::AbstractDistributionVector, acf::AutoCorrelationFunction, ms::MissingStrategy=PassMissing(); kwargs...)
    corrmat = cormatrix_for_acf(length(dv), acf)
    mean(dv, Symmetric(corrmat), ms)
end,
function mean(dv::AbstractDistributionVector, corr::Symmetric, ms::MissingStrategy=PassMissing(); kwargs...)
    error("mean with correlations not defined yet for " *
    "Distributionvector{$(nonmissingtype(eltype(dv)))}")
end


text:
# Simply supported square plate with uniform distributed load
module ribbon_twist_examples

using FinEtools
using FinEtoolsDeforLinear
using FinEtoolsFlexStructures.FESetShellT3Module: FESetShellT3
using FinEtoolsFlexStructures.FESetShellQ4Module: FESetShellQ4
using FinEtoolsFlexStructures.FEMMShellT3FFModule
using FinEtoolsFlexStructures.RotUtilModule: initial_Rfield, update_rotation_field!
using VisualStructures: plot_nodes, plot_midline, render, plot_space_box, plot_midsurface, space_aspectratio, save_to_json

function _execute_dsg_model(formul, n = 2, LW_ratio = 2.0, visualize = true)
    E = 1e7;
    nu = 0.25;
    W = 1.0;
    L = LW_ratio * W
    thickness = 0.05
    analyt_sol = 26e-3 * LW_ratio / 10

    tolerance = thickness/1000
    fens, fes = T3block(L,W,n,1,:b);
    fens.xyz = xyz3(fens)

    mater = MatDeforElastIso(DeforModelRed3D, E, nu)

    # Report
    @info "Ribbon twist, L/W=$(LW_ratio), formulation=$(formul)"

    sfes = FESetShellT3()
    accepttodelegate(fes, sfes)
    femm = formul.make(IntegDomain(fes, TriRule(1), thickness), mater)
    associate = formul.associategeometry!
    stiffness = formul.stiffness

    # Construct the requisite fields, geometry and displacement
    # Initialize configuration variables
    geom0 = NodalField(fens.xyz)
    u0 = NodalField(zeros(size(fens.xyz,1), 3))
    Rfield0 = initial_Rfield(fens)
    dchi = NodalField(zeros(size(fens.xyz,1), 6))

    # Apply EBC's
    # Clamped end
    l1 = selectnode(fens; box = Float64[0 0 -Inf Inf -Inf Inf], inflate = tolerance)
    for i in [1,2,3,4,5,6]
        setebc!(dchi, l1, true, i)
    end

    applyebc!(dchi)
    numberdofs!(dchi);

    # Assemble the system matrix
    associategeometry!(femm, geom0)
    K = stiffness(femm, geom0, u0, Rfield0, dchi);

    # Load
    nl = selectnode(fens; box = Float64[L L 0 0 -Inf Inf], tolerance = tolerance)
    loadbdry = FESetP1(reshape(nl, 1, 1))
    lfemm = FEMMBase(IntegDomain(loadbdry, PointRule()))
    fi = ForceIntensity(FFlt[0, 0, -1, 0, 0, 0]);
    F = distribloads(lfemm, geom0, dchi, fi, 3);
    nl = selectnode(fens; box = Float64[L L W W -Inf Inf], tolerance = tolerance)
    loadbdry = FESetP1(reshape(nl, 1, 1))
    lfemm = FEMMBase(IntegDomain(loadbdry, PointRule()))
    fi = ForceIntensity(FFlt[0, 0, +1, 0, 0, 0]);
    F += distribloads(lfemm, geom0, dchi, fi, 3);

    # @infiltrate
    # Solve
    U = K\F
    scattersysvec!(dchi, U[:])
    targetu =  dchi.values[nl, 3][1]
    @info "Target: $(round(targetu, digits=8)),  $(round(targetu/analyt_sol, digits = 4)*100)%"

    # Visualization
    if !visualize
        return true
    end
    scattersysvec!(dchi, (L/4)/maximum(abs.(U)).*U)
    update_rotation_field!(Rfield0, dchi)
    plots = cat(plot_space_box([[0 0 -L/2]; [L/2 L/2 L/2]]),
        #plot_nodes(fens),
        plot_midsurface(fens, fes; x = geom0.values, u = dchi.values[:, 1:3], R = Rfield0.values);
    dims = 1)
    pl = render(plots)
    return true
end

function test_convergence()
    formul = FEMMShellT3FFModule
    LW_ratio = 2.0
    @info "Twisted Ribbon, L/W=$(LW_ratio), formulation=$(formul)"
    for n in [8, 12, 16, 24]
        _execute_dsg_model(formul, n, LW_ratio, false)
    end
    return true
end

function allrun()
    println("#####################################################")
    println("# test_convergence ")
    test_convergence()
    return true
end # function allrun

@info "All examples may be executed with "
println("using .$(@__MODULE__); $(@__MODULE__).allrun()")

end # module
nothing


text:
using MyExample2
using Test

import MyExample2: f, g
# using ForwardDiff

@testset "MyExample2.jl" begin
    # Write your tests here.
    @test f(2,0) == 4
    @test f(2,1) == 9
    @test g(2,3) == 6
    @test df(x,y) == 2
end


text:
export MtlCommandQueue

const MTLCommandQueue = Ptr{MtCommandQueue}

"""
    MtlCommandQueue(dev::MtlDevice)

A queue that organizes command buffers to be executed by the GPU `MtlDevice`.

A MTLCommandQueue object is used to queue an ordered list of command buffers for a
MTLDevice to execute. Command queues are thread-safe and allow multiple outstanding
command buffers to be encoded simultaneously.

[Metal Docs](https://developer.apple.com/documentation/metal/mtlcommandqueue?language=objc)
"""
mutable struct MtlCommandQueue
    handle::MTLCommandQueue
    device::MtlDevice
end

Base.convert(::Type{MTLCommandQueue}, q::MtlCommandQueue) = q.handle
Base.unsafe_convert(::Type{MTLCommandQueue}, q::MtlCommandQueue) = convert(MTLCommandQueue, q.handle)

Base.:(==)(a::MtlCommandQueue, b::MtlCommandQueue) = a.handle == b.handle
Base.hash(q::MtlCommandQueue, h::UInt) = hash(q.handle, h)

function MtlCommandQueue(dev::MtlDevice)
    queue = mtNewCommandQueue(dev)
    obj = MtlCommandQueue(queue, dev)
    finalizer(unsafe_destroy!, obj)
    return obj
end

function unsafe_destroy!(queue::MtlCommandQueue)
    if queue.handle !== C_NULL
        mtCommandQueueRelease(queue)
    end
end

function label(l::MtlCommandQueue)
    ptr = mtCommandQueueLabel(l)
    return ptr == C_NULL ? "" : unsafe_string(ptr)
end

function show(io::IO, ::MIME"text/plain", q::MtlCommandQueue)
    println(io, "MtlCommandQueue:")
    println(io, " handle  : ", q.handle)
    print(io, " device  : ", q.device)
end


text:
include("./params.jl")

function QQplot(rs, data)
    # Example usage:
    # x = 0.5-0.5*log(rand(1,1e3));
    # w = 0.001;
    # rs = sskernel(x, w);
    # QQplot(rs, x)
    #
    # Input argument
    # rs: Estimated intensity function
    # data: Spike time

    ################# WRITE SOMETHING #################

end

max_stars_repo_name: ganow/julia-point-process

max_issues_repo_name: ganow/julia-point-process

max_forks_repo_name: ganow/julia-point-process


text:
#=
  solve Upper Trianglar Matrices
  solve Lower Triangular Matrices
  Solve via Augmented Guassian Elimination
  Inverse via Augmented Guassian Elimination
  Solve via Inverse Linear Transformation Matrix
  Linear Least Squares Fitting
  Polynomial Least Squares Fitting
=#

function upTri(A, b)
  # accepts upper triangular matrix and vector b
  #TODO: use newtons

function lowTri(A)
  # accepts lower triangular matrix and vector b

function GuassSolve(A, b)
  # square matrix A, vector b
  return upTri(Guass([A b]))

function inverse(A)
  # try all
  r = remotecall(1, ,A)
  s = remotecall(2, ,A)
  t = remotecall(3, ,A)

function InverseSolve(A, b)
  return inverse(A) * b

function augGuassInverse()
  for i = 1:size(A)[1]

function solveviaCramers(A,b)
  B = copy(A)
  x =
  a = determinant(B)
  for i = 1:length(B)[1]
    C = copy(B)
    for j = 1:length(b)
      C[j,i] = b[j]
    end
    x[i] = determinant(C)
  end
  return a / x

function pseudo(A)
  B = transpose(A)
  try
    C = transpose(A) * (A * transpose(A))
    return C
  try
    C = (transpose(A) * A) * transpose(A)
    return C
  catch e
    println e
  end

function linearLS(A,b)
  # accepts 2 by n matrix
  T = [1 transpose(A[2,:])]
  # solve notal equations ATA = ATb
  return pseudo(T) * b

function polyLS(A)
  #TODO: complete

max_stars_repo_name: K-P-S/Julia-Scientific-Computing

max_issues_repo_name: K-P-S/Julia-Scientific-Computing

max_forks_repo_name: K-P-S/Julia-Scientific-Computing


text:
## Return in set strings of infinite sets
# Fallback
function JuMP.in_set_string(print_mode, set::AbstractInfiniteSet)::String
    return string(JuMP._math_symbol(print_mode, :in), " ", set)
end

# Extend to return of in set string for interval sets
function JuMP.in_set_string(print_mode, set::IntervalSet)::String
    if set.lower_bound != set.upper_bound
        return string(JuMP._math_symbol(print_mode, :in), " [",
                      JuMP._string_round(set.lower_bound), ", ",
                      JuMP._string_round(set.upper_bound), "]")
    else
        return string(JuMP._math_symbol(print_mode, :eq), " ",
                      JuMP._string_round(set.lower_bound))
    end
end

# Extend to return of in set string for distribution sets
function JuMP.in_set_string(print_mode, set::DistributionSet)::String
    # get distribution string
    d_string = string(set.distribution)
    # remove number type
    first_bracket = findfirst(isequal('{'), d_string)
    if first_bracket != nothing
        last_bracket = findall(isequal('}'), d_string)[end]
        d_string = d_string[1:first_bracket-1] * d_string[last_bracket+1:end]
    end
    # remove new line characters
    new_lines = findall(isequal('\n'), d_string)
    for i = 1:length(new_lines)
        if new_lines[1] == length(d_string)
            d_string = d_string[1:end-1]
        elseif d_string[new_lines[1]-1] == '(' || d_string[new_lines[1]+1] == ')'
            d_string = d_string[1:new_lines[1]-1] * d_string[new_lines[1]+1:end]
        else
            d_string = d_string[1:new_lines[1]-1] * ", " *
                       d_string[new_lines[1]+1:end]
        end
        new_lines = findall(isequal('\n'), d_string)
    end
    return string(JuMP._math_symbol(print_mode, :in), " ", d_string)
end

# Return contraint parameter bounds as a string
function bound_string(print_mode, bounds::ParameterBounds)::String
    string_list = ""
    for (pref, set) in bounds.intervals
        string_list *= string(JuMP.function_string(print_mode, pref), " ",
                              JuMP.in_set_string(print_mode, set), ", ")
    end
    return string_list[1:end-2]
end

# Return string of a bounded constraint
function JuMP.constraint_string(print_mode,
                                constraint_object::BoundedScalarConstraint)
    # get needed strings
    func_str = JuMP.function_string(print_mode, constraint_object)
    in_set_str = JuMP.in_set_string(print_mode, constraint_object)
    bound_str = bound_string(print_mode, constraint_object.bounds)
    # modify as needed
    if print_mode == JuMP.REPLMode
        lines = split(func_str, '\n')
        lines[1 + div(length(lines), 2)] *= " " * in_set_str * ", " *
                                       JuMP._math_symbol(print_mode, :for_all) *
                                       " " * bound_str
        return join(lines, '\n')
    else
        return func_str * " " * in_set_str * ", " *
               JuMP._math_symbol(print_mode, :for_all) * " " * bound_str
    end
end

# Return a string for a constraint reference
function JuMP.constraint_string(print_mode, ref::GeneralConstraintRef)
    return JuMP.constraint_string(print_mode, JuMP.name(ref),
           JuMP.constraint_object(ref))
end

# return list of constraint
function JuMP.constraints_string(print_mode, model::InfiniteModel)
    strings = String[]
    # Sort by creation order
    constraints = sort(collect(model.constrs), by = c -> c.first)
    # produce a string for each constraint
    for (index, constraint) in constraints
        push!(strings, JuMP.constraint_string(print_mode, constraint))
    end
    # produce the parameters that are used along with their sets
    params = sort(collect(model.params), by = c -> c.first)
    ignore_groups = [] # used to indicate which groups should no longer be used
    for (index, param) in params
        pref = ParameterRef(model, index)
        if is_used(pref) && !(group_id(pref) in ignore_groups)
            if isa(infinite_set(pref), DistributionSet{<:Distributions.MultivariateDistribution})
                # print multivariate distributions with variables grouped together
                push!(ignore_groups, group_id(pref))
                push!(strings, string(_root_name(pref), " ",
                                JuMP.in_set_string(print_mode, param.set)))
            else
                # print parameter individually
                push!(strings, string(JuMP.function_string(print_mode, pref), " ",
                                      JuMP.in_set_string(print_mode, param.set)))
            end
        end
    end
    return strings
end

# Return the objective string
function JuMP.objective_function_string(print_mode, model::InfiniteModel)
    return JuMP.function_string(print_mode, JuMP.objective_function(model))
end

# Show IntervalSets in REPLMode
function Base.show(io::IO, set::IntervalSet)
    print(io, "[", JuMP._string_round(set.lower_bound), ", ",
          JuMP._string_round(set.upper_bound), "]")
    return
end

# Show IntervalSets in IJuliaMode
function Base.show(io::IO, ::MIME"text/latex", set::IntervalSet)
    print(io, "[", JuMP._string_round(set.lower_bound), ", ",
          JuMP._string_round(set.upper_bound), "]")
end

# Show IntervalSets in REPLMode
function Base.show(io::IO, set::DistributionSet)
    print(io, set.distribution)
    return
end

# Show IntervalSets in IJuliaMode
function Base.show(io::IO, ::MIME"text/latex", set::DistributionSet)
    print(io, set.distribution)
end

# TODO make list following group_id and use ... when necessary --> Subdomain bounds: t in [0, 1], x[1] == 0, x[2] == 0, ... x[6] == 0, a in [2, 3]
# Show ParameterBounds in REPLMode
function Base.show(io::IO, bounds::ParameterBounds)
    print(io, "Subdomain bounds (", length(bounds), "): ",
          bound_string(JuMP.REPLMode, bounds))
    return
end

# Show ParameterBounds in IJuliaMode
function Base.show(io::IO, ::MIME"text/latex", bounds::ParameterBounds)
    print(io, "Subdomain bounds (", length(bounds), "): ",
          bound_string(JuMP.IJuliaMode, bounds))
end

# TODO show hold variables better

# Show constraint in REPLMode
function Base.show(io::IO, ref::GeneralConstraintRef)
    print(io, JuMP.constraint_string(JuMP.REPLMode, ref))
    return
end

# Show constraint in IJuliaMode
function Base.show(io::IO, ::MIME"text/latex", ref::GeneralConstraintRef)
    print(io, JuMP.constraint_string(JuMP.IJuliaMode, ref))
end

# Show the backend information associated with the optimizer model
function JuMP.show_backend_summary(io::IO, model::InfiniteModel)
    println(io, "Optimizer model backend information: ")
    JuMP.show_backend_summary(io, optimizer_model(model))
    return
end

# Show the objective function type
function JuMP.show_objective_function_summary(io::IO, model::InfiniteModel)
    println(io, "Objective function type: ",
            JuMP.objective_function_type(model))
    return
end

# Return "s" if n is greater than one
_plural(n) = (isone(n) ? "" : "s")

# Return constraint summary in JuMP like manner
function JuMP.show_constraints_summary(io::IO, model::InfiniteModel)
    for (F, S) in JuMP.list_of_constraint_types(model)
        n_constraints = JuMP.num_constraints(model, F, S)
        println(io, "`$F`-in-`$S`: $n_constraints constraint",
                _plural(n_constraints))
    end
    return
end

# Have the infinitemodel object show similar to Model
function Base.show(io::IO, model::InfiniteModel)
    # indicate is an infinite model
    println(io, "An InfiniteOpt Model")
    # show objective ssnse info
    sense = JuMP.objective_sense(model)
    if sense == MOI.MAX_SENSE
        print(io, "Maximization")
    elseif sense == MOI.MIN_SENSE
        print(io, "Minimization")
    else
        print(io, "Feasibility")
    end
    println(io, " problem with:")
    # show variable info
    println(io, "Variable", _plural(JuMP.num_variables(model)), ": ",
            JuMP.num_variables(model))
    # show obejctive function info
    if sense != MOI.FEASIBILITY_SENSE
        JuMP.show_objective_function_summary(io, model)
    end
    # show constraint info
    JuMP.show_constraints_summary(io, model)
    # show other info
    names_in_scope = sort(collect(keys(JuMP.object_dictionary(model))))
    if !isempty(names_in_scope)
        print(io, "Names registered in the model: ",
              join(string.(names_in_scope), ", "), "\n")
    end
    JuMP.show_backend_summary(io, model)
end


text:
@testset "Test single shooting on logistic map" begin
    # Define logistic map
    function logistic_f(x_next, x, k, θ)
        x_next .=  θ*x - θ*x.^2
        return
    end
    # g(x) = x
    function logistic_g(y, x, k, θ)
        y .= x
        return
    end
    # Define initial conditions
    x0 = [0.5]
    # Initialize buffer
    y = Vector{Vector{Float64}}(100)
    for i = 1:length(y)
        y[i] = zeros(1)
    end
    y2 = Vector{Vector{Float64}}(100)
    for i = 1:length(y)
        y2[i] = zeros(1)
    end
    # Simulate
    x_final = ms.simulate_space_state!(y, logistic_f, logistic_g,
                                       [0.5], (1, 100), (3.78,))
    x_final2 = ms.simulate_space_state!(y2, logistic_f, logistic_g,
                                       [0.5], (1, 100), (3.62,))

    # Define initial values
    k0 = 1
    Ny = 1
    N = 100
    x0 = [0.5]
    θ = [3.78, 3.78, 1.00]
    θ2 = [3.62, 3.62, 1.00]
    # Define polynomial model
    # f(x) = θ1*x - θ2*x^2
    function f(y, dx, dθ, x, k, θ)
        y .= θ[1]*x - θ[2]*x.^2
        dx .= θ[1] - 2*θ[2]*x
        dθ .=  [ x  -x.^2  0]
        return
    end
    # g(x) = x
    function g(y, dx, dθ, x, k, θ)
        y .= θ[3]*x
        dx .= θ[3]
        dθ .=  [ 0 0 x ]
        return
    end

    oss = ms.OneShootSimulation(f, g, x0, y, k0, θ)

    @testset "Test function evaluation" begin
        @test vcat(y...) - vcat(oss.ys) ≈ zeros(100)
        @test oss.x ≈ x_final

        ms.new_simulation!(oss, x0, θ2)

        @test vcat(y2...) - vcat(oss.ys) ≈ zeros(100)
        @test oss.x ≈ x_final2
    end

    @testset "Test jacobian" begin
        function wrapper_simulation(ϕ)
            θ1 = ϕ[1:3]
            x01 = [ϕ[4]]
            ms.new_simulation!(oss, x01, θ1)
            return vcat(oss.ys...)
        end

        function wrapper_jacobian(ϕ)
           θ1 = ϕ[1:3]
           x01 = [ϕ[4]]
           ms.new_simulation!(oss, x01, θ1)
           return hcat(vcat(oss.dydθ...), vcat(oss.dydx0...))
        end

        function finite_difference_jacobian(ϕ)
            return jacobian(wrapper_simulation, ϕ, :central)
        end

       # OBS: finite difference approximation
       # fails when trying parameters in the
       # chaotic region
       for ϕ in ([3.21, 3.2, 1.00, 0.5],
                 [3.21, 1.2, 1.54, 0.5],
                 [2.56, 3.27, 1.54, 0.33])
            @test wrapper_jacobian(ϕ) ≈ finite_difference_jacobian(ϕ)
       end
    end

    @testset "Test cost function" begin
        ms.new_simulation!(oss, x0, θ)
        @test ms.cost_function(oss) ≈ 0
    end

    @testset "Test gradient" begin
        function wrapper_cost(ϕ)
            θ1 = ϕ[1:3]
            x01 = [ϕ[4]]
            ms.new_simulation!(oss, x01, θ1)
            return ms.cost_function(oss)
        end

        function wrapper_gradient(ϕ)
            θ1 = ϕ[1:3]
            x01 = [ϕ[4]]
            ms.new_simulation!(oss, x01, θ1)
            gradθ = zeros(3)
            ms.gradient!(gradθ, oss, "θ")
            gradx0 = zeros(1)
            ms.gradient!(gradx0, oss, "x0")
            return vcat(gradθ, gradx0)
        end

        function finite_difference_gradient(ϕ)
            return Calculus.gradient(wrapper_cost, ϕ)
        end

        # OBS: finite difference approximation
        # fails when trying parameters in the
        # chaotic region
        for ϕ in ([3.21, 3.2, 1.00, 0.5],
                 [3.21, 1.2, 1.54, 0.5],
                 [2.56, 3.27, 1.54, 0.33])
            @test wrapper_gradient(ϕ) ≈ finite_difference_gradient(ϕ)
        end
    end
end

@testset "Test single shooting on linear problem" begin
    # From example 2.11, "Linear System Theory and design" Chen
    # With not very realistic values
    # Define linear function
    function linear_f(x_next, x, k)
        R = 20
        C1 = 10
        C2 = 10
        L = 1
        A = [-1/(R*C1) 0   -1/C1;
             0         0    1/C2;
             1/L       -1/L    0]
        x_next .=  A*x
        return
    end
    # g(x) = x
    function linear_g(y, x, k)
      y .= dot([1, -1, 0], x)
      return
    end
    # Define initial conditions
    x0 = [5, 0.6, 0.5]
    # Initialize buffer
    y = Vector{Vector{Float64}}(10)
    for i = 1:length(y)
       y[i] = zeros(1)
    end
    # Simulate
    x_final = ms.simulate_space_state!(y, linear_f, linear_g,
                                       x0, (1, 10))

    # Define initial values
    k0 = 1
    Ny = 1
    N = 10
    x0 = [5, 0.6, 0.5]
    θ = [-0.005, 0.1, 0.1, 1]
    # Define polynomial model
    # f(x) = θ1*x - θ2*x^2
    function f(y, dx, dθ, x, k, θ)
        A = [θ[1]  0 -θ[2];
             0     0  θ[3];
             θ[4] -θ[4]  0]
         y .= A*x
         dx .= A
         dθ .= [x[1] -x[3]    0      0;
                0      0      x[3]   0;
                0      0       0    x[1]-x[2]]
         return
    end
    # g(x) = x
    function g(y, dx, dθ, x, k, θ)
      y .= dot([1, -1, 0], x)
      dx .= [1 -1 0]
      dθ .=  0
      return
    end

    oss = ms.OneShootSimulation(f, g, x0, y, k0, θ)

    @testset "Test function evaluation" begin
       @test vcat(y...) - vcat(oss.ys) ≈ zeros(10)
       @test oss.x ≈ x_final
    end

    @testset "Test jacobian" begin
        function wrapper_simulation(ϕ)
           θ1 = ϕ[1:4]
           x01 = ϕ[5:7]
           ms.new_simulation!(oss, x01, θ1)
           return vcat(oss.ys...)
        end

        function wrapper_jacobian(ϕ)
           θ1 = ϕ[1:4]
           x01 = ϕ[5:7]
           ms.new_simulation!(oss, x01, θ1)
           return hcat(vcat(oss.dydθ...), vcat(oss.dydx0...))
        end

        function finite_difference_jacobian(ϕ)
            return jacobian(wrapper_simulation, ϕ, :central)
        end

        for ϕ in ([-0.005, 0.1, 0.1, 1, 5, 0.6, 0.5],
                  [-0.03, 0.11, 0.1, 2, 5, 0, 0.5],
                  [-0.5, 1.11, 1, 2, 2, 10, 0.5])
            @test wrapper_jacobian(ϕ) ≈ finite_difference_jacobian(ϕ)
        end
    end

    @testset "Test cost function" begin
       ms.new_simulation!(oss, x0, θ)
       @test ms.cost_function(oss) ≈ 0
    end

    function wrapper_cost(ϕ)
       θ1 = ϕ[1:4]
       x01 = ϕ[5:7]
       ms.new_simulation!(oss, x01, θ1)
       return ms.cost_function(oss)
    end

    @testset "Test gradient" begin
       function wrapper_gradient(ϕ)
          θ1 = ϕ[1:4]
          x01 = ϕ[5:7]
          ms.new_simulation!(oss, x01, θ1)
          gradθ = zeros(4)
          ms.gradient!(gradθ, oss, "θ")
          gradx0 = zeros(3)
          ms.gradient!(gradx0, oss, "x0")
          return vcat(gradθ, gradx0)
       end
       function finite_difference_gradient(ϕ)
          return Calculus.gradient(wrapper_cost, ϕ)
       end
       for ϕ in ([-0.005, 0.1, 0.1, 1, 5, 0.6, 0.5],
                 [-0.03, 0.11, 0.1, 2, 5, 0, 0.5],
                 [-0.5, 1.11, 1, 2, 2, 10, 0.5])
          @test wrapper_gradient(ϕ) ≈ finite_difference_gradient(ϕ) atol=1e-9
       end
    end
end


text:
using SeaPearl
using SeaPearlExtras
using ReinforcementLearning
const RL = ReinforcementLearning
using Flux
using GeometricFlux
using JSON
using BSON: @save, @load
using Dates
using Random

# -------------------
# Experience variables
# -------------------
nbEpisodes = 100
restartPerInstances = 20
evalFreq = 50
nbInstances = 10
nbRandomHeuristics = 1

nbNodes = 20
nbMinColor = 5
density = 0.95
# -------------------
# Generator
# -------------------
coloring_generator = SeaPearl.ClusterizedGraphColoringGenerator(nbNodes, nbMinColor, density)

#include("rewards.jl")
include("features.jl")

# -------------------
# Internal variables
# -------------------
SR = SeaPearl.DefaultStateRepresentation{BetterFeaturization, SeaPearl.DefaultTrajectoryState}
numInFeatures = SeaPearl.feature_length(SR)
#numGlobalFeature = SeaPearl.global_feature_length(SR)

# -------------------
# Agent definition
# -------------------
include("agents.jl")

# -------------------
# Value Heuristic definition
# -------------------

learnedHeuristic=SeaPearl.LearnedHeuristic{SR, SeaPearl.CPReward, SeaPearl.FixedOutput}(agent)
# Basic value-selection heuristic
selectMin(x::SeaPearl.IntVar; cpmodel=nothing) = SeaPearl.minimum(x.domain)
heuristic_min = SeaPearl.BasicHeuristic(selectMin)
function select_random_value(x::SeaPearl.IntVar; cpmodel=nothing)
    selected_number = rand(1:length(x.domain))
    i = 1
    for value in x.domain
        if i == selected_number
            return value
        end
        i += 1
    end
    @assert false "This should not happen"
end

randomHeuristics = []
for i in 1:nbRandomHeuristics
    push!(randomHeuristics, SeaPearl.BasicHeuristic(select_random_value))
end

valueSelectionArray = [learnedHeuristic, heuristic_min]
append!(valueSelectionArray, randomHeuristics)
# -------------------
# Variable Heuristic definition
# -------------------
variableSelection = SeaPearl.MinDomainVariableSelection{false}()

# -------------------
# -------------------
# Core function
# -------------------
# -------------------

function trytrain(nbEpisodes::Int)
    experienceTime = now()
    dir = mkdir(string("exp_",Base.replace("$(round(experienceTime, Dates.Second(3)))",":"=>"-")))
    expParameters = Dict(
        :nbEpisodes => nbEpisodes,
        :restartPerInstances => restartPerInstances,
        :evalFreq => evalFreq,
        :nbInstances => nbInstances,
        :nbRandomHeuristics => nbRandomHeuristics,
        nbNodes => nbNodes,
        nbMinColor => nbMinColor,
        density => density
    )
    open(dir*"/params.json", "w") do file
        JSON.print(file, expParameters)
    end

    metricsArray, eval_metricsArray = SeaPearl.train!(
        valueSelectionArray=valueSelectionArray,
        generator=coloring_generator,
        nbEpisodes=nbEpisodes,
        strategy=SeaPearl.ILDSearch(0),
        variableHeuristic=variableSelection,
        out_solver=true,
        verbose = false,
        evaluator=SeaPearl.SameInstancesEvaluator(valueSelectionArray,coloring_generator; evalFreq = evalFreq, nbInstances = nbInstances),
        restartPerInstances = restartPerInstances
    )

    #saving model weights
    model = agent.policy.learner.approximator
    @save dir*"/model_gc"*string(coloring_generator.n)*".bson" model

    SeaPearlExtras.storedata(metricsArray[1]; filename=dir*"/graph_coloring_$(nbNodes)_training_learned")
    SeaPearlExtras.storedata(metricsArray[2]; filename=dir*"/graph_coloring_$(nbNodes)_training_greedy")
    for i = 1:nbRandomHeuristics
        SeaPearlExtras.storedata(metricsArray[2+i]; filename=dir*"/graph_coloring_$(nbNodes)_training_random$(i)")
    end
    SeaPearlExtras.storedata(eval_metricsArray[:,1]; filename=dir*"/graph_coloring_$(nbNodes)_learned")
    SeaPearlExtras.storedata(eval_metricsArray[:,2]; filename=dir*"/graph_coloring_$(nbNodes)_greedy")
    for i = 1:nbRandomHeuristics
        SeaPearlExtras.storedata(eval_metricsArray[:,i+2]; filename=dir*"/graph_coloring_$(nbNodes)_random$(i)")
    end

    return metricsArray, eval_metricsArray
end

metricsArray, eval_metricsArray = trytrain(nbEpisodes)
nothing

max_stars_repo_name: corail-research/SeaPearlZoo

max_issues_repo_name: corail-research/SeaPearlZoo

max_forks_repo_name: corail-research/SeaPearlZoo


text:
export
    DIRCOLSolver

"$(TYPEDEF) Solver options for the Direct Collocation solver.
Most options are passed to the NLP through the `opts` dictionary
$(FIELDS)"
@with_kw mutable struct DIRCOLSolverOptions{T} <: DirectSolverOptions{T}
    "NLP Solver to use. See MathOptInterface for available NLP solvers"
    nlp::MathOptInterface.AbstractOptimizer = Ipopt.Optimizer()

    "Options dictionary for the nlp solver"
    opts::Dict{Symbol,Any} = Dict{Symbol,Any}()

    "Print output to console"
    verbose::Bool = true

    "Feasibility tolerance"
    constraint_tolerance::T = -1.0
end

function DIRCOLSolverOptions(opts::SolverOptions)
	DIRCOLSolverOptions(
		constraint_tolerance = opts.constraint_tolerance,
		verbose=verbose
	)
end

"""
$(TYPEDEF)
Direct Collocation Solver.
Uses a commerical NLP solver to solve the Trajectory Optimization problem.
Uses the MathOptInterface to interface with the NLP.
"""
struct DIRCOLSolver{Q<:QuadratureRule,L,T,N,M,NM} <: DirectSolver{T}
    opts::SolverOptions{T}
    stats::Dict{Symbol,Any}

    NN::Int
    NP::Int

    dyn_con::DynamicsConstraint{Q,L,T,N,M,NM}
    objective::AbstractObjective
    constraints::ConstraintSet{T}
    constraints_all::ConstraintSet{T}
    Z::Vector{KnotPoint{T,N,M,NM}}
    x0::SVector{N,T}

    optimizer::MOI.AbstractOptimizer

    E::Vector{SizedCostExpansion{T,N,N,M}}

    xinds::Vector{SVector{N,Int}}
    uinds::Vector{SVector{M,Int}}

    linds::Vector{<:Vector{SV} where SV}
    con_inds::Vector{<:Vector{SV} where SV}

    jacobian_structure::Symbol

    function DIRCOLSolver(opts,stats, NN,NP,dyn_con::DynamicsConstraint{Q,L,T,N},
            obj,conSet,conSet_all,Z::Vector{KnotPoint{T,N,M,NM}},
            x0, optimizer, E,
            xinds,uinds,linds,con_inds,jac_structure) where {Q,L,T,N,M,NM}
        new{Q,L,T,N,M,NM}(opts,stats, NN,NP,dyn_con,obj,conSet,conSet_all,Z,x0,optimizer,E,
            xinds,uinds,linds,con_inds,jac_structure)
    end

end

Base.size(solver::DIRCOLSolver{Q,L,T,n,m,NM}) where {Q,L,T,n,m,NM} = n,m,length(solver.Z)

function DIRCOLSolver(prob::Problem{Q},
        opts=SolverOptions(),
        jacobian_structure=:by_knotpoint,
    	opts_optimizer::Dict{Symbol,Any} = Dict{Symbol,Any}(),
		nlp=Ipopt.Optimizer(); integration=Q) where Q

    n,m,N = size(prob)
    Z = prob.Z

    stats = Dict{Symbol,Any}()

    # Add dynamics constraints
    prob = copy(prob)
    add_dynamics_constraints!(prob, integration)
    dyn_con = get_constraints(prob).constraints[2].con::DynamicsConstraint{integration}
    conSet = get_constraints(prob)

    # Store a copy of the constraint set with all constraints
    conSet_all = copy(conSet)

    # Add bounds at infinity if the problem doesn't have any bound constraints
    if !any(is_bound.(conSet))
        bnd = BoundConstraint(n,m)
        add_constraint!(conSet, ConstraintVals(bnd, 1:N))
    end

    # Remove bounds
    zU,zL,gU,gL = get_bounds(conSet)

    # Initialize arrays
    dyn_vals = DynamicsVals(dyn_con)
	E = [SizedCostExpansion{Float64}(n,m) for k = 1:N]

    NN = (n+m)*N
    NP = sum(num_constraints(prob))

    P = Primals(n,m,N)
    xinds, uinds = P.xinds, P.uinds

    blk_len = map(con->length(con.∇c[1]), conSet.constraints)
    con_len = map(con->length(con.∇c), conSet.constraints)
    con_inds = gen_con_inds(get_constraints(prob), jacobian_structure)

    # These get filled in with call to constraint_jacobian_structure
    linds = [[@SVector zeros(Int,blk) for i = 1:len] for (blk,len) in zip(blk_len, con_len)]

    # Create MOI Optimizer
    nlp_opts = Dict(Symbol(key)=>value for (key,val) in pairs(opts_optimizer))
    optimizer = typeof(nlp)(;nlp_opts..., nlp_options(nlp, opts)...)

    # Create Solver
    d = DIRCOLSolver(opts, stats, NN, NP, dyn_con, prob.obj, conSet, conSet_all,
        Z, prob.x0, optimizer, E, xinds, uinds, linds, con_inds, jacobian_structure)

    # Set up MOI problem
    V0 = zeros(NN)
    copyto!(V0, Z, xinds, uinds)

    has_objective = true
    nlp_bounds = MOI.NLPBoundsPair.(gL, gU)
    block_data = MOI.NLPBlockData(nlp_bounds, d, has_objective)

    V = MOI.add_variables(optimizer, NN)

    MOI.add_constraints(optimizer, V, MOI.LessThan.(zU))
    MOI.add_constraints(optimizer, V, MOI.GreaterThan.(zL))

    MOI.set(optimizer, MOI.VariablePrimalStart(), V, V0)

    # Set up NLP problem
    MOI.set(optimizer, MOI.NLPBlock(), block_data)
    MOI.set(optimizer, MOI.ObjectiveSense(), MOI.MIN_SENSE)

    return d
end

@inline AbstractSolver(prob::Problem, opts::DIRCOLSolverOptions) = DIRCOLSolver(prob, opts)

# AbstractSolver Interface
get_initial_state(solver::DIRCOLSolver) = solver.x0
get_model(solver::DIRCOLSolver) = solver.dyn_con.model
get_constraints(solver::DIRCOLSolver) = solver.constraints
get_trajectory(solver::DIRCOLSolver) = solver.Z
get_objective(solver::DIRCOLSolver) = solver.objective
iterations(solver::DIRCOLSolver) = stats(solver)[:iterations]

num_primals(solver::DIRCOLSolver) = solver.NN
num_duals(solver::DIRCOLSolver) =  solver.NP

# Include bounds when calculating max violation on the solver
function max_violation(solver::DIRCOLSolver)
	Z = get_trajectory(solver)
    conSet = solver.constraints_all
    max_violation(conSet, Z)
end

primal_partition(prob::DIRCOLSolver) = prob.xinds, prob.uinds
jacobian_linear_inds(prob::DIRCOLSolver) = prob.linds

@inline copy_gradient!(grad_f, solver::DIRCOLSolver) = copy_gradient!(grad_f, solver.E, solver.xinds, solver.uinds)
@inline Base.copyto!(d::DIRCOLSolver, V::Vector{<:Real}) = copyto!(d.Z, V, primal_partition(d)...)

function initial_controls!(d::DIRCOLSolver, U0)
    Z = get_trajectory(d)
    set_controls!(Z, U0)
    V = [MOI.VariableIndex(i) for i = 1:d.NN]
    V0 = zeros(d.NN)
    copyto!(V0, Z, d.xinds, d.uinds)
    MOI.set(d.optimizer, MOI.VariablePrimalStart(), V, V0)
end

function initial_states!(d::DIRCOLSolver, X0)
    Z = get_trajectory(d)
    set_states!(Z, X0)
    V = [MOI.VariableIndex(i) for i = 1:d.NN]
    V0 = zeros(d.NN)
    copyto!(V0, Z, d.xinds, d.uinds)
    MOI.set(d.optimizer, MOI.VariablePrimalStart(), V, V0)
end

function get_rc(A::SparseMatrixCSC)
    row,col,inds = findnz(A)
    v = sortperm(inds)
    row[v],col[v]
end

function cost(solver::DIRCOLSolver{Q}) where Q<:QuadratureRule
	Z = get_trajectory(solver)
	obj = get_objective(solver)
	cost(obj, solver.dyn_con, Z)
end

function cost_gradient!(solver::DIRCOLSolver{Q}) where Q
	obj = get_objective(solver)
	Z = get_trajectory(solver)
	dyn_con = solver.dyn_con
	E = solver.E
	cost_gradient!(E, obj, dyn_con, Z)
end

# Define MOI Interface
MOI.features_available(d::DIRCOLSolver) = [:Grad, :Jac]
MOI.initialize(d::DIRCOLSolver, features) = nothing

function MOI.jacobian_structure(d::DIRCOLSolver)
    jac_struct = constraint_jacobian_structure(d, d.jacobian_structure)
    r,c = get_rc(jac_struct)
    jac_struct = collect(zip(r,c))
end

MOI.hessian_lagrangian_structure(d::DIRCOLSolver) = []

function MOI.eval_objective(d::DIRCOLSolver, V)
    copyto!(d, V)
    cost(d)
end

function MOI.eval_objective_gradient(d::DIRCOLSolver, grad_f, V)
    copyto!(d, V)
    cost_gradient!(d)
    copy_gradient!(grad_f, d)
    return nothing
end

function MOI.eval_constraint(d::DIRCOLSolver, g, V)
    copyto!(d, V)
    update_constraints!(d)
    copy_constraints!(g, d)
end

function MOI.eval_constraint_jacobian(d::DIRCOLSolver, jac, V)
    copyto!(d, V)
    constraint_jacobian!(d)
    copy_jacobians!(jac, d)
end

MOI.eval_hessian_lagrangian(::DIRCOLSolver, H, x, σ, μ) = nothing

function solve!(d::DIRCOLSolver)
    # Update options
    nlp_opts = nlp_options(d.optimizer, d.opts)
    for (key,val) in nlp_opts
        d.optimizer.options[String(key)] = val
    end

    # Solve with MOI
    MOI.optimize!(d.optimizer)

    # Get result and copy to trajectory
    V = [MOI.VariableIndex(k) for k = 1:d.NN]
    res = MOI.get(d.optimizer, MOI.VariablePrimal(), V)
    copyto!(d, res)

    # Copy stats
    for (key,val) in parse_ipopt_summary()
        d.stats[key] = val
    end
    return nothing
end

function nlp_options(nlp::MOI.AbstractOptimizer, opts::SolverOptions)
    solver_name = optimizer_name(nlp)
	opts_optimizer = Dict{Symbol,Any}()
    if solver_name == :Ipopt
        !opts.verbose ? opts_optimizer[:print_level] = 0 : opts_optimizer[:print_level] = 5
        if opts.constraint_tolerance > 0.
            opts_optimizer[:constr_viol_tol] = opts.constraint_tolerance
            opts_optimizer[:tol] = opts.constraint_tolerance
        end
    elseif solver_name == :SNOPT7
        if !opts.verbose
            opts_optimizer[:Major_print_level] = 0
            opts_optimizer[:Minor_print_level] = 0
        end
        if opts.constraint_tolerance > 0.
            opts_optimizer[:Major_constraint_tolerance] = opts.constraint_tolerance
            opts_optimizer[:Minor_constraint_tolerance] = opts.constraint_tolerance
            opts_optimizer[:Major_optimality_tolerance] = opts.constraint_tolerance
        end
    else
        error("Nonlinear solver not implemented")
    end

    return opts_optimizer
end

function optimizer_name(optimizer::MathOptInterface.AbstractOptimizer)
    nameof(parentmodule(typeof(optimizer)))
end


text:
@doc """ Bi-directional LSTM for language modeling, training procedure:
         Suppose an input sequence of x1-x2-x3,
         The backward lstm gets input as </s>-x3-x2,
         The forward lstm gets input as <s>-x1-x2 """

"""
parameters: holds the weight and bias for kth lstm layer as parameters[2k-1,2k]: weight and bias.
nlayers: Number of stacked lstm layers
states: Holds the cell and hidden states as state[2k-1,2k]: hidden and cell
"""
type LSTM
    parameters
    nlayers::Int
    states::Array
    atype::DataType

    function LSTM(atype, layerconfig, embedsize, batchsize, winit)
        nlayers = length(layerconfig)
        parameters = initweights(atype, layerconfig, embedsize, winit)
        states = initstate(atype, layerconfig, batchsize)
        new(parameters, nlayers, states, atype)
    end
end

"""
Initializes the weights and biases for the stacked LSTM.
atype: Array type of the machine, namely KnetArray or Array regarding of gpu usage
hidden: Array of hidden layer configurations, e.g, [128 256] 2 layers of lstm.
embed: Embedding size of the model
winit: initialization parameter of the model
OUTPUT:
param[2k-1, 2k]: weight and bias for the kth lstm layer
"""
function initweights(atype, hidden, embed, winit)
    param = Array(Any, 2*length(hidden))
    input = embed
    for k = 1:length(hidden)
        param[2k-1] = winit*randn(input+hidden[k], 4*hidden[k])
        param[2k]   = zeros(1, 4*hidden[k])
        param[2k][1:hidden[k]] = 1 # forget gate bias
        input = hidden[k]
    end
    return map(p->convert(atype,p), param)
end

"""
state[2k-1]: hidden, state[2k]: cell
"""
function initstate(atype, hidden, batchsize)
    state = Array(Any, 2*length(hidden))
    for k = 1:length(hidden)
        state[2k-1] = zeros(batchsize,hidden[k])
        state[2k] = zeros(batchsize,hidden[k])
    end
    return map(s->convert(atype,s), state)
end

"""
Embedding matrix initialization it is of the size (vocab x embed)
The input to lstm units would be x * W where W is embedding and x is traditional one-hots.
"""
function initembedding(atype, embedsize, vocabsize, winit)
    embedding = winit * randn(vocabsize, embedsize)
    return convert(atype, embedding)
end

"""
Softmax layer initialization,
param[1]: softmax layer weights (hidden[final], vocab)
param[2]: softmax layer bias (1, vocab)
"""
function initsoftmax(atype, hiddenfinal, vocabsize, winit)
    param = Array(Any, 2)
    param[1] = winit * randn(hiddenfinal, vocabsize)
    param[2] = zeros(1, vocabsize)
    return map(p->convert(atype, p), param)
end

"""
forward implementation for a single layer lstm
"""
function lstm(weight, bias, hidden, cell, input)
    gates   = hcat(input,hidden) * weight .+ bias
    hsize   = size(hidden,2)
    forget  = sigm(gates[:,1:hsize])
    ingate  = sigm(gates[:,1+hsize:2hsize])
    outgate = sigm(gates[:,1+2hsize:3hsize])
    change  = tanh(gates[:,1+3hsize:end])
    cell    = cell .* forget + ingate .* change
    hidden  = outgate .* tanh(cell)
    return (hidden,cell)
end

"""
Go one step forward through the stacked LSTMs, then return the final hidden state of the final lstm layer for a given input.
parameters[2k-1] =  weight, parameters[2k] = bias for the kth layer
states[2k-1] = the hidden and states[2k] = the cell for the kth layer
"""
function sforw(parameters, states, input)
    x = input
    for i=1:2:length(states)
        (states[i], states[i+1]) = lstm(parameters[i], parameters[i+1], states[i], states[i+1], x)
        x = states[i]
    end
    return x
end

"""
Go forward and collect hidden states of the lstm for a sequence
sequence is already padded from begining with token <s> and end with token </s>
OUTPUT
hiddenlayers: contains the final hidden states of the forward and backward lstms in correct order,
"""
function bilsforw(parameters, states, sequence, atype; forwardlstm=true)
    hiddenlayers = Array(Any, length(sequence))

    # if forward first item else the last item will be zeros of the same size with other items
    traverse = (forwardlstm ? (1:length(sequence) - 1) : (length(sequence):-1:2))

    for i=traverse
        result = sforw(parameters, states, sequence[i])
        if forwardlstm
            hiddenlayers[i+1] = result
        else
            hiddenlayers[i-1] = result
        end
    end
    padindex = (forwardlstm ? 1 : length(hiddenlayers))
    padding = convert(atype, zeros(eltype(hiddenlayers[3]), size(hiddenlayers[3])))
    hiddenlayers[padindex] = padding
    return hiddenlayers
end

"""
Embeds the one-hot given sequence into dense sequence.
Inputs: Sequence of one-hots, i.e., sequence[k] is a matrix of batchsize x vocabsize
and embedding matrix of the size vocab x embedsize
Output: Sequence of dense-matrices, i.e., sequence[k] is a matrix of vocabsize x embedsize
"""
function embed_sequence(parameter, sequence, atype)
    embedded_sequence = Array(atype, length(sequence))
    for i=1:length(sequence)
        embedded_sequence[i] = sequence[i] * parameter
    end
    return embedded_sequence
end

"""
Calculated the cross-entropy loss function through a given sequence and returns the loss per token
paramdic[:forw]: holds the forward lstm variables
paramdict[:back]: holds the backward lstm variables
paramdict[:merge][1],[2]: holds the weights for final prediction, and bias respectively.
"""
function loss_old(paramdict, statedict, sequence)
    total = 0.0
    count = 0
    atype = typeof(AutoGrad.getval(paramdict[:forw][1]))
    #embedded_sequence = embed_sequence(model[:embedding], sequence, atype)
    s = sequence[1] * model[:embedding]
    #fhiddens = bilsforw(paramdict[:forw], statedict[:forw], embedded_sequence, atype)
    #bhiddens = bilsforw(paramdict[:back], statedict[:back], embedded_sequence, atype; forwardlstm=false)

    # go through the sequence one token at a time
    for t=1:8 #length(fhiddens)
        fhiddens = sforw(paramdict[:forw], statedict[:forw], s)
        bhiddens = sforw(paramdict[:back], statedict[:back], s)
        ypred = hcat(fhiddens, bhiddens) * paramdict[:merge][1] .+ paramdict[:merge][2] # merge the hidden layers
        ynorm = logp(ypred, 2)
        ygold = convert(atype, sequence[1])  # if KnetArray does not support linear indexing multiplication seems better alternative
        total += sum(ygold .* ynorm)
        count += size(ygold, 1) # each row corresponds to one instance in minibatch
    end
    return -total / count     # TODO: check with the count calculation to test loss per token is calculated
end

# Here are the remaining TODOs:
# - Check with the weight initializations s.t. the embedding matrix is still embedding, (KnetArray support?)
# - Implement loss function to do that flatten the final prediction matrix s.t. when we choose the correct entries of final prediction matrix
# - Based on the above steps implement the loss function

# FUTURE CODE: In case Knet supports table lookup
# """
# embedding : <vocab x embed> size matrix each row contains one of the related words
# minibatch : contains batchsize many nth words of the sequences
# """
# function lookup(embedding, minibatch, atype)
#     result = embedding[minibatch, :]
#     return convert(atype, result)
# end