text:
# The Snopt.jl package is licensed under the MIT "Expat" License:
#
#  Copyright (c) 2018: Andrew Ning. Modified in 2019: Benoit Landry
#
#  Permission is hereby granted, free of charge, to any person obtaining a copy
#  of this software and associated documentation files (the "Software"), to deal
#  in the Software without restriction, including without limitation the rights
#  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
#  copies of the Software, and to permit persons to whom the Software is
#  furnished to do so, subject to the following conditions:
#
#  The above copyright notice and this permission notice shall be included in all
#  copies or substantial portions of the Software.
#
#  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
#  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
#  SOFTWARE.

__precompile__(false)

if Sys.islinux()
    const snoptlib = "/usr/local/lib/libsnopt_c.so"
else
    const snoptlib = "/usr/local/lib/libsnopt_c.dylib"
end
const codes = Dict{Int64, String}()
codes[1] = "Finished successfully: optimality conditions satisfied"
codes[2] = "Finished successfully: feasible point found"
codes[3] = "Finished successfully: requested accuracy could not be achieved"
codes[11] = "The problem appears to be infeasible: infeasible linear constraints"
codes[12] = "The problem appears to be infeasible: infeasible linear equalities"
codes[13] = "The problem appears to be infeasible: nonlinear infeasibilities minimized"
codes[14] = "The problem appears to be infeasible: infeasibilities minimized"
codes[15] = "The problem appears to be infeasible: infeasible linear constraints in QP subproblem"
codes[21] = "The problem appears to be unbounded: unbounded objective"
codes[22] = "The problem appears to be unbounded: constraint violation limit reached"
codes[31] = "Resource limit error: iteration limit reached"
codes[32] = "Resource limit error: major iteration limit reached"
codes[33] = "Resource limit error: the superbasics limit is too small"
codes[41] =  "Terminated after numerical difficulties: current point cannot be improved"
codes[42] =  "Terminated after numerical difficulties: singular basis"
codes[43] =  "Terminated after numerical difficulties: cannot satisfy the general constraints"
codes[44] =  "Terminated after numerical difficulties: ill-conditioned null-space basis"
codes[51] =  "Error in the user-supplied functions: incorrect objective derivatives"
codes[52] =  "Error in the user-supplied functions: incorrect constraint derivatives"
codes[61] =  "Undefined user-supplied functions: undefined function at the first feasible point"
codes[62] =  "Undefined user-supplied functions: undefined function at the initial point"
codes[63] =  "Undefined user-supplied functions: unable to proceed into undefined region"
codes[71] =  "User requested termination: terminated during function evaluation "
codes[74] =  "User requested termination: terminated from monitor routine"
codes[81] =  "Insufficient storage allocated: work arrays must have at least 500 elements"
codes[82] =  "Insufficient storage allocated: not enough character storage"
codes[83] =  "Insufficient storage allocated: not enough integer storage"
codes[84] =  "Insufficient storage allocated: not enough real storage"
codes[91] =  "Input arguments out of range: invalid input argument"
codes[92] =  "Input arguments out of range: basis file dimensions do not match this problem"
codes[141] =  "System error: wrong number of basic variables"
codes[142] =  "System error: error in basis package"
const PRINTNUM = 18
const SUMNUM = 19

# callback function
function objcon_wrapper(status_::Ptr{Clong}, n::Clong, x_::Ptr{Cdouble},
    needf::Clong, nF::Clong, f_::Ptr{Cdouble}, needG::Clong, lenG::Clong,
    G_::Ptr{Cdouble}, cu_::Ptr{Cchar}, lencu::Clong, iu_::Ptr{Clong},
    leniu::Clong, ru_::Ptr{Cdouble}, lenru::Clong)

    status = unsafe_load(status_)

    # # check if solution finished, no need to calculate more
    # if status >= 2
    #     return
    # end

    # unpack design variables
    x = zeros(n)
    for i = 1:n
        x[i] = unsafe_load(x_, i)
    end

    # cu = Array{Char}(lencu)
    # for i = 1:lencu
    #     cu[i] = unsafe_load(cu_, i)
    # end
    #
    # iu = Array{Int}(leniu)
    # for i = 1:leniu
    #     iu[i] = unsafe_load(iu_, i)
    # end
    #
    # ru = Array{Float64}(lenru)
    # for i = 1:lenru
    #     ru[i] = unsafe_load(ru_, i)
    # end
    #
    # dual_add = 348
    # println(ru[iu[dual_add]:iu[dual_add]+nF])

    # call function
    J, ceq, c, gJ, gceq, gc, HJ = objcon(x)
    fail = false
    gradprovided = true

    # callback
    if ~isa(callback_fn_, Nothing)
        callback_fn_(x)
    end

    # copy obj and con values into C pointer
    unsafe_store!(f_, J, 1)
    for i = 2 : nF - length(ceq)
        unsafe_store!(f_, c[i-1], i)
    end
    if !isempty(ceq)
        for i = nF - length(ceq) + 1 : nF
            unsafe_store!(f_, ceq[i-length(c)-1], i)
        end
    end

    # gradients  TODO: separate gradient computation in interface?
    if needG > 0 && gradprovided

        for j = 1:n
            # gradients of f
            unsafe_store!(G_, gJ[j], j)
        end

        k = n+1
        for i = 2 : nF - length(ceq)
            for j = 1:n
                unsafe_store!(G_, gc[i-1, j], k)
                k += 1
            end
        end
        for i = nF - length(ceq) + 1 : nF
            for j = 1:n
                unsafe_store!(G_, gceq[i-length(c)-1, j], k)
                k += 1
            end
        end
    end

    # check if solutions fails
    if fail
        unsafe_store!(status_, -1, 1)
    end
end

# c wrapper to callback function
const usrfun = @cfunction(objcon_wrapper, Cvoid, (Ptr{Clong}, Ref{Clong}, Ptr{Cdouble},
    Ref{Clong}, Ref{Clong}, Ptr{Cdouble}, Ref{Clong}, Ref{Clong}, Ptr{Cdouble},
    Ptr{Cchar}, Ref{Clong}, Ptr{Clong}, Ref{Clong}, Ptr{Cdouble}, Ref{Clong}))

# main call to snopt
function snopt(fun, num_eqs, num_ineqs, x0, options; x_min = nothing, x_max = nothing,
               printfile = "snopt-print.out", sumfile = "snopt-summary.out", callback_fn = nothing)

    # TODO: there is a probably a better way than to use a global
    global objcon = fun
    global callback_fn_ = callback_fn

    # setup
    Start = 0  # cold start  # TODO: allow warm starts
    nF = 1 + num_ineqs + num_eqs  # 1 objective + constraints
    n = length(x0)  # number of design variables
    ObjAdd = 0.0  # no constant term added to objective (user can add themselves if desired)
    ObjRow = 1  # objective is first thing returned, then constraints

    # linear constraints (none for now)
    iAfun = Clong[1]
    jAvar = Clong[1]
    A = [0.0]  # TODO: change later
    lenA = 1
    neA = 0

    # nonlinear constraints (assume dense jacobian for now)
    lenG = nF*n
    neG = lenG
    iGfun = Array{Clong}(undef, lenG)
    jGvar = Array{Clong}(undef, lenG)
    k = 1
    for i = 1:nF
        for j = 1:n
            iGfun[k] = i
            jGvar[k] = j
            k += 1
        end
    end

    # bound constriaints (no infinite bounds for now)
    if isa(x_min, Nothing)
        xlow = -1e19*ones(n)
    else
        xlow = x_min
    end
    if isa(x_max, Nothing)
        xupp = 1e19*ones(n)
    else
        xupp = x_max
    end
    Flow = -1e20*ones(nF)  # TODO: check Infinite Bound size
    Fupp = zeros(nF)  # TODO: currently c <= 0, but perhaps change

    if num_eqs > 0 #equality constraints
        Flow[nF - num_eqs + 1 : nF] .= 0.0
    end

    # names
    Prob = "opt prob"  # problem name TODO: change later
    nxname = 1  # TODO: change later
    xnames = Array{UInt8}(undef, nxname, 8)
    # xnames = ["TODOTODO"]
    nFname = 1  # TODO: change later
    Fnames = Array{UInt8}(undef, nFname, 8)
    # Fnames = ["TODOTODO"]

    # starting info
    x = x0
    xstate = zeros(n)
    xmul = zeros(n)
    F = zeros(nF)
    Fstate = zeros(nF)
    Fmul = zeros(nF)
    # INFO = 0
    INFO = Clong[0]
    mincw = Clong[0]  # TODO: check that these are sufficient
    miniw = Clong[0]
    minrw = Clong[0]
    nS = Clong[0]
    nInf = Clong[0]
    sInf = Cdouble[0]

    # open files for printing
    iprint = PRINTNUM
    isumm = SUMNUM
    printerr = Clong[0]
    sumerr = Clong[0]
    ccall( (:snopenappend_, snoptlib), Cvoid,
        (Ref{Clong}, Cstring, Ptr{Clong}, Clong),
        iprint, printfile, printerr, sizeof(printfile))
    ccall( (:snopenappend_, snoptlib), Cvoid,
        (Ref{Clong}, Cstring, Ptr{Clong}, Clong),
        isumm, sumfile, sumerr, sizeof(sumfile))
    if printerr[1] != 0
        println("failed to open print file")
    end
    if sumerr[1] != 0
        println("failed to open summary file")
    end

    # temporary working arrays

    ltmpcw = 500
    cw = Array{Cchar}(undef, ltmpcw*8)

    ltmpiw = 500
    iw = Array{Clong}(undef, ltmpiw)

    ltmprw = 500
    rw = Array{Cdouble}(undef, ltmprw)

    # --- initialize ----
    ccall( (:sninit_, snoptlib), Cvoid,
        (Ref{Clong}, Ref{Clong}, Ptr{Cchar}, Ref{Clong}, Ptr{Clong},
        Ref{Clong}, Ptr{Cdouble}, Ref{Clong}, Clong),
        iprint, isumm, cw, ltmpcw, iw,
        ltmpiw, rw, ltmprw, sizeof(cw))

    # --- set options ----
    errors = Clong[0]

    for key in keys(options)
        value = options[key]
        buffer = string(key, repeat(" ", 55-length(key)))  # buffer length is 55 so pad with space.

        if length(key) > 55
            println("warning: invalid option, too long")
            continue
        end

        errors[1] = 0

        if typeof(value) == String

            value = string(value, repeat(" ", 72-length(value)))

            ccall( (:snset_, snoptlib), Cvoid,
                (Cstring, Ref{Clong}, Ref{Clong}, Ptr{Clong},
                Ptr{Cchar}, Ref{Clong}, Ptr{Clong}, Ref{Clong}, Ptr{Cdouble}, Ref{Clong}, Clong, Clong),
                value, iprint, isumm, errors,
                cw, ltmpcw, iw, ltmpiw, rw, ltmprw, sizeof(buffer), sizeof(cw))

        elseif isinteger(value)

            ccall( (:snseti_, snoptlib), Cvoid,
                (Cstring, Ref{Clong}, Ref{Clong}, Ref{Clong}, Ptr{Clong},
                Ptr{Cchar}, Ref{Clong}, Ptr{Clong}, Ref{Clong}, Ptr{Cdouble}, Ref{Clong}, Clong, Clong),
                buffer, value, iprint, isumm, errors,
                cw, ltmpcw, iw, ltmpiw, rw, ltmprw, sizeof(buffer), sizeof(cw))

        elseif isreal(value)

            ccall( (:snsetr_, snoptlib), Cvoid,
                (Cstring, Ref{Cdouble}, Ref{Clong}, Ref{Clong}, Ptr{Clong},
                Ptr{Cchar}, Ref{Clong}, Ptr{Clong}, Ref{Clong}, Ptr{Cdouble}, Ref{Clong}, Clong, Clong),
                buffer, value, iprint, isumm, errors,
                cw, ltmpcw, iw, ltmpiw, rw, ltmprw, sizeof(buffer), sizeof(cw))
        end

        # println(errors[1])

    end

    # --- set memory requirements --- #
    ccall( (:snmema_, snoptlib), Cvoid,
        (Ptr{Clong}, Ref{Clong}, Ref{Clong}, Ref{Clong}, Ref{Clong}, Ref{Clong}, Ref{Clong},
        Ptr{Clong}, Ptr{Clong}, Ptr{Clong},
        Ptr{Cchar}, Ref{Clong}, Ptr{Clong}, Ref{Clong}, Ptr{Cdouble}, Ref{Clong}, Clong),
        INFO, nF, n, nxname, nFname,
        neA, neG,
        mincw, miniw, minrw,
        cw, ltmpcw, iw, ltmpiw, rw, ltmprw, sizeof(cw))

    # --- resize arrays to match memory requirements
    lencw = mincw
    resize!(cw,lencw[1]*8)
    leniw = miniw
    resize!(iw,leniw[1])
    lenrw = minrw
    resize!(rw,lenrw[1])

    memkey = ("Total character workspace", "Total integer   workspace",
        "Total real      workspace")
    memvalue = (lencw,leniw,lenrw)
    for (key,value) in zip(memkey,memvalue)
        buffer = string(key, repeat(" ", 55-length(key)))  # buffer length is 55 so pad with space.
        errors[1] = 0
        ccall( (:snseti_, snoptlib), Cvoid,
            (Cstring, Ref{Clong}, Ref{Clong}, Ref{Clong}, Ptr{Clong},
            Ptr{Cchar}, Ref{Clong}, Ptr{Clong}, Ref{Clong}, Ptr{Cdouble}, Ref{Clong}, Clong, Clong),
            buffer, value, iprint, isumm, errors,
            cw, ltmpcw, iw, ltmpiw, rw, ltmprw, sizeof(buffer), sizeof(cw))
    end

    # --- call snopta ----
    ccall( (:snopta_, snoptlib), Cvoid,
        (Ref{Clong}, Ref{Clong}, Ref{Clong}, Ref{Clong}, Ref{Clong},Ref{Cdouble},
        Ref{Clong}, Cstring, Ptr{Cvoid}, Ptr{Clong}, Ptr{Clong}, Ref{Clong},
        Ref{Clong}, Ptr{Cdouble}, Ptr{Clong}, Ptr{Clong}, Ref{Clong}, Ref{Clong},
        Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cchar}, Ptr{Cdouble}, Ptr{Cdouble}, Ptr{Cchar},
        Ptr{Cdouble}, Ptr{Clong}, Ptr{Cdouble},
        Ptr{Cdouble}, Ptr{Clong},
        Ptr{Cdouble}, Ptr{Clong}, Ptr{Clong}, Ptr{Clong}, Ptr{Clong}, Ptr{Clong},
        Ptr{Clong}, Ptr{Cdouble}, Ptr{Cchar}, Ref{Clong}, Ptr{Clong}, Ref{Clong},
        Ptr{Cdouble}, Ref{Clong}, Ptr{Cchar}, Ref{Clong}, Ptr{Clong}, Ref{Clong},
        Ptr{Cdouble}, Ref{Clong},
        Clong, Clong, Clong, Clong, Clong),
        Start, nF, n, nxname, nFname, ObjAdd,
        ObjRow, Prob, usrfun, iAfun, jAvar, lenA,
        neA, A, iGfun, jGvar, lenG, neG,
        xlow, xupp, xnames, Flow, Fupp, Fnames,
        x, xstate, xmul, F, Fstate,
        Fmul, INFO, mincw, miniw, minrw, nS,
        nInf, sInf, cw, lencw, iw, leniw,
        rw, lenrw, cw, lencw, iw, leniw,
        rw, lenrw,
        sizeof(Prob),sizeof(xnames),sizeof(Fnames),sizeof(cw),sizeof(cw))

    # println("done")

    # close output files
    ccall( (:snclose_, snoptlib), Cvoid,
        (Ref{Clong},),
        iprint)
    ccall( (:snclose_, snoptlib), Cvoid,
        (Ref{Clong},),
        isumm)

    # display(F)
    # println(Fmul)

    return x, codes[INFO[1]]  # xstar, info

