text:
using covid19abm
const cv = covid19abm
using Test

const infectiousstates = (cv.LAT, cv.MILD, cv.MISO, cv.INF, cv.IISO, cv.HOS, cv.ICU)

@testset "parameters" begin
    ip = ModelParameters()
    mod_ip = cv.p ## get the module parameters
    ## set random parameters
    ip.β = 99
    ip.prov = :noknownlocation
    ip.τmild = 1 ## days before they self-isolate for mild cases
    ip.fmild = 0.05  ## percent of people practice self-isolation
    ip.fsevere = 0.80 # fixed at 0.80, always within 1 day.
    ip.eldq = 0.0 ## percentage quarantine of 60+ individuals, removes from chain of transmission
    ip.calibration = false
    ip.modeltime = 500
    ip.popsize = 20000
    cv.reset_params(ip)
    for x in propertynames(cv.p)
        @test getfield(ip, x) == getfield(mod_ip, x)
    end
end

@testset "init" begin
    cv.reset_params()
    cv.initialize()

    # test if human array is resized to the correct length
    @test length(cv.humans) == cv.p.popsize

    ## age groups
    inmodel = (1, 2, 3, 4, 5) ## possible age groups in model
    ags = []
    for x in cv.humans
        push!(ags, x.ag)
    end
    @test length(unique(ags)) == length(inmodel) # check if all age groups are sampled

    ## initial human states (check only subset of humans)
    for i in rand(1:100, 100)
        x = cv.humans[i]
        @test x.health == cv.SUS
        @test x.swap == cv.UNDEF
        @test x.sickfrom == cv.UNDEF
        @test x.nextday_meetcnt >= 0
        @test x.tis == 0
        @test x.exp == 999

        @test x.iso == false
        @test x.isovia == :null

        @test minimum(x.dur) != 0 #check if durations are properly set
        @test x.doi == 999
        @test x.tracing == false
        @test x.tracestart == -1
        @test x.traceend == -1
        @test x.tracedxp == 0
    end

    ## insert_infected check if infected person is added in the correct age group
    for ag in inmodel
        cv.initialize() # reset population
        cv.insert_infected(cv.INF, 1, ag) # 1 infected in age group ag
        @test length(findall(x -> x.health == cv.INF && x.ag == ag, cv.humans)) == 1
        cv.insert_infected(cv.REC, 1, ag) # 1 infected in age group ag
        #@test length(findall(x -> x.health == cv.REC, cv.humans)) == 1
    end

    ## check if the initial infected person is NOT IISO
    cv.p.fsevere = 1.0
    cv.initialize() # reset population
    cv.insert_infected(cv.INF, 1, 1) # 1 infected in age group ag
    @test length(findall(x -> x.health == cv.INF && x.swap == cv.REC, cv.humans)) == 1

    ## check if beta values are time dependent or NOT
    cv.p.seasonal = true
    cv.p.β = 1.0
    svalues = cv.td_seasonality()
    cv.init_betas()
    @test sum(svalues) == sum(cv.BETAS)

    cv.p.seasonal = false
    cv.p.β = 1.0
    cv.init_betas()
    @test cv.p.modeltime == sum(cv.BETAS)
end

@testset "transitions" begin
    cv.reset_params()
    cv.initialize()

    ## check if time in state is up by one
    cv.time_update()
    for i in rand(1:100, 100)
        x = cv.humans[i]
        @test x.tis == 1
        @test x.exp == 999
        @test x.health == cv.SUS
        @test x.swap == cv.UNDEF ## shouldn't set a swap
    end

    for x in cv.humans
        x.exp = 0 ## this will trigger a swap, but no swap is set
    end
    @test_throws ErrorException("swap expired, but no swap set.") cv.time_update()

    #check if it goes through all the move compartments and see if health/swap changes

    for h in 2:9  ## latent to ded, ignore susceptible
        cv.initialize()
        rh = cv.HEALTH(h)
        randhumans = rand(1:100, 100)
        for i in randhumans
            x = cv.humans[i]
            x.swap = rh
            x.exp = 0 ## to force the swap
        end
        cv.time_update() ## since tis > exp
        for i in randhumans
            x = cv.humans[i]
            @test x.health == rh
            if rh ∈ infectiousstates  ## for all states, except rec/ded there should be a swap
                @test x.swap != cv.UNDEF
            end
        end
    end

    # check latent function
    cv.initialize()
    x = cv.humans[1]
    x.ag = 1 ## move to the first age group manually.
    myp = cv.ModelParameters()
    cv.reset_params(myp)
    # check basic variables
    cv.move_to_latent(x)
    @test x.swap ∈ (cv.ASYMP, cv.PRE)
    @test x.iso == false
    @test x.doi == 0
    @test x.tis == 0
    @test x.exp == x.dur[1]

    # check split between asymp/pre
    # split to asymptomatic is now fixed in the model... not a parameter, see old commit for test code

    # CHECKING PRE
    cv.initialize()
    x = cv.humans[1]
    cv.reset_params(myp)
    cv.move_to_pre(x)
    @test x.health == cv.PRE
    @test x.swap ∈ (cv.MILD, cv.INF)

    ## check if individual moves through mild, miso through fmild, tmild parameters
    cv.initialize()
    x = cv.humans[1]
    x.ag = 1 ## move to the first age group manually.
    x.iso = false ## turn this off so we can test the effect of fmild, tmild
    myp.fmild = 0.0
    cv.reset_params(myp)
    cv.move_to_mild(x)
    @test x.swap == cv.REC

    myp.fmild = 1.0
    myp.τmild = 1
    cv.reset_params(myp)
    cv.move_to_mild(x)
    @test x.swap == cv.MISO
    @test x.exp == 1

    cv.time_update()
    @test x.health == cv.MISO
    @test x.swap == cv.REC

    # check if x.iso and x.isovia property are set correctly.
    # 1) check if they are isolated through quarantine
    myp = cv.ModelParameters()
    myp.eldq = 1.0  ## turn on eldq
    myp.eldqag = 4
    cv.reset_params(myp)
    cv.initialize()
    for x in cv.humans
        if x.ag  == myp.eldqag
            @test x.iso == true
            @test x.isovia == :qu
        else
            @test x.iso == false
            @test x.isovia == :null
        end
    end
    # 2) check if they are isolated through presymptomatic capture
    # no need to test for mild/inf movement since it's a simple assignment in code
    # more important to check through `main` when all the function dynamics are happening.
    myp = cv.ModelParameters()
    myp.eldq = 0.0  ## turn off eldq
    myp.fsevere = 0.0 ## turn on fsevere
    myp.fmild = 0.0
    myp.fpreiso = 1.0
    cv.reset_params(myp)
    cv.initialize()
    for x in cv.humans
        cv.move_to_pre(x)
        @test x.iso == true && x.isovia == :pi
    end
    # testing contact tracing isolation below

    cv.initialize()
    x = cv.humans[1]
    x.iso = true
    for i = 1:100 # stress test the function
        @test cv.get_nextday_counts(x) <= 3
    end

    x.health = cv.DED
    for i = 1:100 # stress test the function
        @test cv.get_nextday_counts(x) == 0
    end
end

@testset "tranmission" begin
    cv.reset_params()
    cv.initialize()
    grps = cv.get_ag_dist()

    # get an array of everyones meet counts
    ev_meet_cnts = [x.nextday_meetcnt for x in cv.humans]

    # since beta = 0 default, and everyone sus so no inf to loop over
    tm, ti = cv.dyntrans(1, grps)
    @test ti == 0
    @test tm == 0

    # insert some infectious
    cv.insert_infected(cv.PRE, 1, 1) # 1 infected in age group ag
    tm, ti = cv.dyntrans(1, grps)
    @test ti == 0 # still zero cuz beta = 0
    #@test tm > 0 # count still actually be zero but very unlikely.
    updated_ev_meet_cnts = [x.nextday_meetcnt for x in cv.humans]
    lc_by_one = [x.nextday_meetcnt for x in cv.humans]
    @test tm == sum(ev_meet_cnts - lc_by_one) # test algo. ev_meet_cnts is the original counts. lc_by_one should be counts -1 (except that single infectious person could meet the same person twice)
    # but now some contacts should be -1

    # now change beta
    cv.reset_params(cv.ModelParameters(β = 1.0)) # to account for relative transmission
    tm, ti = cv.dyntrans(1, grps)
    #@test ti > 0 ## actually may still be zero because of stochasticity but very unliekly

    cv.reset_params(cv.ModelParameters(β = 1.0, seasonal = false, frelasymp=0.11))
    @test cv._get_betavalue(1, cv.ASYMP) == 1.0*0.11

    cv.reset_params(cv.ModelParameters(β = 1.0, seasonal = false, frelasymp=0.11))
    @test cv._get_betavalue(1, cv.MILD) == 1.0*0.44

    cv.reset_params(cv.ModelParameters(β = 1.0, seasonal = false, frelasymp=0.11))
    @test cv._get_betavalue(1, cv.INF) == 1.0*0.89
end

@testset "calibration" begin
    # to do, run the model and test total number of infections
    myp = cv.ModelParameters()
    myp.β = 1.0
    myp.prov = :ontario
    myp.calibration = true
    myp.fmild = 0.0
    myp.fsevere = 0.0
    myp.fpreiso = 0.0
    myp.initialinf = 1
    cv.reset_params(myp)
    cv.initialize()
    grps = cv.get_ag_dist()
    cv.insert_infected(cv.PRE, 1, 4)
    # find the single insert_presymptomatic person
    h = findall(x -> x.health == cv.PRE, cv.humans)
    x = cv.humans[h[1]]
    @test length(h) == 1
    # @test x.ag == 4  #not true anymore
    @test x.swap ∈ (cv.ASYMP, cv.MILD, cv.INF) ## always true for calibration
    for i = 1:20 ## run for 20 days
        cv.dyntrans(i, grps)
        cv.time_update()
    end
    @test x.health == cv.REC  ## make sure the initial guy recovered
    @test x.exp == 999
    ## everyone should really be in latent (or susceptible) except the recovered guy
    all = findall(x -> x.health ∈ (cv.PRE, cv.ASYMP, cv.MILD, cv.MISO, cv.INF, cv.IISO, cv.HOS, cv.ICU, cv.DED), cv.humans)
    @test length(all) == 0
    ## to do, make sure everyone stays latent

end

@testset "contact trace" begin
    myp = cv.ModelParameters()
    myp.ctstrat = 1  # turn on contact tracing
    myp.fctcapture = 1.0 # force everyone to get captured.
    myp.fcontactst = 1.0 # force all contacts to be isolated.
    myp.cdaysback = 3 # go back 3 days from time of identification
    myp.mild_props = (0, 0, 0, 0, 0) # turn off mild symptomatic individuals
    cv.reset_params(myp)
    cv.initialize()
    x = cv.humans[1]

    grps = cv.get_ag_dist() # for transmission dynamics function

    # test default values of the parameters.
    @test x.tracestart == -1
    @test x.traceend == -1

    # test whether tracestart and end are set correctly
    # test basic movement between compartments as well
    cv.move_to_latent(x)
    x.swap = cv.ASYMP
    @test x.doi == 0
    @test x.tis == 0
    @test x.exp == x.dur[1]
    @test x.swap == cv.ASYMP

    cv.ct_dynamics(x)  # shouldn't do anything for aymptomatic people
    @test x.tracestart == -1
    @test x.traceend == -1

    x.swap = cv.PRE # force to PRE instead of ASYMP
    @test x.swap == cv.PRE
    cv.ct_dynamics(x) # since ct_dynamics will run when move_to_latent will run in time update func
    @test x.tracestart > 0
    @test x.traceend > 0
    @test x.traceend > x.dur[1] + x.dur[3] ## bigger than since day to identification is randomly sampled
    @test x.traceend <=  x.dur[1] + x.dur[3] + x.dur[4]

    # save the tracestart and end in local variables.
    trace_srt = x.tracestart
    trace_end = x.traceend

    # go through entire inf period of the person
    # first, lets make sure the person goes to
    totalinfperiod = x.dur[1] + x.dur[3] + x.dur[4] # we can do this because we forced x to go to PRE
    @test totalinfperiod > 0

    totaltracedays = 0 # how how many days are true, should match cdaysback
    for i = 1:(trace_end + 2) # plus an 2 days to check if tracing turns off
        #println("tis: $(x.tis), doi: $(x.doi), health: $(x.health),")
        cv.dyntrans(1, grps)
        cv.time_update()
        #println("...dyntrans(); timeupdate() => health: $(x.health), tis: $(x.tis), doi: $(x.doi)")
        @test x.doi == i # should increase with the counter.

        # make sure the tracestart/end only are set once when LAT and doi==0
        if x.health == cv.LAT
            @test x.tracestart == trace_srt
            @test x.traceend == trace_end
            @test sum(x.tracecontacts) == 0
        end

        if x.health == cv.ASYMP || x.health == cv.MILD
            error("should never happen since forcing x to PRE/INF")
        end

        traced_cts = findall(x -> x == true, x.tracecontacts) # eventually length(traced_cts) == 0

        #println("...total traced: $(length(traced_cts))")
        # when doi == tracestart, check if x.tracing is set true
        # and remains true until traceend
        if x.doi >= x.tracestart && x.doi < x.traceend
            @test x.tracing == true
            totaltracedays += 1
            # this has to be zero for now. The contacts get a tracedxp at x.traceend time.
            for i in traced_cts
                y = humans[i]
                @test y.tracedxp == 0
            end
        end
    end
    @test totaltracedays == cv.p.cdaysback # check against the parameter setting

    traced_cts = findall(x -> x.iso == true, humans) # includes the x human as well as all y contacts
    for i in traced_cts
        y = humans[i]
        @test y.iso == true
        @test y.isovia == :ct
        @test y.tracedxp == 11 ## since the for loop runs time_update 3 more times, reducing tracedxp by 3.
    end

    for t = 1:10
        cv.time_update()
        for i in traced_cts
            y = humans[i]
            @test y.tracedxp == 11 - t
            @test y.iso == true
            @test y.isovia == :ct
        end
    end
    # do t = 14 manually
    cv.time_update()
    for i in traced_cts
        y = humans[i]
        @test y.tracedxp == 0
        @test y.iso == false
        @test y.isovia == :null
    end

    ## reset everything ... check if cdaysback is a larger number trace start should be zero
end

