text:
using DebuggerFramework

function get_insts(session, modules, ip)
    base, mod = NativeDebugger.find_module(session, modules, UInt(ip))
    modrel = UInt(UInt(ip)-base)
    if isa(mod, NativeDebugger.SyntheticModule)
        bounds = mod.get_proc_bounds(session, ip)
        insts = NativeDebugger.load(session, NativeDebugger.RemotePtr{UInt8}(base+first(bounds)),
          length(bounds))
        return base, first(bounds), insts
    end
    if isnull(mod.xpdata)
        loc, fde = NativeDebugger.Unwinder.find_fde(mod, modrel)
        seekloc = loc
        cie = CallFrameInfo.realize_cie(fde)
        nbytes = UInt(CallFrameInfo.fde_range(fde, cie))
    else
        entry = NativeDebugger.Unwinder.find_seh_entry(mod, modrel)
        loc = entry.start
        # Need to translate from virtual to file addresses. Hardcode 0xa00 for
        # now.
        seekloc = loc - 0xa00
        nbytes = entry.stop - entry.start
    end
    if ObjFileBase.isrelocatable(handle(mod))
        # For JIT frames, base is the start of .text, so we need to add that
        # offset back
        text = first(filter(x->sectionname(x)==
            ObjFileBase.mangle_sname(handle(mod),"text"),Sections(handle(mod))))
        seekloc += sectionoffset(text)
    elseif isa(handle(mod), ELF.ELFHandle)
        # In weird executables, the executable segment may not be at offset 0
        phs = ELF.ProgramHeaders(handle(mod))
        idx = findfirst(p->p.p_type==ELF.PT_LOAD &&
                           ((p.p_flags & ELF.PF_X) != 0), phs)
        seekloc += phs[idx].p_offset
    elseif isa(handle(mod), COFF.COFFHandle)
        text = ObjFileBase.deref(first(filter(x->sectionname(x)==
            ObjFileBase.mangle_sname(handle(mod),"text"),Sections(handle(mod)))))
        seekloc -= text.VirtualAddress
        seekloc += text.PointerToRawData
    end
    seek(handle(mod), seekloc)
    insts = read(handle(mod), UInt8, nbytes)
    base, loc, insts
end

function triple_for_sess(sess)
    if isa(NativeDebugger.getarch(sess), NativeDebugger.X86_64.X86_64Arch)
        return "x86_64-linux-gnu"
    else
        return "i386-linux-gnu"
    end
end

demangle(name) = name
function symbolicate_frame(frame::PhysicalStackFrame)
    found = false
    symb = "Unkown"
    try
        found, symb = NativeDebugger.Unwinder.symbolicate(frame.session,
          frame.modules,
          UInt64(NativeDebugger.ip(frame.RC)))
        symb = demangle(symb)
        !found && (symb = "Most likely $symb")
    catch err
        (!isa(err, ErrorException) || !contains(err.msg, "found")) && rethrow(err)
    end
    found, symb
end

function DebuggerFramework.locdesc(frame::PhysicalStackFrame)
    symbolicate_frame(frame)[2]
end

function DebuggerFramework.print_status_synthtic(io::IO, state, frame::PhysicalStackFrame, lines_before, total_lines)
    ipoffset = 0
    ipbase = NativeDebugger.ip(frame)
    try
        # Try disassembling at the start of the function, highlighting the
        # current ip.
        base, loc, insts = get_insts(frame.session, frame.modules, ipbase)
        ipbase = base+loc
        ipoffset = UInt64(ipbase-loc-base-(isstacktop(frame)?0:1))
        disasm_around_ip(io, insts, ipoffset; ipbase=ipbase, triple = triple_for_sess(frame.session))
    catch e
        warn("Failed to obtain instructions from object files ($e). Incorrect unwind info?")
        # This could have failed for a variety of reasons (missing unwind info,
        # self modifying code, etc). If so, get the instructions directly
        # from the target
        insts = NativeDebugger.load(session, NativeDebugger.RemotePtr{UInt8}(x.ip), 40)
        ipoffset = 0
        disasm_around_ip(io, insts, ipoffset; ipbase=ipbase, triple = triple_for_sess(state.top_interp.session))
    end
    return 0
end

#=
function DebuggerFramework.execute_command(state, stack, ::Val{:disas}, command)
    parts = split(command, ' ')
    stacktop = 0
    if length(parts) > 1
        ip = parse(Int, parts[2], 16)
    else
        x = isa(stack, NativeDebugger.NativeStack) ? stack.stack[end] : stack
        stacktop = (x.stacktop?0:1); ip = x.ip
    end
    base, loc, insts = get_insts(state.top_interp.session, state.top_interp.modules, ip)
    disasm_around_ip(STDOUT, insts, UInt64(ip-loc-base-stacktop);
        ipbase=base+loc, circular = false, triple = triple_for_sess(state.top_interp.session))
    return false
end

task_single_step!(timeline) = NativeDebugger.single_step!(timeline)

function DebuggerFramework.execute_command(state, stack::Union{NativeDebugger.NativeStack,NativeDebugger.CStackFrame}, ::Val{:si}, command)
    task_single_step!(state.top_interp.session)
    update_stack_same_frame!(state)
    return true
end

function compute_current_line_range(state, stack)
    mod, base, ip = NativeDebugger.modbaseip_for_stack(state, stack)
    linetab, lip = NativeDebugger.obtain_linetable(state, stack)
    sm = start(linetab)
    local current_entry
    local newentry
    # Start by finding the entry that we're in
    while true
        newentry, sm = next(linetab, sm)
        newentry.address > lip && break
        current_entry = newentry
    end
    range = origrange = current_entry.address:(newentry.address-1)
    # Merge any subsequent entries at the same line
    while newentry.line == current_entry.line
        newentry, sm = next(linetab, sm)
        range = first(origrange):(newentry.address-1)
    end
    range += UInt64(ip-lip)
    range
end

function DebuggerFramework.execute_command(state, stack::Union{NativeDebugger.NativeStack,NativeDebugger.CStackFrame}, ::Val{:n}, command)
    session = state.top_interp.session
    range = compute_current_line_range(state, stack)
    step_over(session, range)
    update_stack_same_frame!(state)
    return true
end

function DebuggerFramework.execute_command(state, stack::Union{NativeDebugger.NativeStack,NativeDebugger.CStackFrame}, ::Val{:ip}, command)
    x = isa(stack, NativeDebugger.NativeStack) ? stack.stack[end] : stack
    println(x.ip)
    return false
end

function DebuggerFramework.execute_command(state, stack::Union{NativeDebugger.NativeStack,NativeDebugger.CStackFrame}, ::Val{:ip}, command)
    session = state.top_interp.session
    modules = state.top_interp.modules
    found, symb = NativeDebugger.Unwinder.symbolicate(session, modules,
      parse(UInt64, split(command,' ')[2][3:end], 16))
    symb = demangle(symb)
    !found && (symb = "Most likely $symb")
    println(symb)
    return false
end

function DebuggerFramework.execute_command(state, stack::Union{NativeDebugger.NativeStack,NativeDebugger.CStackFrame}, ::Val{:nb}, command)
    session = state.top_interp.session
    # First determine the ip of the next branch
    x = isa(stack, NativeDebugger.NativeStack) ? stack.stack[end] : stack
    base, loc, insts = get_insts(session, state.top_interp.modules, x.ip)
    ctx = DisAsmContext()
    Offset = UInt(x.ip - loc - base)
    branchip = 0
    while Offset <= sizeof(insts)
        (Inst, InstSize) = getInstruction(insts, Offset; ctx = ctx)
        if mayAffectControlFlow(Inst,ctx)
            branchip = base + loc + Offset
            break
        end
        Offset += InstSize
        free(Inst)
    end
    @assert branchip != 0
    bp = NativeDebugger.breakpoint(session, branchip)
    task_step_until_bkpt!(session)
    NativeDebugger.disable(bp)
    update_stack_same_frame!(state)
    return true
end

#cxx"""#include <cxxabi.h>"""
function demangle(name)
    return name
    startswith(name,"_Z") || return name
    status = Ref{Cint}()
    bufsize = Ref{Csize_t}(0)
    str = icxx"""
        abi::__cxa_demangle($(pointer(name)),nullptr,
        &$bufsize, &$status);
    """
    @assert status[] == 0
    ret = unsafe_string(str)
    Libc.free(str)
    ret
end

function symbolicate_frame(session, modules, x)
    found = false
    symb = "Unknown Function"
    try
        found, symb = NativeDebugger.Unwinder.symbolicate(session, modules, UInt64(x.ip))
        symb = demangle(symb)
        !found && (symb = "Most likely $symb")
    catch err
        (!isa(err, ErrorException) || !contains(err.msg, "found")) && rethrow(err)
    end
    found, symb
end

function DebuggerFramework.print_frame(state, io, num, x::NativeDebugger.CStackFrame)
    session = state.top_interp.session
    modules = state.top_interp.modules
    print(io, "[$num] ")
    found, symb = symbolicate_frame(session, modules, x)
    print(io, symb, " ")
    if x.line != 0
      print(io, " at ",x.file,":",x.line)
    end
    println(io)
end

function DebuggerFramework.execute_command(state, stack, ::Val{:c}, command)
    try
        NativeDebugger.continue!(state.top_interp.session; only_current_tgid = true)
    catch err
        !isa(err, InterruptException) && rethrow(err)
    end
    update_stack!(state)
    return true
end

function DebuggerFramework.execute_command(state, stack::Union{NativeDebugger.CStackFrame,NativeDebugger.NativeStack}, ::Val{:reg}, command)
    ns = state.top_interp
    @assert isa(ns, NativeDebugger.NativeStack)
    RC = ns.RCs[end-(state.level-1)]
    regname = Symbol(split(command,' ')[2])
    inverse_map = isa(NativeDebugger.getarch(state.top_interp.session),NativeDebugger.X86_64.X86_64Arch) ?
        NativeDebugger.X86_64.inverse_dwarf : NativeDebugger.X86_32.inverse_dwarf
    if !haskey(inverse_map, regname)
        print_with_color(:red, STDOUT, "No such register\n")
        return false
    end
    show(UInt(NativeDebugger.get_dwarf(RC, inverse_map[regname])))
    println(); println()
    return false
end

function DebuggerFramework.execute_command(state, stack::Union{NativeDebugger.CStackFrame,NativeDebugger.NativeStack}, ::Val{:regs}, command)
    ns = state.top_interp
    @assert isa(ns, NativeDebugger.NativeStack)
    RC = ns.RCs[end-(state.level-1)]
    show(STDOUT, RC)
    println(); println()
    return false
end

function DebuggerFramework.execute_command(state, stack::Union{NativeDebugger.CStackFrame,NativeDebugger.NativeStack}, ::Val{:unwind}, command)
    ns = state.top_interp
    newRC = NativeDebugger.Unwinder.unwind_step(ns.session, ns.modules, ns.RCs[end-(state.level-1)];
      allow_frame_based = false, stacktop=ns.stack[end-(state.level-1)].stacktop)[2]
    @show newRC
    return false
end

function DebuggerFramework.execute_command(state, stack::Union{NativeDebugger.CStackFrame,NativeDebugger.NativeStack}, ::Val{:symbolicate}, command)
    session = state.top_interp.session
    modules = state.top_interp.modules
    x = isa(stack, NativeDebugger.NativeStack) ? stack.stack[end] : stack
    approximate, name = NativeDebugger.Unwinder.symbolicate(session, modules, UInt64(x.ip))
    name = demangle(name)
    println(approximate ? "Most likely " : "", name)
    return false
end

function DebuggerFramework.execute_command(state, stack::Union{NativeDebugger.CStackFrame,NativeDebugger.NativeStack}, ::Val{:modules}, command)
    modules = state.top_interp.modules
    modules = sort(collect(modules), by = x->x[1])
    for (base, mod) in modules
        show(UInt64(base)); print('-'); show(UInt64(base)+mod.sz); println()
    end
    return false
end

function DebuggerFramework.execute_command(state, stack::Union{NativeDebugger.CStackFrame,NativeDebugger.NativeStack}, ::Val{:finish}, command)
    ns = state.top_interp
    session = state.top_interp.session
    @assert isa(ns, NativeDebugger.NativeStack)
    parentRC = ns.RCs[end-(state.level)]
    theip = NativeDebugger.ip(parentRC)
    step_to_address(session, theip)
    update_stack!(state)
    return true
