text:
################################################################################
#  Copyright 2020, Marta Vanin, Tom Van Acker                                  #
################################################################################
# PowerModelsDistributionStateEstimation.jl                                    #
# An extention package of PowerModelsDistribution.jl for Static Power System   #
# State Estimation.                                                            #
################################################################################

@enum ConnConfig WYE DELTA

function constraint_mc_generator_power_se(pm::_PMD.AbstractUnbalancedIVRModel, id::Int; nw::Int=_IM.nw_id_default, report::Bool=true, bounded::Bool=true)
    generator = _PMD.ref(pm, nw, :gen, id)
    bus =  _PMD.ref(pm, nw,:bus, generator["gen_bus"])

    N = length(generator["connections"])
    pmin = get(generator, "pmin", fill(-Inf, N))
    pmax = get(generator, "pmax", fill( Inf, N))
    qmin = get(generator, "qmin", fill(-Inf, N))
    qmax = get(generator, "qmax", fill( Inf, N))

    if get(generator, "configuration", WYE) == _PMD.WYE
        constraint_mc_generator_power_wye_se(pm, nw, id, bus["index"], generator["connections"], pmin, pmax, qmin, qmax; report=report, bounded=bounded)
    else
        constraint_mc_generator_power_delta_se(pm, nw, id, bus["index"], generator["connections"], pmin, pmax, qmin, qmax; report=report, bounded=bounded)
    end
end

"wye connected generator setpoint constraint for IVR formulation - SE adaptation"
function constraint_mc_generator_power_wye_se(pm::_PMD.AbstractUnbalancedIVRModel, nw::Int, id::Int, bus_id::Int, connections::Vector{Int}, pmin::Vector, pmax::Vector, qmin::Vector, qmax::Vector; report::Bool=true, bounded::Bool=true)
    vr =  _PMD.var(pm, nw, :vr, bus_id)
    vi =  _PMD.var(pm, nw, :vi, bus_id)
    crg =  _PMD.var(pm, nw, :crg, id)
    cig =  _PMD.var(pm, nw, :cig, id)

    if bounded
        for (idx, c) in enumerate(connections)
            if pmin[c]>-Inf
                JuMP.@constraint(pm.model, pmin[idx] .<= vr[c]*crg[c]  + vi[c]*cig[c])
            end
            if pmax[c]< Inf
                JuMP.@constraint(pm.model, pmax[idx] .>= vr[c]*crg[c]  + vi[c]*cig[c])
            end
            if qmin[c]>-Inf
                JuMP.@constraint(pm.model, qmin[idx] .<= vi[c]*crg[c]  - vr[c]*cig[c])
            end
            if qmax[c]< Inf
                JuMP.@constraint(pm.model, qmax[idx] .>= vi[c]*crg[c]  - vr[c]*cig[c])
            end
        end
    end
end

"delta connected generator setpoint constraint for IVR formulation - adapted for SE"
function constraint_mc_generator_power_delta_se(pm::_PMD.AbstractUnbalancedIVRModel, nw::Int, id::Int, bus_id::Int, connections::Vector{Int}, pmin::Vector, pmax::Vector, qmin::Vector, qmax::Vector; report::Bool=true, bounded::Bool=true)
    vr =  _PMD.var(pm, nw, :vr, bus_id)
    vi =  _PMD.var(pm, nw, :vi, bus_id)
    crg =  _PMD.var(pm, nw, :crg, id)
    cig =  _PMD.var(pm, nw, :cig, id)

    nph = length(pmin)

    prev = Dict(c=>connections[(idx+nph-2)%nph+1] for (idx,c) in enumerate(connections))
    next = Dict(c=>connections[idx%nph+1] for (idx,c) in enumerate(connections))

    vrg = Dict()
    vig = Dict()
    for c in connections
        vrg[c] = JuMP.@NLexpression(pm.model, vr[c]-vr[next[c]])
        vig[c] = JuMP.@NLexpression(pm.model, vi[c]-vi[next[c]])
    end

    if bounded
        JuMP.@NLconstraint(pm.model, [i in 1:nph], pmin[i] <= vrg[i]*crg[i]+vig[i]*cig[i])
        JuMP.@NLconstraint(pm.model, [i in 1:nph], pmax[i] >= vrg[i]*crg[i]+vig[i]*cig[i])
        JuMP.@NLconstraint(pm.model, [i in 1:nph], qmin[i] <= -vrg[i]*cig[i]+vig[i]*crg[i])
        JuMP.@NLconstraint(pm.model, [i in 1:nph], qmax[i] >= -vrg[i]*cig[i]+vig[i]*crg[i])
    end
end

function constraint_mc_current_balance_se(pm::_PMD.AbstractUnbalancedPowerModel, i::Int; nw::Int=_IM.nw_id_default)
    bus = _PMD.ref(pm, nw, :bus, i)
    bus_arcs = _PMD.ref(pm, nw, :bus_arcs_conns_branch, i)
    bus_arcs_sw = _PMD.ref(pm, nw, :bus_arcs_conns_switch, i)
    bus_arcs_trans = _PMD.ref(pm, nw, :bus_arcs_conns_transformer, i)
    bus_gens = _PMD.ref(pm, nw, :bus_conns_gen, i)
    bus_storage = _PMD.ref(pm, nw, :bus_conns_storage, i)
    bus_loads = _PMD.ref(pm, nw, :bus_conns_load, i)
    bus_shunts = _PMD.ref(pm, nw, :bus_conns_shunt, i)

    constraint_mc_current_balance_se(pm, nw, i, bus["terminals"], bus["grounded"], bus_arcs, bus_arcs_sw, bus_arcs_trans, bus_gens, bus_storage, bus_loads, bus_shunts)
end

function constraint_mc_current_balance_se(pm::_PMD.AbstractUnbalancedIVRModel, nw::Int, i::Int, terminals::Vector{Int}, grounded::Vector{Bool}, bus_arcs::Vector{Tuple{Tuple{Int,Int,Int},Vector{Int}}}, bus_arcs_sw::Vector{Tuple{Tuple{Int,Int,Int},Vector{Int}}}, bus_arcs_trans::Vector{Tuple{Tuple{Int,Int,Int},Vector{Int}}}, bus_gens::Vector{Tuple{Int,Vector{Int}}}, bus_storage::Vector{Tuple{Int,Vector{Int}}}, bus_loads::Vector{Tuple{Int,Vector{Int}}}, bus_shunts::Vector{Tuple{Int,Vector{Int}}})
    #NB only difference with pmd is crd_bus replaced by crd, and same with cid
    vr = _PMD.var(pm, nw, :vr, i)
    vi = _PMD.var(pm, nw, :vi, i)

    cr    = get(_PMD.var(pm, nw),    :cr, Dict()); _PMD._check_var_keys(cr, bus_arcs, "real current", "branch")
    ci    = get(_PMD.var(pm, nw),    :ci, Dict()); _PMD._check_var_keys(ci, bus_arcs, "imaginary current", "branch")
    crd   = get(_PMD.var(pm, nw),   :crd, Dict()); _PMD._check_var_keys(crd, bus_loads, "real current", "load")
    cid   = get(_PMD.var(pm, nw),   :cid, Dict()); _PMD._check_var_keys(cid, bus_loads, "imaginary current", "load")
    crg   = get(_PMD.var(pm, nw),   :crg, Dict()); _PMD._check_var_keys(crg, bus_gens, "real current", "generator")
    cig   = get(_PMD.var(pm, nw),   :cig, Dict()); _PMD._check_var_keys(cig, bus_gens, "imaginary current", "generator")
    crs   = get(_PMD.var(pm, nw),   :crs, Dict()); _PMD._check_var_keys(crs, bus_storage, "real currentr", "storage")
    cis   = get(_PMD.var(pm, nw),   :cis, Dict()); _PMD._check_var_keys(cis, bus_storage, "imaginary current", "storage")
    crsw  = get(_PMD.var(pm, nw),  :crsw, Dict()); _PMD._check_var_keys(crsw, bus_arcs_sw, "real current", "switch")
    cisw  = get(_PMD.var(pm, nw),  :cisw, Dict()); _PMD._check_var_keys(cisw, bus_arcs_sw, "imaginary current", "switch")
    crt   = get(_PMD.var(pm, nw),   :crt, Dict()); _PMD._check_var_keys(crt, bus_arcs_trans, "real current", "transformer")
    cit   = get(_PMD.var(pm, nw),   :cit, Dict()); _PMD._check_var_keys(cit, bus_arcs_trans, "imaginary current", "transformer")

    Gs, Bs = _PMD._build_bus_shunt_matrices(pm, nw, terminals, bus_shunts)

    ungrounded_terminals = [(idx,t) for (idx,t) in enumerate(terminals) if !grounded[idx]]

    for (idx, t) in ungrounded_terminals
        JuMP.@NLconstraint(pm.model,  sum(cr[a][t] for (a, conns) in bus_arcs if t in conns)
                                    + sum(crsw[a_sw][t] for (a_sw, conns) in bus_arcs_sw if t in conns)
                                    + sum(crt[a_trans][t] for (a_trans, conns) in bus_arcs_trans if t in conns)
                                    ==
                                      sum(crg[g][t]         for (g, conns) in bus_gens if t in conns)
                                    - sum(crs[s][t]         for (s, conns) in bus_storage if t in conns)
                                    - sum(crd[d][t]         for (d, conns) in bus_loads if t in conns)
                                    - sum( Gs[idx,jdx]*vr[u] -Bs[idx,jdx]*vi[u] for (jdx,u) in ungrounded_terminals) # shunts
                                    )
        JuMP.@NLconstraint(pm.model, sum(ci[a][t] for (a, conns) in bus_arcs if t in conns)
                                    + sum(cisw[a_sw][t] for (a_sw, conns) in bus_arcs_sw if t in conns)
                                    + sum(cit[a_trans][t] for (a_trans, conns) in bus_arcs_trans if t in conns)
                                    ==
                                      sum(cig[g][t]         for (g, conns) in bus_gens if t in conns)
                                    - sum(cis[s][t]         for (s, conns) in bus_storage if t in conns)
                                    - sum(cid[d][t]         for (d, conns) in bus_loads if t in conns)
                                    - sum( Gs[idx,jdx]*vi[u] +Bs[idx,jdx]*vr[u] for (jdx,u) in ungrounded_terminals) # shunts
                                    )
    end
end
#####
# "KCL including transformer arcs and load variables."
function constraint_mc_power_balance_se(pm::_PMD.AbstractUnbalancedPowerModel, i::Int; nw::Int=_IM.nw_id_default)
    bus = _PMD.ref(pm, nw, :bus, i)
    bus_arcs = _PMD.ref(pm, nw, :bus_arcs_conns_branch, i)
    bus_arcs_sw = _PMD.ref(pm, nw, :bus_arcs_conns_switch, i)
    bus_arcs_trans = _PMD.ref(pm, nw, :bus_arcs_conns_transformer, i)
    bus_gens = _PMD.ref(pm, nw, :bus_conns_gen, i)
    bus_storage = _PMD.ref(pm, nw, :bus_conns_storage, i)
    bus_loads = _PMD.ref(pm, nw, :bus_conns_load, i)
    bus_shunts = _PMD.ref(pm, nw, :bus_conns_shunt, i)

    constraint_mc_power_balance_se(pm, nw, i, bus["terminals"], bus["grounded"], bus_arcs, bus_arcs_sw, bus_arcs_trans, bus_gens, bus_storage, bus_loads, bus_shunts)
end