end


text:
# This file is a part of TypeDBClient.  License is MIT: https://github.com/Humans-of-Julia/TypeDBClient.jl/blob/main/LICENSE

#
# package typedb.client.test.integration;
#
# import typedb.client.TypeDB;
# import typedb.client.api.TypeDBClient;
# import typedb.client.api.TypeDBSession;
# import typedb.client.api.TypeDBTransaction;
# import typedb.client.api.answer.ConceptMap;
# import typedb.common.test.server.TypeDBRunner;
# import graql.lang.Graql;
# import graql.lang.common.GraqlArg;
# import graql.lang.query.GraqlDefine;
# import graql.lang.query.GraqlInsert;
# import graql.lang.query.GraqlMatch;
# import org.junit.AfterClass;
# import org.junit.BeforeClass;
# import org.junit.Test;
# import org.slf4j.Logger;
# import org.slf4j.LoggerFactory;
#
# import java.io.IOException;
# import java.util.List;
# import java.util.concurrent.TimeoutException;
# import java.util.function.Consumer;
# import java.util.stream.Stream;
#
# import static typedb.client.api.TypeDBSession.Type.DATA;
# import static typedb.client.api.TypeDBTransaction.Type.WRITE;
# import static graql.lang.Graql.and;
# import static graql.lang.Graql.rel;
# import static graql.lang.Graql.rule;
# import static graql.lang.Graql.type;
# import static graql.lang.Graql.var;
# import static java.util.stream.Collectors.toList;
#
# @SuppressWarnings("Duplicates")
# public class ClientQueryTest {
#     private static final Logger LOG = LoggerFactory.getLogger(ClientQueryTest.class);
#     private static TypeDBRunner typedb;
#     private static TypeDBClient typedbClient;
#
#     @BeforeClass
#     public static void setUpClass() throws InterruptedException, IOException, TimeoutException {
#         typedb = new TypeDBRunner();
#         typedb.start();
#         typedbClient = TypeDB.coreClient(typedb.address());
#         if (typedbClient.databases().contains("typedb")) typedbClient.databases().get("typedb").delete();
#         typedbClient.databases().create("typedb");
#     }
#
#     @AfterClass
#     public static void closeSession() throws Exception {
#         typedbClient.close();
#         typedb.stop();
#     }
#
#     @Test
#     public void applicationTest() {
#         LOG.info("clientJavaE2E() - starting client-java E2E...");
#
#         localhostTypeDBTx(tx -> {
#             GraqlDefine defineQuery = Graql.define(
#                     type("child-bearing").sub("relation").relates("offspring").relates("child-bearer"),
#                     type("mating").sub("relation").relates("male-partner").relates("female-partner").plays("child-bearing", "child-bearer"),
#                     type("parentship").sub("relation").relates("parent").relates("child"),
#
#                     type("name").sub("attribute").value(GraqlArg.ValueType.STRING),
#                     type("lion").sub("entity").owns("name").plays("mating", "male-partner").plays("mating", "female-partner").plays("child-bearing", "offspring").plays("parentship", "parent").plays("parentship", "child")
#             );
#             GraqlDefine ruleQuery = Graql.define(rule("infer-parentship-from-mating-and-child-bearing")
#                                                          .when(and(
#                                                                  rel("male-partner", var("male")).rel("female-partner", var("female")).isa("mating"),
#                                                                  var("childbearing").rel("child-bearer").rel("offspring", var("offspring")).isa("child-bearing")))
#                                                          .then(rel("parent", var("male"))
#                                                                        .rel("parent", var("female"))
#                                                                        .rel("child", var("offspring")).isa("parentship")));
#             LOG.info("clientJavaE2E() - define a schema...");
#             LOG.info("clientJavaE2E() - '" + defineQuery + "'");
#             tx.query().define(defineQuery);
#             tx.query().define(ruleQuery);
#             tx.commit();
#             LOG.info("clientJavaE2E() - done.");
#         }, TypeDBSession.Type.SCHEMA);
#
#         // TODO: re-enable when match is implemented
# //        localhostTypeDBTx(tx -> {
# //            final GraqlMatch getThingQuery = Graql.match(var("t").sub("thing"));
# //            LOG.info("clientJavaE2E() - assert if schema defined...");
# //            LOG.info("clientJavaE2E() - '" + getThingQuery + "'");
# //            final List<String> definedSchema = tx.query().match(getThingQuery).get()
# //                    .map(answer -> answer.get("t").asType().asThingType().getLabel()).collect(Collectors.toList());
# //            final String[] correctSchema = new String[]{"thing", "entity", "relation", "attribute",
# //                    "lion", "mating", "parentship", "child-bearing", "name"};
# //            assertThat(definedSchema, hasItems(correctSchema));
# //            LOG.info("clientJavaE2E() - done.");
# //        });
#
#         localhostTypeDBTx(tx -> {
#             String[] names = lionNames();
#             GraqlInsert insertLionQuery = Graql.insert(
#                     var().isa("lion").has("name", names[0]),
#                     var().isa("lion").has("name", names[1]),
#                     var().isa("lion").has("name", names[2])
#             );
#             LOG.info("clientJavaE2E() - insert some data...");
#             LOG.info("clientJavaE2E() - '" + insertLionQuery + "'");
#             tx.query().insert(insertLionQuery);
#             tx.commit();
#             LOG.info("clientJavaE2E() - done.");
#         });
#
#         // TODO: uncomment when match is implemented
# //        localhostTypeDBTx(tx -> {
# //            final String[] familyMembers = lionNames();
# //            LOG.info("clientJavaE2E() - inserting mating relations...");
# //            final GraqlInsert insertMatingQuery = Graql.match(
# //                    var("lion").isa("lion").has("name", familyMembers[0]),
# //                    var("lioness").isa("lion").has("name", familyMembers[1]))
# //                    .insert(rel("male-partner", "lion").rel("female-partner", var("lioness")).isa("mating"));
# //            LOG.info("clientJavaE2E() - '" + insertMatingQuery + "'");
# //            final long insertedMating = tx.query().insert(insertMatingQuery).get().count();
# //
# //            LOG.info("clientJavaE2E() - inserting child-bearing relations...");
# //            final GraqlInsert insertChildBearingQuery = Graql.match(
# //                    var("lion").isa("lion").has("name", familyMembers[0]),
# //                    var("lioness").isa("lion").has("name", familyMembers[1]),
# //                    var("offspring").isa("lion").has("name", familyMembers[2]),
# //                    var("mating").rel("male-partner", var("lion")).rel("female-partner", var("lioness")).isa("mating")
# //            )
# //                    .insert(var("childbearing").rel("child-bearer", var("mating")).rel("offspring", var("offspring")).isa("child-bearing"));
# //            LOG.info("clientJavaE2E() - '" + insertChildBearingQuery + "'");
# //            final long insertedChildBearing = tx.query().insert(insertChildBearingQuery).get().count();
# //
# //            tx.commit();
# //
# //            assertEquals(1, insertedMating);
# //            assertEquals(1, insertedChildBearing);
# //            LOG.info("clientJavaE2E() - done.");
# //        });
# //
# //
# //        localhostTypeDBTx(tx -> {
# //            LOG.info("clientJavaE2E() - execute match get on the lion instances...");
# //            final GraqlMatch getLionQuery = Graql.match(var("p").isa("lion").has("name", var("n")));
# //            LOG.info("clientJavaE2E() - '" + getLionQuery + "'");
# //            final Stream<ConceptMap> insertedLions = tx.query().match(getLionQuery).get();
# //            final List<String> insertedNames = insertedLions.map(answer -> answer.get("n").asThing().asAttribute().asString().getValue()).collect(Collectors.toList());
# //            assertThat(insertedNames, containsInAnyOrder(lionNames()));
# //
# //            LOG.info("clientJavaE2E() - execute match get on the mating relations...");
# //            final GraqlMatch getMatingQuery = Graql.match(var("m").isa("mating"));
# //            LOG.info("clientJavaE2E() - '" + getMatingQuery + "'");
# //            final long insertedMating = tx.query().match(getMatingQuery).get().count();
# //            assertEquals(1, insertedMating);
# //
# //            LOG.info("clientJavaE2E() - execute match get on the child-bearing...");
# //            final GraqlMatch getChildBearingQuery = Graql.match(var("cb").isa("child-bearing"));
# //            LOG.info("clientJavaE2E() - '" + getChildBearingQuery + "'");
# //            final long insertedChildBearing = tx.query().match(getChildBearingQuery).get().count();
# //            assertEquals(1, insertedChildBearing);
# //            LOG.info("clientJavaE2E() - done.");
# //        });
# //
# //        localhostTypeDBTx(tx -> {
# //            LOG.info("clientJavaE2E() - match get inferred relations...");
# //            final GraqlMatch getParentship = Graql.match(
# //                    var("parentship")
# //                            .rel("parent", var("parent"))
# //                            .rel("child", var("child"))
# //                            .isa("parentship"));
# //            LOG.info("clientJavaE2E() - '" + getParentship + "'");
# //            final long parentship = tx.query().match(getParentship).get().count();
# //            //2 answers - single answer for each parent
# //            assertEquals(2, parentship);
# //            LOG.info("clientJavaE2E() - done.");
# //        });
#
#         // TODO: uncomment when aggregate is implemented
# //        localhostTypeDBTx(tx -> {
# //            LOG.info("clientJavaE2E() - match aggregate...");
# //            GraqlMatch.Aggregate aggregateQuery = Graql.match(var("p").isa("lion")).count();
# //            LOG.info("clientJavaE2E() - '" + aggregateQuery + "'");
# //            int aggregateCount = tx.query().aggregate(aggregateQuery).get().get(0).number().intValue();
# //            assertThat(aggregateCount, equalTo(lionNames().length));
# //            LOG.info("clientJavaE2E() - done.");
# //        });
#
#         // TODO: uncomment when compute is implemented
# //        localhostTypeDBTx(tx -> {
# //            LOG.info("clientJavaE2E() - compute count...");
# //            final GraqlCompute.Statistics computeQuery = Graql.compute().count().in("lion");
# //            LOG.info("clientJavaE2E() - '" + computeQuery + "'");
# //            int computeCount = tx.execute(computeQuery).get().get(0).number().intValue();
# //            assertThat(computeCount, equalTo(lionNames().length));
# //            LOG.info("clientJavaE2E() - done.");
# //        });
#
#         // TODO: uncomment when match is implemented
# //        localhostTypeDBTx(tx -> {
# //            LOG.info("clientJavaE2E() - match delete...");
# //            GraqlDelete deleteQuery = Graql.match(var("m").isa("mating")).delete(var("m").isa("mating"));
# //            LOG.info("clientJavaE2E() - '" + deleteQuery + "'");
# //            tx.query().delete(deleteQuery).get();
# //            final long matings = tx.query().match(Graql.match(var("m").isa("mating"))).get().count();
# //            assertEquals(0, matings);
# //            LOG.info("clientJavaE2E() - done.");
# //        });
#
#         LOG.info("clientJavaE2E() - client-java E2E test done.");
#     }
#
#     private String[] commitSHAs() {
#         return new String[]{
#                 // queried commits
#                 "VladGan/console@4bdc38acb87f9fd2fbdb7cbcf2bcc93837382cab",
#                 "VladGan/console@b5ecd4707ce425d7d2d4d0b0d53420cb46e8ce52",
#                 "VladGan/console@b16788637949c6b4c2a3a4bacc8da101bf838b38",
#                 "VladGan/console@8e996fdf8d802d270385ac3bc7cbf5fa77ac0583",
#                 "VladGan/console@1ff6651afa7abf43b5bdd3b1903e489d279e3dc6",
#                 "VladGan/console@6d3ceda79eb3e3dc86d266095b613a53fb083d30",
#                 "VladGan/console@23da5b400e32805c29f41671ff3f92ef48eafcf8",
#                 // not queried commits
#                 "VladGan/console@0000000000000000000000000000000000000000",
#                 "VladGan/console@1111111111111111111111111111111111111111",
#                 "VladGan/console@2222222222222222222222222222222222222222",
#         };
#     }
#
#     @Test
#     public void parallelQueriesInTransactionTest() {
#         localhostTypeDBTx(tx -> {
#             GraqlDefine defineQuery = Graql.define(
#                     type("symbol").sub("attribute").value(GraqlArg.ValueType.STRING),
#                     type("name").sub("attribute").value(GraqlArg.ValueType.STRING),
#                     type("status").sub("attribute").value(GraqlArg.ValueType.STRING),
#                     type("latest").sub("attribute").value(GraqlArg.ValueType.BOOLEAN),
#
#                     type("commit").sub("entity")
#                             .owns("symbol")
#                             .plays("pipeline-automation", "trigger"),
#                     type("pipeline").sub("entity")
#                             .owns("name")
#                             .owns("latest")
#                             .plays("pipeline-workflow", "pipeline")
#                             .plays("pipeline-automation", "pipeline"),
#                     type("workflow").sub("entity")
#                             .owns("name")
#                             .owns("status")
#                             .owns("latest")
#                             .plays("pipeline-workflow", "workflow"),
#
#                     type("pipeline-workflow").sub("relation")
#                             .relates("pipeline").relates("workflow"),
#                     type("pipeline-automation").sub("relation")
#                             .relates("pipeline").relates("trigger")
#             );
#
#             LOG.info("clientJavaE2E() - define a schema...");
#             LOG.info("clientJavaE2E() - '" + defineQuery + "'");
#             tx.query().define(defineQuery);
#             tx.commit();
#             LOG.info("clientJavaE2E() - done.");
#         }, TypeDBSession.Type.SCHEMA);
#
#
#         localhostTypeDBTx(tx -> {
#             String[] commits = commitSHAs();
#             GraqlInsert insertCommitQuery = Graql.insert(
#                     var().isa("commit").has("symbol", commits[0]),
#                     var().isa("commit").has("symbol", commits[1]),
#                     var().isa("commit").has("symbol", commits[3]),
#                     var().isa("commit").has("symbol", commits[4]),
#                     var().isa("commit").has("symbol", commits[5]),
#                     var().isa("commit").has("symbol", commits[6]),
#                     var().isa("commit").has("symbol", commits[7]),
#                     var().isa("commit").has("symbol", commits[8]),
#                     var().isa("commit").has("symbol", commits[9])
#             );
#
#             LOG.info("clientJavaE2E() - insert commit data...");
#             LOG.info("clientJavaE2E() - '" + insertCommitQuery + "'");
#             tx.query().insert(insertCommitQuery);
#             tx.commit();
#             LOG.info("clientJavaE2E() - done.");
#         });
#
#         localhostTypeDBTx(tx -> {
#             GraqlInsert insertWorkflowQuery = Graql.insert(
#                     var().isa("workflow")
#                             .has("name", "workflow-A")
#                             .has("status", "running")
#                             .has("latest", true),
#                     var().isa("workflow")
#                             .has("name", "workflow-B")
#                             .has("status", "finished")
#                             .has("latest", false)
#             );
#
#             LOG.info("clientJavaE2E() - insert workflow data...");
#             LOG.info("clientJavaE2E() - '" + insertWorkflowQuery + "'");
#             tx.query().insert(insertWorkflowQuery);
#             tx.commit();
#             LOG.info("clientJavaE2E() - done.");
#         });
#
#         localhostTypeDBTx(tx -> {
#             GraqlInsert insertPipelineQuery = Graql.insert(
#                     var().isa("pipeline")
#                             .has("name", "pipeline-A")
#                             .has("latest", true),
#                     var().isa("pipeline")
#                             .has("name", "pipeline-B")
#                             .has("latest", false)
#             );
#
#             LOG.info("clientJavaE2E() - insert pipeline data...");
#             LOG.info("clientJavaE2E() - '" + insertPipelineQuery + "'");
#             tx.query().insert(insertPipelineQuery);
#             tx.commit();
#             LOG.info("clientJavaE2E() - done.");
#         });
#
#         localhostTypeDBTx(tx -> {
#             String[] commitShas = commitSHAs();
#             LOG.info("clientJavaE2E() - inserting pipeline-automation relations...");
#
#             for (int i = 0; i < commitShas.length / 2; i++) {
#                 GraqlInsert insertPipelineAutomationQuery = Graql.match(
#                         var("commit").isa("commit").has("symbol", commitShas[i]),
#                         var("pipeline").isa("pipeline").has("name", "pipeline-A")
#                 )
#                         .insert(
#                                 rel("pipeline", "pipeline").rel("trigger", "commit").isa("pipeline-automation")
#                         );
#                 LOG.info("clientJavaE2E() - '" + insertPipelineAutomationQuery + "'");
#                 List<ConceptMap> x = tx.query().insert(insertPipelineAutomationQuery).collect(toList());
#             }
#
#
#             for (int i = commitShas.length / 2; i < commitShas.length; i++) {
#                 GraqlInsert insertPipelineAutomationQuery = Graql.match(
#                         var("commit").isa("commit").has("symbol", commitShas[i]),
#                         var("pipeline").isa("pipeline").has("name", "pipeline-B")
#                 )
#                         .insert(
#                                 rel("pipeline", "pipeline").rel("trigger", "commit").isa("pipeline-automation")
#                         );
#                 LOG.info("clientJavaE2E() - '" + insertPipelineAutomationQuery + "'");
#                 List<ConceptMap> x = tx.query().insert(insertPipelineAutomationQuery).collect(toList());
#             }
#
#             tx.commit();
#
#             LOG.info("clientJavaE2E() - done.");
#         });
#
#         localhostTypeDBTx(tx -> {
#             LOG.info("clientJavaE2E() - inserting pipeline-automation relations...");
#
#             GraqlInsert insertPipelineWorkflowQuery = Graql.match(
#                     var("pipelineA").isa("pipeline").has("name", "pipeline-A"),
#                     var("workflowA").isa("workflow").has("name", "workflow-A"),
#                     var("pipelineB").isa("pipeline").has("name", "pipeline-B"),
#                     var("workflowB").isa("workflow").has("name", "workflow-B")
#             )
#                     .insert(
#                             rel("pipeline", "pipelineA").rel("workflow", "workflowA").isa("pipeline-workflow"),
#                             rel("pipeline", "pipelineB").rel("workflow", "workflowB").isa("pipeline-workflow")
#                     );
#             LOG.info("clientJavaE2E() - '" + insertPipelineWorkflowQuery + "'");
#             List<ConceptMap> x = tx.query().insert(insertPipelineWorkflowQuery).collect(toList());
#
#             tx.commit();
#
#             LOG.info("clientJavaE2E() - done.");
#         });
#
#         String[] queries = {
#                 "match\n" +
#                         "$commit isa commit, has symbol \"VladGan/console@4bdc38acb87f9fd2fbdb7cbcf2bcc93837382cab\";\n" +
#                         "$_ (trigger: $commit, pipeline: $pipeline) isa pipeline-automation;\n" +
#                         "$pipeline has name $pipeline-name, has latest true;\n" +
#                         "$_ (pipeline: $pipeline, workflow: $workflow) isa pipeline-workflow;\n" +
#                         "$workflow has name $workflow-name, has status $workflow-status, has latest true;\n" +
#                         "get $pipeline-name, $workflow-name, $workflow-status;",
#                 "match\n" +
#                         "$commit isa commit, has symbol \"VladGan/console@b5ecd4707ce425d7d2d4d0b0d53420cb46e8ce52\";\n" +
#                         "$_ (trigger: $commit, pipeline: $pipeline) isa pipeline-automation;\n" +
#                         "$pipeline has name $pipeline-name, has latest true;\n" +
#                         "$_ (pipeline: $pipeline, workflow: $workflow) isa pipeline-workflow;\n" +
#                         "$workflow has name $workflow-name, has status $workflow-status, has latest true;\n" +
#                         "get $pipeline-name, $workflow-name, $workflow-status;",
#                 "match\n" +
#                         "$commit isa commit, has symbol \"VladGan/console@b16788637949c6b4c2a3a4bacc8da101bf838b38\";\n" +
#                         "$_ (trigger: $commit, pipeline: $pipeline) isa pipeline-automation;\n" +
#                         "$pipeline has name $pipeline-name, has latest true;\n" +
#                         "$_ (pipeline: $pipeline, workflow: $workflow) isa pipeline-workflow;\n" +
#                         "$workflow has name $workflow-name, has status $workflow-status, has latest true;\n" +
#                         "get $pipeline-name, $workflow-name, $workflow-status;",
#                 "match\n" +
#                         "$commit isa commit, has symbol \"VladGan/console@8e996fdf8d802d270385ac3bc7cbf5fa77ac0583\";\n" +
#                         "$_ (trigger: $commit, pipeline: $pipeline) isa pipeline-automation;\n" +
#                         "$pipeline has name $pipeline-name, has latest true;\n" +
#                         "$_ (pipeline: $pipeline, workflow: $workflow) isa pipeline-workflow;\n" +
#                         "$workflow has name $workflow-name, has status $workflow-status, has latest true;\n" +
#                         "get $pipeline-name, $workflow-name, $workflow-status;",
#                 "match\n" +
#                         "$commit isa commit, has symbol \"VladGan/console@1ff6651afa7abf43b5bdd3b1903e489d279e3dc6\";\n" +
#                         "$_ (trigger: $commit, pipeline: $pipeline) isa pipeline-automation;\n" +
#                         "$pipeline has name $pipeline-name, has latest true;\n" +
#                         "$_ (pipeline: $pipeline, workflow: $workflow) isa pipeline-workflow;\n" +
#                         "$workflow has name $workflow-name, has status $workflow-status, has latest true;\n" +
#                         "get $pipeline-name, $workflow-name, $workflow-status;",
#                 "match\n" +
#                         "$commit isa commit, has symbol \"VladGan/console@6d3ceda79eb3e3dc86d266095b613a53fb083d30\";\n" +
#                         "$_ (trigger: $commit, pipeline: $pipeline) isa pipeline-automation;\n" +
#                         "$pipeline has name $pipeline-name, has latest true;\n" +
#                         "$_ (pipeline: $pipeline, workflow: $workflow) isa pipeline-workflow;\n" +
#                         "$workflow has name $workflow-name, has status $workflow-status, has latest true;\n" +
#                         "get $pipeline-name, $workflow-name, $workflow-status;",
#                 "match\n" +
#                         "$commit isa commit, has symbol \"VladGan/console@23da5b400e32805c29f41671ff3f92ef48eafcf8\";\n" +
#                         "$_ (trigger: $commit, pipeline: $pipeline) isa pipeline-automation;\n" +
#                         "$pipeline has name $pipeline-name, has latest true;\n" +
#                         "$_ (pipeline: $pipeline, workflow: $workflow) isa pipeline-workflow;\n" +
#                         "$workflow has name $workflow-name, has status $workflow-status, has latest true;\n" +
#                         "get $pipeline-name, $workflow-name, $workflow-status;"
#         };
#
#         localhostTypeDBTx(tx -> {
#             LOG.info("clientJavaE2E() - inserting pipeline-automation relations...");
#
#             Stream.of(queries).parallel().forEach(x -> {
#                 GraqlMatch q = Graql.parseQuery(x).asMatch();
#                 List<ConceptMap> res = tx.query().match(q).collect(toList());
#             });
#
#             LOG.info("clientJavaE2E() - done.");
#         });
#     }
#
#     private String[] lionNames() {
#         return new String[]{"male-partner", "female-partner", "young-lion"};
#     }
#
#     private void localhostTypeDBTx(Consumer<TypeDBTransaction> fn) {
#         String database = "typedb";
#         try (TypeDBSession session = typedbClient.session(database, DATA)) {
#             try (TypeDBTransaction transaction = session.transaction(WRITE)) {
#                 fn.accept(transaction);
#             }
#         }
#     }
#
#     private void localhostTypeDBTx(Consumer<TypeDBTransaction> fn, TypeDBSession.Type sessionType) {
#         String database = "typedb";
#         try (TypeDBSession session = typedbClient.session(database, sessionType)) {
#             try (TypeDBTransaction transaction = session.transaction(WRITE)) {
#                 fn.accept(transaction);
#             }
#         }
#     }
# }

