module ConstraintSolver

using MatrixNetworks
using JSON

CS = ConstraintSolver

mutable struct Variable
    idx         :: Int
    from        :: Int
    to          :: Int
    first_ptr   :: Int
    last_ptr    :: Int
    values      :: Vector{Int}
    indices     :: Vector{Int}
    offset      :: Int
    min         :: Int
    max         :: Int
    changes     :: Vector{Vector{Tuple{Symbol,Int64,Int64,Int64}}}
end

mutable struct CSInfo
    pre_backtrack_calls :: Int
    backtracked         :: Bool
    backtrack_fixes     :: Int
    in_backtrack_calls  :: Int
    backtrack_reverses  :: Int
end

abstract type Constraint
end

abstract type ObjectiveFunction
end

mutable struct MinMaxObjective <: ObjectiveFunction
    fct                 :: Function
    indices             :: Vector{Int}
end

mutable struct BasicConstraint <: Constraint
    idx                 :: Int
    fct                 :: Function
    indices             :: Vector{Int}
    pvals               :: Vector{Int}
    BasicConstraint() = new()
end

mutable struct LinearVariables
    indices             :: Vector{Int}
    coeffs              :: Vector{Int}
end

mutable struct LinearConstraint <: Constraint
    idx                 :: Int
    fct                 :: Function
    indices             :: Vector{Int}
    pvals               :: Vector{Int}
    coeffs              :: Vector{Int}
    operator            :: Symbol
    rhs                 :: Int
    in_all_different    :: Bool
    LinearConstraint() = new()
end

mutable struct BacktrackObj
    idx                 :: Int
    parent_idx          :: Int
    depth               :: Int
    status              :: Symbol
    variable_idx        :: Int
    pval                :: Int
    constraint_idx      :: Vector{Int}
    best_bound          :: Int

    BacktrackObj() = new()
end

mutable struct TreeLogNode
    id              :: Int
    status          :: Symbol
    best_bound      :: Int
    step_nr         :: Int
    var_idx         :: Int
    set_val         :: Int
    var_states      :: Dict{Int64,Vector{Int64}}
    var_changes     :: Dict{Int64,Vector{Tuple{Symbol, Int64, Int64, Int64}}}
    children        :: Vector{TreeLogNode}
    TreeLogNode() = new()
end

mutable struct CoM
    init_search_space   :: Vector{Variable}
    search_space        :: Vector{Variable}
    subscription        :: Vector{Vector{Int}}
    constraints         :: Vector{Constraint}
    bt_infeasible       :: Vector{Int}
    c_backtrack_idx     :: Int
    backtrack_vec       :: Vector{BacktrackObj}
    sense               :: Symbol
    objective           :: ObjectiveFunction
    best_sol            :: Int # Objective of the best solution
    best_bound          :: Int # Overall best bound
    solutions           :: Vector{Int} # saves only the id to the BacktrackObj
    info                :: CSInfo
    input               :: Dict{Symbol,Any}
    logs                :: Vector{TreeLogNode}

    CoM() = new()
end

include("printing.jl")
include("logs.jl")
include("Variable.jl")
include("objective.jl")
include("linearcombination.jl")
include("all_different.jl")
include("eq_sum.jl")
include("equal.jl")
include("not_equal.jl")

function init()
    com = CoM()
    com.constraints         = Vector{Constraint}()
    com.subscription        = Vector{Vector{Int}}()
    com.init_search_space   = Vector{Variable}()
    com.search_space        = Vector{Variable}()
    com.bt_infeasible       = Vector{Int}()
    com.c_backtrack_idx     = 1
    com.sense               = :None # means no objective
    com.backtrack_vec       = Vector{BacktrackObj}()
    com.solutions           = Vector{Int}()
    com.info                = CSInfo(0, false, 0, 0, 0)
    com.input               = Dict{Symbol, Bool}()
    com.logs                = Vector{TreeLogNode}()
    return com
end

function add_var!(com::CS.CoM, from::Int, to::Int; fix=nothing)
    ind = length(com.search_space)+1
    changes = Vector{Vector{Tuple{Symbol,Int64,Int64,Int64}}}()
    push!(changes, Vector{Tuple{Symbol,Int64,Int64,Int64}}())
    var = Variable(ind, from, to, 1, to-from+1, from:to, 1:to-from+1, 1-from, from, to, changes)
    if fix !== nothing
        fix!(com, var, fix)
    end
    push!(com.search_space, var)
    push!(com.subscription, Int[])
    push!(com.bt_infeasible, 0)
    return var
end

function fulfills_constraints(com::CS.CoM, index, value)
    # variable doesn't have any constraint
    if index > length(com.subscription)
        return true
    end
    constraints = com.constraints[com.subscription[index]]
    feasible = true
    for constraint in constraints
        feasible = constraint.fct(com, constraint, value, index)
        if !feasible
            break
        end
    end
    return feasible
end

"""
    fixed_vs_unfixed(search_space, indices)

Returns the fixed_vals as well as the unfixed_indices
"""
function fixed_vs_unfixed(search_space, indices)
    # get all values which are fixed
    fixed_vals = Int[]
    unfixed_indices = Int[]
    for (i,ind) in enumerate(indices)
        if isfixed(search_space[ind])
            push!(fixed_vals, value(search_space[ind]))
        else
            push!(unfixed_indices, i)
        end
    end
    return fixed_vals, unfixed_indices
end

"""
    set_pvals!(com::CS.CoM, constraint::Constraint)

Computes the possible values inside this constraint and sets it as constraint.pvals
"""
function set_pvals!(com::CS.CoM, constraint::Constraint)
    indices = constraint.indices
    variables = Variable[v for v in com.search_space[indices]]
    pvals_intervals = Vector{NamedTuple}()
    push!(pvals_intervals, (from = variables[1].from, to = variables[1].to))
    for (i,ind) in enumerate(indices)
        extra_from = variables[i].from
        extra_to   = variables[i].to
        comp_inside = false
        for cpvals in pvals_intervals
            if extra_from >= cpvals.from && extra_to <= cpvals.to
                # completely inside the interval already
                comp_inside = true
                break
            elseif extra_from >= cpvals.from && extra_from <= cpvals.to
                extra_from = cpvals.to+1
            elseif extra_to <= cpvals.to && extra_to >= cpvals.from
                extra_to = cpvals.from-1
            end
        end
        if !comp_inside && extra_to >= extra_from
            push!(pvals_intervals, (from = extra_from, to = extra_to))
        end
    end
    pvals = collect(pvals_intervals[1].from:pvals_intervals[1].to)
    for interval in pvals_intervals[2:end]
        pvals = vcat(pvals, collect(interval.from:interval.to))
    end
    constraint.pvals = pvals
end

"""
    add_constraint!(com::CS.CoM, constraint::Constraint)

Add a constraint to the model
"""
function add_constraint!(com::CS.CoM, constraint::Constraint)
    constraint.idx = length(com.constraints)+1
    push!(com.constraints, constraint)
    set_pvals!(com, constraint)
    for (i,ind) in enumerate(constraint.indices)
        push!(com.subscription[ind], constraint.idx)
    end
end

function set_objective!(com::CS.CoM, sense::Symbol, objective::ObjectiveFunction)
    if sense != :Min && sense != :Max
        throw(ErrorException("The objective sense must be :Min or :Max"))
    end
    com.sense = sense
    com.objective = objective
end

function get_weak_ind(com::CS.CoM)
    lowest_num_pvals = typemax(Int)
    biggest_inf = -1
    best_ind = -1
    biggest_dependent = typemax(Int)
    found = false

    for ind in 1:length(com.search_space)
        if !isfixed(com.search_space[ind])
            num_pvals = nvalues(com.search_space[ind])
            inf = com.bt_infeasible[ind]
            if inf >= biggest_inf
                if inf > biggest_inf || num_pvals < lowest_num_pvals
                    lowest_num_pvals = num_pvals
                    biggest_inf = inf
                    best_ind = ind
                    found = true
                end
            end
        end
    end
    return found, best_ind
end

function prune!(com::CS.CoM, prune_steps::Vector{Int})
    search_space = com.search_space
    for backtrack_idx in prune_steps
        for var in search_space
            for change in var.changes[backtrack_idx]
                fct_symbol = change[1]
                val = change[2]
                if fct_symbol == :fix
                    fix!(com, var, val; changes=false)
                elseif fct_symbol == :rm
                    rm!(com, var, val; changes=false)
                elseif fct_symbol == :remove_above
                    remove_above!(com, var, val; changes=false)
                elseif fct_symbol == :remove_below
                    remove_below!(com, var, val; changes=false)
                else
                    throw(ErrorException("There is no pruning function for $fct_symbol"))
                end
            end
        end
    end
end

"""
    prune!(com::CS.CoM; pre_backtrack=false)

Prune based on changes by initial solve or backtracking. The information for it is stored in each variable
Returns whether it's still feasible
"""
function prune!(com::CS.CoM; pre_backtrack=false)
    feasible = true
    prev_var_length = zeros(Int, length(com.search_space))
    constraint_idxs_vec = zeros(Int, length(com.constraints))
    # get all constraints which need to be called (only once)
    current_backtrack_id = com.c_backtrack_idx
    for var in com.search_space
        if length(var.changes[current_backtrack_id]) > 0
            inner_constraints = com.constraints[com.subscription[var.idx]]
            for constraint in inner_constraints
                constraint_idxs_vec[constraint.idx] = 1
            end
        end
    end

    current_level = 1
    level_increased = false
    # while we haven't called every constraint
    while true
        b_open_constraint = false
        # will be changed or b_open_constraint => false
        constraint = com.constraints[1]
        for ci in 1:length(com.constraints)
            if constraint_idxs_vec[ci] == current_level
                constraint = com.constraints[ci]
                constraint_idxs_vec[ci] = 0
                b_open_constraint = true
                break
            end
        end
        if !b_open_constraint && !level_increased
            level_increased = true
            current_level += 1
            continue
        end
        if !b_open_constraint && level_increased
            break
        end
        level_increased = false
        if all(v->isfixed(v), com.search_space[constraint.indices])
            continue
        end
        feasible = constraint.fct(com, constraint; logs = false)
        if !pre_backtrack
            com.info.in_backtrack_calls += 1
        else
            com.info.pre_backtrack_calls += 1
        end

        if !feasible
            break
        end

        # if we changed another variable increase the level of the constraints to call them later
        for var in com.search_space
            new_var_length = length(var.changes[current_backtrack_id])
            if new_var_length > prev_var_length[var.idx]
                prev_var_length[var.idx] = new_var_length
                inner_constraints = com.constraints[com.subscription[var.idx]]
                for constraint in inner_constraints
                    if constraint_idxs_vec[constraint.idx] < current_level
                        constraint_idxs_vec[constraint.idx] = current_level + 1
                    end
                end
            end
        end
    end
    return feasible
end

function single_reverse_pruning!(search_space, index::Int, prune_int::Int, prune_fix::Int)
    if prune_int > 0
        var = search_space[index]
        l_ptr = max(1,var.last_ptr)

        new_l_ptr = var.last_ptr + prune_int
        @views min_val = minimum(var.values[l_ptr:new_l_ptr])
        @views max_val = maximum(var.values[l_ptr:new_l_ptr])
        if min_val < var.min
            var.min = min_val
        end
        if max_val > var.max
            var.max = max_val
        end
        var.last_ptr = new_l_ptr
    end
    if prune_fix > 0
        var = search_space[index]
        f_ptr = max(1,var.first_ptr)
        var.last_ptr = prune_fix

        @views min_val = minimum(var.values[1:prune_fix])
        @views max_val = maximum(var.values[1:prune_fix])
        if min_val < var.min
            var.min = min_val
        end
        if max_val > var.max
            var.max = max_val
        end
        var.first_ptr = 1
    end
end

"""
    reverse_pruning!(com, backtrack_idx)

Reverse the changes made by a specific backtrack object
"""
function reverse_pruning!(com::CS.CoM, backtrack_idx::Int)
    search_space = com.search_space
    for var in search_space
        v_idx = var.idx

        for change in Iterators.reverse(var.changes[backtrack_idx])
            single_reverse_pruning!(search_space, v_idx, change[4], change[3])
        end
    end
