text:
using TOML: Token
import TOML
using Test
using Dates

# Use a small buffer size to detect buffering bugs.
TOML.INIT_BUFFER_SIZE[] = 4

@testset "Token" begin
    # bare key
    @test TOML.keyname(Token(:bare_key, "foo")) === "foo"
    @test TOML.keyname(Token(:bare_key, "foo-bar")) === "foo-bar"

    # quoted key
    @test TOML.keyname(Token(:quoted_key, "\"foo\"")) === "foo"
    @test TOML.keyname(Token(:quoted_key, "'foo'")) === "foo"
    @test TOML.keyname(Token(:quoted_key, "\"\"")) === ""
    @test TOML.keyname(Token(:quoted_key, "''")) === ""
    @test TOML.keyname(Token(:quoted_key, "\"foo キー\"")) === "foo キー"
    @test TOML.keyname(Token(:quoted_key, "'foo キー'")) === "foo キー"

    @test_throws ArgumentError("not a key token") TOML.keyname(Token(:decimal, "1234"))

    # decimal
    @test TOML.value(Token(:decimal, "1234567890")) === 1234567890
    @test TOML.value(Token(:decimal, "-1234")) === -1234
    @test TOML.value(Token(:decimal, "-1_23_4")) === -1234

    # binary
    @test TOML.value(Token(:binary, "0b01")) === UInt(0b01)
    @test TOML.value(Token(:binary, "0b0_1")) === UInt(0b01)

    # octal
    @test TOML.value(Token(:octal, "0o01234567")) === UInt(0o01234567)
    @test TOML.value(Token(:octal, "0o0123_45_67")) === UInt(0o01234567)

    # hexadecimal
    @test TOML.value(Token(:hexadecimal, "0x0123456789ABCDEF")) === UInt(0x0123456789ABCDEF)
    @test TOML.value(Token(:hexadecimal, "0x0123456789abcdef")) === UInt(0x0123456789abcdef)
    @test TOML.value(Token(:hexadecimal, "0x0123_456789ab_cdef")) === UInt(0x0123456789abcdef)

    # float
    @test TOML.value(Token(:float, "1.23")) === 1.23
    @test TOML.value(Token(:float, "1e-2")) === 1e-2
    @test TOML.value(Token(:float, "inf"))  === Inf
    @test TOML.value(Token(:float, "-inf")) === -Inf
    @test TOML.value(Token(:float, "+inf")) === Inf
    @test TOML.value(Token(:float, "nan"))  === NaN
    @test TOML.value(Token(:float, "-nan")) === -NaN
    @test TOML.value(Token(:float, "+nan")) === NaN
    @test TOML.value(Token(:float, "9_224_617.445_991_228_313")) === 9224617.445991228313

    # boolean
    @test TOML.value(Token(:boolean, "true"))  === true
    @test TOML.value(Token(:boolean, "false")) === false

    # basic string
    @test TOML.value(Token(:basic_string, "\"\"")) === ""
    @test TOML.value(Token(:basic_string, "\"foobar\"")) === "foobar"
    @test TOML.value(Token(:basic_string, "\"αβγあいう\"")) === "αβγあいう"
    @test TOML.value(Token(:basic_string, "\"escaping: \\\" \\b \\t \\n \\f \\r \\\" \\\\ \"")) === "escaping: \" \b \t \n \f \r \" \\ "
    @test TOML.value(Token(:basic_string, "\"unicode: \\u3042\"")) === "unicode: あ"
    @test TOML.value(Token(:basic_string, "\"unicode: \\U00003042\"")) === "unicode: あ"

    # multiline basic string
    @test TOML.value(Token(:multiline_basic_string, "\"\"\"\"\"\"")) === ""
    @test TOML.value(Token(:multiline_basic_string, "\"\"\"foobar\"\"\"")) === "foobar"
    @test TOML.value(Token(:multiline_basic_string, "\"\"\"αβγあいう\"\"\"")) === "αβγあいう"
    @test TOML.value(Token(:multiline_basic_string,
    """\"\"\"
    multi
    -line
    文章
    \"\"\"""")) === "multi\n-line\n文章\n"
    @test TOML.value(Token(:multiline_basic_string,
    """\"\"\"
    line 1\r
    line 2\r
    \"\"\"""")) === "line 1\nline 2\n"
    @test TOML.value(Token(:multiline_basic_string,
    """\"\"\"
    foo\\

         bar
    \"\"\"""")) === "foobar\n"

    # literal string
    @test TOML.value(Token(:literal_string, "''")) === ""
    @test TOML.value(Token(:literal_string, "'αβγあいう'")) === "αβγあいう"
    @test TOML.value(Token(:literal_string, "'C:\\Users\\Windows\\Path'")) === raw"C:\Users\Windows\Path"

    # multiline literal string
    @test TOML.value(Token(:multiline_literal_string, "''''''")) === ""
    @test TOML.value(Token(:multiline_literal_string, "'''αβγあいう'''")) === "αβγあいう"
    @test TOML.value(Token(:multiline_literal_string, "'''C:\\Users\\Windows\\Path'''")) === raw"C:\Users\Windows\Path"
    @test TOML.value(Token(:multiline_literal_string, """
    '''
    line 1\r
    line 2\r
    '''""")) === "line 1\nline 2\n"
    @test TOML.value(Token(:multiline_literal_string, """
    '''
    foo\\

         bar
    '''""")) === "foo\\\n\n     bar\n"

    # datetime
    @test TOML.value(Token(:datetime, "1979-05-27T07:32:00Z")) == (DateTime(1979, 5, 27, 7, 32, 00), "Z")
    @test TOML.value(Token(:datetime, "1979-05-27 07:32:00Z")) == (DateTime(1979, 5, 27, 7, 32, 00), "Z")
    @test TOML.value(Token(:datetime, "1979-05-27T00:32:00-07:00")) == (DateTime(1979, 5, 27, 00, 32, 00), "-07:00")
    @test TOML.value(Token(:datetime, "1979-05-27T00:32:00.999-07:00")) == (DateTime(1979, 5, 27, 00, 32, 00, 999), "-07:00")

    @test_throws ArgumentError("not a value token: Token(:bare_key, \"foo\")") TOML.value(Token(:bare_key, "foo"))
end

@testset "Tokenizer" begin
    @test TOML.scanvalue(IOBuffer("\"foo\""), TOML.Buffer()) == (:basic_string, 5)
    @test TOML.scanvalue(IOBuffer("'foo'"), TOML.Buffer()) == (:literal_string, 5)
    @test TOML.scanvalue(IOBuffer("123"), TOML.Buffer()) == (:decimal, 3)
    @test TOML.scanvalue(IOBuffer("123.0"), TOML.Buffer()) == (:float, 5)
    @test TOML.scanvalue(IOBuffer("true"), TOML.Buffer()) == (:boolean, 4)
    @test TOML.scanvalue(IOBuffer("1979-05-27T00:32:00-07:00"), TOML.Buffer()) == (:datetime, 25)
    @test TOML.scanvalue(IOBuffer("1979-05-27T07:32:00"), TOML.Buffer()) == (:local_datetime, 19)
    @test TOML.scanvalue(IOBuffer("1979-05-27"), TOML.Buffer()) == (:local_date, 10)
    @test TOML.scanvalue(IOBuffer("07:32:00"), TOML.Buffer()) == (:local_time, 8)
    @test TOML.scanvalue(IOBuffer("abracadabra"), TOML.Buffer()) == (:novalue, 0)

    for (text, rhs, token) in
        [
            # simple tokens
            ("",        false, TOML.TOKEN_EOF),
            (" ",       false, TOML.TOKEN_WHITESPACE_SPACE),
            ("\t",      false, TOML.TOKEN_WHITESPACE_TAB),
            ("\n",      false, TOML.TOKEN_NEWLINE_LF),
            ("\r\n",    false, TOML.TOKEN_NEWLINE_CRLF),
            ("# a b",   false, TOML.Token(:comment, "# a b")),
            ("=",       false, TOML.TOKEN_EQUAL),
            (".",       false, TOML.TOKEN_DOT),
            (",",       false, TOML.TOKEN_COMMA),
            ("{",       false, TOML.TOKEN_CURLY_BRACE_LEFT),
            ("}",       false, TOML.TOKEN_CURLY_BRACE_RIGHT),
            # tokens depending on the `rhs` parameter
            ("100",     false, TOML.Token(:bare_key, "100")),
            ("100",     true,  TOML.Token(:decimal, "100")),
            ("true",    false, TOML.Token(:bare_key, "true")),
            ("true",    true,  TOML.Token(:boolean, "true")),
            ("false",   false, TOML.Token(:bare_key, "false")),
            ("false",   true,  TOML.Token(:boolean, "false")),
            ("\"key\"", false, TOML.Token(:quoted_key, "\"key\"")),
            ("\"key\"", true,  TOML.Token(:basic_string, "\"key\"")),
            ("'key'",   false, TOML.Token(:quoted_key, "'key'")),
            ("'key'",   true,  TOML.Token(:literal_string, "'key'")),
            ("[[",      false, TOML.TOKEN_DOUBLE_BRACKETS_LEFT),
            ("[[",      true,  TOML.TOKEN_SINGLE_BRACKET_LEFT),
            ("]]",      false, TOML.TOKEN_DOUBLE_BRACKETS_RIGHT),
            ("]]",      true,  TOML.TOKEN_SINGLE_BRACKET_RIGHT),
            # unknown tokens (will result in parse error)
            ("!",       false, TOML.Token(:unknown, "!")),
            ("\rX",     false, TOML.Token(:unknown, "\r")),
            ("\"X",     false, TOML.Token(:unknown, "\"")),
            ("'X",      false, TOML.Token(:unknown, "'")),
        ]
        tokenizer = TOML.Tokenizer(IOBuffer(text))
        @test TOML.readtoken(tokenizer, rhs=rhs) == token
    end

    tokenizer = TOML.Tokenizer(IOBuffer("100"))
    @test TOML.peektoken(tokenizer, rhs=false) == TOML.Token(:bare_key, "100")
    @test TOML.peektoken(tokenizer, rhs=true)  == TOML.Token(:decimal, "100")
    @test TOML.peektoken(tokenizer, rhs=false) == TOML.Token(:bare_key, "100")
    @test TOML.peektoken(tokenizer, rhs=true)  == TOML.Token(:decimal, "100")
    @test TOML.readtoken(tokenizer, rhs=false) == TOML.Token(:bare_key, "100")
    @test TOML.readtoken(tokenizer, rhs=false) == TOML.TOKEN_EOF
end

