using DSGE, ModelConstructors, CSV, DataFrames, Dates, Test

# Data vintages and settings
fcast_date      = DSGE.quartertodate("2021-Q1")
date_fcast_end  = iterate_quarters(fcast_date, 40)
cond_type       = :full

# Flexible AIT Rule
tvcred_n_qtrs          = 24 # length of tv cred period
start_tvcred_level     = 0.
end_tvcred_level       = 1.
φ_π                    = 15.
φ_y                    = 3.
Thalf                  = 10.
ρ_smooth               = .75
pgap_init              = 0.125 # NOTE: positive pgap and ygap values imply
ygap_init              = 12.   # NEGATIVE initial pgap and ygaps (i.e. ygap_init = 12 implies
pgap_ygap_init_date    = Date(2020, 6, 30)  # an initial ygap of -12 in 2020:Q2)
start_zlb_date         = Date(2020, 12, 31) # start and end dates should be inclusive, i.e. first period w/zero rate
end_zlb_date           = Date(2023, 9, 30)  # is start_zlb_date, and the last period is end_zlb_date

# Initialize model objects
m = Model1002("ss62")
m <= Setting(:date_forecast_start,  fcast_date)
m <= Setting(:date_conditional_end, fcast_date)
m <= Setting(:forecast_horizons, subtract_quarters(date_fcast_end, fcast_date))

# Load data
df = CSV.read(joinpath(dirname(@__FILE__), "..", "test", "reference", "ss62_modal_forecast_data.csv"), DataFrame)

# Set up Flexible AIT with a temporary ZLB. Note that df will be altered
start_reg_baseline = 6

θ = (φ_π = φ_π, φ_y = φ_y, Thalf = Thalf, pgap = pgap_init,
     ygap = ygap_init, ρ_smooth = ρ_smooth,
     historical_policy = default_policy()) # NamedTuple with parameters of the Flexible AIT policy

m <= Setting(:flexible_ait_φ_π, θ[:φ_π])
m <= Setting(:flexible_ait_φ_y, θ[:φ_y])
m <= Setting(:ait_Thalf, θ[:Thalf])
m <= Setting(:gdp_Thalf, θ[:Thalf])
m <= Setting(:pgap_value, θ[:pgap])
m <= Setting(:pgap_type, :flexible_ait)
m <= Setting(:ygap_value, θ[:ygap])
m <= Setting(:ygap_type, :flexible_ait)
m <= Setting(:flexible_ait_ρ_smooth, θ[:ρ_smooth])
m <= Setting(:alternative_policies, AltPolicy[θ[:historical_policy]])

# Define a function to help update the mappings
# from model regimes to parameter regimes
function set_regime_vals_fnct(m, n)
    m2p = get_setting(m, :model2para_regime)

    for (k, v) in m2p
        for i in 7:n
            v[i] = 1 # For regimes 7 onward, use parameter regime 1
        end
    end

    m
end

# Set up credibility for ZLB and alternative policy
m <= Setting(:gensys2, true)
m <= Setting(:replace_eqcond, true)
m <= Setting(:uncertain_temporary_altpolicy, true)
m <= Setting(:uncertain_altpolicy, true)

# Find the first regime in which gensys2 should apply
gensys2_first_regime = findfirst([get_setting(m, :regime_dates)[i] .== start_zlb_date for i in 1:get_setting(m, :n_regimes)])
n_zlb_reg = DSGE.subtract_quarters(end_zlb_date, start_zlb_date) + 1
if end_zlb_date < get_setting(m, :regime_dates)[get_setting(m, :n_regimes)]
    set_regime_vals_fnct(m, get_setting(m, :n_regimes))
else
    # technically the second input is `gensys2_first_regime + (n_zlb_reg - 1) + 1` b/c
    # gensys2_first_regime is the first regime for the temporary ZLB, so we need to subtract 1
    # from n_zlb_reg, but we need to also add 1 for the lift-off regime
    set_regime_vals_fnct(m, gensys2_first_regime + n_zlb_reg)
end

# Construct regime_eqcond_info
reg_dates = deepcopy(get_setting(m, :regime_dates))
regime_eqcond_info = Dict{Int, EqcondEntry}()
weights = [end_tvcred_level, 1. - end_tvcred_level] # dummy values for now, time-varying credibility added later
for (regind, date) in zip(gensys2_first_regime:(n_zlb_reg - 1 + gensys2_first_regime), # See comments starting at line 75
                          DSGE.quarter_range(reg_dates[gensys2_first_regime], # looping over the date range of the ZLB
                                                iterate_quarters(reg_dates[gensys2_first_regime], n_zlb_reg - 1)))
    reg_dates[regind] = date # updating the regime date dictionary
    regime_eqcond_info[regind] = EqcondEntry(zlb_rule(), weights) # specify which regimes have a ZLB rule applying
end
reg_dates[n_zlb_reg + gensys2_first_regime] = iterate_quarters(reg_dates[gensys2_first_regime], n_zlb_reg) # add liftoff regime date
regime_eqcond_info[n_zlb_reg + gensys2_first_regime] = EqcondEntry(flexible_ait(), weights) # liftoff policy is flexible AIT
m <= Setting(:regime_dates,               reg_dates)
m <= Setting(:regime_eqcond_info,         regime_eqcond_info)
m <= Setting(:temporary_altpolicy_length, n_zlb_reg) # length of the temporary ZLB, otherwise a guess is made. Since we have time-varying credibility, this guess would be wrong
m <= Setting(:zlb_rule_value,             .1) # (optional) set the interest rate during the ZLB, defaults to forecast_zlb_value(m) if not available
m <= Setting(:temporary_altpolicy_names,  [:zlb_rule]) # Ensures compute_system knows the regimes w/ a zlb_rule in EqcondEntry are gensys2 regimes

# Set up regime indices
setup_regime_switching_inds!(m; cond_type = cond_type)

# Set up time-varying credibility (which extends beyond the ZLB, so we need to add more regimes)
tvcred_dates = (start_zlb_date, DSGE.iterate_quarters(start_zlb_date, tvcred_n_qtrs))
n_tvcred_reg = DSGE.subtract_quarters(tvcred_dates[2], tvcred_dates[1]) + 1 # number of regimes for time-varying credibility
reg_start = DSGE.subtract_quarters(tvcred_dates[1], start_zlb_date) + gensys2_first_regime # first regime for time-varying credibility

# Add extra entries for regime_eqcond_info with dummy credibility weights
for (regind, date) in zip(reg_start:(reg_start + n_tvcred_reg),
                          DSGE.quarter_range(reg_dates[reg_start], DSGE.iterate_quarters(reg_dates[reg_start], n_tvcred_reg - 1)))
    reg_dates[regind] = date
    if date > end_zlb_date # Add extra regimes with flexible AIT as the policy after the ZLB's end b/c credibility is still changing
        regime_eqcond_info[regind] = DSGE.EqcondEntry(flexible_ait(), [end_tvcred_level, 1. - end_tvcred_level])
    end
end
m <= Setting(:regime_dates,             reg_dates)
m <= Setting(:regime_eqcond_info, regime_eqcond_info)
setup_regime_switching_inds!(m; cond_type = cond_type)

# Update mapping of model regimes to parameter regimes
set_regime_vals_fnct(m, get_setting(m, :n_regimes))

# Set up TVIS information set
m <= Setting(:tvis_information_set, vcat([i:i for i in 1:(gensys2_first_regime - 1)],
                                         [i:17 for i in gensys2_first_regime:17],
                                         [i:i for i in 18:get_setting(m, :n_regimes)]))

# Now actually populate the weights in regime_eqcond_info so that
# credibility/awareness increases linearly over time
credvec = collect(range(start_tvcred_level, stop = end_tvcred_level, length = n_tvcred_reg))
for (i, k) in enumerate(sort!(collect(keys(regime_eqcond_info))))
    if tvcred_dates[1] <= reg_dates[k] <= tvcred_dates[2]
        get_setting(m, :regime_eqcond_info)[k].weights = [credvec[i], 1. - credvec[i]]
    end
