text:
# This file roughly corresponds to functions documented in the
# Users API: https://canvas.instructure.com/doc/api/users

# Note: This method does not map directly to an Canvas API endpoint
"""
    Canvas.whoami(; kwargs...) -> User

Return the details about the authenticated users. Return a [`User`](@ref).

**Canvas API documentation**:
[Show user details (`GET /api/v1/users/self`)]
(https://canvas.instructure.com/doc/api/users#method.users.api_show)
"""
function whoami(; kwargs...)
    json = request("GET", "/api/v1/users/self"; kwargs...)
    return User(json)
end

"""
    Canvas.user(u::User; kwargs...) -> User

Return the details about the a user specified by id.
Return a [`User`](@ref).

**Canvas API documentation**:
[Show user details (`GET /api/v1/users/:id`)]
(https://canvas.instructure.com/doc/api/users#method.users.api_show)
"""
function user(u::User; kwargs...)
    json = request("GET", "/api/v1$(Internals.pid(u))"; kwargs...)
    return User(json)
end


text:
using Base

"""
    smld = catsmld(smld1::SMLD2D, smld2::SMLD2D)

Concatenate `smld1` and `smld2`.

# Description
This method concatenates the datafields of `smld1` and `smld2` (as defined by
smld1.datafields) into an output `smld`.  The non-data fields are taken from
`smld1` and copied into `smld`.
"""
function catsmld(smld1::SMLD2D, smld2::SMLD2D)
    # Loop through each field and add it to the output SMLD.  For scalar
    # fields, we don't have to apply the indexing of `subind`.
    fields = fieldnames(SMLMData.SMLD2D)
    nfields = Base.length(fields)
    smld = deepcopy(smld1)
    for ii = 1:nfields
        currentfield1 = getfield(smld1, fields[ii])
        if any(fields[ii] .== smld.datafields)
            # Concatenate this field from both input SMLDs.
            currentfield2 = getfield(smld2, fields[ii])
            setfield!(smld, fields[ii], [currentfield1; currentfield2])
        end
    end

    return smld
end

ext: jl


text:
@testset "Counters" begin
    @testset "SerialMem" begin
        seq = LongSequence{DNAAlphabet{2}}("ACTGAAGAGC")
        @test serial_mem(DNAMer{4}, NONCANONICAL)(seq) == [
            MerCount{DNAMer{4}}(mer"AAGA",1), MerCount{DNAMer{4}}(mer"ACTG",1),
            MerCount{DNAMer{4}}(mer"AGAG",1), MerCount{DNAMer{4}}(mer"CTGA",1),
            MerCount{DNAMer{4}}(mer"GAAG",1), MerCount{DNAMer{4}}(mer"GAGC",1),
            MerCount{DNAMer{4}}(mer"TGAA",1)
        ]
    end

end


text:
import Flux

include("TensorOps.jl")

struct Attention{T}
    W::Flux.TrackedArray{T, 2}
    B::Flux.TrackedArray{T, 2}
    U::Flux.TrackedArray{T, 2}
end

Attention(dims::Integer) = Attention(
    randn(Float32, dims, dims) |> Flux.gpu |> Flux.param,
    zeros(Float32, 1, dims) |> Flux.gpu |> Flux.param,
    zeros(Float32, 1, dims) |> Flux.gpu |> Flux.param)

function (m::Attention)(x::AbstractArray{T, 3} where T)
    # x is the encoded input, the channels are (T, D, N)
    # this is equivalent to einsum("jmi,mk,->jki", x, W) + B
    logit = TensorDot(x, m.W) .+ m.B

    score = tanh.(logit)
    score = sum(m.U .* score, dims=2)
    α = TensorSoftmax(score, dims=2)
    return dropdims(sum(x .* α, dims=1), dims=1)
end

function (m::Attention)(x::AbstractArray{T, 3}, y::AbstractArray{T, 3}) where T
    # x is the encoded input, y the decoded input, the channels are (T, D, N)
    # this is equivalent to einsum("jmi,mn,kni->jki", x, W, y), with softmax on each batch matrix
    score = BatchMatMul(TensorDot(x, m.W), permutedims(y, [2, 1, 3]))
    α = TensorSoftmax(score, dims=2)
    α = sum(α, dims=1)
    return dropdims(sum(y .* permutedims(α, [2, 1, 3]), dims=1), dims=1)
end
Flux.@treelike Attention


text:
struct ReactionPath
    forward::Bool
    spcsinds::Array{Int64,1}
    rxninds::Array{Int64,1}
    spcind::Int64
    branchfracts::Array{Float64,1}
    branchfract::Array{Float64,1}
    branchind::Array{Int64,1}
end
export ReactionPath

"""
Take a single step following the fluxes in the network from a single species
checks whether the flux is moving the forward/reverse direction
based on production/loss < 1/steptol for forward and
production/loss > steptol for reverse direction and stops accordingly
"""
function stepnetwork(sim,statespcind,ropp,ropl,rts;forward=true,i=0,steptol=1e-2)
    @views ratio = sum(ropp[:,statespcind])/sum(ropl[:,statespcind])
    if forward
        if ratio > 1.0/steptol
            return (0,[0,])
        elseif i == 0
            @views rxnind = argmax(ropl[:,statespcind])
        else
            @views rxnind = reverse(sortperm(ropl[:,statespcind]))[i]
        end
        rt = rts[rxnind]
        if rt > 0.0
            spcinds = sim.reactions[rxnind].productinds
        else
            spcinds = sim.reactions[rxnind].reactantinds
        end
        return rxnind,spcinds
    else
        if ratio < steptol
            return (0,[0,])
        elseif i == 0
            @views rxnind = argmax(ropp[:,statespcind])
        else
            @views rxnind = reverse(sortperm(ropp[:,statespcind]))[i]
        end
        rt = rts[rxnind]
        if rt > 0.0
            spcinds = sim.reactions[rxnind].reactantinds
        else
            spcinds = sim.reactions[rxnind].productinds
        end
        return rxnind,spcinds
    end
end

"""
Starting from a given species indicated by spcind attempt to follow the fluxes
in the forward or reverse direction to find a reaction indicated by rxnind
it will follow in all directions until the flux tracked reaches branchtol fraction
of the original flux, the flux stops (based on steptol) or it finds the reaction
if it finds the reaction it will then follow the flux past that reaction using
a flux pairs algorithm until the flux stops
"""
function follow(sim,rxnind,spcind,ropp,ropl,rts,forward;steptol=1e-2,branchtol=5e-2)

    rps = [ReactionPath(forward,[spcind],Array{Int64,1}(),spcind,Array{Float64,1}(),
                    Array{Float64,1}([1.0]),Array{Int64,1}())]
    boo = true
    boo2 = true
    i = 0
    rpouts = Array{ReactionPath,1}()
    newrps = Array{ReactionPath,1}()
    while boo #extend off the center species
        newrps = extendpath(rps[1],sim,ropp,ropl,rts,i=i,steptol=steptol)
        if length(newrps) == 0
            return ReactionPath(forward,Array{Int64,1}(),Array{Int64,1}(),spcind,Array{Float64,1}(),
                    Array{Float64,1}(),Array{Int64,1}())
        elseif newrps[1].rxninds[1] == rxnind
            rpouts = newrps
            boo2 = false
            break
        end
        if newrps[1].branchfract[1] < branchtol
            boo = false
        else
            append!(rps,newrps)
            i += 1
        end
    end
    deleteat!(rps,1)

    while boo2
        ind = argmax([rp.branchfract[1] for rp in rps])
        rp = rps[ind]
        if rp.branchfract[1] < branchtol
            return ReactionPath(forward,Array{Int64,1}(),Array{Int64,1}(),spcind,Array{Float64,1}(),
                    Array{Float64,1}(),Array{Int64,1})
        end
        boo3 = true
        i = 0
        while boo3 #extend off the center species
            newrps = extendpath(rp,sim,ropp,ropl,rts,i=i,steptol=steptol)
            if length(newrps) == 0
                deleteat!(rps,ind)
                break
            end
            if newrps[1].rxninds[end] == rxnind
                rpouts = newrps
                boo2 = false
                break
            end
            if newrps[1].branchfract[1] < branchtol
                deleteat!(rps,ind)
                boo3 = false
            else
                deleteat!(rps,ind)
                append!(rps,newrps)
                i += 1
            end
        end
        if length(rps) == 0
            return ReactionPath(forward,Array{Int64,1}(),Array{Int64,1}(),spcind,Array{Float64,1}(),
                    Array{Float64,1}(),Array{Int64,1}())
        end
    end

    rpout = newrps[argmax([getsimilarity(sim.species[nrp.spcsinds[end-1]],sim.species[nrp.spcsinds[end]]) for nrp in rpouts])]
    while true
        newrps = extendpath(rpout,sim,ropp,ropl,rts,i=0,steptol=steptol)
        if length(newrps) == 0 || newrps[1].branchfract[1] < branchtol
            return rpout
        else
            rpout = newrps[argmax([getsimilarity(sim.species[rpout.spcsinds[end]],
                            sim.domain.phase.species[nrp.spcsinds[end]]) for nrp in newrps])]
        end
    end
end
export follow

"""
expand a ReactionPath object out a single step
"""
function extendpath(rp,sim,ropp,ropl,rts;i=0,steptol=1e-2)
    rxnind,spcinds = stepnetwork(sim,rp.spcsinds[end],ropp,ropl,rts,forward=rp.forward,i=i,steptol=steptol)
    if rxnind == 0
        sind = rp.spcsinds[end]
        name = sim.names[sind]
        forward = rp.forward
        return Array{ReactionPath,1}()
    else
        rps = Array{ReactionPath,1}()
        for spcind in spcinds
            rpnew = deepcopy(rp)
            push!(rpnew.rxninds,rxnind)
            push!(rpnew.spcsinds,spcind)
            push!(rpnew.branchind,i)
            if rpnew.forward
                push!(rpnew.branchfracts,ropl[rxnind,rp.spcsinds[end]]/sum(ropl[:,rp.spcsinds[end]]))
            else
                push!(rpnew.branchfracts,ropp[rxnind,rp.spcsinds[end]]/sum(ropp[:,rp.spcsinds[end]]))
            end
            rpnew.branchfract[1] = rp.branchfract[1]*rpnew.branchfracts[end]
            push!(rps,rpnew)
        end
        return rps
    end
end

struct Branching
    spcind::Int64
    rxninds::Array{Int64,1}
    branchingratios::Array{Float64,1}
end
export Branching

"""
Identifies calculates and stores the reactions and
branching ratios associated with the "reactants" for the input reaction
defined by rxnind as defined by the direction the reaction is occuring in
"""
function getbranching(sim,rxnind,ropl,rts)
    rt = rts[rxnind]
    if rt > 0
        spcs = sim.reactions[rxnind].reactants
    else
        spcs = sim.reactions[rxnind].products
    end
    spcsinds = [findfirst(isequal(spc.name),sim.names) for spc in spcs]
    boospcs = falses(length(spcsinds))
    branches = Array{Branching,1}()
    for (i,spc) in enumerate(spcs)
        ind = findfirst(isequal(spc.name),sim.names)
        @views rinds = reverse(sortperm(ropl[:,ind]))
        @views rvals = (ropl[:,ind]/sum(ropl[:,ind]))[rinds]
        indend = findfirst(isequal(0.0),rvals)
        @views rinds = rinds[1:indend]
        @views rvals = rvals[1:indend]
        push!(branches,Branching(ind,rinds,rvals))
    end
    return branches
end
export getbranching

struct ReactionAnalysis
    branchings::Array{Branching,1}
    paths::Array{ReactionPath,1}
    radprodlossfract::Float64
    spcind::Int64
    spcname::String
    rxnind::Int64
    sens::Float64
end
export ReactionAnalysis

"""
Calculate Transitory Sensitivities for analysis
"""
function gettransitoryadjoint(sim,t,spcname,spcind,transitorysensitivitymethod)
    if applicable(transitorysensitivitymethod,sim,t,spcname)
        return transitorysensitivitymethod(sim,t,spcname)
    else
        dSdt = transitorysensitivitymethod(sim,t)
        return dSdt[spcind+sim.domain.indexes[1]-1,:]
    end
end