end

function DebuggerFramework.execute_command(state, stack::Union{NativeDebugger.CStackFrame,NativeDebugger.NativeStack}, ::Val{:b}, command)
    session = state.top_interp.session
    modules = state.top_interp.modules
    symbol = split(command,' ')[2]
    if startswith(symbol, "0x")
      addr = parse(UInt64, symbol[3:end], 16)
      bp = NativeDebugger.breakpoint(session, addr)
    else
      bp = NativeDebugger.breakpoint(session, modules, symbol)
    end
    show(bp)
    return false
end

function dwarf2Cxx(dbgs, dwarfT)
    if DWARF.tag(dwarfT) == DWARF.DW_TAG_pointer_type ||
            DWARF.tag(dwarfT) == DWARF.DW_TAG_array_type
        dwarfT = get(DWARF.extract_attribute(dwarfT,DWARF.DW_AT_type))
        return Cxx.pointerTo(Cxx.instance(RemoteClang), dwarf2Cxx(dbgs, dwarfT.value))
    else
        name = DWARF.extract_attribute(dwarfT,DWARF.DW_AT_name)
        name = bytestring(get(name).value,StrTab(dbgs.debug_str))
        return cxxparse(Cxx.instance(RemoteClang),name,true)
    end
end

function iterate_frame_variables(state, stack, found_cb, not_found_cb)
    mod, base, theip = NativeDebugger.modbaseip_for_stack(state, stack)
    lip = NativeDebugger.compute_ip(NativeDebugger.dhandle(mod),base,theip)
    dbgs = debugsections(NativeDebugger.dhandle(mod))
    ns = state.top_interp
    @assert isa(ns, NativeDebugger.NativeStack)
    RC = ns.RCs[end-(state.level-1)]

    NativeDebugger.iterate_variables(RC, found_cb, not_found_cb, dbgs, lip)
end

function realize_remote_value(T, val, getreg)
    if isa(val, DWARF.Expressions.MemoryLocation)
        val = NativeDebugger.load(timeline, RemotePtr{T}(val.i))
    elseif isa(val, DWARF.Expressions.RegisterLocation)
        val = reinterpret(T, [getreg(val.i)])[]
    end
    val
end

function DebuggerFramework.execute_command(state, stack::Union{NativeDebugger.CStackFrame,NativeDebugger.NativeStack}, ::Val{:vars}, command)
    function found_cb(dbgs, vardie, getreg, name, val)
        dwarfT = get(DWARF.extract_attribute(vardie,DWARF.DW_AT_type))
        try
            T = Cxx.juliatype(dwarf2Cxx(dbgs, dwarfT.value))
            val = realize_remote_value(T, val, getreg)
        end
        @show (name, val)
    end
    iterate_frame_variables(state, stack, found_cb, (dbgs, vardie, name)->nothing)

    return false
end
=#


text:
module Imp_Util_Test

using Compat
using Compat.Test
using IntervalArithmetic
using EAGOSmoothMcCormickGrad
using StaticArrays

# MC Option functions
x = mc_opts()
y = mc_opts(Float64)
x.lambda = 0.75
y.lambda = 0.75
espv = 0.5
seed1a = 0.5*seed_g(Float64,1,2)

set_default!(x)
@test x.lambda == 0.5
@test y.lambda == 0.75

# Test Cutting Routines
#Final_Cut(x_mc,x_mc_int)
#Final_Cut(x_mc,x_mc_int)
#set_diff_relax(val::Integer)

# Test Rounding Routines
z_mc = [SMCg{2,Float64}(2.0,1.0,seed1a,seed1a,Interval(0.5,3.0),true,
                        [Interval(0.5,3.0),Interval(0.5,3.0)],[1.0, 1.0]) for i=1:2]
z_mc1 = EAGOSmoothMcCormickGrad.Rnd_Out_Z_Intv(z_mc,espv)
z_mc2 = [SMCg{2,Float64}(2.0,1.0,seed1a,seed1a,Interval{Float64}(0.5-Float64(espv),3.0+Float64(espv)),
                        true,[Interval(0.5,3.0),Interval(0.5,3.0)],[1.0, 1.0]) for i=1:2]

@test z_mc1[1].Intv.lo == z_mc2[1].Intv.lo
@test z_mc1[1].Intv.hi == z_mc2[1].Intv.hi
@test z_mc1[2].Intv.lo == z_mc2[2].Intv.lo
@test z_mc1[2].Intv.hi == z_mc2[2].Intv.hi

z_mc = [SMCg{2,Float64}(2.0,1.0,seed1a,seed1a,Interval(0.5,3.0),true,
                        [Interval(0.5,3.0),Interval(0.5,3.0)],[1.0, 1.0]) for i=1:2]
z_mc1 = EAGOSmoothMcCormickGrad.Rnd_Out_Z_All(z_mc,espv)
z_mc2 = [SMCg{2,Float64}(2.0+espv,1.0-espv,seed1a,seed1a,Interval(0.5-espv,3.0+espv),true,
                        [Interval(0.5,3.0),Interval(0.5,3.0)],[1.0, 1.0]) for i=1:2]

@test z_mc1[1].Intv.lo == z_mc2[1].Intv.lo
@test z_mc1[1].Intv.hi == z_mc2[1].Intv.hi
@test z_mc1[2].Intv.lo == z_mc2[2].Intv.lo
@test z_mc1[2].Intv.hi == z_mc2[2].Intv.hi
@test z_mc1[1].cc == z_mc2[1].cc
@test z_mc1[1].cv == z_mc2[1].cv
@test z_mc1[2].cc == z_mc2[2].cc
@test z_mc1[2].cv == z_mc2[2].cv

z_mc = [SMCg{2,Float64}(2.0,1.0,seed1a,seed1a,Interval(0.5,3.0),true,
                        [Interval(0.5,3.0),Interval(0.5,3.0)],[1.0, 1.0]) for i=1:2]
Y_mc = [SMCg{2,Float64}(2.0,1.0,seed1a,seed1a,Interval(0.5,3.0),true,
                        [Interval(0.5,3.0),Interval(0.5,3.0)],[1.0, 1.0]) for i=1:2,j=1:2]
z_mc1, Y_mc1 = EAGOSmoothMcCormickGrad.Rnd_Out_H_Intv(z_mc,Y_mc,espv)
z_mc2 = [SMCg{2,Float64}(2.0,1.0,seed1a,seed1a,Interval(0.5-espv,3.0+espv),true,
                        [Interval(0.5,3.0),Interval(0.5,3.0)],[1.0, 1.0]) for i=1:2]
Y_mc2 = [SMCg{2,Float64}(2.0,1.0,seed1a,seed1a,Interval(0.5-espv,3.0+espv),true,
                        [Interval(0.5,3.0),Interval(0.5,3.0)],[1.0, 1.0]) for i=1:2,j=1:2]

@test z_mc1[1].Intv.lo == z_mc2[1].Intv.lo
@test z_mc1[1].Intv.hi == z_mc2[1].Intv.hi
@test z_mc1[2].Intv.lo == z_mc2[2].Intv.lo
@test z_mc1[2].Intv.hi == z_mc2[2].Intv.hi
@test z_mc1[1].cc == z_mc2[1].cc
@test z_mc1[1].cv == z_mc2[1].cv
@test z_mc1[2].cc == z_mc2[2].cc
@test z_mc1[2].cv == z_mc2[2].cv

@test Y_mc1[1,1].Intv.lo == Y_mc2[1,1].Intv.lo
@test Y_mc1[1,1].Intv.hi == Y_mc2[1,1].Intv.hi
@test Y_mc1[2,1].Intv.lo == Y_mc2[2,1].Intv.lo
@test Y_mc1[2,1].Intv.hi == Y_mc2[2,1].Intv.hi
@test Y_mc1[1,1].cc == Y_mc2[1,1].cc
@test Y_mc1[1,1].cv == Y_mc2[1,1].cv
@test Y_mc1[2,1].cc == Y_mc2[2,1].cc
@test Y_mc1[2,1].cv == Y_mc2[2,1].cv
@test Y_mc1[1,2].Intv.lo == Y_mc2[1,2].Intv.lo
@test Y_mc1[1,2].Intv.hi == Y_mc2[1,2].Intv.hi
@test Y_mc1[2,2].Intv.lo == Y_mc2[2,2].Intv.lo
@test Y_mc1[2,2].Intv.hi == Y_mc2[2,2].Intv.hi
@test Y_mc1[1,2].cc == Y_mc2[1,2].cc
@test Y_mc1[1,2].cv == Y_mc2[1,2].cv
@test Y_mc1[2,2].cc == Y_mc2[2,2].cc
@test Y_mc1[2,2].cv == Y_mc2[2,2].cv

z_mc = [SMCg{2,Float64}(2.0,1.0,seed1a,seed1a,Interval(0.5,3.0),true,
                        [Interval(0.5,3.0),Interval(0.5,3.0)],[1.0, 1.0]) for i=1:2]
Y_mc = [SMCg{2,Float64}(2.0,1.0,seed1a,seed1a,Interval(0.5,3.0),true,
                        [Interval(0.5,3.0),Interval(0.5,3.0)],[1.0, 1.0]) for i=1:2,j=1:2]
z_mc1, Y_mc1 = EAGOSmoothMcCormickGrad.Rnd_Out_H_All(z_mc,Y_mc,espv)
z_mc2 = [SMCg{2,Float64}(2.0+espv,1.0-espv,seed1a,seed1a,Interval(0.5-espv,3.0+espv),true,
                        [Interval(0.5,3.0),Interval(0.5,3.0)],[1.0, 1.0]) for i=1:2]
Y_mc2 = [SMCg{2,Float64}(2.0+espv,1.0-espv,seed1a,seed1a,Interval(0.5-espv,3.0+espv),true,
                        [Interval(0.5,3.0),Interval(0.5,3.0)],[1.0, 1.0]) for i=1:2,j=1:2]

@test z_mc1[1].Intv.lo == z_mc2[1].Intv.lo
@test z_mc1[1].Intv.hi == z_mc2[1].Intv.hi
@test z_mc1[2].Intv.lo == z_mc2[2].Intv.lo
@test z_mc1[2].Intv.hi == z_mc2[2].Intv.hi
@test z_mc1[1].cc == z_mc2[1].cc
@test z_mc1[1].cv == z_mc2[1].cv
@test z_mc1[2].cc == z_mc2[2].cc
@test z_mc1[2].cv == z_mc2[2].cv

@test Y_mc1[1,1].Intv.lo == Y_mc2[1,1].Intv.lo
@test Y_mc1[1,1].Intv.hi == Y_mc2[1,1].Intv.hi
@test Y_mc1[2,1].Intv.lo == Y_mc2[2,1].Intv.lo
@test Y_mc1[2,1].Intv.hi == Y_mc2[2,1].Intv.hi
@test Y_mc1[1,1].cc == Y_mc2[1,1].cc
@test Y_mc1[1,1].cv == Y_mc2[1,1].cv
@test Y_mc1[2,1].cc == Y_mc2[2,1].cc
@test Y_mc1[2,1].cv == Y_mc2[2,1].cv
@test Y_mc1[1,2].Intv.lo == Y_mc2[1,2].Intv.lo
@test Y_mc1[1,2].Intv.hi == Y_mc2[1,2].Intv.hi
@test Y_mc1[2,2].Intv.lo == Y_mc2[2,2].Intv.lo
@test Y_mc1[2,2].Intv.hi == Y_mc2[2,2].Intv.hi
@test Y_mc1[1,2].cc == Y_mc2[1,2].cc
@test Y_mc1[1,2].cv == Y_mc2[1,2].cv
@test Y_mc1[2,2].cc == Y_mc2[2,2].cc
@test Y_mc1[2,2].cv == Y_mc2[2,2].cv