function constraint_mc_power_balance_se(pm::_PMD.AbstractUnbalancedACRModel, nw::Int, i::Int, terminals::Vector{Int}, grounded::Vector{Bool}, bus_arcs::Vector{Tuple{Tuple{Int,Int,Int},Vector{Int}}}, bus_arcs_sw::Vector{Tuple{Tuple{Int,Int,Int},Vector{Int}}}, bus_arcs_trans::Vector{Tuple{Tuple{Int,Int,Int},Vector{Int}}}, bus_gens::Vector{Tuple{Int,Vector{Int}}}, bus_storage::Vector{Tuple{Int,Vector{Int}}}, bus_loads::Vector{Tuple{Int,Vector{Int}}}, bus_shunts::Vector{Tuple{Int,Vector{Int}}})
    #NB only diffeerence is in pd and qd we refer to :qd, :pd instead of :pd_bus, :qd_bus
    vr = _PMD.var(pm, nw, :vr, i)
    vi = _PMD.var(pm, nw, :vi, i)

    p    = get(_PMD.var(pm, nw), :p,      Dict()); _PMD._check_var_keys(p,   bus_arcs,       "active power",   "branch")
    q    = get(_PMD.var(pm, nw), :q,      Dict()); _PMD._check_var_keys(q,   bus_arcs,       "reactive power", "branch")
    pg   = get(_PMD.var(pm, nw), :pg, Dict()); _PMD._check_var_keys(pg,  bus_gens,       "active power",   "generator")
    qg   = get(_PMD.var(pm, nw), :qg, Dict()); _PMD._check_var_keys(qg,  bus_gens,       "reactive power", "generator")
    ps   = get(_PMD.var(pm, nw), :ps,     Dict()); _PMD._check_var_keys(ps,  bus_storage,    "active power",   "storage")
    qs   = get(_PMD.var(pm, nw), :qs,     Dict()); _PMD._check_var_keys(qs,  bus_storage,    "reactive power", "storage")
    psw  = get(_PMD.var(pm, nw), :psw,    Dict()); _PMD._check_var_keys(psw, bus_arcs_sw,    "active power",   "switch")
    qsw  = get(_PMD.var(pm, nw), :qsw,    Dict()); _PMD._check_var_keys(qsw, bus_arcs_sw,    "reactive power", "switch")
    pt   = get(_PMD.var(pm, nw), :pt,     Dict()); _PMD._check_var_keys(pt,  bus_arcs_trans, "active power",   "transformer")
    qt   = get(_PMD.var(pm, nw), :qt,     Dict()); _PMD._check_var_keys(qt,  bus_arcs_trans, "reactive power", "transformer")
    pd   = get(_PMD.var(pm, nw), :pd, Dict()); _PMD._check_var_keys(pd,  bus_loads,      "active power",   "load")
    qd   = get(_PMD.var(pm, nw), :qd, Dict()); _PMD._check_var_keys(pd,  bus_loads,      "reactive power", "load")

    Gs, Bs = _PMD._build_bus_shunt_matrices(pm, nw, terminals, bus_shunts)

    cstr_p = []
    cstr_q = []

    ungrounded_terminals = [(idx,t) for (idx,t) in enumerate(terminals) if !grounded[idx]]

    # pd/qd can be NLexpressions, so cannot be vectorized
    for (idx, t) in ungrounded_terminals
        cp = _PMD.@smart_constraint(pm.model, [p, q, pg, qg, ps, qs, psw, qsw, pt, qt, pd, qd, vr, vi],
              sum(  p[arc][t] for (arc, conns) in bus_arcs if t in conns)
            + sum(psw[arc][t] for (arc, conns) in bus_arcs_sw if t in conns)
            + sum( pt[arc][t] for (arc, conns) in bus_arcs_trans if t in conns)
            ==
              sum(pg[gen][t] for (gen, conns) in bus_gens if t in conns)
            - sum(ps[strg][t] for (strg, conns) in bus_storage if t in conns)
            - sum(pd[load][t] for (load, conns) in bus_loads if t in conns)
            + ( -vr[t] * sum(Gs[idx,jdx]*vr[u]-Bs[idx,jdx]*vi[u] for (jdx,u) in ungrounded_terminals)
                -vi[t] * sum(Gs[idx,jdx]*vi[u]+Bs[idx,jdx]*vr[u] for (jdx,u) in ungrounded_terminals)
            )
        )
        push!(cstr_p, cp)

        cq = _PMD.@smart_constraint(pm.model, [p, q, pg, qg, ps, qs, psw, qsw, pt, qt, pd, qd, vr, vi],
              sum(  q[arc][t] for (arc, conns) in bus_arcs if t in conns)
            + sum(qsw[arc][t] for (arc, conns) in bus_arcs_sw if t in conns)
            + sum( qt[arc][t] for (arc, conns) in bus_arcs_trans if t in conns)
            ==
              sum(qg[gen][t] for (gen, conns) in bus_gens if t in conns)
            - sum(qd[load][t] for (load, conns) in bus_loads if t in conns)
            - sum(qs[strg][t] for (strg, conns) in bus_storage if t in conns)
            + ( vr[t] * sum(Gs[idx,jdx]*vi[u]+Bs[idx,jdx]*vr[u] for (jdx,u) in ungrounded_terminals)
               -vi[t] * sum(Gs[idx,jdx]*vr[u]-Bs[idx,jdx]*vi[u] for (jdx,u) in ungrounded_terminals)
            )
        )
        push!(cstr_q, cq)
    end
end

function constraint_mc_power_balance_se(pm::_PMD.AbstractUnbalancedACPModel, nw::Int, i::Int, terminals::Vector{Int}, grounded::Vector{Bool}, bus_arcs::Vector{Tuple{Tuple{Int,Int,Int},Vector{Int}}}, bus_arcs_sw::Vector{Tuple{Tuple{Int,Int,Int},Vector{Int}}}, bus_arcs_trans::Vector{Tuple{Tuple{Int,Int,Int},Vector{Int}}}, bus_gens::Vector{Tuple{Int,Vector{Int}}}, bus_storage::Vector{Tuple{Int,Vector{Int}}}, bus_loads::Vector{Tuple{Int,Vector{Int}}}, bus_shunts::Vector{Tuple{Int,Vector{Int}}})
    vm   = _PMD.var(pm, nw, :vm, i)
    va   = _PMD.var(pm, nw, :va, i)

    p    = get(_PMD.var(pm, nw), :p,      Dict()); _PMD._check_var_keys(p,   bus_arcs,       "active power",   "branch")
    q    = get(_PMD.var(pm, nw), :q,      Dict()); _PMD._check_var_keys(q,   bus_arcs,       "reactive power", "branch")
    pg   = get(_PMD.var(pm, nw), :pg, Dict()); _PMD._check_var_keys(pg,  bus_gens,       "active power",   "generator")
    qg   = get(_PMD.var(pm, nw), :qg, Dict()); _PMD._check_var_keys(qg,  bus_gens,       "reactive power", "generator")
    ps   = get(_PMD.var(pm, nw), :ps,     Dict()); _PMD._check_var_keys(ps,  bus_storage,    "active power",   "storage")
    qs   = get(_PMD.var(pm, nw), :qs,     Dict()); _PMD._check_var_keys(qs,  bus_storage,    "reactive power", "storage")
    psw  = get(_PMD.var(pm, nw), :psw,    Dict()); _PMD._check_var_keys(psw, bus_arcs_sw,    "active power",   "switch")
    qsw  = get(_PMD.var(pm, nw), :qsw,    Dict()); _PMD._check_var_keys(qsw, bus_arcs_sw,    "reactive power", "switch")
    pt   = get(_PMD.var(pm, nw), :pt,     Dict()); _PMD._check_var_keys(pt,  bus_arcs_trans, "active power",   "transformer")
    qt   = get(_PMD.var(pm, nw), :qt,     Dict()); _PMD._check_var_keys(qt,  bus_arcs_trans, "reactive power", "transformer")
    pd   = get(_PMD.var(pm, nw), :pd, Dict()); _PMD._check_var_keys(pd,  bus_loads,      "active power",   "load")
    qd   = get(_PMD.var(pm, nw), :qd, Dict()); _PMD._check_var_keys(pd,  bus_loads,      "reactive power", "load")

    Gs, Bs = _PMD._build_bus_shunt_matrices(pm, nw, terminals, bus_shunts)

    cstr_p = []
    cstr_q = []

    ungrounded_terminals = [(idx,t) for (idx,t) in enumerate(terminals) if !grounded[idx]]

    for (idx,t) in ungrounded_terminals
        if any(Bs[idx,jdx] != 0 for (jdx, u) in ungrounded_terminals if idx != jdx) || any(Gs[idx,jdx] != 0 for (jdx, u) in ungrounded_terminals if idx != jdx)
            cp = JuMP.@NLconstraint(pm.model,
                  sum(  p[a][t] for (a, conns) in bus_arcs if t in conns)
                + sum(psw[a][t] for (a, conns) in bus_arcs_sw if t in conns)
                + sum( pt[a][t] for (a, conns) in bus_arcs_trans if t in conns)
                - sum( pg[g][t] for (g, conns) in bus_gens if t in conns)
                + sum( ps[s][t] for (s, conns) in bus_storage if t in conns)
                + sum( pd[l][t] for (l, conns) in bus_loads if t in conns)
                + ( # shunt
                    +Gs[idx,idx] * vm[t]^2
                    +sum( Gs[idx,jdx] * vm[t]*vm[u] * cos(va[t]-va[u])
                         +Bs[idx,jdx] * vm[t]*vm[u] * sin(va[t]-va[u])
                        for (jdx,u) in ungrounded_terminals if idx != jdx)
                )
                ==
                0.0
            )
            push!(cstr_p, cp)

            cq = JuMP.@NLconstraint(pm.model,
                  sum(  q[a][t] for (a, conns) in bus_arcs if t in conns)
                + sum(qsw[a][t] for (a, conns) in bus_arcs_sw if t in conns)
                + sum( qt[a][t] for (a, conns) in bus_arcs_trans if t in conns)
                - sum( qg[g][t] for (g, conns) in bus_gens if t in conns)
                + sum( qs[s][t] for (s, conns) in bus_storage if t in conns)
                + sum( qd[l][t] for (l, conns) in bus_loads if t in conns)
                + ( # shunt
                    -Bs[idx,idx] * vm[t]^2
                    -sum( Bs[idx,jdx] * vm[t]*vm[u] * cos(va[t]-va[u])
                         -Gs[idx,jdx] * vm[t]*vm[u] * sin(va[t]-va[u])
                         for (jdx,u) in ungrounded_terminals if idx != jdx)
                )
                ==
                0.0
            )
            push!(cstr_q, cq)
        else
            cp = _PMD.@smart_constraint(pm.model, [p, pg, ps, psw, pt, pd, vm],
                  sum(  p[a][t] for (a, conns) in bus_arcs if t in conns)
                + sum(psw[a][t] for (a, conns) in bus_arcs_sw if t in conns)
                + sum( pt[a][t] for (a, conns) in bus_arcs_trans if t in conns)
                - sum( pg[g][t] for (g, conns) in bus_gens if t in conns)
                + sum( ps[s][t] for (s, conns) in bus_storage if t in conns)
                + sum( pd[l][t] for (l, conns) in bus_loads if t in conns)
                + Gs[idx,idx] * vm[t]^2
                ==
                0.0
            )
            push!(cstr_p, cp)

            cq = _PMD.@smart_constraint(pm.model, [q, qg, qs, qsw, qt, qd, vm],
                  sum(  q[a][t] for (a, conns) in bus_arcs if t in conns)
                + sum(qsw[a][t] for (a, conns) in bus_arcs_sw if t in conns)
                + sum( qt[a][t] for (a, conns) in bus_arcs_trans if t in conns)
                - sum( qg[g][t] for (g, conns) in bus_gens if t in conns)
                + sum( qs[s][t] for (s, conns) in bus_storage if t in conns)
                + sum( qd[l][t] for (l, conns) in bus_loads if t in conns)
                - Bs[idx,idx] * vm[t]^2
                ==
                0.0
            )
            push!(cstr_q, cq)
        end
    end
end

function constraint_mc_power_balance_se(pm::_PMD.SDPUBFPowerModel, nw::Int, i::Int, terminals::Vector{Int}, grounded::Vector{Bool}, bus_arcs::Vector{Tuple{Tuple{Int,Int,Int},Vector{Int}}}, bus_arcs_sw::Vector{Tuple{Tuple{Int,Int,Int},Vector{Int}}}, bus_arcs_trans::Vector{Tuple{Tuple{Int,Int,Int},Vector{Int}}}, bus_gens::Vector{Tuple{Int,Vector{Int}}}, bus_storage::Vector{Tuple{Int,Vector{Int}}}, bus_loads::Vector{Tuple{Int,Vector{Int}}}, bus_shunts::Vector{Tuple{Int,Vector{Int}}})
    Wr = _PMD.var(pm, nw, :Wr, i)
    Wi = _PMD.var(pm, nw, :Wi, i)
    P = get(_PMD.var(pm, nw), :P, Dict()); _PMD._check_var_keys(P, bus_arcs, "active power", "branch")
    Q = get(_PMD.var(pm, nw), :Q, Dict()); _PMD._check_var_keys(Q, bus_arcs, "reactive power", "branch")
    Psw  = get(_PMD.var(pm, nw),  :Psw, Dict()); _PMD._check_var_keys(Psw, bus_arcs_sw, "active power", "switch")
    Qsw  = get(_PMD.var(pm, nw),  :Qsw, Dict()); _PMD._check_var_keys(Qsw, bus_arcs_sw, "reactive power", "switch")
    Pt   = get(_PMD.var(pm, nw),   :Pt, Dict()); _PMD._check_var_keys(Pt, bus_arcs_trans, "active power", "transformer")
    Qt   = get(_PMD.var(pm, nw),   :Qt, Dict()); _PMD._check_var_keys(Qt, bus_arcs_trans, "reactive power", "transformer")

    pd = get(_PMD.var(pm, nw), :pd, Dict()); _PMD._check_var_keys(pd, bus_loads, "active power", "load")
    qd = get(_PMD.var(pm, nw), :qd, Dict()); _PMD._check_var_keys(qd, bus_loads, "reactive power", "load")
    pg = get(_PMD.var(pm, nw), :pg, Dict()); _PMD._check_var_keys(pg, bus_gens, "active power", "generator")
    qg = get(_PMD.var(pm, nw), :qg, Dict()); _PMD._check_var_keys(qg, bus_gens, "reactive power", "generator")
    ps   = get(_PMD.var(pm, nw),   :ps, Dict()); _PMD._check_var_keys(ps, bus_storage, "active power", "storage")
    qs   = get(_PMD.var(pm, nw),   :qs, Dict()); _PMD._check_var_keys(qs, bus_storage, "reactive power", "storage")

    Gs, Bs = _PMD._build_bus_shunt_matrices(pm, nw, terminals, bus_shunts)

    cstr_p = []
    cstr_q = []

    ungrounded_terminals = [(idx,t) for (idx,t) in enumerate(terminals) if !grounded[idx]]

    for (idx,t) in ungrounded_terminals
        cp = JuMP.@constraint(pm.model,
            sum(diag(P[a])[findfirst(isequal(t), conns)] for (a, conns) in bus_arcs if t in conns)
            + sum(diag(Psw[a_sw])[findfirst(isequal(t), conns)] for (a_sw, conns) in bus_arcs_sw if t in conns)
            + sum(diag(Pt[a_trans])[findfirst(isequal(t), conns)] for (a_trans, conns) in bus_arcs_trans if t in conns)
            ==
            sum(pg[g][t] for (g, conns) in bus_gens if t in conns)
            - sum(ps[s][t] for (s, conns) in bus_storage if t in conns)
            - sum(pd[d][t] for (d, conns) in bus_loads if t in conns)
            - diag(Wr*Gs'+Wi*Bs')[idx]
        )
        push!(cstr_p, cp)

        cq = JuMP.@constraint(pm.model,
            sum(diag(Q[a])[findfirst(isequal(t), conns)] for (a, conns) in bus_arcs if t in conns)
            + sum(diag(Qsw[a_sw])[findfirst(isequal(t), conns)] for (a_sw, conns) in bus_arcs_sw if t in conns)
            + sum(diag(Qt[a_trans])[findfirst(isequal(t), conns)] for (a_trans, conns) in bus_arcs_trans if t in conns)
            ==
            sum(qg[g][t] for (g, conns) in bus_gens if t in conns)
            - sum(qs[s][t] for (s, conns) in bus_storage if t in conns)
            - sum(qd[d][t] for (d, conns) in bus_loads if t in conns)
            - diag(-Wr*Bs'+Wi*Gs')[idx]
        )
        push!(cstr_q, cq)
    end
end