end

# We implement a shorter ZLB with a return to the historical
# monetary policy reaction function as the other policy
# agents believe could occur. This is where the `MultiPeriodAltPolicy` comes into play.
altpol_eqcond_info = deepcopy(get_setting(m, :regime_eqcond_info))
for (reg, eq_entry) in altpol_eqcond_info
    if reg >= start_reg_baseline + 5
        altpol_eqcond_info[reg] = DSGE.EqcondEntry(DSGE.default_policy(), [1.])
    end
    altpol_eqcond_info[reg].weights = [1.]
end

# The key for the MultiPeriodAltPolicy can be anything; it is never used in the code
# to infer some property. Of course, it should still be a descriptive name.
altpol = MultiPeriodAltPolicy(:other_altpol, get_setting(m, :n_regimes), altpol_eqcond_info, gensys2 = true,
                              temporary_altpolicy_names = [:zlb_rule],
                              temporary_altpolicy_length = 6,
                              infoset = copy(get_setting(m, :tvis_information_set)))
delete!(m.settings, :alternative_policies) # if this setting already exists, the type may be wrong
m <= Setting(:alternative_policies, DSGE.AbstractAltPolicy[altpol])

# Run forecasts
output_vars = [:histobs, :histpseudo, :forecastpseudo, :forecastobs]
modal_params = map(x -> x.value, m.parameters)
fcast = DSGE.forecast_one_draw(m, :mode, cond_type, output_vars, modal_params, df,
                               regime_switching = true, n_regimes = get_setting(m, :n_regimes))

nothing


# Data topmute
# Author: Philipp Witte, pwitte@eos.ubc.ca
# Date: December 2017
#

export marineTopmute2D, judiMarineTopmute2D
export model_topmute, judiTopmute, find_water_bottom, depth_scaling, judiDepthScaling, laplace

############################################ Data space preconditioners ################################################

function marineTopmute2D(Dobs::judiVector, muteStart::Integer; mute=Array{Any}(undef, 3), flipmask=false)
    # Data topmute for end-on spread marine streamer data
    Din = deepcopy(Dobs)

    for j=1:Din.nsrc

        # Design mute window
        x0 = 1f0
        xend = length(Din[j].geometry.xloc[1])
        nt = Din[j].geometry.nt[1]
        nrec = length(Din[j].geometry.xloc[1])
        drec = abs(Din[j].geometry.xloc[1][1] - Din[j].geometry.xloc[1][2])
        offsetDirectWave = 1.5f0*Din[j].geometry.t[1]
        idxOffset = Int(round(offsetDirectWave/drec))
        dx = round(idxOffset - idxOffset/10f0)

        if j==1 && ~isassigned(mute)
            z0 = muteStart - Int(round(muteStart/10))
            slope = 1.05f0*(nt - z0)/dx
            mute[1] = x0
            mute[2] = z0
            mute[3] = slope
        else#if j==1 && isassigned(mute)
            x0 = mute[1]
            z0 = mute[2]
            slope = mute[3]
        end

        mask = ones(Float32,nt,nrec)
        mask[1:z0,:] .= 0f0

        # Linear mute
        if (nrec-x0 < dx)
            x = nrec
            zIntercept = Int(round(z0+slope*(x-x0)))
            zax = z0+1:1:zIntercept
        else
            x = x0+dx
            zax = z0+1:1:nt
        end
        if length(zax) > 1
            xax = Array{Int}(round.(range(x0,stop=x,length=length(zax))))
        else
            xax = Int(round(x0))
        end
        for k=1:length(zax)
            mask[zax[k],xax[k]:end] .= 0f0
        end
        flipmask == true && (mask = reverse(mask, dims=2))
        Din.data[j] = Din.data[j].*mask
    end
    return Din
end

function judiMarineTopmute2D(muteStart,geometry;params=Array{Any}(undef, 3),flipmask=false)
# JOLI wrapper for the linear depth scaling function
    nsrc = length(geometry.xloc)
    N = 0
    for j=1:nsrc
        N += geometry.nt[j]*length(geometry.xloc[j])
    end
    D = joLinearFunctionFwd_T(N,N,
                             v -> marineTopmute2D(v,muteStart;mute=params, flipmask=flipmask),
                             w -> marineTopmute2D(w,muteStart;mute=params, flipmask=flipmask),
                             Float32,Float32,name="Data mute")
    return D
end

############################################ Model space preconditioners ###############################################

function model_topmute(n::Tuple{Int64,Int64}, mute_end::Array{Integer,1}, length::Int64, x_orig)
# Model domain topmute for a velocity model of dimensions n = [nx, nz].
    x = copy(reshape(x_orig,n))
    for j=1:n[1]
        mute_start = mute_end[j] - length
        filter = zeros(Float32, n[2])
        filter[1:mute_start-1] .= 0f0
        filter[mute_end[j]+1:end] .= 1f0
        taper_length = mute_end[j] - mute_start + 1
        taper = (1f0 .+ sin.((Float32(pi)*(0:taper_length-1))/(taper_length - 1) .- Float32(pi)/2f0))/2f0
        filter[mute_start:mute_end[j]] = taper
        M = diagm(0=>filter)
        global x[j,:] = x[j,:].*filter
    end
    return vec(x)
end

function model_topmute(n::Tuple{Int64,Int64}, mute_end::Int64, length::Int64, x_orig)
# Model domain topmute for a velocity model of dimensions n = [nx, nz].
    x = copy(reshape(x_orig,n))
    mute_start = mute_end - length
    filter = zeros(Float32, n[2])
    filter[1:mute_start-1] .= 0f0
    filter[mute_end+1:end] .= 1f0
    taper_length = mute_end - mute_start + 1
    taper = (1f0 .+ sin.((Float32(pi)*(0:taper_length-1))/(taper_length - 1) .- Float32(pi)/2f0))/2f0
    filter[mute_start:mute_end] = taper
    M = diagm(0=>filter)
    for j=1:n[1]
        global x[j,:] = x[j,:].*filter
    end
    return vec(x)
end

model_topmute(n::Tuple{Int64,Int64}, mute_end::Array{Float32, 2}, length, x) = vec(mute_end) .* vec(x)

function judiTopmute(n, mute_start, length)
    # JOLI wrapper for model domain topmute
    N = prod(n)
    T = joLinearFunctionFwd_T(N,N,
                             v -> model_topmute(n, mute_start, length, v),
                             w -> model_topmute(n, mute_start, length, w),
                             Float32,Float32,name="Model topmute")
    return T
end

function find_water_bottom(m)
    #return the indices of the water bottom of a seismic image
    n = size(m)
    idx = zeros(Integer,n[1])
    eps = 1e-4
    for j=1:n[1]
        k=1
        while true
            if abs(m[j,k]) > eps
                idx[j] = k
                break
            end
            k += 1
        end
    end
    return idx
end

function depth_scaling(m,model)
# Linear depth scaling function for seismic images
    m = reshape(m,model.n)
    filter = sqrt.(0f0:model.d[2]:model.d[2]*(model.n[2]-1))
    F = diagm(0=>filter)
    for j=1:model.n[1]
        m[j,:] = F*m[j,:]
    end
    return vec(m)
end

function judiDepthScaling(model)
# JOLI wrapper for the linear depth scaling function
    N = prod(model.n)
    D = joLinearFunctionFwd_T(N, N,
                             v -> depth_scaling(v,model),
                             w -> depth_scaling(w,model),
                             Float32,Float32,name="Depth scaling")
end