ext: jl


text:
using NamedTupleTools: namedtuple, merge

# -------------- Implementation of Product Basis

struct Product1pBasis{NB, TB <: Tuple} <: OneParticleBasis{Any}
   bases::TB
   indices::Vector{NTuple{NB, Int}}
end

function Product1pBasis(bases)
   NB = length(bases)
   return Product1pBasis(bases, NTuple{NB, Int}[])
end

import Base.*
*(B1::OneParticleBasis, B2::OneParticleBasis) =
      Product1pBasis((B1, B2))
*(B1::Product1pBasis, B2::OneParticleBasis) =
      Product1pBasis((B1.bases..., B2))
*(B1::OneParticleBasis, B2::Product1pBasis) =
      Product1pBasis((B1, B2.bases...))
*(B1::Product1pBasis, B2::Product1pBasis) =
      Product1pBasis((B1.bases..., B2.bases...))
*(B1::Product1pBasis, B2::B1pComponent) =
      Product1pBasis((B1.bases..., B2))

_numb(b::Product1pBasis{NB}) where {NB} = NB

Base.length(basis::Product1pBasis) = length(basis.indices)

function Base.show(io::IO, basis::Product1pBasis)
   print(io, "Product1pBasis")
   print(io, basis.bases)
end

Base.getindex(basis::Product1pBasis, i::Integer) = basis.bases[i]