function constraint_mc_power_balance_se(pm::_PMD.LPUBFDiagModel, nw::Int, i::Int, terminals::Vector{Int}, grounded::Vector{Bool}, bus_arcs::Vector{Tuple{Tuple{Int,Int,Int},Vector{Int}}}, bus_arcs_sw::Vector{Tuple{Tuple{Int,Int,Int},Vector{Int}}}, bus_arcs_trans::Vector{Tuple{Tuple{Int,Int,Int},Vector{Int}}}, bus_gens::Vector{Tuple{Int,Vector{Int}}}, bus_storage::Vector{Tuple{Int,Vector{Int}}}, bus_loads::Vector{Tuple{Int,Vector{Int}}}, bus_shunts::Vector{Tuple{Int,Vector{Int}}})
    w = _PMD.var(pm, nw, :w, i)
    p   = get(_PMD.var(pm, nw),      :p,   Dict()); _PMD._check_var_keys(p,   bus_arcs, "active power", "branch")
    q   = get(_PMD.var(pm, nw),      :q,   Dict()); _PMD._check_var_keys(q,   bus_arcs, "reactive power", "branch")
    psw = get(_PMD.var(pm, nw),    :psw, Dict()); _PMD._check_var_keys(psw, bus_arcs_sw, "active power", "switch")
    qsw = get(_PMD.var(pm, nw),    :qsw, Dict()); _PMD._check_var_keys(qsw, bus_arcs_sw, "reactive power", "switch")
    pt  = get(_PMD.var(pm, nw),     :pt,  Dict()); _PMD._check_var_keys(pt,  bus_arcs_trans, "active power", "transformer")
    qt  = get(_PMD.var(pm, nw),     :qt,  Dict()); _PMD._check_var_keys(qt,  bus_arcs_trans, "reactive power", "transformer")
    pg  = get(_PMD.var(pm, nw),     :pg,  Dict()); _PMD._check_var_keys(pg,  bus_gens, "active power", "generator")
    qg  = get(_PMD.var(pm, nw),     :qg,  Dict()); _PMD._check_var_keys(qg,  bus_gens, "reactive power", "generator")
    ps  = get(_PMD.var(pm, nw),     :ps,  Dict()); _PMD._check_var_keys(ps,  bus_storage, "active power", "storage")
    qs  = get(_PMD.var(pm, nw),     :qs,  Dict()); _PMD._check_var_keys(qs,  bus_storage, "reactive power", "storage")
    pd  = get(_PMD.var(pm, nw),     :pd,  Dict()); _PMD._check_var_keys(pd,  bus_loads, "active power", "load")
    qd  = get(_PMD.var(pm, nw),     :qd,  Dict()); _PMD._check_var_keys(qd,  bus_loads, "reactive power", "load")

    cstr_p = []
    cstr_q = []

    ungrounded_terminals = [(idx,t) for (idx,t) in enumerate(terminals) if !grounded[idx]]

    for (idx,t) in ungrounded_terminals
        cp = JuMP.@constraint(pm.model,
              sum(  p[a][t] for (a, conns) in bus_arcs if t in conns)
            + sum(psw[a][t] for (a, conns) in bus_arcs_sw if t in conns)
            + sum( pt[a][t] for (a, conns) in bus_arcs_trans if t in conns)
            - sum( pg[g][t] for (g, conns) in bus_gens if t in conns)
            + sum( ps[s][t] for (s, conns) in bus_storage if t in conns)
            + sum( pd[d][t] for (d, conns) in bus_loads if t in conns)
            + sum(diag(ref(pm, nw, :shunt, sh, "gs"))[findfirst(isequal(t), conns)]*w[t] for (sh, conns) in bus_shunts if t in conns)
            ==
            0.0
        )
        push!(cstr_p, cp)

        cq = JuMP.@constraint(pm.model,
              sum(  q[a][t] for (a, conns) in bus_arcs if t in conns)
            + sum(qsw[a][t] for (a, conns) in bus_arcs_sw if t in conns)
            + sum( qt[a][t] for (a, conns) in bus_arcs_trans if t in conns)
            - sum( qg[g][t] for (g, conns) in bus_gens if t in conns)
            + sum( qs[s][t] for (s, conns) in bus_storage if t in conns)
            + sum( qd[d][t] for (d, conns) in bus_loads if t in conns)
            - sum(diag(_PMD.ref(pm, nw, :shunt, sh, "bs"))[findfirst(isequal(t), conns)]*w[t] for (sh, conns) in bus_shunts if t in conns)
            ==
            0.0
        )
        push!(cstr_q, cq)
   end
end


text:
using OrdinaryDiffEq, RecursiveArrayTools, LinearAlgebra, Test

# Define the constants for the PDE
const α₂ = 1.0
const α₃ = 1.0
const β₁ = 1.0
const β₂ = 1.0
const β₃ = 1.0
const r₁ = 1.0
const r₂ = 1.0
const D = 100.0
const γ₁ = 0.1
const γ₂ = 0.1
const γ₃ = 0.1
const N = 256
const X = reshape([i for i in 1:N for j in 1:N],N,N)
const Y = reshape([j for i in 1:N for j in 1:N],N,N)
const α₁ = 1.0.*(X.>=4*N/5)

const Mx = Tridiagonal([1.0 for i in 1:N-1],[-2.0 for i in 1:N],[1.0 for i in 1:N-1])
const My = copy(Mx)
Mx[2,1] = 2.0
Mx[end-1,end] = 2.0
My[1,2] = 2.0
My[end,end-1] = 2.0

# Define the initial condition as normal arrays
u0 = zeros(N,N,3)

const MyA = zeros(N,N);
const AMx = zeros(N,N);
const DA = zeros(N,N);
# Define the discretized PDE as an ODE function
function f(du,u,p,t)
   A = @view  u[:,:,1]
   B = @view  u[:,:,2]
   C = @view  u[:,:,3]
  dA = @view du[:,:,1]
  dB = @view du[:,:,2]
  dC = @view du[:,:,3]
  mul!(MyA,My,A)
  mul!(AMx,A,Mx)
  @. DA = D*(MyA + AMx)
  @. dA = DA + α₁ - β₁*A - r₁*A*B + r₂*C
  @. dB = α₂ - β₂*B - r₁*A*B + r₂*C
  @. dC = α₃ - β₃*C + r₁*A*B - r₂*C
end

# Solve the ODE
prob = ODEProblem(f,u0,(0.0,100.0))
sol = solve(prob,BS3(),progress=true,save_everystep=false,save_start=false)
sol = solve(prob,ROCK2(),progress=true,save_everystep=false,save_start=false)

println("CPU Times")
println("BS3")
@time sol = solve(prob,BS3(),progress=true,save_everystep=false,save_start=false)
println("ROCK2")
@time sol = solve(prob,ROCK2(),progress=true,save_everystep=false,save_start=false)

using CuArrays
gu0 = CuArray(Float32.(u0))
const gMx = CuArray(Float32.(Mx))
const gMy = CuArray(Float32.(My))
const gα₁ = CuArray(Float32.(α₁))

const gMyA = CuArray(zeros(Float32,N,N))
const gAMx = CuArray(zeros(Float32,N,N))
const gDA = CuArray(zeros(Float32,N,N))

function gf(du,u,p,t)
   A = @view  u[:,:,1]
   B = @view  u[:,:,2]
   C = @view  u[:,:,3]
  dA = @view du[:,:,1]
  dB = @view du[:,:,2]
  dC = @view du[:,:,3]
  mul!(gMyA,gMy,A)
  mul!(gAMx,A,gMx)
  @. gDA = D*(gMyA + gAMx)
  @. dA = gDA + gα₁ - β₁*A - r₁*A*B + r₂*C
  @. dB = α₂ - β₂*B - r₁*A*B + r₂*C
  @. dC = α₃ - β₃*C + r₁*A*B - r₂*C
end

prob2 = ODEProblem(gf,gu0,(0.0,100.0))
CuArrays.allowscalar(false)
sol = solve(prob2,BS3(),save_everystep=false,save_start=false)
sol = solve(prob2,ROCK2(),save_everystep=false,save_start=false)
@test sol.t[end] == 100.0

println("GPU Times")
println("BS3")
@time sol = solve(prob2,BS3(),progress=true,save_everystep=false,save_start=false)
println("ROCK2")
@time sol = solve(prob2,ROCK2(),progress=true,save_everystep=false,save_start=false)

ext: jl


text:
module ExtendedVariantModeExtendedVariantFadeTrigger

using ..Ahorn, Maple

# obsolete
@mapdef Trigger "ExtendedVariantMode/ExtendedVariantFadeTrigger" ExtendedVariantFadeTrigger(x::Integer, y::Integer, width::Integer=16, height::Integer=16,
	variantChange::String="Gravity", valueA::Integer=10, valueB::Integer=10, positionMode::String="LeftToRight", revertOnDeath::Bool=true)

@mapdef Trigger "ExtendedVariantMode/FloatExtendedVariantFadeTrigger" FloatExtendedVariantFadeTrigger(x::Integer, y::Integer, width::Integer=16, height::Integer=16,
	variantChange::String="Gravity", valueA::Number=1.0, valueB::Number=1.0, positionMode::String="LeftToRight", revertOnDeath::Bool=true)

const placements = Ahorn.PlacementDict(
	"Extended Variant Fade Trigger (Extended Variant Mode)" => Ahorn.EntityPlacement(
		FloatExtendedVariantFadeTrigger,
		"rectangle"
	)
)

Ahorn.editingOptions(trigger::FloatExtendedVariantFadeTrigger) = Dict{String, Any}(
	"variantChange" => Ahorn.ExtendedVariantDictionary.FloatVariants,
	"positionMode" => Maple.trigger_position_modes
)

end


text:
# Input data
data = map(
    digstrs ->
        map(digstr -> map(v -> join(sort(collect(v))), split(digstr, " ")), digstrs),
    map(line -> split(line, " | "), readlines("Day08.input")),
);

# Part 1
basemappings =
    (keys) -> Dict(
        first(filter(v -> length(v) == 2, keys)) => 1,
        first(filter(v -> length(v) == 4, keys)) => 4,
        first(filter(v -> length(v) == 3, keys)) => 7,
        first(filter(v -> length(v) == 7, keys)) => 8,
    )

ans = sum(sum(map(rec -> map(num -> haskey(basemappings(rec[1]), num), rec[2]), data)));
println("Part 1: $ans")

# Part 2
allmappings =
    (keys) -> (
        base = basemappings(keys);
        rev = Dict(v => k for (k, v) in base);
        merge(
            base,
            Dict(
                first(filter(v -> all([length(v) == 6, rev[1] ⊆ v, rev[4] ⊈ v]), keys)) => 0,
                first(filter(v -> all([length(v) == 5, rev[8] ⊆ (rev[4] ∪ v)]), keys)) => 2,
                first(filter(v -> all([length(v) == 5, rev[1] ⊆ v]), keys)) => 3,
                first(
                    filter(
                        v -> all([length(v) == 5, rev[8] ⊈ (rev[4] ∪ v), rev[1] ⊈ v]),
                        keys,
                    ),
                ) => 5,
                first(filter(v -> all([length(v) == 6, rev[1] ⊈ v, rev[4] ⊈ v]), keys)) => 6,
                first(filter(v -> all([length(v) == 6, rev[4] ⊆ v]), keys)) => 9,
            ),
        )
    )

ans = sum(
    map(
        rec -> reduce((acc, v) -> acc * 10 + allmappings(rec[1])[v], rec[2], init = 0),
        data,
    ),
);
println("Part 2: $ans")


text:
using Test
using Dates
using Printf
using EzXML
using XML2011

##### TESTE UNITARIO COM MOCK DATA
 data = Date(2020,1,1)
 cnpj = "12345"
 tipo = XML2011.Inclusao()
 responsavel = XML2011.Responsavel("Fulano", "555-1234", "fulano@banco.com")
 c1 = XML2011.Conta("503000", 5.75)
 @test XML2011.get_valor(c1) == 5.75
 c2 = XML2011.Conta(
     "121000",
     [
         XML2011.DetalheConta([XML2011.Moeda(:USD), XML2011.Posicao(:onshore)], 10.00),
         XML2011.DetalheConta([XML2011.Moeda(:EUR), XML2011.Posicao(:offshore)], 5.10)
     ]
 )
 @test XML2011.get_valor(c2) == 15.10
 @test XML2011.get_valor(c2, [XML2011.Moeda(:USD), XML2011.Posicao(:onshore)]) == 10.00
 @test XML2011.get_valor(c2, [XML2011.Posicao(:offshore), XML2011.Moeda(:EUR)]) ==  5.10 # ordem inversa
 @test_throws ErrorException XML2011.get_valor(c2, [XML2011.Posicao(:offshore), XML2011.Moeda(:JPY)]) # inexiste

 c3 = XML2011.Conta("410400", 5.75)
 c4 = XML2011.Conta("410401", 5.75)
 doc = XML2011.Doc2011(data, cnpj, tipo, responsavel, [c1, c2, c3, c4])

 io = IOBuffer()
 XML2011.write_xml(io, doc)
 xml_written = String(take!(io))
 xml_test = """
 <?xml version="1.0" encoding="UTF-8"?>
 <documentoDDR cnpj="12345" dataBase="2020-01-01" codigoDocumento="2011" tipoEnvio="I">
   <parametros>
     <parametro codigoParametro="31" valorParametro="Fulano"/>
     <parametro codigoParametro="32" valorParametro="555-1234"/>
     <parametro codigoParametro="33" valorParametro="fulano@banco.com"/>
   </parametros>
   <contas>
     <conta codigoConta="503000" valorConta="5.75"/>
     <conta codigoConta="121000" valorConta="15.10">
       <detalhamentosDDR>
         <detalhamentoDDR valorDetalhe="10.00">
           <detalhe codigoElemento="83" valorElemento="USD"/>
           <detalhe codigoElemento="84" valorElemento="1"/>
         </detalhamentoDDR>
         <detalhamentoDDR valorDetalhe="5.10">
           <detalhe codigoElemento="83" valorElemento="EUR"/>
           <detalhe codigoElemento="84" valorElemento="2"/>
         </detalhamentoDDR>
       </detalhamentosDDR>
     </conta>
     <conta codigoConta="410400" valorConta="5.75"/>
     <conta codigoConta="410401" valorConta="5.75"/>
   </contas>
 </documentoDDR>
 """

 function is_equal(x1::EzXML.Document, x2::EzXML.Document)

     function get_parsed_string(x::EzXML.Document)
         b = IOBuffer()
         prettyprint(b, x)
         return String(take!(b))
     end

     s1 = get_parsed_string(x1)
     s2 = get_parsed_string(x2)

     return s1 == s2

 end

 @test is_equal(parsexml(xml_test), parsexml(xml_written))