# function laplace(model::TimeModeling.Model)
# # 2D Laplace operator
#
#     # 2nd derivative in x direction
#     d1 = ones(Float32,model.n[1]-1)*1f0/(model.d[1]^2)
#     d2 = ones(Float32,model.n[1])*-2f0/(model.d[1]^2)
#     d3 = d1
#     d = (d1, d2, d3)
#     position = (-1, 0, 1)
#
#     Dx = spdiagm(d,position,model.n[1],model.n[1])
#     Ix = speye(Float32,model.n[1])
#
#     # 2nd derivative in z direction
#     d1 = ones(Float32,model.n[2]-1)*1f0/(model.d[2]^2)
#     d2 = ones(Float32,model.n[2])*-2f0/(model.d[2]^2)
#     d3 = d1
#     d = (d1, d2, d3)
#     position = (-1, 0, 1)
#
#     Dz = spdiagm(d,position,model.n[2],model.n[2])
#     Iz = speye(Float32,model.n[2])
#
#     # 2D Laplace operator
#     D = kron(Dz,Ix) + kron(Iz,Dx)
# end


# Matrix type
@cenum(MatrixType,
    REAL_SYM            = 1,
    REAL_SYM_POSDEF     = 2,
    REAL_SYM_INDEF      = -2,
    COMPLEX_STRUCT_SYM  = 3,
    COMPLEX_HERM_POSDEF = 4,
    COMPLEX_HERM_INDEF  = -4,
    COMPLEX_SYM         = 6,
    REAL_NONSYM         = 11,
    COMPLEX_NONSYM      = 13,
)

Base.isreal(v::MatrixType) = v in (REAL_SYM, REAL_SYM_POSDEF, REAL_SYM_INDEF, REAL_NONSYM)
issymmetric(v::MatrixType) = v in (REAL_SYM, REAL_SYM_POSDEF, REAL_SYM_INDEF, COMPLEX_STRUCT_SYM,
                                        COMPLEX_HERM_POSDEF, COMPLEX_HERM_INDEF, COMPLEX_SYM)
Base.ishermitian(v::MatrixType) = v in (REAL_SYM_POSDEF, COMPLEX_HERM_POSDEF, COMPLEX_HERM_INDEF)
isposornegdef(v::MatrixType) = v in (REAL_SYM_POSDEF, REAL_SYM_INDEF, COMPLEX_HERM_POSDEF, COMPLEX_HERM_INDEF)

const MATRIX_STRING = Dict{MatrixType, String}(
    REAL_SYM => "Real structurally symmetric",
    REAL_SYM_POSDEF => "Real symmetric positive definite",
    REAL_SYM_INDEF => "Real symmetric indefinite",
    COMPLEX_STRUCT_SYM => "Complex structurally symmetric",
    COMPLEX_HERM_POSDEF => "Complex Hermitian postive definite",
    COMPLEX_HERM_INDEF => "Complex Hermitian indefinite",
    COMPLEX_SYM  => "Complex symmetric",
    REAL_NONSYM  => "Real nonsymmetric",
    COMPLEX_NONSYM  => "Complex nonsymmetric"
)

const REAL_MATRIX_TYPES = [REAL_SYM, REAL_SYM_POSDEF, REAL_SYM_INDEF, REAL_NONSYM]
const COMPLEX_MATRIX_TYPES = [COMPLEX_STRUCT_SYM, COMPLEX_HERM_POSDEF, COMPLEX_HERM_INDEF, COMPLEX_NONSYM]

# Messages
@cenum(MessageLevel,
    MESSAGE_LEVEL_OFF = 0,
    MESSAGE_LEVEL_ON = 1
)

# Solver
@cenum(Solver,
    DIRECT_SOLVER = 0,
    ITERATIVE_SOLVER = 1
)

const SOLVER_STRING = Dict{Solver, String}(
    DIRECT_SOLVER => "Direct solver",
    ITERATIVE_SOLVER => "Iterative solver"
)

# Phase
@cenum(Phase,
    ANALYSIS                             = 11,
    ANALYSIS_NUM_FACT                    = 12,
    ANALYSIS_NUM_FACT_SOLVE_REFINE       = 13,
    NUM_FACT                             = 22,
    SELECTED_INVERSION                   = -22,
    NUM_FACT_SOLVE_REFINE                = 23,
    SOLVE_ITERATIVE_REFINE               = 33,
    SOLVE_ITERATIVE_REFINE_ONLY_FORWARD  = 331,
    SOLVE_ITERATIVE_REFINE_ONLY_DIAG     = 332,
    SOLVE_ITERATIVE_REFINE_ONLY_BACKWARD = 333,
    RELEASE_LU_MNUM                      = 0,
    RELEASE_ALL                           = -1
)

const PHASE_STRING = Dict{Phase, String}(
    ANALYSIS                             => "Analysis",
    ANALYSIS_NUM_FACT                    => "Analysis, numerical factorization",
    ANALYSIS_NUM_FACT_SOLVE_REFINE       => "Analysis, numerical factorization, solve, iterative refinement",
    NUM_FACT                             => "Numerical factorization",
    SELECTED_INVERSION                   => "Selected Inversion",
    NUM_FACT_SOLVE_REFINE                => "Numerical factorization, solve, iterative refinement",
    SOLVE_ITERATIVE_REFINE               => "Solve, iterative refinement",
    RELEASE_LU_MNUM                      => "Release internal memory for L and U matrix number MNUM",
    RELEASE_ALL                           => "Release all internal memory for all matrices",
    SOLVE_ITERATIVE_REFINE_ONLY_FORWARD  => "like phase=33, but only forward substitution",
    SOLVE_ITERATIVE_REFINE_ONLY_DIAG     => "like phase=33, but only diagonal substitution (if available)",
    SOLVE_ITERATIVE_REFINE_ONLY_BACKWARD => "like phase=33, but only backward substitution",
)


"""
Reference: The Annotated Transformer (http://nlp.seas.harvard.edu/2018/04/03/attention.html)
"""

using ArgParse

using Flux
using Flux: onecold, gradient
import Flux.Optimise: update!

using WordTokenizers

using Transformers
using Transformers.Basic
using Transformers.Datasets
using Transformers.Datasets: WMT, IWSLT

function parse_commandline()
    s = ArgParseSettings()

    @add_arg_table s begin
        "--gpu", "-g"
            help = "use gpu"
            action = :store_true
        "task"
            help = "task name"
            required = true
            range_tester = x-> x ∈ ["wmt14", "iwslt2016", "copy"]
    end

    return parse_args(ARGS, s)
end

const args = parse_commandline()

if args["gpu"]
    @eval using CuArrays
end

if args["task"] == "copy"
    const N = 2
    const V = 10
    const Smooth = 1e-6
    const Batch = 32
    const lr = 1e-4

    startsym = 11
    endsym = 12
    unksym = 0
    labels = [unksym, startsym, endsym, collect(1:V)...]

    function gen_data()
        global V
        d = rand(1:V, 10)
        (d,d)
    end

    function train!()
        global Batch
        println("start training")
        i = 1
        for i = 1:320*7
            data = batched([gen_data() for i = 1:Batch])
            x, t = data
            x = mkline.(x)
            t = mkline.(t)
            x_mask = getmask(x)
            t_mask = getmask(t)
            x, t = embed.Vocab(x, t)
            x, t, x_mask, t_mask = todevice(x,t,x_mask,t_mask)
            l = loss(x, t, x_mask, t_mask)
            grad = gradient(()->l, ps)
            i%8 == 0 && @show l
            update!(opt, ps, grad)
        end
    end

    mkline(x) = [startsym, x..., endsym]