end

function get_best_bound(com::CS.CoM; var_idx=0, val=0)
    if com.sense == :None
        return 0
    end
    return com.objective.fct(com, var_idx, val)
end

function checkout_from_to!(com::CS.CoM, from_idx::Int, to_idx::Int)
    backtrack_vec = com.backtrack_vec
    from = backtrack_vec[from_idx]
    to = backtrack_vec[to_idx]
    if to.parent_idx == from.idx
        return
    end
    reverse_pruning!(com, from.idx)

    prune_steps = Vector{Int}()
    # first go to same level if new is higher in the tree
    if to.depth < from.depth
        depth = from.depth
        parent_idx = from.parent_idx
        parent = backtrack_vec[parent_idx]
        while to.depth < depth
            reverse_pruning!(com, parent_idx)
            parent = backtrack_vec[parent_idx]
            parent_idx = parent.parent_idx
            depth -= 1
        end
        if parent_idx == to.parent_idx
            return
        else
            from = parent
        end
    elseif from.depth < to.depth
        depth = to.depth
        parent_idx = to.parent_idx
        parent = backtrack_vec[parent_idx]
        while from.depth < depth
            pushfirst!(prune_steps, parent_idx)
            parent = backtrack_vec[parent_idx]
            parent_idx = parent.parent_idx
            depth -= 1
        end

        to = parent
        if backtrack_vec[prune_steps[1]].parent_idx == from.parent_idx
            prune!(com, prune_steps)
            return
        end
    end
    @assert from.depth == to.depth
    # same diff but different parent
    # => level up until same parent
     while from.parent_idx != to.parent_idx
        reverse_pruning!(com, from.parent_idx)
        from = backtrack_vec[from.parent_idx]

        pushfirst!(prune_steps, to.parent_idx)
        to = backtrack_vec[to.parent_idx]
    end

    prune!(com, prune_steps)
end

function backtrack!(com::CS.CoM, max_bt_steps)
    found, ind = get_weak_ind(com)
    com.info.backtrack_fixes   = 1

    pvals = reverse!(values(com.search_space[ind]))
    dummy_backtrack_obj = BacktrackObj()
    dummy_backtrack_obj.status = :Close
    dummy_backtrack_obj.idx = 1
    dummy_backtrack_obj.variable_idx = 1

    backtrack_vec = com.backtrack_vec
    push!(backtrack_vec, dummy_backtrack_obj)

    # the first solve (before backtrack) has idx 1
    num_backtrack_objs = 1
    step_nr = 1
    for pval in pvals
        backtrack_obj = BacktrackObj()
        num_backtrack_objs += 1
        backtrack_obj.idx = num_backtrack_objs
        backtrack_obj.parent_idx = 1
        backtrack_obj.depth = 1
        backtrack_obj.status = :Open
        backtrack_obj.best_bound = get_best_bound(com; var_idx=ind, val=pval)
        backtrack_obj.variable_idx = ind
        backtrack_obj.pval = pval
        push!(backtrack_vec, backtrack_obj)
        for v in com.search_space
            push!(v.changes, Vector{Tuple{Symbol,Int64,Int64,Int64}}())
        end
        if com.input[:logs]
            push!(com.logs, log_one_node(com, length(com.search_space), num_backtrack_objs, -1))
        end
    end
    last_backtrack_id = 0

    started = true
    just_increased_depth = true

    while length(backtrack_vec) > 0
        step_nr += 1
        # get next open backtrack object
        l = 1
        if com.sense == :None
            backtrack_obj = backtrack_vec[l]
            while backtrack_obj.status != :Open
                l += 1
                if l > length(backtrack_vec)
                    break
                end
                backtrack_obj = backtrack_vec[l]
            end
        else
            sorted_backtrack_idxs_nvals = sortperm(backtrack_vec; by=o->o.depth, rev=true)
            sorted_backtrack_idxs_bb = sortperm(backtrack_vec[sorted_backtrack_idxs_nvals]; by=o->o.best_bound, rev=com.sense == :Max )
            sorted_backtrack_idxs = sorted_backtrack_idxs_nvals[sorted_backtrack_idxs_bb]

            backtrack_obj = backtrack_vec[sorted_backtrack_idxs[l]]
            while backtrack_obj.status != :Open
                l += 1
                if l > length(backtrack_vec)
                    break
                end
                backtrack_obj = backtrack_vec[sorted_backtrack_idxs[l]]
            end
        end

        # no open node => Infeasible
        if l > length(backtrack_vec)
            break
        end
        com.best_bound = backtrack_obj.best_bound
        ind = backtrack_obj.variable_idx

        com.c_backtrack_idx = backtrack_obj.idx

        if !started
            com.c_backtrack_idx = 0
            checkout_from_to!(com, last_backtrack_id, backtrack_obj.idx)
            com.c_backtrack_idx = backtrack_obj.idx
        end

        if com.input[:logs]
            com.logs[backtrack_obj.idx].step_nr = step_nr
        end

        started = false
        last_backtrack_id = backtrack_obj.idx

        just_increased_depth = false
        backtrack_obj.status = :Closed

        pval = backtrack_obj.pval

        # check if this value is still possible
        constraints = com.constraints[com.subscription[ind]]
        feasible = true
        for constraint in constraints
            feasible = constraint.fct(com, constraint, pval, ind)
            if !feasible
                break
            end
        end
        if !feasible
            continue
        end
        # value is still possible => set it
        fix!(com, com.search_space[ind], pval)

        for constraint in constraints
            feasible = constraint.fct(com, constraint; logs = false)
            if !feasible
                feasible = false
                break
            end
        end
        if !feasible
            com.info.backtrack_reverses += 1
            continue
        end

        # prune on changed values
        feasible = prune!(com)

        if !feasible
            com.info.backtrack_reverses += 1
            continue
        end

        found, ind = get_weak_ind(com)
        if !found
            com.best_sol = get_best_bound(com)
            if com.best_sol == com.best_bound
                push!(com.solutions, backtrack_obj.idx)
                return :Solved
            else
                continue
            end
        end

        if com.info.backtrack_fixes + 1 > max_bt_steps
            return :NotSolved
        end

        com.info.backtrack_fixes   += 1

        if com.input[:logs]
            com.logs[backtrack_obj.idx] = log_one_node(com, length(com.search_space), backtrack_obj.idx, step_nr)
        end

        pvals = reverse!(values(com.search_space[ind]))
        last_backtrack_obj = backtrack_vec[last_backtrack_id]
        for pval in pvals
            backtrack_obj = BacktrackObj()
            num_backtrack_objs += 1
            backtrack_obj.parent_idx = last_backtrack_obj.idx
            backtrack_obj.depth = last_backtrack_obj.depth + 1
            backtrack_obj.idx = num_backtrack_objs
            backtrack_obj.status = :Open
            backtrack_obj.best_bound = get_best_bound(com; var_idx=ind, val=pval)
            backtrack_obj.variable_idx = ind
            backtrack_obj.pval = pval
            push!(backtrack_vec, backtrack_obj)
            for v in com.search_space
                push!(v.changes, Vector{Tuple{Symbol,Int64,Int64,Int64}}())
            end
            if com.input[:logs]
                push!(com.logs, log_one_node(com, length(com.search_space), num_backtrack_objs, -1))
            end
        end
        just_increased_depth = true
    end
    if length(com.solutions) > 0
        println("Actually found one")
    end
    return :Infeasible
end

function arr2dict(arr)
    d = Dict{Int,Bool}()
    for v in arr
        d[v] = true
    end
    return d
end

function simplify!(com)
    # check if we have all_different and sum constraints
    # (all different where every value is used)
    b_all_different = false
    b_all_different_sum = false
    b_sum = false
    for constraint in com.constraints
        if nameof(constraint.fct) == :all_different
            b_all_different = true
            if length(constraint.indices) == length(constraint.pvals)
                b_all_different_sum = true
            end
        elseif nameof(constraint.fct) == :eq_sum
            b_sum = true
        end
    end
    # l = 0
    if b_all_different_sum && b_sum
        # for each all_different constraint
        # which can be formulated as a sum constraint
        # check which sum constraints are completely inside all different
        # which are partially inside
        # compute inside sum and total sum
        # println("lc :", length(com.constraints))
        n_constraints_before = length(com.constraints)
        for constraint_idx in 1:length(com.constraints)
            constraint = com.constraints[constraint_idx]

            if nameof(constraint.fct) == :all_different
                add_sum_constraint = true
                if length(constraint.indices) == length(constraint.pvals)
                    all_diff_sum = sum(constraint.pvals)
                    in_sum = 0
                    total_sum = 0
                    outside_indices = Int[]
                    cons_indices_dict = arr2dict(constraint.indices)
                    for variable_idx in keys(cons_indices_dict)
                        found_sum_constraint = false
                        for sub_constraint_idx in com.subscription[variable_idx]
                            # don't mess with constraints added later on
                            if sub_constraint_idx > n_constraints_before
                                continue
                            end
                            sub_constraint = com.constraints[sub_constraint_idx]
                            # it must be an equal constraint and all coefficients must be 1 otherwise we can't add a constraint
                            if nameof(sub_constraint.fct) == :eq_sum && all(c->c==1, sub_constraint.coeffs)
                                found_sum_constraint = true
                                total_sum += sub_constraint.rhs
                                all_inside = true
                                for sub_variable_idx in sub_constraint.indices
                                    if !haskey(cons_indices_dict, sub_variable_idx)
                                        all_inside = false
                                        push!(outside_indices, sub_variable_idx)
                                    else
                                        delete!(cons_indices_dict, sub_variable_idx)
                                    end
                                end
                                if all_inside
                                    in_sum += sub_constraint.rhs
                                end
                                break
                            end
                        end
                        if !found_sum_constraint
                            add_sum_constraint = false
                            break
                        end
                    end

                    # make sure that there are not too many outside indices
                    if add_sum_constraint && length(outside_indices) < 3
                        add_constraint!(com, sum(com.search_space[outside_indices]) == total_sum - all_diff_sum)
                    end
                end
            end
        end
    end
end

function set_in_all_different!(com::CS.CoM)
    for constraint in com.constraints
        if :in_all_different in fieldnames(typeof(constraint))
            if !constraint.in_all_different
                subscriptions_idxs = [[i for i in com.subscription[v]] for v in constraint.indices]
                intersects = intersect(subscriptions_idxs...)

                for i in intersects
                    if nameof(com.constraints[i].fct) == :all_different
                        constraint.in_all_different = true
                        break
                    end
                end
            end
        end
    end
end

function solve!(com::CS.CoM; backtrack=true, max_bt_steps=typemax(Int64), keep_logs=false)
    com.input[:logs] = keep_logs
    if keep_logs
        com.init_search_space = deepcopy(com.search_space)
    end

    set_in_all_different!(com)

    # check for better constraints
    simplify!(com)

    # check if all feasible even if for example everything is fixed
    feasible = true
    for constraint in com.constraints
        com.info.pre_backtrack_calls += 1
        feasible = constraint.fct(com, constraint)
        if !feasible
            feasible = false
            break
        end
    end

    if !feasible
        return :Infeasible
    end

    if all(v->isfixed(v), com.search_space)
        com.best_bound = get_best_bound(com)
        com.best_sol = com.best_bound
        return :Solved
    end
    feasible = prune!(com; pre_backtrack=true)

    com.best_bound = get_best_bound(com)
    if keep_logs
        push!(com.logs, log_one_node(com, length(com.search_space), 1, 1))
    end

    if !feasible
        return :Infeasible
    end

    if all(v->isfixed(v), com.search_space)
        com.best_sol = com.best_bound
        return :Solved
    end
    if backtrack
        com.info.backtracked = true
        return backtrack!(com, max_bt_steps)
    else
        @info "Backtracking is turned off."
        return :NotSolved
    end
end

export add_var!, add_constraint!, solve!, value, set_objective!

end # module


__precompile__()
module POMCP

if VERSION >= v"0.6.0"
    warn("The POMCP package is deprecated. Please use BasicPOMCP (https://github.com/JuliaPOMDP/BasicPOMCP.jl) instead.")
end

import POMDPs