# Test non-sparse preconditioner
h1 = z_mc2
hj1 = Y_mc2
Y1 = [1.0 2.0; 4.0 5.0; 8.0 9.0]
nx = 2
EAGOSmoothMcCormickGrad.Precondition!(h1,hj1,Y1,nx)
@test h1[1].cc == 7.5
@test h1[1].cv == 1.5
@test h1[1].cc_grad == SVector{2,Float64}([1.5, 0.0])
@test h1[1].cv_grad == SVector{2,Float64}([1.5, 0.0])
@test h1[1].Intv.lo == 0.0
@test h1[1].Intv.hi == 10.5
@test h1[2].cc == 42.5
@test h1[2].cv == 8.5
@test h1[2].cc_grad == SVector{2,Float64}([8.5, 0.0])
@test h1[2].cv_grad == SVector{2,Float64}([8.5, 0.0])
@test h1[2].Intv.lo == 0.0
@test h1[2].Intv.hi == 59.5

@test hj1[1,1].cc == 7.5
@test hj1[1,1].cv == 1.5
@test hj1[1,1].cc_grad == SVector{2,Float64}([1.5, 0.0])
@test hj1[1,1].cv_grad == SVector{2,Float64}([1.5, 0.0])
@test hj1[1,1].Intv.lo == 0.0
@test hj1[1,1].Intv.hi == 10.5
@test hj1[2,1].cc == 42.5
@test hj1[2,1].cv == 8.5
@test hj1[2,1].cc_grad == SVector{2,Float64}([8.5, 0.0])
@test hj1[2,1].cv_grad == SVector{2,Float64}([8.5, 0.0])
@test hj1[2,1].Intv.lo == 0.0
@test hj1[2,1].Intv.hi == 59.5

@test hj1[1,2].cc == 7.5
@test hj1[1,2].cv == 1.5
@test hj1[1,2].cc_grad == SVector{2,Float64}([1.5, 0.0])
@test hj1[1,2].cv_grad == SVector{2,Float64}([1.5, 0.0])
@test hj1[1,2].Intv.lo == 0.0
@test hj1[1,2].Intv.hi == 10.5
@test hj1[2,2].cc == 42.5
@test hj1[2,2].cv == 8.5
@test hj1[2,2].cc_grad == SVector{2,Float64}([8.5, 0.0])
@test hj1[2,2].cv_grad == SVector{2,Float64}([8.5, 0.0])
@test hj1[2,2].Intv.lo == 0.0
@test hj1[2,2].Intv.hi == 59.5

end

ext: jl


text:
SubSolver = Union{MPB.AbstractMathProgSolver, Function}
QPSolver = Union{MPB.AbstractMathProgSolver, Function}
get_solver(subsolver::MPB.AbstractMathProgSolver) = subsolver
get_solver(generator::Function)::MPB.AbstractMathProgSolver = generator()


text:
function lazy_lb(
    psi_container::PSIContainer,
    ts_data::Vector{DeviceTimeSeries},
    cons_name::Symbol,
    var_name::Symbol,
)
    time_steps = model_time_steps(psi_container)
    names = (d.name for d in ts_data)
    variable = get_variable(psi_container, var_name)
    lb_name = middle_rename(cons_name, PSI_NAME_DELIMITER, "lb")
    con_lb = add_cons_container!(psi_container, lb_name, names, time_steps)

    for data in ts_data
        if data.range.limits.min > -Inf && !isempty(data.range.additional_terms_lb)
            con_lb[data.name, :].data .= JuMP.AffExpr(0.0)
            continue
        end
        for t in time_steps
            expression_lb = JuMP.AffExpr(0.0, variable[data.name, t] => 1.0)
            for val in data.range.additional_terms_lb
                JuMP.add_to_expression!(
                    expression_lb,
                    get_variable(psi_container, val)[data.name, t],
                    -1.0,
                )
            end
            lb_val = max(0.0, data.range.limits.min)
            con_lb[data.name, t] =
                JuMP.@constraint(psi_container.JuMPmodel, expression_lb >= lb_val)
        end
    end
    return
end

@doc raw"""
    device_timeseries_ub(psi_container::PSIContainer,
                     ts_data::Vector{DeviceTimeSeries},
                     cons_name::Symbol,
                     var_name::Symbol)

Constructs upper bound for given variable and time series data and a multiplier.

# Constraint

```variable[name, t] <= ts_data[name].multiplier * ts_data[name].timeseries[t] ```

# LaTeX

`` x_t \leq r^{val} r_t, \forall t ``

# Arguments
* psi_container::PSIContainer : the psi_container model built in PowerSimulations
* ts_data::Vector{DeviceTimeSeries} : container of device time series data and scaling factors
* cons_name::Symbol : name of the constraint
* var_name::Symbol : the name of the variable
"""
function device_timeseries_ub(
    psi_container::PSIContainer,
    ts_data::Vector{DeviceTimeSeries},
    cons_name::Symbol,
    var_name::Symbol,
)
    time_steps = model_time_steps(psi_container)
    names = (d.name for d in ts_data)
    variable = get_variable(psi_container, var_name)
    ub_name = middle_rename(cons_name, PSI_NAME_DELIMITER, "ub")
    con_ub = add_cons_container!(psi_container, ub_name, names, time_steps)
    lazy_add_lb = false

    for data in ts_data
        for t in time_steps
            expression_ub = JuMP.AffExpr(0.0, variable[data.name, t] => 1.0)
            for val in data.range.additional_terms_ub
                JuMP.add_to_expression!(
                    expression_ub,
                    get_variable(psi_container, val)[data.name, t],
                )
            end
            con_ub[data.name, t] = JuMP.@constraint(
                psi_container.JuMPmodel,
                expression_ub <= data.multiplier * data.timeseries[t]
            )
        end
        if data.range.limits.min > -Inf || !isempty(data.range.additional_terms_lb)
            lazy_add_lb = true
        end
    end
    @debug lazy_add_lb
    lazy_add_lb && lazy_lb(psi_container, ts_data, cons_name, var_name)
    return
end

@doc raw"""
    device_timeseries_lb(psi_container::PSIContainer,
                     ts_data::Vector{DeviceTimeSeries},
                     cons_name::Symbol,
                     var_name::Symbol)

Constructs lower bound for given variable subject to time series data and a multiplier.

# Constraint

``` ts_data[name].multiplier * ts_data[name].timeseries[t] <= variable[name, t] ```

# LaTeX

`` r^{val} r_t \leq x_t, \forall t ``

where (name, data) in range_data.

# Arguments
* psi_container::PSIContainer : the psi_container model built in PowerSimulations
* ts_data::Vector{DeviceTimeSeries} : container of device time series data and scaling factors
* cons_name::Symbol : name of the constraint
* var_name::Symbol : the name of the variable
"""
function device_timeseries_lb(
    psi_container::PSIContainer,
    ts_data::Vector{DeviceTimeSeries},
    cons_name::Symbol,
    var_name::Symbol,
)
    time_steps = model_time_steps(psi_container)
    variable = get_variable(psi_container, var_name)
    lb_name = middle_rename(cons_name, PSI_NAME_DELIMITER, "lb")
    names = (d.name for d in ts_data)
    constraint = add_cons_container!(psi_container, lb_name, names, time_steps)

    for data in ts_data
        for t in time_steps
            expression_lb = JuMP.AffExpr(0.0, variable[data.name, t] => 1.0)
            for val in data.range.additional_terms_lb
                JuMP.add_to_expression!(
                    expression_lb,
                    get_variable(psi_container, val)[data.name, t],
                    -1.0,
                )
            end
            constraint[data.name, t] = JuMP.@constraint(
                psi_container.JuMPmodel,
                expression_lb >= data.multiplier * data.timeseries[t]
            )
        end
    end
    return
end

@doc raw"""
    device_timeseries_param_ub(psi_container::PSIContainer,
                                    ts_data::Vector{DeviceTimeSeries},
                                    cons_name::Symbol,
                                    param_reference::UpdateRef,
                                    var_name::Symbol)

Constructs upper bound for given variable using a parameter. The constraint is
    built with a time series data vector and a multiplier

# Constraint

``` variable[name, t] <= ts_data[name].multiplier * param[name, t] ```

# LaTeX

`` x^{var}_t \leq r^{val} x^{param}_t, \forall t ``

# Arguments
* psi_container::PSIContainer : the psi_container model built in PowerSimulations
* ts_data::Vector{DeviceTimeSeries} : container of device time series data and scaling factors
* cons_name::Symbol : name of the constraint
* param_reference::UpdateRef : UpdateRef to access the parameter
* var_name::Symbol : the name of the variable
"""
function device_timeseries_param_ub(
    psi_container::PSIContainer,
    ts_data::Vector{DeviceTimeSeries},
    cons_name::Symbol,
    param_reference::UpdateRef,
    var_name::Symbol,
)
    time_steps = model_time_steps(psi_container)
    names = (d.name for d in ts_data)
    variable = get_variable(psi_container, var_name)
    ub_name = middle_rename(cons_name, PSI_NAME_DELIMITER, "ub")
    con_ub = add_cons_container!(psi_container, ub_name, names, time_steps)
    container = add_param_container!(psi_container, param_reference, names, time_steps)
    multiplier = get_multiplier_array(container)
    param = get_parameter_array(container)
    lazy_add_lb = false

    for data in ts_data
        for t in time_steps
            expression_ub = JuMP.AffExpr(0.0, variable[data.name, t] => 1.0)
            for val in data.range.additional_terms_ub
                JuMP.add_to_expression!(
                    expression_ub,
                    get_variable(psi_container, val)[data.name, t],
                )
            end
            param[data.name, t] =
                PJ.add_parameter(psi_container.JuMPmodel, data.timeseries[t])
            con_ub[data.name, t] = JuMP.@constraint(
                psi_container.JuMPmodel,
                expression_ub <= data.multiplier * param[data.name, t]
            )
            multiplier[data.name, t] = data.multiplier
        end
        if data.range.limits.min > -Inf || !isempty(data.range.additional_terms_lb)
            lazy_add_lb = true
        end
    end
    @debug lazy_add_lb
    lazy_add_lb && lazy_lb(psi_container, ts_data, cons_name, var_name)
    return
end

@doc raw"""
    device_timeseries_param_lb(psi_container::PSIContainer,
                                    ts_data::Vector{DeviceTimeSeries},
                                    cons_name::Symbol,
                                    param_reference::UpdateRef,
                                    var_name::Symbol)

Constructs lower bound for given variable using a parameter. The constraint is
    built with a time series data vector and a multiplier

# Constraint

``` ts_data[name].multiplier * param[name, t] <= variable[name, t] ```

# LaTeX

`` r^{val} x^{param}_t \leq x^{var}_t, \forall t ``

# Arguments
* psi_container::PSIContainer : the psi_container model built in PowerSimulations
* ts_data::Vector{DeviceTimeSeries} : container of device time series data and scaling factors
* cons_name::Symbol : name of the constraint
* param_reference::UpdateRef : UpdateRef to access the parameter
* var_name::Symbol : the name of the variable
"""
function device_timeseries_param_lb(
    psi_container::PSIContainer,
    ts_data::Vector{DeviceTimeSeries},
    cons_name::Symbol,
    param_reference::UpdateRef,
    var_name::Symbol,
)
    time_steps = model_time_steps(psi_container)
    variable = get_variable(psi_container, var_name)
    lb_name = middle_rename(cons_name, PSI_NAME_DELIMITER, "lb")
    names = (d.name for d in ts_data)
    constraint = add_cons_container!(psi_container, lb_name, names, time_steps)
    container = add_param_container!(psi_container, param_reference, names, time_steps)
    multiplier = get_multiplier_array(container)
    param = get_parameter_array(container)

    for data in ts_data
        for t in time_steps
            expression_lb = JuMP.AffExpr(0.0, variable[data.name, t] => 1.0)
            for val in data.range.additional_terms_lb
                JuMP.add_to_expression!(
                    expression_lb,
                    get_variable(psi_container, val)[data.name, t],
                    -1.0,
                )
            end
            param[data.name, t] =
                PJ.add_parameter(psi_container.JuMPmodel, data.timeseries[t])
            constraint[data.name, t] = JuMP.@constraint(
                psi_container.JuMPmodel,
                expression_lb >= data.multiplier * param[data.name, t]
            )
            multiplier[data.name, t] = data.multiplier
        end
    end

    return

end