elseif args["task"] == "wmt14" || args["task"] == "iwslt2016"
    const N = 6
    const Smooth = 0.4
    const Batch = 8
    const lr = 1e-6
    const MaxLen = 100

    const task = args["task"]

    if task == "wmt14"
        wmt14 = WMT.GoogleWMT()

        datas = dataset(Train, wmt14)
        vocab = get_vocab(wmt14)
    else
        iwslt2016 = IWSLT.IWSLT2016(:en, :de)

        datas = dataset(Train, iwslt2016)
        vocab = get_vocab(iwslt2016)
    end

    startsym = "<s>"
    endsym = "</s>"
    unksym = "</unk>"
    labels = [unksym, startsym, endsym, collect(keys(vocab))...]

    function train!()
        global Batch
        println("start training")
        i = 1
        while (batch = get_batch(datas, Batch)) != []
            x = mkline.(batch[1])
            t = mkline.(batch[2])
            x_mask = getmask(x)
            t_mask = getmask(t)
            x, t = embed.Vocab(x, t)
            x, t, x_mask, t_mask = todevice(x,t,x_mask,t_mask)
            l = loss(x,t, x_mask, t_mask)
            grad = gradient(()->l, ps)
            i+=1
            #i%5 == 0 &&
            (@show l; @time update!(opt, ps, grad))
        end
    end

    if task == "wmt14"
        function mkline(x)
            global MaxLen
            xi = split(x)
            if length(xi) > MaxLen
                xi = xi[1:100]
            end

            [startsym, xi..., endsym]
        end
    else
        function mkline(x)
            global MaxLen
            xi = tokenize(x)
            if length(xi) > MaxLen
                xi = xi[1:100]
            end

            [startsym, xi..., endsym]
        end
    end
else
    error("task not define")
end

vocab = Vocabulary(labels, unksym)
const embed = gpu(Embed(512, vocab))

embedding(x) = embed(x, inv(sqrt(512)))

const encoder = gpu(Stack(
    @nntopo(e → pe:(e, pe) → x → x → $N),
    PositionEmbedding(512),
    (e, pe) -> e .+ pe,
    Dropout(0.1),
    [Transformer(512, 8, 64, 2048) for i = 1:N]...
))

const decoder = gpu(Stack(
    @nntopo((e, m, mask):e → pe:(e, pe) → t → (t:(t, m, mask) → t:(t, m, mask)) → $N:t → c),
    PositionEmbedding(512),
    (e, pe) -> e .+ pe,
    Dropout(0.1),
    [TransformerDecoder(512, 8, 64, 2048) for i = 1:N]...,
    Positionwise(Dense(512, length(labels)), logsoftmax)
))

const ps = params(embed, encoder, decoder)
const opt = ADAM(lr)

function smooth(et)
    global Smooth
    sm = fill!(similar(et, Float32), Smooth/length(embed.Vocab))
    p = sm .* (1 .+ -et)
    label = p .+ et .* (1 - convert(Float32, Smooth))
    label
end

function loss(src, trg, src_mask, trg_mask)
    lab = onehot(embed, trg)

    src = embedding(src)
    trg = embedding(trg)

    if src_mask === nothing || trg_mask === nothing
        mask = nothing
    else
        mask = getmask(src_mask, trg_mask)
    end

    enc = encoder(src)
    dec = decoder(trg, enc, mask)

    #label smoothing
    label = smooth(lab)[:, 2:end, :]

    if mask == nothing
        loss = logkldivergence(label, dec[:, 1:end-1, :])
    else
        loss = logkldivergence(label, dec[:, 1:end-1, :], trg_mask[:, 1:end-1, :])
    end
end

function translate(x)
    ix = todevice(embed.Vocab(mkline(x)))
    seq = [startsym]

    src = embedding(ix)
    enc = encoder(src)

    len = length(ix)
    for i = 1:2len
        trg = embedding(todevice(embed.Vocab(seq)))
        dec = decoder(trg, enc, nothing)
        #move back to gpu due to argmax wrong result on CuArrays
        ntok = onecold(collect(dec), labels)
        push!(seq, ntok[end])
        ntok[end] == endsym && break
    end
    seq
end

train!()


# call with
# ARGS[1] = T
# ARGS[2] = frac err
# ARGS[3] = nameOutput
# ARGS[4] = port1
# ARGS[5] = port2
@show ARGS

using Caesar, IncrementalInference, RoME, HDF5, JLD
d = jldopen("data/VicPrkBasic.jld", "r") do file
  read(file, "dBasic")
end
f = jldopen("data/VicPrkBasic.jld", "r") do file
  read(file, "fBasic")
end
# MM = jldopen("test/data/VicPrkBasic.jld", "r") do file
#   read(file, "MMBasic")
# end
include("data/Vic120MM.jl")
MMr = MMwithRedirects(MM);
frac = parse(Float64,ARGS[2])
MMe = MMRandErrs(MMr,frac=frac)
T=parse(Int,ARGS[1])

isamdict, fgu = doISAMSolve(d,f,toT=T, savejld=false, MM=MMe, retfg=true, port1=parse(Int,ARGS[4]),port2=parse(Int,ARGS[5]))

# isamdict = jldopen("data/fgWithISAMrefERR01.jld", "r") do file
#   read(file, "isamdict")
# end
# fgu = jldopen("data/fgWithISAMrefERR01.jld", "r") do file
#   read(file, "fgu")
# end

fg = emptyFactorGraph();
idx = appendFactorGraph!(fg, d, f, toT=T, lcmode=:mmodal, MM=MMe);
tree = prepBatchTree!(fg, ordering=:qr,drawpdf=true);
inferOverTree!(fg,tree, N=100);
fg1 = deepcopy(fg)
@save "results/fgT$(T)_$(ARGS[3])_perErr$(ARGS[2]).jld" fg1 isamdict fgu MM MMr MMe T frac

tic()
inferOverTree!(fg,tree, N=100);
comptime = toc()
fg2 = deepcopy(fg)
@save "results/fgT$(T)_$(ARGS[3])_perErr$(ARGS[2]).jld" fg1 fg2 isamdict fgu MM MMr MMe T frac comptime

tic()
inferOverTree!(fg,tree, N=100);
comptime = toc()
fg3 = fg
@save "results/fgT$(T)_$(ARGS[3])_perErr$(ARGS[2]).jld" fg1 fg2 fg3 isamdict fgu MM MMr MMe T frac comptime


###
### DynamicHMC backend - https://github.com/tpapp/DynamicHMC.jl
###

"""
    DynamicNUTS

Dynamic No U-Turn Sampling algorithm provided by the DynamicHMC package.

To use it, make sure you have DynamicHMC package (version >= 2) loaded:
```julia
using DynamicHMC
```
"""
struct DynamicNUTS{AD,space} <: Hamiltonian{AD} end

DynamicNUTS(args...) = DynamicNUTS{ADBackend()}(args...)
DynamicNUTS{AD}(space::Symbol...) where AD = DynamicNUTS{AD, space}()

DynamicPPL.getspace(::DynamicNUTS{<:Any, space}) where {space} = space

struct DynamicHMCLogDensity{M<:Model,S<:Sampler{<:DynamicNUTS},V<:AbstractVarInfo}
    model::M
    sampler::S
    varinfo::V
end

function DynamicHMC.dimension(ℓ::DynamicHMCLogDensity)
    return length(ℓ.varinfo[ℓ.sampler])
end

function DynamicHMC.capabilities(::Type{<:DynamicHMCLogDensity})
    return DynamicHMC.LogDensityOrder{1}()
end

function DynamicHMC.logdensity_and_gradient(
    ℓ::DynamicHMCLogDensity,
    x::AbstractVector,
)
    return gradient_logp(x, ℓ.varinfo, ℓ.model, ℓ.sampler)
end

"""
    DynamicNUTSState

State of the [`DynamicNUTS`](@ref) sampler.

# Fields
$(TYPEDFIELDS)
"""
struct DynamicNUTSState{V<:AbstractVarInfo,C,M,S}
    vi::V
    "Cache of sample, log density, and gradient of log density."
    cache::C
    metric::M
    stepsize::S
end

function gibbs_update_state(state::DynamicNUTSState, varinfo::AbstractVarInfo)
    return DynamicNUTSState(varinfo, state.cache, state.metric, state.stepsize)
end