"""
At a given timepoint t targeting a species spcname run transitory sensitivity
analysis to identify important reactions for that species and then run
reaction path, branching, radical rop and other analyses to provide the information
necessary to identify why the reaction is important and flesh out important
pathways in the mechanism
Returns a list of ReactionAnalysis objects associated with each reaction
"""
function analyzespc(sim,spcname,t;N=10,tol=1e-3,branchthreshold=0.9,
        pathbranchthreshold=0.2,branchtol=1e-2,steptol=1e-2,
        transitorysensitivitymethod=transitorysensitivitiesfulltrapezoidal,
        eliminate=true
        )

    spcind = findfirst(isequal(spcname),sim.names)
    @assert spcind != nothing "$spcname not found"
    dSdt = gettransitoryadjoint(sim,t,spcname,spcind,transitorysensitivitymethod)

    rop = rops(sim,t)
    ropp = zeros(size(rop))
    for i in eachindex(rop)
        if rop[i] > 0
            ropp[i] = rop[i]
        end
    end
    ropl = zeros(size(rop))
    for i in eachindex(rop)
        if rop[i] < 0
            ropl[i] = abs(rop[i])
        end
    end
    rts = rates(sim,t)

    @views dSdtspc = dSdt[length(sim.domain.phase.species)+sim.domain.parameterindexes[1]:sim.domain.parameterindexes[2]]

    #find sensitive reactions
    inds = reverse(sortperm(abs.(dSdtspc)))
    dSdtmax = maximum(abs.(dSdtspc))
    maxthresh = dSdtmax*tol

    if N == 0
        N = length(inds)
    elseif N > length(inds)
        N = length(inds)
    end
    inds = inds[1:N]
    mval = abs(dSdtspc[inds[1]])
    minval = mval*tol
    k = 1
    while k < length(inds) && abs(dSdtspc[inds[k]]) >= minval
        k += 1
    end
    sensinds = inds[1:k]

    #Run analyses
    rxnanalysis = Array{ReactionAnalysis,1}()
    for rxnind in sensinds
        branches,rps = getbranchpathinfo(sim,spcind,rxnind,ropp,ropl,rts;steptol=steptol,branchtol=branchtol)
        radprodlossfract = getradprodlossfract(sim,rxnind,rts)
        if eliminate
            branches,rps = eliminatereasons(spcind,rxnind,branches,rps,dSdtspc;branchthreshold=branchthreshold,pathbranchthreshold=pathbranchthreshold)
        end
        push!(rxnanalysis,ReactionAnalysis(branches,rps,radprodlossfract,spcind,spcname,rxnind,dSdtspc[rxnind]))
    end
    return rxnanalysis
end
export analyzespc

"""
Calculate the fraction of the radical production or loss
attributable to the given reaction
"""
function getradprodlossfract(sim,rxnind,rts)
    radrops = rts.*getfield.(sim.reactions,:radicalchange)
    radprod = sum(r for r in radrops if r>0.0)
    radloss = sum(r for r in radrops if r<0.0)
    radflux = radrops[rxnind]
    if radflux > 0.0
        return radflux/radprod
    else
        return radflux/radloss
    end
end
export getradprodlossfract

"""
Identify species coupled to the target species whose pathways could
be important to the target species even if their pathways don't overlap
If the species has non-unimolecular loss reactions with branching greater
than branchtol the other species involved are considered coupled
"""
function getcoupledspecies(sim,spcind,rxnind,ropp,ropl,rts;branchtol=0.2)
    spcsinds = Array{Int64,1}()
    @views totflux = sum(ropl[:,spcind])
    @views inds = reverse(sortperm(ropl[:,spcind]))
    i = 1
    ind = inds[i]
    for ind in inds
        if ropl[ind,spcind]/totflux < branchtol
            break
        elseif length(sim.reactions[ind].reactants) > 1
            for rind in sim.reactions[ind].reactantinds
                if rind != spcind && rind > 0 && !(rind in spcsinds)
                    push!(spcsinds,rind)
                end
            end
        end
    end
    return spcsinds
end

"""
Run branching and pathway analysis for a give species and reaction
"""
function getbranchpathinfo(sim,spcind,rxnind,ropp,ropl,rts;steptol=1e-2,branchtol=5e-2,couplingbranchtol=0.2)
    branches = getbranching(sim,rxnind,ropl,rts)
    targetinds = [spcind]
    append!(targetinds,getcoupledspecies(sim,spcind,rxnind,ropp,ropl,rts;branchtol=couplingbranchtol))
    rps = Array{ReactionPath,1}()
    for spind in targetinds
        for forward in [true,false]
            rp = follow(sim,rxnind,spcind,ropp,ropl,rts,forward;steptol=steptol,branchtol=branchtol)
            if rp.branchfract != 0 && length(rp.spcsinds) > 0
                push!(rps,rp)
            end
        end
    end
    return branches,rps
end

"""
Attempt to remove importance hypotheses that don't make sense
if the reaction constitutes more than branchthreshold fraction of the flux we
consider it to not be sensitive due to branching
If the sensitivity is >0 increasing the rate coefficient increases concentration
so it can't be rate limited downstream so forward reaction paths don't make sense
Same for sensitivity <0 for reverse paths
"""
function eliminatereasons(spcind,rxnind,branches,rps,dSdtspc;branchthreshold=0.9,pathbranchthreshold=0.2)
    branchesout = Array{Branching,1}()
    for branch in branches
        ind = findfirst(isequal(rxnind),branch.rxninds)
        if ind !== nothing
            if branch.branchingratios[ind] < branchthreshold
                push!(branchesout,branch)
            end
        end
    end
    rpouts = Array{ReactionPath,1}()
    for rp in rps
        if rp.forward && dSdtspc[rxnind] > 0 #increasing forward (loss) paths should decrease spc concentration
            continue
        elseif !rp.forward && dSdtspc[rxnind] < 0 #increasing reverse (production) paths should increase spc concentration
            continue
        else
            rind = findfirst(isequal(rxnind),rp.rxninds)
            if rp.branchfracts[rind] > pathbranchthreshold
                push!(rpouts,rp)
            end
        end
    end
    return branchesout,rpouts
end
export eliminatereasons

"""
Generate a string report from the analysis objects
"""
function getrxnanalysisstring(sim,ra;branchingcutoff=1e-2,radbranchfract=0.01)
    spcname = sim.names[ra.spcind]
    rstr = getrxnstr(sim.domain.phase.reactions[ra.rxnind])
    sens = round(ra.sens,sigdigits=6)
    s = "Analyzing $spcname sensitivity to $rstr at a value of $sens \n"
    s *= "\n"
    for branch in ra.branchings
        sname = sim.names[branch.spcind]
        s *= "Key branching for $sname \n"
        for i = 1:length(branch.rxninds)+1
            br = round(branch.branchingratios[i],sigdigits=6)
            if br < branchingcutoff
                break
            else
                rstr = getrxnstr(sim.reactions[branch.rxninds[i]])
                s *= "$rstr had a branching ratio of $br \n"
            end
        end
        s *= "\n"
    end

    for rp in ra.paths
        spname = sim.names[rp.spcind]
        forward = rp.forward
        if rp.forward
            s *= "Associated key reaction path in $spname loss direction \n"
            for i = 1:length(rp.rxninds)
                rstr = getrxnstr(sim.reactions[rp.rxninds[i]])
                br = round(rp.branchfracts[i],sigdigits=6)
                s *= "$rstr at path branching of $br \n"
            end
        else
            s *= "Associated key reaction path in $spname production direction \n"
            revinds = reverse(rp.rxninds)
            for i = 1:length(rp.rxninds)
                rstr = getrxnstr(sim.reactions[revinds[i]])
                br = round(rp.branchfracts[i],sigdigits=6)
                s *= "$rstr at path step branching of $br \n"
            end
        end
        s *= "\n"
    end
    if abs(ra.radprodlossfract) > radbranchfract
        radfract = round(abs(ra.radprodlossfract),sigdigits=6)
        if ra.radprodlossfract > 0
            s *= "Reaction accounts for $radfract of the net radical production \n"
        else
            s *= "Reaction accounts for $radfract of the net radical loss \n"
        end
    end
    s *= "\n"
    return s
end
export getrxnanalysisstring

"""
Print out a string report from the analysis objects
"""
function printrxnanalysis(sim,ra;branchingcutoff=1e-2,radbranchfract=0.01)
    return println(getrxnanalysisstring(sim,ra;branchingcutoff=branchingcutoff,radbranchfract=radbranchfract))
end
export printrxnanalysis


text:
using FuzzySets, Test

step_size = 0.005
levels = collect(0:step_size:1)

A = TriangularFuzzyNumber(levels; number=0.0, width=0.5)
@test A.levels == levels
@test A.grades[1] == Interval(-0.5, 0.5)
@test A[1] == Interval(-0.5, 0.5)
@test support(A) == Interval(-0.5, 0.5)
@test core(A) == Interval(0) == Interval(0.0)
@test height(A) == 1.0
@test A(-0.2) == 0.6
@test A(0) == 1.0
@test A(0.2) == 0.6
@test A(0.5) == 0.0
@test A(1) == 0.0
@test length(A) == 201
@test round(centroid(A, step=0.01), digits=2) ≈ 0.0

B = TriangularFuzzyNumber(levels, number=5.0, width=0.5)
@test B.levels == levels
@test support(B) == Interval(4.5, 5.5)
@test core(B) == Interval(5.0)
@test height(B) == 1

C = TriangularFuzzyNumber(levels, number=5.0, width=0.5)
@test B == C
@test A != C

D = SingletonFuzzyNumber(levels, number=0)
@test support(D) == Interval(0)
@test D.grades == repeat([Interval(0)], length(levels)) # TODO: check with all(...)

A = TriangularFuzzyNumber(levels, number=1, width=0.5)
FuzzySets.vertical_cut!(A, α=0.2)
@test A.grades[1] == A.grades[41]
@test A.grades[2] == A.grades[41]
@test A.grades[40] == A.grades[41]
B = FuzzySets.vertical_cut(A, α=0.2)
@test B.grades[1] == B.grades[41]
@test B.grades[2] == B.grades[41]
@test B.grades[40] == B.grades[41]

A = SingletonFuzzyNumber(levels, number=5)
B = TriangularFuzzyNumber(levels, number=5)
@test FuzzySets.isSingleton(A) == true
@test FuzzySets.isSingleton(B) == false

p = 6
w_l = 2
w_r = 2
a = 2
@test FuzzySets.trapezoid(0; p, w_l, w_r, a) == Interval(2, 10)
@test FuzzySets.trapezoid(0.5; p, w_l, w_r, a) == Interval(3, 9)
@test FuzzySets.trapezoid(1.0; p, w_l, w_r, a) == Interval(4, 8)
@test_throws ErrorException FuzzySets.trapezoid(1.5; p, w_l, w_r, a)
A = TrapezoidalFuzzyNumber(levels; number=p, w_l, w_r, a)
@test A.grades[1] == Interval(2, 10)
@test A.grades[101] == Interval(3, 9)
@test A.grades[end] == Interval(4, 8)
@test round(centroid(A, step=0.01), digits=2) ≈ p

μ = 0
σ = 1
A = GaussianFuzzyNumber(levels; μ=μ, σ=σ)
@test A.grades[1] == A.grades[2]
@test A.grades[1].left == -A.grades[1].right
@test A.grades[101].left == -A.grades[101].right
@test A.levels[101] == 0.5
@test A.grades[101] == Interval(-1.1774100225154747, 1.1774100225154747)
@test A.grades[end] == Interval(μ)

# Specificity
A = TriangularFuzzyNumber(levels, number=0.5, width=0.5)
@test FuzzySets.specificity(A) == 0.5
@test FuzzySets.u_uncertainty(A) == 1

B = TriangularFuzzyNumber(levels, number=0.5, width=0.2)
@test FuzzySets.specificity(B) ≈ 0.8

C = TriangularFuzzyNumber(levels, number=8.0, width=1.0)
@test FuzzySets.specificity(C) ≈ 0.8

S = SingletonFuzzyNumber(levels, number=4.0)
@test FuzzySets.specificity(S) ≈ 1.0
@test FuzzySets.u_uncertainty(S) ≈ 0

step_size = 0.0001
levels = collect(0:step_size:1)

a = 0
b = 5
BoxFuzzyNumber = FuzzySets.BoxFuzzyNumber(levels, a=a, b=b)
w = b - a
@test round(FuzzySets.u_uncertainty(BoxFuzzyNumber), digits=3) ≈ round(log(1+w), digits=3)

A1 = TriangularFuzzyNumber(levels, number=2, width=2.0)
@test round(FuzzySets.u_uncertainty(A1), digits=3) ≈ 1.012

A2 = TriangularFuzzyNumber(levels, number=2, width=1.0)
@test round(FuzzySets.u_uncertainty(A2), digits=3) ≈ 0.648

# A3 = TriangularFuzzyNumber(levels, number=2, width=2.0)
# middle = Int((length(A3.levels) + 1)/2)
# A3.grades[1:middle] = A3.grades[middle:end]
# for i = middle+1:length(A3.levels)
#     A3.grades[i] = Interval(NaN)
# end
# @test round(FuzzySets.u_uncertainty(A3) * 2, digits=3) ≈ 0.648

A4 = TriangularFuzzyNumber(levels, number=2, width=0.5)
@test round(FuzzySets.u_uncertainty(A4), digits=3) ≈ 0.386 # not sure