@testset "StreamReader" begin
    function tokens(str)
        stream = TOML.StreamReader(IOBuffer(str))
        ts = TOML.Token[]
        while (t = TOML.parsetoken(stream)).kind != :eof
            push!(ts, t)
        end
        return ts
    end

    @test_throws ErrorException("invalid UTF-8 sequence") tokens("\x80")

    @test tokens("") == TOML.Token[]
    @test tokens(" ") == [TOML.Token(:whitespace, " ")]
    @test tokens("  ") == [TOML.Token(:whitespace, "  ")]
    @test tokens("\t") == [TOML.Token(:whitespace, "\t")]
    @test tokens("\t  ") == [TOML.Token(:whitespace, "\t  ")]
    @test tokens("\n") == [TOML.Token(:newline, "\n")]
    @test tokens("\r\n") == [TOML.Token(:newline, "\r\n")]
    @test tokens("#comment") == [TOML.Token(:comment, "#comment")]
    @test tokens("#αあ𐍈") == [TOML.Token(:comment, "#αあ𐍈")]

    @test tokens("""
      # comment
    """) == [
        TOML.Token(:whitespace, "  "),
        TOML.Token(:comment, "# comment"),
        TOML.Token(:newline, "\n"),
    ]

    @test tokens("x=10") == [
        TOML.Token(:bare_key, "x"),
        TOML.Token(:equal, "="),
        TOML.Token(:decimal, "10"),
    ]

    @test tokens("x = 10") == [
        TOML.Token(:bare_key, "x"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:equal, "="),
        TOML.Token(:whitespace, " "),
        TOML.Token(:decimal, "10"),
    ]

    @test tokens("x = 10.0") == [
        TOML.Token(:bare_key, "x"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:equal, "="),
        TOML.Token(:whitespace, " "),
        TOML.Token(:float, "10.0"),
    ]

    @test tokens("x = 6.022e23") == [
        TOML.Token(:bare_key, "x"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:equal, "="),
        TOML.Token(:whitespace, " "),
        TOML.Token(:float, "6.022e23"),
    ]

    @test tokens("x = inf") == [
        TOML.Token(:bare_key, "x"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:equal, "="),
        TOML.Token(:whitespace, " "),
        TOML.Token(:float, "inf"),
    ]

    @test tokens("x = nan") == [
        TOML.Token(:bare_key, "x"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:equal, "="),
        TOML.Token(:whitespace, " "),
        TOML.Token(:float, "nan"),
    ]

    @test tokens("""
    foo = true
    bar = false
    """) == [
        TOML.Token(:bare_key, "foo"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:equal, "="),
        TOML.Token(:whitespace, " "),
        TOML.Token(:boolean, "true"),
        TOML.Token(:newline, "\n"),
        TOML.Token(:bare_key, "bar"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:equal, "="),
        TOML.Token(:whitespace, " "),
        TOML.Token(:boolean, "false"),
        TOML.Token(:newline, "\n"),
    ]

    @test tokens("""
    foo = "text"
    """) == [
        TOML.Token(:bare_key, "foo"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:equal, "="),
        TOML.Token(:whitespace, " "),
        TOML.Token(:basic_string, "\"text\""),
        TOML.Token(:newline, "\n"),
    ]

    @test tokens("""
    foo = \"\"\"
    "text1"
    ""text2""
    \"\"\"
    """) == [
        TOML.Token(:bare_key, "foo"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:equal, "="),
        TOML.Token(:whitespace, " "),
        TOML.Token(:multiline_basic_string, "\"\"\"\n\"text1\"\n\"\"text2\"\"\n\"\"\""),
        TOML.Token(:newline, "\n"),
    ]

    @test tokens("""
    foo = 'text'
    """) == [
        TOML.Token(:bare_key, "foo"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:equal, "="),
        TOML.Token(:whitespace, " "),
        TOML.Token(:literal_string, "'text'"),
        TOML.Token(:newline, "\n"),
    ]

    @test tokens("""
    foo = '''
    'text'
    ''text2''
    '''
    """) == [
        TOML.Token(:bare_key, "foo"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:equal, "="),
        TOML.Token(:whitespace, " "),
        TOML.Token(:multiline_literal_string, "'''\n'text'\n''text2''\n'''"),
        TOML.Token(:newline, "\n"),
    ]

    @test tokens("x = []") == [
        TOML.Token(:bare_key, "x"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:equal, "="),
        TOML.Token(:whitespace, " "),
        TOML.Token(:inline_array_begin, ""),
        TOML.Token(:single_bracket_left, "["),
        TOML.Token(:single_bracket_right, "]"),
        TOML.Token(:inline_array_end, ""),
    ]

    @test tokens("x = [1]") == [
        TOML.Token(:bare_key, "x"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:equal, "="),
        TOML.Token(:whitespace, " "),
        TOML.Token(:inline_array_begin, ""),
        TOML.Token(:single_bracket_left, "["),
        TOML.Token(:decimal, "1"),
        TOML.Token(:single_bracket_right, "]"),
        TOML.Token(:inline_array_end, ""),
    ]

    @test tokens("x = [ 1 ]") == [
        TOML.Token(:bare_key, "x"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:equal, "="),
        TOML.Token(:whitespace, " "),
        TOML.Token(:inline_array_begin, ""),
        TOML.Token(:single_bracket_left, "["),
        TOML.Token(:whitespace, " "),
        TOML.Token(:decimal, "1"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:single_bracket_right, "]"),
        TOML.Token(:inline_array_end, ""),
    ]

    @test tokens("x = [1,]") == [
        TOML.Token(:bare_key, "x"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:equal, "="),
        TOML.Token(:whitespace, " "),
        TOML.Token(:inline_array_begin, ""),
        TOML.Token(:single_bracket_left, "["),
        TOML.Token(:decimal, "1"),
        TOML.Token(:comma, ","),
        TOML.Token(:single_bracket_right, "]"),
        TOML.Token(:inline_array_end, ""),
    ]

    @test tokens("x = [ 1 , ]") == [
        TOML.Token(:bare_key, "x"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:equal, "="),
        TOML.Token(:whitespace, " "),
        TOML.Token(:inline_array_begin, ""),
        TOML.Token(:single_bracket_left, "["),
        TOML.Token(:whitespace, " "),
        TOML.Token(:decimal, "1"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:comma, ","),
        TOML.Token(:whitespace, " "),
        TOML.Token(:single_bracket_right, "]"),
        TOML.Token(:inline_array_end, ""),
    ]

    @test tokens("""
    x = [
        1,
        2, # comment
    ] # comment
    """) == [
        TOML.Token(:bare_key, "x"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:equal, "="),
        TOML.Token(:whitespace, " "),
        TOML.Token(:inline_array_begin, ""),
        TOML.Token(:single_bracket_left, "["),
        TOML.Token(:whitespace, "  "),
        TOML.Token(:newline, "\n"),
        TOML.Token(:whitespace, "    "),
        TOML.Token(:decimal, "1"),
        TOML.Token(:comma, ","),
        TOML.Token(:whitespace, "  "),
        TOML.Token(:newline, "\n"),
        TOML.Token(:whitespace, "    "),
        TOML.Token(:decimal, "2"),
        TOML.Token(:comma, ","),
        TOML.Token(:whitespace, " "),
        TOML.Token(:comment, "# comment"),
        TOML.Token(:newline, "\n"),
        TOML.Token(:single_bracket_right, "]"),
        TOML.Token(:inline_array_end, ""),
        TOML.Token(:whitespace, " "),
        TOML.Token(:comment, "# comment"),
        TOML.Token(:newline, "\n"),
    ]

    @test tokens("x = [[]]") == [
        TOML.Token(:bare_key, "x"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:equal, "="),
        TOML.Token(:whitespace, " "),
        TOML.Token(:inline_array_begin, ""),
        TOML.Token(:single_bracket_left, "["),
        TOML.Token(:inline_array_begin, ""),
        TOML.Token(:single_bracket_left, "["),
        TOML.Token(:single_bracket_right, "]"),
        TOML.Token(:inline_array_end, ""),
        TOML.Token(:single_bracket_right, "]"),
        TOML.Token(:inline_array_end, ""),
    ]

    @test tokens("""
    x = [[1,2],[3,4]]
    """) == [
        TOML.Token(:bare_key, "x"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:equal, "="),
        TOML.Token(:whitespace, " "),
        TOML.Token(:inline_array_begin, ""),
        TOML.Token(:single_bracket_left, "["),
        TOML.Token(:inline_array_begin, ""),
        TOML.Token(:single_bracket_left, "["),
        TOML.Token(:decimal, "1"),
        TOML.Token(:comma, ","),
        TOML.Token(:decimal, "2"),
        TOML.Token(:single_bracket_right, "]"),
        TOML.Token(:inline_array_end, ""),
        TOML.Token(:comma, ","),
        TOML.Token(:inline_array_begin, ""),
        TOML.Token(:single_bracket_left, "["),
        TOML.Token(:decimal, "3"),
        TOML.Token(:comma, ","),
        TOML.Token(:decimal, "4"),
        TOML.Token(:single_bracket_right, "]"),
        TOML.Token(:inline_array_end, ""),
        TOML.Token(:single_bracket_right, "]"),
        TOML.Token(:inline_array_end, ""),
        TOML.Token(:newline, "\n"),
    ]

    @test tokens("""
    x = [
        [1,2,3],
        [4,5,6],
    ]
    """) == [
        TOML.Token(:bare_key, "x"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:equal, "="),
        TOML.Token(:whitespace, " "),
        TOML.Token(:inline_array_begin, ""),
        TOML.Token(:single_bracket_left, "["),
        TOML.Token(:newline, "\n"),
        TOML.Token(:whitespace, "    "),
        TOML.Token(:inline_array_begin, ""),
        TOML.Token(:single_bracket_left, "["),
        TOML.Token(:decimal, "1"),
        TOML.Token(:comma, ","),
        TOML.Token(:decimal, "2"),
        TOML.Token(:comma, ","),
        TOML.Token(:decimal, "3"),
        TOML.Token(:single_bracket_right, "]"),
        TOML.Token(:inline_array_end, ""),
        TOML.Token(:comma, ","),
        TOML.Token(:newline, "\n"),
        TOML.Token(:whitespace, "    "),
        TOML.Token(:inline_array_begin, ""),
        TOML.Token(:single_bracket_left, "["),
        TOML.Token(:decimal, "4"),
        TOML.Token(:comma, ","),
        TOML.Token(:decimal, "5"),
        TOML.Token(:comma, ","),
        TOML.Token(:decimal, "6"),
        TOML.Token(:single_bracket_right, "]"),
        TOML.Token(:inline_array_end, ""),
        TOML.Token(:comma, ","),
        TOML.Token(:newline, "\n"),
        TOML.Token(:single_bracket_right, "]"),
        TOML.Token(:inline_array_end, ""),
        TOML.Token(:newline, "\n"),
    ]

    @test_throws TOML.ParseError("unexpected bare key 'y' at line 1") tokens("x = [1,2] y = 10")

    @test tokens("""
    x = {}
    """) == [
        TOML.Token(:bare_key, "x"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:equal, "="),
        TOML.Token(:whitespace, " "),
        TOML.Token(:inline_table_begin, ""),
        TOML.Token(:curly_brace_left, "{"),
        TOML.Token(:curly_brace_right, "}"),
        TOML.Token(:inline_table_end, ""),
        TOML.Token(:newline, "\n"),
    ]

    @test tokens("""
    x = { y = 10 }
    """) == [
        TOML.Token(:bare_key, "x"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:equal, "="),
        TOML.Token(:whitespace, " "),
        TOML.Token(:inline_table_begin, ""),
        TOML.Token(:curly_brace_left, "{"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:bare_key, "y"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:equal, "="),
        TOML.Token(:whitespace, " "),
        TOML.Token(:decimal, "10"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:curly_brace_right, "}"),
        TOML.Token(:inline_table_end, ""),
        TOML.Token(:newline, "\n"),
    ]

    @test tokens("""
    x = { y = 10, z = 20 }
    """) == [
        TOML.Token(:bare_key, "x"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:equal, "="),
        TOML.Token(:whitespace, " "),
        TOML.Token(:inline_table_begin, ""),
        TOML.Token(:curly_brace_left, "{"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:bare_key, "y"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:equal, "="),
        TOML.Token(:whitespace, " "),
        TOML.Token(:decimal, "10"),
        TOML.Token(:comma, ","),
        TOML.Token(:whitespace, " "),
        TOML.Token(:bare_key, "z"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:equal, "="),
        TOML.Token(:whitespace, " "),
        TOML.Token(:decimal, "20"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:curly_brace_right, "}"),
        TOML.Token(:inline_table_end, ""),
        TOML.Token(:newline, "\n"),
    ]

    @test_throws TOML.ParseError("unexpected '}' at line 1") tokens("x = { y }")
    @test_throws TOML.ParseError("unexpected bare key 'z' at line 1") tokens("x = { y = 10 z = 20 }")
    @test_throws TOML.ParseError("unexpected newline at line 1") tokens("x = { y = 10 \n}")

    @test tokens("""
    x = {y=[1,2]}
    """) == [
        TOML.Token(:bare_key, "x"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:equal, "="),
        TOML.Token(:whitespace, " "),
        TOML.Token(:inline_table_begin, ""),
        TOML.Token(:curly_brace_left, "{"),
        TOML.Token(:bare_key, "y"),
        TOML.Token(:equal, "="),
        TOML.Token(:inline_array_begin, ""),
        TOML.Token(:single_bracket_left, "["),
        TOML.Token(:decimal, "1"),
        TOML.Token(:comma, ","),
        TOML.Token(:decimal, "2"),
        TOML.Token(:single_bracket_right, "]"),
        TOML.Token(:inline_array_end, ""),
        TOML.Token(:curly_brace_right, "}"),
        TOML.Token(:inline_table_end, ""),
        TOML.Token(:newline, "\n"),
    ]

    @test tokens("""
    x = {y=[1,2],z=[3,4]}
    """) == [
        TOML.Token(:bare_key, "x"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:equal, "="),
        TOML.Token(:whitespace, " "),
        TOML.Token(:inline_table_begin, ""),
        TOML.Token(:curly_brace_left, "{"),
        TOML.Token(:bare_key, "y"),
        TOML.Token(:equal, "="),
        TOML.Token(:inline_array_begin, ""),
        TOML.Token(:single_bracket_left, "["),
        TOML.Token(:decimal, "1"),
        TOML.Token(:comma, ","),
        TOML.Token(:decimal, "2"),
        TOML.Token(:single_bracket_right, "]"),
        TOML.Token(:inline_array_end, ""),
        TOML.Token(:comma, ","),
        TOML.Token(:bare_key, "z"),
        TOML.Token(:equal, "="),
        TOML.Token(:inline_array_begin, ""),
        TOML.Token(:single_bracket_left, "["),
        TOML.Token(:decimal, "3"),
        TOML.Token(:comma, ","),
        TOML.Token(:decimal, "4"),
        TOML.Token(:single_bracket_right, "]"),
        TOML.Token(:inline_array_end, ""),
        TOML.Token(:curly_brace_right, "}"),
        TOML.Token(:inline_table_end, ""),
        TOML.Token(:newline, "\n"),
    ]

    @test tokens("x=[{foo=10}]") == [
        TOML.Token(:bare_key, "x"),
        TOML.Token(:equal, "="),
        TOML.Token(:inline_array_begin, ""),
        TOML.Token(:single_bracket_left, "["),
        TOML.Token(:inline_table_begin, ""),
        TOML.Token(:curly_brace_left, "{"),
        TOML.Token(:bare_key, "foo"),
        TOML.Token(:equal, "="),
        TOML.Token(:decimal, "10"),
        TOML.Token(:curly_brace_right, "}"),
        TOML.Token(:inline_table_end, ""),
        TOML.Token(:single_bracket_right, "]"),
        TOML.Token(:inline_array_end, ""),
    ]

    @test tokens("""
    x={foo={x=10}, bar={y=20},}
    """) == [
        TOML.Token(:bare_key, "x"),
        TOML.Token(:equal, "="),
        TOML.Token(:inline_table_begin, ""),
        TOML.Token(:curly_brace_left, "{"),
        TOML.Token(:bare_key, "foo"),
        TOML.Token(:equal, "="),
        TOML.Token(:inline_table_begin, ""),
        TOML.Token(:curly_brace_left, "{"),
        TOML.Token(:bare_key, "x"),
        TOML.Token(:equal, "="),
        TOML.Token(:decimal, "10"),
        TOML.Token(:curly_brace_right, "}"),
        TOML.Token(:inline_table_end, ""),
        TOML.Token(:comma, ","),
        TOML.Token(:whitespace, " "),
        TOML.Token(:bare_key, "bar"),
        TOML.Token(:equal, "="),
        TOML.Token(:inline_table_begin, ""),
        TOML.Token(:curly_brace_left, "{"),
        TOML.Token(:bare_key, "y"),
        TOML.Token(:equal, "="),
        TOML.Token(:decimal, "20"),
        TOML.Token(:curly_brace_right, "}"),
        TOML.Token(:inline_table_end, ""),
        TOML.Token(:comma, ","),
        TOML.Token(:curly_brace_right, "}"),
        TOML.Token(:inline_table_end, ""),
        TOML.Token(:newline, "\n"),
    ]

    @test tokens("""
    "foo"=100
    """) == [
        TOML.Token(:quoted_key, "\"foo\""),
        TOML.Token(:equal, "="),
        TOML.Token(:decimal, "100"),
        TOML.Token(:newline, "\n"),
    ]

    @test tokens("""
    [foo]
    """) == [
        TOML.Token(:table_begin, ""),
        TOML.Token(:single_bracket_left, "["),
        TOML.Token(:bare_key, "foo"),
        TOML.Token(:single_bracket_right, "]"),
        TOML.Token(:table_end, ""),
        TOML.Token(:newline, "\n"),
    ]

    @test tokens("""
    [foo.bar]
    """) == [
        TOML.Token(:table_begin, ""),
        TOML.Token(:single_bracket_left, "["),
        TOML.Token(:bare_key, "foo"),
        TOML.Token(:dot, "."),
        TOML.Token(:bare_key, "bar"),
        TOML.Token(:single_bracket_right, "]"),
        TOML.Token(:table_end, ""),
        TOML.Token(:newline, "\n"),
    ]

    @test tokens("""
    [foo.bar.baz]
    """) == [
        TOML.Token(:table_begin, ""),
        TOML.Token(:single_bracket_left, "["),
        TOML.Token(:bare_key, "foo"),
        TOML.Token(:dot, "."),
        TOML.Token(:bare_key, "bar"),
        TOML.Token(:dot, "."),
        TOML.Token(:bare_key, "baz"),
        TOML.Token(:single_bracket_right, "]"),
        TOML.Token(:table_end, ""),
        TOML.Token(:newline, "\n"),
    ]

    @test tokens("""
    ['foo']
    ["foo"]
    """) == [
        TOML.Token(:table_begin, ""),
        TOML.Token(:single_bracket_left, "["),
        TOML.Token(:quoted_key, "'foo'"),
        TOML.Token(:single_bracket_right, "]"),
        TOML.Token(:table_end, ""),
        TOML.Token(:newline, "\n"),
        TOML.Token(:table_begin, ""),
        TOML.Token(:single_bracket_left, "["),
        TOML.Token(:quoted_key, "\"foo\""),
        TOML.Token(:single_bracket_right, "]"),
        TOML.Token(:table_end, ""),
        TOML.Token(:newline, "\n"),
    ]

    @test tokens("""
    [foo."bar.baz"]
    """) == [
        TOML.Token(:table_begin, ""),
        TOML.Token(:single_bracket_left, "["),
        TOML.Token(:bare_key, "foo"),
        TOML.Token(:dot, "."),
        TOML.Token(:quoted_key, "\"bar.baz\""),
        TOML.Token(:single_bracket_right, "]"),
        TOML.Token(:table_end, ""),
        TOML.Token(:newline, "\n"),
    ]

    @test tokens("""
    [ foo . bar ]
    """) == [
        TOML.Token(:table_begin, ""),
        TOML.Token(:single_bracket_left, "["),
        TOML.Token(:whitespace, " "),
        TOML.Token(:bare_key, "foo"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:dot, "."),
        TOML.Token(:whitespace, " "),
        TOML.Token(:bare_key, "bar"),
        TOML.Token(:whitespace, " "),
        TOML.Token(:single_bracket_right, "]"),
        TOML.Token(:table_end, ""),
        TOML.Token(:newline, "\n"),
    ]

    @test_throws TOML.ParseError("unexpected ']' at line 1") tokens("[]")
    @test_throws TOML.ParseError("unexpected ']' at line 1") tokens("[foo.]")
    @test_throws TOML.ParseError("unexpected newline at line 1") tokens("[\nfoo]")
    @test_throws TOML.ParseError("unexpected newline at line 1") tokens("[foo\n]")
    @test_throws TOML.ParseError("unexpected bare key 'bar' at line 1") tokens("[foo bar]")
    @test_throws TOML.ParseError("unexpected bare key 'x' at line 1") tokens("[foo] x = 10")

    @test tokens("""
    [[foo]]
    """) == [
        TOML.Token(:array_begin, ""),
        TOML.Token(:double_brackets_left, "[["),
        TOML.Token(:bare_key, "foo"),
        TOML.Token(:double_brackets_right, "]]"),
        TOML.Token(:array_end, ""),
        TOML.Token(:newline, "\n"),
    ]

    @test tokens("""
    [[foo.bar]]
    """) == [
        TOML.Token(:array_begin, ""),
        TOML.Token(:double_brackets_left, "[["),
        TOML.Token(:bare_key, "foo"),
        TOML.Token(:dot, "."),
        TOML.Token(:bare_key, "bar"),
        TOML.Token(:double_brackets_right, "]]"),
        TOML.Token(:array_end, ""),
        TOML.Token(:newline, "\n"),
    ]

    @test_throws TOML.ParseError("unexpected ']]' at line 1") tokens("[[]]")
    @test_throws TOML.ParseError("unexpected ']' at line 1") tokens("[[foo]")
    @test_throws TOML.ParseError("unexpected ']]' at line 1") tokens("[foo]]")
    @test_throws TOML.ParseError("unexpected bare key 'x' at line 1") tokens("[[foo]] x = 10")

    @test_throws TOML.ParseError("unexpected '!' at line 1") tokens("!")
    @test_throws TOML.ParseError("unexpected '!' at line 3") tokens("\n\n!")
    @test_throws TOML.ParseError("unexpected '\0' at line 1") tokens("# \0")
    @test_throws TOML.ParseError("unexpected end of file at line 1") tokens("x")
    @test_throws TOML.ParseError("unexpected end of file at line 1") tokens("x=")
    @test_throws TOML.ParseError("unexpected ',' at line 1") tokens("x,")
    @test_throws TOML.ParseError("unexpected 'p' at line 1") tokens("x = p")
    @test_throws TOML.ParseError("unexpected ',' at line 1") tokens("x = [,1]")
    @test_throws TOML.ParseError("unexpected ''' at line 1") tokens("x='foo")
    @test_throws TOML.ParseError("unexpected newline at line 1") tokens("x=\n10")
    @test_throws TOML.ParseError("unexpected bare key 'bar' at line 1") tokens("foo=100 bar=200")
    @test_throws TOML.ParseError("unexpected ']' at line 1") tokens("[]")
    @test_throws TOML.ParseError("unexpected ']' at line 1") tokens("[foo.]")
    @test_throws TOML.ParseError("unexpected '=' at line 1") tokens("[foo=]")
    @test_throws TOML.ParseError("unexpected ']]' at line 1") tokens("[foo]]")
    @test_throws TOML.ParseError("unexpected ']' at line 1") tokens("[[hoge]")
    @test_throws TOML.ParseError("unexpected ']]' at line 1") tokens("[[hoge.]]")
    @test_throws TOML.ParseError("unexpected end of file at line 1") tokens("[foo")
    @test_throws TOML.ParseError("unexpected ']' at line 1") tokens("x=[100]]")
    @test_throws TOML.ParseError("unexpected end of file at line 1") tokens("x=[100")
    @test_throws TOML.ParseError("unexpected '}' at line 1") tokens("x={foo}")
    @test_throws TOML.ParseError("unexpected ',' at line 1") tokens("x={,foo=100}")
    @test_throws TOML.ParseError("unexpected '.' at line 1") tokens("x={10.=10}")
    @test_throws TOML.ParseError("unexpected newline at line 1") tokens("x={\nfoo=100}")
    @test_throws TOML.ParseError("unexpected '\r' at line 1") tokens("foo=100\r")
    @test_throws TOML.ParseError("unexpected newline at line 11") tokens(
    """
    foo = '''
    some
    text
    '''

    bar = \"\"\"
    some
    text
    \"\"\"

    error =
    """)