@testset "main run" begin
    ## run model with high beta
    myp = cv.ModelParameters()
    myp.β = 0.0525
    myp.prov = :newyork
    ## run empty scenario

    myp.τmild = 0
    myp.fmild = 0.0
    myp.fsevere = 0.0
    myp.eldq = 0.0
    myp.fpre = 1.0
    myp.fpreiso = 0.0
    myp.tpreiso = 0
    myp.fctcapture = 0.0
    cv.runsim(1, myp) # warm up the functions
    println("time with contact tracing off:")
    @time results = cv.main(myp)

    myp.fctcapture = 1.0
    cv.runsim(1, myp)
    println("time with contact tracing on:")
    @time results = cv.main(myp)

#     function string_or_char(s::SubString)
#         length(s) == 1 && return first(s)
#         return String(s)
#     end
#     arr = Union{Char,String}[]
# for i in split_string
#     push!(arr, string_or_char(i))
# end

    # function addday(n)
    #     strtorep = "---" * "\u00B7"
    #     str = ""
    #     for i = 1:n
    #         str =str * strtorep
    #     end
    #     str = str[1:end-1]
    #     str = " " * str * " "
    #     return str
    # end

    # function printascii_epi(x)
    #     fs = ""
    #     fs = fs * "L" * addday(x.dur[1])

    #     fs = fs * "P" * addday(x.dur[3])
    #     fs = fs * "S" * addday(x.dur[4])
    #     return fs
    # end

end


text:
using CuArrays
using CUDAnative
using LinearAlgebra

# ::: STRUCTURE

# GPU
GlobalArrayType = CuArray
# CPU
# GlobalArrayType = Array

mutable struct NetworkNode
    forward::Function
    apply_derivs!::Function
    params::Dict{String, Any}
    derivs::Dict{String, Function}
end

struct Optimizer
    f::Function
    params::Dict{String, Any}
end

# :: HELPER FUNCTIONS
include("helper_functions.jl")