import POMDPs: action, solve
import Base: rand, mean
import POMDPs: update, updater, initialize_belief
import POMDPToolbox
import StatsBase: WeightVec, sample
import MCTS: RandomActionGenerator, RolloutEstimator, next_action, estimate_value, init_N, convert_to_policy
using ParticleFilters

using Compat

export
    AbstractPOMCPSolver,
    POMCPSolver,
    POMCPDPWSolver,
    POMCPPlanner,
    RootUpdater,
    BeliefNode,
    RootNode,
    ObsNode,
    AbstractActNode,
    ActNode,
    POMCPTreeVisualizer,

    solve,
    action,
    create_policy,
    update,
    updater,
    create_belief,
    init_V,
    init_N,
    sparse_actions,
    estimate_value,
    extract_belief,
    convert_estimator,
    add_N,

    RandomActionGenerator,
    next_action,

    RolloutEstimator,
    PORollout,
    FORollout,
    FOValue,

    ParticleCollection,
    ParticleReinvigorator,
    reinvigorate!,
    handle_unseen_observation,
    DefaultReinvigoratorStub,

    NoDecision,
    AllSamplesTerminal,
    ExceptionRethrow,
    default_action,

    # deprecated
    PORolloutEstimator

include("tree.jl")
include("particle_filter.jl")
include("exceptions.jl")

abstract type AbstractPOMCPSolver <: POMDPs.Solver end

"""
POMCP Solver type

Fields:

    eps::Float64
        Rollouts and tree expansion will stop when discount^depth is less than this.
        default: 0.01

    max_depth::Int
        Rollouts and tree expension will stop when this depth is reached.
        default: 10

    c::Float64
        UCB exploration constant - specifies how much the solver should explore.
        default: 1.0

    tree_queries::Int
        Number of iterations during each action() call.
        default: 100

    rng::AbstractRNG
        Random number generator.
        default: Base.GLOBAL_RNG

    node_belief_updater::Updater
        Calculates the belief for a new belief node (see notebooks/Belief_and_Particle_Filter_Options.ipynb for more info.)
        default: DefaultReinvigoratorStub() - this will simply keep the particles as described in the paper without doing any reinvigoration.

    estimate_value::Any (rollout policy can be specified by setting this to RolloutEstimator(policy))
        Function, object, or number used to estimate the value at the leaf nodes.
        If this is a function `f`, `f(pomdp, s, h::BeliefNode, steps)` will be called to estimate the value.
        If this is an object `o`, `estimate_value(o, pomdp, s, h::BeliefNode, steps)` will be called.
        If this is a number, the value will be set to that number
        default: RolloutEstimator(RandomSolver(rng))

    init_V::Any
        Function, object, or number used to set the initial V(h,a) value at a new node.
        If this is a function `f`, `f(pomdp, h, a)` will be called to set the value.
        If this is an object `o`, `init_V(o, pomdp, h, a)` will be called.
        If this is a number, V will be set to that number
        default: 0.0

    init_N::Any
        Function, object, or number used to set the initial N(s,a) value at a new node.
        If this is a function `f`, `f(pomdp, h, a)` will be called to set the value.
        If this is an object `o`, `init_N(o, pomdp, h, a)` will be called.
        If this is a number, N will be set to that number
        default: 0

    num_sparse_actions::Int
        Number of actions to be considered at each node.
        If <= 0, the entire action space will be considered.
        default: 0

    default_action::Any
        Function, action, or Policy used to determine the action if POMCP fails with exception `ex`.
        If this is a Function `f`, `f(belief, ex)` will be called.
        If this is a Policy `p`, `action(p, belief)` will be called.
        If it is an object `a`, `default_action(a, belief, ex) will be called, and
        if this method is not implemented, `a` will be returned directly.
"""
type POMCPSolver <: AbstractPOMCPSolver
    eps::Float64 # will stop simulations when discount^depth is less than this
    max_depth::Int
    c::Float64 # UCB exploration constant
    tree_queries::Int
    rng::AbstractRNG
    node_belief_updater::Union{POMDPs.Updater, DefaultReinvigoratorStub}

    estimate_value::Any

    init_V::Any
    init_N::Any

    num_sparse_actions::Int # = 0 or less if not used
    default_action::Any
end

"""
POMCP Solver type

Fields:

    eps::Float64
        Rollouts and tree expansion will stop when discount^depth is less than this.
        default: 0.01

    max_depth::Int
        Rollouts and tree expension will stop when this depth is reached.
        default: 10

    c::Float64
        UCB exploration constant - specifies how much the solver should explore.
        default: 1.0

    tree_queries::Int
        Number of iterations during each action() call.
        default: 100

    rng::AbstractRNG
        Random number generator.
        default: Base.GLOBAL_RNG

    node_belief_updater::Updater
        Calculates the belief for a new belief node (see notebooks/Belief_and_Particle_Filter_Options.ipynb for more info.)
        default: DefaultReinvigoratorStub() - this will simply keep the particles as described in the paper without doing any reinvigoration.

    estimate_value::Any (rollout policy can be specified by setting this to RolloutEstimator(policy))
        Function, object, or number used to estimate the value at the leaf nodes.
        If this is a function `f`, `f(pomdp, s, h::BeliefNode, steps)` will be called to estimate the value.
        If this is an object `o`, `estimate_value(o, pomdp, s, h::BeliefNode, steps)` will be called.
        If this is a number, the value will be set to that number
        default: RolloutEstimator(RandomSolver(rng))

    enable_action_pw::Bool
        Turn progressive widening of the number of actions considered on or off.
        If false, all actions will be considered.
        default: true

    k_action::Float64
    alpha_action::Float64
    k_observation::Float64
    alpha_observation::Float64
        These constants control the double progressive widening. A new observation
        or action will be added if the number of children is less than or equal to kN^alpha.
        defaults: k:10, alpha:0.5

    init_V::Any
        Function, object, or number used to set the initial V(h,a) value at a new node.
        If this is a function `f`, `f(pomdp, h, a)` will be called to set the value.
        If this is an object `o`, `init_V(o, pomdp, h, a)` will be called.
        If this is a number, V will be set to that number
        default: 0.0

    init_N::Any
        Function, object, or number used to set the initial N(s,a) value at a new node.
        If this is a function `f`, `f(pomdp, h, a)` will be called to set the value.
        If this is an object `o`, `init_N(o, pomdp, h, a)` will be called.
        If this is a number, N will be set to that number
        default: 0

    next_action::Any
        Function or object used to choose the next action to be considered for progressive widening.
        The next action is determined based on the POMDP, the belief, `b`, and the current `BeliefNode`, `h`.
        If this is a function `f`, `f(pomdp, b, h)` will be called to set the value.
        If this is an object `o`, `next_action(o, pomdp, b, h)` will be called.
        default: RandomActionGenerator(rng)

    default_action::Any
        Function, action, or Policy used to determine the action if POMCP fails with exception `ex`.
        If this is a Function `f`, `f(belief, ex)` will be called.
        If this is a Policy `p`, `action(p, belief)` will be called.
        If it is an object `a`, `default_action(a, belief, ex) will be called, and
        if this method is not implemented, `a` will be returned directly.

For more information on the k and alpha parameters, see Couëtoux, A., Hoock, J.-B., Sokolovska, N., Teytaud, O., & Bonnard, N. (2011). Continuous Upper Confidence Trees. In Learning and Intelligent Optimization. Rome, Italy. Retrieved from http://link.springer.com/chapter/10.1007/978-3-642-25566-3_32
"""
type POMCPDPWSolver <: AbstractPOMCPSolver
    eps::Float64 # will stop simulations when discount^depth is less than this
    max_depth::Int
    c::Float64
    tree_queries::Int
    rng::AbstractRNG
    node_belief_updater::Union{POMDPs.Updater, DefaultReinvigoratorStub}

    estimate_value::Any

    enable_action_pw::Bool

    alpha_observation::Float64
    k_observation::Float64
    alpha_action::Float64
    k_action::Float64
    init_V::Any
    init_N::Any
    next_action::Any
    default_action::Any
end

"""
Policy that builds a POMCP tree to determine an optimal next action.

Note, you should construct this using the create_policy function
"""
type POMCPPlanner{S, A, O, B, SolverType<:AbstractPOMCPSolver} <: POMDPs.Policy
    problem::POMDPs.POMDP{S,A,O}
    solver::SolverType
    node_belief_updater::POMDPs.Updater{B}
    solved_estimate::Any

    #XXX hack
    _tree_ref::Nullable{Any}
end

include("constructor.jl")
include("updater.jl")
include("actions.jl")
include("rollout.jl")
include("solver.jl")
include("visualization.jl")

end # module


module Minijyro

using Distributions

export sample!,
    MinijyroModel,
    AbstractHandler,
    TraceHandler,
    LogJointHandler,
    ConditionHandler,
    ReplayHandler,
    EscapeHandler,
    EscapeException,
    queue,
    @jyro,
    handle!,
    handle,
    discrete_enumeration,
    hmc,
    nuts,
    enter!,
    exit!

include("minijyro_types.jl")
include("dsl.jl")
include("handlers.jl")
include("inference.jl")

"""
    sample!(args...)

Allows sampling from probability distributions with side effects specified by
effect handlers in `handlers_stack`.

# Arguments
- `trace::Dict`: dict that can be used to store things in a global state
- `handlers_stack::Array{AbstractHandler,1}`: stack of effect handlers
- `name::Any`: unique name for the sample site
- `dist::Distributions.Distribution`: distribution to sample from
"""
function sample!(
    trace::Dict,
    handlers_stack::Array{AbstractHandler,1},
    name::Any,
    dist::Distributions.Distribution
)
    if length(handlers_stack) == 0
        return rand(dist)
    end

    # NOTE: In this simplified case we could replace :args with :dist because
    #       our "effect" is always sampling from a distribution.
    # NOTE: This could be a custom type but is a Dict for simplicity.
    initial_msg = Dict(
        :fn => rand,
        :args => (dist, ),
        :name => name,
        :value => nothing,
        :is_observed => false,
        :done => false, # Use if we still want other handlers to be active.
        :stop => false, # Use if we want handlers below on the stack not to be active.
        :continuation => nothing # Function to be called after effect handlers have been applied.
    )
    msg = apply_stack!(trace, handlers_stack, initial_msg)
    return msg[:value]
end

"""
    apply_stack!(trace::Dict, handlers_stack::Array{AbstractHandler,1}, msg::Dict)

Apply all the effect handlers from `handlers_stack` to the effect described in
`msg`.

# Arguments
- `trace::Dict`: dict which can be used to store things in a global state
- `handlers_stack::Array{AbstractHandler,1}`: stack of effect handlers
- `msg::Dict`: dict describing the effectful computation

# Returns
- `Dict`: message after effect handlers are applied
"""
function apply_stack!(
    trace::Dict,
    handlers_stack::Array{AbstractHandler,1},
    msg::Dict
)
    @assert length(handlers_stack) > 0

    handler_counter = 0
    # Loop through handlers from top of the stack to the bottom.
    for handler in handlers_stack[end:-1:1]
        handler_counter += 1
        process_message!(trace, handler, msg)
        if msg[:stop]
            break
        end
    end

    if !(msg[:value] != nothing || msg[:done])
        msg[:value] = msg[:fn](msg[:args]...)
    end

    # Loop through handlers from bottom of the stack to the top.
    # If we exited the first loop early then we will start looping from the
    # handler which caused the loop to exit.
    for handler in handlers_stack[end-handler_counter+1:end]
        postprocess_message!(trace, handler, msg)
    end

    if msg[:continuation] != nothing
        msg[:continuation](trace, msg)
    end

    return msg
end
end # module


# This file is a part of Julia. License is MIT: http://julialang.org/license

#Period types
value(x::Period) = x.value