function Base.getindex(basis::Product1pBasis, label::AbstractString)
   inds = findall(getlabel.(basis.bases) .== label)
   if length(inds) == 0
      error("label not found amongst 1p basis components")
   elseif length(inds) > 1
      error("label not unique amongst 1p basis components")
   end
   return basis.bases[inds[1]]
end

# ------------------------- FIO CODES

==(B1::Product1pBasis, B2::Product1pBasis) =
      ( all(B1.bases .== B2.bases) &&
        B1.indices == B2.indices )

write_dict(B::Product1pBasis) =
      Dict("__id__" => "ACE_Product1pBasis",
            "bases" => write_dict.(B.bases),
          "indices" => B.indices )

function read_dict(::Val{:ACE_Product1pBasis}, D::Dict)
   bases = tuple( read_dict.(D["bases"])... )
   indices = [ tuple(v...) for v in D["indices"] ]
   return Product1pBasis(bases, indices)
end

# ----------------- evaluation of the basis

import Base.Cartesian: @nexprs

function _write_A_code(VA, NB)
   prodBi = Meta.parse("B_1[ϕ[1]]" * prod(" * B_$i[ϕ[$i]]" for i = 2:NB))
   if VA == Nothing
      getVT = "promote_type(" * prod("eltype(B_$i), " for i = 1:NB) * ")"
      getA = Meta.parse("_A = zeros($(getVT), length(basis))")
   else
      getA = :(_A = A)
   end
   return prodBi, getA
end

"""
`add_into_A!` : this is an internal function implementing the main evaluation
for the one-particle basis and possibly add it into the A basis.

There are two ways to call it.
* Use `A = nothing` for the first argument to allocate the necessary memory to
evaluate the 1p basis into it.
* Use `A::Vector{T}` to evaluate the 1p basis and add it into `A` directly
without additional allocation.
"""
@generated function add_into_A!(A::VA, basis::Product1pBasis{NB}, X) where {NB, VA}
   prodBi, getA = _write_A_code(VA, NB)
   quote
      # evaluate the 1p basis components
      @nexprs $NB i -> begin
         bas_i = basis.bases[i]
         B_i = evaluate(bas_i, X)
      end
      # allocate A if necessary or just name _A = A if A is a buffer
      $(getA)
      # evaluate the 1p product basis functions and add/write into _A
      for (iA, ϕ) in enumerate(basis.indices)
         @inbounds _A[iA] += $prodBi
      end
      # release the memory allocated by the 1p basis components they normally
      # use preallocated chache to avoid too many small allocations.
      @nexprs $NB i -> release!(B_i)
      return _A
   end
end

evaluate(basis::Product1pBasis, X::AbstractState) =
      add_into_A!(nothing, basis, X)

function evaluate(basis::Product1pBasis, cfg::UConfig)
   @assert length(cfg) > 0 "Product1pBasis can only be evaluated with non-empty configurations"
   # evaluate the first item "manually", then so we know the output types
   # but then write directly into the allocated array to avoid additional
   # allocations.
   A = evaluate(basis, first(cfg))
   for (i, X) in enumerate(cfg)
      i == 1 && continue;
      add_into_A!(A, basis, X)
   end
   return A
end

function evaluate!(A::AbstractVector, basis::Product1pBasis, X::AbstractState)
   fill!(A, zero(eltype(A)))
   add_into_A!(A, basis, X)
   return A
end

function evaluate!(A, basis::Product1pBasis, cfg::UConfig)
   fill!(A, zero(eltype(A)))
   for X in cfg
      add_into_A!(A, basis, X)
   end
   return A
end

# -------------------- jacobian codes, forward rule

function _write_dA_code(VDA, NB)
   if VDA == Nothing
      getDVT = "promote_type(eltype(_A), " * prod("eltype_dB_$i, " for i = 1:NB) * ")"
      getdA = Meta.parse("_dA = zeros($(getDVT), length(basis))")
   else
      getdA = :(_dA = dA)
   end

end

# args... could be a symbol to enable partial derivatives
# at the moment this is just passed into evaluate_ed!(...)
# to not evaluate 1p basis derivatives that aren't needed, but
# in the future a more complex construction could be envisioned that
# might considerably reduce the computational cost here...
@generated function _add_into_A_dA!(A::VA, dA::VDA, basis::Product1pBasis{NB}, X,
                                    args...) where {VA, VDA, NB}
   prodBi, getA = _write_A_code(VA, NB)
   getdA = _write_dA_code(VDA, NB)
   quote
      Base.Cartesian.@nexprs($NB, i -> begin   # for i = 1:NB
         bas_i = basis.bases[i]
         if !(bas_i isa Discrete1pBasis)
            # only evaluate basis gradients for a continuous basis
            B_i, dB_i = evaluate_ed(bas_i, X, args...)
            eltype_dB_i = eltype(dB_i)
         else
            # we still need the basis values for the discrete basis though
            # TODO: maybe the d part should be a no-op and remove this
            # case distinction ...
            B_i = evaluate(bas_i, X)
            eltype_dB_i = Bool
         end
      end)
      # allocate A if necessary or just name _A = A if A is a buffer
      $(getA)
      $(getdA)

      for (iA, ϕ) in enumerate(basis.indices)
         # evaluate A
         @inbounds _A[iA] += $prodBi

         # evaluate dA
         # TODO: redo this with adjoints!!!!
         #     also reverse order of operations to make fewer multiplications!
         _dA[iA] = zero(eltype(_dA))
         Base.Cartesian.@nexprs($NB, a -> begin  # for a = 1:NB
            if !(basis.bases[a] isa Discrete1pBasis)
               dt = dB_a[ϕ[a]]
               Base.Cartesian.@nexprs($NB, b -> begin  # for b = 1:NB
                  if b != a
                     dt *= B_b[ϕ[b]]
                  end
               end)
               _dA[iA] += dt
            end
         end)
      end
      Base.Cartesian.@nexprs($NB, i -> ( begin   # for i = 1:NB
         # release_B!(bas_i, B_i)
         release!(B_i)
         if !(basis.bases[i] isa Discrete1pBasis)
            release!(dB_i)
         end
      end))
      return _A, _dA
   end
end

# this is a hack to resolve a method ambiguity.
add_into_A_dA!(A, dA, basis::Product1pBasis, X) =
               _add_into_A_dA!(A, dA, basis, X)
add_into_A_dA!(A, dA, basis::Product1pBasis, X, sym::Symbol) =
               _add_into_A_dA!(A, dA, basis, X, sym)

evaluate_ed(basis::Product1pBasis, X::AbstractState) =
         _add_into_A_dA!(nothing, nothing, basis, X)

function evaluate_ed(basis::Product1pBasis, Xs::UConfig)
   A, dA1 = evaluate_ed(basis, first(Xs))
   dA = zeros(eltype(dA1), length(basis), length(Xs))
   dA[:, 1] .= dA1[:]
   release!(dA1)
   for (i, X) in enumerate(Xs)
      i == 1 && continue;
      _add_into_A_dA!(A, (@view dA[:, i]), basis, X)
   end
   return A, dA
end

function evaluate_d(basis::Product1pBasis, X::Union{AbstractState, UConfig})
   A, dA = evaluate_ed(basis, X)
   release!(A)
   return dA
end

# ------------- Partial derivative functionality

_check_args_is_sym() = true
_check_args_is_sym(::Symbol) = true

# args... may be empty or a symbol  for partial derivatives
function evaluate_ed!(A, dA, basis::OneParticleBasis,
                     cfg::UConfig, args...)
   @assert _check_args_is_sym(args...)
   fill!(A, 0)
   for (j, X) in enumerate(cfg)
      add_into_A_dA!(A, (@view dA[:, j]), basis, X, args...)
   end
   return A, dA
end

# args... may be empty or a symbol for partial derivatives
function evaluate_ed!(A, dA, basis::Product1pBasis, X::AbstractState, args...)
   @assert _check_args_is_sym(args...)
   fill!(A, 0)
   add_into_A_dA!(A, dA, basis, X, args...)
   return A, dA
end

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

_symbols_prod(bases) = tuple(union( symbols.(bases)... )...)

symbols(basis::Product1pBasis) = _symbols_prod(basis.bases)

function indexrange(basis::Product1pBasis)
   allsyms = tuple(symbols(basis)...)
   rg = Dict{Symbol, Vector{Any}}([ sym => [] for sym in allsyms]...)
   for b in basis.bases
      rgb = indexrange(b)
      for sym in allsyms
         if haskey(rgb, sym)
            rg[sym] = union(rg[sym], rgb[sym])
         end
      end
   end
   # HACK: fix the m range based on the maximal l-range
   #       this needs to be suitably generalised if we have multiple
   #       (l, m) pairs, e.g. (l1, m1), (l2, m2)
   if haskey(rg, :m)
      maxl = maximum(rg[:l])
      rg[:m] = collect(-maxl:maxl)
   end

   # convert the range into a named tuple so that we remember the order!!
   return NamedTuple{allsyms}(ntuple(i -> rg[allsyms[i]], length(allsyms)))
end

isadmissible(b, basis::Product1pBasis) = all(isadmissible.(Ref(b), basis.bases))

function set_spec!(basis::Product1pBasis{NB}, spec) where {NB}
   empty!(basis.indices)
   for b in spec
      inds = ntuple(i -> get_index(basis.bases[i], b), NB)
      push!(basis.indices, inds)
   end
   return basis
end

get_spec(basis::Product1pBasis) = [ get_spec(basis, i) for i = 1:length(basis) ]

function get_spec(basis::Product1pBasis, i::Integer)
   inds = basis.indices[i]
   specs = get_spec.(basis.bases, inds)
   # TODO: here we should check that we are only merging compatible tuples,
   #       e.g. (n = 5, l = 2), (l = 2, m = -1) is ok
   #       but  (n = 5, l = 2), (l = 3, m = -1) is forbidden!
   return merge(specs...)
end

degree(b, basis::Product1pBasis) = sum( degree(b, B) for B in basis.bases )

degree(b::NamedTuple, basis::Product1pBasis, weight::Dict) =
      sum( degree(b, B, weight) for B in basis.bases )

# TODO: this looks like a horrible hack ...
function rand_radial(basis::Product1pBasis)
   for B in basis.bases
      if B isa ScalarACEBasis
         return rand_radial(B)
      end
   end
   return nothing
end

# -------------- sparsification

function sparsify!(basis1p::Product1pBasis, keep::AbstractVector{<: NamedTuple})
   # spec, keep, new_spec will be lists of named tuples,
   #                    e.g. [ (n = , l = , m = ), ... ]
   spec = get_spec(basis1p)
   new_spec = eltype(spec)[]
   new_inds = Vector{Int}(undef, length(spec))
   for (ib, b) in enumerate(spec)
      if b in keep
         push!(new_spec, b)
         new_inds[ib] = length(new_spec)
      end
   end

   # now we need to recompute the indices array, this can be easily done via
   # set_spec!(basis::Product1pBasis{NB}, spec), but before we do that
   # we should sparsify the basis components as well
   #  .... but it is not so clear that his is a good idea, maybe the
   #       1p basis components should just remain frozen???
   #       => turn this off for now
   # TODO - return to this point?!?!?
   # for bas_i in basis1p.bases
   #    _sparsify_component!(bas_i, new_spec)
   # end

   # finally fix the basis1pspec internally:
   set_spec!(basis1p, new_spec)

   # return the old to new index mapping so that the pibasis can fix itself.
   return basis1p, new_inds
end

using NamedTupleTools: select

# """
# this performs some generic work to sparsify a 1p-basis component.
# but the actual sparsificatin happens in the individual basis implementations
# """
# function _sparsify_component!(basis1p, keep)
#    # if basis1p has no symbols (e.g. a multiplier) then it means it must
#    # be a one-component basis, so there is nothing to sparsify.
#    syms = symbols(basis1p)
#    if isempty(syms)
#       return basis1p
#    end
#    # get rid of all info we don't need
#    keep1 = unique( select.(keep, Ref(syms)) )
#    # double-check that keep1 is compatible
#    spec = get_spec(basis1p)
#    @assert all(b in spec for b in keep1)
#    # now get the basis spec and get the list of indices to keep
#    if length(keep1) < length(spec)
#       # Ikeep = findall( [b in keep1 for b in spec] )
#       # sparsify!(basis1p, Ikeep)
#       sparsify!(basis1p, keep1)
#    end
#    return basis1p
# end

# --------------- AD codes

# import ChainRules: rrule, NoTangent, ZeroTangent

# _evaluate_bases(basis::Product1pBasis{NB}, X::AbstractState) where {NB} =
#       ntuple(i -> evaluate(basis.bases[i], X), NB)

# _evaluate_A(basis::Product1pBasis{NB}, BB) where {NB} =
#       [ prod(BB[i][ϕ[i]] for i = 1:NB) for ϕ in basis.indices ]

# evaluate(basis::Product1pBasis, X::AbstractState) =
#       _evaluate_A(basis, _evaluate_bases(basis, X))

# function _rrule_evaluate(basis::Product1pBasis{NB}, X::AbstractState,
#                          w::AbstractVector{<: Number},
#                          BB = _evaluate_bases(basis, X)) where {NB}
#    VT = promote_type(valtype(basis, X), eltype(w))