text:
function shake(package)
    # raw snooping
    name = randstring(12)

    call = """
    SnoopCompile.@snoopc "/tmp/$name.log" begin
        using Pkg, $package
        Pkg.build("$package")
        include(joinpath(dirname(dirname(pathof($package))), "test", "runtests.jl"))
    end
    """;
    println("Snooping tests/build for $package...")
    eval(Meta.parse(call));

    # process snooping
    data = SnoopCompile.read("/tmp/$name.log");
    pc = SnoopCompile.parcel(reverse!(data[2]));
    rm("/tmp/$name.log")

    # construct deps sets
    println("Diffing dependencies...")
    used = [String(key) for key in keys(pc)];
    ctx = Pkg.Types.Context()
    pkg_ctx = ctx.env.manifest[ctx.env.project.deps[package]]
    listed = keys(pkg_ctx.deps)

    # filtering step, just to reduce false positives
    lowHangingFruit = setdiff(listed, used)
    println("Filtering...")
    filter!(fruit -> ~any(occursin.(Ref(fruit), data[2])), lowHangingFruit)

    println("These project deps are unused in tests:")
    return lowHangingFruit
end


text:
module TestForwardPopGenSimulations

using ForwardPopGenSimulations
using FactCheck

# write your own tests here
testnames = [
    "chromosomes"
]

for test in testnames
    Core.include("$test.jl")
end

end


text:
using Base.Test
using Klara

println("    Testing score functions for acceptance rate:")

println("      Testing logistic_rate_score...")

@test logistic_rate_score(0.25) == 1.7039056039366212
@test logistic_rate_score(0.5, 11) == 1.991859724568208

println("      Testing erf_rate_score...")

@test erf_rate_score(-0.1) == 0.6713732405408726
@test erf_rate_score(0.93, 2) == 1.9914724883356396

println("    Testing AcceptanceRateMCTuner constructors and methods...")

trate = 0.23
tuner = AcceptanceRateMCTuner(trate)

@test tuner.targetrate == trate
@test tuner.period == 100
@test tuner.verbose == false

tune = AcceptanceRateMCTune(1., 30, 100)
rate!(tune)
orate = tune.rate
tune!(tune, tuner)

@test tune.rate == logistic_rate_score(orate-tuner.targetrate)*orate

trate = 0.45
tuner = AcceptanceRateMCTuner(trate, score=erf_rate_score, period=1000, verbose=true)

@test tuner.targetrate == trate
@test tuner.period == 1000
@test tuner.verbose == true

tune = AcceptanceRateMCTune(1., 621, 1000)
rate!(tune)
orate = tune.rate
tune!(tune, tuner)

@test tune.rate == erf_rate_score(orate-tuner.targetrate)*orate


text:
"""
Outline:
1. Construct LDF.Inputs from openDSS file, load dict
2. Add power flow constraints to m,
    - set Pⱼ's = -1 * (dvGridPurchase_j - dvProductionToGrid_j)
3. solve model
"""
# TODO add complementary constraint to UL for dvProductionToGrid_ and dvGridPurchase_ (don't want it in LL s.t. it stays linear)

function LDF.build_ldf!(m::JuMP.AbstractModel, p::LDF.Inputs, ps::Array{REoptInputs, 1};
        make_import_export_complementary::Bool=true
    )
    LDF.add_variables(m, p)
    add_expressions(m, ps)
    LDF.constrain_power_balance(m, p)
    LDF.constrain_substation_voltage(m, p)
    LDF.constrain_KVL(m, p)
    LDF.constrain_loads(m, p, ps)
    LDF.constrain_bounds(m, p)

    if make_import_export_complementary
        add_complementary_constraints(m, ps)
    end
end

function add_expressions(m::JuMP.AbstractModel, ps::Array{REoptInputs, 1})
    for p in ps
        _n = string("_", p.s.site.node)
        m[Symbol("TotalExport"*_n)] = @expression(m, [t in p.time_steps],
            sum(
                m[Symbol("dvProductionToGrid"*_n)][t,u,ts]
                for t in p.techs.elec, u in p.export_bins_by_tech[t]
            )
        )
    end
end

function add_complementary_constraints(m::JuMP.AbstractModel, ps::Array{REoptInputs, 1})
    for p in ps
        _n = string("_", p.s.site.node)
        for (i, e) in zip(m[Symbol("dvGridPurchase"*_n)], m[Symbol("TotalExport"*_n)])
            @constraint(m,
                [i, e] in MOI.SOS1([1.0, 2.0])
            )
        end
    end
end

function LDF.constrain_loads(m::JuMP.AbstractModel, p::LDF.Inputs, ps::Array{REoptInputs, 1})
    reopt_nodes = [p.s.site.node for p in ps]

    Pⱼ = m[:Pⱼ]
    Qⱼ = m[:Qⱼ]
    # positive values are injections

    for j in p.busses
        if j in keys(p.Pload)
            if parse(Int, j) in reopt_nodes
                @constraint(m, [t in 1:p.Ntimesteps],
                    Pⱼ[j,t] == 1e3/p.Sbase * (  # 1e3 b/c REopt values in kW
                        m[Symbol("TotalExport_" * j)][t]
                        - m[Symbol("dvGridPurchase_" * j)][t]
                    )
                )
            else
                @constraint(m, [t in 1:p.Ntimesteps],
                    Pⱼ[j,t] == -p.Pload[j][t]
                )
            end
        elseif j != p.substation_bus
            @constraint(m, [t in 1:p.Ntimesteps],
                Pⱼ[j,t] == 0
            )
        end

        if j in keys(p.Qload)
            if parse(Int, j) in reopt_nodes
                @constraint(m, [t in 1:p.Ntimesteps],
                    Qⱼ[j,t] == 1e3/p.Sbase * p.pf * (  # 1e3 b/c REopt values in kW
                    m[Symbol("TotalExport_" * j)][t]
                        - m[Symbol("dvGridPurchase_" * j)][t]
                    )
                )
            else
                @constraint(m, [t in 1:p.Ntimesteps],
                    Qⱼ[j,t] == -p.Qload[j][t]
                )
            end
        elseif j != p.substation_bus
            @constraint(m, [t in 1:p.Ntimesteps],
                Qⱼ[j,t] == 0
            )
        end
    end
    p.Nequality_cons += 2 * (p.Nnodes - 1) * p.Ntimesteps
end

# TODO add LDF results (here and in LDF package)

function run_reopt(m::JuMP.AbstractModel, p::REoptInputs, ldf::LDF.Inputs)

end

max_stars_repo_name: NREL/REopt

max_issues_repo_name: NREL/REopt

max_forks_repo_name: NREL/REopt


text:
"""
Low-level type which manages the actual meshcat server. See [`Visualizer`](@ref)
for the public-facing interface.
"""
struct CoreVisualizer
    tree::SceneNode
    connections::Set{Any}
    host::IPAddr
    port::Int

    function CoreVisualizer(host::IPAddr = ip"127.0.0.1", default_port=8700)
        connections = Set([])
        tree = SceneNode()
        port = find_open_port(host, default_port, 500)
        core = new(tree, connections, host, port)
        start_server(core)
        return core
    end
end

function find_open_port(host, default_port, max_retries)
    for port in default_port:(default_port + max_retries)
        server = try
            listen(host, port)
        catch e
            if e isa Base.IOError
                continue
            end
        end
        close(server)
        # It is *possible* that a race condition could occur here, in which
        # some other process grabs the given port in between the close() above
        # and the open() below. But it's unlikely and would not be terribly
        # damaging (the user would just have to call open() again).
        return port
    end
end

function start_server(core::CoreVisualizer)
    asset_files = Set(["index.html", "main.min.js", "main.js"])

    function read_asset(file)
        if file in asset_files
            return open(s -> read(s, String), joinpath(VIEWER_ROOT, file))
        else
            return "Not found"
        end
    end
    default = "index.html"
    Mux.@app h = (
        Mux.defaults,
        Mux.page("/index.html", req -> read_asset("index.html")),
        Mux.page("/main.js", req -> read_asset("main.js")),
        Mux.page("/main.min.js", req -> read_asset("main.min.js")),
        Mux.page("/", req -> read_asset(default)),
        Mux.notfound());
    Mux.@app w = (
        Mux.wdefaults,
        Mux.route("/", req -> add_connection!(core, req)),
        Mux.wclose,
        Mux.notfound());
    @async begin
        # Suppress noisy unhelpful log messages from HTTP.jl, e.g.
        # https://github.com/JuliaWeb/HTTP.jl/issues/392
        Logging.with_logger(Logging.NullLogger()) do
            WebSockets.serve(
            WebSockets.ServerWS(
                Mux.http_handler(h),
                Mux.ws_handler(w),
            ), core.host, core.port);
        end
    end
    @info "MeshCat server started. You can open the visualizer by visiting the following URL in your browser:\n$(url(core))"
end

function url(core::CoreVisualizer)
    "http://$(core.host):$(core.port[])"
end

function add_connection!(core::CoreVisualizer, req)
    connection = req[:socket]
    push!(core.connections, connection)
    send_scene(core, connection)
    wait()
end

function update_tree!(core::CoreVisualizer, cmd::SetObject, data)
    core.tree[cmd.path].object = data
end

function update_tree!(core::CoreVisualizer, cmd::SetTransform, data)
    core.tree[cmd.path].transform = data
end

function update_tree!(core::CoreVisualizer, cmd::SetProperty, data)
    core.tree[cmd.path].properties[cmd.property] = data
end

function update_tree!(core::CoreVisualizer, cmd::Delete, data)
    if length(cmd.path) == 0
        core.tree = SceneNode()
    else
        delete!(core.tree, cmd.path)
    end
end

function update_tree!(core::CoreVisualizer, cmd::SetAnimation, data)
    core.tree.animation = data
end

function update_tree!(core::CoreVisualizer, cmd::SaveImage, data)
    nothing
end

function send_scene(core::CoreVisualizer, connection)
    foreach(core.tree) do node
        if node.object !== nothing
            WebSockets.writeguarded(connection, node.object)
        end
        if node.transform !== nothing
            WebSockets.writeguarded(connection, node.transform)
        end
        for data in values(node.properties)
            WebSockets.writeguarded(connection, data)
        end
        if node.animation !== nothing
            WebSockets.writeguarded(connection, node.animation)
        end
    end
end

function Base.write(core::CoreVisualizer, data)
    for connection in core.connections
        if isopen(connection)
            WebSockets.writeguarded(connection, data)
        else
            delete!(core.connections, connection)
        end
    end
end

function send(core::CoreVisualizer, cmd::AbstractCommand)
    data = pack(lower(cmd))
    update_tree!(core, cmd, data)
    write(core, data)

    nothing
end

function Base.wait(core::CoreVisualizer)
    while isempty(core.connections)
        sleep(0.5)
    end
end

"""
    vis = Visualizer()

Construct a new MeshCat visualizer instance.

Useful methods:

    vis[:group1] # get a new visualizer representing a sub-tree of the scene
    setobject!(vis, geometry) # set the object shown by this visualizer's sub-tree of the scene
    settransform!(vis, tform) # set the transformation of this visualizer's sub-tree of the scene
    setvisible!(vis, false) # hide this part of the scene
"""
struct Visualizer <: AbstractVisualizer
    core::CoreVisualizer
    path::Path
end

Visualizer() = Visualizer(CoreVisualizer(), ["meshcat"])

"""
$(TYPEDSIGNATURES)

Wait until at least one browser has connected to the
visualizer's server. This is useful in scripts to delay
execution until the browser window has opened.
"""
Base.wait(v::Visualizer) = wait(v.core)

# IJuliaCell(vis::Visualizer; kw...) = iframe(vis.core; kw...)

Base.show(io::IO, v::Visualizer) = print(io, "MeshCat Visualizer with path $(v.path) at $(url(v.core))")

"""
Set the object at this visualizer's path. This replaces whatever
geometry was presently at that path.

$(TYPEDSIGNATURES)

To draw multiple geometries, place them at different paths by using the slicing notation:

    setobject!(vis[:group1][:box1], geometry1)
    setobject!(vis[:group1][:box2], geometry2)
"""
function setobject!(vis::Visualizer, obj::AbstractObject)
    send(vis.core, SetObject(obj, vis.path))
    vis
end

"""
Set the transform of this visualizer's path. This can be done before or after
adding an object at that path.

$(TYPEDSIGNATURES)

The overall transform of an object is the composition of the transforms of all
of its parents, so setting the transform of `vis[:group1]` affects the poses of
the objects at `vis[:group1][:box1]` and `vis[:group1][:box2]`.
"""
function settransform!(vis::Visualizer, tform::Transformation)
    send(vis.core, SetTransform(tform, vis.path))
    vis
end

"""
Delete the geometry at this visualizer's path and all of its descendants.

$(TYPEDSIGNATURES)
"""
function delete!(vis::Visualizer)
    send(vis.core, Delete(vis.path))
    vis
end

"""
Set a single property for the object at the given path.

$(TYPEDSIGNATURES)

(this is named setprop! instead of setproperty! to avoid confusion
with the Base.setproperty! function introduced in Julia v0.7)
"""
function setprop!(vis::Visualizer, property::AbstractString, value)
    send(vis.core, SetProperty(vis.path, property, value))
    vis
end