# ::: LAYERS
DenseLayer(input_size::Int, output_size::Int; norm = (f,_)->f, w_init = rand_glorot, b_init = CuArrays.zeros) = begin
    w = w_init(input_size, output_size)
    b = b_init(input_size, output_size)

    p = Dict(
        "w" => w,
        "b" => b
    )

    function forward(x::GlobalArrayType, p::Dict)::Tuple{GlobalArrayType, Dict{String, Function}}
        d = Dict("dydx" => (Δ) -> p["w"]' * Δ,
            "dydw" => norm((Δ) -> Δ * x', Dict("params" => p["w"])),
            "dydb" => norm((Δ) -> Δ, Dict("params" => p["b"])))

        return (p["w"] * x .+ p["b"], d)
    end

    function apply_derivs!(n, Δ, opt)::Nothing
        n.params["w"] .-= apply!(opt, Δ, n.derivs["dydw"])
        n.params["b"] .-= apply!(opt, Δ, n.derivs["dydb"])
        return nothing
    end

    NetworkNode(forward, apply_derivs!, p, Δp)
end
ActivationLayer(input_size::Int, f::Function, df::Function) = begin
    p = Dict()
    Δp = Dict("dydx" => (Δ) -> Δ)

    function forward(x::GlobalArrayType, p::Dict)::Tuple{GlobalArrayType, Dict{String, Function}}
        return (f(x),
            Dict("dydx" => (Δ) -> df(x) * Δ)
        )
    end

    function apply_derivs!(_, _, _)::Nothing
        return nothing
    end

    NetworkNode(forward, apply_derivs!, p, Dict())
end
BatchNormLayer(in_out_size, batch_size; α_init = rand_glorot, β_init = CuArrays.zeros) = begin
    p = Dict(
        "α" => α_init(in_out_size, batch_size)
        "β" => β_init(in_out_size, batch_size)
        "ϵ" => 0.00001 # to prevent devision by zero
    )

    function forward(x::GlobalArrayType, p::Dict)::Tuple{GlobalArrayType, Dict{String, Function}}
        μ = sum(x) / length(x)
        σ = sum((x .- μ) .^ 2) / length(x)

        bn = (x .- μ) ./ sqrt(σ ^ 2 + p["ϵ"])
        y = bn .* p["α"] .+ p["β"]

        return (p["w"] * x .+ p["b"],
                Dict("dydx" => (Δ) -> Δ .* p["α"] .* sqrt(σ + p["ϵ"]),
                    "dydα" => (Δ) -> Δ .* bn,
                    "dydβ" => (Δ) -> Δ)
                )
    end

    function apply_derivs!(n, Δ, opt)::Nothing
        n.params["α"] .-= apply!(opt, Δ, n.derivs["dydα"])
        n.params["β"] .-= apply!(opt, Δ, n.derivs["dydβ"])
        return nothing
    end

    NetworkNode(forward, apply_derivs!, p, Dict())
end
ConvolutionalLayer(input_size, out_channels, k_size; padding=(0,0,0), k_stride=(1,1,1), k_dilation=(1,1,1); w_init = rand_glorot, b_init = CuArrays.zeros) = begin

    function forward(x::GlobalArrayType, p::Dict)
        x = pad(x, P)

    end

    function apply_derivs!(Δ)
    end

    NetworkNode(forward, apply_derivs!, p, Dict())
end

RecurrentLayer(input_size, hidden_size, max_buffer_length) = begin
    p = Dict(
        "xh_w"
        "xh_b"
        "hh_w"
        "hh_b"
        "hy_w"
        "hy_b"
    )

    function forward(x::GlobalArrayType, p::Dict)

        pop!(a)
        pushfirst!(a, 4)
    end

    function apply_derivs!(Δ)
    end

    NetworkNode(forward, apply_derivs!, p, Dict())
end

# in order (height, width, channels, batch)
is = (11, 13, 20, 20)
P = (2, 2, 0, 0)
K = (3, 3, 3, 20)
S = (2,2,2,1)
D = (2,2,2,1)

# conv 1: get estimated input size
is_p = is .+ P

# conv 2: get the valid output size estimate
total_KS = K .+ (K .- 1) .* (D .- 1)
max_ind = is_p .- total_KS .+ 1
available_indecies = ((max_ind .* (max_ind .> 0)) ./ S) .|> ceil .|> Int
all_alpha_i = get_valid_indecies(is_p, available_indecies, S)

# conv 3: get flat kernel indecies from any 1D location of 4D tensor from arbitrary location
all_beta_i = flat_kernel_indecies(is_p, K, D)

# conv 4: get all indecies
all_x_indecies = [i .+ all_beta_i for i in all_alpha_i] |> Iterators.flatten |> collect
final_x_ind = reshape(all_x_indecies, (prod(K[1:3]), K[4], length(all_alpha_i)))

function index_assigne!(out, X, I)
    I_ind = (blockIdx().x - 1) * blockDim().x + threadIdx().x

    i = I[I_ind]
    out[I_ind] = X[i]

    return nothing
end

x = CuArray(rand(is_p...))
out = CuArray(similar(final_x_ind))

@cuda threads=length(all_beta_i) blocks=length(all_alpha_i) index_assigne!(out, x, final_x_ind |> CuArray)

# TODO:
# index_assigne => tmp_out
#   => reshape(prod(K[1:3]), K[end], number_of_kernels)
#   => permutedims([prod(K[1:3], #_kernels, K[end])])
#   => multiply pointwise with weights [#_output_channels, prod(K[1:3]), #_kernels, K[end]]
#   => sum(dims=2)
#   => reshape into [#_output_channels, #_kernels, K[end]]
#   => add biases [#_output_channels, #_kernels, K[end]]
#   => reshape into [#_output_size..., #_output_channels, K[end]]

# ::: NORMS
no_norm(df::Function, _) = df
L2_norm(df::Function, p::Dict; λ = 0.005) = (x) -> df(x) .+ (λ .* p["params"])

# ::: OPTIMIZERS
SGD(η) = Optimizer((d, df, p) -> df(d) .* p["lr"], Dict("lr" => η))

# ::: ACTIVATION FUNCTIONS
relu(A::GlobalArrayType) = map(a->a * (a>0), A)
Δrelu(A::GlobalArrayType) =  GlobalArrayType{Float32}(I, length(A), length(A)) .* map(a->(a>0), A)

sigmoid(A::GlobalArrayType) = 1 ./ (1 .+ exp.(.-A))
Δsigmoid(A::GlobalArrayType) = GlobalArrayType{Float32}(I, length(A), length(A)) .* (sigmoid(A) .* (1 .- sigmoid(A)))

diminish(A::GlobalArrayType) = (2 .* A) ./ (1 .+ A .^ 2)
Δdiminish(A::GlobalArrayType) = GlobalArrayType{Float32}(I, length(A), length(A)) .* (2 .* ((1 .- A .^ 2) ./ (1 .+ 2 .* A .^ 2 .+ A .^ 4)))
diminish(A) = (2 .* A) ./ (1 .+ A .^ 2)

softmax(A::GlobalArrayType) = exp.(A) ./ sum(exp.(A))
function d_softmax(A::GlobalArrayType)
    f1 = (.-softmax(A) * softmax(A)') .* (CuArrays.fill(1, length(A), length(A)) .- GlobalArrayType{Float32}(I, length(A), length(A)))
    f2 = (softmax(A) .* (1 .- softmax(A))) .* GlobalArrayType{Float32}(I, length(A), length(A))
    return f1 .+ f2
end

# ::: LOSS FUNCTIONS
crossentropy(p::Array, q::Array; base=ℯ) = -sum(q .* log.(base, p))
d_crossentropy(p::Array, q::Array; base=ℯ) = .-(q .* (1 ./ (p .* log(base))) .* log.(base, p))

mse(x::GlobalArrayType, y::GlobalArrayType) = sum((x .- y).^2) / length(x)
Δmse(x::GlobalArrayType, y::GlobalArrayType) = 2 .* (x .- y) ./ length(x)

diminish(A::GlobalArrayType) = (2 .* A) ./ (1 .+ A .^ 2)
Δdiminish(A::GlobalArrayType) = GlobalArrayType{Float32}(I, length(A), length(A)) .* (2 .* ((1 .- A .^ 2) ./ (1 .+ 2 .* A .^ 2 .+ A .^ 4)))

ext: jl


text:
### A Pluto.jl notebook ###
# v0.16.0

using Markdown
using InteractiveUtils

# ╔═╡ 03688332-bcc4-4c84-8924-7f1610db077e
# imports
begin
	using Dierckx
	using Markdown
	using Measurements
	using Unitful
	using UnitfulAngles
	using UnitfulAstro
end

# ╔═╡ b1450a24-3120-11ec-2c56-d7b2f96fdf46
md"""
# Star Hoppers

This notebook will follow along with some derivations of [Kratter & Perets (2012)](https://ui.adsabs.harvard.edu/abs/2012ApJ...753...91K).
"""

# ╔═╡ 2dbe3b2b-f33b-46f3-98f3-869c539451db
M_A = (2.063 ± 0.023)u"Msun"

# ╔═╡ 6eef835d-49b3-4ffe-bc3a-dc7d5938fe5b
M_B_WD = (1.018 ± 0.011)u"Msun"

# ╔═╡ c79b9551-4ccc-4ea4-8e41-1c56cd5e4c72
md"white dwarf initial-final mass relation (IFMR) from [Cummings et al. (eqn. 6; 2018)](https://ui.adsabs.harvard.edu/abs/2018ApJ...866...21C)"

# ╔═╡ 6062c398-de6a-457f-a7c5-e8ce13d7bd3e
function MIST_highmass_IFMR(M_WD)
	b = (0.471 ± 0.077)u"Msun"
	m = (0.107 ± 0.016)
	return (M_WD - b) / m
end

# ╔═╡ eb6f04d6-596f-497e-93ab-fdcde0e92d8d
M_B_MS = MIST_highmass_IFMR(M_B_WD)

# ╔═╡ df958bae-32e7-47b6-8e9f-95698151113c
md"""
Adiabatic expansion factor

```math
f = \frac{a_f}{a_i} = \frac{M_\mathrm{MS}}{M_\mathrm{WD}}
```
"""

# ╔═╡ f6d29ca4-b536-49a9-9391-81b44d3e7413
f = M_B_MS / M_B_WD

# ╔═╡ c431232c-b181-483a-95a7-25426d865281

# ╔═╡ 8b14cc60-5373-4ca9-bc90-932e1a729e0e
# semimajor axis from GAIA parallax and Bond et al. 2017
begin
	plx = (0.37449 ± 0.00023)
	a_arc = (7.4957 ± 0.0025)
	a = (a_arc / plx) * u"AU"
end

# ╔═╡ b5d25245-2c5c-41ce-a255-8fb5aa7b4ee8
(M_B_WD + M_A) / (M_B_MS + M_A) |> inv

# ╔═╡ 4c3173e9-3599-4bba-864d-354c369a2a3e
a_initial = a * (M_B_WD + M_A) / (M_B_MS + M_A)

# ╔═╡ 83e74886-9204-49ae-8f65-e9550e442c39
@. a_initial * (1 + [0.59, -0.59])

# ╔═╡ a5527e3e-15f1-4f14-9c2a-f77706dfda3e
md"""
Forbidden region defined by minimum separation required to (1) escape engulfment and (2) escape tidal shredding. The first is defined by the max giant branch (GB) radius, and the second is dependent on the companion mass, but for giant planets we assume some flat percentage extra (here, 20% is based off Veras 2016)
"""

# ╔═╡ 2f6a3dfc-ca25-455e-ad61-df8eaa049e08
max_GB_radius = (5.104 ± 0.075)u"AU"

# ╔═╡ 677ad47b-538a-41e4-844b-8fff1a3d6314
tide_factor = 0.2

# ╔═╡ 537a9087-ddb1-45e8-a872-593f6fad93da
stable_radius(M1, M2, a) = (0.464 - 0.38 * M2 / (M1 + M2)) * a

# ╔═╡ 11d5f746-56c5-49dd-971a-fd90e47a7952
R_P = stable_radius(M_B_MS, M_A, a_initial)

# ╔═╡ 1ac684d7-f0c6-426f-8318-075352c92ca0
R_S = stable_radius(M_A, M_B_MS, a_initial)

# ╔═╡ 88f576cf-696d-4d7f-8123-729bdef2d9a6
area_ratio = R_S^2 / (R_P^2 + R_S^2)

# ╔═╡ 2c15891e-c3bf-4a43-b4b1-4567ad7c5ed9
coll_prob_m1 = Spline1D(
	[0.5, 0.55, 0.6, 0.65, 0.7],
	[0.4, 0.49, 0.63, 0.75, 0.8],
	k = 1,
	bc="extrapolate"
)

# ╔═╡ 28f682d4-f9bd-4a00-a2ec-c6be9ceaeb46
coll_prob_m2 = Spline1D(
	[0.5, 0.55, 0.6, 0.65, 0.7],
	[0.53, 0.5, 0.27, 0.22, 0.17],
	k = 1,
	bc="extrapolate"
)

# ╔═╡ a365cdca-14fb-4889-afde-431f8b292ba2
@uncertain coll_prob_m1(area_ratio)

# ╔═╡ 2a550342-7b8c-4dde-a89f-e7e7e63d1bc1
@uncertain coll_prob_m2(area_ratio)

# ╔═╡ 29b9a599-a519-4b72-afd1-3aa1d872a8cf
min_escape_separation_pre = max_GB_radius * (1 + tide_factor)

# ╔═╡ c61024e3-769b-4968-a1ba-a5509de1848a
min_escape_separation_post = min_escape_separation_pre * f

# ╔═╡ c6de59a9-57b2-40ff-8732-d4a01f11db17
(M_B_MS + M_A) / 2.5u"Msun" * ustrip(u"Msun", M_A)^(3/2) / sqrt(1.711) / (a_initial / 100u"AU")^(3/2)

# ╔═╡ 806d35ba-1efb-44eb-9153-64f6992bfe31
(M_B_WD + M_A) / 2.5u"Msun" * ustrip(u"Msun", M_A)^(3/2) / sqrt(1.711) / (a / 100u"AU")^(3/2)

# ╔═╡ 286abf89-2ffc-42e8-8f89-b8b58ffe507e
begin
	μ = M_B_MS / (M_A + M_B_MS)
	eb = 0.59
	ac_ab = (0.464 ± 0.006) + (-0.38 ± 0.01) * μ +
		(-0.631 ± 0.034) * eb + (0.586 ± 0.061) * eb * μ +
		(0.15 ± 0.041) * eb^2 + (-0.198 ± 0.047) * μ * eb^2
	ac = ac_ab * a
end

# ╔═╡ Cell order:
# ╠═03688332-bcc4-4c84-8924-7f1610db077e
# ╟─b1450a24-3120-11ec-2c56-d7b2f96fdf46
# ╟─2dbe3b2b-f33b-46f3-98f3-869c539451db
# ╟─6eef835d-49b3-4ffe-bc3a-dc7d5938fe5b
# ╟─c79b9551-4ccc-4ea4-8e41-1c56cd5e4c72
# ╠═6062c398-de6a-457f-a7c5-e8ce13d7bd3e
# ╠═eb6f04d6-596f-497e-93ab-fdcde0e92d8d
# ╟─df958bae-32e7-47b6-8e9f-95698151113c
# ╠═f6d29ca4-b536-49a9-9391-81b44d3e7413
# ╠═c431232c-b181-483a-95a7-25426d865281
# ╠═8b14cc60-5373-4ca9-bc90-932e1a729e0e
# ╠═b5d25245-2c5c-41ce-a255-8fb5aa7b4ee8
# ╠═4c3173e9-3599-4bba-864d-354c369a2a3e
# ╠═83e74886-9204-49ae-8f65-e9550e442c39
# ╟─a5527e3e-15f1-4f14-9c2a-f77706dfda3e
# ╟─2f6a3dfc-ca25-455e-ad61-df8eaa049e08
# ╟─677ad47b-538a-41e4-844b-8fff1a3d6314
# ╠═537a9087-ddb1-45e8-a872-593f6fad93da
# ╠═11d5f746-56c5-49dd-971a-fd90e47a7952
# ╠═1ac684d7-f0c6-426f-8318-075352c92ca0
# ╠═88f576cf-696d-4d7f-8123-729bdef2d9a6
# ╠═2c15891e-c3bf-4a43-b4b1-4567ad7c5ed9
# ╠═28f682d4-f9bd-4a00-a2ec-c6be9ceaeb46
# ╠═a365cdca-14fb-4889-afde-431f8b292ba2
# ╠═2a550342-7b8c-4dde-a89f-e7e7e63d1bc1
# ╠═29b9a599-a519-4b72-afd1-3aa1d872a8cf
# ╠═c61024e3-769b-4968-a1ba-a5509de1848a
# ╠═c6de59a9-57b2-40ff-8732-d4a01f11db17
# ╠═806d35ba-1efb-44eb-9153-64f6992bfe31
# ╠═286abf89-2ffc-42e8-8f89-b8b58ffe507e

max_stars_repo_name: mileslucas/sirius-b

max_issues_repo_name: mileslucas/sirius-b

max_forks_repo_name: mileslucas/sirius-b


text:
# Peter Norvig's Spelling Corrector
# based off of the python implementation at http://norvig.com/spell-correct.html

# I've attempted to keep this close to the python version, although in some cases
# the python version is tighter because of some syntax that is not available in julia
# (e.g., guards in comprehensions).
#
# The python version, at the time of original submission, is also 5-6x faster than
# the julia version.
#
# @kmsquire, 2013-11-29
#

Pkg.add("DataStructures")

using DataStructures
using HTTPClient.HTTPC

include("../perfutil.jl")

words(text) = eachmatch(r"[a-z]+", lowercase(text))

function train(features)
    model = DefaultDict(String, Int, 1)
    for f in features
        model[f.match] += 1
    end
    return model
end

const NWORDS = train(words(bytestring(get("http://norvig.com/big.txt").body)))

const alphabet = "abcdefghijklmnopqrstuvwxyz"

function edits1(word::String)
    splits     = [(word[1:i], word[i+1:end]) for i=0:length(word) ]
    deletes    = ["$a$(b[2:end])"               for (a,b) in splits[1:end-1]]
    transposes = ["$a$(b[2])$(b[1])$(b[3:end])" for (a,b) in splits[1:end-2]]
    replaces   = ["$a$c$(b[2:end])"             for (a,b) in splits[1:end-1], c in alphabet]
    inserts    = ["$a$c$b"                      for (a,b) in splits,          c in alphabet]
    return Set([deletes; transposes; replaces[:]; inserts[:]])
end

function known_edits2(word::String)
    xs = Set{String}()
    for e1 in edits1(word)
        for e2 in edits1(e1)
            haskey(NWORDS, e2) && push!(xs, e2)
        end
    end
    xs
end

function known(words)
    xs = Set{String}()
    for word in words
        haskey(NWORDS, word) && push!(xs, word)
    end
    xs
end

function correct(word::String)
    candidates = known([word])
    length(candidates) == 0 && (candidates = known(edits1(word)))
    length(candidates) == 0 && (candidates = known_edits2(word) )
    length(candidates) == 0 && (candidates = Set([word])        )

    maximum(x->(get(NWORDS, x, 0),x), candidates)[2]
end

################ Testing code from here on ################

function spelltest(tests; bias=0, verbose=false)
    n, bad, unknown, start = 0, 0, 0, tic()
    if bias > 0
        for target in keys(tests)
            NWORDS[target] += bias
        end
    end
    for (target,wrongs) in tests
        for wrong in split(wrongs)
            n += 1
            w = correct(wrong)
            if w!=target
                bad += 1
                unknown += !(target in NWORDS)
                if verbose
                    @printf("correct(%s) => %s (%d); expected %s (%d)\n",
                            wrong, w, NWORDS[w], target, NWORDS[target])
                end
            end
        end
    end

    return ["bad"=>bad, "n"=>n, "bias"=>bias, "pct"=>int(100. - 100.*bad/n),
            "unknown"=>unknown, "secs"=>toc() ]
end

const tests1 = [ "access"=> "acess", "accessing"=> "accesing", "accommodation"=>
"accomodation acommodation acomodation", "account"=> "acount", "address"=>
"adress adres", "addressable"=> "addresable", "arranged"=> "aranged arrainged",
"arrangeing"=> "aranging", "arrangement"=> "arragment", "articles"=> "articals",
"aunt"=> "annt anut arnt", "auxiliary"=> "auxillary", "available"=> "avaible",
"awful"=> "awfall afful", "basically"=> "basicaly", "beginning"=> "begining",
"benefit"=> "benifit", "benefits"=> "benifits", "between"=> "beetween", "bicycle"=>
"bicycal bycicle bycycle", "biscuits"=>
"biscits biscutes biscuts bisquits buiscits buiscuts", "built"=> "biult",
"cake"=> "cak", "career"=> "carrer",
"cemetery"=> "cemetary semetary", "centrally"=> "centraly", "certain"=> "cirtain",
"challenges"=> "chalenges chalenges", "chapter"=> "chaper chaphter chaptur",
"choice"=> "choise", "choosing"=> "chosing", "clerical"=> "clearical",
"committee"=> "comittee", "compare"=> "compair", "completely"=> "completly",
"consider"=> "concider", "considerable"=> "conciderable", "contented"=>
"contenpted contende contended contentid", "curtains"=>
"cartains certans courtens cuaritains curtans curtians curtions", "decide"=> "descide", "decided"=>
"descided", "definitely"=> "definately difinately", "definition"=> "defenition",
"definitions"=> "defenitions", "description"=> "discription", "desiccate"=>
"desicate dessicate dessiccate", "diagrammatically"=> "diagrammaticaally",
"different"=> "diffrent", "driven"=> "dirven", "ecstasy"=> "exstacy ecstacy",
"embarrass"=> "embaras embarass", "establishing"=> "astablishing establising",
"experience"=> "experance experiance", "experiences"=> "experances", "extended"=>
"extented", "extremely"=> "extreamly", "fails"=> "failes", "families"=> "familes",
"february"=> "febuary", "further"=> "futher", "gallery"=> "galery gallary gallerry gallrey",
"hierarchal"=> "hierachial", "hierarchy"=> "hierchy", "inconvenient"=>
"inconvienient inconvient inconvinient", "independent"=> "independant independant",
"initial"=> "intial", "initials"=> "inetials inistals initails initals intials",
"juice"=> "guic juce jucie juise juse", "latest"=> "lates latets latiest latist",
"laugh"=> "lagh lauf laught lugh", "level"=> "leval",
"levels"=> "levals", "liaison"=> "liaision liason", "lieu"=> "liew", "literature"=>
"litriture", "loans"=> "lones", "locally"=> "localy", "magnificent"=>
"magnificnet magificent magnifcent magnifecent magnifiscant magnifisent magnificant",
"management"=> "managment", "meant"=> "ment", "minuscule"=> "miniscule",
"minutes"=> "muinets", "monitoring"=> "monitering", "necessary"=>
"neccesary necesary neccesary necassary necassery neccasary", "occurrence"=>
"occurence occurence", "often"=> "ofen offen offten ofton", "opposite"=>
"opisite oppasite oppesite oppisit oppisite opposit oppossite oppossitte", "parallel"=>
"paralel paralell parrallel parralell parrallell", "particular"=> "particulaur",
"perhaps"=> "perhapse", "personnel"=> "personnell", "planned"=> "planed", "poem"=>
"poame", "poems"=> "poims pomes", "poetry"=> "poartry poertry poetre poety powetry",
"position"=> "possition", "possible"=> "possable", "pretend"=>
"pertend protend prtend pritend", "problem"=> "problam proble promblem proplen",
"pronunciation"=> "pronounciation", "purple"=> "perple perpul poarple",
"questionnaire"=> "questionaire", "really"=> "realy relley relly", "receipt"=>
"receit receite reciet recipt", "receive"=> "recieve", "refreshment"=>
"reafreshment refreshmant refresment refressmunt", "remember"=> "rember remeber rememmer rermember",
"remind"=> "remine remined", "scarcely"=> "scarcly scarecly scarely scarsely",
"scissors"=> "scisors sissors", "separate"=> "seperate",
"singular"=> "singulaur", "someone"=> "somone", "sources"=> "sorces", "southern"=>
"southen", "special"=> "speaical specail specal speical", "splendid"=>
"spledid splended splened splended", "standardizing"=> "stanerdizing", "stomach"=>
"stomac stomache stomec stumache", "supersede"=> "supercede superceed", "there"=> "ther",
"totally"=> "totaly", "transferred"=> "transfred", "transportability"=>
"transportibility", "triangular"=> "triangulaur", "understand"=> "undersand undistand",
"unexpected"=> "unexpcted unexpeted unexspected", "unfortunately"=>
"unfortunatly", "unique"=> "uneque", "useful"=> "usefull", "valuable"=> "valubale valuble",
"variable"=> "varable", "variant"=> "vairiant", "various"=> "vairious",
"visited"=> "fisited viseted vistid vistied", "visitors"=> "vistors",
"voluntary"=> "volantry", "voting"=> "voteing", "wanted"=> "wantid wonted",
"whether"=> "wether", "wrote"=> "rote wote"]

@timeit1(spelltest(tests1), "spell", "Peter Norvig's spell corrector")

max_stars_repo_name: porterjamesj/julia

max_issues_repo_name: porterjamesj/julia

max_forks_repo_name: porterjamesj/julia


text:
include("../src/PowerFlowNetworks.jl")
using .PowerFlowNetworks
using SQLite
using ArgParse
using JSON

ArgParse.parse_item(::Type{Vector{String}}, x::AbstractString) = split(x, ",")
ArgParse.parse_item(::Type{Vector{Int}}, x::AbstractString) = parse.(Int, split(x, ","))

function parse_commandline()
    s = ArgParseSettings()
    @add_arg_table s begin
        "--dbpath", "-p"
            help = "Path of the sqlite database file"
            default = "data/PowerFlowNetworks.sqlite"
        "--heuristic"
            help = "Algorithm(s) used for the chordal extension"
            arg_type = Vector{String}
            default = ["molzahn"]
        "--heuristic_switch"
            help = "Iteration(s) when to switch the heuristic"
            arg_type = Vector{Int}
            default = [0]
        "--treshold_name"
            help = "Name of the treshold used to stop the merge"
            arg_type = String
            default = "clique_nv_up"
        "--kwargs_path"
            help = "Julia file containing a dict named `merge_kwargs` containing the keyword arguments of the merge."
            arg_type = String
        "--min_nv"
            help = "Minimum number of vertices a network had to have to be processed."
            arg_type = Int
            default = typemin(Int)
        "--max_nv"
            help = "Maximum number of vertices a network had to have to be processed."
            arg_type = Int
            default = typemax(Int)
    end
    return parse_args(s)
end

function main()
    parsed_args = parse_commandline()
    include("../$(parsed_args["kwargs_path"])")
    db = SQLite.DB(parsed_args["dbpath"])
    if "sliwak" in parsed_args["heuristic"]
        error("Sliwak merge not implemented yet. (I need to get the coefficients of the merge for every instance)")
    end
    merge_decompositions!(db;
                          heuristic=parsed_args["heuristic"],
                          heuristic_switch=parsed_args["heuristic_switch"],
                          treshold_name=parsed_args["treshold_name"],
                          merge_kwargs=merge_kwargs,
                          min_nv=parsed_args["min_nv"],
                          max_nv=parsed_args["max_nv"])
end

main()

max_stars_repo_name: charlyalizadeh/PowerFlowNetworks

max_issues_repo_name: charlyalizadeh/PowerFlowNetworks

max_forks_repo_name: charlyalizadeh/PowerFlowNetworks


text:
"""
    Weighted(B::AbstractBasis, w) = [ x-> w(x) * B[k](x) for k = 1:length(B)]
"""
struct Weighted{B,W} <: AbstractBasis
    basis::B
    weight::W
end

Base.length(B::Weighted) = length(B.basis)

approxpoints(B::Weighted) = approxpoints(B.basis)
function approxtransform(B::Weighted, f)
    basis = B.basis
    x = approxpoints(basis)
    w = B.weight.(x)
    return approxtransform(basis, w.\f)
end

function iterate_basis(B::Weighted, x)
    Bx = B.basis|x
    b,state = IterTools.@ifsomething iterate(Bx)
    w = B.weight(x)
    return w*b, (w,Bx,state)
end
function iterate_basis(B::Weighted, x, (w,Bx,state))
    b,state = IterTools.@ifsomething iterate(Bx,state)
    return w*b, (w,Bx,state)
end

function iterate_basis(B::Weighted, x::MatrixVectorWrapper)
    xx = MatrixVectorWrapper(x.matrix,B.weight(x))
    Bxx = B.basis|xx
    b,state = IterTools.@ifsomething iterate(Bxx)
    return b, (xx,Bxx,state)
end
function iterate_basis(B::Weighted, x::MatrixVectorWrapper, (xx,Bxx,state))
    b,state = IterTools.@ifsomething iterate(Bxx,state)
    return b, (xx,Bxx,state)
end

evaluate_basis(B::Weighted,i,x) = B.weight(x)*B.basis[i](x)
evaluate_basis(B::Weighted,i,x::MatrixVectorWrapper) = B.basis[i](MatrixVectorWrapper(x.matrix,B.weight(x)))

evaltransform(B::Weighted, x, c) = grideval(B.weight,(x,)) .* evaltransform(B.basis,x,c)


text:
###############################################################################
#
#   Rings.jl : Generic rings
#
###############################################################################

function isequal(a::RingElem, b::RingElem)
   return parent(a) == parent(b) && a == b
end

################################################################################
#
#   Promotion system
#
# The promote_rule functions are not extending Base.promote_rule. The Nemo
# promotion system is orthogonal to the built-in julia promotion system. The
# julia system assumes that whenever you have a method signature of the form
# Base.promote_rule(::Type{T}, ::Type{S}) = R, then there is also a
# corresponding Base.convert(::Type{R}, ::T) and similar for S. Since we
# cannot use the julia convert system (we need an instance of the type and not
# the type), we cannot use the julia promotion system.
#
# The Nemo promotion system is used to define catch all functions for
# arithmetic between arbitrary ring elements.
#
################################################################################

promote_rule(::Type{T}, ::Type{T}) where T <: RingElement = T

###############################################################################
#
#   Generic catchall functions
#
###############################################################################

function +(x::S, y::T) where {S <: RingElem, T <: RingElem}
   if S == promote_rule(S, T)
      +(x, parent(x)(y))
   else
      +(parent(y)(x), y)
   end
end

+(x::RingElem, y::RingElement) = x + parent(x)(y)

+(x::RingElement, y::RingElem) = parent(y)(x) + y

function -(x::S, y::T) where {S <: RingElem, T <: RingElem}
   if S == promote_rule(S, T)
      -(x, parent(x)(y))
   else
      -(parent(y)(x), y)
   end
end

-(x::RingElem, y::RingElement) = x - parent(x)(y)

-(x::RingElement, y::RingElem) = parent(y)(x) - y

function *(x::S, y::T) where {S <: RingElem, T <: RingElem}
   if S == promote_rule(S, T)
      *(x, parent(x)(y))
   else
      *(parent(y)(x), y)
   end
end

*(x::RingElem, y::RingElement) = x*parent(x)(y)

*(x::RingElement, y::RingElem) = parent(y)(x)*y

function divexact(x::S, y::T) where {S <: RingElem, T <: RingElem}
   if S == promote_rule(S, T)
      divexact(x, parent(x)(y))
   else
      divexact(parent(y)(x), y)
   end
end

divexact(x::RingElem, y::RingElement) = divexact(x, parent(x)(y))

divexact(x::RingElement, y::RingElem) = divexact(parent(y)(x), y)

function divides(x::T, y::T) where {T <: RingElem}
   q, r = divrem(x, y)
   return iszero(r), q
end

function ==(x::S, y::T) where {S <: RingElem, T <: RingElem}
   if S == promote_rule(S, T)
      ==(x, parent(x)(y))
   else
      ==(parent(y)(x), y)
   end
end

==(x::RingElem, y::RingElement) = x == parent(x)(y)

==(x::RingElement, y::RingElem) = parent(y)(x) == y

function addmul!(z::T, x::T, y::T, c::T) where {T <: RingElem}
   c = mul!(c, x, y)
   z = addeq!(z, c)
   return z
end

Base.literal_pow(::typeof(^), x::T, ::Val{p}) where {p, T <: RingElem} = x^p

###############################################################################
#
#   Baby-steps giant-steps powering
#
###############################################################################

function powers(a::T, d::Int) where {T <: RingElement}
   d <= 0 && throw(DomainError())
   S = parent(a)
   A = Array{T}(d + 1)
   A[1] = one(S)
   if d > 1
      c = a
      A[2] = a
      for i = 2:d
         c *= a
         A[i + 1] = c
      end
   end
   return A
end

###############################################################################
#
#   Type traits
#
###############################################################################

# Type can only represent elements of an exact ring
# true unless explicitly specified
isexact_type(R::Type{T}) where T <: RingElem = true

# Type can only represent elements of domains
# false unless explicitly specified
isdomain_type(R::Type{T}) where T <: RingElem = false

###############################################################################
#
#   Exponential function for generic rings
#
###############################################################################

function exp(a::T) where {T <: RingElem}
   a != 0 && error("Exponential of nonzero element")
   return one(parent(a))
end

################################################################################
#
#   Transpose for ring elements
#
################################################################################

transpose(x::T) where {T <: RingElem} = deepcopy(x)

###############################################################################
#
#   One and zero
#
###############################################################################

one(x::T) where {T <: RingElem} = one(parent(x))

zero(x::T) where {T <: RingElem} = zero(parent(x))

###############################################################################
#
#   Coprime bases
#
###############################################################################

# Bernstein, "Factoring into coprimes in essentially linear time"
# ppio(a,b) = (c,n) where v_p(c) = v_p(a) if v_p(b) != 0, 0 otherwise
# c*n = a or c = gcd(a, b^infty), n = div(a, c).
# This is used in various Euclidean domains for Chinese remaindering.

function ppio(a::E, b::E) where E <: RingElem
   c = gcd(a, b)
   n = div(a, c)
   g = gcd(c, n)
   while !isone(g)
      c *= g
      n = div(n, g)
      g = gcd(c, n)
   end
   return c, n
end

###############################################################################
#
#   Generic and specific rings and fields
#
###############################################################################

include("julia/Integer.jl")

include("julia/Rational.jl")

include("julia/Float.jl")

include("flint/fmpz.jl")

include("flint/fmpz_poly.jl")

include("flint/nmod_poly.jl")

include("flint/nmod.jl")

include("flint/fmpz_mod_poly.jl")

#include("flint/fmpz_mpoly.jl")

include("flint/fmpz_rel_series.jl")

include("flint/fmpz_abs_series.jl")

include("flint/nmod_rel_series.jl")

include("flint/fmpz_mod_rel_series.jl")

include("flint/fmpz_mod_abs_series.jl")

include("flint/fmpz_mat.jl")

include("flint/fmpq_mat.jl")

include("flint/nmod_mat.jl")

include("Fields.jl")

include("flint/fmpq_poly.jl")

include("flint/padic.jl")

include("flint/fmpq_rel_series.jl")

include("flint/fmpq_abs_series.jl")

include("flint/fq_rel_series.jl")

include("flint/fq_abs_series.jl")

include("flint/fq_nmod_rel_series.jl")

include("flint/fq_nmod_abs_series.jl")

include("flint/fq_poly.jl")

include("flint/fq_nmod_poly.jl")

include("arb/arb_poly.jl")

include("arb/acb_poly.jl")

include("arb/arb_mat.jl")

include("arb/acb_mat.jl")

include("Factor.jl")

###############################################################################
#
#   Generic functions to be defined after all rings
#
###############################################################################

include("polysubst.jl")


text:
using Transformations
using Base.Test

using LearnBase
import LossFunctions: L2DistLoss
using Transformations.TestTransforms
import Transformations.Testing: check_gradient
using Distributions
using StochasticOptimization.Iteration
import MultivariateStats

# for reproducibility
srand(1)

function verify_gradient(err)
    isempty(err) && return
    if maximum(abs(err)) > 1e-3
        @show err extrema(err)
        @test false
    end
end

function test_gradient(t::Transformation, ϵ=1e-5)
    perr, xerr = check_gradient(t, ϵ=ϵ)
    verify_gradient(perr)
    verify_gradient(xerr)
end

@testset "Concat" begin
    n1, n2 = 2, 3
    n = n1 + n2
    cc = Concat(n)

    node1 = Transformations.OutputNode(n1)
    node1.val[:] = rand(n1)
    link_nodes!(node1, input_node(cc))
    node2 = Transformations.OutputNode(n2)
    node2.val[:] = rand(n2)
    link_nodes!(node2, input_node(cc))
    transform!(cc.input)

    @test transform!(cc) == vcat(node1.val, node2.val)

    ∇y = ones(n)
    grad!(cc, ∇y)
    @test input_grad(cc) == ∇y

    test_gradient(cc)
end

@testset "Gate" begin
    n = 3
    g = Gate(n)

    node1 = Transformations.OutputNode(n)
    node1.val[:] = rand(n)
    link_nodes!(node1, input_node(g))
    node2 = Transformations.OutputNode(n)
    node2.val[:] = rand(n)
    link_nodes!(node2, input_node(g))
    transform!(g.input)

    @test transform!(g) ≈ node1.val .* node2.val

    ∇y = ones(n)
    grad!(g, ∇y)
    @test input_grad(g) == ∇y

    # test_gradient(g)
end

@testset "PCA" begin
    nin, nout = 4, 2
    n = 50
    μ = zeros(nin)
    Λ = UpperTriangular(rand(nin,nin))
    Σ = Λ'*Λ
    mv = MultivariateNormal(μ, Σ)
    x = rand(mv, n)

    # learn the whitening by passing over the data a few times
    # note: since we do this online, it should approach the "true"
    # value in the limit
    totn = 1000n
    w = Whiten(nin, nout, lookback=100n, method=:pca)
    for i=1:totn
        learn!(w, x[:,mod1(i,n)])
    end

    # compute the projected output
    y = zeros(nout,n)
    for (xi,yi) in each_obs(x,y)
        yi[:] = transform!(w,xi)
    end

    # do it using MultivariateStats (the reference)
    wref = MultivariateStats.fit(MultivariateStats.PCA, x, maxoutdim=nout)
    yref = MultivariateStats.transform(wref, x)

    # check that we're close enough to the "true" projection
    @test norm(abs(y./yref)-1) < 1

    # check that the covariance matrices are close
    @test maximum(abs(cov(y') - cov(yref'))) < 1e-2

    # check that the means are close
    @test maximum(abs(mean(y,2) - mean(yref,2))) < 1e-2
end

@testset "Whitened PCA" begin
    nin, nout = 4, 2
    n = 50
    μ = zeros(nin)
    Λ = UpperTriangular(rand(nin,nin))
    Σ = Λ'*Λ
    mv = MultivariateNormal(μ, Σ)
    x = rand(mv, n)

    # learn the whitening by passing over the data a few times
    # note: since we do this online, it should approach the "true"
    # value in the limit
    totn = 50n
    w = Whiten(nin, nout, lookback=totn, method=:whitened_pca)
    for i=1:totn
        learn!(w, x[:,mod1(i,n)])
    end

    # compute the projected output
    y = zeros(nout,n)
    for (xi,yi) in each_obs(x,y)
        yi[:] = transform!(w,xi)
    end

    # check that the covariance is close to the identity
    # @show cov(y')
    @test maximum(abs(cov(y')-I)) < 1e-1

    # check that the mean is close to zero
    # @show mean(y,2)
    @test maximum(abs(mean(y,2))) < 1e-1
end

@testset "Whitened ZCA" begin
    nin, nout = 4, 4
    n = 50
    μ = zeros(nin)
    Λ = UpperTriangular(rand(nin,nin))
    Σ = Λ'*Λ
    mv = MultivariateNormal(μ, Σ)
    x = rand(mv, n)

    # learn the whitening by passing over the data a few times
    # note: since we do this online, it should approach the "true"
    # value in the limit
    totn = 100n
    w = Whiten(nin, nout, lookback=10n, method=:zca)
    for i=1:totn
        learn!(w, x[:,mod1(i,n)])
    end

    # compute the projected output
    y = zeros(nout,n)
    for (xi,yi) in each_obs(x,y)
        yi[:] = transform!(w,xi)
    end

    # check that the mean is close to zero
    # @show mean(y,2)
    @test maximum(abs(mean(y,2))) < 1e-1
end

@testset "Distributions" begin
    n = 4
    μ = rand(n)
    σ = rand(n)
    t = MvNormalTransformation(μ, σ)
    @test typeof(t.dist.Σ) <: Distributions.PDiagMat
    @test t.n == n
    @test t.nμ == n
    @test t.nU == n
    @test input_length(t) == 2n
    @test output_length(t) == n
    newx = rand(2n)
    transform!(t, newx)
    grad!(t)
    @test t.dist.Σ.diag ≈ newx[n+1:end] .^ 2
    @test t.dist.Σ.inv_diag ≈ 1 ./ (newx[n+1:end] .^ 2)

    U = rand(n,n)
    t = MvNormalTransformation(μ, U*U')
    @test typeof(t.dist.μ) <: Vector{Float64}
    @test typeof(t.dist.Σ) <: Distributions.PDMat
    @test t.n == n
    @test t.nμ == n
    @test t.nU == div(n*(n+1), 2)
    @test input_length(t) == n + div(n*(n+1), 2)
    @test output_length(t) == n
    newx = rand(input_length(t))
    transform!(t, newx)
    grad!(t)
    @test t.dist.μ == newx[1:n]
    cf = t.dist.Σ.chol[:UL]
    @test cf[1,2] == newx[n+2]
    @test cf[end] == newx[end]

    # test_gradient(t)
end

@testset "Affine" begin
    nin, nout = 2, 3
    input = rand(nin)
    target = rand(nout)
    a = Affine(nin, nout)
    loss = L2DistLoss()
    w, b = a.params.views
    x, y = input_value(a), output_value(a)
    ∇w, ∇b = a.params.∇_views
    ∇x, ∇y = input_grad(a), output_grad(a)
    nparams = nout*(nin+1)

    for i=1:2
        output = transform!(a, input)

        @test y == w * x + b
        @test size(x) == (nin,)
        @test size(w) == (nout,nin)
        @test size(b) == (nout,)
        @test size(y) == (nout,)
        @test size(∇x) == (nin,)
        @test size(∇w) == (nout,nin)
        @test size(∇b) == (nout,)
        @test size(∇y) == (nout,)
        @test size(output) == (nout,)
        @test size(params(a)) == (nparams,)
        @test size(grad(a)) == (nparams,)

        l = value(loss, target, output)
        dl = deriv(loss, target, output)
        grad!(a, dl)

        @test grad(a.output) == dl
        @test ∇w ≈ repmat(input', nout, 1) .* repmat(dl, 1, nin)
        @test ∇b == dl
        @test ∇x ≈ w' * dl
    end

    test_gradient(a)
end

@testset "Linear" begin
    nin, nout = 2, 3
    input = rand(nin)
    target = rand(nout)
    t = Linear(nin, nout)
    loss = L2DistLoss()
    w = t.params.views[1]
    x, y = input_value(t), output_value(t)
    ∇w = t.params.∇_views[1]
    ∇x, ∇y = input_grad(t), output_grad(t)
    nparams = nout*nin

    for i=1:2
        output = transform!(t, input)

        wx = w * x
        @test y ≈ wx
        @test size(x) == (nin,)
        @test size(w) == (nout,nin)
        @test size(y) == (nout,)
        @test size(∇x) == (nin,)
        @test size(∇w) == (nout,nin)
        @test size(∇y) == (nout,)
        @test size(output) == (nout,)
        @test size(params(t)) == (nparams,)
        @test size(grad(t)) == (nparams,)

        l = value(loss, target, output)
        dl = deriv(loss, target, output)
        grad!(t, dl)

        @test grad(t.output) == dl
        @test ∇w ≈ repmat(input', nout, 1) .* repmat(dl, 1, nin)
        @test ∇x ≈ w' * dl
    end

    test_gradient(t, 1e-5)
end

@testset "LayerNorm" begin
    n = 3
    input = rand(n)
    t = LayerNorm(n, α=1.0)
    g, b = t.params.views
    a, y = input_value(t), output_value(t)
    ∇g, ∇b = t.params.∇_views
    ∇a, ∇y = input_grad(t), output_grad(t)
    nparams = 2n

    for i=1:2
        output = transform!(t, input)

        @test y ≈ g .* (a .- mean(a)) ./ std(a) .+ b
        @test size(a) == (n,)
        @test size(g) == (n,)
        @test size(b) == (n,)
        @test size(y) == (n,)
        @test size(∇a) == (n,)
        @test size(∇g) == (n,)
        @test size(∇b) == (n,)
        @test size(∇y) == (n,)
        @test size(output) == (n,)
        @test size(params(t)) == (nparams,)
        @test size(grad(t)) == (nparams,)
    end

    # first test with g=ones, b=zeros, then test with random number
    test_gradient(t, 1e-5)
    g[:] = rand(n)
    b[:] = rand(n)
    test_gradient(t, 1e-5)
end

@testset "Activations" begin
    n = 2
    input = rand(n)
    for s in Transformations.activations
        f = Activation{s,Float64}(n)
        output = transform!(f, input)
        @test output == map(@eval($s), input)

        grad!(f, ones(2))
        @test f.input.∇ == map(@eval($(Symbol(s,"′"))), input)

        test_gradient(f)
    end
end

@testset "Chain" begin
    n1, n2, n3 = 4,5,2
    input = rand(n1)

    T = Float64
    chain = Chain(T,
        Affine(T, n1, n2),
        Activation{:relu,T}(n2),
        Affine(T, n2, n3),
        Activation{:logistic,T}(n3)
    )

    @test length(chain.ts) == 4
    @test chain.ts[1].input.val === chain.input.val
    @test chain.ts[end].output.val === chain.output.val

    x = input
    for t in chain.ts
        x = transform!(t, x)
        @test x ≈ output_value(t)
    end
    y = transform!(chain, input)
    @test y ≈ output_value(chain.ts[end])
    @test y ≈ x

    # first compute the chain of transformations manually,
    manual_output = input
    for t in chain.ts
        manual_output = transform!(t, manual_output)
    end

    # now do it using the chain, and make sure it matches
    output = transform!(chain, input)
    @test manual_output == output

    test_gradient(chain)
end

@testset "ResidualLayer" begin
    n = 3
    f = nnet(3, 3, [4], layernorm=false)
    t = ResidualLayer(n, f)
    @test f === t.f

    x = rand(n)
    x2 = copy(x)
    y = copy(transform!(t, x))
    @test x == x2
    y2 = transform!(t, x)
    @test y == y2
    @test y ≈ transform!(f, x) + x2

    ∇y = rand(n)
    grad!(t, ∇y)
    @test output_grad(f) == ∇y

    test_gradient(t, 1e-8)
end

@testset "resnet" begin
    nin, nout = 3, 2
    nblocks = 2
    t = resnet(nin, nout, nblocks, nh=[2], layernorm=false)
    # @show t

    @test length(t.ts) == nblocks+1
    @test isa(t.ts[end], Affine)
    for i=1:nblocks
        @test isa(t.ts[i], ResidualLayer)
    end

    x = rand(nin)
    y = copy(transform!(t, x))

    x̂ = copy(transform!(t.ts[1], x))
    for i=2:nblocks+1
        x̂ = transform!(t.ts[i])
    end
    @test y ≈ x̂

    ∇y = rand(nout)
    grad!(t, ∇y)
    @test output_grad(t) == ∇y

    test_gradient(t)
end

@testset "ConvFilter" begin
    nin = (3,3)
    nfilter = (2,2)
    stride = (1,1)
    nrf, ncf = nfilter
    x = reshape(linspace(1,9,9), nin...)
    f = ConvFilter(nin, nfilter, stride)

    @test f.sizein == nin
    @test f.sizefilter == nfilter
    @test f.sizeout == (2,2)

    w, b = f.params.views
    b[1] = 0.5
    @test size(w) == nfilter
    @test size(b) == (1,)

    y = transform!(f, x)
    nr,nc = f.sizeout
    for r=1:nr, c=1:nc
        @test y[r,c] ≈ sum(w .* view(x, r:r+nrf-1, c:c+ncf-1)) + b[1]
    end

    ∇y = rand(f.sizeout...)
    grad!(f, ∇y)

    ∇x = input_grad(f)
    ∇w, ∇b = f.params.∇_views
    for i=1:nrf, j=1:ncf
        @test ∇w[i,j] ≈ sum(view(x, i:nr+i-1, j:nc+j-1) .* ∇y)
    end
    @test ∇b[1] == sum(∇y)
    for r=1:nin[1], c=1:nin[2]
        @test ∇x[r,c] ≈ sum(view(w, nrf:-1:1, ncf:-1:1) .*
            Transformations.TileView{Float64,2,typeof(∇y)}(∇y, (r-nrf+1,c-ncf+1), (nrf,ncf)))
    end

    # TODO: once stride is implemented, test that
end

@testset "Differentiable" begin
    f(θ) = sum(100θ.^2)
    df(θ) = map(θi -> 200θi, θ)
    t = tfunc(f, 2, df)
    θ = params(t)
    θ[:] = [1.,2.]
    @test transform!(t) == f(θ)
    grad!(t)
    @test grad(t) == df(θ)

    t = tfunc(rosenbrock, 4, rosenbrock_gradient)
    @test typeof(t) <: OnceDifferentiable{Float64}
    θ = params(t)
    @test size(θ) == (4,)
    θ[:] = linspace(0,3,4)
    @test transform!(t) == rosenbrock(θ)
    grad!(t)
    @test grad(t) == rosenbrock_gradient(θ)

    test_gradient(t)
end


text:
# Semantic structure of a SQL query.

# Base node type.

"""
A SQL expression.
"""
abstract type AbstractSQLNode
end

# A node that produces tabular output.
abstract type TabularNode <: AbstractSQLNode
end

function dissect(scr::Symbol, NodeType::Type{<:AbstractSQLNode}, pats::Vector{Any})
    scr_core = gensym(:scr_core)
    ex = Expr(:&&, :($scr_core isa $NodeType), Any[dissect(scr_core, pat) for pat in pats]...)
    :($scr isa SQLNode && (local $scr_core = $scr[]; $ex))
end

# Specialization barrier node.

"""
An opaque wrapper over an arbitrary SQL node.
"""
struct SQLNode <: AbstractSQLNode
    core::AbstractSQLNode

    SQLNode(@nospecialize core::AbstractSQLNode) =
        new(core)
end

Base.getindex(n::SQLNode) =
    getfield(n, :core)

Base.convert(::Type{SQLNode}, n::SQLNode) =
    n

Base.convert(::Type{SQLNode}, @nospecialize n::AbstractSQLNode) =
    SQLNode(n)

Base.convert(::Type{SQLNode}, obj) =
    convert(SQLNode, convert(AbstractSQLNode, obj)::AbstractSQLNode)

(n::AbstractSQLNode)(n′) =
    n(convert(SQLNode, n′))

(n::AbstractSQLNode)(n′::SQLNode) =
    rebase(n, n′)

label(n::SQLNode) =
    label(n[])::Symbol

label(::Union{AbstractSQLNode, Nothing}) =
    :_

rebase(n::SQLNode, n′) =
    convert(SQLNode, rebase(n[], n′))

# Generic traversal and substitution.

function visit(f, n::SQLNode)
    visit(f, n[])
    f(n)
    nothing
end

function visit(f, ns::Vector{SQLNode})
    for n in ns
        visit(f, n)
    end
end

visit(f, ::Nothing) =
    nothing

@generated function visit(f, n::AbstractSQLNode)
    exs = Expr[]
    for f in fieldnames(n)
        t = fieldtype(n, f)
        if t === SQLNode || t === Union{SQLNode, Nothing} || t === Vector{SQLNode}
            ex = quote
                visit(f, n.$(f))
            end
            push!(exs, ex)
        end
    end
    push!(exs, :(return nothing))
    Expr(:block, exs...)
end

substitute(n::SQLNode, c::SQLNode, c′::SQLNode) =
    SQLNode(substitute(n[], c, c′))

function substitute(ns::Vector{SQLNode}, c::SQLNode, c′::SQLNode)
    i = findfirst(isequal(c), ns)
    i !== nothing || return ns
    ns′ = copy(ns)
    ns′[i] = c′
    ns′
end

substitute(::Nothing, ::SQLNode, ::SQLNode) =
    nothing

@generated function substitute(n::AbstractSQLNode, c::SQLNode, c′::SQLNode)
    exs = Expr[]
    fs = fieldnames(n)
    for f in fs
        t = fieldtype(n, f)
        if t === SQLNode || t === Union{SQLNode, Nothing}
            ex = quote
                if n.$(f) === c
                    return $n($(Any[Expr(:kw, f′, f′ !== f ? :(n.$(f′)) : :(c′))
                                    for f′ in fs]...))
                end
            end
            push!(exs, ex)
        elseif t === Vector{SQLNode}
            ex = quote
                let cs′ = substitute(n.$(f), c, c′)
                    if cs′ !== n.$(f)
                        return $n($(Any[Expr(:kw, f′, f′ !== f ? :(n.$(f′)) : :(cs′))
                                        for f′ in fs]...))
                    end
                end
            end
            push!(exs, ex)
        end
    end
    push!(exs, :(return n))
    Expr(:block, exs...)
end

# Pretty-printing.

Base.show(io::IO, n::AbstractSQLNode) =
    print(io, quoteof(n, limit = true))

Base.show(io::IO, ::MIME"text/plain", n::AbstractSQLNode) =
    pprint(io, n)

function PrettyPrinting.quoteof(n::SQLNode;
                                limit::Bool = false,
                                unwrap::Bool = false)
    if limit
        ctx = QuoteContext(limit = true)
        ex = quoteof(n[], ctx)
        if unwrap
            ex = Expr(:ref, ex)
        end
        return ex
    end
    tables_seen = OrderedSet{SQLTable}()
    nodes_seen = OrderedSet{SQLNode}()
    nodes_toplevel = Set{SQLNode}()
    visit(n) do n
        core = n[]
        if core isa FromNode
            push!(tables_seen, core.table)
        end
        if core isa TabularNode
            push!(nodes_seen, n)
            push!(nodes_toplevel, n)
        elseif n in nodes_seen
            push!(nodes_toplevel, n)
        else
            push!(nodes_seen, n)
        end
    end
    ctx = QuoteContext()
    defs = Any[]
    if length(nodes_toplevel) >= 2 || (length(nodes_toplevel) == 1 && !(n in nodes_toplevel))
        for t in tables_seen
            def = quoteof(t, limit = true)
            name = t.name
            push!(defs, Expr(:(=), name, def))
            ctx.vars[t] = name
        end
        qidx = 0
        for n in nodes_seen
            n in nodes_toplevel || continue
            def = quoteof(n, ctx)
            qidx += 1
            name = Symbol('q', qidx)
            push!(defs, Expr(:(=), name, def))
            ctx.vars[n] = name
        end
    end
    ex = quoteof(n, ctx)
    if unwrap
        ex = Expr(:ref, ex)
    end
    if !isempty(defs)
        ex = Expr(:let, Expr(:block, defs...), ex)
    end
    ex
end

PrettyPrinting.quoteof(n::AbstractSQLNode; limit::Bool = false) =
    quoteof(convert(SQLNode, n), limit = limit, unwrap = true)

PrettyPrinting.quoteof(n::SQLNode, ctx::QuoteContext) =
    if !ctx.limit
        var = get(ctx.vars, n, nothing)
        if var !== nothing
            var
        else
            quoteof(n[], ctx)
        end
    else
        :…
    end

PrettyPrinting.quoteof(ns::Vector{SQLNode}, ctx::QuoteContext) =
    if isempty(ns)
        Any[]
    elseif !ctx.limit
        Any[quoteof(n, ctx) for n in ns]
    else
        Any[:…]
    end

# Errors.

struct DuplicateLabelError <: FunSQLError
    name::Symbol
    path::Vector{SQLNode}

    DuplicateLabelError(name; path = SQLNode[]) =
        new(name, path)
end

function Base.showerror(io::IO, err::DuplicateLabelError)
    print(io, "FunSQL.DuplicateLabelError: $(err.name) is used more than once")
    showpath(io, err.path)
end

struct IllFormedError <: FunSQLError
    path::Vector{SQLNode}

    IllFormedError(; path = SQLNode[]) =
        new(path)
end

function Base.showerror(io::IO, err::IllFormedError)
    print(io, "FunSQL.IllFormedError")
    showpath(io, err.path)
end

module REFERENCE_ERROR_TYPE

@enum ReferenceErrorType::UInt8 begin
    UNDEFINED_HANDLE
    AMBIGUOUS_HANDLE
    UNDEFINED_NAME
    AMBIGUOUS_NAME
    UNEXPECTED_ROW_TYPE
    UNEXPECTED_SCALAR_TYPE
    UNEXPECTED_AGGREGATE
    AMBIGUOUS_AGGREGATE
end

end

import .REFERENCE_ERROR_TYPE.ReferenceErrorType

struct ReferenceError <: FunSQLError
    type::ReferenceErrorType
    name::Union{Symbol, Nothing}
    path::Vector{SQLNode}

    ReferenceError(type; name = nothing, path = SQLNode[]) =
        new(type, name, path)
end

function Base.showerror(io::IO, err::ReferenceError)
    print(io, "FunSQL.ReferenceError: ")
    if err.type == REFERENCE_ERROR_TYPE.UNDEFINED_HANDLE
        print(io, "node-bound reference failed to resolve")
    elseif err.type == REFERENCE_ERROR_TYPE.AMBIGUOUS_HANDLE
        print(io, "node-bound reference is ambiguous")
    elseif err.type == REFERENCE_ERROR_TYPE.UNDEFINED_NAME
        print(io, "cannot find $(err.name)")
    elseif err.type == REFERENCE_ERROR_TYPE.AMBIGUOUS_NAME
        print(io, "$(err.name) is ambiguous")
    elseif err.type == REFERENCE_ERROR_TYPE.UNEXPECTED_ROW_TYPE
        print(io, "incomplete reference $(err.name)")
    elseif err.type == REFERENCE_ERROR_TYPE.UNEXPECTED_SCALAR_TYPE
        print(io, "unexpected reference after $(err.name)")
    elseif err.type == REFERENCE_ERROR_TYPE.UNEXPECTED_AGGREGATE
        print(io, "aggregate expression requires Group or Partition")
    elseif err.type == REFERENCE_ERROR_TYPE.AMBIGUOUS_AGGREGATE
        print(io, "aggregate expression is ambiguous")
    end
    showpath(io, err.path)
end

function showpath(io, path::Vector{SQLNode})
    if !isempty(path)
        q = highlight(path)
        println(io, " in:")
        pprint(io, q)
    end
end

function highlight(path::Vector{SQLNode}, color = Base.error_color())
    @assert !isempty(path)
    n = Highlight(over = path[1], color = color)
    for k = 2:lastindex(path)
        n = substitute(path[k], path[k-1], n)
    end
    n
end

# Concrete node types.

include("nodes/aggregate.jl")
include("nodes/append.jl")
include("nodes/as.jl")
include("nodes/bind.jl")
include("nodes/define.jl")
include("nodes/from.jl")
include("nodes/function.jl")
include("nodes/get.jl")
include("nodes/group.jl")
include("nodes/highlight.jl")
include("nodes/join.jl")
include("nodes/limit.jl")
include("nodes/literal.jl")
include("nodes/order.jl")
include("nodes/partition.jl")
include("nodes/select.jl")
include("nodes/sort.jl")
include("nodes/variable.jl")
include("nodes/where.jl")


text:
export OperationOnSequences

struct OperationOnSequences <: Sequence
  seq1
  seq2
  op
end

function sample(seq::OperationOnSequences, t::AbstractVector)
  seq.op.(sample(seq.seq1, t), sample(seq.seq2, t))
end

import Base.+, Base.-, Base.*
for op in (:+, :-, :*)
  @eval begin
    $op(s1::Sequence, s2::Sequence) = OperationOnSequences(s1, s2, $op)
    $op(s::Sequence, x::Number) = $op(s, Constant(x))
    $op(x::Number, s::Sequence) = $op(Constant(x), s)
  end
end


text:
module Poisson

export poisson

using PyCall

@pyimport pyamg.gallery as gallery
@pyimport pyamg.gallery.diffusion as diffusion

function poisson(s)
    stencil = diffusion.diffusion_stencil_2d()
    m = gallery.stencil_grid(stencil, s, format="csc")
    plus_one(x) = map(x -> x + 1, x)
    SparseMatrixCSC(m[:shape][1], m[:shape][2], plus_one(m[:indptr]), plus_one(m[:indices]), m[:data])
end

end


text:
function make_useful_fname(exp_dict::Dict,cond_key::Int64)
    exp_name = exp_dict[:name]
    cond_role = exp_dict[:conditions][cond_key][:role]

    fname = exp_name * "_" * string(cond_role) * "_condition_" * string(cond_key)
end
function make_condition_dict(;num_nets::Int64=2,t_bnds::Array{Float64}=[0.0,1.0],
                       e_bnds::Array{Float64}=[1000.],mcts_d::Array{Int64}=[2],
                       deg_bnds::Array{Float64}=[4.],net_type::Symbol=:random,d_bnds::Array{Float64}=[0.95],
                       exit_rwd_bnds::Array{Float64}=[2000.],sense_rwd_bnds::Array{Float64}=[-100.],
                       caught_rwd_bnds::Array{Float64}=[-2000.],mcts_its_bnds::Array{Int64}=[1000],
                       seed::Int64=5,n_bnds::Array{Int64}=[13],repeats::Int64=250,steps::Int64=-1,
                       dis_rwd::Bool=false,log_to_file::Bool=false,role::Symbol=:candidate)
    # this is a default condition, to reduce the need for typing a bunch of repeated information
    # num_nets -- number of networks to produce
    # t_bnds -- bounds for transition probabilities of produced networks,
    #           single parameter indicates that the parameter will be constant across all networks
    # e_bnds -- bounds for mcts exploration constant of produced networks
    # mcts_d -- bounds for mcts depth of produced networks
    # deg_bnds -- bounds for degree of produced networks
    # net_type -- type of network produced -- :random or :original
    # d_bnds -- bound for discount factor of produced networks
    # exit_rwd_bnds -- bound for exit reward of produced networks
    # sense_rwd_bnds -- bound for sense reward of produced networks
    # caught_rwd_bnds -- bound for caught reward of produced networks
    # mcts_its_gnds -- bound for number of mcts iterations of produced networks
    # seed -- seed for random number generator
    # n_bnds -- bound for number of nodes in network, original has 13
    # repeats -- number of simulations that will be run to determine reward distribution
    # steps -- number of steps used in simulation, before simulation terminates
    # dis_rwd -- whether to use discounted reward or not (other possibility is un-discounted reward)
    # log_to_file -- whether to log to a file
    # role -- what role the condition will play in the overall experiment. :trusted, or :candidate

    net_dict = Dict(:num_nets=>num_nets,:t_bnds=>t_bnds,:e_bnds=>e_bnds,:solver_depth=>mcts_d,
                    :deg_bnds=>deg_bnds,:net_type=>net_type,:n_bnds=>n_bnds,:d_bnds=>d_bnds,
                    :exit_rwd_bnds=>exit_rwd_bnds,:sense_rwd_bnds=>sense_rwd_bnds,
                    :caught_rwd_bnds=>caught_rwd_bnds,:mcts_its_bnds=>mcts_its_bnds,:seed=>seed)
    sims_dict = Dict(:repeats=>repeats,:steps=>steps,:dis_rwd=>dis_rwd,:log_to_file=>log_to_file,
                      :log_fname=>"logs/$(now())",:log_lvl=>:debug)
    cond = Dict(:role=>role,:nets=>net_dict,:sims=>sims_dict)

    return cond
end
if experiment_name == "transition_vary"
    include("experiment_params/transition_vary.jl")
elseif experiment_name == "mturk_fast"
    include("experiment_params/mturk_fast.jl")
elseif experiment_name == "mturk"
    include("experiment_params/mturk.jl")
elseif experiment_name == "mturk_supplement"
    include("experiment_params/mturk_supplement.jl")
elseif experiment_name == "x4_test"
    include("experiment_params/x4_test.jl")
else
    # include others if desired, examples in `experiment_params` folder
    error("not implemented yet")
end

max_stars_repo_name: COHRINT/FaMSeC

max_issues_repo_name: COHRINT/FaMSeC

max_forks_repo_name: COHRINT/FaMSeC


text:
# Predefined regularizers
# You may also implement your own regularizer by subtyping
# the abstract type Regularizer.
# Regularizers should implement `evaluate` and `prox`.

import Base: *

export Regularizer, ProductRegularizer, # abstract types
       # concrete regularizers
       QuadReg, QuadConstraint,
       OneReg, ZeroReg, NonNegConstraint, NonNegOneReg, NonNegQuadReg,
       OneSparseConstraint, UnitOneSparseConstraint, SimplexConstraint,
       KSparseConstraint,
       lastentry1, lastentry_unpenalized,
       fixed_latent_features, FixedLatentFeaturesConstraint,
       fixed_last_latent_features, FixedLastLatentFeaturesConstraint,
       OrdinalReg, MNLOrdinalReg,
       RemQuadReg,
       # methods on regularizers
       prox!, prox,
       # utilities
       scale, mul!, *

# numerical tolerance
TOL = 1e-12

# regularizers
# regularizers r should have the method `prox` defined such that
# prox(r)(u,alpha) = argmin_x( alpha r(x) + 1/2 \|x - u\|_2^2)
abstract type Regularizer end
abstract type MatrixRegularizer <: LowRankModels.Regularizer end

# default inplace prox operator (slower than if inplace prox is implemented)
prox!(r::Regularizer,u::AbstractArray,alpha::Number) = (v = prox(r,u,alpha); @simd for i=1:length(u) @inbounds u[i]=v[i] end; u)

# default scaling
scale(r::Regularizer) = r.scale
mul!(r::Regularizer, newscale::Number) = (r.scale = newscale; r)
mul!(rs::Array{Regularizer}, newscale::Number) = (for r in rs mul!(r, newscale) end; rs)
*(newscale::Number, r::Regularizer) = (newr = typeof(r)(); mul!(newr, scale(r)*newscale); newr)

## utilities

function allnonneg(a::AbstractArray)
  for ai in a
    ai < 0 && return false
  end
  return true
end

## Quadratic regularization
mutable struct QuadReg<:Regularizer
    scale::Float64
end
QuadReg() = QuadReg(1)
prox(r::QuadReg,u::AbstractArray,alpha::Number) = 1/(1+2*alpha*r.scale)*u
prox!(r::QuadReg,u::Array{Float64},alpha::Number) = rmul!(u, 1/(1+2*alpha*r.scale))
evaluate(r::QuadReg,a::AbstractArray) = r.scale*sum(abs2, a)

## constrained quadratic regularization
## the function r such that
## r(x) = inf    if norm(x) > max_2norm
##        0      otherwise
## can be used to implement maxnorm regularization:
##   constraining the maxnorm of XY to be <= mu is achieved
##   by setting glrm.rx = QuadConstraint(sqrt(mu))
##   and the same for every element of glrm.ry
mutable struct QuadConstraint<:Regularizer
    max_2norm::Float64
end
QuadConstraint() = QuadConstraint(1)
prox(r::QuadConstraint,u::AbstractArray,alpha::Number) = (r.max_2norm)/norm(u)*u
prox!(r::QuadConstraint,u::Array{Float64},alpha::Number) = mul!(u, (r.max_2norm)/norm(u))
evaluate(r::QuadConstraint,u::AbstractArray) = norm(u) > r.max_2norm + TOL ? Inf : 0
scale(r::QuadConstraint) = 1
mul!(r::QuadConstraint, newscale::Number) = 1

## one norm regularization
mutable struct OneReg<:Regularizer
    scale::Float64
end
OneReg() = OneReg(1)
function softthreshold(x::Number; alpha=1)
 return max(x-alpha,0) + min(x+alpha,0)
end
prox(r::OneReg,u::AbstractArray,alpha::Number) = (st(x) = softthreshold(x; alpha=r.scale*alpha); st.(u))
prox!(r::OneReg,u::AbstractArray,alpha::Number) = (st(x) = softthreshold(x; alpha=r.scale*alpha); map!(st, u, u))
evaluate(r::OneReg,a::AbstractArray) = r.scale*sum(abs,a)

## no regularization
mutable struct ZeroReg<:Regularizer
end
prox(r::ZeroReg,u::AbstractArray,alpha::Number) = u
prox!(r::ZeroReg,u::Array{Float64},alpha::Number) = u
evaluate(r::ZeroReg,a::AbstractArray) = 0
scale(r::ZeroReg) = 0
mul!(r::ZeroReg, newscale::Number) = 0

## indicator of the nonnegative orthant
## (enforces nonnegativity, eg for nonnegative matrix factorization)
mutable struct NonNegConstraint<:Regularizer
end
prox(r::NonNegConstraint,u::AbstractArray,alpha::Number=1) = broadcast(max,u,0)
prox!(r::NonNegConstraint,u::Array{Float64},alpha::Number=1) = (@simd for i=1:length(u) @inbounds u[i] = max(u[i], 0) end; u)
function evaluate(r::NonNegConstraint,a::AbstractArray)
    for ai in a
        if ai<0
            return Inf
        end
    end
    return 0
end
scale(r::NonNegConstraint) = 1
mul!(r::NonNegConstraint, newscale::Number) = 1

## one norm regularization restricted to nonnegative orthant
## (enforces nonnegativity, in addition to one norm regularization)
mutable struct NonNegOneReg<:Regularizer
    scale::Float64
end
NonNegOneReg() = NonNegOneReg(1)
prox(r::NonNegOneReg,u::AbstractArray,alpha::Number) = max.(u.-alpha,0)

function prox!(r::NonNegOneReg,u::AbstractArray,alpha::Number)
  nonnegsoftthreshold = (x::Number -> max(x-alpha,0))
  map!(nonnegsoftthreshold, u, u)
end

function evaluate(r::NonNegOneReg,a::AbstractArray)
    for ai in a
        if ai<0
            return Inf
        end
    end
    return r.scale*sum(a)
end
scale(r::NonNegOneReg) = 1
mul!(r::NonNegOneReg, newscale::Number) = 1

## Quadratic regularization restricted to nonnegative domain
## (Enforces nonnegativity alongside quadratic regularization)
mutable struct NonNegQuadReg
    scale::Float64
end
NonNegQuadReg() = NonNegQuadReg(1)
prox(r::NonNegQuadReg,u::AbstractArray,alpha::Number) = max.(1/(1+2*alpha*r.scale)*u, 0)
prox!(r::NonNegQuadReg,u::AbstractArray,alpha::Number) = begin
  mul!(u, 1/(1+2*alpha*r.scale))
  maxval = maximum(u)
  clamp!(u, 0, maxval)
end
function evaluate(r::NonNegQuadReg,a::AbstractArray)
    for ai in a
        if ai<0
            return Inf
        end
    end
    return r.scale*sumabs2(a)
end

## indicator of the last entry being equal to 1
## (allows an unpenalized offset term into the glrm when used in conjunction with lastentry_unpenalized)
mutable struct lastentry1<:Regularizer
    r::Regularizer
end
lastentry1() = lastentry1(ZeroReg())
prox(r::lastentry1,u::AbstractArray{Float64,1},alpha::Number=1) = [prox(r.r,view(u,1:length(u)-1),alpha); 1]
prox!(r::lastentry1,u::AbstractArray{Float64,1},alpha::Number=1) = (prox!(r.r,view(u,1:length(u)-1),alpha); u[end]=1; u)
prox(r::lastentry1,u::AbstractArray{Float64,2},alpha::Number=1) = [prox(r.r,view(u,1:size(u,1)-1,:),alpha); ones(1, size(u,2))]
prox!(r::lastentry1,u::AbstractArray{Float64,2},alpha::Number=1) = (prox!(r.r,view(u,1:size(u,1)-1,:),alpha); u[end,:]=1; u)
evaluate(r::lastentry1,a::AbstractArray{Float64,1}) = (a[end]==1 ? evaluate(r.r,a[1:end-1]) : Inf)
evaluate(r::lastentry1,a::AbstractArray{Float64,2}) = (all(a[end,:].==1) ? evaluate(r.r,a[1:end-1,:]) : Inf)
scale(r::lastentry1) = scale(r.r)
mul!(r::lastentry1, newscale::Number) = mul!(r.r, newscale)

## makes the last entry unpenalized
## (allows an unpenalized offset term into the glrm when used in conjunction with lastentry1)
mutable struct lastentry_unpenalized<:Regularizer
    r::Regularizer
end
lastentry_unpenalized() = lastentry_unpenalized(ZeroReg())
prox(r::lastentry_unpenalized,u::AbstractArray{Float64,1},alpha::Number=1) = [prox(r.r,u[1:end-1],alpha); u[end]]
prox!(r::lastentry_unpenalized,u::AbstractArray{Float64,1},alpha::Number=1) = (prox!(r.r,view(u,1:size(u,1)-1),alpha); u)
evaluate(r::lastentry_unpenalized,a::AbstractArray{Float64,1}) = evaluate(r.r,a[1:end-1])
prox(r::lastentry_unpenalized,u::AbstractArray{Float64,2},alpha::Number=1) = [prox(r.r,u[1:end-1,:],alpha); u[end,:]]
prox!(r::lastentry_unpenalized,u::AbstractArray{Float64,2},alpha::Number=1) = (prox!(r.r,view(u,1:size(u,1)-1,:),alpha); u)
evaluate(r::lastentry_unpenalized,a::AbstractArray{Float64,2}) = evaluate(r.r,a[1:end-1,:])
scale(r::lastentry_unpenalized) = scale(r.r)
mul!(r::lastentry_unpenalized, newscale::Number) = mul!(r.r, newscale)

## fixes the values of the first n elements of the column to be y
## optionally regularizes the last k-n elements with regularizer r
mutable struct fixed_latent_features<:Regularizer
    r::Regularizer
    y::Array{Float64,1} # the values of the fixed latent features
    n::Int # length of y
end
fixed_latent_features(r::Regularizer, y::Array{Float64,1}) = fixed_latent_features(r,y,length(y))
# standalone use without another regularizer
FixedLatentFeaturesConstraint(y::Array{Float64, 1}) = fixed_latent_features(ZeroReg(),y,length(y))

prox(r::fixed_latent_features,u::AbstractArray,alpha::Number) = [r.y; prox(r.r,u[(r.n+1):end],alpha)]
function prox!(r::fixed_latent_features,u::Array{Float64},alpha::Number)
  	prox!(r.r,u[(r.n+1):end],alpha)
  	u[1:r.n]=y
  	u
end
evaluate(r::fixed_latent_features, a::AbstractArray) = a[1:r.n]==r.y ? evaluate(r.r, a[(r.n+1):end]) : Inf
scale(r::fixed_latent_features) = scale(r.r)
mul!(r::fixed_latent_features, newscale::Number) = mul!(r.r, newscale)

## fixes the values of the last n elements of the column to be y
## optionally regularizes the first k-n elements with regularizer r
mutable struct fixed_last_latent_features<:Regularizer
    r::Regularizer
    y::Array{Float64,1} # the values of the fixed latent features
    n::Int # length of y
end
fixed_last_latent_features(r::Regularizer, y::Array{Float64,1}) = fixed_last_latent_features(r,y,length(y))
# standalone use without another regularizer
FixedLastLatentFeaturesConstraint(y::Array{Float64, 1}) = fixed_last_latent_features(ZeroReg(),y,length(y))

prox(r::fixed_last_latent_features,u::AbstractArray,alpha::Number) = [prox(r.r,u[(r.n+1):end],alpha); r.y]
function prox!(r::fixed_last_latent_features,u::Array{Float64},alpha::Number)
    u[length(u)-r.n+1:end]=y
    prox!(r.r,u[1:length(a)-r.n],alpha)
    u
end
evaluate(r::fixed_last_latent_features, a::AbstractArray) = a[length(a)-r.n+1:end]==r.y ? evaluate(r.r, a[1:length(a)-r.n]) : Inf
scale(r::fixed_last_latent_features) = scale(r.r)
mul!(r::fixed_last_latent_features, newscale::Number) = mul!(r.r, newscale)

## indicator of 1-sparse vectors
## (enforces that exact 1 entry is nonzero, eg for orthogonal NNMF)
mutable struct OneSparseConstraint<:Regularizer
end
prox(r::OneSparseConstraint, u::AbstractArray, alpha::Number=0) = (idx = argmax(u); v=zeros(size(u)); v[idx]=u[idx]; v)
prox!(r::OneSparseConstraint, u::Array, alpha::Number=0) = (idx = argmax(u); ui = u[idx]; mul!(u,0); u[idx]=ui; u)
function evaluate(r::OneSparseConstraint, a::AbstractArray)
    oneflag = false
    for ai in a
        if oneflag
            if ai!=0
                return Inf
            end
        else
            if ai!=0
                oneflag=true
            end
        end
    end
    return 0
end
scale(r::OneSparseConstraint) = 1
mul!(r::OneSparseConstraint, newscale::Number) = 1

## Indicator of k-sparse vectors
mutable struct KSparseConstraint<:Regularizer
  k::Int
end
function evaluate(r::KSparseConstraint, a::AbstractArray)
  k = r.k
  nonzcount = 0
  for ai in a
    if nonzcount == k
      if ai != 0
        return Inf
      end
    else
      if ai != 0
        nonzcount += 1
      end
    end
  end
  return 0
end
function prox(r::KSparseConstraint, u::AbstractArray, alpha::Number)
  k = r.k
  ids = partialsortperm(u, 1:k, by=abs, rev=true)
  uk = zero(u)
  uk[ids] = u[ids]
  uk
end
function prox!(r::KSparseConstraint, u::Array, alpha::Number)
  k = r.k
  ids = partialsortperm(u, 1:k, by=abs, rev=true)
  vals = u[ids]
  mul!(u,0)
  u[ids] = vals
  u
end

## indicator of 1-sparse unit vectors
## (enforces that exact 1 entry is 1 and all others are zero, eg for kmeans)
mutable struct UnitOneSparseConstraint<:Regularizer
end
prox(r::UnitOneSparseConstraint, u::AbstractArray, alpha::Number=0) = (idx = argmax(u); v=zeros(size(u)); v[idx]=1; v)
prox!(r::UnitOneSparseConstraint, u::Array, alpha::Number=0) = (idx = argmax(u); mul!(u,0); u[idx]=1; u)

function evaluate(r::UnitOneSparseConstraint, a::AbstractArray)
    oneflag = false
    for ai in a
        if ai==0
            continue
        elseif ai==1
            if oneflag
                return Inf
            else
                oneflag=true
            end
        else
            return Inf
        end
    end
    return 0
end
scale(r::UnitOneSparseConstraint) = 1
mul!(r::UnitOneSparseConstraint, newscale::Number) = 1

## indicator of vectors in the simplex: nonnegative vectors with unit l1 norm
## (eg for QuadLoss mixtures, ie soft kmeans)
## prox for the simplex is derived by Chen and Ye in [this paper](http://arxiv.org/pdf/1101.6081v2.pdf)
mutable struct SimplexConstraint<:Regularizer
end
function prox(r::SimplexConstraint, u::AbstractArray, alpha::Number=0)
    n = length(u)
    y = sort(u, rev=true)
    ysum = cumsum(y)
    t = (ysum[end]-1)/n
    for i=1:(n-1)
        if (ysum[i]-1)/i >= y[i+1]
            t = (ysum[i]-1)/i
            break
        end
    end
    max.(u .- t, 0)
end
function evaluate(r::SimplexConstraint,a::AbstractArray)
    # check it's a unit vector
    abs(sum(a)-1)>TOL && return Inf
    # check every entry is nonnegative
    for i=1:length(a)
        a[i] < 0 && return Inf
    end
    return 0
end
scale(r::SimplexConstraint) = 1
mul!(r::SimplexConstraint, newscale::Number) = 1

## ordinal regularizer
## a block regularizer which
    # 1) forces the first k-1 entries of each column to be the same
    # 2) forces the last entry of each column to be increasing
    # 3) applies an internal regularizer to the first k-1 entries of each column
## should always be used in conjunction with lastentry1 regularization on x
mutable struct OrdinalReg<:Regularizer
    r::Regularizer
end
OrdinalReg() = OrdinalReg(ZeroReg())
prox(r::OrdinalReg,u::AbstractArray,alpha::Number) = (uc = copy(u); prox!(r,uc,alpha))
function prox!(r::OrdinalReg,u::AbstractArray,alpha::Number)
    um = mean(u[1:end-1, :], dims=2)
    prox!(r.r,um,alpha)
    for i=1:size(u,1)-1
        for j=1:size(u,2)
            u[i,j] = um[i]
        end
    end
    # this enforces rule 2) (increasing last row of u), but isn't exactly the prox function
    # for j=2:size(u,2)
    #     if u[end,j-1] > u[end,j]
    #         m = (u[end,j-1] + u[end,j])/2
    #         u[end,j-1:j] = m
    #     end
    # end
    u
end
evaluate(r::OrdinalReg,a::AbstractArray) = evaluate(r.r,a[1:end-1,1])
scale(r::OrdinalReg) = scale(r.r)
mul!(r::OrdinalReg, newscale::Number) = mul!(r.r, newscale)

# make sure we don't add two offsets cuz that's weird
lastentry_unpenalized(r::OrdinalReg) = r

mutable struct MNLOrdinalReg<:Regularizer
    r::Regularizer
end
MNLOrdinalReg() = MNLOrdinalReg(ZeroReg())
prox(r::MNLOrdinalReg,u::AbstractArray,alpha::Number) = (uc = copy(u); prox!(r,uc,alpha))
function prox!(r::MNLOrdinalReg,u::AbstractArray,alpha::Number; TOL=1e-3)
    um = mean(u[1:end-1, :], dims=2)
    prox!(r.r,um,alpha)
    for i=1:size(u,1)-1
        for j=1:size(u,2)
            u[i,j] = um[i]
        end
    end
    # this enforces rule 2) (decreasing last row of u, all less than 0), but isn't exactly the prox function
    u[end,1] = min(-TOL, u[end,1])
    for j=2:size(u,2)
      u[end,j] = min(u[end,j], u[end,j-1]-TOL)
    end
    u
end
evaluate(r::MNLOrdinalReg,a::AbstractArray) = evaluate(r.r,a[1:end-1,1])
scale(r::MNLOrdinalReg) = scale(r.r)
mul!(r::MNLOrdinalReg, newscale::Number) = mul!(r.r, newscale)
# make sure we don't add two offsets cuz that's weird
lastentry_unpenalized(r::MNLOrdinalReg) = r

## Quadratic regularization with non-zero mean
mutable struct RemQuadReg<:Regularizer
        scale::Float64
        m::Array{Float64, 1}
end
RemQuadReg(m::Array{Float64, 1}) = RemQuadReg(1, m)
prox(r::RemQuadReg, u::AbstractArray, alpha::Number) =
     (u + 2 * alpha * r.scale * r.m) / (1 + 2 * alpha * r.scale)
prox!(r::RemQuadReg, u::Array{Float64}, alpha::Number) = begin
        broadcast!(.+, u, u, 2 * alpha * r.scale * r.m)
        mul!(u, 1 / (1 + 2 * alpha * r.scale))
end
evaluate(r::RemQuadReg, a::AbstractArray) = r.scale * sum(abs2, a - r.m)

## simpler method for numbers, not arrays
evaluate(r::Regularizer, u::Number) = evaluate(r, [u])
prox(r::Regularizer, u::Number, alpha::Number) = prox(r, [u], alpha)[1]
# if step size not specified, step size = 1
prox(r::Regularizer, u) = prox(r, u, 1)


text:
"the effective low frequency bulk modulus and density of a material filled with particles"
function effective_medium(medium::Medium{T}, species::Array{Specie{T}}; dim = 2) where T<:Number
    φ = sum(volume_fraction.(species; dim = dim))

    # calculate effective properties
    β = medium.ρ*medium.c^2 # medium bulk modulus
    if abs(β) == zero(T) || abs(prod(s -> s.ρ*s.c^2, species)) == zero(T)
        β_eff = zero(T)
    else
        β_eff = one(T)/((1-φ)/β + sum(volume_fraction(s; dim = dim)/(s.ρ*s.c^2) for s in species))
    end
    ρ_frac = sum(volume_fraction(s; dim = dim)*(medium.ρ - s.ρ)/(medium.ρ + s.ρ) for s in species)
    ρ_eff = medium.ρ*(1 - ρ_frac)/(1 + T(dim - 1) * ρ_frac)

    return Medium(ρ=ρ_eff, c=sqrt(β_eff/ρ_eff))
end

"the reflection from a homogenious halfspace, which is also the low frequency reflection from a particulate material when using the effective_medium."
function reflection_coefficient_halfspace(incident_medium::Medium{T}, reflect_medium::Medium{T};
        θin::T = zero(T), tol = 1e-8) where T<:Number

    k_in = T(1)/incident_medium.c
    k_r = T(1)/reflect_medium.c

    q = real(reflect_medium.c*reflect_medium.ρ/(incident_medium.c*incident_medium.ρ))
    θ_trans = transmission_angle(k_in, k_r, θin; tol=tol)
    R = (q*cos(θin) - cos(θ_trans))/(q*cos(θin) + cos(θ_trans))
end


text:
module ValueDispatch

using ExpressionUtils

export @dispatch, @on, register

const dispatchtable = Dict{Function, Dict}()

macro dispatch(ex)
    if ex.head == :(=)
        template = quote
            _ESC_name_(_SPLAT_params_) = _SPLAT_body_
        end
    elseif ex.head == :function
        template = quote
            function _ESC_name_(_SPLAT_params_)
                _SPLAT_body_
            end
        end
    else
        error(ArgumentError, "@dispatch must be given a generic function literal")
    end

    @gensym method

    out = quote
        local methodtable = {:default => false}
        function _name_(_UNSPLAT_params_)
            val = begin _UNSPLAT_body_ end
            method = get(methodtable, val, methodtable[:default])
            if method != false
                method(_UNSPLAT_params_)
            else
                error("No implementation for dispatch value $val")
            end
        end
        dispatchtable[_name_] = methodtable
        _name_
    end

    expr_replace(ex, template, out)
end

function register(method::Function, dispatchfn::Function, val)
    methodtable = get(dispatchtable, dispatchfn, false)
    if methodtable != false
        methodtable[val] = method
    else
        error("$dispatchfn is not a multimethod dispatch function")
    end
    dispatchfn
end

macro on(val, ex)
    if ex.head == :(=)
        template = :(_ESC_name_(_SPLAT_params_) = _SPLAT_body_)
    elseif ex.head == :function
        template = quote
            function _ESC_name_(_SPLAT_params_)
                _SPLAT_body_
            end
        end
    else
        error(ArgumentError, "@value must be given a generic function literal")
    end

    out = quote
        register(function (_UNSPLAT_params_)
            _UNSPLAT_body_
        end, _name_, $(esc(val)))
    end

    expr_replace(ex, template, out)
end

end # module ValueDispatch


text:
"""
		sparsify(x,nnzrate)

		replace matrices with at most `nnzrate` fraction of non-zeros with SparseMatrixCSC

```juliadoctest
julia> x = ProductNode((
				ProductNode((
					MatrixNode(randn(5,5)),
					MatrixNode(zeros(5,5))
						)),
				MatrixNode(zeros(5,5))
				))
julia> mapdata(i -> sparsify(i,0.05),x)

```
"""
sparsify(x,nnzrate) = x
sparsify(x::Matrix,nnzrate) = (mean(x .!= 0) <nnzrate) ? sparse(x) : x

# can be removed when https://github.com/FluxML/Flux.jl/pull/1357 is merged
function Base.:*(A::AbstractMatrix, B::Adjoint{Bool,<: Flux.OneHotMatrix})
    m = size(A,1)
    Y = similar(A, m, size(B,2))
    Y .= 0
    BT = B'
    for (j,ohv) in enumerate(BT.data)
        ix = ohv.ix
        for i in 1:m
            @inbounds Y[i,ix] += A[i,j]
        end
    end
    Y
end


text:
using Documenter
using AnotherParser

#=
See this document
https://juliadocs.github.io/Documenter.jl/stable/man/hosting/
for somewhat rambling instructions about how to set up a
documentation generation workflow.
=#

# Temporary hack:
push!(LOAD_PATH,"../src/")

#=
# One time setup of deployment secret
DocumenterTools.genkeys(; user="MarkNahabedian",
                        repo="AnotherParser.jl")
=#

makedocs(;
         modules=[AnotherParser],
         format=Documenter.HTML(),
         pages=[
             "Home" => "index.md",
         ],
         sitename="AnotherParser.jl",
         authors="Mark Nahabedian"
)

deploydocs(;
    repo="github.com/MarkNahabedian/AnotherParser.jl",
)


text:
using Test, Distributed

# TEMPORARY worker processes for test
workers = addprocs(4)

using DataLoaders

include("sampler.jl")
include("loader.jl")

# MUST remove workers
rmprocs(workers)


text:
import DSP.resample
export resample, resample!

function cheap_resample!(t::Array{Int64, 2}, x::FloatArray, fs_new::Float64, fs_old::Float64)
  r = fs_new/fs_old
  n_seg = size(t,1)-1
  gap_inds = zeros(Int64, n_seg+1)

  # resize S.x if we're upsampling
  if (r > 1.0)
    resize!(x, ceil(Int64, length(x)*r))
  end

  for k = n_seg:-1:1

    # indexing
    si            = t[k,1]
    ei            = t[k+1,1] - (k == n_seg ? 0 : 1)
    nx_in         = ei-si+1

    xr = DSP.resample(x[si:ei], r)

    # indexing
    nx_out        = min(floor(Int64, nx_in*r), length(xr))
    gap_inds[k+1] = nx_out

    # resample and copy
    copyto!(x, si, xr, 1, nx_out)

    # resize S.x if we downsampled
    (fs_new < fs_old) && (deleteat!(x, si+nx_out:ei))
  end

  for k = 2:n_seg+1
    gap_inds[k] += gap_inds[k-1]
  end
  copyto!(t, 2, gap_inds, 2, n_seg)

  # ensure length(S.x[i]) == S.t[i][end,1] if upsampled
  (r > 1.0) && resize!(x, t[end, 1])
  return nothing
end

@doc """
    resample!(S::SeisData [, chans=CC, fs=FS])
    resample(S::SeisData [, chans=CC, fs=FS])

Resample data in S to `FS`. If keyword `fs` is not specified, data are resampled
to the lowest non-zero value in `S.fs[CC]`.Note that a poor choice of `FS` can
lead to upsampling and other undesirable behavior.

Use keyword `chans=CC` to only resample channel numbers `CC`. By default, all
channels `i` with `S.fs[i] > 0.0` are resampled.

    resample!(C::SeisChannel, fs::Float64)
    resample(C::SeisChannel, fs::Float64)

Resample `C.x` to `fs`.
""" resample!
function resample!(S::GphysData;
  chans::ChanSpec=Int64[],
  fs::Float64=0.0)

  chans     = mkchans(chans, S, keepirr=false)
  f0        = fs == 0.0 ? minimum(S.fs[S.fs .> 0.0]) : fs
  proc_str  = string("resample!(S, chans=", chans, ", fs=",
              repr("text/plain", f0, context=:compact=>true), ")")

  for i = 1:S.n
    (S.fs[i] == 0.0)  && continue
    (S.fs[i] == f0)   && continue
    cheap_resample!(S.t[i], S.x[i], f0, S.fs[i])
    desc_str  = string("resampled from ", S.fs[i], " to ", f0, "Hz")
    proc_note!(S, i, proc_str, desc_str)
    S.fs[i]   = f0
  end
  return nothing
end

function resample!(C::GphysChannel, f0::Float64)
  C.fs > 0.0 || error("Can't resample non-timeseries data!")
  (C.fs == f0) && return nothing
  @assert f0 > 0.0
  proc_str = string("resample!(C, fs=",
              repr("text/plain", f0, context=:compact=>true), ")")
  cheap_resample!(C.t, C.x, f0, C.fs)
  desc_str = string("resampled from ", C.fs, " to ", f0, "Hz")
  proc_note!(C, proc_str, desc_str)
  C.fs = f0
  return nothing
end

@doc (@doc resample!)
function resample(S::GphysData;
  chans::ChanSpec=Int64[],
  fs::Float64=0.0)

  U = deepcopy(S)
  resample!(U, chans=chans, fs=fs)
  return U
end

function resample(C::GphysChannel, f0::Float64)
  U = deepcopy(C)
  resample!(U, f0)
  return U
end


text:
#= Julia Wrapper for MAFFT (http://mafft.cbrc.jp/alignment/software/)
=#
module Mafft
export mafft, mafft_from_string, mafft_from_fasta, mafft_profile, mafft_profile_from_string, mafft_profile_from_fasta, print_fasta, alignment_length, to_aminoacids

    using FastaIO
    using BioSeq

    # preconfigurations
    default = ["--auto"]

    # Accuracy-oriented methods

    #= L-INS-i (probably most accurate; recommended for <200 sequences;
           iterative refinement method incorporating local pairwise alignment
           information)
    =#
    linsi = ["--localpair", "--maxiterate", "1000"]

    #= G-INS-i (suitable for sequences of similar lengths; recommended for
           <200 sequences; iterative refinement method incorporating global
           pairwise alignment information)
    =#
    ginsi = ["--globalpair", "--maxiterate", "1000"]

    #= E-INS-i (suitable for sequences containing large unalignable regions;
           recommended for <200 sequences)
    =#
    einsi = ["--ep", "0", "--genafpair", "--maxiterate", "1000"]

    # Speed-oriented methods
    # FFT-NS-i (iterative refinement method; two cycles only)
    fftnsi = ["--retree", "2", "--maxiterate", "2"]

    # FFT-NS-2 (fast; progressive method)
    fftns = ["--retree", "2", "--maxiterate", "0"]

    #= NW-NS-i (iterative refinement method without FFT approximation; two
           cycles only)
    =#
    nwnsi = ["--retree", "2", "--maxiterate", "2", "--nofft"]

    # NW-NS-2 (fast; progressive method without the FFT approximation)
    nwns = ["--retree", "2", "--maxiterate", "0", "--nofft"]

    #= calls MAFFT and returns aligned FASTA
       fasta_in: path to FASTA file
       preconfiguration: optional commandline arguments for MAFFT (array of strings)
    =#
    function mafft(fasta_in::AbstractString, preconfiguration=:default)
        try success(`mafft --version`)
        catch
            error("MAFFT is not installed.")
        end

        fasta = readall(`mafft '--quiet' $(eval(preconfiguration)) $fasta_in`)
        fr = readall(FastaReader(IOBuffer(fasta)))
        return fr
    end

    #= calls MAFFT with the given FASTA string as input and returns aligned FASTA
       fasta_in: FASTA string
       preconfiguration: optional commandline arguments for MAFFT (array of strings)
    =#
    function mafft_from_string(fasta_in::AbstractString, preconfiguration=:default)
        # write to tempfile because mafft can not read from stdin
        tempfile_path, tempfile_io = mktemp()
        write(tempfile_io, fasta_in)
        close(tempfile_io)
        return mafft(tempfile_path, preconfiguration)
    end

    #= calls MAFFT with the given FASTA in FastaIO format
       fasta_in: FASTA in FastaIO format
       preconfiguration: optional commandline arguments for MAFFT (array of strings)
    =#
    function mafft_from_fasta(fasta_in, preconfiguration=:default)
        io = IOBuffer()
        writefasta(io, fasta_in)
        return mafft_from_string(takebuf_string(io), preconfiguration)
    end

    #= Group-to-group alignments with input as paths to FASTA files
       group1 and group2 have to be files with alignments
    =#
    function mafft_profile(group1::AbstractString, group2::AbstractString)
        try success(`mafft --version`)
        catch
            error("MAFFT is not installed.")
        end

        fasta = readall(`mafft '--quiet' --maxiterate 1000 --seed $group1 --seed $group2 /dev/null`)
        fr = readall(FastaReader(IOBuffer(fasta)))
        return fr
    end

    #= Group-to-group alignments with input strings in FASTA format
       group1 and group2 have to be strings with alignments in FASTA format
    =#
    function mafft_profile_from_string(group1::AbstractString, group2::AbstractString)
        # write to tempfiles because mafft can not read from stdin
        tempfile1_path, tempfile1_io = mktemp()
        write(tempfile1_io, group1)
        close(tempfile1_io)

        tempfile2_path, tempfile2_io = mktemp()
        write(tempfile2_io, group2)
        close(tempfile2_io)

        return mafft_profile(tempfile1_path, tempfile2_path)
    end

    #= Group-to-group alignments with input in FastaIO format
       group1 and group2 have to be in FastaIO format and have to be alignments
    =#
    function mafft_profile_from_fasta(group1, group2)
        io = IOBuffer()
        writefasta(io, group1)
        group1_string = takebuf_string(io)
        writefasta(io, group2)
        group2_string = takebuf_string(io)

        return mafft_profile_from_string(group1_string, group2_string)
    end

    # helper methods for aligned FASTA

    function print_fasta(fasta)
        for f in fasta
            println(">", f[1])
            println(f[2])
        end
    end

    # returns the length of the alignment (FastaIO-format as input)
    function alignment_length(fasta)
        return length(first(fasta)[2])
    end

    # converts a FastaIO-formatted array into an array of BioSeq AminoAcid
    function to_aminoacids(fasta)
        aminoacids = Array{AminoAcid,1}[]
        for f in fasta
            push!(aminoacids, aminoacid(f[2]))
        end
        return aminoacids
    end
end


text:
using DelayedKalmanFilter
using Documenter

makedocs(;
    modules=[DelayedKalmanFilter],
    authors="Joshua Burton",
    repo="https://github.com/burtonjosh/DelayedKalmanFilter.jl/blob/{commit}{path}#L{line}",
    sitename="DelayedKalmanFilter.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://burtonjosh.github.io/DelayedKalmanFilter.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
    ],
)

deploydocs(;
    repo="github.com/burtonjosh/DelayedKalmanFilter.jl",
)