#    # dB = evaluate_d(basis, X)
#    # return sum( (real(w) * real(db) + imag(w) * imag(db))
#    #             for (w, db) in zip(w, dB) )

#    # Compute the differentials for the individual sub-bases
#    Wsub = ntuple(i -> zeros(VT, length(BB[i])), NB)
#    for (ivv, vv) in enumerate(basis.indices)
#       for t = 1:NB
#          _A = one(VT)
#          for s = 1:NB
#             if s != t
#                _A *= BB[s][vv[s]]
#             end
#          end
#          Wsub[t][vv[t]] += w[ivv] * conj(_A)
#       end
#    end

#    # now these can be propagated into the inner basis
#    #  -> type instab to be fixed here
#    g = sum( _rrule_evaluate(basis.bases[t], X, Wsub[t] )
#             for t = 1:NB )
#    return g
# end

# function rrule(::typeof(evaluate), basis::Product1pBasis, X::AbstractState)
#    BB = _evaluate_bases(basis, X)
#    A = _evaluate_A(basis, BB)
#    return A,
#       w -> (NoTangent(), NoTangent(), _rrule_evaluate(basis, X, w, BB))
# end

#    function _rrule_evaluate(basis::Scal1pBasis, X::AbstractState,
#       w::AbstractVector{<: Number})
# x = _val(X, basis)
# a = _rrule_evaluate(basis.P, x, w)
# TDX = ACE.dstate_type(a, X)
# return TDX( NamedTuple{(_varsym(basis),)}( (a,) ) )
# end


text:
module WeaponSystems
using ExternalBallistics
using ReferenceFrameRotations
using DataFrames, Distances
# Write your package code here.
include("types.jl")
include("tank.jl")
include("transformations.jl")
include("ballisticComputer.jl")

export  TargetRect, TargetCirc, TargetSphe, Hull, Turret, Canon, Sight, Tank, targetPos, muzzlePos, muzzleVel, wind,adjustedFire!,
Fuze
end


text:
module NetworkBasedInference

	using LinearAlgebra

	export NBI,
			NWNBI,
			SDTNBI

	"""

	"""
	function NBI(F₀::AbstractMatrix)
		M,N = size(F₀)
		R = diagm([ sum(F₀[i,:]) for i in 1:M ])
		H = diagm([ sum(F₀[:,j]) for j in 1:N ])
		W = (F₀ * H^-1)' * (R^-1 * F₀)

		return W
    end

	NBI(SymF₀::Symmetric, Nd::Int)::Matrix = NBI(SymF₀[1:Nd, Nd+1:end])

	"""

	"""
	function NWNBI(F₀::Matrix, β::AbstractFloat)::Matrix
		M,N = size(F₀)
        R  = diagm([ sum(F₀[i,:])   for i in 1:M ])
        H  = diagm([ sum(F₀[:,j])   for j in 1:N ])
        H´ = diagm([ sum(F₀[:,j])^β for j in 1:N ])
        W´ = (F₀ * H^-1)' * (R^-1 * F₀ * H´^-1)

        return W´
    end

	NWNBI(SymF₀::Symmetric, Nd::Int, β::AbstractFloat) = NWNBI(SymF₀[1:Nd, Nd+1:end], β)

	"""

	"""
	function SDTNBI(F₀::AbstractMatrix)::Matrix
		k(A, x) = count(r->(r != 0), A[x,:])	# Degree helper function
		n_nodes = size(F₀, 1)					# Number of nodes
		W	    = zeros(n_nodes, n_nodes)		# Transfer matrix

		for idx in 1:n_nodes
			W[idx,:] = k(F₀, idx) > 0 ? F₀[idx,:] ./ k(F₀, idx) : zeros(n_nodes, 1)
		end

		return W
	end
end


text:
using Multilane
using MCTS
using GenerativeModels
using POMDPToolbox
using JLD
using POMDPs
using DataFrames
using Plots

# NOTE: assumes 4 lanes, desired lane is lane 4
desired_lane_reward = 10.

lambdas = Float64[0.1, 1., 10., 100., 1000.]

nb_lanes = 4 # XXX assumption

rmodels = Multilane.NoCrashRewardModel[
                Multilane.NoCrashRewardModel(desired_lane_reward*lambda,
                                             desired_lane_reward,1.5,nb_lanes)
                for lambda in lambdas]

nb_lanes = 4
pp = PhysicalParam(nb_lanes,lane_length=100.) #2.=>col_length=8\n",
_discount = 1.
nb_cars=10
dmodel = NoCrashIDMMOBILModel(nb_cars, pp)

N = 1

mdp = NoCrashMDP(dmodel, rmodels[1], _discount);
isrng = MersenneTwister(123)
initial_states = MLState[initial_state(mdp, isrng) for i in 1:N]

# Calculate single points
point_solvers = Dict{String, Solver}(
    "random" => RandomSolver(),
    "heuristic" => SimpleSolver()
)

point_results = evaluate([mdp], initial_states, point_solvers, parallel=true)

dpws = DPWSolver(depth=10,
                 n_iterations=1000,
                 exploration_constant=100.0,
                 rollout_solver=SimpleSolver())

curve_solvers = Dict{String, Solver}(
    "dpw" => dpws,
    # "robust" =>
    #     Robust
    "single_behavior" =>
        SingleBehaviorSolver(dpws,
                 IDMMOBILBehavior("normal", 30.0, 10.0, 1))
)

curve_problems = []
for i in 1:length(lambdas)
    rmodel = rmodels[i]
    mdp = NoCrashMDP(dmodel, rmodel, _discount);
    push!(curve_problems, mdp)
end

curve_results = evaluate(curve_problems, initial_states, curve_solvers, parallel=true)

results = merge_results!(curve_results, point_results)

println(results["stats"])

filename = string("results_", Dates.format(Dates.now(),"u_d_HH_MM"), ".jld")
results["histories"] = nothing
save(filename, results)
println("results saved to $filename")

stats = results["stats"]
mean_performance = by(stats, :solver_key) do df
    by(df, :lambda) do df
        DataFrame(steps_in_lane=mean(df[:steps_in_lane]),
                  nb_brakes=mean(df[:nb_brakes])
                  )
    end
end

unicodeplots()
for g in groupby(mean_performance, :solver_key)
    if size(g,1) > 1
        plot!(g, :steps_in_lane, :nb_brakes, group=:solver_key)
    else
        scatter!(g, :steps_in_lane, :nb_brakes, group=:solver_key)
    end
end
gui()


text:
@testset "Memory: Static" begin

function memory_static_kernel(a,b)
    # Local
    ptr_local = alloc_special(Val(:local), Float32, Val(AS.Local), Val(1))
    unsafe_store!(ptr_local, a[1])
    b[1] = unsafe_load(ptr_local)

    # Region
    #= TODO: AMDGPU target cannot select
    ptr_region = alloc_special(Val(:region), Float32, Val(AS.Region), Val(1))
    unsafe_store!(ptr_region, a[2])
    b[2] = unsafe_load(ptr_region)
    =#

    # Private
    #= TODO
    ptr_private = alloc_special(Val(:private), Float32, Val(AS.Private), Val(1))
    unsafe_store!(ptr_private, a[3])
    b[3] = unsafe_load(ptr_private)
    =#

    nothing
end

A = ones(Float32, 1)
B = zeros(Float32, 1)

HA = HSAArray(A)
HB = HSAArray(B)

wait(@roc memory_static_kernel(HA, HB))

@test Array(HA) ≈ Array(HB)

end

@testset "Memory: Dynamic" begin

function malloc_kernel(X)
    ptr = AMDGPU.malloc(Csize_t(4))
    X[1] = convert(AMDGPU.DevicePtr{UInt64,AS.Global}, ptr)
    AMDGPU.free(ptr)
    nothing
end

HA = HSAArray(zeros(UInt64, 1))

wait(@roc malloc_kernel(HA))

@test Array(HA)[1] != 0

end

@testset "Memcpy/Memset" begin

function memcpy_kernel(X,Y)
    AMDGPU.memcpy!(Y.ptr, X.ptr, sizeof(Float32)*length(X))
    nothing
end

A = rand(Float32, 4)
B = zeros(Float32, 4)
HA, HB = HSAArray.((A,B))

wait(@roc memcpy_kernel(HA,HB))

@test A == collect(HA) == collect(HB)

function memset_kernel(X,y)
    AMDGPU.memset!(X.ptr, y, div(length(X),2))
    nothing
end

A = zeros(UInt8, 4)
HA = HSAArray(A)

wait(@roc memset_kernel(HA,0x3))

@test all(collect(HA)[1:2] .== 0x3)
@test all(collect(HA)[3:4] .== 0x0)

end


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

struct SimpleSMatrix{N,M,T} <: AbstractMatrix{T}
    m::Matrix{T}
end

SimpleSMatrix{N,M}(m::AbstractMatrix{T}) where {N,M,T} =
    size(m) == (N, M) ? SimpleSMatrix{N,M,T}(m) : throw(error("Wrong matrix size"))

Base.:*(a::SimpleSMatrix{N,O}, b::SimpleSMatrix{O,M}) where {N,O,M} =
    SimpleSMatrix{N,M}(a.m * b.m)

Base.:*(a::LinearAlgebra.Adjoint{<:Any, <:SimpleSMatrix{O,N}}, b::SimpleSMatrix{O,M}) where {N,O,M} =
    SimpleSMatrix{N,M}(adjoint(a.parent.m) * b.m)

Base.:*(a::SimpleSMatrix{N,O}, b::LinearAlgebra.Adjoint{<:Any, <:SimpleSMatrix{M,O}}) where {N,O,M} =
    SimpleSMatrix{N,M}(a.m * adjoint(b.parent.m))

Base.:+(a::SimpleSMatrix{N,M}, b::SimpleSMatrix{N,M}) where {N,M} =
    SimpleSMatrix{N,M}(a.m + b.m)

Base.:+(a::LinearAlgebra.Adjoint{<:Any, <:SimpleSMatrix{M,N}}, b::SimpleSMatrix{N,M}) where {N,M} =
    SimpleSMatrix{N,M}(adjoint(a.parent.m) + b.m)

Base.:+(a::LinearAlgebra.Adjoint{<:Any, <:SimpleSMatrix}, b::LinearAlgebra.Adjoint{<:Any, <:SimpleSMatrix}) =
    (a' + b')'

Base.:+(a::SimpleSMatrix{N,M}, b::LinearAlgebra.Adjoint{<:Any, <:SimpleSMatrix{M,N}}) where {N,M} =
    SimpleSMatrix{N,M}(a.m + adjoint(b.parent.m))

Base.:-(a::SimpleSMatrix{N,M}, b::SimpleSMatrix{N,M}) where {N,M} =
    SimpleSMatrix{N,M}(a.m - b.m)

Base.:-(a::LinearAlgebra.Adjoint{<:Any, <:SimpleSMatrix{M,N}}, b::SimpleSMatrix{N,M}) where {N,M} =
    SimpleSMatrix{N,M}(adjoint(a.parent.m) - b.m)

Base.:-(a::LinearAlgebra.Adjoint{<:Any, <:SimpleSMatrix}, b::LinearAlgebra.Adjoint{<:Any, <:SimpleSMatrix}) =
    (a' - b')'

Base.:-(a::SimpleSMatrix{N,M}, b::LinearAlgebra.Adjoint{<:Any, <:SimpleSMatrix{M,N}}) where {N,M} =
    SimpleSMatrix{N,M}(a.m - adjoint(b.parent.m))

Base.size(a::SimpleSMatrix{N,M}) where {N,M} = (N, M)

Base.length(a::SimpleSMatrix{N,M}) where {N,M} = N * M

Base.zero(::Type{S}) where {N,M,T,S<:SimpleSMatrix{N,M,T}} = SimpleSMatrix{N,M}(zeros(T, N, M))

Base.getindex(s::SimpleSMatrix, inds...) = getindex(s.m, inds...)

Base.convert(::Type{S}, value::Matrix) where {N,M,T,S<:SimpleSMatrix{N,M,T}} =
    SimpleSMatrix{N,M}(T.(value))

max_stars_repo_name: TimoLarson/julia

max_issues_repo_name: TimoLarson/julia

max_forks_repo_name: TimoLarson/julia


text:
import StatsBase