# The default constructors for Periods work well in almost all cases
# P(x) = new((convert(Int64,x))
# The following definitions are for Period-specific safety
for period in (:Year, :Month, :Week, :Day, :Hour, :Minute, :Second, :Millisecond)
    period_str = string(period)
    accessor_str = lowercase(period_str)
    # Convenience method for show()
    @eval _units(x::$period) = " " * $accessor_str * (abs(value(x)) == 1 ? "" : "s")
    # periodisless
    @eval periodisless(x::$period,y::$period) = value(x) < value(y)
    # AbstractString parsing (mainly for IO code)
    @eval $period(x::AbstractString) = $period(Base.parse(Int64,x))
    # Period accessors
    typ_str = period in (:Hour, :Minute, :Second, :Millisecond) ? "DateTime" : "TimeType"
    description = typ_str == "TimeType" ? "`Date` or `DateTime`" : "`$typ_str`"
    reference = period == :Week ? " For details see [`$accessor_str(::$typ_str)`](:func:`$accessor_str`)." : ""
    @eval begin
        @doc """
            $($period_str)(dt::$($typ_str)) -> $($period_str)

        The $($accessor_str) part of a $($description) as a `$($period_str)`.$($reference)
        """ ->
        $period(dt::$(Symbol(typ_str))) = $period($(Symbol(accessor_str))(dt))

        @doc """
            $($period_str)(v)

        Construct a `$($period_str)` object with the given `v` value. Input must be
        losslessly convertible to an `Int64`.
        """ $period(v)
    end
end
# Now we're safe to define Period-Number conversions
# Anything an Int64 can convert to, a Period can convert to
Base.convert{T<:Number}(::Type{T},x::Period) = convert(T,value(x))
Base.convert{T<:Period}(::Type{T},x::Real) = T(x)

#Print/show/traits
Base.string{P<:Period}(x::P) = string(value(x),_units(x))
Base.show(io::IO,x::Period) = print(io,string(x))
Base.zero{P<:Period}(::Union{Type{P},P}) = P(0)
Base.one{P<:Period}(::Union{Type{P},P}) = P(1)
Base.typemin{P<:Period}(::Type{P}) = P(typemin(Int64))
Base.typemax{P<:Period}(::Type{P}) = P(typemax(Int64))

# Default values (as used by TimeTypes)
"""
    default(p::Period) -> Period

Returns a sensible "default" value for the input Period by returning `one(p)` for Year,
Month, and Day, and `zero(p)` for Hour, Minute, Second, and Millisecond.
"""
function default end

default{T<:DatePeriod}(p::Union{T,Type{T}}) = one(p)
default{T<:TimePeriod}(p::Union{T,Type{T}}) = zero(p)

(-){P<:Period}(x::P) = P(-value(x))
Base.isless{P<:Period}(x::P,y::P) = isless(value(x),value(y))
=={P<:Period}(x::P,y::P) = value(x) == value(y)

# Period Arithmetic, grouped by dimensionality:
import Base: div, fld, mod, rem, gcd, lcm, +, -, *, /, %, .+, .-, .*, .%
for op in (:+,:-,:lcm,:gcd)
    @eval ($op){P<:Period}(x::P,y::P) = P(($op)(value(x),value(y)))
end

for op in (:/,:div,:fld)
    @eval begin
        ($op){P<:Period}(x::P,y::P) = ($op)(value(x),value(y))
        ($op){P<:Period}(x::P,y::Real) = P(($op)(value(x),Int64(y)))
    end
end

for op in (:rem,:mod)
    @eval begin
        ($op){P<:Period}(x::P,y::P) = P(($op)(value(x),value(y)))
        ($op){P<:Period}(x::P,y::Real) = P(($op)(value(x),Int64(y)))
    end
end

/{P<:Period}(X::StridedArray{P}, y::P) = X ./ y
%{P<:Period}(X::StridedArray{P}, y::P) = X .% y
*{P<:Period}(x::P,y::Real) = P(value(x) * Int64(y))
*(y::Real,x::Period) = x * y
.*{P<:Period}(y::Real, X::StridedArray{P}) = X .* y
for (op,Ty,Tz) in ((:.*,Real,:P),
                   (:./,:P,Float64), (:./,Real,:P),
                   (:div,:P,Int64), (:div,Integer,:P),
                   (:.%,:P,:P),
                   (:mod,:P,:P))
    sop = string(op)
    op_ = sop[1] == '.' ? Symbol(sop[2:end]) : op
    @eval begin
        function ($op){P<:Period}(X::StridedArray{P},y::$Ty)
            Z = similar(X, $Tz)
            for (Idst, Isrc) in zip(eachindex(Z), eachindex(X))
                @inbounds Z[Idst] = ($op_)(X[Isrc],y)
            end
            return Z
        end
    end
end

# intfuncs
Base.gcdx{T<:Period}(a::T,b::T) = ((g,x,y)=gcdx(value(a),value(b)); return T(g),x,y)
Base.abs{T<:Period}(a::T) = T(abs(value(a)))

periodisless(::Period,::Year)        = true
periodisless(::Period,::Month)       = true
periodisless(::Year,::Month)         = false
periodisless(::Period,::Week)        = true
periodisless(::Year,::Week)          = false
periodisless(::Month,::Week)         = false
periodisless(::Period,::Day)         = true
periodisless(::Year,::Day)           = false
periodisless(::Month,::Day)          = false
periodisless(::Week,::Day)           = false
periodisless(::Period,::Hour)        = false
periodisless(::Minute,::Hour)        = true
periodisless(::Second,::Hour)        = true
periodisless(::Millisecond,::Hour)   = true
periodisless(::Period,::Minute)      = false
periodisless(::Second,::Minute)      = true
periodisless(::Millisecond,::Minute) = true
periodisless(::Period,::Second)      = false
periodisless(::Millisecond,::Second) = true
periodisless(::Period,::Millisecond) = false

# return (next coarser period, conversion factor):
coarserperiod{P<:Period}(::Type{P}) = (P,1)
coarserperiod(::Type{Millisecond}) = (Second,1000)
coarserperiod(::Type{Second}) = (Minute,60)
coarserperiod(::Type{Minute}) = (Hour,60)
coarserperiod(::Type{Hour}) = (Day,24)
coarserperiod(::Type{Day}) = (Week,7)
coarserperiod(::Type{Month}) = (Year,12)

# Stores multiple periods in greatest to least order by type, not values,
# canonicalized to eliminate zero periods, merge equal period types,
# and convert more-precise periods to less-precise periods when possible
"""
    CompoundPeriod

A `CompoundPeriod` is useful for expressing time periods that are not a fixed multiple of
smaller periods. For example, \"a year and a  day\" is not a fixed number of days, but can
be expressed using a `CompoundPeriod`. In fact, a `CompoundPeriod` is automatically
generated by addition of different period types, e.g. `Year(1) + Day(1)` produces a
`CompoundPeriod` result.
"""
type CompoundPeriod <: AbstractTime
    periods::Array{Period,1}
    function CompoundPeriod(p::Vector{Period})
        n = length(p)
        if n > 1
            sort!(p, rev=true, lt=periodisless)
            # canonicalize p by merging equal period types and removing zeros
            i = j = 1
            while j <= n
                k = j+1
                while k <= n
                    if typeof(p[j]) == typeof(p[k])
                        p[j] += p[k]
                        k += 1
                    else
                        break
                    end
                end
                if p[j] != zero(p[j])
                    p[i] = p[j]
                    i += 1
                end
                j = k
            end
            n = i - 1 # new length
        elseif n == 1 && value(p[1]) == 0
            p = Period[]
            n = 0
        end
        # canonicalize Periods by pushing "overflow" into a coarser period.
        if n > 0
            pc = sizehint!(Period[], n)
            P = typeof(p[n])
            v = value(p[n])
            i = n - 1
            while true
                Pc, f = coarserperiod(P)
                if i > 0 && typeof(p[i]) == P
                    v += value(p[i])
                    i -= 1
                end
                v0 = f == 1 ? v : rem(v, f)
                v0 != 0 && push!(pc, P(v0))
                if v != v0
                    P = Pc
                    v = div(v - v0, f)
                elseif i > 0
                    P = typeof(p[i])
                    v = value(p[i])
                    i -= 1
                else
                    break
                end
            end
            p = reverse!(pc)
            n = length(p)
        else
            return new(resize!(p, n))
        end
        # reduce the amount of mixed positive/negative Periods.
        if n > 0
            pc = sizehint!(Period[], n)
            i = n
            while i > 0
                j = i

                # Determine sign of the largest period in this group which
                # can be converted into via coarserperiod.
                last = Union{}
                current = typeof(p[i])
                while i > 0 && current != last
                    if typeof(p[i]) == current
                        i -= 1
                    end
                    last, current = current, coarserperiod(current)[1]
                end
                s = sign(value(p[i + 1]))

                # Adjust all the periods in the group based upon the
                # largest period sign.
                P = typeof(p[j])
                v = 0
                while j > i
                    Pc, f = coarserperiod(P)
                    if j > 0 && typeof(p[j]) == P
                        v += value(p[j])
                        j -= 1
                    end
                    v0 = f == 1 ? v : mod(v, f * s)
                    v0 != 0 && push!(pc, P(v0))
                    if v != v0
                        P = Pc
                        v = div(v - v0, f)
                    elseif j > 0
                        P = typeof(p[j])
                        v = 0
                    else
                        break
                    end
                end
            end
            p = reverse!(pc)
        end
        return new(p)
    end
end

"""
    CompoundPeriod(periods) -> CompoundPeriod

Construct a `CompoundPeriod` from a `Vector` of `Period`s. The constructor will
automatically simplify the periods into a canonical form according to the following rules:

* All `Period`s of the same type will be added together
* Any `Period` large enough be partially representable by a coarser `Period` will be broken
  into multiple `Period`s (eg. `Hour(30)` becomes `Day(1) + Hour(6)`)
* `Period`s with opposite signs will be combined when possible
  (eg. `Hour(1) - Day(1)` becomes `-Hour(23)`)

Due to the canonicalization, `CompoundPeriod` is also useful for converting time periods
into more human-comprehensible forms.

# Examples
```julia
julia> Dates.CompoundPeriod([Dates.Hour(12), Dates.Hour(13)])
1 day, 1 hour

julia> Dates.CompoundPeriod([Dates.Hour(-1), Dates.Minute(1)])
-59 minutes

julia> Dates.CompoundPeriod([Dates.Month(1), Dates.Week(-2)])
1 month, -2 weeks

julia> Dates.CompoundPeriod(Dates.Minute(50000)))
4 weeks, 6 days, 17 hours, 20 minutes
```
"""
CompoundPeriod{P<:Period}(p::Vector{P}) = CompoundPeriod(Array{Period}(p))

Base.convert(::Type{CompoundPeriod}, x::Period) = CompoundPeriod(Period[x])
function Base.string(x::CompoundPeriod)
    if isempty(x.periods)
        return "empty period"
    else
        s = ""
        for p in x.periods
            s *= ", " * string(p)
        end
        return s[3:end]
    end
end
Base.show(io::IO,x::CompoundPeriod) = print(io,string(x))

# E.g. Year(1) + Day(1)
(+)(x::Period,y::Period) = CompoundPeriod(Period[x,y])
(+)(x::CompoundPeriod,y::Period) = CompoundPeriod(vcat(x.periods,y))
(+)(y::Period,x::CompoundPeriod) = x + y
(+)(x::CompoundPeriod,y::CompoundPeriod) = CompoundPeriod(vcat(x.periods,y.periods))
# E.g. Year(1) - Month(1)
(-)(x::Period,y::Period) = CompoundPeriod(Period[x,-y])
(-)(x::CompoundPeriod,y::Period) = CompoundPeriod(vcat(x.periods,-y))
(-)(x::CompoundPeriod) = CompoundPeriod(-x.periods)
(-)(y::Union{Period,CompoundPeriod},x::CompoundPeriod) = (-x) + y

GeneralPeriod = Union{Period,CompoundPeriod}
(+)(x::GeneralPeriod) = x
(+){P<:GeneralPeriod}(x::StridedArray{P}) = x

for op in (:.+, :.-)
    op_ = Symbol(string(op)[2:end])
    @eval begin
        function ($op){P<:GeneralPeriod}(X::StridedArray{P},y::GeneralPeriod)
            Z = similar(X, CompoundPeriod)
            for (Idst, Isrc) in zip(eachindex(Z), eachindex(X))
                @inbounds Z[Idst] = ($op_)(X[Isrc],y)
            end
            return Z
        end
        ($op){P<:GeneralPeriod}(x::GeneralPeriod,Y::StridedArray{P}) = ($op)(Y,x) |> ($op_)
        ($op_){P<:GeneralPeriod}(x::GeneralPeriod,Y::StridedArray{P}) = ($op)(Y,x) |> ($op_)
        ($op_){P<:GeneralPeriod}(Y::StridedArray{P},x::GeneralPeriod) = ($op)(Y,x)
        ($op_){P<:GeneralPeriod, Q<:GeneralPeriod}(X::StridedArray{P}, Y::StridedArray{Q}) =
            reshape(CompoundPeriod[($op_)(x,y) for (x,y) in zip(X, Y)], promote_shape(size(X),size(Y)))
    end