#### TESTES READ
# Parse String
doc_read1 = XML2011.parse_xml(xml_test)

# Read File
file = "ddr.xml"
XML2011.write_xml(file, doc_read1)
doc_read2 = XML2011.read_xml(file)
rm(file)

# Comparar
b1 = IOBuffer(); XML2011.write_xml(b1, doc_read1); s1 = String(take!(b1))
b2 = IOBuffer(); XML2011.write_xml(b2, doc_read2); s2 = String(take!(b2))
@test s1 == s2

 #### TESTES GENERICOS
 @test XML2011.Moeda(:USD) == XML2011.Moeda(:USD)
 @test XML2011.Moeda(:USD) != XML2011.Moeda(:EUR)
 @test XML2011.Moeda(:USD) != XML2011.Posicao(:onshore)

## Validacao de formula
# "310000" => doc -> get_valor(doc, "310100") * (get_valor(doc, "310105")/100)
v = Vector{XML2011.Conta}()
c1 = XML2011.Conta("310100", 1000.0)
c2 = XML2011.Conta("310105", 90.0)
c3 = XML2011.Conta("310000", 1000.0 * (90/100))
c3_wrong = XML2011.Conta("310000", 1000.0 * (89/100))
c4 = XML2011.Conta("310101", 700.0)
c5 = XML2011.Conta("310104", 300.0)

doc = XML2011.Doc2011(data, cnpj, tipo, responsavel, [c1, c2, c3, c4, c5])

@test_throws AssertionError XML2011.Doc2011(data, cnpj, tipo, responsavel, [c1, c2, c3_wrong, c4, c5])


text:
# This file includes code that was formerly a part of Julia.
# License is MIT: LICENSE.md

using MurmurHash3

@static VERSION < v"0.7.0-DEV" ? (using Base.Test) : (using Test)

_memhash(siz, ptr, seed) =
    ccall(Base.memhash, UInt, (Ptr{UInt8}, Csize_t, UInt32), ptr, siz, seed % UInt32)

p1 = SubString("--hello--",3,7)
p2 = "hello"

mmhash(str::String) = mmhash128_a(sizeof(str), pointer(str), 0%UInt32)
@static if sizeof(Int) == 8
    mmhashc(str::AbstractString) = mmhash128_c(str, 0%UInt32)
else
    mmhashc(str::AbstractString) = (s = string(str); mmhash128_c(sizeof(s), pointer(s), 0%UInt32))
end
memhash(str) = _memhash(sizeof(str), pointer(str), 0%UInt32)

@testset "MurmurHash3" begin
    @test mmhashc(p1) == mmhash(p2)
    @test last(mmhashc(p1)) == memhash(p1)
    @test last(mmhashc(p2)) == memhash(p2)
    @test last(mmhash(p2))  == memhash(p1)
end


text:
"""
Define the `hadwick_james` find_circuits algorithm.
"""
find_circuits(g::DiGraph, ::Type{Val{:hadwick_james}}) = hadwick_james(g)

"""
Find circuits using circuit! for each vertex in the graph `g`.
"""
function hadwick_james(g::DiGraph)
    nvg = nv(g)
    T = eltype(g)
    B = [T[] for i in 1:nvg]
    blocked = [false for i in 1:nvg]
    stack = T[]
    cycles = Vector{Vector{T}}()
    for v in vertices(g)
        circuit!(g, v, v, blocked, B, stack, cycles)
        resetblocked!(blocked)
        resetB!(B)
    end
    return cycles
end

"""
Reset B work structure.
"""
resetB!(B) = map!(empty!, B, B)

"""
Reset vector of "blocked" vertices.
"""
resetblocked!(blocked) = fill!(blocked, false)

"""
Find circuits starting from v1 recursively.
"""
function circuit!(g::DiGraph, v1::T, v2::T, blocked::Vector{Bool}, B::Vector{Vector{T}}, stack::Vector{T}, cycles::Vector{Vector{T}}) where T
    f = false
    push!(stack, v2)
    blocked[v2] = true

    Av = outneighbors(g, v2)
    for w in Av
        (w < v1) && continue
        if w == v1 # Found a circuit
            push!(cycles, copy(stack))
            f = true
        elseif !blocked[w]
            f = circuit!(g, v1, w, blocked, B, stack, cycles)
        end
    end
    if f
        unblock!(v2, blocked, B)
    else
        for w in Av
            (w < v1) && continue
            if !(v2 in B[w])
                push!(B[w], v2)
            end
        end
    end
    pop!(stack)
    return f
end

"""
Simultaneously count and remove occurences of a value `val` in the array `list`.
"""
function countAndFilter!(val::T, list::AbstractArray{T}) where T
    nocc = 0
    function doFilter(v)
        if v == val
            nocc += 1
            false
        else
            true
        end
    end
    filter!(doFilter, list)
    return nocc
end

"""
Unblock the value `v` from the blocked list and remove from `B`.
"""
function unblock!(v::T, blocked::Vector{Bool}, B::Vector{Vector{T}}) where T
    blocked[v] = false
    wPos = 1
    Bv = B[v]
    while wPos <= length(Bv)
        w = Bv[wPos]
        wPos += 1 - countAndFilter!(w, Bv)
        blocked[w] && unblock!(w, blocked, B)
    end
    return nothing
end


text:
function nlp_cvx_107_010(
    optimizer,
    objective_tol,
    primal_tol,
    dual_tol,
    termination_target = TERMINATION_TARGET_LOCAL,
    primal_target = PRIMAL_TARGET_LOCAL,
)
    # Test Goals:
    # - convex objective
    # - single simple constraint
    # Variants
    #   010 - no binding constraints
    #   011 - binding constraint
    #   012 - binding constraint, different starting point

    model = Model(optimizer)

    @variable(model, x)
    @variable(model, y)

    @objective(model, Min, (x - 0.5)^2 + (y - 0.5)^2)
    @NLconstraint(model, x^2 + y^2 <= 1)

    optimize!(model)

    check_status(model, FEASIBLE_PROBLEM, termination_target, primal_target)
    check_objective(model, 0, tol = objective_tol)
    return check_solution([x, y], [1 / 2, 1 / 2], tol = primal_tol)
end


text:
using IterativelyHypergraphEmbedding
using Documenter

DocMeta.setdocmeta!(IterativelyHypergraphEmbedding, :DocTestSetup, :(using IterativelyHypergraphEmbedding); recursive=true)

makedocs(;
    modules=[IterativelyHypergraphEmbedding],
    authors="Inazuma110 <c011703534@edu.teu.ac.jp> and contributors",
    repo="https://github.com/Inazuma110/IterativelyHypergraphEmbedding.jl/blob/{commit}{path}#{line}",
    sitename="IterativelyHypergraphEmbedding.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://Inazuma110.github.io/IterativelyHypergraphEmbedding.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
    ],
)

deploydocs(;
    repo="github.com/Inazuma110/IterativelyHypergraphEmbedding.jl",
)


text:
include("failure_model.jl")

N_iter = 20 # number of UKI iterations
# need to choose regularization factor α ∈ (0,1],
# when you have enough observation data α=1: no regularization
α_reg = 1.0
# update_freq 1 : approximate posterior covariance matrix with an uninformative prior
#             0 : weighted average between posterior covariance matrix with an uninformative prior and prior
update_freq = 0

process = Unscented(prior_mean, prior_cov, α_reg, update_freq)
ukiobj = EnsembleKalmanProcessModule.EnsembleKalmanProcess(truth_sample, truth.obs_noise_cov, process)

# UKI iterations
err = zeros(N_iter)
for i in 1:N_iter

    params_i = get_u_final(ukiobj)

    g_ens = run_G_ensemble(params_i)

    # analysis step
    EnsembleKalmanProcessModule.update_ensemble!(ukiobj, g_ens)

    err[i] = get_error(ukiobj)[end] #mean((params_true - mean(params_i,dims=2)).^2)
    println(
        "Iteration: " *
        string(i) *
        ", Error: " *
        string(err[i]) *
        " norm(Cov): " *
        string(norm(ukiobj.process.uu_cov[i])),
    )
end

println("True parameters: ")
println(params_true)

println("\nUKI results:")
println(get_u_mean_final(ukiobj))

####
θ_mean_arr = hcat(ukiobj.process.u_mean...)
N_θ = length(ukiobj.process.u_mean[1])
θθ_std_arr = zeros(Float64, (N_θ, N_iter + 1))
for i in 1:(N_iter + 1)
    for j in 1:N_θ
        θθ_std_arr[j, i] = sqrt(ukiobj.process.uu_cov[i][j, j])
    end
end

ites = Array(LinRange(1, N_iter + 1, N_iter + 1))
plot(ites, grid = false, θ_mean_arr[1, :], yerror = 3.0 * θθ_std_arr[1, :], label = "θ1")
plot!(ites, fill(params_true[1], N_iter + 1), linestyle = :dash, linecolor = :grey, label = nothing)

plot!(ites, grid = false, θ_mean_arr[2, :], yerror = 3.0 * θθ_std_arr[2, :], label = "θ2", xaxis = "Iterations")
plot!(ites, fill(params_true[2], N_iter + 1), linestyle = :dash, linecolor = :grey, label = nothing)

#plots
gr(size = (600, 600))
dθ = 1.0
for i in 1:N_iter
    θ_mean, θθ_cov = ukiobj.process.u_mean[i], ukiobj.process.uu_cov[i]
    xx = Array(LinRange(params_true[1] - dθ, params_true[1] + dθ, 100))
    yy = Array(LinRange(params_true[2] - dθ, params_true[2] + dθ, 100))
    xx, yy, Z = Gaussian_2d(θ_mean, θθ_cov, 100, 100, xx = xx, yy = yy)

    p = contour(xx, yy, Z)
    plot!(
        [params_true[1]],
        xaxis = "θ₁",
        yaxis = "θ₂",
        seriestype = "vline",
        linestyle = :dash,
        linecolor = :red,
        label = false,
        title = "UKI iteration = " * string(i),
    )
    plot!([params_true[2]], seriestype = "hline", linestyle = :dash, linecolor = :red, label = "optimum")
    display(p)
    sleep(0.5)
end


text:
type FunctionalCPD{D} <: CPD{D}
    target::NodeName
    parents::NodeNames
    accessor::Function # calling this gives you the distribution from the assignment

    function FunctionalCPD{D}(target::NodeName, accessor::Function) where D
        new(target, NodeName[], accessor)
    end
    function FunctionalCPD{D}(target::NodeName, parents::NodeNames, accessor::Function) where D
        new(target, parents, accessor)
    end
end

name(cpd::FunctionalCPD) = cpd.target
parents(cpd::FunctionalCPD) = cpd.parents
(cpd::FunctionalCPD)(a::Assignment) = cpd.accessor(a)
(cpd::FunctionalCPD)() = (cpd)(Assignment()) # cpd()
(cpd::FunctionalCPD)(pair::Pair{NodeName}...) = (cpd)(Assignment(pair)) # cpd(:A=>1)


text:
using ..Transformers: batchedmul, batched_triu!

# GPT2 specific initializers

function FakeTHLinear(config::HGFGPT2Config, hi, ho; bias=true)
  weight = randn(Float32, ho, hi) .* config.initializer_range
  if bias
    bias = zeros(Float32, ho)
  else
    bias = nothing
  end
  FakeTHLinear(weight, bias)
end

function FakeTHEmbedding(config::HGFGPT2Config, num, dims)
  weight = randn(Float32, dims, num) .* config.initializer_range
  FakeTHEmbedding(0, weight)
end

function FakeTHLayerNorm(config::HGFGPT2Config, dims; eps::Float32=1e-05)
  weight = ones(Float32, dims)
  bias = zeros(Float32, dims)
  FakeTHLayerNorm(eps, weight, bias)
end

function FakeHGFConv1D(config::HGFGPT2Config, nf, nx)
  weight = randn(Float32, nx, nf) .* config.initializer_range
  bias = zeros(Float32, nf)
  FakeHGFConv1D(weight, bias)
end

function FakeHGFSequenceSummary(config::HGFGPT2Config)
  global ACT2FN
  summary_type = SummaryType(Symbol(config.summary_type))
  if config.summary_use_proj
    num_classes = config.num_labels
  else
    num_classes = config.n_embd
  end
  summary = FakeTHLinear(config, config.n_embd, num_classes)
  if isnothing(config.summary_activation)
    activation = identity
  else
    activation = ACT2FN[Symbol(config.summary_activation)]
  end
  return FakeHGFSequenceSummary{summary_type}(summary, activation)
end

# HGF GPT2 compounts

# attention

struct HGFGPT2Attention{P<:FakeHGFConv1D, O<:FakeHGFConv1D} <: THModule
  num_attention_heads::Int
  c_attn::P
  c_proj::O
end

Functors.functor(::Type{<:HGFGPT2Attention}, a) = (c_attn=a.c_attn, c_proj=a.c_proj), y->HGFGPT2Attention(a.num_attention_heads, y...)

function _split_qkv(x)
  h = size(x, 1)
  nh = div(h, 3)
  return (
    @view(x[1:nh, :, :]),
    @view(x[nh+1:2nh, :, :]),
    @view(x[2nh+1:3nh, :, :]),
  )
end

function (self::HGFGPT2Attention)(x, attention_mask,
                                  _output_attentions::Val{output_attentions},
                                  _use_cache::Val{use_cache}
                                  ) where {output_attentions, use_cache}
  x = self.c_attn(x)
  query, key, value = _split_qkv(x)
  query = _split_tranpose_for_scores(query, self.num_attention_heads)
  key = _split_tranpose_for_scores(key, self.num_attention_heads)
  value = _split_tranpose_for_scores(value, self.num_attention_heads)
  self(query, key, value, attention_mask, _output_attentions, _use_cache)