end

@testset "Parser" begin
    @test TOML.parse("") == Dict()
    @test TOML.parse("foo = 100") == Dict("foo" => 100)
    @test TOML.parse("foo = 100\nbar = 1.23") == Dict("foo" => 100, "bar" => 1.23)
    @test TOML.parse("sushi = \"🍣\"\nbeer = \"🍺\"") == Dict("sushi" => "🍣", "beer" => "🍺")
    @test TOML.parse("name = { first = \"Tom\", last = \"Preston-Werner\" }") == Dict("name" => Dict("first" => "Tom", "last" => "Preston-Werner"))
    @test TOML.parse("x = {foo = 10, bar = { a = 1, b = 2 }, baz = {}}") == Dict("x" => Dict("foo" => 10, "bar" => Dict("a" => 1, "b" => 2), "baz" => Dict()))

    data = TOML.parse("x = [1]\ny = [1,2,3]\nz = []")
    @test data == Dict("x" => [1], "y" => [1,2,3], "z" => [])
    @test data["x"] isa Vector{Int}
    @test data["y"] isa Vector{Int}
    @test data["z"] isa Vector{Any}

    data = TOML.parse("x = [[1,2,3], [1,2]]\ny = [[1,2]]\nz = [[[]]]")
    @test data == Dict("x" => [[1,2,3], [1,2]], "y" => [[1,2]], "z" => [[[]]])
    @test data["x"] isa Vector{Vector}

    data = TOML.parse(
    """
    key1 = 0

    [foo]
    key1 = 1
    key2 = 2

    [bar]
    key1 = 10
    key2 = 20

    [foo.bar.baz]
    qux = 100
    """)
    @test data == Dict(
        "key1" => 0,
        "foo" => Dict("key1" => 1, "key2" => 2, "bar" => Dict("baz" => Dict("qux" => 100))),
        "bar" => Dict("key1" => 10, "key2" => 20))

    data = TOML.parse(
    """
    [[foo]]
    x = 1
    y = 2

    [[foo]]

    [[foo]]
    x = 10
    y = 20
    z = 30

    [[bar.baz]]
    qux = 100
    """)
    @test data == Dict(
        "foo" => [
            Dict("x" => 1, "y" => 2),
            Dict(),
            Dict("x" => 10, "y" => 20, "z" => 30),
        ],
        "bar" => Dict("baz" => [Dict("qux" => 100)]))
    @test data["foo"] isa Vector{Dict{String,Any}}

    data = TOML.parse(
    """
    [[fruit]]
      name = "apple"

      [fruit.physical]
        color = "red"
        shape = "round"

      [[fruit.variety]]
        name = "red delicious"

      [[fruit.variety]]
        name = "granny smith"

    [[fruit]]
      name = "banana"

      [[fruit.variety]]
        name = "plantain"
    """)
    @test data == Dict(
        "fruit" => [
            Dict(
                "name" => "apple",
                "physical" => Dict("color" => "red", "shape" => "round"),
                "variety" => [Dict("name" => "red delicious"), Dict("name" => "granny smith")]),
            Dict(
                 "name" => "banana",
                 "variety" => [Dict("name" => "plantain")])
        ])

    data = TOML.parse(IOBuffer(
    """
    x = 100
    """))
    @test data == Dict("x" => 100)

    data = TOML.parse("""
    [a.b]
    c = 1

    [a]
    d = 2
    """)
    @test data == Dict("a" => Dict("b" => Dict("c" => 1), "d" => 2))

    @test_throws TOML.ParseError("mixed array types at line 1") TOML.parse("x = [1, 1.0]")
    @test_throws TOML.ParseError("mixed array types at line 1") TOML.parse("x = [[1,2], 3]")
    @test_throws TOML.ParseError("found a duplicated definition at line 4") TOML.parse(
    """
    [a]
    b = 1

    [a]
    c = 2
    """)
    @test_throws TOML.ParseError("found a duplicated definition at line 4") TOML.parse(
    """
    [a]
    b = 1

    [a.b]
    c = 2
    """)