@doc raw"""
    device_timeseries_ub_bin(psi_container::PSIContainer,
                                    ts_data::Vector{DeviceTimeSeries},
                                    var_name::Symbol,
                                    binvar_name::Symbol)

Constructs upper bound for variable and time series or confines to 0 depending on binary variable.
    The upper bound is defined by a time series and a multiplier.

# Constraints

``` varcts[name, t] <= varbin[name, t]* ts_data[name].multiplier * ts_data[name].timeseries[t] ```

where (name, data) in range_data.

# LaTeX

`` x^{cts}_t \leq r^{val} r_t x^{bin}_t, \forall t ``

# Arguments
* psi_container::PSIContainer : the psi_container model built in PowerSimulations
* ts_data::Vector{DeviceTimeSeries} : container of device time series data and scaling factors
* cons_name::Symbol : name of the constraint
* var_name::Symbol :  name of the variable
* binvar_name::Symbol : name of binary variable
"""
function device_timeseries_ub_bin(
    psi_container::PSIContainer,
    ts_data::Vector{DeviceTimeSeries},
    cons_name::Symbol,
    var_name::Symbol,
    binvar_name::Symbol,
)
    time_steps = model_time_steps(psi_container)
    ub_name = middle_rename(cons_name, PSI_NAME_DELIMITER, "ub")
    varcts = get_variable(psi_container, var_name)
    varbin = get_variable(psi_container, binvar_name)
    names = (d.name for d in ts_data)
    con_ub = add_cons_container!(psi_container, ub_name, names, time_steps)
    for data in ts_data
        for t in time_steps
            forecast = data.timeseries[t]
            multiplier = data.multiplier
            expression_ub = JuMP.AffExpr(0.0, varcts[data.name, t] => 1.0)
            for val in data.range.additional_terms_ub
                JuMP.add_to_expression!(
                    expression_ub,
                    get_variable(psi_container, val)[data.name, t],
                )
            end
            con_ub[data.name, t] = JuMP.@constraint(
                psi_container.JuMPmodel,
                expression_ub <= varbin[data.name, t] * multiplier * forecast
            )
        end
    end
    return
end

@doc raw"""
    device_timeseries_ub_bigM(psi_container::PSIContainer,
                                    ts_data::Vector{DeviceTimeSeries},
                                    cons_name::Symbol,
                                    var_name::Symbol,
                                    param_reference::UpdateRef,
                                    binvar_name::Symbol,
                                    M_value::Float64 = 1e6)

Constructs upper bound for variable and time series and a multiplier or confines to 0 depending on binary variable.
    Uses BigM constraint type to allow for parameter since ParameterJuMP doesn't support var*parameter

# Constraints

``` varcts[name, t] - ts_data[name].multipliers * param[name, t] <= (1 - varbin[name, t]) * M_value ```

``` varcts[name, t] <= varbin[name, t]*M_value ```

# LaTeX

`` x^{cts}_t - r^{val} x^{param}_t \leq M(1 - x^{bin}_t ), forall t ``

`` x^{cts}_t \leq M x^{bin}_t, \forall t ``

# Arguments
* psi_container::PSIContainer : the psi_container model built in PowerSimulations
* ts_data::Vector{DeviceTimeSeries} : container of device time series data and scaling factors
* cons_name::Symbol : name of the constraint
* var_name::Symbol :  name of the variable
param_reference::UpdateRef : UpdateRef of access the parameters
* binvar_name::Symbol : name of binary variable
* M_value::Float64 : bigM
"""
function device_timeseries_ub_bigM(
    psi_container::PSIContainer,
    ts_data::Vector{DeviceTimeSeries},
    cons_name::Symbol,
    var_name::Symbol,
    param_reference::UpdateRef,
    binvar_name::Symbol,
    M_value::Float64 = 1e6,
)
    time_steps = model_time_steps(psi_container)
    ub_name = middle_rename(cons_name, PSI_NAME_DELIMITER, "ub")
    key_status = middle_rename(cons_name, PSI_NAME_DELIMITER, "status")

    varcts = get_variable(psi_container, var_name)
    varbin = get_variable(psi_container, binvar_name)
    names = (d.name for d in ts_data)
    con_ub = add_cons_container!(psi_container, ub_name, names, time_steps)
    con_status = add_cons_container!(psi_container, key_status, names, time_steps)
    container = add_param_container!(psi_container, param_reference, names, time_steps)
    multiplier = get_multiplier_array(container)
    param = get_parameter_array(container)

    for data in ts_data
        for t in time_steps
            expression_ub = JuMP.AffExpr(0.0, varcts[data.name, t] => 1.0)
            for val in data.range.additional_terms_ub
                JuMP.add_to_expression!(
                    expression_ub,
                    get_variable(psi_container, val)[data.name, t],
                )
            end
            param[data.name, t] =
                PJ.add_parameter(psi_container.JuMPmodel, data.timeseries[t])
            con_ub[data.name, t] = JuMP.@constraint(
                psi_container.JuMPmodel,
                expression_ub - param[data.name, t] * data.multiplier <=
                (1 - varbin[data.name, t]) * M_value
            )
            con_status[data.name, t] = JuMP.@constraint(
                psi_container.JuMPmodel,
                expression_ub <= varbin[data.name, t] * M_value
            )
            multiplier[data.name, t] = data.multiplier
        end
    end
    return
end


text:
@testset "getindex" begin
    @testset "getindex(::Matrix{<:Number}, ...)" begin
        x = [1.0 2.0 3.0; 10.0 20.0 30.0]

        @testset "forward mode" begin
            test_frule(getindex, x, 2)
            test_frule(getindex, x, 2, 1)
            test_frule(getindex, x, CartesianIndex(2, 3))

            test_frule(getindex, x, 2:3)
            test_frule(getindex, x, (:), 2:3)
        end

        @testset "single element" begin
            test_rrule(getindex, x, 2)
            test_rrule(getindex, x, 2, 1)
            test_rrule(getindex, x, 2, 2)

            test_rrule(getindex, x, CartesianIndex(2, 3))
        end

        @testset "slice/index postions" begin
            test_rrule(getindex, x, 2:3)
            test_rrule(getindex, x, 3:-1:2)
            test_rrule(getindex, x, [3,2])
            test_rrule(getindex, x, [2,3])

            test_rrule(getindex, x, 1:2, 2:3)
            test_rrule(getindex, x, (:), 2:3)

            test_rrule(getindex, x, 1:2, 1)
            test_rrule(getindex, x, 1, 1:2)

            test_rrule(getindex, x, 1:2, 2:3)
            test_rrule(getindex, x, (:), 2:3)

            test_rrule(getindex, x, (:), (:))
            test_rrule(getindex, x, (:))
        end

        @testset "masking" begin
            test_rrule(getindex, x, trues(size(x)))
            test_rrule(getindex, x, trues(length(x)))

            mask = falses(size(x))
            mask[2,3] = true
            mask[1,2] = true
            test_rrule(getindex, x, mask)

            test_rrule(getindex, x, [true, false], (:))
        end

        @testset "By position with repeated elements" begin
            test_rrule(getindex, x, [2, 2])
            test_rrule(getindex, x, [2, 2, 2])
            test_rrule(getindex, x, [2,2], [3,3])
        end
    end
end

@testset "view" begin
    test_frule(view, rand(3, 4), :, 1)
    test_frule(view, rand(3, 4), 2, [1, 1, 2])
    test_frule(view, rand(3, 4), 3, 4)
end

@testset "setindex!" begin
    test_frule(setindex!, rand(3, 4), rand(), 1, 2)
    test_frule(setindex!, rand(3, 4), [1,10,100.0], :, 3)
end

@testset "eachslice" begin
    # Testing eachrow not collect∘eachrow leads to errors, e.g.
    # test_rrule: eachrow on Vector{Float64}: Error During Test at /Users/me/.julia/packages/ChainRulesTestUtils/8dFTY/src/testers.jl:195
    #   Got exception outside of a @test
    #   DimensionMismatch("second dimension of A, 6, does not match length of x, 5")
    # Probably similar to https://github.com/JuliaDiff/ChainRulesTestUtils.jl/issues/234 (about Broadcasted not Generator)

    test_rrule(collect∘eachrow, rand(5))
    test_rrule(collect∘eachrow, rand(3, 4))

    test_rrule(collect∘eachcol, rand(3, 4))
    @test_skip test_rrule(collect∘eachcol, Diagonal(rand(5)))  # works locally!

    if VERSION >= v"1.7"
        # On 1.6, ComposedFunction doesn't take keywords. Only affects this testing strategy, not real use.
        test_rrule(collect∘eachslice, rand(3, 4, 5); fkwargs = (; dims = 3))
        test_rrule(collect∘eachslice, rand(3, 4, 5); fkwargs = (; dims = (2,)))
    end

    # Make sure pulling back an array that mixes some AbstractZeros in works right
    _, back = rrule(eachcol, rand(3, 4))
    @test back([1:3, ZeroTangent(), 7:9, NoTangent()]) == (NoTangent(), [1 0 7 0; 2 0 8 0; 3 0 9 0])
    @test back([1:3, ZeroTangent(), 7:9, NoTangent()])[2] isa Matrix{Float64}
    @test back([ZeroTangent(), ZeroTangent(), NoTangent(), NoTangent()]) == (NoTangent(), [0 0 0 0; 0 0 0 0; 0 0 0 0])

    # Second derivative rule
    test_rrule(ChainRules.∇eachslice, [rand(4) for _ in 1:3], rand(3, 4), Val(1))
    test_rrule(ChainRules.∇eachslice, [rand(3) for _ in 1:4], rand(3, 4), Val(2))
    test_rrule(ChainRules.∇eachslice, [rand(2, 3) for _ in 1:4], rand(2, 3, 4), Val(3), check_inferred=false)
end


text:
mutable struct TrackingGradient{G} <: Function
    grad!::G
    counter::Int
end

TrackingGradient(grad!) = TrackingGradient(grad!, 0)

function (tg::TrackingGradient)(storage, x)
    tg.counter += 1
    return tg.grad!(storage, x)
end

mutable struct TrackingObjective{F} <: Function
    f::F
    counter::Int
end

function (tf::TrackingObjective)(x)
    tf.counter += 1
    return tf.f(x)
end

function compute_extreme_point(lmo::TrackingLMO, x; kwargs...)
    lmo.counter += 1
    return compute_extreme_point(lmo.inner, x; kwargs...)
end

function build_tracking_callback(f_values, dual_values, function_calls, gradient_calls, lmo_calls, time_vec)
    function tracking_callback(state)
        push!(function_calls, state.f.counter)
        push!(gradient_calls, state.grad!.counter)
        push!(gradient_calls, state.lmo_calls.counter)
        push!(f_values, state.primal)
        push!(f_values, state.primal)
        push!(dual_values, state.dual_gap)
        push!(lmo_calls, state.lmo.counter)
        push!(time_vec, state.time)
    end
end


text:
"""
    mincost_flow(graph, capacity, demand, cost, solver, [, source][, sink])

Find a flow satisfying the `demand` and `capacity` constraints for each edge
while minimizing the `sum(cost.*flow)`.

- If `source` and `sink` are specified, they are allowed a net flow production,
consumption respectively. All other nodes must respect the flow conservation
property.

- The problem can be seen as a linear programming problem and uses a LP
solver under the hood. We use Clp in the examples and tests.

Returns a flow matrix, flow[i,j] corresponds to the flow on the (i,j) arc.

### Usage Example:

```julia
julia> using Clp: ClpSolver # use your favorite LP solver here
julia> g = lg.DiGraph(6) # Create a flow-graph
julia> lg.add_edge!(g, 5, 1)
julia> lg.add_edge!(g, 5, 2)
julia> lg.add_edge!(g, 3, 6)
julia> lg.add_edge!(g, 4, 6)
julia> lg.add_edge!(g, 1, 3)
julia> lg.add_edge!(g, 1, 4)
julia> lg.add_edge!(g, 2, 3)
julia> lg.add_edge!(g, 2, 4)
julia> w = zeros(6,6)
julia> w[1,3] = 10.
julia> w[1,4] = 5.
julia> w[2,3] = 2.
julia> w[2,4] = 2.
julia> # v2 -> sink have demand of one
julia> demand = spzeros(6,6)
julia> demand[3,6] = 1
julia> demand[4,6] = 1
julia> capacity = ones(6,6)
julia> flow = mincost_flow(g, capacity, demand, w, ClpSolver(), 5, 6)
```
"""
function mincost_flow(g::lg.DiGraph,
		capacity::AbstractMatrix,
		demand::AbstractMatrix,
		cost::AbstractMatrix,
		solver::AbstractMathProgSolver,
		source::Int = -1, # if source and/or sink omitted or not in nodes, circulation problem
		sink::Int = -1)
	flat_cap = collect(Iterators.flatten(capacity))
	flat_dem = collect(Iterators.flatten(demand))
	flat_cost = collect(Iterators.flatten(cost))
	n = lg.nv(g)
	b = zeros(n+n*n)
	A = spzeros(n+n*n,n*n)
	sense = ['=' for _ in 1:n]
	append!(sense, ['>' for _ in 1:n*n])
	for node in 1:n
		if sink == node
			sense[sink] = '>'
		elseif source == node
			sense[source] = '<'
		end
		col_idx = (node-1)*n+1:node*n
		line_idx = node:n:n*n
		for jdx in col_idx
			A[node,jdx] = A[node,jdx]+1.0
		end
		for idx in line_idx
			A[node,idx] = A[node,idx]-1.0
		end
	end
	for src in 1:n
		for dst in 1:n
			if lg.Edge(src,dst) ∉ lg.edges(g)
				A[n+src+n*(dst-1),src+n*(dst-1)] = 1
				sense[n+src+n*(dst-1)] = '<'
			end
		end
	end
	sol = linprog(flat_cost, A, sense, b, flat_dem, flat_cap, solver)
	[sol.sol[idx + n*(jdx-1)] for idx=1:n,jdx=1:n]