DynamicPPL.initialsampler(::Sampler{<:DynamicNUTS}) = SampleFromUniform()

function DynamicPPL.initialstep(
    rng::AbstractRNG,
    model::Model,
    spl::Sampler{<:DynamicNUTS},
    vi::AbstractVarInfo;
    kwargs...
)
    # Ensure that initial sample is in unconstrained space.
    if !DynamicPPL.islinked(vi, spl)
        DynamicPPL.link!(vi, spl)
        model(rng, vi, spl)
    end

    # Perform initial step.
    results = DynamicHMC.mcmc_keep_warmup(
        rng,
        DynamicHMCLogDensity(model, spl, vi),
        0;
        initialization = (q = vi[spl],),
        reporter = DynamicHMC.NoProgressReport(),
    )
    steps = DynamicHMC.mcmc_steps(results.sampling_logdensity, results.final_warmup_state)
    Q, _ = DynamicHMC.mcmc_next_step(steps, results.final_warmup_state.Q)

    # Update the variables.
    vi[spl] = Q.q
    DynamicPPL.setlogp!(vi, Q.ℓq)

    # If a Gibbs component, transform the values back to the constrained space.
    if spl.selector.tag !== :default
        DynamicPPL.invlink!(vi, spl)
    end

    # Create first sample and state.
    sample = Transition(vi)
    state = DynamicNUTSState(vi, Q, steps.H.κ, steps.ϵ)

    return sample, state
end

function AbstractMCMC.step(
    rng::AbstractRNG,
    model::Model,
    spl::Sampler{<:DynamicNUTS},
    state::DynamicNUTSState;
    kwargs...
)
    # Compute next sample.
    vi = state.vi
    ℓ = DynamicHMCLogDensity(model, spl, vi)
    steps = DynamicHMC.mcmc_steps(
        rng,
        DynamicHMC.NUTS(),
        state.metric,
        ℓ,
        state.stepsize,
    )
    Q = if spl.selector.tag !== :default
        # When a Gibbs component, transform values to the unconstrained space
        # and update the previous evaluation.
        DynamicPPL.link!(vi, spl)
        DynamicHMC.evaluate_ℓ(ℓ, vi[spl])
    else
        state.cache
    end
    newQ, _ = DynamicHMC.mcmc_next_step(steps, Q)

    # Update the variables.
    vi[spl] = newQ.q
    DynamicPPL.setlogp!(vi, newQ.ℓq)

    # If a Gibbs component, transform the values back to the constrained space.
    if spl.selector.tag !== :default
        DynamicPPL.invlink!(vi, spl)
    end

    # Create next sample and state.
    sample = Transition(vi)
    newstate = DynamicNUTSState(vi, newQ, state.metric, state.stepsize)

    return sample, newstate
end


module ALPA
    using StatsBase

export addnode!, delnode!, addedge!, deledge!, ne,
    lpa_addnode!, lpa_deletenode!, lpa_addedge!, lpa_deleteedge!, loadgml, savegml, modularity, graphdiff, Graph, NeighComm, NodeStatus, getctnf
# package code goes here
include("wgraph.jl")
include("wevolve.jl")
include("graphio.jl")
include("utils.jl")
# package code goes here
N = 10000
c = NeighComm(collect(1:N), fill(-1,N), 1)
ns = NodeStatus(collect(1:N), fill(false,N), 1)
random_order = Array{Int}(undef,N)
#m = zeros(Int,N)
a = Set{Int}()

    # package code goes here
    include("wgraph.jl")
    include("wevolve.jl")
    include("graphio.jl")
    include("utils.jl")

    # package code goes here
    N = 10000
    c = NeighComm(collect(1:N), fill(-1,N), 1)
    ns = NodeStatus(collect(1:N), fill(false,N), 1)
    random_order = fill(0, N)
    #m = zeros(Int,N)
    a = BitSet()

    lpa_deletenode!(g, m, u) = lpa_deletenode!(g, m, c, u, a, random_order, ns)

    lpa_addedge!(g,m,u,v) = lpa_addedge!(g,m,c,u,v,a,random_order,ns)

    lpa_deleteedge!(g,m,u,v) = lpa_deleteedge!(g,m,c,u,v,a,random_order,ns)

end # module


module RateLimiter

using Dates

export AbstractRateLimiter, @rate_limit, NoLimitRateLimiter, TokenBucketRateLimiter

#TODO: should the Manifest be committed?

""" Abstract base type for all rate limiters. """
abstract type AbstractRateLimiter end

"""
    !(limiter, cost::T) where {T <: Real}

Wait until the method can be called again.  This method is thread safe.

# Arguments
- `limiter`: rate limiter
- `cost`: cost of executing the method.
"""
wait!(limiter::AbstractRateLimiter, cost::T) where {T <: Real} = error("Function is not defined for AbstractRateLimiter: type=$(typeof(limiter)) function=wait!")

"""
    @rate_limit(limiter, cost, expr)

A macro to limit the rate an expression can be called.

# Arguments
    - `limiter`: rate limiter
    - `cost`: cost of executing the method.
    - `expr`: expression to be evaluated
"""
macro rate_limit(limiter, cost, expr)
    return esc(quote
        RateLimiter.wait!($limiter, $cost)
        $expr
    end)
end

#######################################################################

"""
A rate limiter which does not limit the rate.
"""
struct NoLimitRateLimiter <: AbstractRateLimiter end

"""
    !(limiter, cost::T) where {T <: Real}

Wait until the method can be called again.  This method is thread safe.

# Arguments
- `limiter`: rate limiter
- `cost`: cost of executing the method.
"""
wait!(limiter::NoLimitRateLimiter, cost::T) where {T <: Real} = return

#######################################################################

"""
Token-Bucket rate limiter.  This limiter is used to control for constraints such as bandwidth and burstiness.

See:  https://en.wikipedia.org/wiki/Token_bucket#Algorithm
"""
mutable struct TokenBucketRateLimiter <: AbstractRateLimiter
    tokens_per_second::Float64
    max_tokens::Float64
    lock::ReentrantLock

    last_update::DateTime
    tokens::Float64

    """
    Token-Bucket rate limiter.  This limiter is used to control for constraints such as bandwidth and burstiness.

    See:  https://en.wikipedia.org/wiki/Token_bucket#Algorithm

    # Arguments
    - `tokens_per_second`: number of tokens added per second.
    - `max_tokens`: maximum number of tokens.
    - `initial_tokens`: initial number of tokens.
    - `sleep_seconds`: number of seconds to sleep before checking to see if execution is possible.
    """
    function TokenBucketRateLimiter(tokens_per_second, max_tokens, initial_tokens)
        tokens_per_second <= 0 && error("tokens_per_second must be positive: tokens_per_second=$tokens_per_second")
        max_tokens <= 0 && error("max_tokens must be positive: max_tokens=$max_tokens")
        return new(tokens_per_second, max_tokens, ReentrantLock(), now(), initial_tokens)
    end
end

"""
    !(limiter, cost::T) where {T <: Real}

Wait until the method can be called again.  This method is thread safe.

# Arguments
- `limiter`: rate limiter
- `cost`: cost of executing the method.
"""
function wait!(limiter::TokenBucketRateLimiter, cost::T) where {T <: Real}
    # error check
    cost < 0 && error("Cost must be greater than zero: cost=$cost")
    cost > limiter.max_tokens && error("Cost must be less than the max number of tokens: cost=$cost max_tokens=$(limiter.max_tokens)")

    lock(limiter.lock) do
        while true
            # fill
            t = now()
            limiter.tokens = min(limiter.max_tokens, limiter.tokens + Dates.toms(t-limiter.last_update) / Dates.toms(Second(1)) * limiter.tokens_per_second)
            limiter.last_update = t

            # take
            if limiter.tokens >= cost
                limiter.tokens -= cost
                return
            end

            # retry
            sleep((cost-limiter.tokens)/limiter.tokens_per_second)
        end
    end
end