end

(==)(x::CompoundPeriod, y::Period) = x == CompoundPeriod(y)
(==)(y::Period, x::CompoundPeriod) = x == y
(==)(x::CompoundPeriod, y::CompoundPeriod) = x.periods == y.periods

# Capture TimeType+-Period methods
(+)(a::TimeType,b::Period,c::Period) = (+)(a,b+c)
(-)(a::TimeType,b::Period,c::Period) = (-)(a,b-c)
(+)(a::TimeType,b::Period,c::Period,d::Period...) = (+)((+)(a,b+c),d...)
(-)(a::TimeType,b::Period,c::Period,d::Period...) = (-)((-)(a,b-c),d...)

function (+)(x::TimeType,y::CompoundPeriod)
    for p in y.periods
        x += p
    end
    return x
end
(+)(x::CompoundPeriod,y::TimeType) = y + x

function (-)(x::TimeType,y::CompoundPeriod)
    for p in y.periods
        x -= p
    end
    return x
end
(-)(x::CompoundPeriod,y::TimeType) = y - x

# Fixed-value Periods (periods corresponding to a well-defined time interval,
# as opposed to variable calendar intervals like Year).
typealias FixedPeriod Union{Week,Day,Hour,Minute,Second,Millisecond}

# like div but throw an error if remainder is nonzero
function divexact(x,y)
    q,r = divrem(x, y)
    r == 0 || throw(InexactError())
    return q
end

# FixedPeriod conversions and promotion rules
const fixedperiod_conversions = [(Week,7),(Day,24),(Hour,60),(Minute,60),(Second,1000),(Millisecond,1)]
for i = 1:length(fixedperiod_conversions)
    (T,n) = fixedperiod_conversions[i]
    N = 1
    for j = i-1:-1:1 # less-precise periods
        (Tc,nc) = fixedperiod_conversions[j]
        N *= nc
        vmax = typemax(Int64) ÷ N
        vmin = typemin(Int64) ÷ N
        @eval function Base.convert(::Type{$T}, x::$Tc)
            $vmin ≤ value(x) ≤ $vmax || throw(InexactError())
            return $T(value(x)*$N)
        end
    end
    N = n
    for j = i+1:length(fixedperiod_conversions) # more-precise periods
        (Tc,nc) = fixedperiod_conversions[j]
        @eval Base.convert(::Type{$T}, x::$Tc) = $T(divexact(value(x), $N))
        @eval Base.promote_rule(::Type{$T},::Type{$Tc}) = $Tc
        N *= nc
    end
end
# have to declare thusly so that diagonal dispatch above takes precedence:
(==){T<:FixedPeriod,S<:FixedPeriod}(x::T,y::S) = (==)(promote(x,y)...)
Base.isless{T<:FixedPeriod,S<:FixedPeriod}(x::T,y::S) = isless(promote(x,y)...)

# other periods with fixed conversions but which aren't fixed time periods
typealias OtherPeriod Union{Month,Year}
let vmax = typemax(Int64) ÷ 12, vmin = typemin(Int64) ÷ 12
    @eval function Base.convert(::Type{Month}, x::Year)
        $vmin ≤ value(x) ≤ $vmax || throw(InexactError())
        Month(value(x)*12)
    end
end
Base.convert(::Type{Year}, x::Month) = Year(divexact(value(x),12))
Base.promote_rule(::Type{Year}, ::Type{Month}) = Month
(==){T<:OtherPeriod,S<:OtherPeriod}(x::T,y::S) = (==)(promote(x,y)...)
Base.isless{T<:OtherPeriod,S<:OtherPeriod}(x::T,y::S) = isless(promote(x,y)...)

# truncating conversions to milliseconds and days:
toms(c::Millisecond) = value(c)
toms(c::Second)      = 1000*value(c)
toms(c::Minute)      = 60000*value(c)
toms(c::Hour)        = 3600000*value(c)
toms(c::Day)         = 86400000*value(c)
toms(c::Week)        = 604800000*value(c)
toms(c::Month)       = 86400000.0*30.436875*value(c)
toms(c::Year)        = 86400000.0*365.2425*value(c)
toms(c::CompoundPeriod) = isempty(c.periods)?0.0 : Float64(sum(toms,c.periods))
days(c::Millisecond) = div(value(c),86400000)
days(c::Second)      = div(value(c),86400)
days(c::Minute)      = div(value(c),1440)
days(c::Hour)        = div(value(c),24)
days(c::Day)         = value(c)
days(c::Week)        = 7*value(c)
days(c::Year)        = 365.2425*value(c)
days(c::Month)       = 30.436875*value(c)
days(c::CompoundPeriod) = isempty(c.periods)?0.0 : Float64(sum(days,c.periods))


module Data

using Configurations: Configurations

export TMaxtrix, MSTMParameters

Configurations.@option "sphere" struct SphereConfig
    size::Float64
    pos::Array{Float64,1}
    real_ref_index::Union{Float64,Nothing} = nothing
    imag_ref_index::Union{Float64,Nothing} = nothing
    real_chiral_factor::Union{Float64,Nothing} = nothing
    imag_chiral_factor::Union{Float64,Nothing} = nothing
    t_matrix::Union{String,Nothing} = nothing
end

Configurations.@option "run" struct MSTMConfig
    sphere_position_file::Union{String,Nothing} = nothing
    output_file::String = "mstm_default_out.dat"
    run_print_file::Union{String,Nothing} = nothing
    write_sphere_data::Bool = true
    length_scale_factor::Float64 = 1.0
    real_ref_index_scale_factor::Float64 = 1.0
    imag_ref_index_scale_factor::Float64 = 1.0
    real_chiral_factor::Float64 = 0.0
    imag_chiral_factor::Float64 = 0.0
    medium_real_ref_index::Float64 = 1.0
    medium_imag_ref_index::Float64 = 0.0
    medium_real_chiral_factor::Float64 = 0.0
    medium_imag_chiral_factor::Float64 = 0.0
    target_euler_angles_deg::Array{Float64,1} = [0.0, 0.0, 0.0]
    mie_epsilon::Float64 = 1e-6
    translation_epsilon::Float64 = 1e-6
    solution_epsilon::Float64 = 1e-8
    max_number_iterations::Int64 = 5000
    store_translation_matrix::Bool = false
    sm_number_processors::Int64 = 10
    near_field_distance::Float64 = 1e8
    iterations_per_correction::Int64 = 20
    min_scattering_angle_deg::Float64 = 0.0
    max_scattering_angle_deg::Float64 = 180.0
    min_scattering_plane_angle_deg::Float64 = 0.0
    max_scattering_plane_angle_deg::Float64 = 0.0
    delta_scattering_angle_deg::Float64 = 1.0
    frame::String = "target"
    normalize_scattering_matrix::Bool = true
    gaussian_beam_constant::Float64 = 0.0
    gaussian_beam_focal_point::Array{Float64,1} = [0.0, 0.0, 0.0]
    orientation::String = "fixed"
    incident_azimuth_angle_deg::Float64 = 0.0
    incident_polar_angle_deg::Float64 = 0.0
    calculate_scattering_coefficients::Bool = true
    scattering_coefficient_file::String = ""
    track_iterations::Bool = true
    calculate_near_field::Bool = true
    near_field_plane_coord::Int64 = 1
    near_field_plane_vertices::Array{Float64,1} = [10.0, 10.0]
    spacial_step_size::Float64 = 0.5
    polarization_angle_deg::Float64 = 0.0
    near_field_output_file::String = "nf_default_out.dat"
    near_field_output_data::Int64 = 2
    plane_wave_epsilon::Float64 = 1e-4
    calculate_t_matrix::Bool = true
    t_matrix_file::String = "tm_default.dat"
    t_matrix_convergence_epsilon::Float64 = 1e-7
    spheres::Array{SphereConfig}
end

"""
    MSTMConfig(filename::String)

Read configurations from a .toml file.
"""
function MSTMConfig(filename::String)
    return Configurations.from_toml(MSTMConfig, filename)
end

struct Sphere end

function Sphere(config::MSTMConfig, sphere_config::SphereConfig) end

struct MSTMParameters end

function MSTMParameters(config::MSTMConfig) end

struct TMaxtrix end

"""
    TMaxtrix(filename::String)

Read a previously calculated T matrix.
"""
function TMaxtrix(filename::String) end

end # module Data


isdefined(Base, :__precompile__) && __precompile__()

module SHA

using Compat

# Export convenience functions, context types, update!() and digest!() functions
export sha1, SHA1_CTX, update!, digest!
export sha224, sha256, sha384, sha512
export sha2_224, sha2_256, sha2_384, sha2_512
export sha3_224, sha3_256, sha3_384, sha3_512
export SHA224_CTX, SHA256_CTX, SHA384_CTX, SHA512_CTX
export SHA2_224_CTX, SHA2_256_CTX, SHA2_384_CTX, SHA2_512_CTX
export SHA3_224_CTX, SHA3_256_CTX, SHA3_384_CTX, SHA3_512_CTX

include("constants.jl")
include("types.jl")
include("base_functions.jl")
include("sha1.jl")
include("sha2.jl")
include("sha3.jl")
include("common.jl")

# Create data types and convenience functions for each hash implemented
for (f, ctx) in [(:sha1, :SHA1_CTX),
                 (:sha224, :SHA224_CTX),
                 (:sha256, :SHA256_CTX),
                 (:sha384, :SHA384_CTX),
                 (:sha512, :SHA512_CTX),
                 (:sha2_224, :SHA2_224_CTX),
                 (:sha2_256, :SHA2_256_CTX),
                 (:sha2_384, :SHA2_384_CTX),
                 (:sha2_512, :SHA2_512_CTX),
                 (:sha3_224, :SHA3_224_CTX),
                 (:sha3_256, :SHA3_256_CTX),
                 (:sha3_384, :SHA3_384_CTX),
                 (:sha3_512, :SHA3_512_CTX),]
    @eval begin
        # Our basic function is to process arrays of bytes
        function $f(data::Array{UInt8,1})
            ctx = $ctx()
            update!(ctx, data)
            return digest!(ctx)
        end

        # AbstractStrings are a pretty handy thing to be able to crunch through
        $f(str::AbstractString) = $f(convert(Array{UInt8,1}, str))

        # Convenience function for IO devices, allows for things like:
        # open("test.txt") do f
        #     sha256(f)
        # done
        $f(io::IO) = $f(read(io))
    end
end

end #module SHA


module AudioSchedules

import Base:
    close,
    getindex,
    iterate,
    IteratorEltype,
    IteratorSize,
    show,
    size,
    wait,
    write
using Base: EltypeUnknown, SizeUnknown
using Base.FastMath: sin_fast
using Base.Meta: ParseError
using Base.Threads: nthreads
using DataStructures: SortedSet
using InfiniteArrays: Fill, ∞
using Interpolations: CubicSplineInterpolation
using NLsolve: nlsolve
using PortAudio: Messanger, PortAudioStream, Scribe, write_buffer
import PortAudio: get_input_type, get_output_type
using RegularExpressions: capture, CONSTANTS, of, pattern, raw, short
import SampledSignals: nchannels, SampleBuf, samplerate
using Unitful: dB, Hz, μPa, s

const FREQUENCY = typeof(1.0Hz)
const RATE = typeof(1.0 / s)
const TIME = typeof(0.0s)

include("miscellaneous.jl")
include("series.jl")

export AudioSchedule

const ORCHESTRA = Vector{Tuple{Any, TIME}}
const ACTIVATED = Vector{Bool}
const TRIGGERS = SortedSet{Tuple{TIME, Int}}

struct AudioSchedule
    sample_rate::FREQUENCY
    orchestra::ORCHESTRA
    activated::ACTIVATED
    triggers::TRIGGERS