end

function (self::HGFGPT2Attention)(x, past, attention_mask,
                                  _output_attentions::Val{output_attentions},
                                  _use_cache::Val{use_cache}
                                  ) where {output_attentions, use_cache}
  x = self.c_attn(x)
  query, key, value = _split_qkv(x)
  query = _split_tranpose_for_scores(query, self.num_attention_heads)
  key = _split_tranpose_for_scores(key, self.num_attention_heads)
  value = _split_tranpose_for_scores(value, self.num_attention_heads)

  past_key, past_value = past
  key = hcat(past_key, key)
  value = hcat(past_value, value)

  self(query, key, value, attention_mask, _output_attentions, _use_cache)
end

struct ShiftAttentionMask{T}
  mask::T
  past_length::Int
end

function apply_shift_mask(scores, mask)
    out = copy(scores)
    @view(out[1:(end-mask.past_length), :, :, :]) .+= mask.mask
    return out
end

@adjoint function apply_shift_mask(scores, mask)
    out = apply_shift_mask
    return out, Δ -> (Δ, nothing)
end

function _compute_attention_scores(query, key, attention_mask::ShiftAttentionMask)
  w_wo_mask = _compute_attention_scores(query, key, nothing)
  return apply_shift_mask(w_wo_mask, attention_mask)
end

function _attn(query, key, value, attention_mask)
  w = _compute_attention_scores(query, key, attention_mask)
  w = softmax(w; dims=1)
  a = batchedmul(value, w)
  return a, w
end

function (self::HGFGPT2Attention)(query, key, value, attention_mask,
                                  ::Val{output_attentions},
                                  ::Val{use_cache}
                                  ) where {output_attentions, use_cache}

  if use_cache
    present = (key, value)
  end

  a, w = _attn(query, key, value, attention_mask)
  a = _merge_transpose_for_output(a, self.num_attention_heads)
  a = self.c_proj(a)

  if output_attentions && use_cache
    return a, present, w
  elseif output_attentions
    return a, w
  elseif use_cache
    return a, present
  else
    return a
  end
end

function HGFGPT2Attention(config::HGFGPT2Config, nx, n_ctx)
  @assert nx % config.n_head == 0
  c_attn = FakeHGFConv1D(config, 3nx, nx)
  c_proj = FakeHGFConv1D(config, nx, nx)
  return HGFGPT2Attention(config.n_head, c_attn, c_proj)
end

# positionwise

struct HGFGPT2MLP{F, P1<:FakeHGFConv1D, P2<:FakeHGFConv1D} <: THModule
  act::F
  c_fc::P1
  c_proj::P2
end

Functors.functor(::Type{<:HGFGPT2MLP}, mlp) = (c_fc = mlp.c_fc, c_proj = mlp.c_proj), y->HGFGPT2MLP(mlp.act, y...)

(mlp::HGFGPT2MLP)(x) = mlp.c_proj(mlp.act.(mlp.c_fc(x)))

function HGFGPT2MLP(config::HGFGPT2Config, n_state)
  global ACT2FN
  act = ACT2FN[Symbol(config.activation_function)]
  c_fc = FakeHGFConv1D(config, n_state, config.n_embd)
  c_proj = FakeHGFConv1D(config, config.n_embd, n_state)
  return HGFGPT2MLP(act, c_fc, c_proj)
end

# pre-ln transformer block

struct HGFGPT2Block{A<:HGFGPT2Attention, L1<:FakeTHLayerNorm, L2<:FakeTHLayerNorm, M<:HGFGPT2MLP} <: THModule
  ln_1::L1
  attn::A
  ln_2::L2
  mlp::M
end

@functor HGFGPT2Block

function (self::HGFGPT2Block)(x, attention_mask,
                              _output_attentions::Val{output_attentions},
                              _use_cache::Val{use_cache}) where {output_attentions, use_cache}
  _output_attn = self.attn(self.ln_1(x), attention_mask, _output_attentions, _use_cache)

  if output_attentions && use_cache
    output_attn, present, attention_prob = _output_attn
  elseif output_attentions
    output_attn, attention_prob = _output_attn
  elseif use_cache
    output_attn, present = _output_attn
  else
    output_attn = _output_attn
  end

  x = x + output_attn
  m = self.mlp(self.ln_2(x))
  x = x + m

  if output_attentions && use_cache
    return x, present, attention_prob
  elseif output_attentions
    return x, attention_prob
  elseif use_cache
    return x, present
  else
    return x
  end
end

function (self::HGFGPT2Block)(x, past, attention_mask,
                              _output_attentions::Val{output_attentions},
                              _use_cache::Val{use_cache}) where {output_attentions, use_cache}
  _output_attn = self.attn(self.ln_1(x), past, attention_mask, _output_attentions, _use_cache)

  if output_attentions && use_cache
    output_attn, present, attention_prob = _output_attn
  elseif output_attentions
    output_attn, attention_prob = _output_attn
  elseif use_cache
    output_attn, present = _output_attn
  else
    output_attn = _output_attn
  end

  x = x + output_attn
  m = self.mlp(self.ln_2(x))
  x = x + m

  if output_attentions && use_cache
    return x, present, attention_prob
  elseif output_attentions
    return x, attention_prob
  elseif use_cache
    return x, present
  else
    return x
  end
end

function HGFGPT2Block(config::HGFGPT2Config, n_ctx)
  nx = config.n_embd
  ln_1 = FakeTHLayerNorm(config, nx; eps=config.layer_norm_epsilon)
  attn = HGFGPT2Attention(config, nx, n_ctx)
  ln_2 = FakeTHLayerNorm(config, nx; eps=config.layer_norm_epsilon)
  mlp = HGFGPT2MLP(config, 4nx)
  return HGFGPT2Block(ln_1, attn, ln_2, mlp)
end

# gpt2 model

abstract type HGFGPT2PreTrainedModel <: HGFPreTrainedModel end

struct HGFGPT2Model{N, L<:FakeTHModuleList{N},
                    WE<:FakeTHEmbedding,
                    PE<:FakeTHEmbedding,
                    LN<:FakeTHLayerNorm
                    } <: HGFGPT2PreTrainedModel
  wte::WE
  wpe::PE
  h::L
  ln_f::LN
end

@functor HGFGPT2Model

@inline get_past_length(past_key_values) = size(past_key_values[1][1], 2)
@inline get_past_length(::Nothing) = 0

@inline get_word_emb(word_embeddings::FakeTHEmbedding, input_ids::AbstractArray{<:Integer}) = word_embeddings(input_ids)
@inline get_word_emb(word_embeddings::FakeTHEmbedding, input_embed::AbstractArray{T}) where T = input_embed

@inline maybe_past_arange(inputs_embeds, ::Nothing) = _arange(inputs_embeds, size(inputs_embeds)[end-1])
@inline maybe_past_arange(inputs_embeds, past_key_values) = _arange(inputs_embeds, get_past_length(past_key_values), size(inputs_embeds)[end-1])

@inline get_position_emb(position_embeddings::FakeTHEmbedding, inputs_embeds, ::Nothing, past_key_values) = get_position_emb(position_embeddings, inputs_embeds, maybe_past_arange(inputs_embeds, past_key_values), past_key_values)
@inline get_position_emb(position_embeddings::FakeTHEmbedding, inputs_embeds, position_ids, past_key_values) = position_embeddings(position_ids)

@inline get_token_type_emb(token_type_embeddings::FakeTHEmbedding, token_type_ids) = token_type_embeddings(token_type_ids)

@inline apply_token_type_emb(token_type_embeddings::FakeTHEmbedding, embeddings, ::Nothing) = embeddings
@inline apply_token_type_emb(token_type_embeddings::FakeTHEmbedding, embeddings, token_type_ids) = embeddings + get_token_type_emb(token_type_embeddings, token_type_ids)

function (self::HGFGPT2Model)(input, position_ids, token_type_ids, past_key_values)
  inputs_embeds = get_word_emb(self.wte, input)
  position_embeds = get_position_emb(self.wpe, inputs_embeds, position_ids, past_key_values)
  embeddings = inputs_embeds .+ position_embeds
  embeddings = apply_token_type_emb(self.wte, embeddings, token_type_ids)
  return embeddings
end

function (self::HGFGPT2Model)(input; position_ids=nothing, token_type_ids=nothing,
                              past_key_values=nothing,
                              attention_mask=nothing,
                              output_attentions=false,
                              output_hidden_states=false,
                              use_cache=false)
  return self(input, position_ids, token_type_ids,
              past_key_values, attention_mask,
              Val(output_attentions),
              Val(output_hidden_states),
              Val(use_cache))
end

function (self::HGFGPT2Model)(input, position_ids, token_type_ids, attention_mask,
                              _output_attentions::Val,
                              _output_hidden_states::Val,
                              _use_cache::Val
                              )
  return self(input, position_ids, token_type_ids, nothing, attention_mask,
              _output_attentions,
              _output_hidden_states,
              _use_cache)
end

@generated function (self::HGFGPT2Model{N})(input, position_ids, token_type_ids, past, attention_mask,
                                            _output_attentions::Val{output_attentions},
                                            _output_hidden_states::Val{output_hidden_states},
                                            _use_cache::Val{use_cache}
                                            ) where {N, output_attentions, output_hidden_states, use_cache}

  if output_attentions
    all_attentions = Expr(:tuple)
  end

  if output_hidden_states
    all_hidden_states = Expr(:tuple, :hidden_state)
  end

  if use_cache
    presents = Expr(:tuple)
  end

  body = Expr[]

  for i = 1:N
    previous = i == 1 ? :hidden_state : Symbol(:hidden_state, i-1)
    current = Symbol(:hidden_state, i)

    if output_attentions && use_cache
      current_attention = Symbol(:attention_, i)
      current_present = Symbol(:present_, i)
      current_output = :($current, $current_present, $current_attention)
    elseif output_attentions
      current_attention = Symbol(:attention_, i)
      current_output = :($current, $current_attention)
    elseif use_cache
      current_present = Symbol(:present_, i)
      current_output = :($current, $current_present)
    else
      current_output = current
    end

    if past <: Nothing
      expr = :($current_output = self.h[$i]($previous, attention_mask, _output_attentions, _use_cache))
    else
      expr = :($current_output = self.h[$i]($previous, past[$i], attention_mask, _output_attentions, _use_cache))
    end
    push!(body, expr)

    if output_attentions
      push!(all_attentions.args, current_attention)
    end

    if output_hidden_states
      push!(all_hidden_states.args, current)
    end

    if use_cache
      push!(presents.args, current_present)
    end
  end

  current = Symbol(:hidden_state, N)
  push!(body, :($current = self.ln_f($current)))

  if output_attentions
    push!(body, :(all_attentions = $all_attentions))
  else
    push!(body, :(all_attentions = nothing))
  end

  if output_hidden_states
    push!(body, :(all_hidden_states = $all_hidden_states))
  else
    push!(body, :(all_hidden_states = nothing))
  end

  if use_cache
    push!(body, :(presents = $presents))
  else
    push!(body, :(presents = nothing))
  end

  mask_expr = [:(attention_mask = create_causal_attention_mask(hidden_state, attention_mask))]
  if !(past <: Nothing)
    push!(mask_expr, :(attention_mask = ShiftAttentionMask(attention_mask, get_past_length(past))))
  end

  return quote
    hidden_state = self(input, position_ids, token_type_ids, past)
    $(mask_expr...)
    $(body...)

    return (
      last_hidden_state = $current,
      past_key_values = presents,
      hidden_states = all_hidden_states,
      attentions = all_attentions
    )
  end
end

function HGFGPT2Model(config::HGFGPT2Config)
  wte = FakeTHEmbedding(config, config.vocab_size, config.n_embd)
  wpe = FakeTHEmbedding(config, config.n_positions, config.n_embd)
  h = FakeTHModuleList(
    [HGFGPT2Block(config, config.n_ctx) for _ in 1:config.n_layer]
  )
  ln_f = FakeTHLayerNorm(config, config.n_embd; eps=config.layer_norm_epsilon)
  return HGFGPT2Model(wte, wpe, h, ln_f)
end

#

struct HGFGPT2LMHeadModel{T<:HGFGPT2Model, L<:FakeTHLinear} <: HGFGPT2PreTrainedModel
  transformer::T
  lm_head::L
end

@functor HGFGPT2LMHeadModel

function (self::HGFGPT2LMHeadModel)(input; position_ids=nothing, token_type_ids=nothing,
                                    past_key_values=nothing,
                                    attention_mask=nothing,
                                    output_attentions=false,
                                    output_hidden_states=false,
                                    use_cache=false)
  return self(input, position_ids, token_type_ids,
              past_key_values, attention_mask,
              Val(output_attentions),
              Val(output_hidden_states),
              Val(use_cache))
end

function (self::HGFGPT2LMHeadModel)(input, position_ids, token_type_ids, attention_mask,
                                    _output_attentions::Val,
                                    _output_hidden_states::Val,
                                    _use_cache::Val
                                    )
  return self(input, position_ids, token_type_ids,
              nothing, attention_mask,
              _output_attentions,
              _output_hidden_states,
              _use_cache)
end

function (self::HGFGPT2LMHeadModel)(input, position_ids, token_type_ids, past::Union{Tuple, Nothing}, attention_mask,
                                    _output_attentions::Val,
                                    _output_hidden_states::Val,
                                    _use_cache::Val
                                    )
  transformer_outputs = self.transformer(input, position_ids, token_type_ids, past, attention_mask,
                                         _output_attentions,
                                         _output_hidden_states,
                                         _use_cache)
  hidden_states = transformer_outputs.last_hidden_state
  lm_logits = self.lm_head(hidden_states)
  loss = nothing

  return (
    loss = loss,
    logits = lm_logits,
    past_key_values = transformer_outputs.past_key_values,
    hidden_states = transformer_outputs.hidden_states,
    attentions = transformer_outputs.attentions
  )
end