#######################################################################

end # module


convertdepth(depth::Real) = depth * u"m"
convertdepth(depth::Quantity{U, Unitful.𝐋, V}) where {U,V} = depth
convertdepth(x) = error("Not a valid depth")

"""
    plotscattertransect(t)

Plots a scatter of the discrete obs of `t` in (lat,depth) space.
"""
@userplot PlotScatterTransect
@recipe function f(p::PlotScatterTransect)
    transect, = p.args
    x, y, v = scattertransect(transect)
    @series begin
        seriestype := :scatter
        yflip := true
        marker_z --> v
        markershape --> :circle
        clims --> (0, 1) .* maximum(transect)
        label --> ""
        yguide --> "Depth"
        xguide --> "Distance"
        x, y
    end
end
function scattertransect(t::Transect)
    depths = reduce(vcat, pro.depths for pro in t.profiles)
    values = reduce(vcat, pro.values for pro in t.profiles)
    distances = reduce(vcat, [fill(d, length(pro)) for (d,pro) in zip(vcat(0.0u"km", cumsum(diff(t))), t.profiles)])
    return distances, convertdepth.(depths), values
end
export scattertransect

"""
    plotcruisetrack(t)

Plots the cruise track of `t` in (lat,lon) space.
"""
@userplot PlotCruiseTrack
@recipe function f(p::PlotCruiseTrack, central_longitude=200°)
    wlon = central_longitude - 180°
    ct, = p.args
    ctlon, ctlat = [s.lon for s in ct.stations]°, [s.lat for s in ct.stations]°
    @series begin
        label := ""
        xguide --> ""
        yguide --> ""
        markershape --> :circle
        markersize --> 4
        color_palette --> :tab10
        markercolor --> 4
        markerstrokewidth --> 0
        [mod(ctlon[1] - wlon, 360°) + wlon], [ctlat[1]]
    end
    @series begin
        color_palette --> :tab10
        seriescolor --> 4
        xguide --> ""
        yguide --> ""
        title --> ct.name
        markershape --> :circle
        markersize --> 1
        markercolor --> :black
        titlefontsize --> 10
        linewidth --> 2
        mod.(ctlon .- wlon, 360°) .+ wlon, ctlat
    end
end


function create_funnel_graph(domain_name, problem_name; max_levels=10)
    graph = Graph()
    domain = load_domain(domain_name)
    problem = load_problem(problem_name)
    init_propositions = get_init_propositions(domain, problem)
    ex_constraints = get_external_constraints(domain, problem)
    graph.props[1] = init_propositions
    graph.initprops = init_propositions
    graph.goalprops = get_goal_propositions(domain, problem)
    graph.num_levels = 1

    for _=1:max_levels
        expand!(graph, domain, problem, ex_constraints)
        if goal_reached(graph, domain, problem) break end
    end
    return graph
end

function expand!(graph, domain, problem, constraints)
    k = graph.num_levels
    if k ≥ 2
        graph.μprops[k] = get_proposition_mutexes(graph, k)
    end
    actions = get_all_actions(domain, problem)
    action_list=[]
    graph.props[k+1] = Dict()
    graph.props[k+1][:discrete]=[]
    graph.props[k+1][:continuous]=[]
    graph.acts[k]=[]
    for act in actions
        if is_applicable(act, graph, constraints, k)
            instantiated_action = compute_funnel(act, graph, constraints, k)
            propositions = get_action_propositions(instantiated_action)
            push!(graph.acts[k], instantiated_action)
            if !isempty(propositions[:discrete])
                push!(graph.props[k+1][:discrete], propositions[:discrete]...)
            end
            if !isempty(propositions[:continuous])
                push!(graph.props[k+1][:continuous], propositions[:continuous]...)
            end
            push!(action_list, instantiated_action)
        end
    end
    for prop in graph.props[k][:discrete]
            noop = get_noop_action(prop)
            push!(graph.acts[k], noop)
            push!(graph.props[k+1][:discrete], prop)
    end
    for region in graph.props[k][:continuous]
        if !(region in graph.props[k+1][:continuous]) && region.name == :b1
            mn = get_maintain_funnel(region)
            push!(graph.acts[k], mn)
            push!(graph.props[k+1][:continuous], region)
        end
    end
    graph.props[k+1][:discrete] = collect(Set(graph.props[k+1][:discrete]))
    graph.props[k+1][:continuous] = get_unique_regions(graph.props[k+1][:continuous])
    graph.μacts[k] = get_action_mutexes(graph, k)
    graph.num_levels+=1
    # println(" ")
    return graph
end

function get_unique_regions(regions)
    unique_regions = []
    for reg in regions
        if !any([reg.r[1].ϕ₁ == x.r[1].ϕ₁ && reg.r[1].ϕ₂ == x.r[1].ϕ₂ && reg.r[1].ϕ₃ == x.r[1].ϕ₃ && reg.r[1].þ == x.r[1].þ &&
            reg.r[2].ϕ₁ == x.r[2].ϕ₁ && reg.r[2].ϕ₂ == x.r[2].ϕ₂ && reg.r[2].ϕ₃ == x.r[2].ϕ₃ && reg.r[2].þ == x.r[2].þ &&
            reg.r[3].ϕ₁ == x.r[3].ϕ₁ && reg.r[3].ϕ₂ == x.r[3].ϕ₂ && reg.r[3].ϕ₃ == x.r[3].ϕ₃ && reg.r[3].þ == x.r[3].þ &&
            reg.r[4].ϕ₁ == x.r[4].ϕ₁ && reg.r[4].ϕ₂ == x.r[4].ϕ₂ && reg.r[4].ϕ₃ == x.r[4].ϕ₃ && reg.r[4].þ == x.r[4].þ
            for x in unique_regions])
            push!(unique_regions, reg)
        end
    end
    return unique_regions
end

function fill_proposition(proposition, objs)
    if isempty(objs) objs=Term[] end
    prop = Compound(Symbol(proposition.name), objs)
    return prop
end

function term_props(domain, problem; init=true)
    termprops=[]
    if init props = collect(initstate(domain, problem).facts) else props = collect(goalstate(domain, problem).facts) end
    for prop in props
        objs = prop.args
        push!(termprops, fill_proposition(prop, objs))
    end
    return termprops

end

function get_init_propositions(domain, problem)
    init_props = Dict()
    init_props[:discrete] =  term_props(domain, problem)

    robot_region = Region(:robot, [Ineq(1,0,0,0), Ineq(-1,0,0,0.5), Ineq(0,1,0,0), Ineq(0,-1,0,0.5)], 0.0)
    b1_region = Region(:b1, [Ineq(1,0,0,10), Ineq(-1,0,0,11), Ineq(0,1,0,10), Ineq(0,-1,0,11)], 0.0)
    init_props[:continuous] = [robot_region, b1_region]
    # init_props[:robot_position] = (0.25,0.25,0.0)
    # init_props[:b1_position] = [10.5, 10.5, 0.0]

    return init_props
end

function get_goal_propositions(domain, problem)
    goal_props = Dict()
    goal_props[:discrete] = term_props(domain, problem;init=false)

    # robot_region = Region(:robot, [Ineq(1,0,0,0), Ineq(-1,0,0,0.5), Ineq(0,1,0,0), Ineq(0,-1,0,0.5)], 0.0)
    b1_region = Region(:b1, [Ineq(1,0,0,19), Ineq(-1,0,0,20), Ineq(0,1,0,19), Ineq(0,-1,0,20)], 0.0)
    robot_region = Region(:robot, [Ineq(1,0,0,40), Ineq(-1,0,0,41), Ineq(0,1,0,40), Ineq(0,-1,0,41)], 0.0)
    goal_props[:continuous] = [robot_region, b1_region]

    return goal_props
end

function get_external_constraints(domain, problem)
    constraints = []
    return constraints
end