end

export AudioSchedule

"""
    AudioSchedule(; sample_rate = 44100Hz)

Create an empty `AudioSchedule`.
You can [`add!`](@ref) new synthesizers to the a_schedule.
Specify a `sample_rate` with units per time, like `1/s` or `Hz`.

```jldoctest audio_schedule
julia> using AudioSchedules

julia> using Unitful: s, Hz

julia> a_schedule = AudioSchedule()
0.0 s 44100.0 Hz AudioSchedule

julia> add!(
           a_schedule,
           Map(sin, Cycles(440Hz)),
           0s,
           0,
           Line => 1s,
           0.05,
           Line => 1s,
           0.05,
           Line => 1s,
           0,
       )

julia> add!(
            a_schedule,
            Map(sin, Cycles(660Hz)),
            2s,
            0,
            Line => 1s,
            0.05,
            Line => 1s,
            0.05,
            Line => 1s,
            0,
        )

julia> a_schedule
5.0 s 44100.0 Hz AudioSchedule
```

You can use write the a_schedule directly to a `PortAudioStream`.
However, to do so, the `PortAudioStream` must have a [`Weaver`](@ref) writer.
The `PortAudioStream` must have exactly 0 input channels, 1 output channel, and a matching
sample rate.

```jldoctest audio_schedule
julia> using PortAudio: PortAudioStream

julia> PortAudioStream(0, 1, writer = Weaver(), warn_xruns = false) do stream
           write(stream, a_schedule)
       end

```

You can save an `AudioSchedule` as a `SampledSignals.SampleBuf`.

```jldoctest audio_schedule
julia> using SampledSignals: SampleBuf

julia> saved = SampleBuf(a_schedule)
220500-frame, 1-channel SampleBuf{Float64, 2}
5.0s sampled at 44100.0Hz
▂▂▃▃▃▃▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▄▃▃▃▃▂▂

julia> PortAudioStream(0, 1, warn_xruns = false) do stream
           write(stream, saved)
       end
220500
```
"""
function AudioSchedule(; sample_rate = 44100Hz)
    AudioSchedule(sample_rate, ORCHESTRA(), ACTIVATED(), TRIGGERS())
end

function IteratorSize(::Type{AudioSchedule})
    SizeUnknown()
end
function IteratorEltype(::Type{AudioSchedule})
    EltypeUnknown()
end

function samplerate(a_schedule::AudioSchedule)
    a_schedule.sample_rate / Hz
end
function nchannels(a_schedule::AudioSchedule)
    1
end

function duration(a_schedule::AudioSchedule)
    triggers = a_schedule.triggers
    if length(triggers) > 0
        last(triggers)[1] - first(triggers)[1]
    else
        0.0s
    end
end

function show(io::IO, a_schedule::AudioSchedule)
    print(io, "$(duration(a_schedule)) $(a_schedule.sample_rate) AudioSchedule")
end

function make_envelope(start_level, (shape, duration), end_level, more_segments...)
    segments(shape, start_level, duration, end_level)...,
    make_envelope(end_level, more_segments...)...
end
function make_envelope(_)
    ()
end

function add_series!(a_schedule::AudioSchedule, series, start_time, duration)
    triggers = a_schedule.triggers
    orchestra = a_schedule.orchestra
    push!(orchestra, (series, start_time))
    push!(a_schedule.activated, false)
    instrument_index = length(orchestra)
    push!(triggers, (start_time, instrument_index))
    push!(triggers, (start_time + duration, instrument_index))
    nothing
end

function second(something)
    something[2]
end

function all_but_last(thing1, thing2, rest...)
    thing1, all_but_last(thing2, rest...)...
end

function all_but_last(_)
    ()
end

function skip(sample_rate, wave, skip_time)
    wave[(round(Int, skip_time * sample_rate)+1):end]
end

function make_wave(sample_rate, wave, shape, segment_start_time, segment_skip_time, duration)
    broadcast(
        *,
        # skip ahead to time in wave, then the one after
        skip(sample_rate, wave, segment_skip_time),
        make_series(shape, sample_rate),
    ),
    segment_start_time,
    duration
end

function triples(sample_rate, synthesizer, start_time, envelope...)
    wave = make_series(synthesizer, sample_rate)
    shapes_durations = make_envelope(envelope...)
    shapes = map(first, shapes_durations)
    durations = map(second, shapes_durations)
    segment_skip_times = all_but_last(cumsum((0.0s, durations...))...)
    segment_start_times = start_time .+ segment_skip_times
    map(
        let sample_rate = sample_rate, wave = wave
            function (shape, segment_start_time, segment_skip_time, duration)
                make_wave(sample_rate, wave, shape, segment_start_time, segment_skip_time, duration)
            end
        end,
        shapes,
        segment_start_times,
        segment_skip_times,
        durations,
    )
end

"""
    add!(a_schedule::AudioSchedule, synthesizer, start_time,
        start_level, shape => duration, end_level, more_segments...
    )

Add a synthesizer to a [`AudioSchedule`](@ref), where `synthesizer` is anything that supports [`make_series`](@ref),
`start_time` has units of time (like `s`), and the rest of the arguments specify the shape of the envelope.

For all envelope [`segments`](@ref), call

```
segments(shape, start_level, duration, end_level)
```

`duration` should have units of time (like `s`). For example,

```
add!(a_schedule, synthesizer, start_time, 0, Line => 1s, 1, Line => 1s, 0)
```

will call segments twice:

```
segments(Line, 0, 1s, 1)
segments(Line, 1, 1s, 0)
```
"""
function add!(a_schedule::AudioSchedule, synthesizer, start_time, piece_1, rest...)
    sample_rate = a_schedule.sample_rate
    map(
        let a_schedule = a_schedule
            function ((shaped, time, duration),)
                add_series!(a_schedule, shaped, time, duration)
            end
        end,
        triples(sample_rate, synthesizer, start_time, piece_1, rest...),
    )
    nothing
end

export add!

function iterate(a_schedule::AudioSchedule)
    iteration = iterate(a_schedule.triggers)
    # Done if no more triggers
    if iteration === nothing
        nothing
    else
        # Start at 0s
        trigger_iterate(a_schedule, 0.0s, iteration...)
    end
end

function iterate(a_schedule::AudioSchedule, (schedule_time, trigger_state))
    iteration = iterate(a_schedule.triggers, trigger_state)
    if iteration === nothing
        nothing
    else
        trigger_iterate(a_schedule, schedule_time, iteration...)
    end
end

function trigger_iterate(
    a_schedule,
    schedule_time,
    (trigger_time, instrument_index),
    trigger_state,
)
    activated = a_schedule.activated
    triggers = a_schedule.triggers
    # pull the rest of the triggers at the current time
    while schedule_time == trigger_time
        activated[instrument_index] = !activated[instrument_index]
        iteration = iterate(triggers, trigger_state)
        if iteration === nothing
            return nothing
        end
        (trigger_time, instrument_index), trigger_state = iteration
    end
    # adjust these based on how long ago they went on
    active = view(a_schedule.orchestra, activated)
    result = (
        if isempty(active)
            Fill(0.0, ∞)
        else
            broadcast(+, Iterators.map(
                let sample_rate = a_schedule.sample_rate, schedule_time = schedule_time
                    function ((instrument, instrument_start_time),)
                        skip(sample_rate, instrument, schedule_time - instrument_start_time)
                    end
                end,
                active
            )...)
        end,
        round(Int, (trigger_time - schedule_time) * a_schedule.sample_rate),
    ),
    (trigger_time, trigger_state)
    # pull the next trigger
    activated[instrument_index] = !activated[instrument_index]
    result
end

function trigger_iterate(_, __, ::Nothing)
    nothing
end

# TODO: make a SampledSignal without sacrificing performance
# TODO: figure out a way to avoid peaking !!!

"""
    struct Weaver <: Scribe end

A special `PortAudio` `Scribe`.
You can use a `Weaver` to write an [`AudioSchedule`](@ref) directly to your speakers.
"""
struct Weaver <: Scribe end

get_input_type(::Weaver, _) = Nothing
get_output_type(::Weaver, _) = Nothing

export Weaver

struct DividedRange <: AbstractVector{UnitRange{Int}}
    previous::Int
    share::Float64
    divisions::Int
end

function size(divided_range::DividedRange)
    (divided_range.divisions,)
end

function getindex(divided_range::DividedRange, index)
    share = divided_range.share
    divided_range.previous .+
    ((round(Int, (index - 1) * share)+1):round(Int, index * share))
end

function add_to(divided_range::DividedRange, increment)
    DividedRange(
        divided_range.previous + increment,
        divided_range.share,
        divided_range.divisions,
    )
end

function DividedRange(a_range::UnitRange, divisions)
    DividedRange(first(a_range) - 1, length(a_range) / divisions, divisions)
end

struct DualSubRanges <: AbstractVector{Tuple{UnitRange{Int},UnitRange{Int}}}
    frames_per_buffer::Int
    the_length::Int
    left_in_buffer::Int
    first_item::Tuple{DividedRange,DividedRange,Bool}
    last_item::Tuple{DividedRange,DividedRange,Bool}
    full_buffer_ranges::DividedRange
end

function DualSubRanges(
    first_buffer_at,
    frames_per_buffer,
    series_total,
    number_of_subsections,
)
    left_in_buffer = frames_per_buffer - first_buffer_at
    middle_buffers, last_buffer_at =
        fldmod(series_total - left_in_buffer, frames_per_buffer)
    is_last_complete = last_buffer_at == 0
    if is_last_complete
        middle_buffers = middle_buffers - 1
        last_buffer_at = frames_per_buffer
    end
    if series_total <= left_in_buffer
        first_series_sub_ranges = DividedRange(1:series_total, number_of_subsections)
        is_first_complete = is_last_complete
    else
        first_series_sub_ranges = DividedRange(1:left_in_buffer, number_of_subsections)
        is_first_complete = true
    end
    last_buffer_sub_ranges = DividedRange(1:last_buffer_at, number_of_subsections)
    DualSubRanges(
        frames_per_buffer,
        middle_buffers + 2,
        left_in_buffer,
        (
            first_series_sub_ranges,
            add_to(first_series_sub_ranges, first_buffer_at),
            is_first_complete,
        ),
        (
            add_to(
                last_buffer_sub_ranges,
                left_in_buffer + middle_buffers * frames_per_buffer,
            ),
            last_buffer_sub_ranges,
            is_last_complete,
        ),
        DividedRange(1:frames_per_buffer, number_of_subsections),
    )
end

function size(dual_sub_ranges::DualSubRanges)
    (dual_sub_ranges.the_length,)
end

function getindex(dual_sub_ranges::DualSubRanges, index::Int)
    the_length = dual_sub_ranges.the_length
    @boundscheck if index < 0 || index > the_length
        throw(BoundsError(dual_sub_ranges, index))
    end
    if index == 1
        dual_sub_ranges.first_item
    elseif index == the_length
        dual_sub_ranges.last_item
    else
        full_buffer_ranges = dual_sub_ranges.full_buffer_ranges
        (
            add_to(
                full_buffer_ranges,
                dual_sub_ranges.left_in_buffer +
                (index - 2) * dual_sub_ranges.frames_per_buffer,
            ),
            full_buffer_ranges,
            true,
        )
    end
end

struct TaskIO
    task::Task
    series_buffers_channel::Channel{Tuple{Any,Int}}
    sub_ranges_channel::Channel{Tuple{UnitRange{Int},UnitRange{Int},Bool}}
    output_channel::Channel{Nothing}
end

function close(task_io::TaskIO)
    close(task_io.series_buffers_channel)
    close(task_io.sub_ranges_channel)
    wait(task_io.task)
    close(task_io.output_channel)
end

function fill_buffer_signal!(
    buffer,
    series,
    sub_ranges_channel,
    output_channel,
)
    (series_sub_range, buffer_sub_range, is_complete) = take!(sub_ranges_channel)
    buffer.data[buffer_sub_range] .= view(series, series_sub_range)
    if is_complete
        put!(output_channel, nothing)
    end
end