function (self::HGFGPT2LMHeadModel)(input, labels; position_ids=nothing, token_type_ids=nothing,
                                    past_key_values=nothing,
                                    attention_mask=nothing,
                                    output_attentions=false,
                                    output_hidden_states=false,
                                    use_cache=false)
  return self(input, labels, position_ids, token_type_ids,
              past_key_values, attention_mask,
              Val(output_attentions),
              Val(output_hidden_states),
              Val(use_cache))
end

function (self::HGFGPT2LMHeadModel)(input, labels, position_ids, token_type_ids, attention_mask,
                                    _output_attentions::Val,
                                    _output_hidden_states::Val,
                                    _use_cache::Val
                                    )
  return self(input, labels, position_ids, token_type_ids,
              nothing, attention_mask,
              _output_attentions,
              _output_hidden_states,
              _use_cache)
end

function (self::HGFGPT2LMHeadModel)(input, labels, position_ids, token_type_ids,
                                    past::Union{Tuple, Nothing}, attention_mask,
                                    _output_attentions::Val,
                                    _output_hidden_states::Val,
                                    _use_cache::Val
                                    )
  outputs = self(input, position_ids, token_type_ids, past, attention_mask,
                 _output_attentions,
                 _output_hidden_states,
                 _use_cache)
  shift_logits = outputs.logits[:, 1:end-1, :]
  shift_labels = labels[:, 2:end, :]
  loss = Flux.logitcrossentropy(shift_logits, shift_labels)

  return (
    loss = loss,
    logits = outputs.logits,
    past_key_values = outputs.past_key_values,
    hidden_states = outputs.hidden_states,
    attentions = outputs.attentions
  )
end

function HGFGPT2LMHeadModel(config::HGFGPT2Config)
  transformer = HGFGPT2Model(config)
  lm_head = FakeTHLinear(config, config.n_embd, config.vocab_size; bias=false)
  return HGFGPT2LMHeadModel(transformer, lm_head)
end

#

struct HGFGPT2DoubleHeadsModel{T<:HGFGPT2Model, L<:FakeTHLinear, S<:FakeHGFSequenceSummary} <: HGFGPT2PreTrainedModel
  transformer::T
  lm_head::L
  multiple_choice_head::S
end

@functor HGFGPT2DoubleHeadsModel

function (self::HGFGPT2DoubleHeadsModel)(input; position_ids=nothing, token_type_ids=nothing,
                                          past_key_values=nothing,
                                          mc_token_ids=nothing,
                                          attention_mask=nothing,
                                          output_attentions=false,
                                          output_hidden_states=false,
                                          use_cache=false)
  return self(input, position_ids, token_type_ids,
              past_key_values, mc_token_ids, attention_mask,
              Val(output_attentions),
              Val(output_hidden_states),
              Val(use_cache))
end

function (self::HGFGPT2DoubleHeadsModel)(input, position_ids, token_type_ids, mc_token_ids, attention_mask,
                                    _output_attentions::Val,
                                    _output_hidden_states::Val,
                                    _use_cache::Val
                                    )
  return self(input, position_ids, token_type_ids,
              nothing, mc_token_ids, attention_mask,
              _output_attentions,
              _output_hidden_states,
              _use_cache)
end

function (self::HGFGPT2DoubleHeadsModel)(input, position_ids, token_type_ids, past, mc_token_ids, attention_mask,
                                    _output_attentions::Val,
                                    _output_hidden_states::Val,
                                    _use_cache::Val
                                         )

  num_choices = size(input, ndims(input)-1)
  flat_choice(x) = reshape(x, size(x)[1:end-2]..., :)
  flat_choice(::Nothing) = nothing
  map_flat_choice2(x) = map(x) do (k, v)
    (flat_choice(k), flat_choice(v))
  end
  map_flat_choice2(::Nothing) = nothing

  flat_input = flat_choice(input)
  flat_position_ids = flat_choice(position_ids)
  flat_token_type_ids = flat_choice(token_type_ids)
  flat_attention_mask = flat_choice(attention_mask)
  flat_past = map_flat_choice2(past)

  transformer_outputs = self.transformer(flat_input, flat_position_ids, flat_token_type_ids,
                                         flat_past, flat_attention_mask,
                                         _output_attentions,
                                         _output_hidden_states,
                                         _use_cache)
  hidden_states = transformer_outputs.last_hidden_state |> Base.Fix2(reshape, (num_choices, :))
  lm_logits = self.lm_head(hidden_states)
  mc_logits = self.multiple_choice_head(hidden_states, mc_token_ids)
  lm_loss = nothing
  mc_loss = nothing

  return (
    lm_loss = lm_loss,
    mc_loss = mc_loss,
    lm_logits = lm_logits,
    mc_logits = mc_logits,
    past_key_values = transformer_outputs.past_key_values,
    hidden_states = transformer_outputs.hidden_states,
    attentions = transformer_outputs.attentions
  )
end

function (self::HGFGPT2DoubleHeadsModel)(input, labels, mc_labels; position_ids=nothing, token_type_ids=nothing,
                                          past_key_values=nothing, mc_token_ids=nothing,
                                          attention_mask=nothing,
                                          output_attentions=false,
                                          output_hidden_states=false,
                                          use_cache=false)
  return self(input, labels, mc_labels, position_ids, token_type_ids,
              past_key_values, mc_token_ids, attention_mask,
              Val(output_attentions),
              Val(output_hidden_states),
              Val(use_cache))
end

function (self::HGFGPT2DoubleHeadsModel)(input, labels, mc_labels, position_ids, token_type_ids,
                                          mc_token_ids, attention_mask,
                                          _output_attentions::Val,
                                          _output_hidden_states::Val,
                                          _use_cache::Val
                                          )
  return self(input, labels, mc_labels, position_ids, token_type_ids,
              nothing, mc_token_ids, attention_mask,
              _output_attentions,
              _output_hidden_states,
              _use_cache)
end

function (self::HGFGPT2DoubleHeadsModel)(input, labels, mc_labels, position_ids, token_type_ids,
                                          past, mc_token_ids, attention_mask,
                                          _output_attentions::Val,
                                          _output_hidden_states::Val,
                                          _use_cache::Val
                                          )
  outputs = self(input, position_ids, token_type_ids, past, attention_mask,
                 _output_attentions,
                 _output_hidden_states,
                 _use_cache)
  shift_logits = outputs.lm_logits[:, 1:end-1, :, :]
  shift_labels = labels[:, 2:end, :, :]
  lm_loss = Flux.logitcrossentropy(shift_logits, shift_labels)

  mc_loss = Flux.logitcrossentropy(outputs.mc_logits, mc_labels)

  return (
    lm_loss = lm_loss,
    mc_loss = mc_loss,
    lm_logits = outputs.lm_logits,
    mc_logits = outputs.mc_logits,
    past_key_values = outputs.past_key_values,
    hidden_states = outputs.hidden_states,
    attentions = outputs.attentions
  )
end

function HGFGPT2DoubleHeadsModel(config::HGFGPT2Config)
  transformer = HGFGPT2Model(config)
  lm_head = FakeTHLinear(config, config.n_embd, config.vocab_size; bias=false)
  multiple_choice_head = FakeHGFSequenceSummary(config)
  return HGFGPT2DoubleHeadsModel(transformer, lm_head, multiple_choice_head)
end

# load model utils

basemodelkey(::HGFGPT2PreTrainedModel) = :transformer
basemodel(m::HGFGPT2PreTrainedModel) = getproperty(m, basemodelkey(m))
basemodel(m::HGFGPT2Model) = m

isbasemodel(m::HGFGPT2Model) = true
isbasemodel(m::HGFGPT2PreTrainedModel) = false

get_model_type(::Val{:gpt2}) = (
  :model => HGFGPT2Model,
  :lmheadmodel => HGFGPT2LMHeadModel,
  :doubleheadsmodel => HGFGPT2DoubleHeadsModel,
)


text:
@testset "Offset Types" begin
    @test center((1,2), (1,1)) == (1,1)
    @test center((100,2), (5,1)) == (5,1)

    @test center((1,2), 2) == (2,2)
    @test center((100,2,2,1), 5) == (5,5,5,5)

    @test center((1,2,3,4), CenterFirst) == (1,1,1,1)
    @test center((1,2,3,4), CenterLast) == (1,2,3,4)
    @test center((1,2,3,4), CenterMiddle) == (1,1.5,2,2.5)
    @test center((1,2,3,4), CenterFT) == (1,2,2,3)
end


text:
function build_globalmin{T <: Union(ASCIIString, UTF8String)}(list::Array{T})

  stocks  = stock_returns(list)
  params  = modelparams(stocks)
  weights = rowsums(DataArray(inv(params[5]))) / sum(inv(params[5]))

  port = Portfolio(params[1],
                   params[2],
                   params[3],
                   params[4],
                   params[5],
                   weights)
  port
end

function build_efficient(x)
  #code here
end

function build_tangency(x)
  #code here
end


text:
export WNN, train!, classify

@doc raw"""
    WNN{S, T}(d::Int, n::Int; seed::Union{Int, Nothing}=nothing) where {S <: Any, T <: Union{Unsigned, BigInt}}
    WNN{T}(d::Int, n::Int; seed::Union{Int, Nothing}=nothing) where {T <: Union{Unsigned, BigInt}}

## References:
 * [1] Carvalho, Danilo & Carneiro, Hugo & França, Felipe & Lima, Priscila. (2013). B-bleaching : Agile Overtraining Avoidance in the WiSARD Weightless Neural Classifier.
"""
struct WNN{S <: Any, T <: Union{Unsigned, BigInt}}
    d::Int
    n::Int
    cls::Dict{S, Vector{Dict{T, Int}}}
    map::Vector{Int}

    function WNN{S, T}(d::Int, n::Int; seed::Union{Int, Nothing}=nothing) where {S <: Any, T <: Union{Unsigned, BigInt}}
        if n <= 0 || d <= 0
            error("Values for 'd' and 'n' must be positive")
        end

        if T !== BigInt && d > (T.size * 8)
            error("'$T' is insufficient to guarantee 'd' addressing bits")
        end

        if seed === nothing
            seed = trunc(Int, time())
        end

        rng = Random.MersenneTwister(seed)

        new{S, T}(
            d,
            n,
            Dict{S, Vector{Dict{T, Int}}}(),
            Random.shuffle(rng, 1:n * d)
        )
    end

    function WNN{T}(d::Int, n::Int; seed::Union{Int, Nothing}=nothing) where {T <: Union{Unsigned, BigInt}}
        WNN{Symbol, T}(d, n; seed=seed)
    end

end

Base.show(io::IO, wnn::WNN{S, T}) where {S <: Any, T <: BigInt}  = print(io, "WNN[∞ bits, $(wnn.d) × $(wnn.n)]")
Base.show(io::IO, wnn::WNN{S, T}) where {S <: Any, T <: Unsigned} = print(io, "WNN[$(T.size * 8) bits, $(wnn.d) × $(wnn.n)]")

function address(wnn::WNN, x::AbstractArray, i::Int)
    @inbounds sum(iszero(x[wnn.map[(i - 1) * wnn.d + j]]) ? 0 : 1 << (j - 1) for j = 1:wnn.d)
end

@doc raw"""
    train!(wnn::WNN{S, T}, x::S, y::Vector{Bool}) where {S, T}

Train model with a single pair (class `x`, sample `y`)
"""
function train!(wnn::WNN{S, T}, y::S, x::AbstractArray) where {S, T}
    if !haskey(wnn.cls, y)
        wnn.cls[y] = Dict{T, Int}[Dict{T, Int}() for i = 1:wnn.n]
    end

    cls = wnn.cls[y]

    for i = 1:wnn.n
        k = address(wnn, x, i)
        @inbounds cls[i][k] = get(cls[i], k, 0) + 1
    end

    nothing
end

@doc raw"""
    classify(wnn::WNN, y::Vector{Bool}; bleach=0::Int, gamma=0.5::Float)

Classifies input `y` returning some label `x`. If no training happened, `nothing` will be returned instead.
"""
function classify(wnn::WNN, x::AbstractArray; bleach=0::Int, gamma=0.5::Float64)

    r₁ = r₂ = 0
    y₁ = y₂ = nothing

    for y₀ in keys(wnn.cls)
        r₀ = rate(wnn, y₀, x)
        if r₀ >= r₁
            r₁, r₂ = r₀, r₁
            y₁, y₂ = y₀, y₁
        elseif r₀ > r₂
            r₂ = r₀
        end
    end

    return if bleach == 0
        y₁
    else
        γ = (r₁ - r₂) / r₁

        if γ >= gamma
            y₁
        else
            r₂ = rate(wnn, y₂, x; bleach=bleach)
            r₁ = rate(wnn, y₁, x; bleach=bleach)

            if r₁ > r₂
                y₁
            else
                y₂
            end
        end
    end
end

function rate(wnn::WNN{S, T}, y::Union{S, Nothing}, x::AbstractArray; bleach=0::Int) where {S, T}
    return if !haskey(wnn.cls, y)
        0.0
    else
        @inbounds sum(
            get(wnn.cls[y][i], address(wnn, x, i), 0) > bleach ? 1.0 : 0.0
            for i = 1:wnn.n
        )
    end
end

function train!(wnn::WNN{S, T}, y::Vector{S}, x::Vector{<:AbstractArray}) where {S, T}
    for (yᵢ, xᵢ) in zip(y, x)
        train!(wnn, yᵢ, xᵢ)
    end

    nothing
end

function classify(wnn::WNN, x::Vector{<:AbstractArray}; bleach::Int=0, gamma::Float64=0.5)
    [classify(wnn, xᵢ; bleach=bleach, gamma=gamma) for xᵢ in x]
end


text:
# This file is part of the IntervalArithmetic.jl package; MIT licensed

using IntervalArithmetic
using Test

const eeuler = Base.MathConstants.e