end


text:
using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:2] <= 1)

                     @variable(m, -7.9284408379108005 <= q <= 14.377266175441108)

                     add_NL_constraint(m, :(sigmoid(0.4343485145630286 + -0.5735104656090533*sigmoid(-0.47832195115882303 + 0.6358617391786257*$(x[1]) + -0.4037373919488174*$(x[2])) + -0.40833038450807946*sigmoid(0.4113265195541662 + -0.7994001067331897*$(x[1]) + -0.3401124060181533*$(x[2])) + 0.09714914815406805*sigmoid(-0.3249145933093094 + 0.9313960721440679*$(x[1]) + 0.2643917347767659*$(x[2])) + -0.19623354791767778*sigmoid(0.34085895078435247 + -0.33650259644885194*$(x[1]) + -0.9436140496760688*$(x[2])) + -0.3237276355524821*sigmoid(0.8342164127508305 + -0.6721074785776127*$(x[1]) + 0.009028597341863787*$(x[2]))) + sigmoid(-0.11905941659917563 + 0.17224590981590104*sigmoid(-0.47832195115882303 + 0.6358617391786257*$(x[1]) + -0.4037373919488174*$(x[2])) + 0.8406566605626429*sigmoid(0.4113265195541662 + -0.7994001067331897*$(x[1]) + -0.3401124060181533*$(x[2])) + 0.24719777641241691*sigmoid(-0.3249145933093094 + 0.9313960721440679*$(x[1]) + 0.2643917347767659*$(x[2])) + -0.24308200319021767*sigmoid(0.34085895078435247 + -0.33650259644885194*$(x[1]) + -0.9436140496760688*$(x[2])) + -0.39557038452022475*sigmoid(0.8342164127508305 + -0.6721074785776127*$(x[1]) + 0.009028597341863787*$(x[2]))) + sigmoid(0.967203829156619 + 0.3785230261408685*sigmoid(-0.47832195115882303 + 0.6358617391786257*$(x[1]) + -0.4037373919488174*$(x[2])) + -0.23118999327456224*sigmoid(0.4113265195541662 + -0.7994001067331897*$(x[1]) + -0.3401124060181533*$(x[2])) + -0.5197073892915105*sigmoid(-0.3249145933093094 + 0.9313960721440679*$(x[1]) + 0.2643917347767659*$(x[2])) + 0.6005003791404913*sigmoid(0.34085895078435247 + -0.33650259644885194*$(x[1]) + -0.9436140496760688*$(x[2])) + 0.15463985115916623*sigmoid(0.8342164127508305 + -0.6721074785776127*$(x[1]) + 0.009028597341863787*$(x[2]))) + sigmoid(-0.18340968019959902 + 0.01980109053282808*sigmoid(-0.47832195115882303 + 0.6358617391786257*$(x[1]) + -0.4037373919488174*$(x[2])) + 0.8767310690037822*sigmoid(0.4113265195541662 + -0.7994001067331897*$(x[1]) + -0.3401124060181533*$(x[2])) + 0.42243485881005904*sigmoid(-0.3249145933093094 + 0.9313960721440679*$(x[1]) + 0.2643917347767659*$(x[2])) + -0.1694649191201152*sigmoid(0.34085895078435247 + -0.33650259644885194*$(x[1]) + -0.9436140496760688*$(x[2])) + 0.5278959282338329*sigmoid(0.8342164127508305 + -0.6721074785776127*$(x[1]) + 0.009028597341863787*$(x[2]))) + sigmoid(0.7043455967116397 + -0.48930630508951234*sigmoid(-0.47832195115882303 + 0.6358617391786257*$(x[1]) + -0.4037373919488174*$(x[2])) + 0.7375930151357304*sigmoid(0.4113265195541662 + -0.7994001067331897*$(x[1]) + -0.3401124060181533*$(x[2])) + -0.9188673755139103*sigmoid(-0.3249145933093094 + 0.9313960721440679*$(x[1]) + 0.2643917347767659*$(x[2])) + 0.5801272291202126*sigmoid(0.34085895078435247 + -0.33650259644885194*$(x[1]) + -0.9436140496760688*$(x[2])) + 0.06743755305487698*sigmoid(0.8342164127508305 + -0.6721074785776127*$(x[1]) + 0.009028597341863787*$(x[2]))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m

max_stars_repo_name: PSORLab/RSActivationFunctions

max_issues_repo_name: PSORLab/RSActivationFunctions

max_forks_repo_name: PSORLab/RSActivationFunctions


text:
export ForceField, System

# TODO list:
# 1. Implement Dissulfide bond assignment
# 2. Check why graph canonicalization isn't working for nucleic acids with CHARMM

import LightXML
import Chemfiles
using DataFrames
using OrderedCollections

struct ResidueTemplate
    atoms::Vector{Chemfiles.Atom}
    adjacency::BitArray{2}

    function ResidueTemplate(residue, type_masses)
        natoms = length(residue.atoms)
        atom_index = LittleDict(Chemfiles.name.(residue.atoms) .=> 1:natoms)
        adjmat = falses(natoms, natoms)
        for atom_pair in residue.bonds
            i, j = [atom_index[atom] for atom in atom_pair]
            adjmat[i, j] = adjmat[j, i] = true
        end
        atom_masses = [type_masses[Chemfiles.type(atom)] for atom in residue.atoms]
        order, adjmat = canonical_form(adjmat, atom_masses)
        return new(residue.atoms[order], adjmat)
    end
end

struct ForceField
    atom_types::DataFrame
    bond_types::DataFrame
    angle_types::DataFrame
    dihedral_types::DataFrame
    improper_types::DataFrame
    nonbonded::DataFrame
    templates::OrderedDict{String,ResidueTemplate}
    lj₁₋₄::Float64
    coulomb₁₋₄::Float64
end

const ATOM_TYPE = LittleDict(
    :name => String, :class => String, :element => String, :mass => Float64,
)

const HARMONIC_BOND = LittleDict(
    :type1  => String, :type2  => String,
    :class1 => String, :class2 => String,
    :length => Float64,
    :k => Float64,
)

const HARMONIC_ANGLE = LittleDict(
    :type1  => String, :type2  => String, :type3  => String,
    :class1 => String, :class2 => String, :class3 => String,
    :angle => Float64,
    :k => Float64,
)

const PERIODIC_TORSION = LittleDict(
    :type1  => String, :type2  => String, :type3  => String, :type4  => String,
    :class1 => String, :class2 => String, :class3 => String, :class4 => String,
    :periodicity1 => Int, :phase1 => Float64, :k1 => Float64,
    :periodicity2 => Int, :phase2 => Float64, :k2 => Float64,
    :periodicity3 => Int, :phase3 => Float64, :k3 => Float64,
    :periodicity4 => Int, :phase4 => Float64, :k4 => Float64,
    :periodicity5 => Int, :phase5 => Float64, :k5 => Float64,
    :periodicity6 => Int, :phase6 => Float64, :k6 => Float64,
)

const NONBONDED = LittleDict(
    :type => String, :charge => Float64, :sigma => Float64, :epsilon => Float64,
)

mutable struct Residue
    atoms
    bonds
    external_bonds
    Residue() = new([], [], [])
    Residue(atoms, bonds, external_bonds) = new(atoms, bonds, external_bonds)
end

Base.copy(r::Residue) = Residue(copy(r.atoms), copy(r.bonds), copy(r.external_bonds))

sanitized(str) = replace(replace(replace(str, "-" => "_"), "'" => "p"), "*" => "a")

function AddAtom!(residue, attributes)
    atom = Chemfiles.Atom(sanitized(attributes["name"]))
    charge = convert(Float64, get(attributes, "charge", 0))
    Chemfiles.set_charge!(atom, charge)
    Chemfiles.set_type!(atom, attributes["type"])
    push!(residue.atoms, atom)
end

function AddBond!(residue, attributes)
    push!(residue.bonds, Set(sanitized.(values(attributes))))
end

function AddExternalBond!(residue, attributes)
    push!(residue.external_bonds, sanitized(attributes["atomName"]))
end

function ChangeAtom!(residue, attributes)
    name = sanitized(attributes["name"])
    for atom in residue.atoms
        if Chemfiles.name(atom) == name
            charge = convert(Float64, get(attributes, "charge", 0))
            Chemfiles.set_charge!(atom, charge)
            Chemfiles.set_type!(atom, attributes["type"])
            return nothing
        end
    end
end

function RemoveAtom!(residue, attributes)
    atom = sanitized(attributes["name"])
    residue.atoms = filter(x -> Chemfiles.name(x) ≠ atom, residue.atoms)
end

function RemoveBond!(residue, attributes)
    bond = Set(sanitized.(attributes[a] for a in ["atomName1", "atomName2"]))
    residue.bonds = filter(x -> x ≠ bond, residue.bonds)
end

function RemoveExternalBond!(residue, attributes)
    atom = sanitized(attributes["atomName"])
    residue.external_bonds = filter(x -> x ≠ atom, residue.external_bonds)
end

Base.convert(::Type{T}, x::S) where {T<:Number,S<:AbstractString} = parse(T, x)

Base.zero(::Type{String}) = ""

function DataFrame(category, element_list, key)
    df = DataFrame((collect ∘ values)(category), (collect ∘ keys)(category))
    zeros = LittleDict(string(key) => zero(value) for (key, value) in category)
    attribute_dicts(list, key) = [LightXML.attributes_dict(a) for e in list for a in e[key]]
    for dict in attribute_dicts(element_list, key)
        append!(df, merge(zeros, dict))
    end
    # TODO: remove all empty columns
    return df
end

function ForceField(xml_file)
    xroot = LightXML.root(LightXML.parse_file(xml_file))
    patches = LittleDict(
        LightXML.attribute(item, "name") => [
            Symbol(LightXML.name(child)*"!") => LightXML.attributes_dict(child)
            for child in LightXML.child_elements(item)
        ]
        for elem in xroot["Patches"] for item in elem["Patch"]
    )
    atom_types = DataFrame(ATOM_TYPE, xroot["AtomTypes"], "Type")
    type_index = LittleDict(atom_types.name .=> collect(1:nrow(atom_types)))
    type_masses = LittleDict(atom_types.name .=> atom_types.mass)
    templates = OrderedDict{String,ResidueTemplate}()
    for elem in xroot["Residues"]
        for (residue_index, residue_item) in enumerate(elem["Residue"])
            residue = Residue()
            names = []
            for (index, atom_item) in enumerate(residue_item["Atom"])
                attributes = LightXML.attributes_dict(atom_item)
                push!(names, attributes["name"])
                AddAtom!(residue, attributes)
            end
            for bond in residue_item["Bond"]
                atom_names = [
                    key ∈ ["to", "from"] ? names[parse(Int, value)+1] : value
                    for (key, value) in LightXML.attributes_dict(bond)
                ]
                AddBond!(residue, Dict(["atomName1", "atomName2"] .=> atom_names))
            end
            for bond in residue_item["ExternalBond"]
                attributes = LightXML.attributes_dict(bond)
                if haskey(attributes, "from")
                    attributes["atomName"] = names[parse(Int, attributes["from"])+1]
                end
                AddExternalBond!(residue, attributes)
            end
            residue_name = LightXML.attribute(residue_item, "name")
            templates[residue_name] = ResidueTemplate(residue, type_masses)
            for item in residue_item["AllowPatch"]
                patch = LightXML.attribute(item, "name")
                patched_residue = copy(residue)
                for (action, attributes) in patches[patch]
                    getfield(@__MODULE__, action)(patched_residue, attributes)
                end
                templates["$residue_name($patch)"] = ResidueTemplate(patched_residue, type_masses)
            end
        end
    end
    bonds = DataFrame(HARMONIC_BOND, xroot["HarmonicBondForce"], "Bond")
    angles = DataFrame(HARMONIC_ANGLE, xroot["HarmonicAngleForce"], "Angle")
    dihedrals = DataFrame(PERIODIC_TORSION, xroot["PeriodicTorsionForce"], "Proper")
    impropers = DataFrame(PERIODIC_TORSION, xroot["PeriodicTorsionForce"], "Improper")
    nonbonded = DataFrame(NONBONDED, xroot["NonbondedForce"], "Atom")
    scaling_factors = LightXML.attributes_dict(xroot["NonbondedForce"][1])
    lj₁₋₄ = get(scaling_factors, "lj14scale", 1.0)
    coulomb₁₋₄ = get(scaling_factors, "coulomb14scale", 1.0)
    return ForceField(atom_types, bonds, angles, dihedrals, impropers,
                      nonbonded, templates, lj₁₋₄, coulomb₁₋₄)
end

function pdb_aliases(pdb_aliases_file)
    name(element) = LightXML.attribute(element, "name")
    mass(element) = parse(Float64, LightXML.attribute(element, "mass"))
    regex(item) = Regex(LightXML.attribute(item, "code"))
    ids(item) = parse.(Int, values(LightXML.attributes_dict(item)))
    xroot = LightXML.root(LightXML.parse_file(pdb_aliases_file))
    masses = LittleDict(name(e) => mass(e) for e in xroot["Elements"][1]["Element"])
    regex_codes = regex.(xroot["RegularExpressions"][1]["Regex"])
    std_bonds = LittleDict(name(e) => ids.(e["Bond"]) for e in xroot["Residue"])
    return masses, regex_codes, std_bonds
end

const PDB_MASSES, PDB_REGEX_CODES, PDB_STD_BONDS =
    pdb_aliases(joinpath(@__DIR__, "data", "pdb_aliases.xml"))

if isdefined(Chemfiles, :atoms)
    Chemfiles_atoms(residue::Chemfiles.Residue) = Chemfiles.atoms(residue)
else
    function Chemfiles_atoms(residue::Chemfiles.Residue)
        count = size(residue)
        result = Array{UInt64}(undef, count)
        Chemfiles.__check(Chemfiles.lib.chfl_residue_atoms(
            Chemfiles.__const_ptr(residue),
            pointer(result),
            count,
        ))
        return result
    end
end

function System(file, force_field; disambiguation=Dict())
    trajectory = Chemfiles.Trajectory(file)
    frame = read(trajectory)
    unit_cell = Chemfiles.UnitCell(frame)
    # cell_lenghts = Chemfiles.lengths(unit_cell)
    # cell_angles = Chemfiles.angles(unit_cell)
    positions = Chemfiles.positions(frame)
    has_velocities = Chemfiles.has_velocities(frame)
    velocities = has_velocities ? Chemfiles.velocities(frame) : zeros(size(positions))
    topology = Chemfiles.Topology(frame)

    num_atoms = size(topology)
    num_residues = Chemfiles.count_residues(topology)
    atoms = [Chemfiles.Atom(topology, index) for index = 0:num_atoms-1]
    Chemfiles.set_name!.(atoms, (sanitized∘Chemfiles.name).(atoms))
    residues = [Chemfiles.Residue(topology, index) for index = 0:num_residues-1]
    residue_atoms = map((x -> x .+ 1) ∘ Chemfiles_atoms, residues)
    atom_residues = Vector{Int}(undef, num_atoms)
    internal_map = Vector{Int}(undef, num_atoms)
    for (index, atoms_list) in enumerate(residue_atoms)
        atom_residues[atoms_list] .= index
        internal_map[atoms_list] .= collect(1:length(atoms_list))
    end

    is_std_pdb = Chemfiles.property.(residues, "is_standard_pdb")
    for (atom, residue_index) in zip(atoms, atom_residues)
        if is_std_pdb[residue_index]
            element = match(r"[HCNOPS]", Chemfiles.type(atom)).match
            Chemfiles.set_mass!(atom, PDB_MASSES[element])
        end
    end

    bonds = [
        collect(bond)
        for bond in eachcol(Chemfiles.bonds(topology) .+ 1)
        if !all(is_std_pdb[atom_residues[bond]])
    ]
    chain_id = ""
    previous_indices = previous_names = []
    for index = 1:num_residues
        residue = residues[index]
        atom_indices = residue_atoms[index]
        atom_names = Chemfiles.name.(atoms[atom_indices])
        residue_name = Chemfiles.name(residue)
        if is_std_pdb[index]
            new_chain = Chemfiles.property(residue, "chainid") != chain_id
            if new_chain
                chain_id = Chemfiles.property(residue, "chainid")
                previous_indices = previous_names = []
            end
            combined_indices = vcat(previous_indices, atom_indices)
            combined_names = vcat(previous_names, atom_names)
            for (atom_1, atom_2) in PDB_STD_BONDS[residue_name]
                from = broadcast(occursin, PDB_REGEX_CODES[atom_1], combined_names) |> findfirst
                to = broadcast(occursin, PDB_REGEX_CODES[atom_2], combined_names) |> findfirst
                from === nothing || to === nothing || push!(bonds, combined_indices[[from, to]])
            end
            previous_indices = atom_indices
            previous_names = "_" .* atom_names
        end
    end

    adjacency_matrices = map(n -> falses(n, n), length.(residue_atoms))
    for (atom_1, atom_2) in bonds
        index = atom_residues[atom_1]
        if atom_residues[atom_2] == index
            i, j = internal_map[[atom_1, atom_2]]
            adjacency_matrices[index][i, j] = adjacency_matrices[index][j, i] = true
        end
    end

    for resid = 1:num_residues
        residue = residues[resid]
        indices = residue_atoms[resid]
        matrix = adjacency_matrices[resid]
        atom_masses = Chemfiles.mass.(atoms[indices])
        canonical_order, canonical_matrix = canonical_form(matrix, atom_masses)
        matches = [n for (n, t) in force_field.templates if t.adjacency == canonical_matrix]
        name = Chemfiles.name(residue)
        length(matches) == 0 &&
            error("No force field templates matched residue $resid ($(name))")
        if length(matches) > 1
            haskey(disambiguation, resid) ||
                error("Multiple force field templates $(matches) matched residue $resid ($(name))")
            disambiguation[resid] in matches ||
                error("Provided disambiguation for residue $resid ($(name)) is not in $(matches)")
            matches = [disambiguation[resid]]
        end
        template_atoms = force_field.templates[matches[1]].atoms
        for (i, a) in zip(indices[canonical_order], template_atoms)
            Chemfiles.set_property!(atoms[i], "ff.type", Chemfiles.type(a))
            Chemfiles.set_property!(atoms[i], "ff.charge", Chemfiles.charge(a))
        end
    end

    system = Chemfiles.Frame()
    Chemfiles.set_cell!(system, unit_cell)
    has_velocities && Chemfiles.add_velocities!(system)
    location = Vector{Int}(undef, num_atoms)
    atom_index = 0
    for (index, residue) in enumerate(residues)
        new_residue = Chemfiles.Residue(Chemfiles.name(residue), index)
        for i in residue_atoms[index]
            Chemfiles.add_atom!(system, atoms[i], positions[:, i], velocities[:, i])
            Chemfiles.add_atom!(new_residue, atom_index)
            location[i] = atom_index
            atom_index += 1
        end
        Chemfiles.add_residue!(system, new_residue)
    end
    for bond in bonds
        Chemfiles.add_bond!(system, location[bond]...)
    end
    return system
end


text:
using Gadfly

set_default_plot_size(6inch, 3inch)

closure(a) = [x -> a[i]*x for i in 1:length(a)]
plot(closure(1:2), 0, 2π, Scale.y_log10())


text:
export PreconditionerDataK2
"""
Abstract type that defines preconditioners for the K2 formulation.
The available preconditioners are:
- Identity
- Jacobi
- Equilibration (K2 only)
"""
abstract type PreconditionerDataK2{T <: Real, S} end

include("identity.jl")
include("jacobi.jl")
include("equilibration.jl")


text:
open(ARGS[1]) do input
	open(ARGS[2], "w") do output
		for ln in eachline(input)
			ln = chomp(ln);
			if contains(ln, "Number=Plur")
				ln = replace(ln, "|Case=Nom", "");
				ln = replace(ln, "|Case=Gen", "");
				ln = replace(ln, "|Case=Acc", "");
				if contains(ln, "Case=Dat")
					array = split(ln, '|');
					if ismatch(r"s$", array[2])
						array[2] = replace(array[2], r"(.*)s$", s"\1");
						ln = join(array, "|");
						ln = replace(ln, "|Case=Dat", "");
					else
						continue;
					end
				end
			elseif contains(ln, "Number=Sing")
				ln = replace(ln, "|Case=Nom", "");
				ln = replace(ln, "|Case=Dat", "");
				ln = replace(ln, "|Case=Acc", "");
				if contains(ln, "Case=Gen") && (contains(ln, "Gender=Masc") || contains(ln, "Gender=Neut"))
					array = split(ln, '|');
					if ismatch(r"es$", array[2])
						array[2] = replace(array[2], r"(.*)es$", s"\1");
					elseif ismatch(r"s$", array[2])
						array[2] = replace(array[2], r"(.*)s$", s"\1");
					end
					ln = join(array, "|");
				end
				ln = replace(ln, "|Case=Gen", "");
			else
				continue;
			end
			write(output, ln * "\n");
		end
	end
end

max_stars_repo_name: GenaBitu/POPJ-zapocet

max_issues_repo_name: GenaBitu/POPJ-zapocet

max_forks_repo_name: GenaBitu/POPJ-zapocet


text:
function dissipator!(dρ_L::Array{ComplexF64,2},
              ρ::Array{ComplexF64,2},
              γs::Vector{Array{ComplexF64,2}},
              γTs::Vector{Array{ComplexF64,2}},
              γSqs::Vector{Array{ComplexF64,2}},
              A::Array{ComplexF64, 2},
              B::Array{ComplexF64, 2},
              α::Vector{Float64}) where {T <: Number}

    for ii = 1:length(γs)
        mul!(B, γSqs[ii], ρ )            # B = γSq ρ
        mul!(B, ρ, γSqs[ii], 1, 1)       # B = ρ γSq + γSq ρ
        mul!(A, ρ, γTs[ii])              # A = ρ γ'
        mul!(B, γs[ii], A, α[ii], -α[ii])        # B = γ (ρ γ') - (ρ γSq + γSq ρ)
        dρ_L .+= B
    end
    nothing
end

function dissipator!(dρ_L::Array{ComplexF64,2},
              ρ::Array{ComplexF64,2},
              γs::Vector{Array{ComplexF64,2}},
              γTs::Vector{Array{ComplexF64,2}},
              γSqs::Vector{Array{ComplexF64,2}},
              A::Array{ComplexF64, 2},
              B::Array{ComplexF64, 2}) where {T <: Number}

    for ii = 1:length(γs)
        mul!(B, γSqs[ii], ρ )            # B = γSq ρ
        mul!(B, ρ, γSqs[ii], 1, 1)       # B = ρ γSq + γSq ρ
        mul!(A, ρ, γTs[ii])              # A = ρ γ'
        mul!(B, γs[ii], A, 1, -1)        # B = γ (ρ γ') - (ρ γSq + γSq ρ)
        dρ_L .+= B
    end
    nothing
end

# Time independent
# function dRho(rho::Array{ComplexF64,2},H::Array{ComplexF64,2},Gamma::Array{ComplexF64,3},rates::Array{Float64,1},dRho_L::Array{ComplexF64,2})
#
#     fill!(dRho_L,0.)
#     for ii = 1:size(Gamma,3)
#         GammaT = Gamma[:,:,ii]'*Gamma[:,:,ii]
#         dRho_L = dRho_L + rates[ii]*(Gamma[:,:,ii]*rho*Gamma[:,:,ii]' - (GammaT*rho + rho*GammaT)./2)
#     end
#
#     out = -1im*(H*rho-rho*H) + dRho_L
#     out = 1im*(rho*H-H*rho) + dRho_L
# end

"""
In place time independent
"""
function dRho!(dρ::Array{T,2},
              ρ::Array{ComplexF64,2},
              H::Array{ComplexF64,2},
              γs::Vector{Array{ComplexF64,2}},
              γTs::Vector{Array{ComplexF64,2}},
              γSqs::Vector{Array{ComplexF64,2}},
              dρ_L::Array{ComplexF64,2},
              A::Array{ComplexF64, 2},
              B::Array{ComplexF64, 2}) where {T <: Number}

    # Dissipation. Rates and numerical pre-factors have been absorbed into γ terms
    fill!(dρ_L,0.0im)
    dissipator!(dρ_L, ρ, γs, γTs, γSqs, A, B)

    # # Rates and numerical pre-factors have been absorbed into γ terms
    # for ii = 1:length(γs)
    #     mul!(B, γSqs[ii], ρ )            # B = γSq ρ
    #     mul!(B, ρ, γSqs[ii], 1, 1)       # B = ρ γSq + γSq ρ
    #     mul!(A, ρ, γTs[ii])              # A = ρ γ'
    #     mul!(B, γs[ii], A, 1, -1)        # B = γ (ρ γ') - (ρ γSq + γSq ρ)
    #     dρ_L .+= B
    # end

    # Hamiltonian evolution
    mul!(A, H, ρ)                        # A = H ρ
    mul!(A, ρ, H, 1, -1)                 # A = ρ H - H ρ

    dρ .= 1.0im*A .+ dρ_L                # dρ = -i( H ρ - ρ H ) + dρ_L
    nothing
end

# Time dependent Hamiltonian
# function dRho(rho::Array{ComplexF64,2},H::Function,Gamma::Array{ComplexF64,3},rates::Array{Float64,1},dRho_L::Array{ComplexF64,2},t::AbstractFloat)
#
#     fill!(dRho_L,0.)
#     for ii = 1:size(Gamma,3)
#         GammaT = Gamma[:,:,ii]'*Gamma[:,:,ii]
#         dRho_L = dRho_L + rates[ii]*(Gamma[:,:,ii]*rho*Gamma[:,:,ii]' - (GammaT*rho + rho*GammaT)./2)
#     end
#
#     out = -1im*(H(t)*rho-rho*H(t)) + dRho_L
# end

"""
# In place time dependent Hamiltonian with matrix function input
"""
function dRho!(dρ::Array{T,2},
              ρ::Array{ComplexF64,2},
              H::Function,
              H_temp::Array{ComplexF64,2},
              γs::Vector{Array{ComplexF64,2}},
              γTs::Vector{Array{ComplexF64,2}},
              γSqs::Vector{Array{ComplexF64,2}},
              dρ_L::Array{ComplexF64,2},
              A::Array{ComplexF64, 2},
              B::Array{ComplexF64, 2},
              t::AbstractFloat) where {T <: Number}

    # Dissipation. Rates and numerical pre-factors have been absorbed into γ terms
    fill!(dρ_L,0.0im)
    dissipator!(dρ_L, ρ, γs, γTs, γSqs, A, B)

    # Hamiltonian evolution
    fill!(H_temp,0.0im)
    H(H_temp,t)
    mul!(A, H_temp, ρ)                        # A = H ρ
    mul!(A, ρ, H_temp, 1, -1)                 # A = ρ H - H ρ

    dρ .= 1.0im*A .+ dρ_L                # dρ = -i( H ρ - ρ H ) + dρ_L
    nothing
end

"""
# In place time dependent Hamiltonian with matrix basis and vector function inputs
"""

function dRho!(dρ::Array{T,2},
              ρ::Array{ComplexF64,2},
              Hops::Array{ComplexF64,3},
              Hfuncs::Function,
              H_temp::Array{ComplexF64,2},
              Hf_temp::Array{ComplexF64,1},
              γs::Vector{Array{ComplexF64,2}},
              γTs::Vector{Array{ComplexF64,2}},
              γSqs::Vector{Array{ComplexF64,2}},
              dρ_L::Array{ComplexF64,2},
              A::Array{ComplexF64, 2},
              B::Array{ComplexF64, 2},
              t::AbstractFloat) where {T <: Number}

    # Dissipation. Rates and numerical pre-factors have been absorbed into γ terms
    fill!(dρ_L,0.0im)
    dissipator!(dρ_L, ρ, γs, γTs, γSqs, A, B)

    # Hamiltonian evolution
    fill!(H_temp,0.0im)
    Hfuncs(Hf_temp,t)

    for jj = 1:size(Hops,3)
        mul!(H_temp,Hf_temp[jj],Hops[:,:,jj],1,1)
        # H_temp .+= Hf_temp[jj] * Hops[:,:,jj]
    end
    mul!(A, H_temp, ρ)                        # A = H ρ
    mul!(A, ρ, H_temp, 1, -1)                 # A = ρ H - H ρ

    dρ .= 1.0im*A .+ dρ_L                # dρ = -i( H ρ - ρ H ) + dρ_L
    nothing
end

# # Time dependent dissipator
# function dRho(rho::Array{ComplexF64,2},H::Array{ComplexF64,2},Gamma::Array{ComplexF64,3},rates::Function,dRho_L::Array{ComplexF64,2},rates_t::Array{Float64,1},t::AbstractFloat)
#
#     rates_t = rates(t)
#
#     fill!(dRho_L,0.)
#     for ii = 1:size(Gamma,3)
#         GammaT = Gamma[:,:,ii]'*Gamma[:,:,ii]
#         dRho_L = dRho_L + rates_t[ii]*(Gamma[:,:,ii]*rho*Gamma[:,:,ii]' - (GammaT*rho + rho*GammaT)./2)
#     end
#
#     out = -1im*(H*rho-rho*H) + dRho_L
# end

"""
# In place time dependent dissipator
"""
function dRho!(dρ::Array{T,2},
              ρ::Array{ComplexF64,2},
              H::Array{ComplexF64,2},
              rates::Function,
              rates_t::Array{Float64,1},
              γs::Vector{Array{ComplexF64,2}},
              γTs::Vector{Array{ComplexF64,2}},
              γSqs::Vector{Array{ComplexF64,2}},
              dρ_L::Array{ComplexF64,2},
              A::Array{ComplexF64, 2},
              B::Array{ComplexF64, 2},
              t::AbstractFloat) where {T <: Number}

    # Dissipation. Rates and numerical pre-factors have NOT been absorbed into γ terms
    rates_t = rates(t)
    fill!(dρ_L,0.0im)
    dissipator!(dρ_L, ρ, γs, γTs, γSqs, A, B, rates_t)

    # Hamiltonian evolution
    mul!(A, H, ρ)                        # A = H ρ
    mul!(A, ρ, H, 1, -1)                 # A = ρ H - H ρ

    dρ .= 1.0im*A .+ dρ_L                # dρ = -i( H ρ - ρ H ) + dρ_L
    nothing
end

# # Time dependent Hamiltonian and dissipator
# function dRho(rho::Array{ComplexF64,2},H::Function,Gamma::Array{ComplexF64,3},rates::Function,dRho_L::Array{ComplexF64,2},rates_t::Array{Float64,1},t::AbstractFloat)
#
#     rates_t = rates(t)
#
#     fill!(dRho_L,0.)
#     for ii = 1:size(Gamma,3)
#         GammaT = Gamma[:,:,ii]'*Gamma[:,:,ii]
#         dRho_L = dRho_L + rates_t[ii]*(Gamma[:,:,ii]*rho*Gamma[:,:,ii]' - (GammaT*rho + rho*GammaT)./2)
#     end
#
#     out = -1im*(H(t)*rho-rho*H(t)) + dRho_L
# end

"""
# In place time dependent Hamiltonian and dissipator
"""
function dRho!(dρ::Array{T,2},
              ρ::Array{ComplexF64,2},
              H::Function,
              H_temp::Array{ComplexF64,2},
              rates::Function,
              rates_t::Array{Float64,1},
              γs::Vector{Array{ComplexF64,2}},
              γTs::Vector{Array{ComplexF64,2}},
              γSqs::Vector{Array{ComplexF64,2}},
              dρ_L::Array{ComplexF64,2},
              A::Array{ComplexF64, 2},
              B::Array{ComplexF64, 2},
              t::AbstractFloat) where {T <: Number}

    # Dissipation. Rates and numerical pre-factors have NOT been absorbed into γ terms
    rates_t = rates(t)
    fill!(dρ_L,0.0im)
    dissipator!(dρ_L, ρ, γs, γTs, γSqs, A, B, rates_t)

    # Hamiltonian evolution
    fill!(H_temp,0.0im)
    H(H_temp,t)
    mul!(A, H_temp, ρ)                        # A = H ρ
    mul!(A, ρ, H_temp, 1, -1)                 # A = ρ H - H ρ

    dρ .= 1.0im*A .+ dρ_L                # dρ = -i( H ρ - ρ H ) + dρ_L
    nothing
end

# # For vectorized input
# function dRho_vec(rho_v::Array{ComplexF64,1},H::Array{ComplexF64,2},Gamma::Array{ComplexF64,3},rates::Array{Float64,1},dRho_L::Array{ComplexF64,2},rho::Array{ComplexF64,2})
#
#     fill!(rho,0.)
#     rho[:] = rho_v
#
#     fill!(dRho_L,0.)
#     for ii = 1:size(Gamma,3)
#         GammaT = Gamma[:,:,ii]'*Gamma[:,:,ii]
#         dRho_L = dRho_L + rates[ii]*(Gamma[:,:,ii]*rho*Gamma[:,:,ii]' - (GammaT*rho + rho*GammaT)./2)
#     end
#
#     out = -1im*(H*rho-rho*H) + dRho_L
#     out_vec = out[:]
# end

# LUKE: I think we should remove this functionality.
# """
# # In place for vectorized input
# """
# function dRho_vec(out_vec::Array{T,1},
#                   ρ_v::Array{ComplexF64,1},
#                   H::Array{ComplexF64,2},
#                   γ::Array{ComplexF64,3},
#                   rates::Array{Float64,1},
#                   dρ_L::Array{ComplexF64,2},
#                   ρ::Array{ComplexF64,2},
#                   ut::Array{ComplexF64,2}) where {T <: Number}
#
#     fill!(ρ,0.0im)
#     ρ .= ρ_v
#
#     fill!(dρ_L,0.0im)
#     for ii = 1:size(γ,3)
#         #γT = γ[:,:,ii]'*γ[:,:,ii]
#         dρ_L = dρ_L + rates[ii]*(γ[:,:,ii]*ρ*γ[:,:,ii]' - (γT*ρ + ρ*γT)./2)
#     end
#
#     out_vec .= -1.0im .* (H .* ρ .- ρ .* H) .+ dρ_L
#     nothing
# end

"""
# In place time independent covariance matrix
"""
function dCV!(out::Array{T,2},
             C::Array{ComplexF64,2},
             M::Array{Float64,2},
             Z::Array{ComplexF64,2},
             t::AbstractFloat) where {T <: Number}
    out .= M * C + C * transpose(M) .+ Z
    nothing
end

"""
# In place time independent average vector
"""
function d_av!(out::Vector{Float64},
              av_vec::Vector{Float64},
              M::Array{Float64,2},
              drv::Array{Float64,1},
              t::AbstractFloat)

    out .= M * av_vec .+ drv
    nothing
end

"""
# In place time dependent covariance matrix
"""
# function dCV(out::Array{T,2},C::Array{ComplexF64,2},h::Function,t::AbstractFloat,h_tempL::Array{Float64,2}) where {T <: Number}
#     h_tempL[:,:] = h(t)
#     out[:,:] = h_tempL*C + C*transpose(h_tempL[:,:])
#     nothing
# end

# New version
function dCV!(out::Array{T,2},
             C::Array{ComplexF64,2},
             M::Function,
             Z::Array{ComplexF64,2},
             t::AbstractFloat,
             M_temp::Array{Float64,2}) where {T <: Number}

    M(M_temp,t)
    out .= M_temp * C .+ C * transpose(M_temp[:,:]) .+ Z
    nothing
end

# function d_av(out::Vector{Float64},av_vec::Vector{Float64},M::Function,t::AbstractFloat,drv::Function,M_temp::Array{Float64,2},drv_temp::Array{Float64,1})
#     M_temp[:,:] = M(t)
#     drv_temp[:] = drv(t)
#     out[:] = M_temp*av_vec - drv_temp
#     nothing
# end

"""
# In place time dependent average vector
"""
function d_av!(out::Vector{Float64},
              av_vec::Vector{Float64},
              M::Array{Float64,2},
              drv::Function,
              t::AbstractFloat,
              drv_temp::Array{Float64,1})

    drv(drv_temp,t)
    out .= M * av_vec .+ drv_temp
    nothing
end

function d_av!(out::Vector{Float64},
              av_vec::Vector{Float64},
              M::Function,
              drv::Function,
              t::AbstractFloat,
              M_temp::Array{Float64,2},
              drv_temp::Array{Float64,1})

    M(M_temp,t)
    drv(drv_temp,t)
    out .= M_temp * av_vec .+ drv_temp
    nothing
end


text:
macro logprob_str(str)
    expr1, expr2 = get_exprs(str)
    return :(logprob($(esc(expr1)), $(esc(expr2))))
end
macro prob_str(str)
    expr1, expr2 = get_exprs(str)
    return :(exp.(logprob($(esc(expr1)), $(esc(expr2)))))
end

function get_exprs(str::String)
    substrings = split(str, '|'; limit = 2)
    length(substrings) == 2 || error("Invalid expression.")
    str1, str2 = substrings

    expr1 = Meta.parse("($str1,)")
    expr1 = Expr(:tuple, expr1.args...)

    expr2 = Meta.parse("($str2,)")
    expr2 = Expr(:tuple, expr2.args...)

    return expr1, expr2
end

function logprob(ex1, ex2)
    ptype, model, vi = probtype(ex1, ex2)
    if ptype isa Val{:prior}
        return logprior(ex1, ex2, model, vi)
    elseif ptype isa Val{:likelihood}
        return loglikelihood(ex1, ex2, model, vi)
    end
end

function probtype(ntl::NamedTuple{namesl}, ntr::NamedTuple{namesr}) where {namesl, namesr}
    if :chain in namesr
        if isdefined(ntr.chain.info, :model)
            model = ntr.chain.info.model
        elseif isdefined(ntr, :model)
            model = ntr.model
        else
            throw("The model is not defined. Please make sure the model is either saved in the chain or passed on the RHS of |.")
        end
        @assert model isa Model
        if isdefined(ntr.chain.info, :vi)
            _vi = ntr.chain.info.vi
            @assert _vi isa VarInfo
            vi = TypedVarInfo(_vi)
        elseif isdefined(ntr, :varinfo)
            _vi = ntr.varinfo
            @assert _vi isa VarInfo
            vi = TypedVarInfo(_vi)
        else
            vi = nothing
        end
        defaults = model.defaults
        @assert all(getargnames(model)) do arg
            isdefined(ntl, arg) || isdefined(ntr, arg) ||
                isdefined(defaults, arg) && getfield(defaults, arg) !== missing
        end
        return Val(:likelihood), model, vi
    else
        @assert isdefined(ntr, :model)
        model = ntr.model
        @assert model isa Model
        if isdefined(ntr, :varinfo)
            _vi = ntr.varinfo
            @assert _vi isa VarInfo
            vi = TypedVarInfo(_vi)
        else
            vi = nothing
        end
        return probtype(ntl, ntr, model), model, vi
    end
end

function probtype(
    left::NamedTuple{leftnames},
    right::NamedTuple{rightnames},
    model::Model{_F,argnames,defaultnames}
) where {leftnames,rightnames,argnames,defaultnames,_F}
    defaults = model.defaults
    prior_rhs = all(n -> n in (:model, :varinfo) ||
        n in argnames && getfield(right, n) !== missing, rightnames)
    function get_arg(arg)
        if arg in leftnames
            return getfield(left, arg)
        elseif arg in rightnames
            return getfield(right, arg)
        elseif arg in defaultnames
            return getfield(defaults, arg)
        else
            return nothing
        end
    end
    function valid_arg(arg)
        a = get_arg(arg)
        return a !== nothing && a !== missing
    end
    valid_args = all(valid_arg, argnames)

    # Uses the default values for model arguments not provided.
    # If no default value exists, use `nothing`.
    if prior_rhs
        return Val(:prior)
    # Uses the default values for model arguments not provided.
    # If no default value exists or the default value is missing, then error.
    elseif valid_args
        return Val(:likelihood)
    else
        for argname in argnames
            if !valid_arg(argname)
                throw(ArgumentError(missing_arg_error_msg(argname, get_arg(argname))))
            end
        end
    end
end

missing_arg_error_msg(arg, ::Missing) = """Variable $arg has a value of `missing`, or is not defined and its default value is `missing`. Please make sure all the variables are either defined with a value other than `missing` or have a default value other than `missing`."""
missing_arg_error_msg(arg, ::Nothing) = """Variable $arg is not defined and has no default value. Please make sure all the variables are either defined with a value other than `missing` or have a default value other than `missing`."""

function logprior(
    left::NamedTuple,
    right::NamedTuple,
    _model::Model,
    _vi::Union{Nothing, VarInfo}
)
    # For model args on the LHS of |, use their passed value but add the symbol to
    # model.missings. This will lead to an `assume`/`dot_assume` call for those variables.
    # Let `p::PriorContext`. If `p.vars` is `nothing`, `assume` and `dot_assume` will use
    # the values of the random variables in the `VarInfo`. If `p.vars` is a `NamedTuple`
    # or a `Chain`, the value in `p.vars` is input into the `VarInfo` and used instead.

    # For model args not on the LHS of |, if they have a default value, use that,
    # otherwise use `nothing`. This will lead to an `observe`/`dot_observe`call for
    # those variables.
    # All `observe` and `dot_observe` calls are no-op in the PriorContext

    # When all of model args are on the lhs of |, this is also equal to the logjoint.
    model = make_prior_model(left, right, _model)
    vi = _vi === nothing ? VarInfo(deepcopy(model), PriorContext()) : _vi
    foreach(keys(vi.metadata)) do n
        @assert n in keys(left) "Variable $n is not defined."
    end
    model(vi, SampleFromPrior(), PriorContext(left))
    return getlogp(vi)
end

@generated function make_prior_model(
    left::NamedTuple{leftnames},
    right::NamedTuple{rightnames},
    model::Model{_F,argnames,defaultnames}
) where {leftnames,rightnames,argnames,defaultnames,_F}
    argvals = []
    missings = []
    warnings = []

    for argname in argnames
        if argname in leftnames
            push!(argvals, :(deepcopy(left.$argname)))
            push!(missings, argname)
        elseif argname in rightnames
            push!(argvals, :(right.$argname))
        elseif argname in defaultnames
            push!(argvals, :(model.defaults.$argname))
        else
            push!(warnings, :(@warn($(warn_msg(argname)))))
            push!(argvals, :(nothing))
        end
    end

    # `args` is inserted as properly typed NamedTuple expression;
    # `missings` is splatted into a tuple at compile time and inserted as literal
    return quote
        $(warnings...)
        Model{$(Tuple(missings))}(
            model.name, model.f, $(to_namedtuple_expr(argnames, argvals)), model.defaults,
        )
    end
end

warn_msg(arg) = "Argument $arg is not defined. A value of `nothing` is used."

function Distributions.loglikelihood(
    left::NamedTuple,
    right::NamedTuple,
    _model::Model,
    _vi::Union{Nothing, VarInfo},
)
    model = make_likelihood_model(left, right, _model)
    vi = _vi === nothing ? VarInfo(deepcopy(model)) : _vi
    if isdefined(right, :chain)
        # Element-wise likelihood for each value in chain
        chain = right.chain
        ctx = LikelihoodContext(right)
        iters = Iterators.product(1:size(chain, 1), 1:size(chain, 3))
        logps = map(iters) do (sample_idx, chain_idx)
            setval!(vi, chain, sample_idx, chain_idx)
            model(vi, SampleFromPrior(), ctx)
            return getlogp(vi)
        end
        return reshape(logps, size(chain, 1), size(chain, 3))
    else
        # Likelihood without chain
        # Rhs values are used in the context
        ctx = LikelihoodContext(right)
        model(vi, SampleFromPrior(), ctx)
        return getlogp(vi)
    end
end

@generated function make_likelihood_model(
    left::NamedTuple{leftnames},
    right::NamedTuple{rightnames},
    model::Model{_F,argnames,defaultnames},
) where {leftnames,rightnames,argnames,defaultnames,_F}
    argvals = []
    missings = []

    for argname in argnames
        if argname in leftnames
            push!(argvals, :(left.$argname))
        elseif argname in rightnames
            push!(argvals, :(right.$argname))
            push!(missings, argname)
        elseif argname in defaultnames
            push!(argvals, :(model.defaults.$argname))
        else
            throw("This point should not be reached. Please open an issue in the DynamicPPL.jl repository.")
        end
    end

    # `args` is inserted as properly typed NamedTuple expression;
    # `missings` is splatted into a tuple at compile time and inserted as literal
    return :(Model{$(Tuple(missings))}(
        model.name, model.f, $(to_namedtuple_expr(argnames, argvals)), model.defaults,
    ))
end


text:
using Compat.SparseArrays: nonzeroinds, _spdot

function A_mul_Bf(x::SparseVector{Tv,Ti},y::SparseVector{Tv,Ti},f::Function) where {Tv,Ti}
    nnzx = nnz(x)
    nnzy = nnz(y)
    nnzz = nnzx*nnzy # number of nonzeros in new matrix
    I = Vector{Ti}(undef,nnzz) # the indices of nonzeros
    J = Vector{Ti}(undef,nnzz) # the indices of nonzeros
    V  = Vector{Tv}(undef,nnzz) # the values of nonzeros
    @inbounds for i = 1:nnzx, j = 1:nnzy
        this_ind = (i-1)*nnzy+j
        I[this_ind] = x.nzind[i]
        J[this_ind] = y.nzind[j]
        V[this_ind] = x.nzval[i] * f(y.nzval[j])
    end
    return sparse(I,J,V,x.n,y.n)
end
A_mul_Bt(x::SparseVector{Tv,Ti},y::SparseVector{Tv,Ti}) where {Tv,Ti} = A_mul_Bf(x,y,identity)
A_mul_Bc(x::SparseVector{Tv,Ti},y::SparseVector{Tv,Ti}) where {Tv,Ti} = A_mul_Bf(x,y,conj)

function dotu(x::AbstractSparseVector{Tx}, y::AbstractSparseVector{Ty}) where {Tx<:Number,Ty<:Number}
    x === y && return sum(abs2,x)
    n = length(x)
    length(y) == n || throw(DimensionMismatch())
    xnzind = nonzeroinds(x)
    ynzind = nonzeroinds(y)
    xnzval = nonzeros(x)
    ynzval = nonzeros(y)
    return _spdot(*,1,length(xnzind),xnzind,xnzval,1,length(ynzind),ynzind,ynzval)
end

function dotu(x::StridedVector{Tx}, y::AbstractSparseVector{Ty}) where {Tx<:Number,Ty<:Number}
    n = length(x)
    length(y) == n || throw(DimensionMismatch())
    nzind = nonzeroinds(y)
    nzval = nonzeros(y)
    s = zero(Tx) * zero(Ty)
    @inbounds for i = 1:length(nzind)
        s += x[nzind[i]] * nzval[i]
    end
    return s
end

function dotu(x::AbstractSparseVector{Tx}, y::StridedVector{Ty}) where {Tx<:Number,Ty<:Number}
    n = length(y)
    length(x) == n || throw(DimensionMismatch())
    nzind = nonzeroinds(x)
    nzval = nonzeros(x)
    s = zero(Tx) * zero(Ty)
    @inbounds for i = 1:length(nzind)
        s += nzval[i] * y[nzind[i]]
    end
    return s
end