"""
    TransferOperatorGridTest(k::Int = 1, l::Int = 1, m::Int = 1, n::Int = 1,
        τ::Int = 1, b = 2, estimator::TransferOperatorGrid = TransferOperatorGrid(),
        binning_summary_statistic::Function = StatsBase.mean,
        binning::RectangularBinning, ηs)

The parameters for a transfer entropy test using the `TransferOperatorGrid` estimator [1].

## Mandatory keyword arguments

- **`binning::RectangularBinning`**: An instance of a [`RectangularBinning`](@ref) that dictates
    how the delay embedding is discretized.

- **`ηs`**: The prediction lags (that gos into the ``T_{f}`` component of the embedding).

## Optional keyword arguments

- **`k::Int`**: The dimension of the ``T_{f}`` component of the embedding.

- **`l::Int`**: The dimension of the ``T_{pp}`` component of the embedding.

- **`m::Int`**: The dimension of the ``S_{pp}`` component of the embedding.

- **`n::Int`**: The dimension of the ``C_{pp}`` component of the embedding.

- **`τ::Int`**: The embedding lag. Default is `τ = 1`.

- **`b`**: Base of the logarithm. The default (`b = 2`) gives the TE in bits.

- **`estimator::TransferOperatorGrid`**: A `TransferOperatorGrid` estimator instance.

- **`binning_summary_statistic::Function`**: A summary statistic to summarise the
    transfer entropy values if multiple binnings are provided.

## Estimation of the invariant measure

With a `TransferOperatorGridTest`, the first step is to compute an approximation to the
[transfer operator](@ref transfer_operator_rectangular)
over the partition elements of a rectangular [discretization](@ref discretization)
of an appropriate [delay reconstruction](@ref custom_delay_reconstruction) of the time
series to be analysed. Transfer entropy is then computed from the [invariant
distribution](@ref invariant_measure_rectangular) arising from the transfer operator.

## About the delay reconstruction for transfer entropy analysis

Denote the time series for the source process ``S`` as ``S(t)``, and the time series for
the target process ``T`` as ``T(t)``, and ``C_i(t)`` as the time series for any conditional
processes ``C_i`` that also may influence ``T``. To compute (conditional) TE, we need a
generalised embedding [3, 4] incorporating all of these processes.

For convenience, define the state vectors

```math
\\begin{align}
T_f^{(k)} &= \\{(T(t+\\eta_k), \\ldots, T(t+\\eta_2), T(t+\\eta_1))\\}, \\label{eq:Tf} \\\\
T_{pp}^{(l)} &= \\{ (T(t), T(t-\\tau_1), T(t-\\tau_2), \\ldots, T(t - \\tau_{l - 1})) \\\\
S_{pp}^{(m)} &= \\{(S(t), S(t-\\tau_1), S(t-\\tau_2), \\ldots, S(t-\\tau_{m - 1}))\\},\\\\
C_{pp}^{(n)} &= \\{ (C_1(t), C_1(t-\\tau_1), \\ldots,  C_2(t), C_2(t-\\tau_1) \\},
\\end{align}
```

where the state vectors ``T_f^{(k)}`` contain ``k`` future values of the target
variable, ``T_{pp}^{(l)}`` contain ``l`` present and past values of the target
variable, ``S_{pp}^{(m)}`` contain ``m`` present and past values of the source
variable, ``C_{pp}^{(n)}`` contain a total of ``n`` present and past values of any
conditional variable(s).  Combining all variables, we have the generalised embedding

```math
\\begin{align}
\\mathbb{E} = (T_f^{(k)}, T_{pp}^{(l)}, S_{pp}^{(m)}, C_{pp}^{(n)})
\\end{align}
```

with a total embedding dimension of ``k + l + m + n``.
Here, ``\\tau`` indicates the [embedding](@ref custom_delay_reconstruction) lag
(in the `VisitationFrequencyTest`, we set ``\\tau_1 = \\tau_2 = \\tau_3 = \\ldots``), and
``\\eta`` indicates the prediction lag (the lag of the influence the source
has on the target).

Hence, in the generalised embedding, only ``T_f`` depends on the prediction lag ``\\eta``,
which is to be determined by the analyst. For transfer entropy analysis, ``\\eta`` is chosen
to be some positive integer, while for
[predictive asymmetry analysis](@ref predictive_asymmetry),
symmetric negative and positive ``\\eta``s are used for computing ``\\mathbb{A}``.

In terms of this generalised embedding, transfer entropy from a source variable ``S`` to a
target variable ``T`` with conditioning on variable(s) ``C`` is thus defined as

```math
\\begin{align}
TE_{S \\rightarrow T|C} = \\int_{\\mathbb{E}} P(T_f, T_{pp}, S_{pp}, C_{pp}) \\log_{2}{\\left( \\frac{P(T_f | T_{pp}, S_{pp}, C_{pp})}{P(T_f | T_{pp}, C_{pp})}\\right)}
\\end{align}
```

Without conditioning, we have

```math
\\begin{align}
TE_{S \\rightarrow T} = \\int_{\\mathbb{E}} P(T_f, T_{pp}, S_{pp}) \\log_{2}{\\left(\\frac{P(T_f | T_{pp}, S_{pp})}{P(T_f | T_{pp})}\\right)}
\\end{align}
```

## Low-level estimator

This test uses the `TransferOperatorGrid` estimator on the following low-level method
under the hood.

- [`transferentropy(::Any, ::TEVars, ::RectangularBinning, ::TransferEntropyEstimator)`](@ref)

In this estimator, the mapping between variables of the
[generalised embedding](@ref custom_delay_reconstruction) and the marginals during
transfer entropy computation is controlled using a [`TEVars`](@ref)
instance. It is *highly* recommended that you check the documentation for this
method, because it describes the transfer entropy estimation procedure in detail.

## Notes:

- Use `causality(source, target, params::TransferOperatorGridTest)` for regular
    transfer entropy analysis. This method uses only the `k`, `l`, `m` and ignores `n`
    when constructing the delay reconstruction.

- Use `causality(source, target, cond, params::TransferOperatorGridTest)` for conditional
    transfer entropy analysis. This method uses the `k`, `l`, `m` *and* `n` when constructing
    the delay reconstruction.

## Example

```julia
# Prediction lags
ηs = 1:10
binning = RectangularBinning(10)

# Use defaults, binning and prediction lags are required.
# Note that `binning` and `ηs` are *mandatory* keyword arguments.
TransferOperatorGridTest(binning = binning, ηs = ηs)

# The other keywords can also be adjusted
TransferOperatorGridTest(k = 1, l = 2, binning = binning, ηs = ηs)
```

## References

1. Diego, David, Kristian Agasøster Haaga, and Bjarte Hannisdal. "Transfer entropy computation
    using the Perron-Frobenius operator." Physical Review E 99.4 (2019): 042212.
    [https://journals.aps.org/pre/abstract/10.1103/PhysRevE.99.042212](https://journals.aps.org/pre/abstract/10.1103/PhysRevE.99.042212)
"""
Base.@kwdef struct TransferOperatorGridTest{N} <: TransferEntropyCausalityTest{N}
    """ The delay reconstruction parameter k (controls dimension of ``T_{f}`` component of embedding). """
    k::Int = 1

    """ The delay reconstruction parameter l (controls dimension of ``T_{pp}`` component of embedding). """
    l::Int = 1

    """ The delay reconstruction parameter m (controls dimension of ``S_{pp}`` component of embedding). """
    m::Int = 1

    """ The delay reconstruction parameter n (controls dimension of ``C_{pp}`` component of embedding). """
    n::Int = 1

    """ The delay reconstruction lag for the ``T_{pp}`` component of the embedding. """
    τ::Int = 1

    """ The base of the logarithm for computing TE. """
    b::Number = 2

    """ The transfer entropy estimator. """
    estimator::TransferOperatorGrid = TransferOperatorGrid()

    """
    If there are several binnings provided, what is the statistic used to summarise the
    transfer entropy values to a single value?
    """
    binning_summary_statistic::Function = StatsBase.mean

    """
    The binning scheme(s). If more than one is provided, the `binning_summary_statistic` is
    applied to the computed transfer entropy values, and a single value is returned.
    """
    binning::Union{RectangularBinning, Vector{RectangularBinning}}

    """ The prediction lags"""
    ηs

    function TransferOperatorGridTest(k::Int, l::Int, m::Int, n::Int, τ::Int, b::Number,
            estimator::TransferOperatorGrid,
            binning_summary_statistic::Function,
            binning::Union{RectangularBinning, Vector{RectangularBinning}},
            ηs)

        N = length(ηs) # length of return vector when used with `causality`
        return new{N}(k, l, m, n, τ, b, estimator, binning_summary_statistic, binning, ηs)
    end
end

function causality(source::AbstractVector{T}, target::AbstractVector{T}, p::TransferOperatorGridTest)  where {T <: Real}
    [p.binning_summary_statistic(
        transferentropy(source, target, p.binning,
            p.k, p.l, p.m, η = η, τ = p.τ,
            estimator = p.estimator)) for η in p.ηs]
end

export TransferOperatorGridTest


text:
using PalmerPenguins, DataFrames
ENV["DATADEPS_ALWAYS_ACCEPT"] = true
penguins = dropmissing(DataFrame(PalmerPenguins.load()))

using DashiBoard, Sockets
set_aog_theme!()
# TODO: should the font size depend on pixel ratio?
update_theme!(fontsize=24)

(@isdefined server) && close(server)

pipelinetabs = (
    :Load,
    :Filter,
    :Process => (options=[:Predict, :Cluster, :Project, :Wildcard],),
)

visualizationtabs = (:Spreadsheet, :Chart, :Pipelines)

server = DashiBoard.serve(penguins; pipelinetabs, visualizationtabs,
    url=Sockets.localhost, port=9000, verbose=true)

nothing


text:
module NEPSolver
    using ..NEPCore
    using ..NEPTypes
    using ..LinSolvers
    using ..NEPTransformations
    using LinearAlgebra
    using SparseArrays
    using Random

    include("inner_solver.jl");
    include("compute_rf_wrapper.jl");

    ## NEP-Methods

    include("method_newton.jl")
    include("method_iar.jl")
    include("method_iar_chebyshev.jl")
    include("method_tiar.jl")
    include("method_infbilanczos.jl")
    include("method_mslp.jl")
    include("method_companion.jl")
    include("method_nlar.jl")
    include("method_sgiter.jl")
    include("method_rfi.jl")
    include("method_jd.jl")
    include("method_contour_common.jl")
    include("method_beyncontour.jl")
    include("method_block_SS.jl")
    include("method_blocknewton.jl")
    include("method_broyden.jl")
    include("method_ilan.jl")
    include("method_nleigs.jl")

    """
        σ = closest_to(λ_vec::Array{T,1},  λ::T) where {T<:Number}

    Finds the value `σ` in the vector `λ_vec` that is closes to the value `λ`.
    """
    function closest_to(λ_vec::Array{T,1}, λ::T) where {T<:Number}
        idx = argmin(abs.(λ_vec .- λ))
        return λ_vec[idx]
    end

end #End module


text:
module CRUDType

    @enum CRUD begin
         create = 1
         retrieve = 2
         update = 3
         delete = 4
    end

end


text:
using Gadfly

set_default_plot_size(9inch, 3inch)

y, size = [0.55, 0.7, 0.9, 0.99, 0.9], [0.4, 0.5, 0.6, 0.68, 0.63]
juliaclrs = Gadfly.parse_colorant(["forestgreen", "brown3", "mediumorchid"])
theme = Theme(key_swatch_shape=Shape.circle, alphas=[0.1], discrete_highlight_color=identity)

p1 = plot(x=[2,1.13,2.87], y=[3,1.5,1.5], size=[0.75], alpha=[0.8],
    color=juliaclrs, Coord.cartesian(fixed=true))
p2 = plot(x=1:5, y=y, size=10^5*size, Scale.size_radius(maxvalue=10^5), theme)
p3 = plot(x=1:5, y=y, size=10^5*size, Scale.size_area(maxvalue=10^5), theme)

hstack(p1, p2, p3)


text:
module VAN

using StatsBase
import ChainRulesCore: @non_differentiable, NoTangent, Tangent, rrule
using Optimisers: update!
import Zygote

export AutoRegressiveModel, get_logp, model_dispatch!, gen_samples, model_parameters
export exact_free_energy, free_energy
export AbstractSampler
export PSAModel, get_logp, gen_samples, model_parameters, model_dispatch!

include("sampler.jl")
include("arm.jl")
include("psa.jl")
include("exact.jl")
include("loss.jl")
include("utils.jl")
include("zygote_patch.jl")

end # module


text:
"""
    Datatype

A `Datatype` represents the layout of the data in memory.

# Usage

    Datatype(T)

Either return the predefined `Datatype` corresponding to `T`, or create a new `Datatype`
for the Julia type `T`, calling [`Types.commit!`](@ref) so that it can be used for
communication operations.

Note that this can only be called on types for which `isbitstype(T)` is `true`.
"""
@mpi_handle Datatype MPI_Datatype

const DATATYPE_NULL = _Datatype(MPI_DATATYPE_NULL)

const MPI_Datatype_default = MPI_Datatype == Cint ? MPI_DATATYPE_NULL : C_NULL
Datatype() = Datatype(MPI_Datatype_default)

function free(dt::Datatype)
    if dt.val != DATATYPE_NULL.val && !Finalized()
        @mpichk ccall((:MPI_Type_free, libmpi), Cint, (Ptr{MPI_Datatype},), dt)
    end
    return nothing
end

# attributes
function _type_null_copy_fn(oldtype::MPI_Datatype, type_keyval::Cint,
                            extra_state::Ptr{Cvoid},
                            attribute_val_in::Ptr{Cvoid}, attribute_val_out::Ptr{Cvoid},
                            flag::Ptr{Cint})
    unsafe_store!(flag, Cint(0))
    return MPI_SUCCESS
end
function _type_null_delete_fn(oldtype::MPI_Datatype, type_keyval::Cint,
                           attribute_val_in::Ptr{Cvoid}, extra_state::Ptr{Cvoid})
    return MPI_SUCCESS
end

function create_keyval(::Type{Datatype})
    ref = Ref(Cint(0))
    @mpichk ccall((:MPI_Type_create_keyval, libmpi), Cint,
                  (Ptr{Cvoid},
                   Ptr{Cvoid},
                   Ptr{Cint},
                   Ptr{Cvoid}),
                  @cfunction(_type_null_copy_fn,Cint,(MPI_Datatype, Cint, Ptr{Cvoid},
                                                      Ptr{Cvoid}, Ptr{Cvoid}, Ptr{Cint})),
                  @cfunction(_type_null_delete_fn,Cint,(MPI_Datatype, Cint, Ptr{Cvoid},
                                                      Ptr{Cvoid})),
                  ref, C_NULL)
    return ref[]
end
function set_attr!(datatype::Datatype, keyval::Cint, attrval::Ptr{Cvoid})
    @mpichk ccall((:MPI_Type_set_attr, libmpi), Cint,
                  (MPI_Datatype, Cint, Ptr{Cvoid}), datatype, keyval, attrval)
    return nothing