@testset "Constructing intervals" begin
    setprecision(Interval, 53)
    @test IntervalArithmetic.parameters.precision == 53

    setprecision(Interval, Float64)
    @test IntervalArithmetic.parameters.precision == 53

    @test precision(BigFloat) == 53
    @test precision(Interval) == (Float64, 53)

    # Checks for parameters
    @test IntervalArithmetic.parameters.precision_type == Float64
    @test IntervalArithmetic.parameters.precision == 53
    @test IntervalArithmetic.parameters.rounding == :narrow
    @test IntervalArithmetic.parameters.pi == @biginterval(pi)

    # Naive constructors, with no conversion involved
    @test Interval(1) == Interval(1.0, 1.0)
    @test size(Interval(1)) == (1,)
    @test Interval(big(1)) == Interval(1.0, 1.0)
    @test Interval(1//10) == Interval{Rational{Int}}(1//10, 1//10)
    @test Interval(BigInt(1)//10) == Interval{Rational{BigInt}}(1//10, 1//10)
    @test Interval( (1.0, 2.0) ) == Interval(1.0, 2.0)

    @test Interval{Rational{Int}}(1) == Interval(1//1)
    #@test Interval{Rational{Int}}(pi) == Interval(rationalize(1.0*pi))

    @test Interval{BigFloat}(1) == Interval{BigFloat}(big(1.0), big(1.0))

    # Irrational
    for irr in (π, eeuler)
        @test @interval(-irr, irr).hi == (-irr..irr).hi
        @test 0..irr == hull(interval(0), Interval{Float64}(irr))
        @test 1.2..irr == @interval(1.2, irr)
        @test irr..irr == Interval{Float64}(irr)
        @test Interval(irr) == @interval(irr)
        @test Interval(irr, irr) == Interval(irr)
        @test Interval{Float32}(irr, irr) == Interval{Float32}(irr)
    end

    # The following error on windows due to a missing method in MPFR
    @test_broken eeuler..big(4) == hull(Interval{BigFloat}(π), interval(4))
    @test π..big(4) == hull(Interval{BigFloat}(π), interval(4))

    @test eeuler..pi == hull(@interval(eeuler), Interval{Float64}(π))
    @test big(eeuler) in Interval(eeuler, π)
    @test big(π) in Interval(eeuler, π)
    @test big(eeuler) in Interval(0, eeuler)
    @test big(π) in Interval(π, 4)

    @test big(eeuler) in Interval{Float32}(eeuler, π)
    @test big(π) in Interval{Float32}(eeuler, π)
    @test big(eeuler) in Interval{Float32}(0, eeuler)
    @test big(π) in Interval{Float32}(π, 4)

    # a < Inf and b > -Inf
    @test @interval(1e300) == Interval(9.999999999999999e299, 1.0e300)
    @test @interval(-1e307) == Interval(-1.0000000000000001e307, -1.0e307)
    @test @interval(Inf) == IntervalArithmetic.wideinterval(Inf)
    @test IntervalArithmetic.wideinterval(-big(Inf)) == Interval(-Inf, nextfloat(big(-Inf)))

    # Disallowed conversions with a > b

    @test_throws ArgumentError interval(2, 1)
    @test_throws ArgumentError interval(big(2), big(1))
    @test_throws ArgumentError interval(BigInt(1), 1//10)
    @test_throws ArgumentError interval(1, 0.1)
    @test_throws ArgumentError interval(big(1), big(0.1))

    @test_throws ArgumentError @interval(2, 1)
    @test_throws ArgumentError @interval(big(2), big(1))
    @test_throws ArgumentError @interval(big(1), 1//10)
    @test_throws ArgumentError @interval(1, 0.1)
    @test_throws ArgumentError @interval(big(1), big(0.1))
    @test_throws ArgumentError interval(Inf)
    @test_throws ArgumentError interval(-Inf, -Inf)
    @test_throws ArgumentError interval(Inf, Inf)

    # Conversion to Interval without type
    @test convert(Interval, 1) == Interval(1.0)
    @test convert(Interval, pi) == @interval(pi)
    @test convert(Interval, eeuler) == @interval(eeuler)
    @test convert(Interval, BigInt(1)) == Interval(BigInt(1))
    @test convert(Interval, 1//10) == @interval(1//10)
    @test convert(Interval, Interval(0.1, 0.2)) === Interval(0.1, 0.2)

    @test convert(Interval{Rational{Int}}, 0.1) == Interval(1//10)
    # @test convert(Interval{Rational{BigInt}}, pi) == Interval{Rational{BigInt}}(pi)

    ## promotion
    @test promote(Interval(2//1,3//1), Interval(1, 2)) ==
        (Interval(2.0,3.0), Interval(1.0,2.0))
    @test promote(Interval(1.0), pi) == (Interval(1.0), @interval(pi))

    # Constructors from the macros @interval, @floatinterval @biginterval
    setprecision(Interval, 53)

    a = @interval(0.1)
    b = @interval(pi)

    @test nextfloat(a.lo) == a.hi
    @test typeof(a) == Interval{BigFloat}
    @test a == @biginterval("0.1")
    @test convert(Interval{Float64}, a) == @floatinterval(0.1)
    @test nextfloat(b.lo) == b.hi

    @test b == @biginterval(pi)
    x = 10238971209348170283710298347019823749182374098172309487120398471029837409182374098127304987123049817032984712039487
    @test @interval(x) == @biginterval(x)
    @test isthin(@interval(x)) == false

    x = 0.1
    a = @interval(x)
    @test nextfloat(a.lo) == a.hi

    setprecision(Interval, Float64)
    a = @interval(0.1)
    b = @interval(pi)

    @test a == @floatinterval("0.1")
    @test typeof(a) == Interval{Float64}
    @test nextfloat(a.lo) == a.hi
    @test b == @floatinterval(pi)
    @test nextfloat(b.lo) == b.hi
    @test convert(Interval{Float64}, @biginterval(0.1)) == a
    x = typemax(Int64)
    @test @interval(x) == @floatinterval(x)
    @test isthin(@interval(x)) == false
    x = rand()
    c = @interval(x)
    @test nextfloat(c.lo) == c.hi

    a = @interval("[0.1, 0.2]")
    b = @interval(0.1, 0.2)

    @test a == b

    @test_throws ArgumentError @interval("[0.1, 0.2")

    for precision in (64, Float64)
        setprecision(Interval, precision)
        d = big(3)
        f = @interval(d, 2d)
        @test @interval(3, 6) ⊆ f
    end

    for rounding in (:wide, :narrow)
        a = @interval(0.1, 0.2)
        @test a ⊆ Interval(0.09999999999999999, 0.20000000000000004)

        b = @interval(0.1)
        @test b ⊆ Interval(0.09999999999999999, 0.10000000000000002)

        b = setprecision(Interval, 128) do
            @interval(0.1, 0.2)
        end
        @test b ⊆ Interval(0.09999999999999999, 0.20000000000000004)

        @test float(b) ⊆ a

        c = @interval("0.1", "0.2")
        @test c ⊆ a   # c is narrower than a
        @test Interval(1//2) == Interval(0.5)
        @test Interval(1//10).lo == rationalize(0.1)
    end

    @test string(emptyinterval()) == "∅"

    params = IntervalArithmetic.IntervalParameters()
    @test params.precision_type == BigFloat
    @test params.precision == 256
    @test params.rounding == :narrow

    setprecision(Interval, 53)
    a = big(1)//3
    @test @interval(a) == Interval(big(3.3333333333333331e-01), big(3.3333333333333337e-01))
end

@testset "Big intervals" begin
    a = @floatinterval(3)
    @test typeof(a)== Interval{Float64}
    @test typeof(big(a)) == Interval{BigFloat}

    @test @floatinterval(123412341234123412341241234) == Interval(1.234123412341234e26, 1.2341234123412342e26)
    @test @interval(big"3") == @floatinterval(3)

    @test_skip @floatinterval(big"1e10000") == Interval(1.7976931348623157e308, ∞)

    a = big(10)^10000
    @test @floatinterval(a) == Interval(1.7976931348623157e308, ∞)
    setprecision(Interval, 53)
    @test @biginterval(a) == Interval(big"9.9999999999999994e+9999", big"1.0000000000000001e+10000")
end

@testset "Complex intervals" begin
    a = @floatinterval(3 + 4im)
    @test typeof(a)== Complex{Interval{Float64}}
    @test a == Interval(3) + im*Interval(4)

    # TODO; Uncomment these tests
    # b = exp(a)
    # @test real(b) == Interval(-13.12878308146216, -13.128783081462153)
    # @test imag(b) == Interval(-15.200784463067956, -15.20078446306795)
end

@testset "Typed intervals" begin
    setprecision(Interval, Float64)

    @test typeof(@interval Float64 1 2) == Interval{Float64}
    @test typeof(@interval         1 2) == Interval{Float64}
    @test typeof(@interval Float32 1 2) == Interval{Float32}
    @test typeof(@interval Float16 1 2) == Interval{Float16}
end

@testset ".. tests" begin

    a = 0.1..0.3
    @test a == Interval(0.09999999999999999, 0.30000000000000004)
    @test big"0.1" ∈ a
    @test big"0.3" ∈ a

    # part of issue #172:

    a = big(0.1)..2
    @test typeof(a) == Interval{BigFloat}
end

@testset "± tests" begin
    setprecision(Interval, Float64)

    @test 3 ± 1 == Interval(2.0, 4.0)
    @test 3 ± 0.5 == 2.5..3.5
    @test 3 ± 0.1 == 2.9..3.1
    @test 0.5 ± 1 == -0.5..1.5

    # issue 172:
    @test (1..1) ± 1 == 0..2

end

@testset "Conversion to interval of same type" begin
    x = 3..4
    @test convert(Interval{Float64}, x) === x

    x = big(3)..big(4)
    @test convert(Interval{BigFloat}, x) === x
end

@testset "Conversions between different types of interval" begin
    a = convert(Interval{BigFloat}, 3..4)
    @test typeof(a) == Interval{BigFloat}

    a = convert(Interval{Float64}, @biginterval(3, 4))
    @test typeof(a) == Interval{Float64}
end

@testset "Conversion to Interval" begin
    a = convert(Interval, 3)
    @test a == Interval(3.0)
    @test typeof(a) == Interval{Float64}

    a = convert(Interval, big(3))
    @test typeof(a) == Interval{BigFloat}

end

@testset "Interval{T} constructor" begin
    @test Interval{Float64}(1) == 1..1
    @test Interval{Float64}(1.1) == Interval(1.1, 1.1)  # no rounding

    @test Interval{BigFloat}(1) == @biginterval(1, 1)
    @test Interval{BigFloat}(big"1.1") == Interval(big"1.1", big"1.1")
end

# issue 192:
@testset "Disallow a single NaN in an interval" begin
    @test_throws ArgumentError interval(NaN, 2)
    @test_throws ArgumentError interval(Inf, NaN)
end

# issue 206:

setprecision(Interval, Float64)

@testset "Interval strings" begin
    @test I"[1, 2]" == @interval("[1, 2]")
    @test I"[2/3, 1.1]" == @interval("[2/3, 1.1]") == Interval(0.6666666666666666, 1.1)
    @test I"[1]" == @interval("[1]") == Interval(1.0, 1.0)
    @test I"[-0x1.3p-1, 2/3]" == @interval("[-0x1.3p-1, 2/3]") == Interval(-0.59375, 0.6666666666666667)
end

@testset "setdiff tests" begin
    x = 1..3
    y = 2..4
    @test setdiff(x, y) == [1..2]
    @test setdiff(y, x) == [3..4]

    @test setdiff(x, x) == Interval{Float64}[]

    @test setdiff(x, emptyinterval(x)) == [x]

    z = 0..5
    @test setdiff(x, z) == Interval{Float64}[]
    @test setdiff(z, x) == [0..1, 3..5]
end

@testset "Interval{T}(x::Interval)" begin
    @test Interval{Float64}(3..4) == Interval(3.0, 4.0)
    @test Interval{BigFloat}(3..4) == Interval{BigFloat}(3, 4)
end

@testset "@interval with fields" begin
    a = 3..4
    x = @interval(a.lo, 2*a.hi)
    @test x == Interval(3, 8)
end

@testset "@interval with user-defined function" begin
    f(x) = x==Inf ? one(x) : x/(1+x)  # monotonic

    x = 3..4
    @test @interval(f(x.lo), f(x.hi)) == Interval(0.75, 0.8)
end

@testset "a..b with a > b" begin
    @test_throws ArgumentError 3..2
end

@testset "Hashing of Intervals" begin
    x = Interval{Float64}(1, 2)
    y = Interval{BigFloat}(1, 2)
    @test isequal(x, y)
    @test isequal(hash(x), hash(y))

    x = @interval(0.1)
    y = IntervalArithmetic.bigequiv(x)
    @test isequal(x, y)
    @test isequal(hash(x), hash(y))

    x = 0.1
    y = IntervalArithmetic.bigequiv(x)
    @test isequal(x, y)
    @test isequal(hash(x), hash(y))

    x = interval(1, 2)
    y = interval(1, 3)
    @test !isequal(x, y)
    @test !isequal(hash(x), hash(y))

end

import IntervalArithmetic: force_interval
@testset "force_interval" begin
    @test force_interval(4, 3) == Interval(3, 4)
    @test force_interval(4, Inf) == Interval(4, Inf)
    @test force_interval(Inf, 4) == Interval(4, Inf)
    @test force_interval(Inf, -Inf) == Interval(-Inf, Inf)
    @test_throws ArgumentError force_interval(NaN, 3)
end


text:
function summer(A)
   s = zero(eltype(A))
   for a in A
       s += a
   end
   return s
end

#----------------------------
function winter(A)
   s = zero(eltype(A))
   return winter_s1(s,A)
end
function winter_s1(s, A)
   for a in A
       s += exp(a)
   end
   return winter_s1s1(s)
end
function winter_s1s1(s)
   return s + s
end
#----------------------------------------------------------------

function autumn(A)
   s = zero(eltype(A))
   for a in A
       s += autumn_s1(a)
   end
   return autumn_s2(s)
end

function autumn_s2(s)
    for ii in 15
        s+one(s+one(s))
    end
    return autumn_s2s1(s)
end
autumn_s2s1(s) = s-one(s)

function autumn_s1(a)
   s = zero(a)
   i = zero(a)
   while(i<a)
      i += one(a)
      s += autumn_s1s1(a)
   end
   autumn_s1s2(s)
   return s
end

autumn_s1s2(s) = s*one(s)

autumn_s1s1(a) = exp(a-1)

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

function eg1()
    println("1a")
    eg2()
    println("1b")
end
function eg2()
    println("  2a")
    eg3()
    println("  2b")
end
function eg3()
    x=1+1+1
    print("    ")
    println(x)
end


text:
using Pkg
pkg"activate .."
push!(LOAD_PATH,"../src/")
using Documenter, IssueReporter

makedocs(sitename = "IssueReporter Documentation")

max_stars_repo_name: PacktPublishing/Julia-Projects

max_issues_repo_name: PacktPublishing/Julia-Projects

max_forks_repo_name: PacktPublishing/Julia-Projects


text:
using BinaryBuilder

# Collection of pre-build pandoc binaries
name = "pandoc_crossref"

version = v"0.3.9"
crossref_ver = "0.3.9.1"
pandoc_ver = "2.11.4"

url_prefix = "https://github.com/lierdakil/pandoc-crossref/releases/download/v$(crossref_ver)/pandoc-crossref"
sources = [
    ArchiveSource("$(url_prefix)-Linux.tar.xz", "f5efb834779d514e5b8096329efae60426af596ba5f2affc38eac62b9b469641"; unpack_target = "x86_64-linux-gnu"),
    ArchiveSource("$(url_prefix)-macOS.tar.xz", "b72310fcbc25c7f3bcbb4fcc2b512ba591786ddd95a57d8fa007c1b1aff15504"; unpack_target = "x86_64-apple-darwin14"),
    FileSource("$(url_prefix)-Windows.7z", "fe779c649f4f9bee6bb29a1701be9ac7176547ba8e03dcfe8777dcbdfac0d3b7"; filename = "x86_64-w64-mingw32"),
    FileSource("https://raw.githubusercontent.com/lierdakil/pandoc-crossref/v$(crossref_ver)/LICENSE", "39db8f9acf036595a2566ea3fe560bc7bd65d8749f088e0f4a4ef2f8a6cb4b34"),
]

# Bash recipe for building across all platforms
script = raw"""
cd ${WORKSPACE}/srcdir/
mkdir -p "${bindir}"
dirprefix="${target}/"
if [[ "${target}" == *-mingw* ]]; then
    7z x "${target}"
    dirprefix=""
fi
install -m 755 "${dirprefix}pandoc-crossref${exeext}" "${bindir}/pandoc-crossref${exeext}"
install_license LICENSE
"""

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
platforms = [
    Platform("x86_64", "linux"),
    Platform("x86_64", "macos"),
    Platform("x86_64", "windows"),
]

# The products that we will ensure are always built
products = [
    ExecutableProduct("pandoc-crossref", :pandoc_crossref),
]

# Dependencies that must be installed before this package can be built
dependencies = [
    HostBuildDependency("p7zip_jll"),
]

# Build the tarballs, and possibly a `build.jl` as well.
build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies)

max_stars_repo_name: waralex/Yggdrasil

max_issues_repo_name: waralex/Yggdrasil

max_forks_repo_name: waralex/Yggdrasil


text:
# These tests do not require an FT device which supports D2XX to be connected
#
# By Reuben Hill 2019, Gowerlabs Ltd, reuben@gowerlabs.co.uk
#
# Copyright (c) Gowerlabs Ltd.

module TestWrapper

using Test
using LibFTD2XX.Wrapper
using LibFTD2XX.Util

@testset "wrapper" begin

  # FT_CreateDeviceInfoList tests...
  numdevs = FT_CreateDeviceInfoList()
  @test numdevs == 0

  # FT_GetDeviceInfoList tests...
  devinfolist, numdevs2 = FT_GetDeviceInfoList(numdevs)
  @test numdevs2 == numdevs == 0
  @test length(devinfolist) == numdevs == 0

  # FT_GetDeviceInfoDetail tests...
  @test_throws FT_STATUS_ENUM FT_GetDeviceInfoDetail(0)

  # FT_ListDevices tests...
  numdevs2 = Ref{UInt32}()
  retval = FT_ListDevices(numdevs2, Ref{UInt32}(), FT_LIST_NUMBER_ONLY)
  @test retval == nothing
  @test numdevs2[] == numdevs
  devidx = Ref{UInt32}(0)
  buffer = pointer(Vector{Cchar}(undef, 64))
  @test_throws ErrorException FT_ListDevices(devidx, buffer, FT_LIST_BY_INDEX|FT_OPEN_BY_SERIAL_NUMBER)

  # FT_Open tests...
  @test_throws FT_STATUS_ENUM FT_Open(0)

  # FT_OpenEx tests...
  # by description
  @test_throws FT_STATUS_ENUM FT_OpenEx("", FT_OPEN_BY_DESCRIPTION)
  # by serialnumber
  @test_throws FT_STATUS_ENUM FT_OpenEx("", FT_OPEN_BY_SERIAL_NUMBER)

  # FT_Close tests...
  handle = FT_HANDLE() # create with invalid handle...
  @test_throws FT_INVALID_HANDLE FT_Close(handle)

  # FT_Read tests...
  buffer = zeros(UInt8, 5)
  @test_throws FT_INVALID_HANDLE FT_Read(handle, buffer, 0) # read 0 bytes
  @test buffer == zeros(UInt8, 5)
  @test_throws AssertionError FT_Read(handle, buffer, 6) # read 5 bytes
  @test_throws AssertionError FT_Read(handle, buffer, -1) # read -1 bytes

  # FT_Write tests...
  buffer = ones(UInt8, 5)
  @test_throws FT_INVALID_HANDLE FT_Write(handle, buffer, 0) # write 0 bytes
  @test buffer == ones(UInt8, 5)
  @test_throws AssertionError FT_Write(handle, buffer, 6) # write 6 bytes
  @test_throws AssertionError FT_Write(handle, buffer, -1) # write -1 bytes

  # FT_SetDataCharacteristics tests...
  @test_throws FT_INVALID_HANDLE FT_SetDataCharacteristics(handle, FT_BITS_8, FT_STOP_BITS_1, FT_PARITY_NONE)
  # Bad values
  @test_throws AssertionError FT_SetDataCharacteristics(handle, ~(FT_BITS_8 | FT_BITS_7), FT_STOP_BITS_1, FT_PARITY_NONE)
  @test_throws AssertionError FT_SetDataCharacteristics(handle, FT_BITS_8, ~(FT_STOP_BITS_1 | FT_STOP_BITS_2), FT_PARITY_NONE)
  @test_throws AssertionError FT_SetDataCharacteristics(handle, FT_BITS_8, FT_STOP_BITS_1, ~(FT_PARITY_NONE | FT_PARITY_EVEN))
  # closed handle

  # FT_SetTimeouts tests...
  timeout_read, timeout_wr = 200, 100 # milliseconds
  @test_throws FT_INVALID_HANDLE FT_SetTimeouts(handle, timeout_read, timeout_wr)
  @test_throws InexactError FT_SetTimeouts(handle, timeout_read, -1)
  @test_throws InexactError FT_SetTimeouts(handle, -1, timeout_wr)

  # FT_GetModemStatus tests
  @test_throws FT_INVALID_HANDLE FT_GetModemStatus(handle)

  # FT_GetQueueStatus tests
  @test_throws FT_INVALID_HANDLE FT_GetQueueStatus(handle)

  # FT_GetDeviceInfo tests
  @test_throws FT_INVALID_HANDLE FT_GetDeviceInfo(handle)

  # FT_GetDriverVersion tests
  if Sys.iswindows()
    @test_throws FT_INVALID_HANDLE FT_GetDriverVersion(handle)
  else
    @test_throws UndefVarError FT_GetDriverVersion(handle)
  end

  # FT_GetLibraryVersion tests
  if Sys.iswindows()
    version = FT_GetLibraryVersion()
    @test version isa DWORD
    @test version > 0
    @test (version >> 24) & 0xFF == 0x00 # 4th byte should be 0 according to docs
  else
    @test_throws UndefVarError FT_GetLibraryVersion()
  end

  # FT_GetStatus tests
  @test_throws FT_INVALID_HANDLE FT_GetStatus(handle)

  # FT_SetBreakOn tests
  @test_throws FT_INVALID_HANDLE FT_SetBreakOn(handle)

  # FT_SetBreakOff tests
  @test_throws FT_INVALID_HANDLE FT_SetBreakOff(handle)

  # FT_Purge tests
  @test_throws FT_INVALID_HANDLE FT_Purge(handle, FT_PURGE_RX|FT_PURGE_TX)
  @test_throws AssertionError FT_Purge(handle, ~(FT_PURGE_RX))
  @test_throws AssertionError FT_Purge(handle, ~(FT_PURGE_TX))
  @test_throws AssertionError FT_Purge(handle, ~(FT_PURGE_RX | FT_PURGE_TX))

  # FT_StopInTask and FT_RestartInTask tests
  @test_throws FT_INVALID_HANDLE FT_StopInTask(handle)
  @test_throws FT_INVALID_HANDLE FT_RestartInTask(handle)

end

end # module TestWrapper


text:
module Kalman

using JuMP
using Ipopt
using LinearAlgebra
using MKL

include("types.jl")
include("load_data.jl")
include("filter.jl")

export create_filter

end # module


text:
using Graphs

include("common.jl")

example_1 = """
start-A
start-b
A-c
A-b
b-d
A-end
b-end
"""

example_2 = """
dc-end
HN-start
start-kj
dc-start
dc-HN
LN-dc
HN-end
kj-sa
kj-HN
kj-dc
"""

example_3 = """
fs-end
he-DX
fs-he
start-DX
pj-DX
end-zg
zg-sl
zg-pj
pj-he
RW-he
fs-DX
pj-RW
zg-RW
start-pj
he-WI
zg-he
pj-fs
start-RW
"""

struct Caves
    cave_names::Vector{String}
    is_big::Vector{Bool}
    graph::SimpleGraph{Int64}
end

start_index(caves::Caves) = findfirst(==("start"), caves.cave_names)
end_index(caves::Caves) = findfirst(==("end"), caves.cave_names)
Base.@propagate_inbounds is_big(caves::Caves, index::Int64) = caves.is_big[index]

function Base.parse(::Type{Caves}, input::AbstractString)
    graph = SimpleGraph{Int64}()
    cave_names = Vector{String}()
    is_big = Vector{Bool}()
    cave_to_index = Dict{String,Int64}()

    function get_cave_index(cave::AbstractString)
        next_index = nv(graph) + 1
        index = get!(cave_to_index, cave, next_index)
        if index == next_index
            push!(cave_names, cave)
            push!(is_big, isuppercase(first(cave)))
            @assert add_vertex!(graph)
        end
        return index
    end

    for line in split(input, '\n')
        isempty(line) && continue
        cave_1, cave_2 = split(line, '-')

        i1 = get_cave_index(cave_1)
        i2 = get_cave_index(cave_2)

        @assert add_edge!(graph, i1, i2)
    end
    return Caves(cave_names, is_big, graph)
end

"""
    Path

Represent a path through the caves, starting in the correct location.
"""
mutable struct Path
    caves::Caves
    cave_order::Vector{Int64}
    visit_count::Vector{Int64}
    n_additional_small::Int64

    function Path(caves::Caves, n_additional_small::Int64)
        i_start = start_index(caves)
        visit_count = zeros(Int64, nv(caves.graph))
        @inbounds visit_count[i_start] = 1
        return new(caves, [i_start], visit_count, n_additional_small)
    end

    function Path(path::Path)
        return new(
            path.caves,
            copy(path.cave_order),
            copy(path.visit_count),
            path.n_additional_small,
        )
    end
end

function Base.show(io::IO, path::Path)
    repr = join((path.caves.cave_names[i] for i in path.cave_order), '-')
    print(io, repr)
    return nothing
end

start_index(path::Path) = start_index(path.caves)
end_index(path::Path) = end_index(path.caves)

is_finished(path::Path) = @inbounds last(path.cave_order) == end_index(path)

"""
Return true iff the given node is visitable, given the visit counts.
"""
Base.@propagate_inbounds function is_visitable(path::Path, index::Int64)
    return if is_big(path.caves, index)
        true
    elseif index == start_index(path) || index == end_index(path)
        path.visit_count[index] < 1
    else
        path.visit_count[index] < (1 + path.n_additional_small)
    end
end

"""
Get a vector of the next cave indices that are visitable.
Returns an empty vector if the path is finished, or we cannot take further steps.
"""
function next_caves(path::Path)
    is_finished(path) && return Int64[]
    index = @inbounds last(path.cave_order)
    return @inbounds filter(i -> is_visitable(path, i), neighbors(path.caves.graph, index))
end

"""
Add a cave to the path.
Does NOT do any checking to ascertain whether this is a valid move.
"""
Base.@propagate_inbounds function add_cave!(path::Path, index::Int64)
    push!(path.cave_order, index)
    path.visit_count[index] += 1
    @inbounds if !is_big(path.caves, index)
        # Decrement the number of additional small cave visits we can make if we have used.
        if path.visit_count[index] > 1
            path.n_additional_small -= 1
        end
    end
    return path
end

function get_paths(caves::Caves, n_additional_small::Int64)
    finished_paths = Path[]
    paths = Path[Path(caves, n_additional_small)]
    while !isempty(paths)
        path = pop!(paths)

        if is_finished(path)
            # The path is finished, bung it on the finished pile
            push!(finished_paths, path)
            continue
        end

        # Create new paths with every option.
        for index in next_caves(path)
            new_path = Path(path)
            add_cave!(new_path, index)
            push!(paths, new_path)
        end
    end
    return finished_paths
end

get_ans(input::AbstractString) = length(get_paths(parse(Caves, input), 0))

@assert get_ans(example_1) == 10
@assert get_ans(example_2) == 19
@assert get_ans(example_3) == 226
get_ans(load_input(12))

# Part 2

get_ans2(input::AbstractString) = length(get_paths(parse(Caves, input), 1))

@assert get_ans2(example_1) == 36
@assert get_ans2(example_2) == 103
@assert get_ans2(example_3) == 3509
get_ans2(load_input(12))