end


text:
using PkgSRA, Plots, StatsPlots
using Turing, AxisArrays, DataFrames, LinearAlgebra, DataStructures
using BSON: @load
pyplot()

#####
##### Load the data
#####
this_dat_name = DAT_FOLDERNAME*"dat_neuron_"

## Basic spiking model
# Raw data
fname = this_dat_name*"raw.bson";
@load fname dat numerical_grad #grad_true U_true

# Controlled model
fname = this_dat_name*"controlled_model.bson";
@load fname ctr_final

## Spikes with varying input
# Raw data
fname = this_dat_name*"raw2.bson";
@load fname dat2 numerical_grad2 U_true2 #grad_true2

# Controlled model
fname = this_dat_name*"controlled_model2.bson";
@load fname ctr_final2

#####
##### Produce the plots
#####
plot_opt = Dict(:titlefontsize=>24,
        :yticks=>false, :fontfamily=>:serif,
        :legendfontsize=>16)

## Two density plots
# Two panels: data, then control
plot_data = plot(ts, dat[1,:], lw=5,
                color=COLOR_DICT["data"],
                legend=false, xticks=false;
                plot_opt...);
    xlabel!("");
    title!("Neuron With Constant Input ")

# Second: Controller
plot_control = plot(ts, ctr_final[1,:], legend=false,
                color=COLOR_DICT["control_time"], lw=3; plot_opt...);
    xlabel!("Time", guidefontsize=14, tickfontsize=14);
    title!("Learned Controller")

# Create the layout and plot
my_layout = @layout [p1; p2];
    p_final = plot(plot_data, plot_control, layout = my_layout)

# Save
fname = FIGURE_FOLDERNAME * "fig_neuron.png";
savefig(p_final, fname)

##
##
## Part 2: with external control changes
plot_data2 = plot(ts, dat2[1,:], lw=5,
                color=COLOR_DICT["data"],
                legend=false, xticks=false;
                plot_opt...);
    xlabel!("");
    title!("Neuron With Varying Input")

# Second: Learned and true Controller
inset_ind = 3200:4200; # Chosen by hand

plot_control2 = plot(ts, ctr_final2[1,:], legend=false,
                color=COLOR_DICT["control_time"], lw=3; plot_opt...);
    xlabel!("Time", guidefontsize=14, tickfontsize=14);
    title!("Learned Controller")
    plot!([ts[inset_ind[1]], ts[inset_ind[1]]],
        [-1000, 1000], color=:black, lw=3)
    plot!([ts[inset_ind[end]], ts[inset_ind[end]]],
        [-1000, 1000], color=:black, lw=3)

# Create what will be an inset
plot_inset = plot(U_true2[1,inset_ind], lw=5,
                    color=COLOR_DICT["control_true"], # label="True",
                    xticks=false,
                    legendfontsize=48, legend=false) #legend=:bottomleft)
    plot!(ctr_final2[1,inset_ind], lw=5,# label="Learned",
        color=COLOR_DICT["control_time"],
        ylims=[-100, 300]; plot_opt...)

# Create the layout and plot
my_layout = @layout [p1; p2];
    p_final2 = plot(plot_data2, plot_control2, layout = my_layout)

# Save
fname = FIGURE_FOLDERNAME * "fig_neuron2.png";
savefig(p_final2, fname)

fname = FIGURE_FOLDERNAME * "fig_neuron_inset.png";
savefig(plot_inset, fname)

ext: jl


text:
using SharedArrays
include("transport.jl")
const eps=1e-8

struct Basic
    m::Model
    x::JuMP.Containers.DenseAxisArray
    z::JuMP.Containers.DenseAxisArray
    in_flow::Array{VariableRef,1}
    out_flow::Array{VariableRef,1}
    n::Int64
    max_cycle_length::Int64
    max_chain_length::Int64
    graph::DiGraph
    time_param::String
    time_factor::Float64
    start_time::Float64

    function Basic(input_data::Input)
        # start_time = time()
        start_time = input_data.start_time
        graph, weights, ndds, max_cycle_length, max_chain_length, solver_instance = input_data.graph, input_data.weights, input_data.ndds, input_data.max_cycle_length, input_data.max_chain_length, input_data.solver_instance
        m = Model(solver_instance)
        n = nv(graph)
        graph_edges = collect(edges(graph))
        patients = setdiff(1:n, ndds)

        @variable(m, 1 >= x[graph_edges] >= 0, Int)
        @variable(m, 1 >= y[graph_edges] >= 0, Int)
        @variable(m, z[graph_edges] >= 0, Int)
        @variable(m, 1 >= in_flow[1:n] >= 0, Int)
        @variable(m, 1 >= out_flow[1:n] >= 0, Int)

        @objective(m, Max, sum(z[e]*weights[e.src, e.dst] for e in graph_edges))
        @constraint(m, c0[e=graph_edges], x[e]+y[e] == z[e])
        @constraint(m, c1[k=1:n], sum(x[Edge((i,k))] for i=inneighbors(graph, k)) == in_flow[k])
        @constraint(m, c2[k=1:n], sum(x[Edge((k,j))] for j=outneighbors(graph, k)) == out_flow[k])
        @constraint(m, c3[k=patients], out_flow[k] <= in_flow[k])
        @constraint(m, c4[k=ndds], in_flow[k] == 0)

        if max_chain_length < n-1
            @constraint(m, c5[p=patients], out_flow[p] == in_flow[p])
            @constraint(m, c6[k=ndds], out_flow[k] == 0)

            @variable(m, 1>= x_ndds[ndds, graph_edges] >= 0, Int)
            @variable(m, 1>= in_ndds[ndds, 1:n] >= 0, Int)
            @variable(m, 1>= out_ndds[ndds, 1:n] >= 0, Int)

            @constraint(m, c7[k=ndds], sum(x_ndds[k, e] for e in graph_edges) <= max_chain_length)
            @constraint(m, c8[v=ndds, k=1:n], sum(x_ndds[v, Edge((i,k))] for i=inneighbors(graph, k)) == in_ndds[v, k])
            @constraint(m, c9[v=ndds, k=1:n], sum(x_ndds[v, Edge((k,j))] for j=outneighbors(graph, k)) == out_ndds[v, k])

            @constraint(m, c10[v=ndds, p=patients], out_ndds[v, p] <= in_ndds[v, p])
            @constraint(m, c11[k=1:n], in_flow[k] + sum(in_ndds[v, k] for v in ndds) <= 1)
            @constraint(m, c13[k=1:n], out_flow[k] + sum(out_ndds[v, k] for v in ndds) <= 1)
            @constraint(m, c14[v=ndds, k=ndds], in_ndds[v, k] == 0)

            @constraint(m, c12[e=graph_edges], y[e] == sum(x_ndds[v, e] for v in ndds))
        else
            @constraint(m, c12[e=graph_edges], y[e] == 0)
        end

        time_param = input_data.time_param
        time_factor = input_data.time_factor
        new(m, x, z, in_flow, out_flow, n, max_cycle_length, max_chain_length, graph, time_param, time_factor, start_time)
    end
end

function solve(problem::Basic, signal::SharedArray{Bool,1}=SharedArray{Bool}(1))
    found = true
    cycles = []
    output = Output()
    while found
        if signal[1]
            println("Stopping solve")
            get_best_solution(problem, output)
            return output
        end
        for element in cycles
            cycle_length = element[1]
            cycle = element[2]
            # println(cycle_length, " - ", cycle)
            # @constraint(problem.m, sum(problem.z[e] for e in cycle) <= cycle_length-1)
            @constraint(problem.m, sum(problem.z[e] for e in cycle) <= cycle_length-1)
        end
        tf = time()
        remaining_time = max(timeout-(tf-problem.start_time), 1)*problem.time_factor
        remaining_time = convert(Int32, round(remaining_time, digits=0))
        set_parameter(problem.m, problem.time_param, remaining_time)
        optimize!(problem.m)
        # println(JuMP.objective_value(problem.m))
        if termination_status(problem.m) == MOI.TIME_LIMIT
            println("Basic timed out")
            return nothing
        end
        z_val = JuMP.value.(problem.z)
        found, cycles = check_cycle_lengths(z_val, problem.n, problem.max_cycle_length, problem.max_chain_length, problem.graph, true)  # only find the first violating cycle, seems to be faster
        # found, cycles = check_cycle_lengths(z_val, problem.n, problem.max_cycle_length, problem.max_chain_length, problem.graph, false)  # find all violating cycles in the current solution
    end
    return get_output(problem)
end

function get_output(problem::Basic)
    match_vertices = JuMP.value.(problem.in_flow)  # todo: correct to total in_flow
    match_edges = JuMP.value.(problem.z)
    match_cycles = nothing
    graph_cycles = nothing
    value = JuMP.objective_value(problem.m)
    optimal = true
    output = Output(match_vertices, match_edges, match_cycles, graph_cycles, value, optimal)
    return output
end

function get_best_solution(problem::Basic, output::Output)
    return
end

function print_values(x_val)
    for key in keys(x_val)
        value = x_val[key]
        if value > eps
            println(key, value)
        end
    end
end

function check_cycle_lengths(x_val, n::Int64, max_cycle_length::Int64, max_chain_length::Int64, graph::DiGraph, early_stop::Bool=false)
    # print_values(x_val)
    flag = zeros(Int64, n)
    cycles = []
    for i in 1:n
        if flag[i] != 0
            continue
        end
        current_vertex = i
        cycle_length = 0
        cycle = []
        while true
            flag[current_vertex] = 1
            next_vertex = find_recipient(current_vertex, x_val, graph)
            if next_vertex == 0 # no recipient, end of a chain
                if cycle_length > max_chain_length
                    push!(cycles, (cycle_length, cycle))
                    if early_stop
                        return true, cycles
                    end
                end
                break
            end
            push!(cycle, Edge((current_vertex, next_vertex)))
            flag[next_vertex] = 1
            cycle_length += 1
            if next_vertex == i # found a cycle
                if cycle_length > max_cycle_length
                    push!(cycles, (cycle_length, cycle))
                    if early_stop
                        return true, cycles
                    end
                end
                break
            end
            current_vertex = next_vertex
        end
    end
    return length(cycles)>0, cycles
end

function find_recipient(i::Int64, x_val, graph::DiGraph)
    for j in outneighbors(graph, i)
        if x_val[Edge((i,j))] > 0.5
            return j
        end
    end
    return 0
end