"""
Set the currently playing animation in the visualizer.

$(TYPEDSIGNATURES)

"""
function setanimation!(vis::Visualizer, anim::Animation; play::Bool=true, repetitions::Integer=1)
    cmd = SetAnimation(anim, play, repetitions)
    send(vis.core, cmd)
end

"""
Takes a screenshot of the current visualizer and saves it as an image to the disk.

$(TYPEDSIGNATURES)
"""
function save_image(vis::Visualizer)
    send(vis.core, SaveImage())
    vis
end

"""
Get a new `Visualizer` representing a sub-tree of the same scene.

$(TYPEDSIGNATURES)

For example, if you have `vis::Visualizer` with path `/meshcat/foo`, you can do
`vis[:bar]` to get a new `Visualizer` with path `/meshcat/foo/bar`.
"""
Base.getindex(vis::Visualizer, path...) =
    Visualizer(vis.core, joinpath(vis.path, path...))


text:
"""
    share!(d::Dict, set::Dict; save_build = true, overwrite = false)

# Arguments
- `d::Dict` of DataFrames containing the model data.
- `set::Dict` of Arrays describing region, sector, final demand, etc.

# Keywords
- `save_build = true`
- `overwrite = false`
See [`SLiDE.build`](@ref) for keyword argument descriptions.

# Returns
- `d::Dict` of DataFrames containing the model data at the sharing step.
"""
function share(
    dataset::String,
    d::Dict,
    set::Dict;
    save_build::Bool=DEFAULT_SAVE_BUILD,
    overwrite::Bool=DEFAULT_OVERWRITE,
)
    CURR_STEP = "share"

    # If there is already sharing data, read it and return.
    d_read = read_build(dataset, CURR_STEP; overwrite=overwrite)
    !(isempty(d_read)) && (return d_read, set)

    # Add CFS regional sets for filtering.
    [set[k] = set[:r] for k in [:orig,:dest]]

    # Read sharing input data.
    d_read = read_from(joinpath("src", "build", "readfiles", "input", "share.yml"))
    d_read = Dict(k => sort(dropmissing(edit_with(
        filter_with(df, set; extrapolate=true),
        Deselect([:units], "==")
    ))) for (k, df) in d_read)

    merge!(d, d_read)

    share_pce!(d)
    share_sgf!(d)
    share_utd!(d, set)
    share_region!(d, set)
    share_labor!(d, set)
    share_rpc!(d, set)

    write_build!(dataset, CURR_STEP, d; save_build=save_build)
    write_build!(dataset, SET_DIR, Dict(k => set[k] for k in [:notrd,:ng]))

    return (d, set)
end

max_stars_repo_name: NREL/SLiDE

max_issues_repo_name: NREL/SLiDE

max_forks_repo_name: NREL/SLiDE


text:
import SimplePlots: plot, plot!, scatter, scatter!

# some plotting utilities
export plotif, trimplot, signchart
export arrow!, vectorfieldplot!
export newton_vis

###
include("demos.jl")

#  plot "recipe" for functions
function SimplePlots.plot(f::Function, a::Number, b::Number, args...; kwargs...)
    xs = range(float(a), float(b), length=500)
    plot(xs, f.(xs), args...; kwargs...)
end

function  SimplePlots.plot!(plt::SimplePlots.SimplePlot, f::Function, a::Number, b::Number, args...;  kwargs...)
    xs = range(float(a), float(b), length=500)
    plot!(plt, xs, f.(xs);  kwargs...)
end
SimplePlots.plot!(f::Function, args...; kwargs...) = SimplePlots.plot!(SimplePlots._plot, f, args...; kwargs...)

function  SimplePlots.plot!(plt::SimplePlots.SimplePlot, f::Function;  kwargs...)
    xs = plt.data[1]["x"]
    plot!(plt, xs, f.(xs);  kwargs...)
end

#  plot "recipe" for parametric functions
function SimplePlots.plot(f::Function, g::Function, a::Number, b::Number, args...; kwargs...)
    xs = range(float(a), float(b), length=500)
    plot(f.(xs), g.(xs), args...; kwargs...)
end

function  SimplePlots.plot!(plt::SimplePlots.SimplePlot, f::Function, g::Function, a::Number, b::Number, args...; kwargs...)
    xs = range(a, stop=b, length=500)
    plot!(plt, xs, f.(xs), args...;  kwargs...)
end
SimplePlots.plot!(f::Function, g::Function, a::Number, b::Number, args...; kwargs...) = SimplePlots.plot!(SimplePlots._plot, f, g, a, b, args...; kwargs...)

#  plot reciple  for SymPy objects
SimplePlots.plot(ex::Sym, a::Number, b::Number; kwargs...) = plot(lambdify(ex), a, b; kwargs...)
SimplePlots.plot!(ex::Sym; kwargs...)    = plot!(lambdify(ex); kwargs...)

SimplePlots.plot(ex1::Sym, ex2::Sym,  a::Number, b::Number; kwargs...) = plot(lambdify(ex1), lambdify(ex2), a, b; kwargs...)
SimplePlots.plot!(ex1::Sym, ex2::Sym,  a::Number, b::Number; kwargs...) =
    plot!(lambdify(ex1), lambdify(ex2), a,b; kwargs...)

# Plot of tuple
function SimplePlots.plot(fs::Tuple, a::Number, b::Number; kwargs...)

    xs = range(a, stop=b, length=500)
    ys = [ [f(x) for x in xs] for f in fs]
    if length(ys) == 1
        plot(xs, ys[1]; kwargs...)
    else
        plot(ys...; kwargs...)
    end
end

function SimplePlots.plot!(fs::Tuple, args...; kwargs...)
    if length(fs) == 1
        xs = plt.data[1]["x"]
        plot!(xs, fs[1].(xs); kwargs...)
    else
        xs = range(args[1], stop=args[2], length=500)
        ys = [ [f(x) for x in xs] for f in fs]
        plot!(ys...; kwargs...)
    end
end

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

"""
   trimplot(f, a, b, c=20; kwargs...)

Plot f over [a,b] but break graph if it exceeds c in absolute value.
"""
function trimplot(f, a, b, c=20; kwargs...)
  xs = range(a, stop=b, length=251)
  ys = f.(xs)

  us, vs = Real[NaN], Real[NaN]
  p = plot(us, vs, xlim=(a, b), legend=false, kwargs...)
  for (x,y) in zip(xs, ys)
    if abs(y) <= c
       push!(us, x); push!(vs, y)
    else
      length(us) > 0 && plot!(p, us, vs, color=1)
      empty!(us); empty!(vs)
    end
 end
 length(us) > 0 && plot!(p, us, vs, color=1)
 p
end

"""
    plotif(f, g, a, b)

Plot f colored depending on g >= 0 or not.
"""
function plotif(f, g, a, b, args...; colors=(1,2), linewidth=5, legend=false,  kwargs... )

    xs = a:(b-a)/251:b
    zs = f.(xs)
    p = plot(xs, f.(xs), args...; color=colors[1], linewidth=linewidth, legend=legend, kwargs...)

    ys = g.(xs)
    ys[ys .< 0] .= NaN

    us,vs = Float64[], Float64[]
    for (i,y) in enumerate(ys)
        if isnan(y)
            if length(vs) > 1
                plot!(p, us, vs, color=colors[2], linewidth=5)
            end
            empty!(us)
            empty!(vs)
        else
            push!(us, xs[i])
            push!(vs, zs[i])
        end
    end
    if length(vs) > 1
        plot!(p, us, vs, color=colors[2], linewidth=5)
    end
    p
end

"""
   signchart(f, a, b)

Plot f over a,b with different color when negative.
"""
function signchart(f, a, b)
    p = plotif(f, f, a, b)
    plot!(p, zero)
    p
end

"""
   `arrow!(p, v)`

Add the vector `v` to the plot anchored at `p`.

This would just be a call to `quiver`, but there is no 3-D version of that. As well, the syntax for quiver is a bit awkward for plotting just a single arrow. (Though efficient if plotting many).

```
using SimplePlots
r(t) = [sin(t), cos(t), t]
rp(t) = [cos(t), -sin(t), 1]
plot(unzip(r, 0, 2pi)...)
t0 = 1
arrow!(r(t0), r'(t0))
```
"""
function arrow!(plt, p, v; kwargs...)
#function arrow!(plt::SimplePlots.plot, p, v; kwargs...)
  if length(p) == 2
     quiver!(plt, unzip([p])..., quiver=Tuple(unzip([v])); kwargs...)
  elseif length(p) == 3
    # 3d quiver needs support
    # https://github.com/JuliaPlots/Plots.jl/issues/319#issue-159652535
    # headless arrow instead
    plot!(plt, unzip(p, p+v)...; kwargs...)
	end
end
#arrow!(p,v;kwargs...) = arrow!(Plots.current(), p, v; kwargs...)

"""

    vectorfieldplot(V; xlim=(-5,5), ylim=(-5,5), n=10; kwargs...)

V is a function that takes a point and returns a vector (2D dimensions), such as `V(x) = x[1]^2 + x[2]^2`.

The grid `xlim × ylim` is paritioned into (n+1) × (n+1) points. At each point, `pt`, a vector proportional to `V(pt)` is drawn.

This is written to add to an existing plot.

```
plot()  # make a plot
V(x,y) = [x, y-x]
vectorfield_plot!(p, V)
p
```
"""
function vectorfieldplot!(plt, V; xlim=(-5,5), ylim=(-5,5), n=10, kwargs...)

    dx, dy = (xlim[2]-xlim[1])/n, (ylim[2]-ylim[1])/n
    xs, ys = xlim[1]:dx:xlim[2], ylim[1]:dy:ylim[2]

    ps = [[x,y] for x in xs for y in ys]
    vs = V.(ps)
    λ = 0.9 * min(dx, dy) /maximum(norm.(vs))

    quiver!(plt, unzip(ps)..., quiver=unzip(λ * vs))

end
vectorfieldplot!(V; kwargs...) = vectorfieldplot!(SimplePlots.current(), V; kwargs...)


text:
__precompile__()

module Convex
using OrderedCollections: OrderedDict
using LinearAlgebra
using SparseArrays

global DEFAULT_SOLVER = nothing
### modeling framework
include("dcp.jl")
include("expressions.jl")
# need to define `Variable` before `UniqueConicForms`
include("variable.jl")
include("conic_form.jl")
# need to define `conic_form!` for `Variable`s after `UniqueConicForms`
include("variable_conic_form.jl")
include("constant.jl")
include("constraints/constraints.jl")
include("constraints/signs_and_sets.jl")
include("constraints/soc_constraints.jl")
include("constraints/exp_constraints.jl")
include("constraints/sdp_constraints.jl")
include("solver_info.jl")
include("problems.jl")
include("solution.jl")

### affine atoms
include("atoms/affine/add_subtract.jl")
include("atoms/affine/multiply_divide.jl")
include("atoms/affine/sum.jl")
include("atoms/affine/transpose.jl")
include("atoms/affine/index.jl")
include("atoms/affine/diag.jl")
include("atoms/affine/diagm.jl")
include("atoms/affine/stack.jl")
include("atoms/affine/dot.jl")
include("atoms/affine/reshape.jl")
include("atoms/affine/trace.jl")
include("atoms/affine/partialtrace.jl")
include("atoms/affine/partialtranspose.jl")
include("atoms/affine/conv.jl")
include("atoms/affine/real_imag.jl")
include("atoms/affine/inner_product.jl")
include("atoms/affine/conjugate.jl")
include("atoms/affine/kron.jl")

### lp atoms
include("atoms/lp_cone/abs.jl")
include("atoms/lp_cone/maximum.jl")
include("atoms/lp_cone/minimum.jl")
include("atoms/lp_cone/max.jl")
include("atoms/lp_cone/min.jl")
include("atoms/lp_cone/sumlargest.jl")
include("atoms/lp_cone/dotsort.jl")

### SOC atoms
include("atoms/second_order_cone/norm.jl") # also includes some lp atoms
include("atoms/second_order_cone/norm2.jl")
include("atoms/second_order_cone/quadoverlin.jl")
include("atoms/second_order_cone/qol_elementwise.jl")
include("atoms/second_order_cone/geomean.jl")
include("atoms/second_order_cone/quadform.jl")
include("atoms/second_order_cone/power_to_socp.jl")
include("atoms/second_order_cone/rationalnorm.jl")
include("atoms/second_order_cone/huber.jl")

### SDP atoms
include("atoms/sdp_cone/nuclearnorm.jl")
include("atoms/sdp_cone/operatornorm.jl")
include("atoms/sdp_cone/lambda_min_max.jl")
include("atoms/sdp_cone/matrixfrac.jl")
include("atoms/sdp_cone/sumlargesteigs.jl")

### exponential atoms
include("atoms/exp_cone/exp.jl")
include("atoms/exp_cone/log.jl")
include("atoms/exp_cone/logsumexp.jl")
include("atoms/exp_cone/entropy.jl")
include("atoms/exp_cone/relative_entropy.jl")

### exp + sdp atoms
include("atoms/exp_+_sdp_cone/logdet.jl")