function fill_series_signal!(buffer, (series, buffers), sub_ranges_channel, output_channel)
    foreach(
        let buffer = buffer, series = series, sub_ranges_channel = sub_ranges_channel, output_channel = output_channel
            function (_)
                fill_buffer_signal!(
                    buffer,
                    series,
                    sub_ranges_channel,
                    output_channel,
                )
            end
        end,
        1:buffers,
    )
end

function fill_all_signal!(
    buffer,
    series_buffers_channel,
    sub_ranges_channel,
    output_channel,
)
    foreach(
        let buffer = buffer,
            sub_ranges_channel = sub_ranges_channel,
            output_channel = output_channel

            function (series_buffers)
                fill_series_signal!(
                    buffer,
                    series_buffers,
                    sub_ranges_channel,
                    output_channel,
                )
            end
        end,
        series_buffers_channel,
    )
end

function fill_all_task_io(buffer)
    series_buffers_channel = Channel{Tuple{Any,Int}}(Inf)
    sub_ranges_channel = Channel{Tuple{UnitRange{Int},UnitRange{Int},Bool}}(0)
    output_channel = Channel{Nothing}(0)
    task = Task(
        let buffer = buffer,
            series_buffers_channel = series_buffers_channel,
            sub_ranges_channel = sub_ranges_channel,
            output_channel = output_channel

            function ()
                fill_all_signal!(
                    buffer,
                    series_buffers_channel,
                    sub_ranges_channel,
                    output_channel,
                )
            end
        end,
    )
    task.sticky = false
    schedule(task)
    TaskIO(task, series_buffers_channel, sub_ranges_channel, output_channel)
end

function fill_all_task_ios(buffer; number_of_tasks = nthreads() - 1)
    map(let buffer = buffer
        function (_)
            fill_all_task_io(buffer)
        end
    end, 1:number_of_tasks)
end

function feed_write!(task_ios, buffer, (series_sub_ranges, buffer_sub_ranges, is_complete))
    foreach(
        let is_complete = is_complete
            function ((task_io, series_sub_range, output_sub_range),)
                put!(
                    task_io.sub_ranges_channel,
                    (series_sub_range, output_sub_range, is_complete),
                )
            end
        end,
        zip(task_ios, series_sub_ranges, buffer_sub_ranges),
    )
    if is_complete
        foreach(function (task_io)
            take!(task_io.output_channel)
        end, task_ios)
        write_buffer(buffer)
    end
end

function write_series!(task_ios, series, series_total, buffer, buffer_at)
    dual_sub_ranges =
        DualSubRanges(buffer_at, buffer.frames_per_buffer, series_total, length(task_ios))
    foreach(let series_buffers = (series, length(dual_sub_ranges))
        function (task_io)
            put!(task_io.series_buffers_channel, series_buffers)
        end
    end, task_ios)
    foreach(let task_ios = task_ios, buffer = buffer
        function (ranges)
            feed_write!(task_ios, buffer, ranges)
        end
    end, dual_sub_ranges)
    last(last(last(dual_sub_ranges)[2]))
end

function write(
    stream::PortAudioStream{<:Messanger{<:Any,Weaver},<:Any},
    a_schedule::AudioSchedule;
    number_of_tasks = nthreads() - 1,
)
    if number_of_tasks == 0
        throw(ArgumentError("Must write with at least 1 task"))
    end
    if nchannels(stream.source_messanger) != 0
        throw(ArgumentError("$stream does not have 0 input channels"))
    end
    sink_messanger = stream.sink_messanger
    if nchannels(sink_messanger) != 1
        throw(ArgumentError("$stream does not have 1 output channel"))
    end
    if samplerate(stream) != samplerate(a_schedule)
        throw(ArgumentError("Sample rates of $stream and $a_schedule do not match"))
    end
    buffer = sink_messanger.buffer
    buffer_at = 0
    # TODO: Move task_ios inside of the Weaver object and reuse them
    task_ios = fill_all_task_ios(buffer; number_of_tasks = number_of_tasks)
    try
        for (series, series_total) in a_schedule
            buffer_at = write_series!(task_ios, series, series_total, buffer, buffer_at)
        end
    finally
        foreach(close, task_ios)
    end
    # This might make the previous TODO difficult. Do we need to pass-back a shorter-than-expected result?
    if buffer_at > 0
        write_buffer(buffer, buffer_at)
    end
    nothing
end

struct TaskInput{Input}
    task::Task
    input_channel::Channel{Input}
end

function close(task_input::TaskInput)
    close(task_input.input_channel)
    wait(task_input.task)
end

function feed!(task_inputs, series, series_total, output_at)
    number_of_tasks = length(task_inputs)
    series_ranges = DividedRange(1:series_total, number_of_tasks)
    foreach(
        let series = series
            function feed_task_input!((task_input, series_range, output_range))
                put!(task_input.input_channel, (series, series_range, output_range))
            end
        end,
        zip(task_inputs, series_ranges, add_to(series_ranges, output_at)),
    )
    output_at + series_total
end

function fill_output!(output_data, (series, series_sub_range, output_range))
    output_data[output_range] .= view(series, series_sub_range)
end

function fill_all!(input_channel, output_data)
    foreach(let output_data = output_data
        function (series_ranges)
            fill_output!(output_data, series_ranges)
        end
    end, input_channel)
end

function fill_all_task_input(output_data)
    input_channel = Channel{Tuple{Any,UnitRange{Int},UnitRange{Int}}}(Inf)
    task = Task(let input_channel = input_channel, output_data = output_data
        function ()
            fill_all!(input_channel, output_data)
        end
    end)
    task.sticky = false
    schedule(task)
    TaskInput(task, input_channel)
end

function SampleBuf(a_schedule::AudioSchedule; number_of_tasks = nthreads() - 1)
    if number_of_tasks == 0
        throw(ArgumentError("Must write with at least 1 task"))
    end
    number = sum(Iterators.map(second, a_schedule))
    output = SampleBuf(Float64, samplerate(a_schedule), number, 1)
    task_inputs = map(let output_data = output.data
        function (_)
            fill_all_task_input(output_data)
        end
    end, 1:number_of_tasks)
    try
        output_at = 0
        for (series, series_total) in a_schedule
            output_at = feed!(task_inputs, series, series_total, output_at)
        end
    finally
        foreach(close, task_inputs)
    end
    output
end

include("equal_loudness.jl")

end


py"""
import torch.nn as nn
import torch
import torch.nn.functional as F
import numpy as np
from torch.nn import init

class Down_Conv(nn.Module):

    def __init__(self, in_channels, out_channels, mid_channels=None, mode = 'double_conv'):
        super().__init__()
        assert mode in {'double_conv', 'single_conv'}
        if not mid_channels:
            mid_channels = out_channels
        if mode == 'double_conv':
          self.conv = nn.Sequential(
              nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias = False),
              nn.BatchNorm2d(mid_channels),
              nn.LeakyReLU(0.1, inplace=True),
              nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias = False),
              nn.BatchNorm2d(out_channels),
              nn.LeakyReLU(0.1, inplace=True)
          )
        else:
          self.conv = nn.Sequential(
              nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias = False),
              nn.BatchNorm2d(mid_channels),
              nn.LeakyReLU(0.1, inplace=True)
          )

    def forward(self, x):
        return self.conv(x)

class Up_Conv(nn.Module):

    def __init__(self, in_channels, out_channels, mid_channels=None, mode = 'double_conv'):
        super().__init__()
        assert mode in {'double_conv', 'single_conv'}
        if not mid_channels:
            mid_channels = out_channels
        if mode == 'double_conv':
          self.conv = nn.Sequential(
              nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias = False),
              nn.BatchNorm2d(mid_channels),
              nn.ReLU(inplace=True),
              nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias = False),
              nn.BatchNorm2d(out_channels),
              nn.ReLU(inplace=True)
          )
        else:
          self.conv = nn.Sequential(
              nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias = False),
              nn.BatchNorm2d(out_channels),
              nn.ReLU(inplace=True)
          )

    def forward(self, x):
        return self.conv(x)

class Down(nn.Module):

    def __init__(self, in_channels, out_channels, mode = 'double_conv'):
        super().__init__()
        self.maxpool_conv = nn.Sequential(
            nn.MaxPool2d(2),
            Down_Conv(in_channels, out_channels, mode = mode)
        )

    def forward(self, x):
        return self.maxpool_conv(x)

class Up(nn.Module):

    def __init__(self, in_channels, out_channels, nearest=True, mode = 'double_conv'):
        super().__init__()

        # if nearest, use the normal convolutions to reduce the number of channels
        if nearest:
            self.up = nn.Upsample(scale_factor=2, mode='nearest')
            self.conv = Up_Conv(in_channels, out_channels // 2, in_channels // 2, mode = mode)
        else:
            self.up = nn.ConvTranspose2d(in_channels , in_channels // 2, kernel_size=2, stride=2)
            self.conv = Up_Conv(in_channels, out_channels, mode = mode)

    def forward(self, x1, x2):
        x1 = self.up(x1)
        # input is CHW
        diffX = x2.size()[2] - x1.size()[2]
        diffY = x2.size()[3] - x1.size()[3]
        # print('sizes',x1.size(),x2.size(),diffX // 2, diffX - diffX//2, diffY // 2, diffY - diffY//2)
        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
                        diffY // 2, diffY - diffY // 2])
        # if you have padding issues, see
        # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
        # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd
        x = torch.cat([x2, x1], dim=1)
        return self.conv(x)

class OutConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(OutConv, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)

    def forward(self, x):
        return self.conv(x)

class UNet(nn.Module):
    def __init__(self, n_channels, n_classes, init_filter=64, num_down=4, nearest=True, is_single=True):
        super(UNet, self).__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes
        self.nearest = nearest
        self.num_down = num_down
        factor = 2 if nearest else 1
        filters = []
        for i in range(num_down + 1):
            filters.append(init_filter * (2 ** i))
        # print(filters)
        self.inc = Down_Conv(n_channels, filters[0])
        down_list = []
        up_list = []
        if is_single:
            for i in range(num_down):
                if i == num_down - 1:
                    down_list.append(Down(filters[i], filters[i + 1] // factor, mode='single_conv'))
                else:
                    down_list.append(Down(filters[i], filters[i + 1], mode='single_conv'))
            self.down = nn.ModuleList(down_list)
            for i in range(num_down):
                if i == num_down - 1:
                    up_list.append(
                        Up(filters[num_down - i], filters[num_down - i - 1] * factor, nearest, mode='double_conv'))
                else:
                    up_list.append(Up(filters[num_down - i], filters[num_down - i - 1], nearest, mode='single_conv'))
                self.up = nn.ModuleList(up_list)
        else:
            for i in range(num_down):
                if i == num_down - 1:
                    down_list.append(Down(filters[i], filters[i + 1] // factor, mode='double_conv'))
                else:
                    down_list.append(Down(filters[i], filters[i + 1], mode='double_conv'))
            self.down = nn.ModuleList(down_list)
            for i in range(num_down):
                if i == num_down - 1:
                    up_list.append(
                        Up(filters[num_down - i], filters[num_down - i - 1] * factor, nearest, mode='double_conv'))
                else:
                    up_list.append(Up(filters[num_down - i], filters[num_down - i - 1], nearest, mode='double_conv'))
                self.up = nn.ModuleList(up_list)

        self.outc = OutConv(filters[0], n_classes)

    def forward(self, x):
        down = []
        x = self.inc(x)
        down.append(x)
        for i, down_module in enumerate(self.down):
            x = down_module(x)
            down.append(x)
        for i, up_module in enumerate(self.up):
            if i == 0:
                x = up_module(down[-1], down[-2])
            else:
                x = up_module(x, down[self.num_down - 1 - i])
        logits = self.outc(x)
        return logits

def init_weights(net, init_type='normal', mean = 0.0, init_gain=0.02):
    def init_func(m):  # define the initialization function
        classname = m.__class__.__name__
        # print(classname)
        if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
            if init_type == 'normal':
                init.normal_(m.weight.data, mean, init_gain)
            elif init_type == 'xavier':
                init.xavier_normal_(m.weight.data, gain=init_gain)
            elif init_type == 'kaiming':
                init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
            elif init_type == 'orthogonal':
                init.orthogonal_(m.weight.data, gain=init_gain)
            else:
                raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
            if hasattr(m, 'bias') and m.bias is not None:
                init.constant_(m.bias.data, mean)

        elif classname.find('BatchNorm') != -1:  # BatchNorm Layer's weight is not a matrix; only normal distribution applies.
                init.normal_(m.weight.data, 1.0, init_gain)
                init.constant_(m.bias.data, mean)

    print('initialize network with %s' % init_type)
    net.apply(init_func)  # apply the initialization function <init_func>

class My_model(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(My_model, self).__init__()
        self.UNet = UNet(in_channels, out_channels, init_filter=8, num_down=3, nearest= True, is_single= True)
        init_weights(self.UNet, init_type= 'kaiming', mean= 0, init_gain= 0.02)
    def forward(self, img):
        out = self.UNet(img)
        return out

"""