text:
# Functions for running single-locus non-additive and additive gwas models
function run_model(model)
    m1 = 1
    m2 = model.setsize
    while model.notcomplete
        if m2 >= model.m
            m2 = model.m
            model.notcomplete = false
        end
        curridx = 1
        snpset = get_bed_values(model.snps, m1, m2)
        snpset_maxbp = parse(Float64, snpset.markerInfo[end,2])
        makesubmodels = true
        while makesubmodels
            if snpset_maxbp <= model.layerlimit
                snpset = get_bed_values(model.snps, m1 + curridx - 1, m2)
                job = @spawn run_association(SubModel(size(snpset.markerInfo, 1), model.n, model.k, snpset, model.phenotypes, model.phenotypesfilename, nothing, nothing, nothing, nothing, model.outputfilename, model.printchain, model.submodelcount, model.weights, model.type))
                push!(model.futures, job)
                sleep(1)
                model.submodelcount += 1
                makesubmodels = false
            else
                bpidx_less = findall(x -> parse(Float64, x) <= model.layerlimit, snpset.markerInfo[curridx:end,2])
                try
                    snpset_less = get_bed_values(model.snps, m1 + curridx - 1, m1 + bpidx_less[end] - 1)
                    job = @spawn run_association(SubModel(size(snpset_less.markerInfo,1), model.n, model.k, snpset_less, model.phenotypes, model.phenotypesfilename, nothing, nothing, nothing, nothing, model.outputfilename, model.printchain, model.submodelcount, model.weights, model.type))
                    push!(model.futures, job)
                    sleep(1)
                    curridx = 1 + bpidx_less[end]
                    model.submodelcount += 1
                catch error
                    if !isa(error, BoundsError)
                        println("ERROR: Unknown error")
                        println(error)
                    else
                        println("BoundsError: Switching to next phenotypes interval")
                        println(model.submodelcount)
                        next_phenotype_interval!(model)
                        continue
                    end
                end

                if curridx <= (m2-m1+1)
                    # switch phenotype
                    next_phenotype_interval!(model)
                else
                    # you reach the end of setsize and still in same pheno, so move on to next set.
                    makesubmodels = false
                end
            end
        end
        m1 += model.setsize
        m2 += model.setsize
    end
end

function run_association(submodel)
    inititate_submodel!(submodel)
    if submodel.type == "additive"
        for i in 1:submodel.m
            testsnp = select_snp(submodel, i)
            calc_additive_only_stats!(submodel, testsnp, i)
        end
    else
        for i in 1:submodel.m
            testsnp = select_snp(submodel, i)
            calc_dominance_stats!(submodel, testsnp, i)
        end

    end
    calc_summaries(submodel)
    write_summaries(submodel)
end