end
function get_attr(datatype::Datatype, keyval::Cint)
    flagref = Ref(Cint(0))
    attrref = Ref{Ptr{Cvoid}}(C_NULL)
    @mpichk ccall((:MPI_Type_get_attr, libmpi), Cint,
                  (MPI_Datatype, Cint, Ptr{Ptr{Cvoid}}, Ptr{Cint}), datatype, keyval, attrref, flagref)
    if flagref[] == 0
        return nothing
    end
    return attrref[]
end
function del_attr!(datatype::Datatype, keyval::Cint)
    flagref = Ref(Cint(0))
    attrref = Ref{Ptr{Cvoid}}(C_NULL)
    @mpichk ccall((:MPI_Type_delete_attr, libmpi), Cint,
                  (MPI_Datatype, Cint), datatype, keyval)
    return nothing
end

# names
function get_name(datatype::Datatype)
    MPI.Initialized() || return ""
    buffer = Array{UInt8}(undef, MPI_MAX_OBJECT_NAME)
    lenref = Ref{Cint}()
    @mpichk ccall((:MPI_Type_get_name, libmpi), Cint,
                  (MPI_Datatype, Ptr{UInt8}, Ptr{Cint}),
                  datatype, buffer, lenref)
    return String(resize!(buffer, lenref[]))
end

# datatype attribute to store Julia type
const JULIA_TYPE_PTR_ATTR = Ref(Cint(0))
push!(mpi_init_hooks, () -> JULIA_TYPE_PTR_ATTR[] = create_keyval(Datatype))

"""
    to_type(datatype::Datatype)

Return the Julia type corresponding to the MPI [`Datatype`](@ref) `datatype`, or `nothing`
if it doesn't correspond directly.
"""
function to_type(datatype::Datatype)
    if MPI.Initialized()
        ptr = get_attr(datatype, JULIA_TYPE_PTR_ATTR[])
        if !isnothing(ptr)
            return unsafe_pointer_to_objref(ptr)
        end
    end
    return nothing
end

# predefined
for (mpiname, T) in [
    :INT8_T            => Int8
    :UINT8_T           => UInt8
    :INT16_T           => Int16
    :UINT16_T          => UInt16
    :INT32_T           => Int32
    :UINT32_T          => UInt32
    :INT64_T           => Int64
    :UINT64_T          => UInt64
    :BYTE              => UInt8
    :SHORT             => Cshort
    :UNSIGNED_SHORT    => Cushort
    :INT               => Cint
    :UNSIGNED          => Cuint
    :LONG              => Clong
    :UNSIGNED_LONG     => Culong
    :CHAR              => Cchar
    :SIGNED_CHAR       => Cchar
    :UNSIGNED_CHAR     => Cuchar
    :WCHAR             => Cwchar_t
    :FLOAT             => Float32
    :DOUBLE            => Float64
    :C_FLOAT_COMPLEX   => ComplexF32
    :C_DOUBLE_COMPLEX  => ComplexF64]

    @eval if @isdefined($(Symbol(:MPI_,mpiname)))
        const $mpiname = _Datatype($(Symbol(:MPI_,mpiname)))
        if !hasmethod(Datatype, Tuple{Type{$T}})
            Datatype(::Type{$T}) = $mpiname
            push!(mpi_init_hooks, () -> set_attr!($mpiname, JULIA_TYPE_PTR_ATTR[], pointer_from_objref($T)))
        end
    end
end

@generated function Datatype(::Type{T}) where {T}
    newtype = Datatype()
    quote
        datatype = $newtype
        # lazily initialize so that it can be safely precompiled
        if datatype.val === MPI.MPI_Datatype_default
            Types.create!(datatype, T)
            Types.commit!(datatype)
            set_attr!(datatype, JULIA_TYPE_PTR_ATTR[], pointer_from_objref(T))
        end
        return datatype
    end
end

function Base.show(io::IO, datatype::Datatype)
    show(io, Datatype)
    print(io, '(')
    juliatype = to_type(datatype)
    if isnothing(juliatype)
        show(io, datatype.val)
    else
        show(io, juliatype)
    end
    print(io, ')')
    name = get_name(datatype)
    if name != ""
        print(io, ": ")
        print(io, name)
    end
end

module Types

import MPI
import MPI: @mpichk, libmpi, _doc_external,
    Datatype, MPI_Datatype, MPI_Aint,
    free

"""
    lb, extent = MPI.Types.extent(dt::MPI.Datatype)

Gets the lowerbound `lb` and the extent `extent` in bytes.

# External links
$(_doc_external("MPI_Type_get_extent"))
"""
function extent(dt::Datatype)
    lb = Ref{MPI_Aint}()
    extent = Ref{MPI_Aint}()
    # int MPI_Type_get_extent(MPI_Datatype datatype, MPI_Aint *lb,
    #          MPI_Aint *extent)
    @mpichk ccall((:MPI_Type_get_extent, libmpi), Cint,
                  (MPI_Datatype, Ptr{MPI_Aint}, Ptr{MPI_Aint}),
                  dt, lb, extent)
    return lb[], extent[]
end

"""
    MPI.Types.create_contiguous(count::Integer, oldtype::MPI.Datatype)

Create a derived [`Datatype`](@ref) that replicates `oldtype` into `count` contiguous locations.

Note that [`MPI.Types.commit!`](@ref) must be used before the datatype can be used for
communication.

# External links
$(_doc_external("MPI_Type_contiguous"))
"""
function create_contiguous(count::Integer, oldtype::Datatype)
    finalizer(free, create_contiguous!(Datatype(), count, oldtype))
end

function create_contiguous!(newtype::Datatype, count::Integer, oldtype::Datatype)
    @mpichk ccall((:MPI_Type_contiguous, libmpi), Cint,
                  (Cint, MPI_Datatype, Ptr{MPI_Datatype}),
                  count, oldtype, newtype)
    return newtype
end

"""
    MPI.Types.create_vector(count::Integer, blocklength::Integer, stride::Integer, oldtype::MPI.Datatype)

Create a derived [`Datatype`](@ref) that replicates `oldtype` into locations that
consist of equally spaced blocks.

Note that [`MPI.Types.commit!`](@ref) must be used before the datatype can be used for
communication.

# Example

```julia
datatype = MPI.Types.create_vector(3, 2, 5, MPI.Datatype(Int64))
MPI.Types.commit!(datatype)
```
will create a datatype with the following layout
```
|<----->|  block length

+---+---+---+---+---+---+---+---+---+---+---+---+---+---+---+
| X | X |   |   |   | X | X |   |   |   | X | X |   |   |   |
+---+---+---+---+---+---+---+---+---+---+---+---+---+---+---+

|<---- stride ----->|
```
where each segment represents an `Int64`.

(image by Jonathan Dursi, <https://stackoverflow.com/a/10788351/392585>)

# External links
$(_doc_external("MPI_Type_vector"))
"""
function create_vector(count::Integer, blocklength::Integer, stride::Integer, oldtype::Datatype)
    finalizer(free, create_vector!(Datatype(), count, blocklength, stride, oldtype))
end
function create_vector!(newtype::Datatype, count::Integer, blocklength::Integer, stride::Integer, oldtype::Datatype)
    # int MPI_Type_vector(int count, int blocklength, int stride,
    #          MPI_Datatype oldtype, MPI_Datatype *newtype)
    @mpichk ccall((:MPI_Type_vector, libmpi), Cint,
                  (Cint, Cint, Cint, MPI_Datatype, Ptr{MPI_Datatype}),
                  count, blocklength, stride, oldtype, newtype)
    return newtype
end

"""
    MPI.Types.create_subarray(sizes, subsizes, offset, oldtype::Datatype;
                              rowmajor=false)

Creates a derived [`Datatype`](@ref) describing an `N`-dimensional subarray of size
`subsizes` of an `N`-dimensional array of size `sizes` and element type `oldtype`, with
the first element offset by `offset` (i.e. the 0-based index of the first element).

Column-major indexing (used by Julia and Fortran) is assumed; use the keyword
`rowmajor=true` to specify row-major layout (used by C and numpy).

Note that [`MPI.Types.commit!`](@ref) must be used before the datatype can be used for
communication.

# External links
$(_doc_external("MPI_Type_create_subarray"))
"""
function create_subarray(sizes, subsizes, offset, oldtype::Datatype;
                         rowmajor=false)
    finalizer(free, create_subarray!(Datatype(), sizes, subsizes, offset, oldtype; rowmajor=rowmajor))
end
function create_subarray!(newtype::Datatype, sizes, subsizes, offset, oldtype::Datatype;
                          rowmajor=false)
    @assert (N = length(sizes)) == length(subsizes) == length(offset)
    sizes = sizes isa Vector{Cint} ? sizes : Cint[s for s in sizes]
    subsizes = subsizes isa Vector{Cint} ? subsizes : Cint[s for s in subsizes]
    offset = offset isa Vector{Cint} ? offset : Cint[s for s in offset]

    @mpichk ccall((:MPI_Type_create_subarray, libmpi), Cint,
                  (Cint, Ptr{Cint}, Ptr{Cint}, Ptr{Cint}, Cint, MPI_Datatype, Ptr{MPI_Datatype}),
                  N, sizes, subsizes, offset,
                  rowmajor ? MPI.MPI_ORDER_C : MPI.MPI_ORDER_FORTRAN,
                  oldtype, newtype)
    return newtype
end

"""
    MPI.Types.create_struct(blocklengths, displacements, types)

Creates a derived [`Datatype`](@ref) describing a struct layout.

Note that [`MPI.Types.commit!`](@ref) must be used before the datatype can be used for
communication.

# External links
$(_doc_external("MPI_Type_create_struct"))
"""
function create_struct(blocklengths, displacements, types)
    finalizer(free, create_struct!(Datatype(), blocklengths, displacements, types))
end
function create_struct!(newtype::Datatype, blocklengths, displacements, types)
    @assert (N = length(blocklengths)) == length(displacements) == length(types)
    blocklengths = blocklengths isa Vector{Cint} ? blocklengths : Cint[s for s in blocklengths]
    displacements = displacements isa Vector{MPI_Aint} ? displacements : MPI_Aint[s for s in displacements]
    # int MPI_Type_create_struct(int count,
    #                            const int array_of_blocklengths[],
    #                            const MPI_Aint array_of_displacements[],
    #                            const MPI_Datatype array_of_types[],
    #                            MPI_Datatype *newtype)
    GC.@preserve types begin
        mpi_types = [t.val for t in types]
        @mpichk ccall((:MPI_Type_create_struct, libmpi), Cint,
                      (Cint, Ptr{Cint}, Ptr{MPI_Aint}, Ptr{MPI_Datatype}, Ptr{MPI_Datatype}),
                      N, blocklengths, displacements, mpi_types, newtype)
    end
    return newtype
end

"""
    MPI.Types.create_resized(oldtype::Datatype, lb::Integer, extent::Integer)

Creates a new [`Datatype`](@ref) that is identical to `oldtype`, except that the lower
bound of this new datatype is set to be `lb`, and its upper bound is set to be `lb +
extent`.

Note that [`MPI.Types.commit!`](@ref) must be used before the datatype can be used for
communication.

# See also
- [`MPI.Types.extent`](@ref)

# External links
$(_doc_external("MPI_Type_create_resized"))
"""
function create_resized(oldtype::Datatype, lb::Integer, extent::Integer)
    finalizer(free, create_resized!(Datatype(), oldtype, lb, extent))
end
function create_resized!(newtype::Datatype, oldtype::Datatype, lb::Integer, extent::Integer)
    # int MPI_Type_create_resized(MPI_Datatype oldtype, MPI_Aint lb,
    #              MPI_Aint extent, MPI_Datatype *newtype)
    @mpichk ccall((:MPI_Type_create_resized, libmpi), Cint,
                  (MPI_Datatype, Cptrdiff_t, Cptrdiff_t, Ptr{MPI_Datatype}),
                  oldtype, lb, extent, newtype)
    return newtype
end

function duplicate!(newtype::Datatype, oldtype::Datatype)
    # int MPI_Type_dup(MPI_Datatype oldtype, MPI_Datatype * newtype)
    @mpichk ccall((:MPI_Type_dup, libmpi), Cint,
                  (MPI_Datatype, Ptr{MPI_Datatype}),
                  oldtype, newtype)
    return newtype
end
"""
    MPI.Types.duplicate(oldtype::Datatype)

Duplicates the datatype `oldtype`.

# External links
$(_doc_external("MPI_Type_dup"))
"""
duplicate(oldtype::Datatype) = duplicate!(Datatype(), oldtype::Datatype)

"""
    MPI.Types.commit!(newtype::Datatype)

Commits the [`Datatype`](@ref) `newtype` so that it can be used for communication.
Returns `newtype`.

# External links
$(_doc_external("MPI_Type_commit"))
"""
function commit!(newtype::Datatype)
    # int MPI_Type_commit(MPI_Datatype *datatype)
    @mpichk ccall((:MPI_Type_commit, libmpi), Cint,
                  (Ptr{MPI_Datatype},), newtype)
    return newtype
end

function create!(newtype::Datatype, ::Type{T}) where {T}
    if !isbitstype(T)
        throw(ArgumentError("Type must be isbitstype"))
    end
    blocklengths = Cint[]
    displacements = MPI_Aint[]
    types = Datatype[]

    if isprimitivetype(T)
        # primitive type
        szrem = sz = sizeof(T)
        disp = 0
        for (i,basetype) in (8 => Datatype(UInt64), 4 => Datatype(UInt32), 2 => Datatype(UInt16), 1 => Datatype(UInt8))
            if sz == i
                return MPI.Types.duplicate!(newtype, basetype)
            end
            blk, szrem = divrem(szrem, i)
            if blk != 0
                push!(blocklengths, blk)
                push!(displacements, disp)
                push!(types, basetype)
                disp += i * blk
            end
        end
    else
        # struct
        Fprev = nothing
        for i in 1:fieldcount(T)
            F = fieldtype(T,i)
            offset = fieldoffset(T,i)
            if sizeof(F) == 0
                continue
            elseif F == Fprev
                blocklengths[end] += 1
            else
                push!(blocklengths, 1)
                push!(displacements, offset)
                push!(types, Datatype(F))
                Fprev = F
            end
        end
    end
    create_struct!(newtype, blocklengths, displacements, types)