function goal_reached(graph, domain, problem)
    goals = get_goal_propositions(domain, problem)
    goal_set = goals[:discrete]
    index = graph.num_levels - 1
    props = graph.props[index]
    μprops = graph.μprops[index]
    goal_found = false
    if issubset(goal_set, props[:discrete])
        goal_found = true
        for goal_pair in collect(permutations(goal_set, 2))
            if goal_pair in μprops
                goal_found = false
                break
            end
        end
    elseif index > 1 && graph.props[index-1] == props
        graph.leveled = true
    end
    if goal_found
        if !in_region(goals[:continuous], props[:continuous])
            goal_found = false
        end
    end
    return goal_found
end

function get_min_max(region)
    x=[0.,0.]; y=[0.,0.]
    x[1] = region.r[1].þ
    x[2] = region.r[2].þ
    y[1] = region.r[3].þ
    y[2] = region.r[4].þ
    return x,y

end

function overlaps(x1r, y1r, x2r, y2r)
    if x1r[2] < x2r[1] || x2r[2] < x1r[1] || y1r[2] < y2r[1] || y2r[2] < y1r[1]
        return false
    end
    return true
end

function intersects(reg1, reg2)
    x1r, y1r = get_min_max(reg1)
    x2r, y2r = get_min_max(reg2)
    return overlaps(x1r, y1r, x2r, y2r)
end

function regions_intersect(regions1, regions2)
    for r2 in regions2
        x2r, y2r = get_min_max(r2)
        for r1 in regions1
            if r2.name == r1.name
                x1r, y1r = get_min_max(r1)
                if !overlaps(x1r, y1r, x2r, y2r) return false end
            end
        end
    end
    return true
end

# function in_region(regions1, regions2)
#     for r2 in regions2
#         x2r, y2r = get_min_max(r2)
#         for r1 in regions1
#             if r2.name == r1.name
#                 x1r, y1r = get_min_max(r1)
#                 if overlaps(x1r, y1r, x2r, y2r) return true end
#             end
#         end
#     end
#     return false
# end

function in_region(regions1, regions2)
    allbool = []
    for r1 in regions1
        x1r, y1r = get_min_max(r1)
        isin = false
        for r2 in regions2
            if r1.name == r2.name
                x2r, y2r = get_min_max(r2)
                if overlaps(x1r, y1r, x2r, y2r) isin = true; break end
            end
        end
        push!(allbool, isin)
    end
    return all(allbool)
end

function get_preconditions(domain, act, vars)
    action = PDDL.get_action(domain, act.name)
    vardict = Dict()
    [vardict[action.args[i]] = vars[i] for i=1:length(action.args)]
    pos, neg = [], []
    for pre in action.precond.args
        if pre.name == :not
            vs = []
            [push!(vs, vardict[arg]) for arg in pre.args[1].args]
            prop = fill_proposition(pre.args[1], vs)
            push!(neg, prop)
        else
            vs = []
            [push!(vs, vardict[arg]) for arg in pre.args]
            prop = fill_proposition(pre, vs)
            push!(pos, prop)
        end
    end
    return pos, neg
end

function get_effects(domain, act, vars)
    action = PDDL.get_action(domain, act.name)
    vardict = Dict()
    [vardict[action.args[i]] = vars[i] for i=1:length(action.args)]
    pos, neg = [], []
    for eff in action.effect.args
        if eff.name == :not
            vs = []
            [push!(vs, vardict[arg]) for arg in eff.args[1].args]
            prop = fill_proposition(eff.args[1], vs)
            push!(neg, prop)
        else
            vs = []
            [push!(vs, vardict[arg]) for arg in eff.args]
            prop = fill_proposition(eff, vs)
            push!(pos, prop)
        end
    end
    return pos, neg
end

function get_all_actions(domain, problem)
    actions = []
    obs = collect(PDDL.get_objects(problem))
    for act in values(PDDL.get_actions(domain))
        vars = act.args
        # if act.name != :move
        for vs in collect(permutations(obs, length(vars)))
            a = Funnel(act.name)
            a.pos_prec, a.neg_prec = get_preconditions(domain, act, vs)
            a.pos_eff, a.neg_eff = get_effects(domain, act, vs)
            a.params = vs
            if act.name == :pick #make more general. preferrably  from prob.pddl
                push!(a.continuous_prec, Region(:robot, [Ineq(1,0,0,20), Ineq(-1,0,0,21), Ineq(0,1,0,20), Ineq(0,-1,0,21)], 0.0))
                a.is_continuous = false

            elseif act.name == :place #make more general later
                push!(a.continuous_prec, Region(:b1, [Ineq(1,0,0,0), Ineq(-1,0,0,0.5), Ineq(0,1,0,0), Ineq(0,-1,0,0.5)], 0.0))
                push!(a.continuous_prec, Region(:robot, [Ineq(1,0,0,0), Ineq(-1,0,0,0.5), Ineq(0,1,0,0), Ineq(0,-1,0,0.5)], 0.0))
                a.is_continuous = false

            elseif act.name == :move_holding #compute funnel at instantitation
                push!(a.continuous_prec, Region(:b1, [Ineq(1,0,0,-Inf), Ineq(-1,0,0,Inf), Ineq(0,1,0,-Inf), Ineq(0,-1,0,Inf)], 0.0))
                a.dynamics = Dynamics(I(2), I(2), [-5.,5], [-5, 5], 1)
                a.is_continuous = true
            elseif act.name == :move
                push!(a.continuous_prec, Region(:b1, [Ineq(1,0,0,-Inf), Ineq(-1,0,0,Inf), Ineq(0,1,0,-Inf), Ineq(0,-1,0,Inf)], 0.0))
                a.dynamics = Dynamics(I(2), I(2), [-5.,5], [-5, 5], 1)
                a.is_continuous = true
            end
            push!(actions, a)
        end
    end
    return actions
end

function apply_constraints(xmin, xmax, ymin, ymax, constraints)
    isempty(constraints) && return xmin, xmax, ymin, ymax
end

#TODO: get init position ranges from problem.pddl
function compute_funnel(action, graph, constraints, level)
    if !action.is_continuous return action end
    regions = graph.props[level][:continuous]
    xyr_init = get_min_max(regions[1])

    a = action
    d = a.dynamics.d
    if action.name == :move
        xmin = xyr_init[1][1] + d*a.dynamics.vx_range[1]
        xmax = xyr_init[1][2] + d*a.dynamics.vx_range[2]
        ymin = xyr_init[2][1] + d*a.dynamics.vy_range[1]
        ymax = xyr_init[2][2] + d*a.dynamics.vy_range[2]
        xmin, xmax, ymin, ymax = apply_constraints(xmin, xmax, ymin, ymax, constraints)
        push!(action.end_region, Region(:robot, [Ineq(1,0,0,xmin), Ineq(-1,0,0,xmax), Ineq(0,1,0,ymin), Ineq(0,-1,0,ymax)], 0.0))

    elseif action.name == :move_holding
        xyo_init = get_min_max(regions[2])
        xmin = xyr_init[1][1] + d*a.dynamics.vx_range[1]
        xmax = xyr_init[1][2] + d*a.dynamics.vx_range[2]
        ymin = xyr_init[2][1] + d*a.dynamics.vy_range[1]
        ymax = xyr_init[2][2] + d*a.dynamics.vy_range[2]

        xomin = xyo_init[1][1] + d*a.dynamics.vx_range[1]
        xomax = xyo_init[1][2] + d*a.dynamics.vx_range[2]
        yomin = xyo_init[2][1] + d*a.dynamics.vy_range[1]
        yomax = xyo_init[2][2] + d*a.dynamics.vy_range[2]
        xomin, xomax, yomin, yomax = apply_constraints(xomin, xomax, yomin, yomax, constraints)

        push!(action.end_region, Region(:robot, [Ineq(1,0,0,xmin), Ineq(-1,0,0,xmax), Ineq(0,1,0,ymin), Ineq(0,-1,0,ymax)], 0.0))
        push!(action.end_region, Region(:b1, [Ineq(1,0,0,xomin), Ineq(-1,0,0,xomax), Ineq(0,1,0,yomin), Ineq(0,-1,0,yomax)], 0.0))
    end
    return action