function calc_dominance_stats!(submodel, testsnp, i)
    cellsums, cellcounts, zero_cells = calc_betas(testsnp.genotypes[:], submodel.phenotype.trait, submodel.n, submodel.weights.rinverse)
    if zero_cells
        submodel.output_stats[i,:] = [-9 -9 -9 -9 -9 -9 -9 -9 -9 -9]
        return
    end
    invcellcounts = 1 ./ cellcounts
    betahat = cellsums .* invcellcounts
    mse = diag(submodel.yty .- betahat'cellsums) ./ (submodel.n-3)
    sigma2 = invcellcounts*mse'
    sigma = sqrt.(sigma2)
    z = randn(submodel.rng, Float64, (3,submodel.k))
    beta_plaus = (z .* sigma) .+ betahat
    if submodel.printchain
	writedlm(submodel.output_filename .* "." .* submodel.snpset.markerInfo[i,3] .* ".tsv", beta_plaus', '\t')
    end
    beta = submodel.K*beta_plaus
    maf = sum(Array{Int64,1}(testsnp.genotypes[:]))/(2*submodel.n)
    aview = view(beta, 2, :)
    dview = view(beta, 3, :)
    mean_a = mean(aview)
    var_a = var(aview)
    mean_d = mean(dview)
    var_d = var(dview)
    submodel.output_stats[i,:] = [maf submodel.k mean_a var_a mean_d var_d 0 0 0 0]
end

function calc_additive_only_stats!(submodel, testsnp, i)
    X = [ones(submodel.n) testsnp.genotypes[:]]
    try
        invXTX = inv(X' * (submodel.weights.rinverse .* X))
    catch
        submodel.output_stats[i,:] = [maf submodel.k mean_a var_a 0 0]
        return
    end
    XTY = (submodel.weights.rinverse .* X)' * submodel.phenotype.trait
    betahat = invXTX * XTY
    mse = diag(submodel.yty .- betahat'XTY) ./ (submodel.n-2)
    sigma2 = diag(invXTX)*mse'
    sigma = sqrt.(sigma2)
    z = randn(submodel.rng, Float64, (2,submodel.k))
    beta = (z .* sigma) .+ betahat # 3 x k
    if submodel.printchain
        writedlm(submodel.output_filename .* "." .* submodel.snpset.markerInfo[i,3] .* ".tsv", beta_plaus', '\t')
    end
    maf = sum(Array{Int64,1}(testsnp.genotypes[:]))/(2*submodel.n)
    aview = view(beta, 2, :)
    mean_a = mean(aview)
    var_a = var(aview)
    submodel.output_stats[i,:] = [maf submodel.k mean_a var_a 0 0]
end

max_stars_repo_name: egmreynolds/MANA

max_issues_repo_name: egmreynolds/MANA

max_forks_repo_name: egmreynolds/MANA


text:
#==============================================================================
 = Julia Wrapper to the ros_interface.h Header
 =
 = Maintainer: Sidney Carvalho - sydney.rdc@gmail.com
 = Last Change: 2020 Abr 20 17:46:26
 = Info: Maps all functions of ros_interface into native julia code.
 =============================================================================#

module ROS_INTERFACE

using Cxx           # package to wrapper C++ code (demands Julia v0.6 or later)

# default C++ lib and include paths
global const path_to_lib = "../lib/"
global const path_to_include = "../include/"

# define the path_to_include as default C++ headers folder
addHeaderDir(path_to_include, kind=C_System)

# load ROS interface shared library
Libdl.dlopen(path_to_lib * "libros_interface.so", Libdl.RTLD_GLOBAL)

# include cpp headers
cxxinclude("ros_interface.hpp")                 # ROS interface header
cxxinclude("vector")                            # C++ vector header

# public functions
export ros_interface, add_bot, set_vel, get_pose, get_laser, clock, ros_ok, shutdown, get_time

# initialize ROS communication interface
function ros_interface(node_name::String)
    # instantiate a 'ros_interface' object
    global ros_com = @cxxnew ros_interface(pointer(node_name))
end

# add robots to the interface
function add_bot(bot_id::Int64, bot_type::Int64, bot_name::String, bot_pose::Array{Float64, 1})
    # position array of ros_interface
    ros_pose = @cxxnew space_t()

    # set robot initial position at ROS
    @cxx ros_pose->set_x(bot_pose[1])
    @cxx ros_pose->set_y(bot_pose[2])
    @cxx ros_pose->set_yaw(bot_pose[3])

    # add the robot to the ROS environment
    @cxx ros_com->add_node(bot_id, bot_type, pointer(bot_name), ros_pose)

end

# set the velocity of a specific node
function set_vel(bot_id::Int64, bot_ctrl::Array{Float64, 1})
    # control array of ros_interface
    ros_vel = @cxxnew space_t()

    # translate julia array to ros_interface datatype
    @cxx ros_vel->set_x(bot_ctrl[1])
    @cxx ros_vel->set_y(bot_ctrl[2])
    @cxx ros_vel->set_yaw(bot_ctrl[3])

    # send control actions to ROS
    @cxx ros_com->node_vel(bot_id, ros_vel)

end

# get the position of a specific node
function get_pose(bot_id::Int64)
    # get robot position from ROS
    ros_pose = @cxx ros_com->node_pose(bot_id)

    # translate data from type on ros_interface to julia
    return [@cxx ros_pose->x; @cxx ros_pose->y; @cxx ros_pose->yaw]
end

# get the laser readings of a specific node
function get_laser()
    # TODO: finish this
end

# set data sampling time (seconds) and execute callback functions (must be included in the main loop)
function clock(dt::Float64)
    @cxx ros_com->clock(dt)
end

# check if ros is running
function ros_ok()
    return @cxx ros_com->ros_ok()
end

# remove the instance of ros_interface and all its nodes
function shutdown()
    @cxx ros_com->shutdown()
end

# get ROS current time (seconds)
function get_time()
    return @cxx ros_com->get_time()
end

end


text:
function show(io::IO, p::TAMQPFrameProperties)
    print(io, "Channel $(p.channel), Size $(p.payloadsize) bytes")
end

function show(io::IO, p::TAMQPMethodPayload)
    print(io, displayname(p.class, p.method))
end

function show(io::IO, m::TAMQPMethodFrame)
    print(io, "MethodFrame ", m.payload, "(", length(m.payload.fields), " fields...)")
    if isa(io, IOContext)
        if !((:limit => true) in io)
            print(io, '\n')
            show(io, m.payload.fields)
        end
    end
end

function show(io::IO, f::TAMQPFieldValue)
    show(io, f.fld)
end

function show(io::IO, f::TAMQPFieldValuePair)
    indent = isa(io, IOContext) ? get(io, :indent, "") : ""
    print(io, indent)
    show(io, f.name)
    print(io, " => ")
    show(io, f.val)
end

function show(io::IO, f::TAMQPFieldTable)
    indent = isa(io, IOContext) ? get(io, :indent, "") : ""
    println(io, "FieldTable")
    ioc = IOContext(io, :indent => (indent * "    "))
    idx = 1
    for fpair in f.data
        (idx > 1) && print(ioc, '\n')
        show(ioc, fpair)
        idx += 1
    end
end

function show(io::IO, s::T) where {T<:Union{TAMQPShortStr,TAMQPLongStr}}
    print(io, convert(String, s))
end

function show(io::IO, fields::Vector{Pair{Symbol,TAMQPField}})
    indent = isa(io, IOContext) ? get(io, :indent, "") : ""
    println(io, indent, "Fields:")
    indent = indent * "    "
    ioc = IOContext(io, :indent => indent)
    idx = 1
    for fld in fields
        (idx > 1) && print(ioc, '\n')
        print(ioc, indent)
        show(ioc, fld.first)
        print(ioc, " => ")
        show(ioc, fld.second)
        idx += 1
    end
end


text:
function foo(username)
  @join begin
    account = login(username)
    last = getlastlogin(account)
    dms = getdms(account)
    who = getuser(account, dms[end])
    followers = publicinfo(username)[:followers]
  end

  # Do something with followers, last, who etc.
end

followers = @fork publicinfo(username)[:followers]
account = @fork login(username)
last = @fork getlastlogin(fetch(account))
dms = getdms(fetch(account))
who = getuser(fetch(account), dms[end])
fetchall(last, who, followers)


text:
# Analyze the data using Lasso Path pacakge I wrote.
cd("C:\\Users\\victo\\source\\repos\\Silly_Julia_stuff\\Julia Packages\\LassoPath\\usage_examples\\crime_data")
using DelimitedFiles, CSV, DataFrames
using LassoPath
using BenchmarkTools

Df = CSV.read("process_data.csv", DataFrame)
Df = Df[!, Not([:state, :fold])]

# Extrac Numerics, predictors, predictants.
A, b = Matrix(Df[:, Not(end)]), Df[:, end]

PrintTitle("The Data is ready, we are ready for analyais now")

function ShowResults(constructor::Function, fname::String)
    lassoIntance = constructor(A, b)
    lassoIntance.Tol = -1
    GetLassoPath(lassoIntance, 0.99)
    VisualizeLassoPath(lassoIntance, fname=fname)
    predictorIndices, λ, ws =
        CaptureImportantWeights(lassoIntance, 10, 1e-2)
    PrintTitle("Important Predictors are")
    predictors = names(Df[:, Not(end)])[predictorIndices]
    impacts = map(ws) do x
        if x >= 0
            return "+"
        else
            return "-"
        end
    end
    display([reshape(predictors, :) reshape(impacts, :)])
end

# @time ShowResults((x, y) -> LassoSCOP(x, y), "Lasso SCOP.png")
@time ShowResults((x, y) -> Lasso2NormProx(x, y), "lasso_proximal.png")


text:
module FieldDocTables

using DocStringExtensions, PrettyTables
import DocStringExtensions: Abbreviation, format

export FieldDocTable

struct FieldDocTable{L,T,TR,F} <: Abbreviation
    labels::L
    functions::T
    truncation::TR
    tableformat::F
    fenced::Bool
end
function FieldDocTable(
   labels::L, functions::T;
   truncation=((100 for f in functions)...,), tf=tf_markdown, fenced=false
) where {L,T}
    FieldDocTable{L,T,typeof(truncation), typeof(tableformat)}(
        labels, functions, truncation, tableformat, fenced
    )
end
FieldDocTable(nt::NamedTuple; kwargs...) = FieldDocTable(keys(nt), (nt...,); kwargs...)

function format(doctable::FieldDocTable, buf, doc)
    local docs = get(doc.data, :fields, Dict())
    local binding = doc.data[:binding]
    local object = Docs.resolve(binding)
    # On 0.7 fieldnames() on an abstract type throws an error. We then explicitly return
    # an empty vector to be consistent with the behaviour on v0.6.
    local fields = isabstracttype(object) ? Symbol[] : fieldnames(object)

    if !isempty(fields)

        # Fieldnames and passed in functions
        colnames = [:Field, doctable.labels...]
        funcdata = ([safestring.(f(object), doctable.truncation[i])...] for (i, f) in enumerate(doctable.functions))
        data = hcat([fields...], funcdata...)

        # Only add a fielddocs column if there are fielddocs
        fielddocs = []
        for field in fields
            if haskey(docs, field) && isa(docs[field], AbstractString)
                for line in split(docs[field], "\n")
                    fielddoc = isempty(line) ? "" : rstrip(line)
                end
            else
                fielddoc = ""
            end
            push!(fielddocs, fielddoc)
        end
        if any(d -> d != "", fielddocs)
            data = hcat(data, fielddocs)
            colnames = [colnames..., :Docs]
        end

        println(buf)
        doctable.fenced && println(buf, "```")
        pretty_table(buf, data, colnames; tf=doctable.tableformat)
        doctable.fenced && println(buf, "```")
        println(buf)
    end
end

safestring(::Nothing, n) = "nothing"
safestring(s, n) = truncate_utf8(string(s), n)

# Is there simpler way to do this?
function truncate_utf8(s, n)
    eo = lastindex(s)
    neo = 0
    for i = 1:n
      if neo < eo
          neo = nextind(s, neo)
      else
          break
      end
    end
    SubString(s, 1, neo)
end

end # module


text:
using DataFrames
using DataStructures: OrderedDict
using JSON
using Random

function generate_data1()
    Random.seed!(0)
    n = 10
    df = DataFrame([rand(1:10, n), 0.5 .+ randn(n), 10 .+ 3randn(n), rand(["X", "Y", "Z"], n), rand(Int8[0,1], n), zeros(Bool, n)], [:A1, :A2, :B, :C, :Da, :D2])
    df.D2 .= 1 .- df.Da
    df
end

function transform_data1(X, sid=:default)
    dsrc = datasource(X; sid=sid)
    re = replacevals(:C => dsrc, "X" => "x", "Z" => "z")
    oh = onehot(:C => re)
    re2 = replacevals(:C_Y => oh, false => "no", true => "yes")
    oc = onecold([:Da => dsrc, :D2 => dsrc], :D, OrderedDict(:Da => "a", :D2 => 2))
    sd = standardize(:A1 => dsrc)
    sd2 = standardize(:B => dsrc)
    sd3 = standardize(:C_z => oh)
    ddst = mergecols(:A1 => sd, :A2 => dsrc, :B => sd2, :C_Y => re2, :C_z => sd3, :C_x => oh, :D => oc)
    dsrc, ddst
end

function generate_data2()
    Random.seed!(0)
    n = 10
    DataFrame([rand(0:100, n), rand(["M", "F"], n)], [:O, :S])
end

@testset "transform and inverse" begin
    X = generate_data1()
    dsrc, ddst = transform_data1(X, :sid1)
    smap = Dict(:sid1 => X)
    fit!(ddst, smap)

    Y = transform(ddst, smap, DataFrame)
    T = DataFrame(
        A1 = [-1.1423471897815947, -0.5565281180987257, 1.4938386327913158, 0.32220048942557783, -1.1423471897815947, -1.1423471897815947, -0.2636185822572912, 1.2009290969498814, 0.9080195611084468, 0.32220048942557783],
        A2 = [-1.1072563241277753, -1.9807927306599402, 2.7762328327845243, 0.7196934675425409, 0.3828624692670648, -0.1012535941853836, 1.6422764389413194, 0.41138368179952706, 0.7794662520898984, 0.6114216780915173],
        B = [-1.0510047730459515, 0.28255507796868556, 0.004359925255237256, -1.700136355232938, 1.8049184082894767, 0.17779513965723276, 0.46682135785382906, -0.9522129757959658, 0.6311744633370341, 0.33572973171335546],
        C_Y = ["no", "no", "no", "no", "yes", "no", "yes", "no", "no", "no"],
        C_z = [1.1618950038622249, -0.7745966692414833, -0.7745966692414833, 1.1618950038622249, -0.7745966692414833, 1.1618950038622249, -0.7745966692414833, 1.1618950038622249, -0.7745966692414833, -0.7745966692414833],
        C_x = Bool[false, true, true, false, false, false, false, false, true, true],
        D = Any["a", "a", "a", "a", 2, "a", "a", "a", "a", "a"]
    )
    @test Y ≈ T

    smap2 = Dict(:sid2 => Y)
    dsrc2 = datasource(Y; sid=:sid2)
    ddst2 = invert(ddst, dsrc2)
    Xt = transform(ddst2, smap2, DataFrame)
    @test X ≈ Xt

    @test dsrc == getdatasource(ddst)
end

@testset "tonumerical" begin
    X = generate_data1()
    ddst = tonumerical(X)
    fit!(ddst, X)
    Ym = transform(ddst, X, Matrix{Float64})
    dsrcm = datasource(Ym, outids(ddst))
    ddstm = invert(ddst, dsrcm)
    Xtm = transform(ddstm, Ym, DataFrame)
    @test X ≈ Xtm

    Xtm2 = invtransform(ddst, Ym, DataFrame)
    @test X ≈ Xtm2
end

@testset "serialization" begin
    X = generate_data1()
    dsrc, ddst = transform_data1(X, :sid1)
    smap = Dict(:sid1 => X)
    fit!(ddst, smap)

    str = JSON.json(ddst, 3)
    dict = JSON.parse(str, dicttype = OrderedDict)
    ddst2 = deserialize(dict)

    Y = transform(ddst, smap, DataFrame)
    Y2 = transform(ddst2, smap, DataFrame)

    @test Y ≈ Y2
end

max_stars_repo_name: dhonza/MLGoodies


text:
# Provides functions to parse model equations. Array access syntax is used to indicate the time period of functions.

# List of mathematical operators supported in model equations
operators = [:-
            :+
            :*
            :/
            :^
            :exp
            :log
            :Expect
            :abs
            :max
            :min
            :∫
            :<
            :>
            :≤
            :≥]

"""
    subs!(x,d)

Replaces instances of the keys of dictionary d with the values in expression x.
"""
function subs!(x::Expr, d::Dict)
    for i = 1:length(x.args)
        if in(x.args[i], keys(d))
            x.args[i] = d[x.args[i]]
        elseif isa(x.args[i], Expr)
            subs!(x.args[i], d)
        end
    end
end

function subs(x1::Expr, d::Dict)
    x = deepcopy(x1)
    subs!(x, d)
    return x
end

subs!(x::Expr, p::Pair) = subs!(x, Dict(p))
subs(x::Expr, p::Pair) = subs(x, Dict(p))

# function subs!(x::Expr,s::Pair)
#     for i = 1:length(x.args)
#         if x.args[i]==s.first
#             x.args[i] = s.second
#         elseif isa(x.args[i],Expr)
#             subs!(x.args[i],s)
#         end
#     end
# end
#
# function subs(x1::Expr,s::Pair)
#     x = deepcopy(x1)
#     subs!(x,s)
#     return x
# end

"""
    addindex!(x)

Adds a time index [0] to unindexed variables in x.
e.g. c^-2.5 -> c[0]^-2.5
"""
function addindex!(x, ignore = operators)
    if typeof(x) == Expr
        if x.head == :ref
            return x
        else
            for i = 1:length(x.args)
                x.args[i] = addindex!(x.args[i], ignore)
            end
        end
    elseif typeof(x) == Symbol
        if !in(x, ignore)
            x = :($x[0])
            return x
        end
    end
    return x
end

addindex(x, ignore = operators) = addindex!(deepcopy(x), ignore)

"""
    removeindex!(x)
Removes all time indices from model expression x.
"""
function removeindex!(x)
    if isa(x, Expr)
        if x.head == :ref
            return x.args[1]
        else
            for i = 1:length(x.args)
                x.args[i] = removeindex!(x.args[i])
            end
        end
    end
    return x
end

removeindex(x, ignore = operators) = removeindex!(deepcopy(x), ignore)

"""
    tchange!(x,t)
Either shifts the time indices in model expression x forward or backwards depending on
whether t = 1 or t = -1 respectively.
"""
function tchange!(x::Expr, t::Int, ignore = operators)
    t != -1 && t != 1 && return
    if x.head == :ref
        if (x.args[2] == -1 && t == 1) || (x.args[2] == 1 && t == -1)
            x = :($(x.args[1])[0])
        elseif x.args[2] == 0
            x.args[2] = t
        end
    else
        for i = 1:length(x.args)
            if isa(x.args[i], Expr)
                x.args[i] = tchange!(x.args[i], t, ignore)
            end
        end
    end
    x
end

tchange(x::Expr, t::Int, ignore = operators) = tchange!(deepcopy(x), t, ignore)

"""
    getv(x)
Returns all symbols in the expression x, excluding mathematical operators.
"""
function getv(x, list = Expr[], ignore = operators)
    if isa(x, Expr)
        if x.head == :ref
            push!(list, x)
        else
            for i = 1:length(x.args)
                list = getv(x.args[i], list, ignore)
            end
        end
    end
    return unique(list)
end

# function vecind!(x::Expr,n::Int)
#     if x.head == :ref
#         if length(x.args)==3
#             x.args[2]=x.args[2]+(pop!(x.args)-1)*n
#         end
#     else
#         for i = 1:length(x.args)
#             if isa(x.args[i],Expr)
#                 x.args[i] = vecind!(x.args[i],n)
#             end
#         end
#     end
#     x
# end

"""
    simplifyindices!(x)
Simplifies expressions in array access portions of an arbitrary expression x.
"""
function simplifyindices!(x)
    if isa(x, Expr)
        if x.head == :ref
            for i = 2:length(x.args)
                x.args[i] = simplify(x.args[i])
            end
        else
            for a in x.args
                a = simplifyindices!(a)
            end
        end
    end
    return x
end

"""
    addpweights!(x,nP)
Replaces instances of ProbWeights*(...) in model equations x with a weighted sum over the probabilities of future realisations of the models stochastic variables. nP defines the number of possible future realisations.
"""
function addpweights!(x, nP)
    if typeof(x) == Expr
        if (x.head == :call) && (x.args[1] == :*) && (x.args[2] == :ProbWeights)
            return sum([subs(x.args[3], :j => j) * :(M.ProbWeights[i,$j]) for j = 1:nP])
        # return :(sum($(x.args[3])*M.ProbWeights[i,j] for j = 1:$nP))
        else
            for i = 1:length(x.args)
                x.args[i] = addpweights!(x.args[i], nP)
            end
        end
    end
    return x
end

addpweights(x, nP) = addpweights!(deepcopy(x), nP)

"""
    findrepeated(x)
Returns a list of repeated sub-expressions in the model equations x.
"""
function findrepeated(x, list = Dict())
    if isa(x, Expr) && x.head != :ref
        if in(x, keys(list))
            list[x] = list[x] + 1
        else
            list[x] = 1
        end

        for i = 1:length(x.args)
            list = findrepeated(x.args[i], list)
        end
    end
    return (list)
end

function Base.length(ex, n = 0)
    if isa(ex, Expr)
        n += sum(map(length, ex.args))
    else
        n += 1
    end
    return n
end

# The following section allows arithmetic between expressions, symbols and numbers simplifying julia function generation from the parsed model equations.

import Base:+,-,*,/,vec

# When applied to an expression that defines a matrix will return the elements as an expression defining the column vector of all elements.
function vec(ex::Expr)
    nex = :([])
    nc = length(ex.args)
    nr = length(ex.args[1].args)
    @assert ex.head == :vcat

    for r in ex.args
        @assert r.head == :row && length(r.args) == nr
    end

    for ri = 1:nr
        for ci = 1:nc
            push!(nex.args, ex.args[ci].args[ri])
        end
    end
    return nex
end

+{T<:Union{Expr, Symbol, Number}}(x::T) = x
+{T<:Union{Expr, Symbol, Number}, S <: Union{Expr, Symbol, Number}}(x1::T, x2::S) = :(+($x1, $x2))

function +{T<:Union{Expr, Symbol}, S <: Float64}(x1::T, x2::Vector{S})
    out = :([])
    for i = 1:length(x2)
        push!(out.args, simplify(x2[i] + x1))
    end
    return out
end

function +{T <: Union{Expr, Symbol}, S <: Float64}(x2::Vector{S}, x1::T)
    out = :([])
    for i = 1:length(x2)
        push!(out.args, simplify(x2[i] + x1))
    end
    return out
end

-{T<:Union{Expr, Symbol, Number}}(x::T) = -1 * x

function -{T <: Union{Expr, Symbol, Number}, S <: Union{Expr, Symbol, Number}}(x1::T, x2::S)
    :(-($x1, $x2))
end

function -{T <: Union{Expr, Symbol}, S <: Float64}(x1::T, x2::Vector{S})
    out = :([])
    for i = 1:length(x2)
        push!(out.args, simplify(x2[i] - x1))
    end
    return out
end

function -{T <: Union{Expr, Symbol}, S <: Float64}(x2::Vector{S}, x1::T)
    out = :([])
    for i = 1:length(x2)
        push!(out.args, simplify(x2[i] - x1))
    end
    return out
end

function *{T <: Union{Expr, Symbol, Number}, S <: Union{Expr, Symbol, Number}}(x1::T, x2::S)
    :(*($x1, $x2))
end

function *{T <: Union{Expr, Symbol}, S <: Float64}(x1::T, x2::Vector{S})
    out = :([])
    for i = 1:length(x2)
        push!(out.args, simplify(x2[i] * x1))
    end
    return out
end

function *{T <: Union{Expr, Symbol}, S <: Float64}(x2::Vector{S}, x1::T)
    out = :([])
    for i = 1:length(x2)
        push!(out.args, simplify(x2[i] * x1))
    end
    return out
end

function /{T <: Union{Expr, Symbol, Number}, S<:Union{Expr, Symbol, Number}}(x1::T, x2::S)
    :(/($x1, $x2))
end

function /{T <: Union{Expr, Symbol}, S <: Float64}(x1::T, x2::Vector{S})
    out = :([])
    for i = 1:length(x2)
        push!(out.args, simplify(x1 / x2[i]))
    end
    return out
end

function /{T <: Union{Expr, Symbol}, S <: Float64}(x2::Vector{S}, x1::T)
    out = :([])
    for i = 1:length(x2)
        push!(out.args, simplify(x2[i] / x1))
    end
    return out
end

max_stars_repo_name: ZacLN/HADSGE

max_issues_repo_name: ZacLN/HADSGE

max_forks_repo_name: ZacLN/HADSGE


text:
# This file includes modified code from https://github.com/odow/SDDP.jl as at
# commit a6662eecff6d5de5c499b760773a8569a0979b9b.

#  Copyright 2017, Oscar Dowson and contributors
#  This Source Code Form is subject to the terms of the Mozilla Public
#  License, v. 2.0. If a copy of the MPL was not distributed with this
#  file, You can obtain one at http://mozilla.org/MPL/2.0/.
#############################################################################

#=
    The air conditioning example from Anthony Papavasiliou
    https://perso.uclouvain.be/anthony.papavasiliou/public_html/SDDP.pdf

    Consider the following problem
        Produce air conditioners for 3 months
        200 units/month at 100 $/unit
        Overtime costs 300 $/unit
        Known demand of 100 units for period 1
        Equally likely demand, 100 or 300 units, for periods 2, 3
        Storage cost is 50 $/unit
        All demand must be met

    Optimal bound $62,500
=#
using SDDP, JuMP, GLPKMathProgInterface, Base.Test, SDDiP

function airconditioningmodel()
    m = SDDPModel(
                 stages = 3,
                 # example for issue #64
        objective_bound = (t, i) -> 0.0,
                  sense = :Min,
                 # solver = GLPKSolverMIP()
                 solver = GLPKSolverMIP()
                            ) do sp, stage
        # number of units
        @binarystate(sp, 0 <= stored_production <= 900, incoming_storage == 0, Int)
        # @state(sp, stored_production >= 0, incoming_storage == 0)
        @variables(sp, begin
            # number of units produced during normal production
            # circumstances
            0 <= production <= 200
            # number of units produced during overtime production
            # circumstances
            overtime   >= 0
        end)
        # balance constraints
        if stage == 1
            @constraint(sp, stored_production == incoming_storage + production + overtime - 100)
        else
            @rhsnoise(sp, demand=[100, 300], stored_production == incoming_storage + production + overtime - demand)
        end
        @stageobjective(sp, 100 * production + 300 * overtime + 50 * stored_production)
        # Set integer solver
        setSDDiPsolver!(sp, method=KelleyMethod(), LPsolver=GLPKSolverLP(), pattern=Pattern(lagrangian=5, benders=1))
    end
end

srand(1234)
m = airconditioningmodel()
solve(m, max_iterations=10)
@test isapprox(getbound(m), 62_500.0)


text:
using PyCall
if PyCall.pyversion.major != 3 || PyCall.pyversion.minor != 7
    using Conda
    const Pkg = Base.require(Base.PkgId(Base.UUID("44cfe95a-1eb2-52ea-b672-e2afdf69b78f"), "Pkg"))
    Conda.add("python=3.7")
    try
        Conda.rm("numpy") #get around MKL problem
    catch e
    end
    Conda.add("nomkl")
    Conda.add("numpy")
    Conda.add_channel("rmg")
    Conda.add("rmg")
    Pkg.build("PyCall")
end


text:
push!(LOAD_PATH,"../src/")

using Documenter
using Jack

makedocs(
    sitename = "Jack.jl",
    authors = "Stéphane Laurent",
    format = Documenter.HTML(;
        prettyurls = get(ENV, "CI", "false") == "true",
        canonical = "https://stla.github.io/Jack.jl",
        assets = String[],
    ),
    modules = [Jack],
    pages = ["Documentation"  => "index.md"],
    repo = "https://github.com/stla/Jack.jl/blob/{commit}{path}#{line}"
)

deploydocs(;
    branch = "gh-pages",
    devbranch = "main",
    repo = "github.com/stla/Jack.jl",
)


text:
module BedsData

using CSV
using DataFrames

export n_beds

basepath = normpath(@__DIR__, "../../")

function n_beds(locations::Array; level::Symbol=:state, source::Symbol=:definitivehc, bed_type::Symbol=:all, pct_beds_available::Real=0.25)
	if     level == :state    && source == :definitivehc
		beds_data = load_definitivehc()
		beds_data = definitivehc_filter!(beds_data, locations, level=level)
		beds_data = definitivehc_select_type!(beds_data, bed_type=bed_type)
		beds = definitivehc_aggregate(beds_data)
	elseif level == :state    && source == :gov
		@assert bed_type == :all
		beds_data = load_gov()
		beds = compute_beds_gov(beds_data, locations)
	elseif level == :county   && source == :definitivehc
		beds_data = load_definitivehc()
		beds_data = definitivehc_filter!(beds_data, locations, level=level)
		beds_data = definitivehc_select_type!(beds_data, bed_type=bed_type)
		beds = definitivehc_aggregate(beds_data, locations=locations)
	elseif level == :hospital && source == :definitivehc
		beds_data = load_definitivehc()
		beds_data = definitivehc_filter!(beds_data, locations, level=level)
		beds_data = definitivehc_select_type!(beds_data, bed_type=bed_type)
		sort!(beds_data, :selected_location)
		beds = beds_data.selected_beds
	else
		error("Invalid parameters to compute_beds.")
	end

	return beds * Float32(pct_beds_available)
end

########################
###### data.gov ########
########################

function load_gov()
	beds_data = DataFrame(CSV.File(joinpath(basepath, "data/hospitals/gov.csv")), copycols=true)

	filter!(row -> row.BEDS > 0, beds_data)
	filter!(row -> row.TYPE == "GENERAL ACUTE CARE", beds_data)

	beds_data = combine(groupby(beds_data, :STATE), :BEDS => sum)
	sort!(beds_data, :STATE)

	return beds_data
end

function compute_beds_gov(beds_data::DataFrame, states::Array{String,1})
	@assert states == sort(states)
	beds_df = filter(row -> row.STATE in states, beds_data)
	return Float32.(beds_df.BEDS_sum)
end

###########################
## Definitive Healthcare ##
###########################

function load_definitivehc()
	_beds_data = DataFrame(CSV.File(joinpath(basepath, "data/hospitals/definitivehc.csv")))
	beds_data = filter(row -> !(row.HOSPITAL_TYPE in ["Psychiatric Hospital", "Rehabilitation Hospital"]), _beds_data)
	return beds_data
end

function definitivehc_filter!(beds_data::DataFrame, locations::Array; level::Symbol=:state)
	if level == :state
		filter!(row -> !ismissing(row.HQ_STATE), beds_data)
		beds_data.selected_location = convert(Array{Union{String,Missing}}, deepcopy(beds_data.HQ_STATE))
	elseif level == :county
		filter!(row -> !ismissing(row.FIPS), beds_data)
		beds_data.selected_location = convert(Array{Union{Int,Missing}}, deepcopy(beds_data.FIPS))
	elseif level == :hospital
		filter!(row -> !ismissing(row.FIPS), beds_data)
		beds_data.selected_location = convert(Array{Union{Int,Missing}}, deepcopy(beds_data.FIPS))
		return beds_data
	else
		error("Invalid parameters.")
	end
	filter!(row -> row.selected_location in locations, beds_data)
	return beds_data
end

function definitivehc_select_type!(beds_data::DataFrame; bed_type::Symbol=:all)
	bed_type_cvt = Dict(
		:all       => :NUM_STAFFED_BEDS,
		:staffed   => :NUM_STAFFED_BEDS,
		:licensed  => :NUM_LICENSED_BEDS,
		:icu       => :NUM_ICU_BEDS,
		:adult_icu => :ADULT_ICU_BEDS,
		:regular   => :NUM_NON_ICU_BEDS,
	)
	col = bed_type_cvt[bed_type]

	if bed_type == :regular
		beds_data.NUM_NON_ICU_BEDS = beds_data.NUM_STAFFED_BEDS - beds_data.NUM_ICU_BEDS
	end

	filter!(row -> !ismissing(row[col]), beds_data)
	filter!(row -> row[col] > 0, beds_data)
	beds_data.selected_beds = convert(Array{Union{Int,Missing}}, deepcopy(beds_data[:,col]))
	return beds_data
end

function definitivehc_aggregate(beds_data::DataFrame; locations::Array=[])
	beds_data_agg = combine(groupby(beds_data, :selected_location), :selected_beds => sum => :selected_beds_agg)
	if !isempty(locations)
		missing_locs = setdiff(locations, beds_data_agg.selected_location)
		for loc in missing_locs
			push!(beds_data_agg, Dict(:selected_location => loc, :selected_beds_agg => 0), cols=:subset)
		end
	end
	sort!(beds_data_agg, :selected_location)
	return Float32.(beds_data_agg.selected_beds_agg)
end

end;

max_stars_repo_name: flixpar/covid-resource-allocation

max_issues_repo_name: flixpar/covid-resource-allocation

max_forks_repo_name: flixpar/covid-resource-allocation


text:
function feasibility_search!(
    solution::Solution,
    mapf::MAPF;
    neighborhood_size=1,
    conflict_price=Inf,
    conflict_price_increase=0.0,
    progress=true,
)
    A = nb_agents(mapf)
    pathless_agents = shuffle([a for a in 1:A if length(solution[a]) == 0])
    cooperative_astar!(solution, pathless_agents, mapf; conflict_price=conflict_price)
    cp = colliding_pairs(solution, mapf)
    prog = ProgressUnknown("LNS2 steps: "; enabled=progress)
    while !is_feasible(solution, mapf)
        next!(prog; showvalues=[(:colliding_pairs, cp)])
        neighborhood_agents = random_neighborhood(mapf, neighborhood_size)
        backup = remove_agents!(solution, neighborhood_agents)
        cooperative_astar!(
            solution, neighborhood_agents, mapf; conflict_price=conflict_price
        )
        new_cp = colliding_pairs(solution, mapf)
        if is_feasible(solution, mapf) || (new_cp <= cp)  # keep
            cp = new_cp
        else  # revert
            for a in neighborhood_agents
                solution[a] = backup[a]
            end
        end
        conflict_price *= (1 + conflict_price_increase)
    end
    return solution
end

function feasibility_search(
    mapf::MAPF;
    neighborhood_size=1,
    conflict_price=Inf,
    conflict_price_increase=0.0,
    progress=true,
)
    solution = independent_dijkstra(mapf)
    feasibility_search!(
        solution,
        mapf;
        neighborhood_size=neighborhood_size,
        conflict_price=conflict_price,
        conflict_price_increase=conflict_price_increase,
        progress=progress,
    )
    return solution
end


text:
using DashBootstrapComponents

dropdown = html_div(
    [
        dbc_label("Dropdown", html_for = "dropdown"),
        dcc_dropdown(
            id = "dropdown",
            options = [
                Dict("label" => "Option 1", "value" => 1),
                Dict("label" => "Option 2", "value" => 2),
            ],
        ),
    ],
    className = "mb-3",
);

slider = html_div(
    [
        dbc_label("Slider", html_for = "slider"),
        dcc_slider(id = "slider", min = 0, max = 10, step = 0.5, value = 3),
    ],
    className = "mb-3",
);

range_slider = html_div(
    [
        dbc_label("RangeSlider", html_for = "range-slider"),
        dcc_rangeslider(id = "range-slider", min = 0, max = 10, value = [3, 7]),
    ],
    className = "mb-3",
);

form = dbc_form([dropdown, slider, range_slider]);

max_issues_repo_name: sthagen/dash-bootstrap-components

max_forks_repo_name: sthagen/dash-bootstrap-components


text:
# Nonlinear constraints (objective not supported)

MOI.supports(::Optimizer, ::MOI.NLPBlock) = true

function MOI.set(o::Optimizer, ::MOI.NLPBlock, data::MOI.NLPBlockData)
    # We don't actually store the data (to be queried later during the solve
    # process). Instead, we extract the expression graphs and add the
    # corresponding constraints to the model directly.
    if data.has_objective
        error("Nonlinear objective not supported by SCIP.jl!")
    end

    # go back to problem stage
    allow_modification(o::Optimizer)

    MOI.initialize(data.evaluator, [:ExprGraph])
    for i in 1:length(data.constraint_bounds)
        expr = MOI.constraint_expr(data.evaluator, i)
        bounds = data.constraint_bounds[i]
        cr = add_nonlinear_constraint(o.mscip, expr, bounds.lower, bounds.upper)
        # Not registering or returning constraint reference!
    end

    return nothing
end


text:
using SimpleGraphs, SimpleGraphAlgorithms

const DEFAULT_FILE_NAME = "tree_codes.jl"

# initialize the table of unlabeled trees

"""
    init_trees_table()
Create a new table of distinct trees on 1 and 2 vertices.
"""
function init_trees_table()
    TT = Dict{Int,Vector{SimpleGraph{Int}}}()

    TT[1] = [IntGraph(1)]
    T = IntGraph(2)
    add!(T, 1, 2)
    TT[2] = [T]

    return TT
end

"""
    check_in(G,S)
See if the set `S` contains a graph isomorphic to `G`. Return `true` if so.
"""
function check_in(G::SimpleGraph{Int}, S::Set{SimpleGraph{Int}})
    if isempty(S)
        return false
    end
    for H ∈ S
        if uhash(G) == uhash(H) #is_iso(G, H)
            return true
        end
    end
    return false
end

"""
    extend_trees_table!(TT)
Given a table of distinct trees up to size `n`, extend that table to include
all distinct trees of size `n+1`.
"""
function extend_trees_table!(TT::Dict{Int,Vector{SimpleGraph{Int}}})::Nothing
    n = maximum(keys(TT))
    outset = Set{SimpleGraph{Int}}()  # set of trees with n+1 vertices
    for T ∈ TT[n]
        for w = 1:n
            X = deepcopy(T)
            add!(X, w, n + 1)
            if check_in(X, outset)
                continue
            end
            push!(outset, X)
        end
    end

    TT[n+1] = collect(outset)

    @info "Added $(length(TT[n+1])) new trees with $(n+1) vertices"
    nothing
end

function build_trees_table(nmax::Int)
    TT = init_trees_table()
    while maximum(keys(TT)) < nmax
        @info "Adding trees of size $(maximum(keys(TT))+1)"
        extend_trees_table!(TT)
    end
    return TT
end

"""
    create_codes_table(TT)
Given a table of distinct trees, convert that into a table of Prufer codes.
This is used by `save_trees_table` and not useful to be called directly.
"""
function create_codes_table(TT::Dict{Int64,Vector{SimpleGraph{Int}}})
    codes = Dict{Int64,Vector{Vector{Int}}}()
    codes[2] = [Int[]]
    for n = 3:maximum(keys(TT))
        codes[n] = prufer_code.(TT[n])
    end
    return codes
end

"""
    save_tree_table(TT, filename)
Save a trees table into a file specified by `filename`.
If the file name is omitted, use `codes.jl`.
"""
function save_trees_table(
    TT::Dict{Int64,Vector{SimpleGraph{Int}}},
    filename::String = DEFAULT_FILE_NAME,
)
    outfile = open(filename, "w")
    codes = create_codes_table(TT)
    println(outfile, "codes = Dict{Int,Vector{Vector{Int}}}()")
    for n = 2:maximum(keys(TT))
        print(outfile, "codes[$n] = ")
        println(outfile, codes[n])
    end
    close(outfile)
    nothing
end

"""
    load_trees_table(filename)
Create a table of distinct trees by reading in a file that has been
precomputed (and presumably saved using `save_tree_table`). If `filename`
is omitted, use `codes.jl`.
"""
function load_trees_table(filename::String = DEFAULT_FILE_NAME)
    include(filename)
    TT = init_trees_table()
    nmax = maximum(keys(codes))

    for n = 3:nmax
        TT[n] = prufer_restore.(codes[n])
    end
    @info "Read in table of distinct trees up to $nmax vertices"
    return TT

end


text:
using Compat
import Base.getindex
import Compat.IteratorSize, Base.length, Compat.IteratorEltype, Base.eltype

import Phylo.API: _nodetype, _branchtype
import Phylo.API: _getnodenames, _getbranchnames, _getleafnames

mutable struct TreeSet{LABEL, NL, BL, TREE <: AbstractTree{NL, BL}} <: AbstractTree{NL, BL}
    trees::Dict{LABEL, TREE}
    treeinfo::Dict{String, Dict{String, Any}}
end

TreeSet(trees::Dict{LABEL, TREE}, treeinfo) where
        {LABEL, NL, BL, TREE <: AbstractTree{NL, BL}} =
    TreeSet{LABEL, NL, BL, TREE}(trees, treeinfo)

function IteratorSize(::Type{TreeSet})
    return HasLength()
end

function IteratorEltype(::Type{TreeSet})
    return HasEltype()
end

_nodetype(ts::TreeSet) = _nodetype(first(treeiter(ts)))
_branchtype(ts::TreeSet) = _branchtype(first(treeiter(ts)))
_branchtype(::TreeSet{LABEL, NL, BL, TREE}) where {LABEL, NL, BL, TREE <: AbstractBranchTree{NL, BL}} = Branch{NL}
ntrees(ts::TreeSet) = length(ts.trees)

struct TreeIterator{LABEL, NL, BL, TREE <: AbstractTree{NL, BL},
                    TREESET <: TreeSet{LABEL, NL, BL, TREE}} <: AbstractTreeIterator{TREE}
    ts::TREESET
end
treeiter(ts::TREESET) where {LABEL, NL, BL, TREE <: AbstractTree{NL, BL},
                    TREESET <: TreeSet{LABEL, NL, BL, TREE}} =
    TreeIterator{LABEL, NL, BL, TREE, TREESET}(ts)

struct TreeNameIterator{LABEL, NL, BL, TREE <: AbstractTree{NL, BL},
                    TREESET <: TreeSet{LABEL, NL, BL, TREE}} <: AbstractTreeIterator{TREE}
    ts::TREESET
end
treenameiter(ts::TREESET) where {LABEL, NL, BL, TREE <: AbstractTree{NL, BL},
                    TREESET <: TreeSet{LABEL, NL, BL, TREE}} =
    TreeNameIterator{LABEL, NL, BL, TREE, TREESET}(ts)

struct TreeInfoIterator{LABEL, NL, BL, TREE <: AbstractTree{NL, BL},
                    TREESET <: TreeSet{LABEL, NL, BL, TREE}} <: AbstractTreeIterator{TREE}
    ts::TREESET
end
treeinfoiter(ts::TREESET) where {LABEL, NL, BL, TREE <: AbstractTree{NL, BL},
                    TREESET <: TreeSet{LABEL, NL, BL, TREE}} =
    TreeInfoIterator{LABEL, NL, BL, TREE, TREESET}(ts)

eltype(::TreeSet{LABEL, NL, BL, TREE}) where {LABEL, NL, BL, TREE <: AbstractTree{NL, BL}} = TREE
eltype(::TreeIterator{LABEL, NL, BL, TREE, TREESET}) where {LABEL, NL, BL, TREE <: AbstractTree{NL, BL}, TREESET <: TreeSet{LABEL, NL, BL, TREE}} = TREE
eltype(tni::TreeNameIterator) = eltype(keys(tni.ts.trees))
eltype(tii::TreeInfoIterator) = eltype(values(tii.ts.treeinfo))

length(ts::TreeSet) = length(ts.trees)
length(ti::TreeIterator) = length(ti.ts.trees)
length(tni::TreeNameIterator) = length(tni.ts.trees)
length(tii::TreeInfoIterator) = length(tii.ts.treeinfo)

if VERSION >= v"0.7.0-"
import Base.iterate
iterate(ts::TreeSet) = iterate(treeiter(ts))
iterate(ts::TreeSet, state) = iterate(treeiter(ts), state)
iterate(ti::TreeIterator) = iterate(values(ti.ts.trees))
iterate(ti::TreeIterator, state) = iterate(values(ti.ts.trees), state)
iterate(tni::TreeNameIterator) = iterate(keys(tni.ts.trees))
iterate(tni::TreeNameIterator, state) = iterate(keys(tni.ts.trees), state)
iterate(tii::TreeInfoIterator) = iterate(values(tii.ts.treeinfo))
iterate(tii::TreeInfoIterator, state) = iterate(values(tii.ts.treeinfo), state)
else
import Base.start, Base.next, Base.done
start(ts::TreeSet) = start(treeiter(ts))
start(ti::TreeIterator) = start(ti.ts.trees)
start(tni::TreeNameIterator) = start(keys(tni.ts.trees))
start(tii::TreeInfoIterator) = start(tii.ts.treeinfo)

next(ts::TreeSet, state) = next(treeiter(ts), state)
function next(ti::TreeIterator, state)
    v, s = next(ti.ts.trees, state)
    return v[2], s
end
next(tni::TreeNameIterator, state) = next(keys(tni.ts.trees), state)

function next(tii::TreeInfoIterator, state)
    v, s = next(tii.ts.treeinfo, state)
    return v[2], s
end

done(ts::TreeSet, state) = done(treeiter(ts), state)
done(ti::TreeIterator, state) = done(ti.ts.trees, state)
done(ti::TreeNameIterator, state) = done(keys(ti.ts.trees), state)
done(ti::TreeInfoIterator, state) = done(ti.ts.treeinfo, state)
end
getindex(ts::TreeSet, idx) = ts.trees[idx]

function _getleafnames(ts::TREESET) where {LABEL, NL, BL, TREE <: AbstractTree{NL, BL}, TREESET <: TreeSet{LABEL, NL, BL, TREE}}
    lns = unique(map(t -> _getleafnames(t), values(ts.trees)))
    if length(lns) > 1
        error("Inconsistent leaf names in TreeSet")
    elseif isempty(lns)
        return NL[]
    end
    return first(lns)
end

function _getnodenames(ts::TREESET) where {LABEL, NL, BL, TREE <: AbstractTree{NL, BL}, TREESET <: TreeSet{LABEL, NL, BL, TREE}}
    nns = unique(map(t -> _getnodenames(t), values(ts.trees)))
    if length(nns) > 1
        error("Inconsistent node names in TreeSet")
    elseif isempty(nns)
        return NL[]
    end
    return first(nns)
end

function _getbranchnames(ts::TREESET) where {LABEL, NL, BL, TREE <: AbstractTree{NL, BL}, TREESET <: TreeSet{LABEL, NL, BL, TREE}}
    bns = unique(map(t -> _getbranchnames(t), values(ts.trees)))
    if length(bns) > 1
        error("Inconsistent branch names in TreeSet")
    elseif isempty(bns)
        return BL[]
    end
    return first(bns)
end

function _getleafinfo(ts::TREESET) where {LABEL, NL, BL, TREE <: AbstractTree{NL, BL}, TREESET <: TreeSet{LABEL, NL, BL, TREE}}
    lis = unique(map(t -> _getleafinfo(t), values(ts.trees)))
    if length(lis) > 1
        error("Inconsistent leafinfo in TreeSet")
    elseif isempty(lis)
        return eltype(lis)()
    end
    return first(lis)
end

function _nleaves(ts::TREESET) where {LABEL, NL, BL, TREE <: AbstractTree{NL, BL}, TREESET <: TreeSet{LABEL, NL, BL, TREE}}
    nls = unique(map(t -> _nleaves(t), values(ts.trees)))
    if length(nls) > 1
        error("Inconsistent leafinfo in TreeSet")
    elseif isempty(nls)
        return 0
    end
    return first(nls)
end


text:
if length(ARGS) == 0
    Base.warn("No solver is given!")
    @printf "Please provide option for solver to use: Ipopt, ParPips, or Pips \n"
    exit(0);
end
if ARGS[1] == "Ipopt"
    using SerialIpoptInterface
    structJuMPSolve = SerialIpoptInterface.structJuMPSolve
elseif ARGS[1] == "ParPips"
    using ParPipsNlpInterface
    structJuMPSolve = ParPipsNlpInterface.structJuMPSolve
elseif ARGS[1] == "Pips"
    using SerialPipsNlpInterface
    structJuMPSolve = SerialPipsNlpInterface.structJuMPSolve
else
    Base.error("Known solvers are: Ipopt, ParPips, or Pips")
end