# http://www.neos-server.org/neos/NEOS-API.html
module NEOS

@warn("Make sure you comply with the NEOS terms of use: http://www.neos-server.org/neos/termofuse.html")

using LightXML
using HTTP
using CodecZlib
using SparseArrays
using Base64

using AmplNLWriter
using MathProgBase
const MPB = MathProgBase

include("NEOSServer.jl")
include("NEOSSolver.jl")
include("NEOSSolverInterface.jl")
include("MPSWriter.jl")
include("nl.jl")
include("mps.jl")

include("solvers/CPLEX_MPS.jl")
include("solvers/MOSEK_MPS.jl")
include("solvers/SYMPHONY_MPS.jl")
include("solvers/FICOXpress_MPS.jl")
include("solvers/CPLEX_NL.jl")

export NEOSServer, NEOSSolver,
    NEOSCPLEXSolver, NEOSMOSEKSolver, NEOSSYMPHONYSolver, NEOSXpressSolver,
    addparameter!, addemail!,

    # NEOS API functions
    neosHelp, emailHelp, welcome, version, ping, printQueue,
    listAllSolvers, listCategories,
    getSolverTemplate,
    listSolversInCategory,
    submitJob,
    getJobStatus, killJob, getFinalResults, getFinalResultsNonBlocking,
    getJobInfo,
    getIntermediateResults, getIntermediateResultsNonBlocking,
    createNEOSJob, getNEOSJobXML
end


x_shift = deepcopy(x_sol)
for x in x_shift
    x[3] += 0.01
end
z = [minimal_to_maximal(env.mechanism, x) for x in x_shift]

t = 1 #10, 20, 30, 41
set_robot(env.vis, env.mechanism, z[41])

# ## visualize
x_view = [[x_shift[1] for t = 1:15]..., x_shift..., [x_shift[end] for t = 1:15]...]
visualize(env, x_view)

set_camera!(vis, cam_pos=[0,-50,0], zoom=30)
set_floor!(env.vis, z=0.01)


using TMLEEpistasis

include(joinpath(pkgdir(TMLEEpistasis), "test", "runtests.jl"))


using CUDA
using CUDA: curand_rng
using CUDA.CUSPARSE: CuSparseMatrix, CuSparseMatrixCSR, CuSparseMatrixCOO
using CUDA.CUSOLVER: CuQR

export cuda_gc, gpu

const CuBaseField{B,M,T,A<:CuArray} = BaseField{B,M,T,A}

# printing
typealias(::Type{CuArray{T,N}}) where {T,N} = "CuArray{$T,$N}"
Base.print_array(io::IO, X::Diagonal{<:Any,<:CuBaseField}) = Base.print_array(io, cpu(X))

# a function version of @cuda which can be referenced before CUDA is
# loaded as long as it exists by run-time (unlike the macro @cuda which must
# exist at compile-time)
function cuda(f, args...; threads=256)
    @cuda threads=threads f(args...)
end

is_gpu_backed(::BaseField{B,M,T,A}) where {B,M,T,A<:CuArray} = true
global_rng_for(::Type{<:CuArray}) = curand_rng()

# handy conversion functions and macros
@doc doc"""

    gpu(x)

Recursively moves x to GPU, but unlike `CUDA.cu`, doesn't also convert
to Float32. Equivalent to `adapt_structure(CuArray, x)`. Returns nothing.
"""
gpu(x) = adapt_structure(CuArray, x)

adapt_structure(::CUDA.Float32Adaptor, proj::ProjLambert) = adapt_structure(CuArray{Float32}, proj)

function Cℓ_to_2D(Cℓ, proj::ProjLambert{T,<:CuArray}) where {T}
    # todo: remove needing to go through cpu here:
    gpu(Complex{T}.(nan2zero.(Cℓ.(cpu(proj.ℓmag)))))
end

### misc
# the generic versions of these trigger scalar indexing of CUDA, so provide
# specialized versions:
pinv(D::Diagonal{T,<:CuBaseField}) where {T} = Diagonal(@. ifelse(isfinite(inv(D.diag)), inv(D.diag), $zero(T)))
inv(D::Diagonal{T,<:CuBaseField}) where {T} = any(Array((D.diag.==0)[:])) ? throw(SingularException(-1)) : Diagonal(inv.(D.diag))
fill!(f::CuBaseField, x) = (fill!(f.arr,x); f)
==(a::CuBaseField, b::CuBaseField) = (==)(promote(a.arr, b.arr)...)
sum(f::CuBaseField; dims=:) = (dims == :) ? sum(f.arr) : (1 in dims) ? error("Sum over invalid dims of CuFlatS0.") : f

# adapting of SparseMatrixCSC ↔ CuSparseMatrixCSR (otherwise dense arrays created)
adapt_structure(::Type{<:CuArray}, L::SparseMatrixCSC)   = CuSparseMatrixCSR(L)
adapt_structure(::Type{<:Array},   L::CuSparseMatrixCSR) = SparseMatrixCSC(L)
adapt_structure(::Type{<:CuArray}, L::CuSparseMatrixCSR) = L
adapt_structure(::Type{<:Array},   L::SparseMatrixCSC)   = L

# CUDA somehow missing this one
# see https://github.com/JuliaGPU/CuArrays.jl/issues/103
# and https://github.com/JuliaGPU/CuArrays.jl/pull/580
ldiv!(qr::CuQR, x::CuVector) = qr.R \ (CuMatrix(qr.Q)' * x)

# some Random API which CUDA doesn't implement yet
Random.randn(rng::CUDA.CURAND.RNG, T::Random.BitFloatType) =
    cpu(randn!(rng, CuVector{T}(undef,1)))[1]

# perhaps minor type-piracy, but this lets us simulate into a CuArray using the
# CPU random number generator
Random.randn!(rng::MersenneTwister, A::CuArray) =
    (A .= adapt(CuArray, randn!(rng, adapt(Array, A))))

"""
    cuda_gc()

Gargbage collect and reclaim GPU memory (technically should never be
needed to do this by hand, but sometimes helps with GPU OOM errors)
"""
function cuda_gc()
    isdefined(Main,:Out) && empty!(Main.Out)
    GC.gc(true)
    CUDA.reclaim()
end


# differentiated shock capturing functions

# main entry point
"""
  Main function for assembling the shock capturing terms into the Jacobian.

  **Inputs**

   * mesh
   * sbp
   * eqn
   * opts
   * capture: an [`AbstractShockCapturing`](@ref)
   * assem: an [`AssembleElementData`](@ref)
"""
function applyShockCapturing_diff(mesh::AbstractMesh, sbp::AbstractOperator,
                             eqn::EulerData, opts,
                             capture::AbstractVolumeShockCapturing,
                             assem::AssembleElementData)

  # nothing to do here, call the implementation
  sensor = getShockSensor(capture)
  calcShockCapturing_diff(mesh, sbp, eqn, opts, sensor, capture, assem)

  return nothing
end

function applyShockCapturing_diff(mesh::AbstractMesh, sbp::AbstractOperator,
                             eqn::EulerData{Tsol, Tres}, opts,
                             capture::AbstractFaceShockCapturing,
                             assem::AssembleElementData) where {Tsol, Tres}

  sensor = getShockSensor(capture)
  _applyShockCapturing_diff(mesh, sbp, eqn, opts, sensor, capture, assem)

  return nothing
end

function _applyShockCapturing_diff(mesh::AbstractMesh, sbp::AbstractOperator,
                             eqn::EulerData{Tsol, Tres}, opts,
                             sensor::AbstractShockSensor,
                             capture::AbstractFaceShockCapturing,
                             assem::AssembleElementData) where {Tsol, Tres}

  if mesh.commsize > 1 &&
    getParallelData(eqn.shared_data) != PARALLEL_DATA_ELEMENT

    error("shock capturing requires PARALLEL_DATA_ELEMENT")
  end

  assertReceivesWaited(eqn.shared_data)  # parallel communication for the regular
                                     # face integrals should already have
                                     # finished parallel communication

  # re-using the shockmesh from one iteration to the next makes allows
  # the algorithm to re-use existing arrays that depend on the number of
  # elements with the shock in it.
  shockmesh = eqn.params.shockmesh
  reset(shockmesh)
  for i=1:mesh.numEl
    q_i = ro_sview(eqn.q, :, :, i)
    coords_i = ro_sview(mesh.coords, :, :, i)
    dxidx_i = ro_sview(mesh.dxidx, :, :, :, i)
    jac_i = ro_sview(mesh.jac, :, i)

    val = isShockElement(eqn.params, sbp, sensor, q_i, i, coords_i, dxidx_i,
                            jac_i)
    if val
      # push to shockmesh
      push!(shockmesh, i)
    end
  end

  completeShockElements(mesh, shockmesh)

  allocateArrays(capture, mesh, shockmesh)

  # call shock capturing scheme
  calcShockCapturing_diff(mesh, sbp, eqn, opts, capture, shockmesh, assem)

  return nothing
end

#------------------------------------------------------------------------------
# revq

"""
  Reverse mode of `applyShockCapturing` wrt q

  **Inputs**

   * mesh
   * sbp
   * eqn
   * opts
   * data: an [`AbstractShockCaputring`](@ref) object
"""
function applyShockCapturing_revq(mesh::AbstractMesh, sbp::AbstractOperator,
                             eqn::EulerData, opts,
                             capture::AbstractVolumeShockCapturing)

  # unlike AbstractFaceShockCapturing, nothing to do here

  sensor = getShockSensor(capture)
  calcShockCapturing_revq(mesh, sbp, eqn, opts, sensor, capture)

  return nothing
end

# Face shock capturing version
function applyShockCapturing_revq(mesh::AbstractMesh, sbp::AbstractOperator,
                             eqn::EulerData, opts,
                             capture::AbstractFaceShockCapturing)

  sensor = getShockSensor(capture)
  shockmesh = eqn.params.shockmesh
  setupShockCapturing(mesh, sbp, eqn, opts, sensor, capture, shockmesh)

  # call shock capturing scheme
  calcShockCapturing_revq(mesh, sbp, eqn, opts, capture, shockmesh)

  return nothing
end

#------------------------------------------------------------------------------
# revm

"""
  Reverse mode of `applyShockCapturing` wrt metrics

  **Inputs**

   * mesh
   * sbp
   * eqn
   * opts
   * data: an [`AbstractShockCaputring`](@ref) object
"""
function applyShockCapturing_revm(mesh::AbstractMesh, sbp::AbstractOperator,
                             eqn::EulerData, opts,
                             capture::AbstractVolumeShockCapturing)

  # unlike AbstractFaceShockCapturing, nothing to do here

  sensor = getShockSensor(capture)
  initForRevm(sensor)
  calcShockCapturing_revm(mesh, sbp, eqn, opts, sensor, capture)
  finishRevm(mesh, sbp, eqn, opts, sensor)

  return nothing
end

function applyShockCapturing_revm(mesh::AbstractMesh, sbp::AbstractOperator,
                             eqn::EulerData, opts,
                             capture::AbstractFaceShockCapturing)

  # unlike AbstractFaceShockCapturing, nothing to do here

  sensor = getShockSensor(capture)
  shockmesh = eqn.params.shockmesh
  initForRevm(sensor)
  setupShockCapturing(mesh, sbp, eqn, opts, sensor, capture, shockmesh)
  calcShockCapturing_revm(mesh, sbp, eqn, opts, capture, shockmesh)
  finishRevm(mesh, sbp, eqn, opts, sensor)

  return nothing
end