end

end # module

function Get_address(location)
    addr = Ref{Cptrdiff_t}(0)
    @mpichk ccall((:MPI_Get_address, libmpi), Cint, (Ptr{Cvoid}, Ref{MPI_Aint}), location, addr)
    return addr[]
end


text:
# vim: set et ts=2 sw=2;

# ## Uses
println("# lib.jl ...")
using Random
using Parameters

# -------------------------------------------------------------------
# ## Misc Utils
# ### One-liners.
within(m,n,x) = x<m ? m : (x>n ? n : x)
same(s)  = s                                  #noop
int(x)   = floor(Int,x)                       #round
per(a,n) = a[within(1,length(a), int(length(a)*n))]  #percentile
thing(x) = try parse(Float64,x) catch _ x end #coerce
say(i)   = println(o(i))                      #print+nl
any(a)   = a[ int(length(a) * rand()) + 1 ]   #pick any one
few(a,n=it.divs.few) =                        #pick many
  length(a)<n ? a : [any(a) for _ in 1:n]

# ### How to print a struct
# Skips any fields starting with "`_`".
oo(x) = println(o(x))

o(i::String)     = i
o(i::SubString)  = i
o(i::Char)       = string(i)
o(i::Number)     = string(i)
o(i::Array)      = "["*join(map(o,i),", ")*"]"
o(i::NamedTuple) = "("*join(map(o,i),", ")*")"
o(i::Dict)       = "{"*join(["$k="*o(v) for (k,v) in i],", ")*"}"
o(i::Any) = begin
  s, pre="$(typeof(i)){", ""
  for f in sort([x for x in fieldnames(typeof(i))
                  if ("$x"[1] != '_')])
    s = s * pre * "$f=$(o(getfield(i,f)))"
    pre=", " end
  return s * "}" end

# ### How to read a CSV File
# Skip blank lines. Coerce numeric strings to numbers.
using ResumableFunctions
@resumable function csv(file;zap=r"(\s+|#.*)") #iterate on file
  b4=""
  for line in eachline(file)
    line = replace(line,zap =>"")
    if length(line) != 0
      if line[end] == ',' # if line ends with ",",
        b4 = b4 * line    # join it to next
      else
        @yield [thing(x) for x in split(b4*line,",")]
        b4 = "" end end end end

max_stars_repo_name: timm/peek

max_issues_repo_name: timm/peek

max_forks_repo_name: timm/peek


text:
module PenPlots

# Types

include("geomtypes.jl")
include("metatypes.jl")

# Functions

include("extent.jl")
include("svg.jl")
include("rotation.jl")
include("noise.jl")

end


text:
using Distributions, Random
using Statistics
using SpecialFunctions
const pg_t = 0.64
const pg_inv_t = inv(pg_t)

"""
    PolyaGamma(b::Real, c::Real)

## Arguments
- `b::Real`
- `c::Real` exponential tilting

## Keyword Arguments
- `nmax::Int` Sampling parameter
- `trunc::Int` Sampling parameter

Create a PolyaGamma sampler with parameters `b` and `c`
"""
struct PolyaGamma{Tb,Tc} <: Distributions.ContinuousUnivariateDistribution
    # For sum of Gammas.
    b::Tb
    c::Tc
end

Base.eltype(::PolyaGamma{T,Tc}) where {T,Tc} = Tc

Distributions.params(d::PolyaGamma) = (d.b, d.c)

Statistics.mean(d::PolyaGamma) = d.b / (2 * d.c) * tanh(d.c / 2)

Base.minimum(d::PolyaGamma) = zero(eltype(d))
Base.maximum(::PolyaGamma) = Inf
Distributions.insupport(::PolyaGamma, x::Real) = zero(x) <= x < Inf

function Distributions.pdf(d::PolyaGamma, x::Real)
    b, c = params(d)
    iszero(x) && return zero(x)
    return _tilt(x, b, c) * 2^(b - 1) / gamma(b) * sum(0:200) do n
        ifelse(iseven(n), 1, -1) * exp(
            loggamma(n + b) - loggamma(n + 1) + log(2n + b) - log(twoπ * x^3) / 2 -
            (2n + b)^2 / (8x),
        )
    end
end

function _tilt(ω, b, c)
    return cosh(c / 2)^b * exp(-c^2 / 2 * ω)
end

function Distributions.rand(rng::AbstractRNG, d::PolyaGamma)
    if iszero(d.b)
        return zero(eltype(d))
    end
    return draw_sum(rng, d)
end

## Sampling when `b` is an integer
function draw_sum(rng::AbstractRNG, d::PolyaGamma{<:Int})
    return sum(Base.Fix1(sample_pg1, rng), d.c * ones(d.b))
end

function draw_sum(rng::AbstractRNG, d::PolyaGamma{<:Real})
    if d.b < 1
        return rand_gamma_sum(rng, d, d.b)
    end
    trunc_b = floor(Int, d.b)
    res_b = d.b - trunc_b
    trunc_term = sum(Base.Fix1(sample_pg1, rng), d.c * ones(trunc_b))
    res_term = rand_gamma_sum(rng, d, res_b)
    return trunc_term + res_term
end

## Utility functions
function a(n::Int, x::Real)
    k = (n + 0.5) * π
    if x > pg_t
        return k * exp(-k^2 * x / 2)
    elseif x > 0
        expnt = -3 / 2 * (log(halfπ) + log(x)) + log(k) - 2 * (n + 1//2)^2 / x
        return exp(expnt)
    else
        error("x should be a positive real")
    end
end

function mass_texpon(z::Real)
    t = pg_t

    K = π^2 / 8 + z^2 / 2
    b = sqrt(inv(t)) * (t * z - 1)
    a = -sqrt(inv(t)) * (t * z + 1)

    x0 = log(K) + K * t
    xb = x0 - z + logcdf(Distributions.Normal(), b)
    xa = x0 + z + logcdf(Distributions.Normal(), a)

    qdivp = fourinvπ * (exp(xb) + exp(xa))

    return 1 / (1 + qdivp)
end

# Sample from a truncated inverse gaussian
function rand_truncated_inverse_gaussian(rng::AbstractRNG, z::Real)
    μ = inv(z)
    x = one(z) + pg_t
    if μ > pg_t
        d_exp = Exponential()
        while true
            E = rand(rng, d_exp)
            E′ = rand(rng, d_exp)
            while E^2 > 2E′ / pg_t
                E = rand(rng, d_exp)
                E′ = rand(rng, d_exp)
            end
            x = pg_t / (1 + E * pg_t)^2
            α = exp(-z^2 * x / 2)
            α >= rand(rng) && break
        end
    else
        while (x > pg_t)
            Y = randn(rng)^2
            μY = μ * Y
            x = μ + μ * μY / 2 - μ / 2 * sqrt(4 * μY + μY^2)
            if rand(rng) > μ / (μ + x)
                x = μ^2 / x
            end
            x > pg_t && break
        end
    end
    return x
end

# Sample from PG(1, z)
# Algorithm 1 from "Bayesian Inference for logistic models..." p. 26
function sample_pg1(rng::AbstractRNG, z::Real)
    # Change the parameter.
    z = abs(z) / 2

    # Now sample 0.25 * J^*(1, Z := Z/2).
    K = π^2 / 8 + z^2 / 2
    t = pg_t

    r = mass_texpon(z)

    while true
        if r > rand(rng) # sample from truncated exponential
            x = t + rand(rng, Exponential()) / K
        else # sample from truncated inverse Gaussian
            x = rand_truncated_inverse_gaussian(rng,z)
        end
        s = a(0, x)
        y = rand(rng) * s
        n = 0
        while true
            n = n + 1
            if isodd(n)
                s = s - a(n, x)
                y <= s && return x / 4
            else
                s = s + a(n, x)
                y > s && break
            end
        end
    end
end # Sample PG(1, c)

# Sample ω as the series of Gamma variables (truncated at 200)
function rand_gamma_sum(rng::AbstractRNG, d::PolyaGamma, e::Real)
    C = inv2π / π
    c = d.c
    w = (c * inv2π)^2
    d = Gamma(e, 1)
    return C * sum(1:200) do k
        rand(rng, d) / ((k - 0.5)^2 + w)
    end
end


text:
# ---
# title: 219. Contains Duplicate II
# id: problem219
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: Array, Hash Table
# link: <https://leetcode.com/problems/contains-duplicate-ii/description/>
# hidden: true
# ---
#
# Given an array of integers and an integer _k_ , find out whether there are two
# distinct indices _i_ and _j_ in the array such that **nums[i] = nums[j]** and
# the **absolute** difference between _i_ and _j_ is at most _k_.
#
# **Example 1:**
#
#
#
#     Input: nums = [1,2,3,1], k = 3
#     Output: true
#
#
# **Example 2:**
#
#
#
#     Input: nums = [1,0,1,1], k = 1
#     Output: true
#
#
# **Example 3:**
#
#
#
#     Input: nums = [1,2,3,1,2,3], k = 2
#     Output: false
#
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


text:
using GLVisualize, GeometryTypes, Colors
using Reactive, GLAbstraction

if !isdefined(:runtests)
	window = glscreen()
	timesignal = loop(linspace(0f0, 1f0, 360))
end

description = """
Example demoing the linesegment visualization and how to use indices to connect
points with line segments. Indices can be used like this on other primitives,
E.g. particles and the normal line type.
"""

large_sphere = HyperSphere(Point3f0(0), 1f0)
rotation_angle = const_lift(*, timesignal, 2f0*pi)
rotation = map(rotationmatrix_z, rotation_angle)

positions = decompose(Point3f0, large_sphere)
indices = rand(range(Cuint(0), Cuint(length(positions))), 1000)

color = map(large_sphere->RGBA{Float32}(large_sphere, 0.9f0), colormap("Blues", length(positions)))
color2 = map(large_sphere->RGBA{Float32}(large_sphere, 1f0), colormap("Blues", length(positions)))

lines = visualize(
	positions, :linesegment, thickness=0.5f0,
	color=color, indices=indices, model=rotation
)
spheres = visualize(
	(Sphere{Float32}(Point3f0(0.0), 1f0), positions),
	color=color2, scale=Vec3f0(0.05), model=rotation
)
_view(lines, window, camera=:perspective)
_view(spheres, window, camera=:perspective)

if !isdefined(:runtests)
	renderloop(window)
end


text:
using Knet
include("../bilstm.jl")
function test_convert()
    atype = (gpu() >= 0 ? KnetArray{Float32} : Array{Float32})
    batchsize = 25; hiddens = [280]; embedding = 256; vocabsize = 10000
    m = initmodel(atype, hiddens, embedding, vocabsize)
    mco = convertmodel(m)
    for k in keys(mco)
        @assert(mco[k] == m[k])
    end
    info("Conversion tast passed")
end
!isinteractive() && test_convert()

max_stars_repo_name: kirnap/wordsubs

max_issues_repo_name: kirnap/wordsubs

max_forks_repo_name: kirnap/wordsubs


text:
using StatsPlots,Dates,JLD2,Statistics,Optim,Parameters,Distributions,DataFrames,CSV
using Plots.PlotMeasures,OrdinaryDiffEq,DiffEqCallbacks
import KenyaCoVSD
include("fitting_methods.jl");
include("plotting_methods.jl");
## Load the collated fits
@load("data/N_kenya.jld2")
@load("data/linelist_data_with_pos_neg_20feb_to_27apr_c_age.jld2")
@load("data/serological_data_with_20feb_to_31dec_age.jld2")
@load("data/serological_data_with_20feb_to_10Mar2021_age.jld2")
@load("data/cleaned_linelist20210521_deaths_c__date_of_lab_confirmation.jld2")
@load("data/p_ID.jld2")
@load("forecasts/condensed_county_forecasts.jld2")
@load("data/rel_sero_detection_after_infection.jld2")
@load("data/cleaned_linelist20210521_c.jld2")#<--- Positive tests only looking ahead of linelist we fitted to

## Loop over fitted counties to get county specific plots of incidence rates, population exposure, deaths and positive case rates

for fit in condensed_county_forecasts
    inc_plt = plot_group_incidence(fit);
    sero_plt = plot_pop_exposure(fit,serological_data,serology_data,N_kenya);
    deaths_plt = plot_deaths(fit,linelist_data_with_pos_neg);
    PCR_plt = plot_PCR(fit,linelist_data_with_pos_neg,linelist_data);
    savefig(inc_plt,"plots/county_plots/infection_rates/$(fit.name)_infection_rate.png")
    savefig(sero_plt,"plots/county_plots/pop_exposure/$(fit.name)_pop_exposure.png")
    savefig(deaths_plt,"plots/county_plots/deaths_rates/$(fit.name)_deaths_rate.png")
    savefig(PCR_plt,"plots/county_plots/PCR_positives/$(fit.name)_pos_cases_rate.png")
end

##Loop over county fits to get R(t) plots for every county

fitfiles = readdir("modelfits",join=true)

for filename in fitfiles
    fit_dict = load(filename)
    model = fit_dict[first(keys(fit_dict))]
    plt_Rt = plot_Rt_both_SES_groups(model,Date(2021,6,1))
    savefig(plt_Rt,"plots/county_plots/Rt/$(model.areaname)_Rt.png")
end