### utilities
include("utilities/tree_print.jl")
include("utilities/tree_interface.jl")
include("utilities/show.jl")
include("utilities/iteration.jl")
include("utilities/broadcast.jl")

# Deprecated workaround for memory leak (https://github.com/JuliaOpt/Convex.jl/issues/83)
function clearmemory()
    Base.depwarn("Convex.clearmemory() is deprecated, as the memory leak it works around has been closed (in https://github.com/JuliaOpt/Convex.jl/pull/322). This function no longer does anything and will be removed in a future Convex.jl release.", :clearmemory )
end

end


text:
export SolitaryWaveWhithamBoussinesq, SolitaryWB

"""
    SolitaryWaveWhithamBoussinesq(param; kwargs)

Compute the Whitham-Boussinesq solitary wave with prescribed velocity.

# Argument
- `param :: NamedTuple`: parameters of the problem containing velocity `c` \
and dimensionless parameters `ϵ` and `μ`,
and mesh size `L` and number of collocation points `N`;

## Keywords (optional)
- `guess :: Vector{Real}`: initial guess for the surface deformation (if not provided, the exact formula for SGN is used);
- `x₀ :: Real`: center of solitary wave (if guess is not provided);
- `α :: Real`: determines the model used (typically `1` or `1/2`, default is 1);
- `Boussinesq`: if `true` (default is `false`), compute the standard Boussinesq system with parameters `a` (defaut `-1//3`), `b=d` (defaut `1//3`), and `c=0`);
- `iterative :: Bool`: inverts Jacobian through GMRES if `true`, LU decomposition if `false`;
- `verbose :: Bool`: prints numerical errors at each step if `true`;
- `max_iter :: Int`: maximum number of iterations of the Newton algorithm;
- `tol :: Real`: general tolerance (default is `1e-10`);
- `ktol :: Real`: tolerance of the Krasny filter (default is `0`, i.e. no filtering);
- `gtol :: Real`: relative tolerance of the GMRES algorithm (default is `1e-10`);
- `dealias :: Int`: dealiasing with Orlicz rule `1-dealias/(dealias+2)` (default is `0`, i.e. no dealiasing);
- `q :: Real`: Newton algorithm modified with
`u_{n+1}=q*(u_n+du)+(1-q)*u_n`
(default is `1`);
- `β :: Real`: adds `β` times spectral projection onto the Kernel to the Jacobian.β

# Return values
`(η,v,mesh)` with
- `η :: Vector{Float64}`: surface deformation;
- `v :: Vector{Float64}`: velocity (derivative of the trace of the velocity potential at the surface);
- `mesh :: Mesh`: mesh collocation points.

"""
function SolitaryWaveWhithamBoussinesq(
        param :: NamedTuple;
        guess = zeros(0) :: Vector{Float64},
        x₀ = 0 :: Real,
        α = 1 :: Real,
		a = -1//3, b = 1//3,
		Boussinesq = false :: Bool,
        iterative = false :: Bool,
        verbose = false :: Bool,
        max_iter = 20 :: Int,
        tol = 1e-10 :: Real,
        ktol = 0 :: Real,
        gtol = 1e-10 :: Real,
        dealias = 0 :: Int,
        q=1 :: Real,
        β=0 :: Real)

        c = param.c
        ϵ = param.ϵ
        μ = param.μ
        if abs(c)<1
                @error("The velocity must be greater than 1 (in absolute value).")
        end

        mesh=Mesh(param)
        if guess == zeros(0)
                # Using the exact formula for the SGN solitary wave as initial guess (if not provided)
                η = (c^2-1)/ϵ*sech.(sqrt(3*(c^2-1)/(c^2)/μ)/2*(mesh.x.-x₀)).^2
                h = 1 .+ param.ϵ*η
                u = c*η./h
                k = mesh.k
        	F₀ = sqrt(param.μ)*1im * k
        	DxF(v) = real.(ifft(F₀ .* fft(v)))
        	guess = u - 1/3 ./h .* (DxF(h.^3 .*DxF(u)))
        end

        k = mesh.k
	if Boussinesq==false
		F₁ 	= tanh.(sqrt(μ)*abs.(k))./(sqrt(μ)*abs.(k))
		F₁[1] 	= 1
		F₂ = F₁.^α
	else
		F₂ = 1 ./(1 .+μ*b*abs.(k).^2)
		F₁ 	= (1 .-μ*a*abs.(k).^2).*(F₂.^2)
	end
        Dx       =  1im * sqrt(μ)*k

        if dealias == 0
                Π = ones(size(k)) # no dealiasing (Π=Id)
        else
                K = (mesh.kmax-mesh.kmin)/(2+dealias)
                Π = abs.(k) .<= K # Dealiasing low-pass filter
        end
        krasny(k) = (abs.(k).>= ktol ).*k
        krasny!(k) = k[abs.(k).< ktol ].=0

        function filt( v :: Vector{Float64} )
                real.(ifft(krasny(Π.*fft(v))))
        end
        function filt( v :: Vector{Complex{Float64}} )
                ifft(krasny(Π.*fft(v)))
        end
        function Four1( v  )
                ifft(F₁.*fft(v))
        end
        function Four2( v  )
                ifft(F₂.*fft(v))
        end
        function FF( v )
                w = Four2(v)
                return real.(-c^2*v .+ Four1(v) .+ c*ϵ/2*w.^2 .+ ϵ*Four2(c*w.*v .- ϵ/2*w.^3))
        end
        function Fabs( v )
                w = Four2(v)
                return real.(c^2*abs.(v) .+ abs.(Four1(v)) .+ c*ϵ/2*w.^2 .+ ϵ*Four2(c*abs.(w.*v) .+ ϵ/2*abs.(w).^3))
        end

        if iterative == false
                k = mesh.k
                x = mesh.x
                x₀ = mesh.x[1]
                FFT = exp.(-1im*k*(x.-x₀)');
                IFFT = exp.(1im*k*(x.-x₀)')/length(x);
                Id = Diagonal( ones(size(x)));
                M₁ = IFFT*(Diagonal( F₁)*FFT)
                M₂ = IFFT*(Diagonal( F₂)*FFT)
                function JacF( v ,dxv )
                        M₀ = Diagonal( Four2(v))
                        Symmetric(real.(-c^2 .* Id .+ M₁ .+ c*ϵ.*(M₀*M₂ .+ M₂*M₀) .+ ϵ*M₂*(Diagonal(c*v -3*ϵ/2*Four2(v).^2))*M₂ .+ β*dxv*dxv'))
                end
        else
                function JacFfast( v , dxv )
                        dF(φ) = real.(-c^2*φ+Four1(φ)+c*ϵ*Four2(v).*Four2(φ)+ϵ*Four2(c*Four2(v).*φ+(c*v -3*ϵ/2*(Four2(v).^2)).*Four2(φ))+ β*dot(dxv,φ)*dxv)
                        return LinearMap(dF, length(v); issymmetric=true, ismutating=false)
                end
        end

        flag=0
        iter = 0
        u = filt(guess)
        du = similar(u)
        fu = similar(u)
        dxu = similar(u)
        for i in range(0, stop=max_iter)
                dxu .= real.(ifft(Dx.*fft(u)))
                dxu ./= norm(dxu,2)
                fu .= FF(u)
                relerr = norm(fu,Inf)/norm(Fabs(u),Inf)
                abserr = norm(fu,Inf)
                if relerr < tol
    			@info string("Converged : relative error ",relerr," in ",i," steps\n")
    			break
    		elseif verbose == true
                        print(string("absolute error at step ",i,": ",abserr,"\n"))
                        print(string("relative error at step ",i,": ",relerr,"\n"))

    		end
                if i == max_iter
                        @warn string("The algorithm did not converge after ",i," steps: final relative error is ",relerr,"\n")
                        break
                end

                if iterative == false
                        du .= JacF(u,dxu) \ fu
                else
                        du .= gmres( JacFfast(u,dxu) , fu ; verbose = verbose, reltol = gtol )
                end
    		u .-= q*filt(du)
        end
        return (c*u-ϵ/2*real.(Four2(u).^2),u,mesh)
end

"""
    SolitaryWB(param; kwargs)

Build the initial data associated with `SolitaryWaveWhithamBoussinesq(param; kwargs)`, of type `InitialData`,
to be used in initial-value problems `Problem(model, initial::InitialData, param)`.

---
    SolitaryWB(c; ϵ=1,μ=1,N=2^12,kwargs)
Build the initial data with velocity `c`, dimensionless parameters `ϵ` and `μ`, and number of collocation points `N`, \
and `kwargs` the other (optional) keyword arguments as above.
"""
struct SolitaryWB <: InitialData

	η
	v
	label :: String
	info  :: String

	function SolitaryWB(param::NamedTuple;
					guess = zeros(0) :: Vector{Float64},
			        x₀ = 0 :: Real,
			        α = 1 :: Real,
					a = -1//3, b = 1//3,
					Boussinesq = false :: Bool,
			        iterative = false :: Bool,
			        verbose = false :: Bool,
			        max_iter = 20 :: Int,
			        tol = 1e-10 :: Real,
			        ktol = 0 :: Real,
			        gtol = 1e-10 :: Real,
			        dealias = 0 :: Int,
			        q=1 :: Real,
			        β=0 :: Real)

        (η,v,mesh)=SolitaryWaveWhithamBoussinesq(param;
                        guess=guess,x₀=x₀,α=α,a=a,b=b,Boussinesq=Boussinesq,
                        iterative=iterative,verbose=verbose,max_iter=max_iter,
                        tol=tol,ktol=ktol,gtol=gtol,dealias=dealias,q=q,β=β)
		init = Init(mesh,η,v)
		if Boussinesq == false
        		model = "Whitham-Boussinesq"
		else
			model = "Boussinesq"
		end
                label = "$model solitary wave"
		info = "Solitary travelling wave for the $model model.\n\
		├─velocity c = $(param.c)\n\
		└─maximum h₀ = $(maximum(η)) (from rest state)."

		new( init.η,init.v,label,info  )
	end

	function SolitaryWB(
                c::Real; ϵ=1::Real,μ=1::Real,N=2^12::Int,
				guess = zeros(0) :: Vector{Float64},
		        x₀ = 0 :: Real,
		        α = 1 :: Real,
				a = -1//3, b = 1//3,
				Boussinesq = false :: Bool,
		        iterative = false :: Bool,
		        verbose = false :: Bool,
		        max_iter = 20 :: Int,
		        tol = 1e-10 :: Real,
		        ktol = 0 :: Real,
		        gtol = 1e-10 :: Real,
		        dealias = 0 :: Int,
		        q=1 :: Real,
		        β=0 :: Real)

		L=200/sqrt(3*(c^2-1)/(c^2)/μ)/2
		xmin,xmax = x₀-L,x₀+L;
		param=(ϵ=ϵ,μ=μ,c=c,L=L,xmin=xmin,xmax=xmax,N=N)
		sol=SolitaryWB(param;
                        guess=guess,x₀=x₀,α=α,a=a,b=b,Boussinesq=Boussinesq,
                        iterative=iterative,verbose=verbose,max_iter=max_iter,
                        tol=tol,ktol=ktol,gtol=gtol,dealias=dealias,q=q,β=β)
		new( sol.η,sol.v,sol.label,sol.info  )
	end

end


text:
using HetaSimulator
using Test

function test_show(t)
  try
    display(t)
    return true
  catch e
    return false
  end
end

@testset "HetaSimulator" begin
  @testset "Heta build" begin include("heta_test.jl") end
  @testset "Single-compartment model without events" begin include("single_comp_test.jl") end
end


text:
module ScalarArrays

using Base.Broadcast: Broadcasted, DefaultArrayStyle

export ScalarArray

mutable struct ScalarArray{T} <: AbstractArray{T, 0}
    val::T
    function ScalarArray{T}() where {T}
        new{T}()
    end
    function ScalarArray{T}(val) where {T}
        new{T}(val)
    end
end
ScalarArray(val::T) where {T} = ScalarArray{T}(val)

Base.axes(::ScalarArray) = ()
Base.size(::ScalarArray) = ()
(Base.getindex(arr::ScalarArray{T})::T) where {T} = arr.val
Base.setindex!(arr::ScalarArray, x) = (arr.val = x)

end


text:
"""
```
measurement(m::JPT{T}, TTT::Matrix{T}, RRR::Matrix{T},
            CCC::Vector{T}) where {T<:AbstractFloat}
```

Assign measurement equation

```
y_t = ZZ*s_t + DD + u_t
```

where

```
Var(ϵ_t) = QQ
Var(u_t) = EE
Cov(ϵ_t, u_t) = 0
```
"""
function measurement(m::JPT{T}, TTT::Matrix{T}, RRR::Matrix{T}, CCC::Vector{T}) where {T<:AbstractFloat}

    endo      = m.endogenous_states
    endo_addl = m.endogenous_states_augmented
    exo       = m.exogenous_shocks
    obs       = m.observables

    _n_observables      = n_observables(m)
    _n_states           = n_states_augmented(m)
    _n_shocks_exogenous = n_shocks_exogenous(m)

    ZZ = zeros(_n_observables, _n_states)
    DD = zeros(_n_observables)
    EE = zeros(_n_observables, _n_observables)
    QQ = zeros(_n_shocks_exogenous, _n_shocks_exogenous)

    ## Output growth - Quarterly!
    ZZ[obs[:obs_gdp], endo[:gdp_t]]  = 1.0
    ZZ[obs[:obs_gdp], endo[:gdp_t1]] = -1.0
    ZZ[obs[:obs_gdp], endo[:z_t]]    = 1.0
    DD[obs[:obs_gdp]]                = m[:γ100]

    ## Consumption Growth
    ZZ[obs[:obs_consumption], endo[:c_t]]  = 1.0
    ZZ[obs[:obs_consumption], endo[:c_t1]] = -1.0
    ZZ[obs[:obs_consumption], endo[:z_t]]  = 1.0
    DD[obs[:obs_consumption]]              = m[:γ100]

    ## Investment Growth
    ZZ[obs[:obs_investment], endo[:i_t]]  = 1.0
    ZZ[obs[:obs_investment], endo[:i_t1]] = -1.0
    ZZ[obs[:obs_investment], endo[:z_t]]  = 1.0
    DD[obs[:obs_investment]]              = m[:γ100]

    ## Hours growth
    ZZ[obs[:obs_hours], endo[:L_t]] = 1.0
    DD[obs[:obs_hours]]             = m[:L_ss]

    ## Real wage growth
    ZZ[obs[:obs_wages], endo[:w_t]]  = 1.0
    ZZ[obs[:obs_wages], endo[:w_t1]] = -1.0
    ZZ[obs[:obs_wages], endo[:z_t]]  = 1.0
    DD[obs[:obs_wages]]              = m[:γ100]

    ## Inflation (GDP Deflator)
    ZZ[obs[:obs_gdpdeflator], endo[:π_t]] = 1.0
    DD[obs[:obs_gdpdeflator]]             = m[:π_ss100]

    ## Nominal interest rate
    ZZ[obs[:obs_nominalrate], endo[:R_t]] = 1.0
    DD[obs[:obs_nominalrate]]             = m[:π_ss100] + m[:r_ss100]

    # Variance of innovations
    QQ[exo[:mp_sh], exo[:mp_sh]]   = m[:σ_mp]^2
    QQ[exo[:z_sh], exo[:z_sh]]     = m[:σ_z]^2
    QQ[exo[:g_sh], exo[:g_sh]]     = m[:σ_g]^2
    QQ[exo[:μ_sh], exo[:μ_sh]]     = m[:σ_μ]^2
    QQ[exo[:λ_p_sh], exo[:λ_p_sh]] = m[:σ_λ_p]^2
    QQ[exo[:λ_w_sh], exo[:λ_w_sh]] = m[:σ_λ_w]^2
    QQ[exo[:b_sh], exo[:b_sh]]     = m[:σ_b]^2

    # Adjustment to DD because measurement equation assumes CCC is the zero vector
    if any(CCC .!= 0)
        DD += ZZ * ((UniformScaling(1) - TTT) \ CCC)
    end

    return Measurement(ZZ, DD, QQ, EE)
end


text:
module Compiler

import ..Enzyme: Const, Active, Duplicated, DuplicatedNoNeed
import ..Enzyme: API, TypeTree, typetree, TypeAnalysis, FnTypeInfo, Logic

using LLVM, GPUCompiler, Libdl
import Enzyme_jll

import GPUCompiler: CompilerJob, FunctionSpec, codegen
using LLVM.Interop
import LLVM: Target, TargetMachine

# We have one global JIT and TM
const jit = Ref{OrcJIT}()
const tm  = Ref{TargetMachine}()

function __init__()
    opt_level = Base.JLOptions().opt_level
    if opt_level < 2
        optlevel = LLVM.API.LLVMCodeGenLevelNone
    elseif opt_level == 2
        optlevel = LLVM.API.LLVMCodeGenLevelDefault
    else
        optlevel = LLVM.API.LLVMCodeGenLevelAggressive
    end

    tm[] = GPUCompiler.JITTargetMachine(optlevel=optlevel)
    LLVM.asm_verbosity!(tm[], true)

    jit[] = OrcJIT(tm[]) # takes ownership of tm
    atexit() do
        dispose(jit[])
    end
end

# Define EnzymeTarget
Base.@kwdef struct EnzymeTarget <: AbstractCompilerTarget
end
GPUCompiler.llvm_triple(::EnzymeTarget) = Sys.MACHINE

# GPUCompiler.llvm_datalayout(::EnzymeTarget) =  nothing

function GPUCompiler.llvm_machine(::EnzymeTarget)
    return tm[]
end

module Runtime
    # the runtime library
    signal_exception() = return
    malloc(sz) = ccall("extern malloc", llvmcall, Csize_t, (Csize_t,), sz)
    report_oom(sz) = return
    report_exception(ex) = return
    report_exception_name(ex) = return
    report_exception_frame(idx, func, file, line) = return
end

abstract type AbstractEnzymeCompilerParams <: AbstractCompilerParams end
struct EnzymeCompilerParams <: AbstractEnzymeCompilerParams
    adjoint::FunctionSpec
    split::Bool
    run_enzyme::Bool
end

struct PrimalCompilerParams <: AbstractEnzymeCompilerParams
end

## job

# TODO: We shouldn't blanket opt-out
GPUCompiler.check_invocation(job::CompilerJob{EnzymeTarget}, entry::LLVM.Function) = nothing

GPUCompiler.runtime_module(::CompilerJob{<:Any,<:AbstractEnzymeCompilerParams}) = Runtime
# GPUCompiler.isintrinsic(::CompilerJob{EnzymeTarget}, fn::String) = true
# GPUCompiler.can_throw(::CompilerJob{EnzymeTarget}) = true

# TODO: encode debug build or not in the compiler job
#       https://github.com/JuliaGPU/CUDAnative.jl/issues/368
GPUCompiler.runtime_slug(job::CompilerJob{EnzymeTarget}) = "enzyme"

include("compiler/optimize.jl")

"""
Create the `FunctionSpec` pair, and lookup the primal return type.
"""
@inline function fspec(f::F, tt::TT) where {F, TT}
    # Entry for the cache look-up
    adjoint = FunctionSpec(f, tt, #=kernel=# false, #=name=# nothing)

    # primal function. Inferred here to get return type
    _tt = (tt.parameters...,)

    primal_tt = Tuple{map(eltype, _tt)...}
    primal = FunctionSpec(f, primal_tt, #=kernel=# false, #=name=# nothing)

    return primal, adjoint
end

##
# Enzyme compiler step
##

function annotate!(mod)
    inactive = LLVM.StringAttribute("enzyme_inactive", "", context(mod))
    fns = functions(mod)
    for inactivefn in ["jl_gc_queue_root", "gpu_report_exception", "gpu_signal_exception", "julia.ptls_states", "julia.write_barrier", "julia.typeof", "jl_box_int64", "jl_subtype"]
        if haskey(fns, inactivefn)
            fn = fns[inactivefn]
            push!(function_attributes(fn), inactive)
        end
    end

    if haskey(fns, "julia.ptls_states")
        fn = fns["julia.ptls_states"]
        push!(function_attributes(fn), LLVM.EnumAttribute("readnone", 0, context(mod)))
    end

    for boxfn in ["jl_box_int64"]
        if haskey(fns, boxfn)
            fn = fns[boxfn]
            push!(return_attributes(fn), LLVM.EnumAttribute("noalias", 0, context(mod)))
            push!(function_attributes(fn), LLVM.EnumAttribute("inaccessiblememonly", 0, context(mod)))
        end
    end

end

function enzyme!(job, mod, primalf, adjoint, split, parallel)
    primal = job.source
    rt = Core.Compiler.return_type(primal.f, primal.tt)
    ctx     = context(mod)
    dl      = string(LLVM.datalayout(mod))

    tt = [adjoint.tt.parameters...,]

    if rt === Union{}
        error("return type is Union{}, giving up.")
    end

    args_activity     = API.CDIFFE_TYPE[]
    uncacheable_args  = Bool[]
    args_typeInfo     = TypeTree[]
    args_known_values = API.IntList[]

    ctx = LLVM.context(mod)
    if !GPUCompiler.isghosttype(typeof(adjoint.f)) && !Core.Compiler.isconstType(typeof(adjoint.f))
        push!(args_activity, API.DFT_CONSTANT)
        typeTree = typetree(typeof(adjoint.f), ctx, dl)
        push!(args_typeInfo, typeTree)
        if split
            push!(uncacheable_args, true)
        else
            push!(uncacheable_args, false)
        end
        push!(args_known_values, API.IntList())
    end

    for T in tt
        source_typ = eltype(T)
        if GPUCompiler.isghosttype(source_typ) || Core.Compiler.isconstType(source_typ)
            if !(T <: Const)
                error("Type of ghost or constant is marked as differentiable.")
            end
            continue
        end
        isboxed = GPUCompiler.deserves_argbox(source_typ)

        if T <: Const
            push!(args_activity, API.DFT_CONSTANT)
        elseif T <: Active

            if isboxed
                push!(args_activity, API.DFT_DUP_ARG)
            else
                push!(args_activity, API.DFT_OUT_DIFF)
            end
        elseif  T <: Duplicated
            push!(args_activity, API.DFT_DUP_ARG)
        elseif T <: DuplicatedNoNeed
            push!(args_activity, API.DFT_DUP_NONEED)
        else
            @assert("illegal annotation type")
        end
        T = source_typ
        if isboxed
            T = Ptr{T}
        end
        typeTree = typetree(T, ctx, dl)
        push!(args_typeInfo, typeTree)
        if split
            push!(uncacheable_args, true)
        else
            push!(uncacheable_args, false)
        end
        push!(args_known_values, API.IntList())
    end

    # TODO ABI returned
    # The return of createprimal and gradient has this ABI
    #  It returns a struct containing the following values
    #     If requested, the original return value of the function
    #     If requested, the shadow return value of the function
    #     For each active (non duplicated) argument
    #       The adjoint of that argument
    if rt <: Integer || rt <: DataType
        retType = API.DFT_CONSTANT
    elseif rt <: AbstractFloat
        retType = API.DFT_OUT_DIFF
    elseif rt <: Complex{<:AbstractFloat}
        retType = API.DFT_OUT_DIFF
    elseif GPUCompiler.isghosttype(rt) || Core.Compiler.isconstType(rt)
        retType = API.DFT_CONSTANT
    else
        error("Unhandled return type $rt")
    end

    TA = TypeAnalysis(triple(mod))
    logic = Logic()

    if GPUCompiler.isghosttype(rt)|| Core.Compiler.isconstType(rt)
        retTT = typetree(Nothing, ctx, dl)
    else
        retTT = typetree(rt, ctx, dl)
    end

    typeInfo = FnTypeInfo(retTT, args_typeInfo, args_known_values)

    if split
        augmented = API.EnzymeCreateAugmentedPrimal(
            logic, primalf, retType, args_activity, TA, #=returnUsed=# true,
            typeInfo, uncacheable_args, #=forceAnonymousTape=# false, #=atomicAdd=# parallel, #=postOpt=# false)

        # 2. get new_primalf
        augmented_primalf = LLVM.Function(API.EnzymeExtractFunctionFromAugmentation(augmented))

        # TODOs:
        # 1. Handle mutable or !pointerfree arguments by introducing caching
        #     + specifically by setting uncacheable_args[i] = true
        # 2. Forward tape from augmented primalf to adjoint (as last arg)
        # 3. Make creation of augumented primalf vs joint forward and reverse optional

        tape = API.EnzymeExtractTapeTypeFromAugmentation(augmented)
        data = Array{Int64, 1}(undef, 3)
        existed = Array{UInt8, 1}(undef, 3)

        API.EnzymeExtractReturnInfo(augmented, data, existed)

        adjointf = LLVM.Function(API.EnzymeCreatePrimalAndGradient(
            logic, primalf, retType, args_activity, TA,
            #=returnValue=#false, #=dretUsed=#false, #=topLevel=#false,
            #=additionalArg=#tape, typeInfo,
            uncacheable_args, augmented, #=atomicAdd=# parallel, #=postOpt=#false))
    else
        adjointf = LLVM.Function(API.EnzymeCreatePrimalAndGradient(
            logic, primalf, retType, args_activity, TA,
            #=returnValue=#false, #=dretUsed=#false, #=topLevel=#true,
            #=additionalArg=#C_NULL, typeInfo,
            # uncacheable_args, #=augmented=#C_NULL, #=atomicAdd=# parallel, #=postOpt=#false))
            uncacheable_args, #=augmented=#C_NULL, #=atomicAdd=# parallel, #=postOpt=#false))
        augmented_primalf = nothing
    end
    return adjointf, augmented_primalf
end

function Base.in(attr::LLVM.EnumAttribute, iter::LLVM.FunctionAttrSet)
    elems = Vector{LLVM.API.LLVMAttributeRef}(undef, length(iter))
    if length(iter) > 0
      # FIXME: this prevents a nullptr ref in LLVM similar to D26392
      LLVM.API.LLVMGetAttributesAtIndex(iter.f, iter.idx, elems)
    end
    for eattr in elems
        at = Attribute(eattr)
        if isa(at, LLVM.EnumAttribute)
            if kind(at) == kind(attr)
                return true
            end
        end
    end
    return false
end

# Modified from GPUCompiler/src/irgen.jl:365 lower_byval
function lower_convention(@nospecialize(job::CompilerJob), mod::LLVM.Module, entry_f::LLVM.Function)
    ctx = context(mod)
    entry_ft = eltype(llvmtype(entry_f)::LLVM.PointerType)::LLVM.FunctionType

    RT = return_type(entry_ft)
    args = GPUCompiler.classify_arguments(job, entry_f)
    filter!(args) do arg
        arg.cc != GPUCompiler.GHOST
    end

    # generate the wrapper function type & definition
    wrapper_types = LLVM.LLVMType[]
    sret = 0
    if !isempty(parameters(entry_f)) && EnumAttribute("sret") in parameter_attributes(entry_f, 1)
        RT = eltype(llvmtype(first(parameters(entry_f))))
        sret = 1
    end
    for (parm, arg) in zip(collect(parameters(entry_f))[1+sret:end], args)
        typ = if !GPUCompiler.deserves_argbox(arg.typ) && arg.cc == GPUCompiler.BITS_REF
            eltype(arg.codegen.typ)
        else
            llvmtype(parm)
        end
        push!(wrapper_types, typ)
    end
    wrapper_fn = LLVM.name(entry_f)
    LLVM.name!(entry_f, wrapper_fn * ".inner")
    wrapper_ft = LLVM.FunctionType(RT, wrapper_types)
    wrapper_f = LLVM.Function(mod, wrapper_fn, wrapper_ft)

    # emit IR performing the "conversions"
    let builder = Builder(ctx)
        entry = BasicBlock(wrapper_f, "entry", ctx)
        position!(builder, entry)

        wrapper_args = Vector{LLVM.Value}()

        if !isempty(parameters(entry_f)) && EnumAttribute("sret") in parameter_attributes(entry_f, 1)
            sretPtr = alloca!(builder, llvmtype(parameters(wrapper_f)[1]))
            push!(wrapper_args, sretPtr)
        end

        # perform argument conversions
        for (parm, arg) in zip(collect(parameters(entry_f))[1+sret:end], args)
            if !GPUCompiler.deserves_argbox(arg.typ) && arg.cc == GPUCompiler.BITS_REF
                # copy the argument value to a stack slot, and reference it.
                ty = llvmtype(parm)
                ptr = alloca!(builder, eltype(ty))
                if LLVM.addrspace(ty) != 0
                    ptr = addrspacecast!(builder, ptr, ty)
                end
                store!(builder, parameters(wrapper_f)[arg.codegen.i], ptr)
                push!(wrapper_args, ptr)
            else
                push!(wrapper_args, parameters(wrapper_f)[arg.codegen.i])
                for attr in collect(parameter_attributes(entry_f, arg.codegen.i+sret))
                    push!(parameter_attributes(wrapper_f, arg.codegen.i), attr)
                end
            end
        end
        res = call!(builder, entry_f, wrapper_args)

        if sret == 1
            ret!(builder, load!(builder, sretPtr))
        elseif return_type(entry_ft) == LLVM.VoidType(ctx)
            ret!(builder)
        else
            ret!(builder, res)
        end

        dispose(builder)
    end

    # early-inline the original entry function into the wrapper
    push!(function_attributes(entry_f), EnumAttribute("alwaysinline", 0, ctx))
    linkage!(entry_f, LLVM.API.LLVMInternalLinkage)

    # copy debug info
    sp = LLVM.get_subprogram(entry_f)
    if sp !== nothing
        LLVM.set_subprogram!(wrapper_f, sp)
    end

    GPUCompiler.fixup_metadata!(entry_f)
    ModulePassManager() do pm
        always_inliner!(pm)
        run!(pm, mod)
    end
    ModulePassManager() do pm
        # Kill the temporary staging function
        global_dce!(pm)
        global_optimizer!(pm)
        run!(pm, mod)
    end
    if haskey(globals(mod), "llvm.used")
        unsafe_delete!(mod, globals(mod)["llvm.used"])
        for u in user.(collect(uses(entry_f)))
            if isa(u, LLVM.GlobalVariable) && endswith(LLVM.name(u), "_slot") && startswith(LLVM.name(u), "julia")
                unsafe_delete!(mod, u)
            end
        end
    end
    return wrapper_f
end

function GPUCompiler.codegen(output::Symbol, job::CompilerJob{<:EnzymeTarget};
                 libraries::Bool=true, deferred_codegen::Bool=true, optimize::Bool=true,
                 strip::Bool=false, validate::Bool=true, only_entry::Bool=false, parent_job::Union{Nothing, CompilerJob} = nothing)
    params  = job.params
    split   = params.split
    adjoint = params.adjoint
    primal  = job.source

    if parent_job === nothing
        primal_target = GPUCompiler.NativeCompilerTarget()
        primal_params = Compiler.PrimalCompilerParams()
        primal_job    = CompilerJob(primal_target, primal, primal_params)
    else
        primal_job = similar(parent_job, job.source)
    end
    mod, meta = GPUCompiler.codegen(:llvm, primal_job, optimize=false, validate=false, parent_job=parent_job)
    primalf = meta.entry
    check_ir(job, mod)

    ctx = context(mod)
    custom = []
    must_wrap = false

    # Julia function to LLVM stem and arity
    known_ops = Dict(
        Base.sin => (:sin, 1),
        Base.cos => (:cos, 1),
        Base.tan => (:tan, 1),
        Base.exp => (:exp, 1),
        Base.log => (:log, 1),
        Base.asin => (:asin, 1),
        Base.tanh => (:tanh, 1)
    )
    for (mi, k) in meta.compiled
        meth = mi.def
        name = meth.name
        jlmod  = meth.module

        Base.isbindingresolved(jlmod, name) && isdefined(jlmod, name) || continue
        func = getfield(jlmod, name)
        func ∈ keys(known_ops) || continue

        sparam_vals = mi.sparam_vals
        name, arity = known_ops[func]

        length(sparam_vals) == arity || continue
        T = first(sparam_vals)
        T ∈ (Float32, Float64) && all(==(T), sparam_vals) || continue
        name = string(name)
        name = T == Float32 ? name*"f" : name

        llvmfn = functions(mod)[k.specfunc]
        push!(custom, llvmfn)

        attributes = function_attributes(llvmfn)
        push!(attributes, EnumAttribute("noinline", 0, ctx))
        push!(attributes, StringAttribute("enzyme_math", name, ctx))

        # Need to wrap the code when outermost
        must_wrap |= llvmfn == primalf
    end

    if must_wrap
        llvmfn = primalf
        FT = eltype(llvmtype(llvmfn)::LLVM.PointerType)::LLVM.FunctionType

        wrapper_f = LLVM.Function(mod, LLVM.name(llvmfn)*"wrap", FT)

        let builder = Builder(ctx)
            entry = BasicBlock(wrapper_f, "entry", ctx)
            position!(builder, entry)

            res = call!(builder, llvmfn, collect(parameters(wrapper_f)))

            if return_type(FT) == LLVM.VoidType(ctx)
                ret!(builder)
            else
                ret!(builder, res)
            end

            dispose(builder)
        end
        primalf = wrapper_f
    end

    if primal_job.target isa GPUCompiler.NativeCompilerTarget
        target_machine = tm[]
    else
        target_machine = GPUCompiler.llvm_machine(primal_job.target)
    end

    parallel = false
    process_module = false
    if parent_job !== nothing
        if parent_job.target isa GPUCompiler.PTXCompilerTarget ||
           parent_job.target isa GPUCompiler.GCNCompilerTarget
            parallel = true
        end
        if parent_job.target isa GPUCompiler.GCNCompilerTarget
           process_module = true
        end
    end

    primalf = lower_convention(job, mod, primalf)
    flush(stderr)
    flush(stdout)

    # annotate
    annotate!(mod)

    # Run early pipeline
    optimize!(mod, target_machine)

    if process_module
        GPUCompiler.optimize_module!(parent_job, mod)
    end

    if params.run_enzyme
        # Generate the adjoint
        adjointf, augmented_primalf = enzyme!(job, mod, primalf, adjoint, split, parallel)
    else
        adjointf = primalf
        augmented_primalf = nothing
    end

    for f in custom
        iter = function_attributes(f)
        elems = Vector{LLVM.API.LLVMAttributeRef}(undef, length(iter))
        LLVM.API.LLVMGetAttributesAtIndex(iter.f, iter.idx, elems)
        for eattr in elems
            at = Attribute(eattr)
            if isa(at, LLVM.EnumAttribute)
                if kind(at) == "noinline"
                    delete!(iter, at)
                end
            end
        end
    end

    linkage!(adjointf, LLVM.API.LLVMExternalLinkage)
    adjointf_name = name(adjointf)

    if augmented_primalf !== nothing
        linkage!(augmented_primalf, LLVM.API.LLVMExternalLinkage)
        augmented_primalf_name = name(augmented_primalf)
    end

    if parent_job !== nothing
        post_optimze!(mod, target_machine)
    end

    adjointf = functions(mod)[adjointf_name]
    # API.EnzymeRemoveTrivialAtomicIncrements(adjointf)

    if process_module
        GPUCompiler.process_module!(parent_job, mod)
    end

    adjointf = functions(mod)[adjointf_name]
    push!(function_attributes(adjointf), EnumAttribute("alwaysinline", 0, context(mod)))
    if augmented_primalf === nothing
        return mod, adjointf
    else
        augmented_primalf = functions(mod)[augmented_primalf_name]
        return mod, (adjointf, augmented_primalf)
    end
end

##
# Thunk
##

struct CombinedAdjointThunk{F, RT, TT}
    fn::F
    # primal::Ptr{Cvoid}
    adjoint::Ptr{Cvoid}
end

@inline (thunk::CombinedAdjointThunk{F, RT, TT})(args...) where {F, RT, TT} =
   enzyme_call(thunk.adjoint, thunk.fn, TT, RT, args...)

@generated function enzyme_call(fptr::Ptr{Cvoid}, f::F, tt::Type{T}, rt::Type{RT}, args::Vararg{Any, N}) where {F, T, RT, N}
    argtt    = tt.parameters[1]
    rettype  = rt.parameters[1]
    argtypes = DataType[argtt.parameters...]
    argexprs = Union{Expr, Symbol}[:(args[$i]) for i in 1:N]
    @assert length(argtypes) == length(argexprs)

    types = DataType[]

    if rettype === Union{}
        error("return type is Union{}, giving up.")
    end

    LLVM.Interop.JuliaContext() do ctx
        T_void = convert(LLVMType, Nothing, ctx)
        ptr8 = LLVM.PointerType(LLVM.IntType(8, ctx))
        T_jlvalue = LLVM.StructType(LLVMType[], ctx)
        T_prjlvalue = LLVM.PointerType(T_jlvalue, #= AddressSpace::Tracked =# 10)

        # Create Enzyme calling convention
        T_wrapperargs = LLVMType[] # Arguments of the wrapper
        T_EnzymeSRet = LLVMType[] # Struct returns of Active variables in the enzyme call
                                  # Equal to all Active vars passed by value
        T_JuliaSRet = LLVMType[]  # Struct return of all Active variables (includes all of T_EnzymeSRet)
        sret_types  = DataType[]  # Julia types of all Active variables
        inputexprs = Union{Expr, Symbol}[]
  # By ref values we create and need to preserve
        ccexprs = Union{Expr, Symbol}[] # The expressions passed to the `llvmcall`

        if !GPUCompiler.isghosttype(F) && !Core.Compiler.isconstType(F)
            isboxed = GPUCompiler.deserves_argbox(F)
            llvmT = isboxed ? T_prjlvalue : convert(LLVMType, F, ctx)
            argexpr = :(f)
            if isboxed
                push!(types, Any)
            else
                push!(types, F)
            end

            push!(ccexprs, argexpr)
            push!(T_wrapperargs, llvmT)
        end

        for (i, T) in enumerate(argtypes)
            source_typ = eltype(T)
            if GPUCompiler.isghosttype(source_typ) || Core.Compiler.isconstType(source_typ)
                @assert T <: Const
                continue
            end
            expr = argexprs[i]

            isboxed = GPUCompiler.deserves_argbox(source_typ)
            llvmT = isboxed ? T_prjlvalue : convert(LLVMType, source_typ, ctx)
            argexpr = Expr(:., expr, QuoteNode(:val))
            if isboxed
                push!(types, Any)
            else
                push!(types, source_typ)
            end

            push!(ccexprs, argexpr)
            push!(T_wrapperargs, llvmT)

            T <: Const && continue

            if T <: Active
                # Use deserves_argbox??
                llvmT = convert(LLVMType, source_typ, ctx)
                push!(sret_types, source_typ)
                push!(T_JuliaSRet, llvmT)
                if !isboxed # XXX: Not consistent
                    push!(T_EnzymeSRet, llvmT)
                end
            elseif T <: Duplicated || T <: DuplicatedNoNeed
                argexpr =  Expr(:., expr, QuoteNode(:dval))
                if isboxed
                    push!(types, Any)
                else
                    push!(types, source_typ)
                end
                push!(ccexprs, argexpr)
                push!(T_wrapperargs, llvmT)
            else
                error("calling convention should be annotated, got $T")
            end
        end

        # API.DFT_OUT_DIFF
        if rettype <: AbstractFloat || rettype <: Complex{<:AbstractFloat}
            push!(types, rettype)
            push!(T_wrapperargs, convert(LLVMType, rettype, ctx))
            push!(ccexprs, :(one($rettype)))
        end
        # XXX: What if not `Nothing`/`Missing` what if struct or array or...

        if !isempty(T_EnzymeSRet)
            ret = LLVM.StructType(T_EnzymeSRet)
        else
            ret = T_void
        end

        # pointer to call
        pushfirst!(T_wrapperargs, convert(LLVMType, Int, ctx))

        # sret argument
        if !isempty(sret_types)
            pushfirst!(T_wrapperargs, convert(LLVMType, Int, ctx))
        end

        llvm_f, _ = LLVM.Interop.create_function(T_void, T_wrapperargs)
        mod = LLVM.parent(llvm_f)
        dl = datalayout(mod)

        params = [parameters(llvm_f)...]
        target =  !isempty(sret_types) ? 2 : 1

        intrinsic_typ = LLVM.FunctionType(T_void, [ptr8, LLVM.IntType(8, ctx), LLVM.IntType(64, ctx), LLVM.IntType(1, ctx)])
        memsetIntr = LLVM.Function(mod, "llvm.memset.p0i8.i64", intrinsic_typ)
        LLVM.Builder(ctx) do builder
            entry = BasicBlock(llvm_f, "entry", ctx)
            position!(builder, entry)

            realparms = LLVM.Value[]
            i = target+1

            if !isempty(T_JuliaSRet)
                sret = inttoptr!(builder, params[1], LLVM.PointerType(LLVM.StructType(T_JuliaSRet)))
            end

            activeNum = 0

            if !GPUCompiler.isghosttype(F) && !Core.Compiler.isconstType(F)
                push!(realparms, params[i])
                i+=1
            end

            for T in argtypes
                T′ = eltype(T)

                if GPUCompiler.isghosttype(T′) || Core.Compiler.isconstType(T′)
                    continue
                end
                push!(realparms, params[i])
                i+=1
                if T <: Const
                elseif T <: Active
                    isboxed = GPUCompiler.deserves_argbox(T′)
                    if isboxed
                        ptr = gep!(builder, sret, [LLVM.ConstantInt(LLVM.IntType(64, ctx), 0), LLVM.ConstantInt(LLVM.IntType(32, ctx), activeNum)])
                        cst = pointercast!(builder, ptr, ptr8)
                        push!(realparms, ptr)

                        cparms = LLVM.Value[cst,
                        LLVM.ConstantInt(LLVM.IntType(8, ctx), 0),
                        LLVM.ConstantInt(LLVM.IntType(64, ctx), LLVM.storage_size(dl, Base.eltype(LLVM.llvmtype(ptr)) )),
                        LLVM.ConstantInt(LLVM.IntType(1, ctx), 0)]
                        call!(builder, memsetIntr, cparms)
                    end
                    activeNum+=1
                elseif T <: Duplicated || T <: DuplicatedNoNeed
                    push!(realparms, params[i])
                    i+=1
                end
            end

            # Primal Differential Return type
            if rettype <: AbstractFloat || rettype <: Complex{<:AbstractFloat}
                push!(realparms, params[i])
            end

            E_types = LLVM.LLVMType[]
            for p in realparms
                push!(E_types, LLVM.llvmtype(p))
            end
            ft = LLVM.FunctionType(ret, E_types)

            ptr = inttoptr!(builder, params[target], LLVM.PointerType(ft))
            val = call!(builder, ptr, realparms)
            if !isempty(T_JuliaSRet)
                activeNum = 0
                returnNum = 0
                for T in argtypes
                    T′ = eltype(T)
                    isboxed = GPUCompiler.deserves_argbox(T′)
                    if T <: Active
                        if !isboxed
                            eval = extract_value!(builder, val, returnNum)
                            store!(builder, eval, gep!(builder, sret, [LLVM.ConstantInt(LLVM.IntType(64, ctx), 0), LLVM.ConstantInt(LLVM.IntType(32, ctx), activeNum)]))
                            returnNum+=1
                        end
                        activeNum+=1
                    end
                end
            end
            ret!(builder)
        end

        ir = string(mod)
        fn = LLVM.name(llvm_f)

        if !isempty(T_JuliaSRet)
            quote
                Base.@_inline_meta
                sret = Ref{$(Tuple{sret_types...})}()
                GC.@preserve sret begin
                    tptr = Base.unsafe_convert(Ptr{$(Tuple{sret_types...})}, sret)
                    tptr = Base.unsafe_convert(Ptr{Cvoid}, tptr)
                    Base.llvmcall(($ir,$fn), Cvoid,
                        $(Tuple{Ptr{Cvoid}, Ptr{Cvoid}, types...}),
                        tptr, fptr, $(ccexprs...))
                end
                sret[]
            end
        else
            quote
                Base.@_inline_meta
                Base.llvmcall(($ir,$fn), Cvoid,
                    $(Tuple{Ptr{Cvoid}, types...}),
                    fptr, $(ccexprs...))
            end
        end
    end
end

##
# JIT
##

function resolver(name, ctx)
    name = unsafe_string(name)
    ptr = try
        ## Step 0: Should have already resolved it iff it was in the
        ##         same module
        ## Step 1: See if it's something known to the execution enging
        # TODO: Do we need to do this?
        # address(jit[], name)

        ## Step 2: Search the program symbols
        #
        # SearchForAddressOfSymbol expects an unmangled 'C' symbol name.
        # Iff we are on Darwin, strip the leading '_' off.
        @static if Sys.isapple()
            if name[1] == '_'
                name = name[2:end]
            end
        end
        LLVM.API.LLVMSearchForAddressOfSymbol(name)
        ## Step 4: Lookup in libatomic
        # TODO: Do we need to do this?
    catch ex
        @error "Enzyme: Lookup failed" jl_name exception=(ex, Base.catch_backtrace())
        C_NULL
    end
    if ptr === C_NULL
        @show name
        error("Enzyme: Symbol lookup failed. Aborting!")
    end

    return UInt64(reinterpret(UInt, ptr))
end

function _link(job, (mod, adjoint_name, primal_name))
    params = job.params
    adjoint = params.adjoint
    split = params.split

    primal = job.source
    rt = Core.Compiler.return_type(primal.f, primal.tt)

    # Now invoke the JIT
    orc = jit[]
    jitted_mod = compile!(orc, mod, @cfunction(resolver, UInt64, (Cstring, Ptr{Cvoid})))

    adjoint_addr = addressin(orc, jitted_mod, adjoint_name)
    adjoint_ptr  = pointer(adjoint_addr)
    if adjoint_ptr === C_NULL
        throw(GPUCompiler.InternalCompilerError(job, "Failed to compile Enzyme thunk, adjoint not found"))
    end
    if primal_name === nothing
        primal_ptr = C_NULL
    else
        primal_addr = addressin(orc, jitted_mod, primal_name)
        primal_ptr  = pointer(primal_addr)
        if primal_ptr === C_NULL
            throw(GPUCompiler.InternalCompilerError(job, "Failed to compile Enzyme thunk, primal not found"))
        end
    end

    @assert primal_name === nothing
    return CombinedAdjointThunk{typeof(adjoint.f), rt, adjoint.tt}(adjoint.f, #=primal_ptr,=# adjoint_ptr)
end

# actual compilation
function _thunk(job)
    params = job.params

    mod, fns = codegen(:llvm, job, optimize=false)

    if fns isa Tuple
        adjointf, augmented_primalf = fns
    else
        adjointf = fns
        augmented_primalf = nothing
    end

    adjoint_name = name(adjointf)

    if augmented_primalf !== nothing
        primal_name = name(augmented_primalf)
    else
        primal_name = nothing
    end

    # Run post optimization pipeline
    post_optimze!(mod, tm[])

    return (mod, adjoint_name, primal_name)
end

const cache = Dict{UInt, Dict{UInt, Any}}()

function thunk(f::F,tt::TT=Tuple{},::Val{Split}=Val(false)) where {F<:Core.Function, TT<:Type, Split}
    primal, adjoint = fspec(f, tt)
    # We need to use primal as the key, to lookup the right method
    # but need to mixin the hash of the adjoint to avoid cache collisions
    # This is counter-intuitive since we would expect the cache to be split
    # by the primal, but we want the generated code to be invalidated by
    # invalidations of the primal, which is managed by GPUCompiler.
    local_cache = get!(Dict{Int, Any}, cache, hash(adjoint, UInt64(Split)))

    target = Compiler.EnzymeTarget()
    params = Compiler.EnzymeCompilerParams(adjoint, Split, true)
    job    = Compiler.CompilerJob(target, primal, params)

    rt = Core.Compiler.return_type(primal.f, primal.tt)

    GPUCompiler.cached_compilation(local_cache, job, _thunk, _link)::CombinedAdjointThunk{F,rt,tt}
end

import GPUCompiler: deferred_codegen_jobs

mutable struct CallbackContext
    job::Compiler.CompilerJob
    stub::String
    compiled::Bool
end

const outstanding = IdDict{CallbackContext, Nothing}()

# Setup the lazy callback for creating a module
function callback(orc_ref::LLVM.API.LLVMOrcJITStackRef, callback_ctx::Ptr{Cvoid})
    orc = OrcJIT(orc_ref)
    cc = Base.unsafe_pointer_to_objref(callback_ctx)::CallbackContext

    @assert !cc.compiled
    job = cc.job

    thunk = Compiler._link(job, Compiler._thunk(job))
    cc.compiled = true
    delete!(outstanding, cc)

    # 4. Update the stub pointer to point to the recently compiled module
    set_stub!(orc, cc.stub, thunk.adjoint)

    # 5. Return the address of tie implementation, since we are going to call it now
    ptr = thunk.adjoint
    return UInt64(reinterpret(UInt, ptr))
end

@generated function deferred_codegen(::Val{f}, ::Val{tt}) where {f,tt}
    primal, adjoint = fspec(f, tt)
    target = EnzymeTarget()
    params = EnzymeCompilerParams(adjoint, false, true)
    job    = CompilerJob(target, primal, params)

    cc = CallbackContext(job, String(gensym(:trampoline)), false)
    outstanding[cc] = nothing

    c_callback = @cfunction(callback, UInt64, (LLVM.API.LLVMOrcJITStackRef, Ptr{Cvoid}))

    orc = Compiler.jit[]
    initial_addr = callback!(orc, c_callback, pointer_from_objref(cc))
    create_stub!(orc, cc.stub, initial_addr)
    addr = address(orc, cc.stub)
    id = Base.reinterpret(Int, pointer(addr))

    deferred_codegen_jobs[id] = job
    trampoline = reinterpret(Ptr{Cvoid}, id)

    quote
        ccall("extern deferred_codegen", llvmcall, Ptr{Cvoid}, (Ptr{Cvoid},), $trampoline)
    end
end

include("compiler/reflection.jl")
include("compiler/validation.jl")

end


text:
"""
    output_type(f, arg_types...)

Return the output type of the specified function. Tries to be fast where possible.
"""
output_type(f, arg_types...) = Union{Base.return_types(f, arg_types)...}

output_type(::typeof(-), x) = x
# TODO more efficient versions for common exp & log cases

output_type(::typeof(+), x, y) = promote_type(x, y)
output_type(::typeof(-), x, y) = promote_type(x, y)
output_type(::typeof(*), x, y) = promote_type(x, y)

output_type(::typeof(/), ::Type{T}, ::Type{T}) where {T<:AbstractFloat} = T
function output_type(::typeof(/), ::Type{X}, ::Type{Y}) where {X<:Real,Y<:Real}
    T = promote_type(X, Y)
    return output_type(/, T, T)
end
output_type(::typeof(/), ::Type{T}, ::Type{T}) where {T<:Integer} = Float64


text:
mutable struct Track{V} <: State{V}
    value::V
end

Track(; unit, _value, _type, _...) = begin
    U = value(unit)
    V = valuetype(_type, U)
    v = _value
    #V = promote_type(V, typeof(v))
    Track{V}(v)
end

constructortags(::Val{:Track}) = (:unit,)

genvartype(v::VarInfo, ::Val{:Track}; V, _...) = @q Track{$V}

gendefault(v::VarInfo, ::Val{:Track}) = gendefaultvalue(v, init=true)