end

function get_noop_action(proposition)
    noop = Funnel(:noop)
    noop.pos_prec = [proposition]
    noop.pos_eff = [proposition]
    noop.is_continuous = false
    return noop
end

function get_maintain_funnel(region)
    maintain = Funnel(:maintain)
    maintain.continuous_prec=[region]
    maintain.end_region=[region]
    return maintain
end

function is_applicable(action, graph, constraints, level)
    props = graph.props[level][:discrete]
    μprops = graph.μprops[level]
    if issubset(action.pos_prec, props) #&& isdisjoint(action.neg_prec, props)
        app = true
        # println("Level ",level," action ",action.name," is subset to ",props)
        if !isempty(μprops)
            for precondition in collect(permutations(action.pos_prec, 2))
                if precondition in μprops
                    app = false
                    break
                end
            end
        end

    else
        # println("Level ",level," action ",action.name, " ",action.pos_prec," is NOT subset to ",props)
        app = false
    end
    if app
        regions2 = graph.props[level][:continuous]
        regions1 = action.continuous_prec

        if in_region(regions1, regions2)
            app = true
            # println("Level ",level," action ",action.name," intersects prop region")
        else
            # println("Level ",level," action ",action.name," DOESN'T intersect prop region")
            app = false
        end
    end
    return app
end

#for continuous, resolved region might have something to say about it
function get_action_propositions(action)
    propositions = Dict()
    propositions[:discrete]=[]
    propositions[:continuous]=[]
    for prop in action.pos_eff
        push!(propositions[:discrete], prop)
        # println("pushed ",action.name,"'s ",prop)
    end
    for region in action.end_region
        push!(propositions[:continuous], region)
    end
    return propositions
end

function get_proposition_mutexes(graph, level)
    props = vcat(graph.props[level][:discrete], graph.props[level][:continuous])
    μprops=[]
    actions_with_p = Set()
    actions_with_q = Set()
    action_list = graph.acts[level-1]
    for a in action_list
        for (p,q) in collect(permutations(props, 2))
            if p in a.pos_eff && q in a.neg_eff
                push!(μprops, [p,q])
            end
            if p in a.pos_eff push!(actions_with_p, a) end
            if q in a.pos_eff push!(actions_with_q, a) end
        end

    end
    for (p,q) in collect(permutations(graph.props[level][:continuous], 2))
        if !regions_intersect([p],[q]) push!(μprops, [p,q]) end
    # end
        μacts = graph.μacts[level-1]
        for p_action in actions_with_p
            for q_action in actions_with_q
                if p_action != q_action
                    if [p_action, q_action] in μacts
                        push!(μprops, [p,q])
                    end
                end
            end
        end
    end
    return μprops
end

function action_pairs_are_mutex(a, b, μprops)
    #1
    if !isempty(intersect(a.neg_eff, union(b.pos_prec, b.pos_eff))) ||
        !isempty(intersect(b.neg_eff, union(a.pos_prec, a.pos_eff)))
         return true
    end
    if !isempty(μprops)
        for μ in μprops
            p = μ[1]
            q = μ[2]
            if typeof(p) != Region && typeof(q) != Region
                if (p in a.pos_prec && q in b.pos_prec)
                    return true
                end
            end
        end
    end
    #2
    μall = true
    for r1 in a.continuous_prec
        for r2 in b.continuous_prec
            if r1 != r2
                if !([r1, r2] in μprops) μall = false end
            end
        end
    end
    if !μall return false end
    #3
    for r1 in a.continuous_prec
        μall = false
        for d2 in b.pos_prec
            if [r1, d2] in μprops  μall = true end
        end
        if !μall return false end
    end
    #have to add resolved_intersection mutex conditions at some point
    return false
end

function get_action_mutexes(graph, level)
    actions = graph.acts[level]
    μacts = []
    μprops = graph.μprops[level]
    for (a, b) in collect(permutations(actions, 2))
        if action_pairs_are_mutex(a, b, μprops)
            push!(μacts, [a, b])
        end
    end
    return μacts
end


using Gadfly, DataFrames

i = indexin([Float64], [eltypes...])[1]
tconstr = tconstr[:,:,i]
teval = teval[:,:,i]

# Convert to dataframes, and for teval use points/second
dim = ones(Int,size(tconstr,1))*(1:size(tconstr,2))'
namelist = collect(take(cycle(testnames), length(dim)))
ordlist =  collect(take(cycle(testord), length(dim)))
dfconstr = DataFrame(t = vec(tconstr), name = namelist, ord = ordlist, dim = vec(dim))
dfeval   = DataFrame(rate = vec(10^6./teval), name = namelist, ord = ordlist, dim = vec(dim))

flagBspline = (dfconstr[:name] .== "IBSpline") | (dfconstr[:name] .== "Grid")
dfBconstr = dfconstr[flagBspline, :]
dfBeval   = dfeval[flagBspline, :]
dfGconstr = dfconstr[!flagBspline, :]
dfGeval   = dfeval[!flagBspline, :]

yscale = Scale.y_continuous

# Plot the B-spline methods
plotsize = (6inch, 8inch)
set_default_plot_size(plotsize...)
pBc = plot(dfBconstr, xgroup="dim", ygroup="ord", x="name", y="t", Geom.subplot_grid(Geom.point), yscale, Guide.ylabel("Time (s) by order"))
draw(PNG("constructionB.png", plotsize...), pBc)

pBr = plot(dfBeval, xgroup="dim", ygroup="ord", x="name", y="rate", Geom.subplot_grid(Geom.point), yscale, Guide.ylabel("Throughput (pts/s) by order"))
draw(PNG("rateB.png", plotsize...), pBr)

# Plot the gridded methods
plotsize = (10inch, 8inch)
set_default_plot_size(plotsize...)

pGc = plot(dfGconstr, xgroup="dim", ygroup="ord", x="name", y="t", Geom.subplot_grid(Geom.point), yscale, Guide.ylabel("Time (s) by order"))
draw(PNG("constructionG.png", plotsize...), pGc)

pGr = plot(dfGeval, xgroup="dim", ygroup="ord", x="name", y="rate", Geom.subplot_grid(Geom.point), Scale.y_sqrt, Guide.ylabel("Throughput (pts/s) by order"))
draw(PNG("rateG.png", plotsize...), pGr)


"""
A [Julia](https://julialang.org/) threaded audio output interface to [PortAudio](https://github.com/JuliaAudio/PortAudio.jl) that allows dynamic real-time changes to your audio code while it is playing.  This is useful for quickly iterating your code in Julia's REPL or a [Jupyter](https://jupyter.org/) [notebook](https://github.com/JuliaLang/IJulia.jl).

See [documentation and examples](https://github.com/cduck/InteractiveAudio.jl).
"""
module InteractiveAudio

using Base.Threads
using Base: lock, unlock, trylock, islocked, ReentrantLock

using SampledSignals
using PortAudio

export BackgroundPlayer, play, pause, is_playing, lock_player
export AbstractAudioGenerator, NullAudioGenerator, SampleAudioGenerator,
    SineAudioGenerator, SumAudioGenerator, MutateAudioGenerator

include("audio_generator.jl")
include("background_player.jl")
include("sample_audio_generator.jl")
include("sine_audio_generator.jl")
include("sum_audio_generator.jl")
include("mutate_audio_generator.jl")

end


@testset "101.symmetric-tree.jl" begin
    @test is_symmetric_tree(TreeNode{Int}([1, 2, 2, 3, 4, 4, 3])) == true
    @test is_symmetric_tree(TreeNode{Int}([1, 2, 2, nothing, 3, nothing, 3])) == false
end