"""
Create `Evaluator` instance for a given `problem`.
"""
function make_evaluator(problem::OptimizationProblem, archive=nothing, params::Parameters=ParamsDict())
    workers = get(params, :Workers, Vector{Int}())
    nthreads = get(params, :NThreads, 0)
    if archive===nothing
        # make the default archive
        archiveCapacity = get(params, :ArchiveCapacity, 10)
        archive = TopListArchive(fitness_scheme(problem), numdims(problem), archiveCapacity)
    end
    if length(workers) > 0
        return ParallelEvaluator(problem, archive, pids=workers)
    elseif nthreads > 0
        return MultithreadEvaluator(problem, archive, nworkers=nthreads)
    else
        return ProblemEvaluator(problem, archive)
    end
end

"""
Optimization Run Controller.
Manages problem optimization using the specified method.

See `OptController`.
"""
mutable struct OptRunController{O<:Optimizer, E<:Evaluator}
    optimizer::O   # optimization algorithm
    evaluator::E   # problem evaluator

    trace_mode::Symbol # controller state trace mode (:verbose, :compact, :silent)
                       # :silent makes tr() generate no output)
    save_trace::Bool # FIXME if traces should be saved to a file
    trace_interval::Float64 # periodicity of calling trace_progress()

    callback_function::Function
    callback_interval::Float64

    # termination conditions
    max_steps::Int      # maximal number of steps
    max_fevals::Int # maximal number of fitness evaluations
    max_steps_without_fevals::Int # stop optimization if no func evals in this many steps (indicates a converged/degenerate search)
    max_steps_without_progress::Int # stop optimization if no improvement in this many steps (indicates a converged/degenerate search)
    max_time::Float64   # maximal time, 0 to ignore

    min_delta_fitness_tol::Float64 # minimal difference between current best fitness and second-best one
    fitness_tol::Float64  # minimal difference between current best fitness and the known optmimum

    num_steps::Int       # optimization steps done
    num_better::Int      # number of steps that improved best fitness

    num_better_since_last_report::Int # ditto, since last call to trace_progress()
    num_steps_since_last_report::Int # number of steps since last call to trace_progress()

    last_num_fevals::Int # the number of function evals on the previous step
    num_steps_without_fevals::Int # the number of steps without the function evals

    start_time::Float64 # time optimization started, 0 if not running yet
    stop_time::Float64  # time optimization stopped, 0 if still running
    last_report_time::Float64 # last time trace_progress() was called
    last_callback_time::Float64 # last time callback function was called

    stop_reason::String # the reason for algorithm termination, empty if it's not terminated
    shutdown::Bool # shutdown optimization run iff true
end

"""
Create `OptRunController` for a given problem using specified `optimizer`.

#Arguments
    * `optimizer` initialized optimization method
    * `evaluator` the evaluator of the problem fitness
    * `params` controller settings, see `DefaultParameters` for the default values:
        * `:MaxTime` max time in seconds (takes precedence over the other budget-related params if specified), 0.0 disables the check
        * `:MaxFuncEvals` max fitness evals (takes precedence over max iterations, but not max time), 0 disables the check
        * `:MaxSteps` max iterations gives the least control since different optimizers have different "size" of their "iterations"
        * `:MaxStepsWithoutProgress` max iterations without fitness improvement
        * `:MinDeltaFitnessTolerance` minimum delta fitness (difference between the two consecutive best fitness improvements) we can accept before terminating
        * `:FitnessTolerance` stop the optimization when the best fitness found is within this distance of the actual optimum (if known)
        * `:MaxNumStepsWithoutFuncEvals` stop optimization if no new fitness evals in this many steps (indicates a converged/degenerate search)
        * `:NumRepetitions` number of repetitions to run for each optimizer for each problem
        * `:TraceMode` how the optimizer state is traced to the STDOUT during the optimization (one of `:silent`, `:verbose`)
        * `:TraceInterval` the trace interval (in seconds)
        * `:SaveTrace` whether to save it to a file (defaults to `false`)
        * `:SaveFitnessTraceToCsv` whether the history of fitness changes during optimization should be save to a csv file
        * `:SaveParameters` save method/controller parameters to a JSON file
"""
function OptRunController(optimizer::O, evaluator::E, params) where {O<:Optimizer, E<:Evaluator}
    OptRunController{O,E}(optimizer, evaluator,
        [params[key] for key in Symbol[:TraceMode, :SaveTrace, :TraceInterval,
                      :CallbackFunction, :CallbackInterval,
                      :MaxSteps, :MaxFuncEvals, :MaxNumStepsWithoutFuncEvals, :MaxStepsWithoutProgress, :MaxTime,
                      :MinDeltaFitnessTolerance, :FitnessTolerance]]...,
        0, 0, 0, 0, 0, 0, 0.0, 0.0, 0.0, -1.0, "", false)
end

# stepping optimizer has it's own evaluator, get a reference
OptRunController(optimizer::SteppingOptimizer, problem::OptimizationProblem, params) =
    OptRunController(optimizer, evaluator(optimizer), params)
# all other optimizers are using make_evaluator() method to create evaluator by default
OptRunController(optimizer::Optimizer, problem::OptimizationProblem, params) =
    OptRunController(optimizer, make_evaluator(problem, nothing, params), params)
# We reuse evaluator and optimizer from a previous run if set up from such a run.
OptRunController(orc::OptRunController, params) =
    OptRunController(optimizer(orc), evaluator(orc), params)

# logging/tracing
function trace(ctrl::OptRunController, msg::AbstractString, obj = nothing)
    if ctrl.trace_mode != :silent
        print(msg)
        if obj !== nothing
            if isa(obj, AbstractString)
                print(obj)
            else
                show(IOContext(stdout, :compact => true), obj)
            end
        end
    end
    if ctrl.save_trace
        # No saving for now
    end
end

optimizer(ctrl::OptRunController) = ctrl.optimizer
evaluator(ctrl::OptRunController) = ctrl.evaluator
problem(ctrl::OptRunController) = problem(evaluator(ctrl))
isstarted(ctrl::OptRunController) = ctrl.start_time > 0
isrunning(ctrl::OptRunController) = isstarted(ctrl) && ctrl.stop_time == 0
isstopped(ctrl::OptRunController) = ctrl.stop_time > 0

num_steps(ctrl::OptRunController) = ctrl.num_steps
num_func_evals(ctrl::OptRunController) = num_evals(ctrl.evaluator)
stop_reason(ctrl::OptRunController) = ctrl.stop_reason
start_time(ctrl::OptRunController) = ctrl.start_time

best_candidate(ctrl::OptRunController) = best_candidate(ctrl.evaluator.archive)
best_fitness(ctrl::OptRunController) =  best_fitness(ctrl.evaluator.archive)

"""
    show_fitness(io, fit, [problem::OptimizationProblem])

Output fitness to the given IO stream.
`show_fitness()` method could be overridden for a specific problem, e.g.
to print the names of each objective.
"""
show_fitness(io::IO, fit::Number) = @printf(io, "%.9f", fit)

function show_fitness(io::IO, fit::NTuple{N}) where N
    print(io, "(")
    for i in 1:N
        if i > 1
            print(io, ", ")
        end
        @printf(io, "%.5f", fit[i])
    end
    print(io, ")")
end

function show_fitness(io::IO, fit::IndexedTupleFitness{N}) where N
    show_fitness(io, fit.orig)
    @printf(io, " agg=%.5f", fit.agg)
end

# problem-specific method defaults to problem-agnostic method
# note that the fitness type of the problem could be different from FA,
# because show_fitness() typically is called for the archived fitness
show_fitness(io::IO, fit::FA, problem::OptimizationProblem) where {FA} = show_fitness(io, fit)

"""
    format_fitness(fit, [problem::OptimizationProblem])

Format fitness into a string.
Calls `show_fitness()` under the hood.
"""
function format_fitness(fit::Any, problem::OptimizationProblem)
    buf = IOBuffer(false, true)
    show_fitness(buf, fit, problem)
    return bytestring(buf.data)
end

function format_fitness(fit::Any)
    buf = IOBuffer(false, true)
    show_fitness(buf, fit)
    return bytestring(buf.data)
end

function elapsed_time(ctrl::OptRunController)
    isrunning(ctrl) ? time() - ctrl.start_time : ctrl.stop_time - ctrl.start_time
end

function check_stop_condition(ctrl::OptRunController)
    if ctrl.shutdown
        return "Run explicitly stopped via shutdown method"
    end

    if ctrl.max_time > 0 && elapsed_time(ctrl) > ctrl.max_time
        return "Max time ($(ctrl.max_time) s) reached"
    end

    if ctrl.max_fevals > 0 && num_func_evals(ctrl) > ctrl.max_fevals
        return "Max number of function evaluations ($(ctrl.max_fevals)) reached"
    end

    if ctrl.max_steps_without_fevals > 0 && ctrl.num_steps_without_fevals > ctrl.max_steps_without_fevals
        return "Too many steps ($(ctrl.num_steps_without_fevals)) without any function evaluations (probably search has converged)"
    end

    if ctrl.max_steps > 0 && num_steps(ctrl) > ctrl.max_steps
        return "Max number of steps ($(ctrl.max_steps)) reached"
    end

    return check_stop_condition(ctrl.evaluator, ctrl)

    return "" # empty string, no termination
end

@inline function callback(ctrl::OptRunController)
    if ctrl.callback_interval >= 0.0
        ctrl.callback_function(ctrl)
        ctrl.last_callback_time = time()
    end
end

function trace_progress(ctrl::OptRunController)
    # update the counters
    ctrl.last_report_time = time()
    total_improvement_rate = ctrl.num_better/num_steps(ctrl)
    recent_improvement_rate = ctrl.num_better_since_last_report/ctrl.num_steps_since_last_report
    ctrl.num_better_since_last_report = 0
    ctrl.num_steps_since_last_report = 0

    if ctrl.trace_mode == :silent
        return
    end

    # Always print step number, num fevals and elapsed time
    trace(ctrl, @sprintf("%.2f secs, %d evals, %d steps",
            elapsed_time(ctrl), num_func_evals(ctrl), num_steps(ctrl)))

    # Only print if this optimizer reports on number of better. They return 0
    # if they do not.
    if total_improvement_rate > 0.0
        trace(ctrl, @sprintf(", improv/step: %.3f (last = %.4f)",
                total_improvement_rate, recent_improvement_rate))
    end

    # Always print fitness if num_evals > 0
    if num_func_evals(ctrl) > 0
        trace(ctrl, ", fitness=")
        show_fitness(stdout, best_fitness(ctrl), problem(ctrl))
    end

    trace(ctrl, "\n")

    trace_state(stdout, ctrl.optimizer, ctrl.trace_mode)
end

function step!(ctrl::OptRunController{<:AskTellOptimizer})
    # The ask()/tell() interface is more general since you can mix and max
    # elements from several optimizers using it. However, in this top-level
    # execution function we do not make use of this flexibility...
    candidates = ask(ctrl.optimizer)
    rank_by_fitness!(ctrl.evaluator, candidates)
    return tell!(ctrl.optimizer, candidates)
end

# step for SteppingOptimizers
function step!(ctrl::OptRunController{<:SteppingOptimizer})
    step!(ctrl.optimizer)
    return 0
end

setup_optimizer!(ctrl::OptRunController{<:SteppingOptimizer}) =
    setup!(ctrl.optimizer)
setup_optimizer!(ctrl::OptRunController{<:AskTellOptimizer}) =
    setup!(ctrl.optimizer, ctrl.evaluator)

shutdown!(ctrl::OptRunController) = ctrl.shutdown = true

function shutdown_optimizer!(ctrl::OptRunController{<:SteppingOptimizer})
    shutdown!(ctrl.optimizer)
    shutdown!(ctrl)
end

function shutdown_optimizer!(ctrl::OptRunController{<:AskTellOptimizer})
    shutdown!(ctrl.optimizer)
    shutdown!(ctrl.evaluator)
    shutdown!(ctrl)
end

"""
    run!(ctrl::OptRunController)

Run optimization until one of the stopping conditions are satisfied.
"""
function run!(ctrl::OptRunController)
    trace(ctrl, "Starting optimization with optimizer $(name(ctrl.optimizer))\n")
    try
        setup_optimizer!(ctrl)

        ctrl.start_time = time()
        ctrl.num_steps = 0
        while isempty(ctrl.stop_reason)
            # Report on progress every now and then...
            if (time() - ctrl.last_report_time) > ctrl.trace_interval
                trace_progress(ctrl)
            end

            # Take the step and then update the counters
            nstep_better = step!(ctrl)
            ctrl.num_better += nstep_better
            ctrl.num_better_since_last_report += nstep_better
            ctrl.num_steps += 1
            ctrl.num_steps_since_last_report += 1
            if num_evals(ctrl.evaluator) == ctrl.last_num_fevals
                ctrl.num_steps_without_fevals += 1
            else
                ctrl.num_steps_without_fevals = 0
            end
            ctrl.last_num_fevals = num_func_evals(ctrl)

            # Callback every now and then (if a callback interval has been set)...
            if ctrl.callback_interval >= 0.0
                if ctrl.callback_interval == 0.0 || (ctrl.last_callback_time <= 0.0 ||
                        (time() - ctrl.last_callback_time) > ctrl.callback_interval)
                    callback(ctrl)
                end
            end

            # Check if we have a reason to stop on next loop
            ctrl.stop_reason = check_stop_condition(ctrl)
        end
        ctrl.stop_time = time()
    finally
        shutdown_optimizer!(ctrl)
    end
    trace(ctrl, @sprintf "\nOptimization stopped after %d steps and %.2f seconds\n" ctrl.num_steps elapsed_time(ctrl))

    return ctrl.stop_reason
end

function show_report(ctrl::OptRunController, population_stats=false)
    final_elapsed_time = elapsed_time(ctrl)
    trace(ctrl, "Termination reason: $(ctrl.stop_reason)\n")
    trace(ctrl, @sprintf "Steps per second = %.2f\n" num_steps(ctrl)/final_elapsed_time)
    trace(ctrl, @sprintf "Function evals per second = %.2f\n" num_func_evals(ctrl)/final_elapsed_time)
    trace(ctrl, @sprintf "Improvements/step = %.5f\n" ctrl.num_better/ctrl.max_steps)
    trace(ctrl, "Total function evaluations = $(num_func_evals(ctrl))\n")

    if population_stats && isa(ctrl.optimizer, PopulationOptimizer)
        trace(ctrl, "\nMean value (in population) per position:",  mean(population(ctrl.optimizer),1))
        trace(ctrl, "\n\nStd dev (in population) per position:", std(population(ctrl.optimizer),1))
    end

    trace(ctrl, "\n\nBest candidate found: ", best_candidate(ctrl))
    trace(ctrl, "\n\nFitness: ")
    show_fitness(stdout, best_fitness(ctrl), problem(ctrl))
    trace(ctrl, "\n\n")
end

function write_result(ctrl::OptRunController, filename = "")
    if isempty(filename)
        timestamp = Libc.strftime("%y%m%d_%H%M%S", floor(Int, ctrl.start_time))
        filename = "$(timestamp)_$(problem_summary(ctrl.evaluator))_$(name(ctrl.optimizer)).csv"
        filename = replace(replace(filename, r"\s+", "_"), r"/", "_")
    end
    save_fitness_history_to_csv_file(ctrl.evaluator.archive, filename;
            header_prefix = "Problem,Dimension,Optimizer",
            line_prefix = "$(name(problem(ctrl.evaluator))),$(numdims(ctrl.evaluator)),$(name(ctrl.optimizer))",
            bestfitness = opt_value(problem(ctrl.evaluator)))
end

"""
Optimization Controller.

Applies specific optimization method to a given problem.
Supports restarts and modifying parameter of the method between runs.
`runcontrollers` field maintains the list of `OptRunController` instances applied so far.

See `OptRunController`.
"""
mutable struct OptController{O<:Optimizer, P<:OptimizationProblem}
    optimizer::O   # optimization algorithm
    problem::P     # opt problem
    parameters::ParamsDictChain
    runcontrollers::Vector{OptRunController{O}}
end

"""
Create `OptController` for a given `optimizer` and a `problem`.

`params` are any of `OptRunController` parameters plus
    * `:RngSeed` and `:RandomizeRngSeed` params for controlling the random seed
    * `:RecoverResults` if intermediate results are returned upon `InterruptException()` (on by default)
"""
OptController(
    optimizer::O, problem::P,
    params::ParamsDictChain) where {O<:Optimizer, P<:OptimizationProblem} =
    OptController{O, P}(optimizer, problem, params, OptRunController{O}[])

problem(oc::OptController) = oc.problem
numruns(oc::OptController) = length(oc.runcontrollers)
lastrun(oc::OptController) = oc.runcontrollers[end]

"""
    update_parameters!(oc::OptController, parameters::Associative)

Update the `OptController` parameters.
"""
function update_parameters!(oc::OptController, parameters::Parameters = EMPTY_DICT)
    parameters = convert(ParamsDict, parameters)

    # Most parameters cannot be changed since the problem and optimizer has already
    # been setup.
    valid_params = Set([:MaxTime, :MaxSteps, :MaxFuncEvals, :TraceMode])
    for k in keys(parameters)
        if k ∉ valid_params
            throw(ArgumentError("It is currently not supported to change parameters that can affect the original opt problem or optimizer (here: $(k))"))
        end
    end

    # Add new params in front if any are specified and they are valid.
    if length(parameters) > 0
        oc.parameters = chain(oc.parameters, parameters)
        check_valid!(oc.parameters) # We must recheck that new param settings are valid
    end
end

function init_rng!(parameters::Parameters)
    if parameters[:RandomizeRngSeed]
        parameters[:RngSeed] = rand(1:1_000_000)
        Random.seed!(parameters[:RngSeed])
    end
end

"""
    run!(oc::OptController)

Start a new optimization run, possibly with new parameters and report on results.
"""
function run!(oc::OptController)
    # If this is not the first run we create the controller based on the previous one
    # to reuse the archive etc, otherwise create it based on the problem.
    ctrl = if numruns(oc) > 0
        OptRunController(lastrun(oc), oc.parameters)
    else
        # No previous run controller so create a new one from optimizer and problem.
        OptRunController(oc.optimizer, oc.problem, oc.parameters)
    end
    push!(oc.runcontrollers, ctrl)

    # If this is the first run we might have to init the RNG.
    if numruns(oc) == 1
        init_rng!(oc.parameters)
    end

    try
        run!(ctrl)
        if ctrl.trace_mode != :silent
            show_report(ctrl)
        end

        if oc.parameters[:SaveFitnessTraceToCsv]
            write_result(ctrl)
        end

        return OptimizationResults(ctrl, oc)
    catch ex
        # If it was a ctrl-c interrupt we try to make a result and return it...
        if get(oc.parameters, :RecoverResults, true) && isa(ex, InterruptException)
            @warn("Optimization interrupted, recovering intermediate results...")
            ctrl.stop_reason = @sprintf("%s", ex)
            return OptimizationResults(ctrl, oc)
        else
            rethrow(ex)
        end
    end
end


function ProperFractionFormat()
    return ProperFractionFormat(())
end

function ProperFractionFormat(arg0::NumberFormat)
    return ProperFractionFormat((NumberFormat,), arg0)
end

function ProperFractionFormat(arg0::NumberFormat, arg1::NumberFormat, arg2::NumberFormat)
    return ProperFractionFormat((NumberFormat, NumberFormat, NumberFormat), arg0, arg1, arg2)
end

function format(obj::ProperFractionFormat, arg0::Fraction, arg1::StringBuffer, arg2::FieldPosition)
    return jcall(obj, "format", StringBuffer, (Fraction, StringBuffer, FieldPosition), arg0, arg1, arg2)
end

function get_whole_format(obj::ProperFractionFormat)
    return jcall(obj, "getWholeFormat", NumberFormat, ())
end

function parse(obj::ProperFractionFormat, arg0::JString, arg1::ParsePosition)
    return jcall(obj, "parse", Fraction, (JString, ParsePosition), arg0, arg1)
end


# https://juliaci.github.io/PkgBenchmark.jl/stable/run_benchmarks/
# https://github.com/JuliaCI/PkgBenchmark.jl/blob/master/benchmark/benchmarks.jl
using BenchmarkTools
using Distributions
using HMMBase
using PyCall

include("helpers.jl")

# hmmlearn
hmmlearn = pyimport("hmmlearn.hmm")

# pyhsmm
pyhsmm_viterbi = pyimport("pyhsmm.internals.hmm_messages_interface").viterbi
HMMStatesPython = pyimport("pyhsmm.internals.hmm_states").HMMStatesPython

const SUITE = BenchmarkGroup()

function randhmm(K)
    A = randtransmat(K)
    B = [Normal(rand() * 100, rand() * 10) for _ = 1:K]
    HMM(A, B)
end

HMMs = Dict([K => randhmm(K) for K = 2:2:10])
Ys = Dict([K => rand(HMMs[K], 5000) for K in keys(HMMs)])

# HMMBase

for f in (forward, backward, viterbi), (K, hmm) in HMMs
    LL = loglikelihoods(hmm, Ys[K])
    mkbench!(SUITE, ["hmmbase", f, K]) do
        @benchmarkable ($f)($hmm.a, $hmm.A, $LL)
    end
end

# pyhsmm

fs = [
    HMMStatesPython._messages_forwards_normalized,
    HMMStatesPython._messages_backwards_normalized,
]

for f in fs, (K, hmm) in HMMs
    LL = loglikelihoods(hmm, Ys[K])
    mkbench!(SUITE, ["pyhsmm", f, K]) do
        @benchmarkable ($f)($hmm.A, $hmm.a, $LL)
    end
end

for (K, hmm) in HMMs
    LL = loglikelihoods(hmm, Ys[K])
    buf = zeros(Int32, size(LL, 1))

    a = hmm.a
    A = PyReverseDims(permutedims(hmm.A))
    LL = PyReverseDims(permutedims(LL))

    mkbench!(SUITE, ["pyhsmm", "viterbi", K]) do
        @benchmarkable pyhsmm_viterbi($A, $LL, $a, $buf)
    end
end

# hmmlearn

function tohmmlearn(hmm::AbstractHMM)
    K = size(hmm, 1)
    model = hmmlearn.GaussianHMM(n_components = K, covariance_type = :diag)
    model.startprob_ = hmm.a
    model.transmat_ = hmm.A # permutedims(...) ?
    model.means_ = [d.μ for d in hmm.B]
    model._covars_ = reshape([d.σ^2 for d in hmm.B], (1, 1, :))
    model
end

fs = [:_do_forward_pass, :_do_backward_pass, :_do_viterbi_pass]

for f in fs, (K, hmm) in HMMs
    model = tohmmlearn(hmm)
    LL = loglikelihoods(hmm, Ys[K])
    mkbench!(SUITE, ["hmmlearn", f, K]) do
        @benchmarkable $(model).$(f)($LL)
    end
end


export BehaviorCloningPolicy

mutable struct BehaviorCloningPolicy{A} <: AbstractPolicy
    approximator::A
    explorer::AbstractExplorer
    sampler::BatchSampler{(:state, :action)}
    min_reservoir_history::Int
end

"""
    BehaviorCloningPolicy(;kw...)

# Keyword Arguments

- `approximator`: calculate the logits of possible actions directly
- `explorer=GreedyExplorer()`
- `batch_size::Int = 32`
- `min_reservoir_history::Int = 100`, number of transitions that should be experienced before updating the `approximator`.
- `rng = Random.GLOBAL_RNG`
"""
function BehaviorCloningPolicy(;
    approximator::A,
    explorer::AbstractExplorer = GreedyExplorer(),
    batch_size::Int = 32,
    min_reservoir_history::Int = 100,
    rng = Random.GLOBAL_RNG,
) where {A}
    sampler = BatchSampler{(:state, :action)}(batch_size; rng = rng)
    BehaviorCloningPolicy(approximator, explorer, sampler, min_reservoir_history)
end

function (p::BehaviorCloningPolicy)(env::AbstractEnv)
    s = state(env)
    s_batch = Flux.unsqueeze(s, ndims(s) + 1)
    s_batch = send_to_device(device(p.approximator), s_batch)
    logits = p.approximator(s_batch) |> vec |> send_to_host # drop dimension
    typeof(ActionStyle(env)) == MinimalActionSet ? p.explorer(logits) :
    p.explorer(logits, legal_action_space_mask(env))
end

function RLBase.update!(p::BehaviorCloningPolicy, batch::NamedTuple{(:state, :action)})
    s = send_to_device(device(p.approximator), batch.state)
    a = send_to_device(device(p.approximator), batch.action)
    m = p.approximator
    gs = gradient(params(m)) do
        ŷ = m(s)
        y = Flux.OneHotMatrix(a, size(ŷ, 1))
        logitcrossentropy(ŷ, y)
    end
    update!(m, gs)
end

function RLBase.update!(p::BehaviorCloningPolicy, t::Any)
    (length(t) <= p.min_reservoir_history || length(t) <= p.sampler.batch_size) && return

    _, batch = p.sampler(t)
    update!(p, batch)
end

function RLBase.prob(p::BehaviorCloningPolicy, env::AbstractEnv)
    s = state(env)
    m = p.approximator
    s_batch = send_to_device(device(m), Flux.unsqueeze(s, ndims(s) + 1))
    values = m(s_batch) |> vec |> send_to_host
    typeof(ActionStyle(env)) == MinimalActionSet ? prob(p.explorer, values) :
    prob(p.explorer, values, legal_action_space_mask(env))
end

function RLBase.prob(p::BehaviorCloningPolicy, env::AbstractEnv, action)
    A = action_space(env)
    P = prob(p, env)
    @assert length(A) == length(P)
    if A isa Base.OneTo
        P[action]
    else
        for (a, p) in zip(A, P)
            if a == action
                return p
            end
        end
        @error "action[$action] is not found in action space[$(action_space(env))]"
    end
end


using VisualRegressionTests
using Plots
using Random
using BinaryProvider
using Test
using FileIO

include("imgcomp.jl")
# don't actually show the plots
Random.seed!(1234)
default(show=false, reuse=true)
img_tol = isinteractive() ? 1e-2 : 10e-2

@testset "GR" begin
    ENV["PLOTS_TEST"] = "true"
    ENV["GKSwstype"] = "100"
    @test gr() == Plots.GRBackend()
    @test backend() == Plots.GRBackend()

    @static if Sys.islinux()
        image_comparison_facts(:gr, tol=img_tol, skip = [25, 30])
    end
end

@testset "UnicodePlots" begin
    @test unicodeplots() == Plots.UnicodePlotsBackend()
    @test backend() == Plots.UnicodePlotsBackend()

    # lets just make sure it runs without error
    p = plot(rand(10))
    @test isa(p, Plots.Plot) == true
    @test isa(display(p), Nothing) == true
    p = bar(randn(10))
    @test isa(p, Plots.Plot) == true
    @test isa(display(p), Nothing) == true
end

@testset "Axes" begin
    p = plot()
    axis = p.subplots[1][:xaxis]
    @test typeof(axis) == Plots.Axis
    @test Plots.discrete_value!(axis, "HI") == (0.5, 1)
    @test Plots.discrete_value!(axis, :yo) == (1.5, 2)
    @test Plots.ignorenan_extrema(axis) == (0.5,1.5)
    @test axis[:discrete_map] == Dict{Any,Any}(:yo  => 2, "HI" => 1)

    Plots.discrete_value!(axis, ["x$i" for i=1:5])
    Plots.discrete_value!(axis, ["x$i" for i=0:2])
    @test Plots.ignorenan_extrema(axis) == (0.5, 7.5)
end

@testset "NoFail" begin
    histogram([1, 0, 0, 0, 0, 0])
end


using DFTShims
using Base.Test
using Unitful
using AxisArrays

const DD = DFTShims.Dispatch.Dimensions
const DH = DFTShims.Dispatch.Hartree
const UH = DFTShims.UnitfulHartree

@testset "Scalars" begin
    @test 1u"ρ" isa DD.Scalars.ρ
    @test 1u"nm^-3" isa DD.Scalars.ρ
    @test 1u"ρ" isa DH.Scalars.ρ
    @test !(1u"nm^-3" isa DH.Scalars.ρ)

    @test DH.Scalars.ρ <: DD.Scalars.ρ
    @test DH.Scalars.ρ{Float32} <: DD.Scalars.ρ
    @test DH.Scalars.ρ{Float32} <: DD.Scalars.ρ{Float32}
    @test DH.Scalars.ρ{Float32} === DD.Scalars.ρ{Float32, typeof(UH.ρ)}
    @test DH.Scalars.ρ{Float32} <: DD.Scalars.All{Float32}
    @test !(DH.Scalars.ρ{Int32} <: DD.Scalars.All{Float32})
    @test !(DH.Scalars.ρ <: DD.Scalars.All{Float32})

    f = (::DH.Scalars.ρ{<: Integer}) -> :hartree_with_integer
    (::typeof(f))(::DH.Scalars.ρ) = :hartree_with_any_number
    (::typeof(f))(::DD.Scalars.ρ{<: Integer}) = :density_with_integer
    (::typeof(f))(::DD.Scalars.ρ) = :density_with_any_number
    @test f(1u"ρ") == :hartree_with_integer
    @test f(1.0u"ρ") == :hartree_with_any_number
    @test f(1u"nm^-3") == :density_with_integer
    @test f(1.0u"nm^-3") == :density_with_any_number
end

@testset "Arrays" begin
    @test [1u"ρ", 2.0u"ρ"] isa DFTShims.DFTArray
    @test [1u"ρ", 2.0u"ρ"] isa DD.DensityArray
    @test [1u"ρ", 2.0u"ρ"] isa DD.Arrays.ρ
    @test [1u"ρ", 2.0u"ρ"] isa DD.Arrays.ρ{Float64}
    @test [1u"ρ", 2.0u"ρ"] isa DH.DensityDenseArray{Float64, 1}
    @test [1u"ρ", 2.0u"ρ"] isa DH.DenseArrays.ρ{Float64, 1}
    @test [1u"nm^-3", 2.0u"nm^-3"] isa DD.DenseArrays.ρ{Float64, 1}
    @test !([1u"nm^-3", 2.0u"nm^-3"] isa DH.DenseArrays.ρ{Float64, 1})

    @test DH.DensityDenseArray{Int8, 2} === DD.DensityDenseArray{Int8, 2, typeof(u"ρ")}
    @test DH.Arrays.ρ{Float32} <: DD.Arrays.All{Float32}
    @test !(DH.Arrays.ρ{Int32} <: DD.Arrays.All{Float32})
    @test !(DH.Arrays.ρ <: DD.Arrays.All{Float32})
    @test DH.Arrays.ρ{Float32, 2} <: DD.Arrays.All{Float32, 2}
    @test [1u"nm^-3", 2.0u"nm^-3"] isa DD.Arrays.ρ{Float64, 1}
    @test !([1u"nm^-3", 2.0u"nm^-3"] isa DD.Arrays.ρ{Float64, 2})

    f = (::DH.Arrays.ρ{<: Integer}) -> :hartree_with_integer
    (::typeof(f))(::DH.Arrays.ρ) = :hartree_with_any_number
    (::typeof(f))(::DD.Arrays.ρ{<: Integer}) = :density_with_integer
    (::typeof(f))(::DD.Arrays.ρ) = :density_with_any_number
    (::typeof(f))(::DD.AxisArrays.ρ) = :axis_density_with_any_number
    @test f([1u"ρ"]) == :hartree_with_integer
    @test f([1.0u"ρ"]) == :hartree_with_any_number
    @test f([1u"nm^-3"]) == :density_with_integer
    @test f([1.0u"nm^-3"]) == :density_with_any_number
    @test f(AxisArray([1.0u"nm^-3"])) == :axis_density_with_any_number
end


# https://github.com/JuliaPlots/UnicodePlots.jl

# don't warn on unsupported... there's just too many warnings!!
warn_on_unsupported_args(::UnicodePlotsBackend, plotattributes::KW) = nothing

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

_canvas_map() = (
    ascii = UnicodePlots.AsciiCanvas,
    block = UnicodePlots.BlockCanvas,
    braille = UnicodePlots.BrailleCanvas,
    density = UnicodePlots.DensityCanvas,
    dot = UnicodePlots.DotCanvas,
    heatmap = UnicodePlots.HeatmapCanvas,
    lookup = UnicodePlots.LookupCanvas,
)

# do all the magic here... build it all at once, since we need to know about all the series at the very beginning
function unicodeplots_rebuild(plt::Plot{UnicodePlotsBackend})
    plt.o = UnicodePlots.Plot[]
    canvas_map = _canvas_map()
    for sp in plt.subplots
        xaxis = sp[:xaxis]
        yaxis = sp[:yaxis]
        xlim = collect(axis_limits(sp, :x))
        ylim = collect(axis_limits(sp, :y))

        # we set x/y to have a single point, since we need to create the plot with some data.
        # since this point is at the bottom left corner of the plot, it shouldn't actually be shown
        x = Float64[xlim[1]]
        y = Float64[ylim[1]]

        # create a plot window with xlim/ylim set, but the X/Y vectors are outside the bounds
        canvas_type = if (ct = _canvas_type[]) == :auto
            isijulia() ? :ascii : :braille
        else
            ct
        end

        kw = (
            title = sp[:title],
            xlim = xlim,
            ylim = ylim,
            border = isijulia() ? :ascii : :solid,
            xlabel = xaxis[:guide],
            ylabel = yaxis[:guide],
            xscale = xaxis[:scale],
            yscale = yaxis[:scale],
        )

        o = UnicodePlots.Plot(x, y, canvas_map[canvas_type]; kw...)
        for series in series_list(sp)
            o = addUnicodeSeries!(sp, o, kw, series, sp[:legend] != :none)
        end

        for ann in sp[:annotations]
            x, y, val = locate_annotation(sp, ann...)
            o = UnicodePlots.annotate!(
                o, x, y, val.str;
                color = up_color(val.font.color), halign = val.font.halign, valign = val.font.valign
            )
        end

        push!(plt.o, o)  # save the object
    end
end

function up_color(col)
    if typeof(col) <: UnicodePlots.UserColorType
        color = col
    elseif typeof(col) <: RGBA
        col = convert(ARGB32, col)
        color = map(Int, (red(col).i, green(col).i, blue(col).i))
    else
        color = :auto
    end
    color
end

# add a single series
function addUnicodeSeries!(
    sp::Subplot{UnicodePlotsBackend},
    up::UnicodePlots.Plot,
    kw, series, addlegend::Bool,
)
    st = series[:seriestype]

    # get the series data and label
    x, y = if st == :straightline
        straightline_data(series)
    elseif st == :shape
        shape_data(series)
    else
        float(series[:x]), float(series[:y])
    end

    # special handling (src/interface)
    if st == :histogram2d
        kw[:xlim][:] .= kw[:ylim][:] .= 0
        return UnicodePlots.densityplot(x, y; kw...)
    elseif st == :heatmap
        rng = range(0, 1, length = length(UnicodePlots.COLOR_MAP_DATA[:viridis]))
        cmap = [(red(c), green(c), blue(c)) for c in get(get_colorgradient(series), rng)]
        return UnicodePlots.heatmap(
            series[:z].surf;
            zlabel = sp[:colorbar_title],
            colormap = cmap,
            kw...
        )
    elseif st == :spy
        return UnicodePlots.spy(series[:z].surf; kw...)
    end

    # now use the ! functions to add to the plot
    if st in (:path, :straightline, :shape)
        func = UnicodePlots.lineplot!
    elseif st == :scatter || series[:markershape] != :none
        func = UnicodePlots.scatterplot!
    else
        error("Series type $st not supported by UnicodePlots")
    end

    label = addlegend ? series[:label] : ""

    for (n, segment) in enumerate(series_segments(series, st; check = true))
        i, rng = segment.attr_index, segment.range
        lc = get_linecolor(series, i)
        up = func(up, x[rng], y[rng]; color = up_color(lc), name = n == 1 ? label : "")
    end

    for (xi, yi, str, fnt) in EachAnn(series[:series_annotations], x, y)
        up = UnicodePlots.annotate!(
            up, xi, yi, str;
            color = up_color(fnt.color), halign = fnt.halign, valign = fnt.valign
        )
    end

    return up
end

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

# since this is such a hack, it's only callable using `png`... should error during normal `show`
function png(plt::Plot{UnicodePlotsBackend}, fn::AbstractString)
    fn = addExtension(fn, "png")

    @static if Sys.isapple()
        # make some whitespace and show the plot
        println("\n\n\n\n\n\n")
        gui(plt)

        # BEGIN HACK

        # wait while the plot gets drawn
        sleep(0.5)

        # use osx screen capture when my terminal is maximized and cursor starts at the bottom (I know, right?)
        # TODO: compute size of plot to adjust these numbers (or maybe implement something good??)
        run(`screencapture -R50,600,700,420 $fn`)

        # END HACK (phew)
        return
    elseif Sys.islinux()
        run(`clear`)
        gui(plt)
        run(`import -window $(ENV["WINDOWID"]) $fn`)
        return
    end

    error(
        "Can only savepng on MacOS or Linux with UnicodePlots (though even then I wouldn't do it)",
    )
end

# -------------------------------
Base.show(plt::Plot{UnicodePlotsBackend}) = _show(stdout, MIME("text/plain"), plt)

function _show(io::IO, ::MIME"text/plain", plt::Plot{UnicodePlotsBackend})
    unicodeplots_rebuild(plt)
    nr, nc = size(plt.layout)
    lines_colored = Array{Union{Nothing,Vector{String}}}(undef, nr, nc)
    lines_uncolored = copy(lines_colored)
    l_max = zeros(Int, nr)
    buf = IOBuffer()
    cbuf = IOContext(buf, :color => true)
    sps = wmax = 0
    for r in 1:nr
        lmax = 0
        for c in 1:nc
            l = plt.layout[r, c]
            if l isa GridLayout && size(l) != (1, 1)
                @error "UnicodePlots: complex nested layout is currently unsupported !"
            else
                if get(l.attr, :blank, false)
                    lines_colored[r, c] = lines_uncolored[r, c] = nothing
                else
                    sp = plt.o[sps += 1]
                    show(cbuf, sp)
                    colored = String(take!(buf))
                    uncolored = replace(colored, r"\x1B\[[0-9;]*[a-zA-Z]" => "")
                    lines_colored[r, c] = lc = split(colored, "\n")
                    lines_uncolored[r, c] = lu = split(uncolored, "\n")
                    lmax = max(length(lc), lmax)
                    wmax = max(maximum(length.(lu)), wmax)
                end
            end
        end
        l_max[r] = lmax
    end
    empty = ' '^wmax
    for r in 1:nr
        for n in 1:l_max[r]
            for c in 1:nc
                pre = c == 1 ? '\0' : ' '
                lc = lines_colored[r, c]
                if lc === nothing || length(lc) < n
                    print(io, pre, empty)
                else
                    lu = lines_uncolored[r, c]
                    print(io, pre, lc[n], ' '^(wmax - length(lu[n])))
                end
            end
            println(io)
        end
        r < nr && println(io)
    end
    nothing
end

function _display(plt::Plot{UnicodePlotsBackend})
    unicodeplots_rebuild(plt)
    map(display, plt.o)
    nothing
end


module MLJTime

# IMPORTS
using IndexedTables, CSVFiles, ZipFile
using  MLJBase: fit!, predict, machine, partition, l1, CV, range, cross_entropy,
       predict_mode, accuracy, predict_mean
import MLJModelInterface
import MLJModelInterface: @mlj_model, fit, predict, matrix
using StableRNGs

# EXPORTS
export RandomForestClassifierFit, InvFeatureGen, predict_new, predict_mean,
       fit!, predict, machine, fitted_params, predict_mode, accuracy

export univariate_datasets, load_dataset, load_from_tsfile_to_NDArray, dwt,
       load_ts_file, ts_dataset, partition, matrix, dtw_distance, load_NDdataset

export l1, StableRNG, CV, range, cross_entropy
#CONSTANTS
# the directory containing this file: (.../src/)
const MODULE_DIR = dirname(@__FILE__)
const MMI = MLJModelInterface

# Includes
include("interval_based_forest.jl")
include("datapath.jl")
include("interface.jl")
include("distances_based_classifier.jl")
include("rocket.jl")

end # module


module lmmsim

using DataFrames, DataStructures, MixedModels
export
    dat,
    makedata,
    makeθ,
    mods,
    simulate

include("makedata.jl")
include("simulate.jl")

end # module


mutable struct scratch_data
    n_particles::Int
    polygon::Vector{Point2D{Float64}}
    extremePoint::Point2D{Float64}
    pointOnWall::Point2D{Float64}
    w::Vector{Float64}
    uEl::Vector{Float64}
    vEl::Vector{Float64}
    a::Vector{Float64}
    b::Vector{Float64}
    c::Vector{Float64}
    d::Vector{Float64}
    cfX::Vector{Float64}
    cfY::Vector{Float64}
    afX::Vector{Float64}
    afY::Vector{Float64}
    gfX::Vector{Float64}
    gfY::Vector{Float64}
    sfX::Vector{Float64}
    sfY::Vector{Float64}

    function scratch_data(n_particles::Int)
        polygon      = [Point2D(0.0,0.0),Point2D(0.0,0.0),Point2D(0.0,0.0),Point2D(0.0,0.0)]
        extremePoint = Point2D(100_000.0,0.0)
        pointOnWall  = Point2D(0.0,0.0)
        w   = zeros(4)
        uEl = zeros(4)
        vEl = zeros(4)
        a   = zeros(3)
        b   = zeros(3)
        c   = zeros(3)
        d   = zeros(3)
        cfX = zeros(n_particles)
        cfY = zeros(n_particles)
        afX = zeros(n_particles)
        afY = zeros(n_particles)
        gfX = zeros(n_particles)
        gfY = zeros(n_particles)
        sfX = zeros(n_particles)
        sfY = zeros(n_particles)
        return new(n_particles,polygon,extremePoint,pointOnWall,
                    w,uEl,vEl,a,b,c,d,
                    cfX,cfY,afX,afY,gfX,gfY,sfX,sfY)
    end
end


using PSSFSS
using PSSFSS.Modes: choose_layer_modes!
#using PSSFSS.Elements: rectstrip, polyring
using PSSFSS.Layers: Layer
using PSSFSS.Sheets: RWGSheet
#using PSSFSS.PSSFSSLen: mm, cm, inch, mil
using PSSFSS.GSMs: Gblock, choose_gblocks
using PSSFSS.Constants: c₀, twopi
using Test
using Logging: Error, ConsoleLogger, default_metafmt, global_logger

testlogger = ConsoleLogger(stderr, Error,
                       meta_formatter=default_metafmt, show_limited=true,
                       right_justify=0)
oldlogger = global_logger(testlogger)

@testset "choose_layer_modes!1" begin
    strata =
    [
        Layer(name="Vacuum")
        Layer(name="Duroid 6035", ϵᵣ=3.6, tanδ=0.0013, width=20mil)
        polyring(sides=4, orient=45, a=[0.1902], b=[0.2169], units=inch,
                    s1=[0.5117, 0], s2=[0, 0.5117], ntri=100)
        Layer(name="Duroid 6035", ϵᵣ=3.6, tanδ=0.0013, width=20mil)
        Layer(name="Vacuum", width=0.2inch)
        Layer(name="Duroid 6035", ϵᵣ=3.6, tanδ=0.0013, width=20mil)
        polyring(sides=4, orient=45, a=[0.1902], b=[0.2169], units=inch,
                    s1=[0.5117, 0], s2=[0, 0.5117], ntri=100)
        Layer(name="Duroid 6035", ϵᵣ=3.6, tanδ=0.0013, width=20mil)
        Layer(name="Vacuum")
    ]

    islayer = map(x -> x isa Layer, strata)
    issheet = map(x -> x isa RWGSheet, strata)
    layers = convert(Vector{Layer}, strata[islayer])
    sheets = convert(Vector{RWGSheet}, strata[issheet])
    nl = length(layers)
    nj = nl - 1
    ns = length(sheets)
    sint = cumsum(islayer)[issheet] # sint[k] contains dielectric interface number of k'th sheet
    junc = zeros(Int, nj)
    junc[sint] = 1:ns #  junc[i] is the sheet number present at interface i, or 0 if no sheet is there
    fmin = 2.0 * 1e9  # minimum frequency [Hz]
    k0min = twopi * fmin / c₀
    fmax = 8.0 * 1e9  # minimum frequency [Hz]
    k0max = twopi * fmax / c₀
    dbmin = 25.0

    gbl = choose_gblocks(layers, sheets, junc, k0min)
    choose_layer_modes!(layers, sheets, junc, gbl, k0max, dbmin)
    @test [length(l.P) for l in strata if l isa Layer] == [2,0,0,18,0,0,2]
end

@testset "choose_layer_modes!2" begin
    strata =
    [
        Layer(name="Vacuum")
        Layer(name="Duroid 6035", ϵᵣ=3.6, tanδ=0.0013, width=20mil)
        polyring(sides=4, orient=45, a=[0.1902], b=[0.2169], units=inch,
                    s1=[0.5117, 0], s2=[0, 0.5117], ntri=100)
        Layer(name="Duroid 6035", ϵᵣ=3.6, tanδ=0.0013, width=20mil)
        Layer(name="Vacuum", width=0.1inch)
        Layer(name="Vacuum", width=0.1inch)
        Layer(name="Duroid 6035", ϵᵣ=3.6, tanδ=0.0013, width=20mil)
        polyring(sides=4, orient=45, a=[0.1902], b=[0.2169], units=inch,
                    s1=[0.9005, 0], s2=[0, 0.9005], ntri=100)
        Layer(name="Duroid 6035", ϵᵣ=3.6, tanδ=0.0013, width=20mil)
        Layer(name="Vacuum")
    ]

    islayer = map(x -> x isa Layer, strata)
    issheet = map(x -> x isa RWGSheet, strata)
    layers = convert(Vector{Layer}, strata[islayer])
    sheets = convert(Vector{RWGSheet}, strata[issheet])
    nl = length(layers)
    nj = nl - 1
    ns = length(sheets)
    sint = cumsum(islayer)[issheet] # sint[k] contains dielectric interface number of k'th sheet
    junc = zeros(Int, nj)
    junc[sint] = 1:ns #  junc[i] is the sheet number present at interface i, or 0 if no sheet is there
    fmin = 2.0 * 1e9  # minimum frequency [Hz]
    k0min = twopi * fmin / c₀
    fmax = 8.0 * 1e9  # minimum frequency [Hz]
    k0max = twopi * fmax / c₀
    dbmin = 25.0

    gbl = choose_gblocks(layers, sheets, junc, k0min)
    choose_layer_modes!(layers, sheets, junc, gbl, k0max, dbmin)
    @test [length(l.P) for l in strata if l isa Layer] == [2,0,0,2,0,0,0,2]
end

@testset "choose_layer_modes!3" begin
    strata =
    [
        Layer(name="Vacuum")
        Layer(name="Duroid 6035", ϵᵣ=3.6, tanδ=0.0013, width=20mil)
        polyring(sides=4, orient=45, a=[0.1902], b=[0.2169], units=inch,
                    s1=[0.5117, 0], s2=[0, 0.5117], ntri=100)
        Layer(name="Duroid 6035", ϵᵣ=3.6, tanδ=0.0013, width=20mil)
        Layer(name="Vacuum", width=0.1inch)
        Layer(name="Vacuum", width=0.1inch)
        Layer(name="Duroid 6035", ϵᵣ=3.6, tanδ=0.0013, width=20mil)
        polyring(sides=4, orient=45, a=[0.1902], b=[0.2169], units=inch,
                    s1=[0.5117, 0], s2=[0, 0.5117], ntri=100)
        Layer(name="Duroid 6035", ϵᵣ=3.6, tanδ=0.0013, width=20mil)
        Layer(name="Vacuum")
    ]

    fmin = 2.0 * 1e9  # minimum frequency [Hz]
    k0min = twopi * fmin / c₀
    fmax = 8.0 * 1e9  # minimum frequency [Hz]
    k0max = twopi * fmax / c₀
    dbmin = 25.0
    islayer = map(x -> x isa Layer, strata)
    issheet = map(x -> x isa RWGSheet, strata)
    layers = convert(Vector{Layer}, strata[islayer])
    sheets = convert(Vector{RWGSheet}, strata[issheet])
    nl = length(layers)
    nj = nl - 1
    ns = length(sheets)
    sint = cumsum(islayer)[issheet] # sint[k] contains dielectric interface number of k'th sheet
    junc = zeros(Int, nj)
    junc[sint] = 1:ns #  junc[i] is the sheet number present at interface i, or 0 if no sheet is there

    gbl = choose_gblocks(layers, sheets, junc, k0min)
    choose_layer_modes!(layers, sheets, junc, gbl, k0max, dbmin)
    #@test_logs [(:info,),(:info,),(:info,)] choose_layer_modes!(strata, gbl, k0max, dbmin)
    @test [length(l.P) for l in strata if l isa Layer] == [2,0,0,42,0,0,0,2]

end

global_logger(oldlogger)
nothing


"
    MPSComoving(leftstate,window,rightstate)

    muteable window of tensors on top of an infinite chain
"
mutable struct MPSComoving{A<:GenericMPSTensor,B<:MPSBondTensor} <: AbstractMPS
    left_gs::InfiniteMPS{A,B}
    window::FiniteMPS{A,B}
    right_gs::InfiniteMPS{A,B}

    function MPSComoving(left::InfiniteMPS{A,B},window::FiniteMPS{A,B},right::InfiniteMPS{A,B}) where {A<:GenericMPSTensor,B<:MPSBondTensor}
        #todo:insert checks
        return new{A,B}(left,window,right);
    end
end

#allow construction with an array of tensors
MPSComoving(left::InfiniteMPS,site_tensors::Array,right::InfiniteMPS) = MPSComoving(left,FiniteMPS(site_tensors),right)

# allow construction with one large tensorkit space
MPSComoving(f, elt,P::ProductSpace, args...; kwargs...) = MPSComoving(f, elt, collect(P), args...; kwargs...)

# allow construction given only a physical space and length
MPSComoving(f,elt, N::Int, V::VectorSpace, args...; kwargs...) = MPSComoving(f, elt,fill(V, N), args...; kwargs...)

function MPSComoving(f, elt, physspaces::Vector{<:Union{S,CompositeSpace{S}}}, maxvirtspace::S,
                    leftgs::M,
                    rightgs::M) where {S<:ElementarySpace,M<:InfiniteMPS}

    left = virtualspace(leftgs,0);
    right = virtualspace(rightgs,length(physspaces));
    window = FiniteMPS(f,elt,physspaces,maxvirtspace,left=left,right=right);
    MPSComoving(leftgs,window,rightgs)
end
function MPSComoving(f,elt,
                    physspaces::Vector{<:Union{S,CompositeSpace{S}}},
                    virtspaces::Vector{S},
                    leftgs::M,
                    rightgs::M) where {S<:ElementarySpace,M<:InfiniteMPS}
    N = length(physspaces)
    length(virtspaces) == N+1 || throw(DimensionMismatch())

    tensors = [TensorMap(f, elt,virtspaces[n] ⊗ physspaces[n], virtspaces[n+1]) for n=1:N]

    return MPSComoving(leftgs,FiniteMPS(tensors,overwrite=true),rightgs)
end

#take a window from an infinite mps, and use that as state
function MPSComoving(state::InfiniteMPS{A,B},len::Int) where {A,B}
    CLs = Vector{Union{Missing,B}}(missing,len+1)
    ALs = Vector{Union{Missing,A}}(missing,len)
    ARs = Vector{Union{Missing,A}}(missing,len)
    ACs = Vector{Union{Missing,A}}(missing,len)

    ALs.= state.AL[1:len];
    ARs.= state.AR[1:len];
    ACs.= state.AC[1:len];
    CLs.= state.CR[0:len];

    MPSComoving(state,FiniteMPS(ALs,ARs,ACs,CLs),state)
end

Base.copy(state::MPSComoving{A,B}) where {A,B} = MPSComoving(state.left_gs,copy(state.window),state.right_gs);

Base.length(state::MPSComoving) = length(state.window)
Base.size(psi::MPSComoving, i...) = size(psi.window, i...)

Base.eltype(::Type{MPSComoving{Mtype,Vtype}}) where {Mtype<:GenericMPSTensor,Vtype<:MPSBondTensor} = Mtype
site_type(::Type{MPSComoving{Mtype,Vtype}}) where {Mtype<:GenericMPSTensor,Vtype<:MPSBondTensor} = Mtype
bond_type(::Type{MPSComoving{Mtype,Vtype}}) where {Mtype<:GenericMPSTensor,Vtype<:MPSBondTensor} = Vtype

TensorKit.space(psi::MPSComoving{<:MPSTensor}, n::Integer) = space(psi.AC[n], 2)
virtualspace(psi::MPSComoving, n::Integer) =
    n < length(psi) ? _firstspace(psi.AC[n+1]) : dual(_lastspace(psi.AC[n]))

r_RR(state::MPSComoving)=r_RR(state.right_gs,length(state))
l_LL(state::MPSComoving)=l_LL(state.left_gs,1)

function Base.getproperty(psi::MPSComoving,prop::Symbol)
    if prop == :AL
        return ALView(psi)
    elseif prop == :AR
        return ARView(psi)
    elseif prop == :AC
        return ACView(psi)
    elseif prop == :CR
        return CRView(psi)
    else
        return getfield(psi,prop)
    end
end

max_Ds(f::MPSComoving) = max_Ds(f.window)

Base.:*(psi::MPSComoving, a::Number) = rmul!(copy(psi),a)
Base.:*(a::Number, psi::MPSComoving) = lmul!(a,copy(psi))

function TensorKit.lmul!(a::Number, psi::MPSComoving)
    psi.window.ACs .*=a;
    psi.window.CLs .*=a;
    return psi
end

function TensorKit.rmul!(psi::MPSComoving,a::Number)
    psi.window.ACs .*=a;
    psi.window.CLs .*=a;
    return psi
end

function TensorKit.dot(psi1::MPSComoving, psi2::MPSComoving)
    length(psi1) == length(psi2) || throw(ArgumentError("MPS with different length"))
    psi1.left_gs == psi2.left_gs || throw(ArgumentError("left InfiniteMPS is different"))
    psi1.right_gs == psi2.right_gs || throw(ArgumentError("right InfiniteMPS is different"))

    ρr = transfer_right(r_RR(psi2),psi2.AR[2:end],psi1.AR[2:end]);
    return tr(_transpose_front(psi1.AC[1])' * _transpose_front(psi2.AC[1]) * ρr)
end

TensorKit.norm(psi::MPSComoving) = norm(psi.AC[1])

TensorKit.normalize!(psi::MPSComoving) = rmul!(psi, 1/norm(psi))
TensorKit.normalize(psi::MPSComoving) = normalize!(copy(psi))


# Juliaを使った株価のアラート通知
using JSON, Parameters, TimeZones, Logging, CSV, DataFrames, BusinessDays, Dates
# cd(raw"C:\workspace\StockPrice.jl") # for debug
include("SendGmail.jl")
include("GetStockPrice.jl")

# -----------------------------------------
# アラート設定
# -----------------------------------------
@with_kw mutable struct AlertSetting
    stock::String   # 株価ティッカー
    func::String    # アラート関数
    value::Array{Float64,1}  # アラートのしきい値とする値
    enable::Bool  # アラートの有効化フラグ
    alerted::Bool # アラートがその日に上がったことを示すフラグ
end

# 設定ファイルの作成
function initalertsetting()
    settings = Dict(
        "stock" => "SPXL",
        "func" => "checklowerlimit",
        "value" => "120.0",
        "enable" => true,
        "alerted" => false
    )
    json = JSON.json(settings)
    open("./setting/SPXL_alert_setting.json", "w") do f
        print(f, json)
    end
end

## 設定をDictで読み取りstructにマッピングする-----------------
# 参考：https://matsueushi.github.io/posts/julia-mapping-nested-dict/
# キーを文字列からシンボルに変換する関数
keytosymbol(x) = Dict(Symbol(k) => v for (k, v) in pairs(x))
# アラート設定を記載したJSONファイルを読み込みstructのインスタンスを返す
function readalertsetting(jsonfile)
    str = open(f -> read(f, String), jsonfile)
    json = JSON.parse(str)
    symbol = keytosymbol(json)
    symbol[:value] = parse.(Float64, split(symbol[:value], ","))
    setting = AlertSetting(;symbol...)
    return setting
end

# 指定したdirからkeyを含むファイル名からsettingを読み込みリスト化する．
# dir: 設定読み込み先ディレクトリ名
# key: 読み込む設定に該当するキー
function readalertsettings(;key=".json", dir="./")
    settings_json = filter(f -> contains(f, key), readdir(dir))
    settings = readalertsetting.(dir .* settings_json)
    return settings
end

# -----------------------------------------
# アラート関数：動作検証済み
# -----------------------------------------
# 設定下限値を下回った場合にアラート
function checklowerlimit(alert::AlertSetting, prices)
    message = ""
    lowerlimit = alert.value[1]
    if ( sum(prices."close" .< lowerlimit) > 0 )
        alert.alerted = true
        meet = prices[prices."close" .< lowerlimit,:] # 条件に当てはまったデータ
        price = string(round(meet[1,"close"], digits=1))
        dt = meet[1, "timestamp"]
        dtfmt = "YYYY/mm/dd HH:MM:SS"
        message = alert.stock * "の株価が下限値設定を下回りました．\r\n" *
            "下限値: " * string(lowerlimit) * "\r\n" *
            "時刻: " * Dates.format(dt, dtfmt) * "\r\n" *
            "株価: " * price * "\r\n"
    end
    return message
end

# 設定上限値を上回った場合にアラート
function checkupperlimit(alert::AlertSetting, prices)
    message = ""
    upperlimit = alert.value[1]
    if ( sum(prices."close" .> upperlimit) > 0 )
        alert.alerted = true
        meet = prices[prices."close" .> upperlimit, :] # 条件に当てはまったデータ
        price = string(round(meet[1, "close"], digits=1))
        dt = meet[1, "timestamp"]
        dtfmt = "YYYY/mm/dd HH:MM:SS"
        message = alert.stock * "の株価が上限を超過しました．\r\n" *
            "上限値: " * string(upperlimit) * "\r\n" *
            "時刻: " * Dates.format(dt, dtfmt) * "\r\n" *
            "株価: " * price * "\r\n"
    end
    return message
end

# -----------------------------------------
# アラート実行
# -----------------------------------------
# 指定されたalertsリストのアラートをチェックして，条件を満たせばアラートメールを発信する
# alerts Array{AlertSetting, 1}
function exealert(alerts)

    for alert in alerts
        # 開場時間判定
        if( contains(alert.stock, ".T") )
            now_tk = now() # 日本時間
            open_tk = (Time(now_tk) >= Time(9, 00)) && (Time(now_tk) <= Time(15, 0)) # 東京証券取引所
            market_open = open_tk && isbday("JPTSE", Date(now_tk))
        else
            now_ny = now(tz"America/New_York") # 米国東部時間
            open_ny = (Time(now_ny) >= Time(9, 30)) && (Time(now_ny) <= Time(16, 0)) # NY証券取引所
            market_open =  open_ny && isbday("USNYSE", Date(now_ny))
        end

        # 開場時にアラート有効かつアラート未発報
        if market_open && alert.enable && !alert.alerted
            info("株価を確認します: " * alert.stock)
            prices = getstockprice(alert.stock)
            info("price: " * string(prices[end,:].close))
            message = eval(:( $(Symbol(alert.func))($alert, $prices) )) # 文字列で与えた関数の評価
            if alert.alerted
                title = "【株価アラート通知】" * alert.stock
                info("アラートメールを発信します．\n" * message)
                resp = sendmail(title, message) # アラートメールの発信
                info("アラートメールを発信しました．\n"*string(resp))
            end
        end
    end

end

# 東京証券取引所の祝日を追加定義
struct JPTSE <: BusinessDays.HolidayCalendar end
function BusinessDays.isholiday(::JPTSE, dt::Dates.Date)
    # [東証の祝日](https://www.jpx.co.jp/corporate/about-jpx/calendar/index.html)
    # 上記は取得しにくいので内閣府からCSVをダウンロード．ただし 1/2, 1/3, 12/31が不足しているので追記必須．年1回程度確認する
    # download("https://www8.cao.go.jp/chosei/shukujitsu/syukujitsu.csv", filename)
    filename = "./data/holidays_JPTSE.csv"
    dfmt = dateformat"yyyy/mm/dd"
    df = DataFrame(CSV.File(filename, dateformat=dfmt))
    return sum(df[:,1] .== Date(dt)) > 0
end

# settingリストからアラートを1つずつ3時まで定期実行する
# それぞれのアラートを個別チェックし，最大日に1回ずつアラートするようにする
function managealerts()
    # アラート定義の取得
    global alerts = readalertsettings(key="alert_setting.json", dir="./setting/")
    [info(alert) for alert in alerts] # アラートの内容出力

    # アラートを9時に起動して翌6時まで10分おき想定で実行
    int = 60 * 10 # アラートチェック間隔=10分の秒数
    period = 21 * 3600 # 朝9時～翌6時の秒数
    info("株価を" * string(int / 60) * "分間隔で" * string(period / 60 / 60) * "時間監視します．")
    cb(timer) = (exealert(alerts))
    t = Timer(cb, 0, interval=int)
    wait(t)
    sleep(period)
    close(t)
end

# ConsoleLoggerに時刻出力を追加
function timed_metafmt(level, _module, group, id, file, line)
    color, prefix, suffix = Logging.default_metafmt(level, _module, group, id, file, line)
    timestamp = Dates.now()
    prefix2 = "$timestamp : $prefix"
    return color, prefix2, suffix
end

# flushを付加
function info(x)
    @info x
    flush(io)
end
function warn(x)
    @warn x
    flush(io)
end

function main()
    # Loggingの設定: log.txtに追加書き込みでInfo以上のログを書き出し．debug, info, warn, errorがあり．
    global io = open("log.txt", "a+")
    logger = Logging.ConsoleLogger(io, Logging.Info, meta_formatter=timed_metafmt)
    Logging.global_logger(logger)
    info(logger)

    managealerts()

    close(io) # close log file
end

main()


facts("Exchange responses for backtesting") do
  context("Software-implemented limit order") do
    plim = 150.0
    buyord = TradingLogic.Order("buysome", 10, plim,
                                :buy, :limit, :pending, "aux")
    sellord = TradingLogic.Order("sellsome", 10, plim,
                                 :sell, :limit, :pending, "aux")
    @fact TradingLogic.plimitcheck(buyord, 175.0) --> false
    @fact TradingLogic.query_orderstatus(buyord, 175.0) --> :pending
    @fact TradingLogic.plimitcheck(buyord, 125.0) --> true
    @fact TradingLogic.query_orderstatus(buyord, 125.0) --> :complete
    @fact TradingLogic.plimitcheck(sellord, 175.0) --> true
    @fact TradingLogic.query_orderstatus(sellord, 175.0) --> :complete
    @fact TradingLogic.plimitcheck(sellord, 125.0) --> false
    @fact TradingLogic.query_orderstatus(sellord, 125.0) --> :pending
  end
  context("Market order and order submission") do
    ordid = "markord"
    mord = TradingLogic.Order(ordid, 10, NaN,
                              :buy, :market, :pending, "aux")
    # buy no matter how high
    @fact TradingLogic.query_orderstatus(mord, 1e7) --> :complete
    # simulated order submission
    @fact TradingLogic.submit_ordernew(mord, false) --> "FAIL"
    @fact TradingLogic.submit_ordernew(mord, true) --> ordid
  end
  context("Cancel pending limit-order") do
    limord = TradingLogic.Order("buysome", 10, 100.0,
                                :buy, :limit, :pending, "aux")
    @fact TradingLogic.submit_ordercancel(limord) --> true
    # can not cancel non-pending order
    @fact TradingLogic.submit_ordercancel(
      TradingLogic.emptyorder()) --> false
  end
end


Pkg.add("DataFrames")
Pkg.add("Gadfly")
Pkg.add("Cairo")
Pkg.add("Fontconfig")
using DataFrames
using Gadfly
using Cairo
using Fontconfig

# Plot all validities

style = Theme(
    major_label_font_size=35pt,
    minor_label_font_size=30pt,
    background_color="white",
    major_label_font="CMU Serif",
    minor_label_font="CMU Serif",
    bar_spacing=150px
)

numbers = ["377", "11177", "127413", "140054", "7506"];

for number in numbers
    println("Creating " , number , "...")
    table = readtable(string("output/" , number , "_answers.csv"));
    output = plot(
        table,
        x=:answer,
        y=:count,
        Geom.bar(position=:dodge),
        Coord.Cartesian(),
        Guide.xlabel("Answer"),
        Guide.ylabel("# of answers"),
        style
    );
    draw(PDF(string("graphs/" , number , "-answers.pdf"), 1600px, 800px), output);
end


module MolecularPotentials
    using StaticArrays
    using Parameters
    include("AbstractMolecule.jl")
    include("AbstractPotential.jl")
    include("/home/heindelj/dev/julia_development/wally/src/molecule_tools/read_xyz.jl")
end


function UniformRealDistribution()
    return UniformRealDistribution(())
end

function UniformRealDistribution(arg0::jdouble, arg1::jdouble)
    return UniformRealDistribution((jdouble, jdouble), arg0, arg1)
end

function cumulative_probability(obj::UniformRealDistribution, arg0::jdouble)
    return jcall(obj, "cumulativeProbability", jdouble, (jdouble,), arg0)
end

function density(obj::UniformRealDistribution, arg0::jdouble)
    return jcall(obj, "density", jdouble, (jdouble,), arg0)
end

function get_numerical_mean(obj::UniformRealDistribution)
    return jcall(obj, "getNumericalMean", jdouble, ())
end

function get_numerical_variance(obj::UniformRealDistribution)
    return jcall(obj, "getNumericalVariance", jdouble, ())
end

function get_support_lower_bound(obj::UniformRealDistribution)
    return jcall(obj, "getSupportLowerBound", jdouble, ())
end

function get_support_upper_bound(obj::UniformRealDistribution)
    return jcall(obj, "getSupportUpperBound", jdouble, ())
end

function inverse_cumulative_probability(obj::UniformRealDistribution, arg0::jdouble)
    return jcall(obj, "inverseCumulativeProbability", jdouble, (jdouble,), arg0)
end

function is_support_connected(obj::UniformRealDistribution)
    return jcall(obj, "isSupportConnected", jboolean, ())
end


function check_open_sell(config, open_orders, closed_orders, bot_path, client, usr)
    if length(open_orders["SELL"]) == 0
        return open_orders, closed_orders
    end
    fees = haskey(config, "current_fees") ? config["current_fees"][config["symbol"]] : client.get_trade_fee(symbol=config["symbol"])["tradeFee"][1]

    #
    ## Go through all open SELL orders
    for (k,v) in open_orders["SELL"]
        id_str = split(k, "_")[end]
        ##------------------------------------------------------------------------##
        ## Pinguu in STOP mode----------------------------------------------------##
        ##------------------------------------------------------------------------##
        if config["status"] == "stopped" && v["status"] != "stopped"
            check_stopped_orders(k, v, config, open_orders, closed_orders, bot_path, client, usr)
            continue
        elseif config["status"] == "active" && v["status"] == "stopped"
            check_orders_to_revive(k, v, config, open_orders, closed_orders, bot_path, client, usr)
            continue
        end

        ##------------------------------------------------------------------------##
        ## PLACE ORDER command----------------------------------------------------##
        ##------------------------------------------------------------------------##
        if v["status"] == "place order"
            place_sell_order_command(k, v, id_str, config, open_orders, closed_orders, bot_path, client, usr, fees)
            continue
        end

        ##------------------------------------------------------------------------##
        ## Load SELL updates------------------------------------------------------##
        ##------------------------------------------------------------------------##
        tmp = try client.get_order(
            symbol=config["symbol"],
            orderId=v["orderId"])
        catch
            @info("Skipped..")
            open(joinpath(bot_path, "trigger.json"), "w") do f
                JSON.print(f, Dict(), 4)
            end
            continue
        end

        ##------------------------------------------------------------------------##
        ## SELL Filled------------------------------------------------------------##
        ##------------------------------------------------------------------------##
        if tmp["status"] == client.ORDER_STATUS_FILLED
            sell_filled(tmp, k, v, id_str, config, open_orders, closed_orders, bot_path, client, usr, fees)
            continue
        end

        ##------------------------------------------------------------------------##
        ## SELL Canceled----------------------------------------------------------##
        ##------------------------------------------------------------------------##
        if tmp["status"] == client.ORDER_STATUS_CANCELED
            sell_canceled(tmp, k, v, id_str, config, open_orders, closed_orders, bot_path, client, usr, fees)
            continue
        end

        ##------------------------------------------------------------------------##
        ## Cancel SELL------------------------------------------------------------##
        ##------------------------------------------------------------------------##
        if tmp["status"] == "cancel order"
            cancel_sell(tmp, k, v, id_str, config, open_orders, closed_orders, bot_path, client, usr, fees)
            continue
        end

        ##------------------------------------------------------------------------##
        ## SELL status changed----------------------------------------------------##
        ##------------------------------------------------------------------------##
        if tmp["status"] != v["status"]
            sell_status_changed(tmp, k, v, id_str, config, open_orders, closed_orders, bot_path, client, usr, fees)
            continue
        end
    end
    return open_orders, closed_orders
end
export check_open_sell


"""
`module  JAC.MultiPhotonDeExcitation`
    ... a submodel of JAC that contains all methods for computing multi-photon excitation and decay rates.
"""
module MultiPhotonDeExcitation

    using Printf, QuadGK, ..AngularMomentum, ..AtomicState, ..Basics, ..Defaults, ..ManyElectron, ..Nuclear,
                          ..PhotoEmission, ..Radial, ..TableStrings

    """
    `abstract type MultiPhotonDeExcitation.AbstractMultiPhotonProperty`
        ... defines an abstract type to distinguish different multi-photon properties; however, by far not all of these properties
            are usually relevant for all the process::MultiPhotonDeExcitation.AbstractMultiPhotonProcess; cf. the JAC User Guide:

        + struct EnergyDiffCs
                    ... Energy differential cross section for initially unpolarized atoms (TwoPhotonEmission, ...)
        + struct TotalCsLinear
                    ... Total cross section for linerarly-polarized incident radiation (TwoPhotonAbsorptionMonochromatic, ...)
        + struct TotalCsRightCircular
                    ... Total cross section for right-circularly polarized incident radiation (TwoPhotonAbsorptionMonochromatic, ...)
        + struct TotalCsUnpolarized
                    ... Total cross section for unpolarized incident radiation (TwoPhotonAbsorptionMonochromatic, ...)
        + struct TotalCsDensityMatrix
                    ... Total cross section for incident radiation with user-defined density matrix (not yet).
                        (TwoPhotonAbsorptionMonochromatic, ...)
    """
    abstract type  AbstractMultiPhotonProperty       end

    """
    `struct  MultiPhotonDeExcitation.EnergyDiffCs  <:  MultiPhotonDeExcitation.AbstractMultiPhotonProperty`
        ... a struct to request for the calculation of energy-diffrential cross sections in a two-photon emission process.
    """
    struct   EnergyDiffCs  <:  MultiPhotonDeExcitation.AbstractMultiPhotonProperty      end

    function Base.show(property::EnergyDiffCs)
        sa = "energy-differential cross sections for multi-photon emission";  print(io, sa, "\n")
    end

    """
    `struct  MultiPhotonDeExcitation.TotalCsLinear  <:  MultiPhotonDeExcitation.AbstractMultiPhotonProperty`
        ... a struct to request for the calculation of the total (absorption) cross sections for linearly-polarized
            incident radiation.
    """
    struct   TotalCsLinear  <:  MultiPhotonDeExcitation.AbstractMultiPhotonProperty      end

    function Base.show(property::TotalCsLinear)
        sa = "total (absorption) cross sections for linearly-polarized radiation";  print(io, sa, "\n")
    end

    """
    `struct  MultiPhotonDeExcitation.TotalCsRightCircular  <:  MultiPhotonDeExcitation.AbstractMultiPhotonProperty`
        ... a struct to request for the calculation of the total (absorption) cross sections for right-circularly polarized
            incident radiation.
    """
    struct   TotalCsRightCircular  <:  MultiPhotonDeExcitation.AbstractMultiPhotonProperty      end

    function Base.show(property::TotalCsRightCircular)
        sa = "total (absorption) cross sections for right-circularly polarized radiation";  print(io, sa, "\n")
    end

    """
    `struct  MultiPhotonDeExcitation.TotalCsUnpolarized  <:  MultiPhotonDeExcitation.AbstractMultiPhotonProperty`
        ... a struct to request for the calculation of the total (absorption) cross sections for unpolarized
            incident radiation.
    """
    struct   TotalCsUnpolarized  <:  MultiPhotonDeExcitation.AbstractMultiPhotonProperty      end

    function Base.show(property::TotalCsUnpolarized)
        sa = "total (absorption) cross sections for unpolarized radiation";  print(io, sa, "\n")
    end

    """
    `abstract type MultiPhotonDeExcitation.AbstractMultiPhotonProcess`
        ... defines an abstract type to distinguish different multi-photon processes; see also:

        + struct TwoPhotonAbsorptionMonochromatic
                    ... to specify a two-photon absorption process with monochromatic and equally-polarized photons,
                        usually from the same beam.
        + struct TwoPhotonAbsorptionBichromatic
                    ... to specify a two-photon absorption process by photons with well-defined frequency, propagation
                        direction and helicity, usually from different (oriented) beams.
        + struct TwoPhotonEmission
                    ... to specify a two-photon emission process, often for initially-unpolarized atoms.
    """
    abstract type  AbstractMultiPhotonProcess       end

    """
    `struct  MultiPhotonDeExcitation.TwoPhotonAbsorptionMonochromatic  <:  MultiPhotonDeExcitation.AbstractMultiPhotonProcess`
        ... a struct to specify a two-photon absorption process with monochromatic and equally-polarized photons,
            usually from the same beam.

        + properties        ::Array{MultiPhotonDeExcitation.AbstractMultiPhotonProperty,1}
            ... List of the multi-photon processes that are to be calculated for this absorption process.
    """
    struct   TwoPhotonAbsorptionMonochromatic  <:  MultiPhotonDeExcitation.AbstractMultiPhotonProcess
        properties          ::Array{MultiPhotonDeExcitation.AbstractMultiPhotonProperty,1}
    end

    """
    `MultiPhotonDeExcitation.TwoPhotonAbsorptionMonochromatic()`
        ... constructor for an 'default' instance of a MultiPhotonDeExcitation.TwoPhotonAbsorptionMonochromatic() process.
    """
    function TwoPhotonAbsorptionMonochromatic()
        TwoPhotonAbsorptionMonochromatic( [TotalCsLinear(), TotalCsUnpolarized()] )
    end

    function Base.string(process::MultiPhotonDeExcitation.TwoPhotonAbsorptionMonochromatic)
        sa = "Two-photon absorption process with mono-chromatic photons (from the same beam):"
        return( sa )
    end

    function Base.show(io::IO, process::MultiPhotonDeExcitation.TwoPhotonAbsorptionMonochromatic)
        sa = Base.string(process);               print(io, sa, "\n")
        println(io, "properties:                 $(process.properties)  ")
    end

    """
    `struct  MultiPhotonDeExcitation.TwoPhotonAbsorptionBichromatic  <:  MultiPhotonDeExcitation.AbstractMultiPhotonProcess`
        ... a struct to specify a two-photon absorption process by photons with well-defined frequency, propagation direction
            and helicity, usually from different (oriented) beams.

        + properties        ::Array{MultiPhotonDeExcitation.AbstractMultiPhotonProperty,1}
            ... List of the multi-photon processes that are to be calculated for this absorption process.
        + omegaLess         ::Float64
            ... energy of the photon with smaller frequency; the larger frequency is derived from energy conservation.
    """
    struct   TwoPhotonAbsorptionBichromatic  <:  MultiPhotonDeExcitation.AbstractMultiPhotonProcess
        properties          ::Array{MultiPhotonDeExcitation.AbstractMultiPhotonProperty,1}
        omegaLess           ::Float64
    end

    """
    `MultiPhotonDeExcitation.TwoPhotonAbsorptionBichromatic()`
        ... constructor for an 'default' instance of a MultiPhotonDeExcitation.TwoPhotonAbsorptionBichromatic() process.
    """
    function TwoPhotonAbsorptionBichromatic()
        TwoPhotonAbsorptionBichromatic( [TotalCsUnpolarized()], 0.05 )
    end

    function Base.string(process::MultiPhotonDeExcitation.TwoPhotonAbsorptionBichromatic)
        sa = "Two-photon absorption process with bi-chromatic photons (from different beams):"
        return( sa )
    end

    function Base.show(io::IO, process::MultiPhotonDeExcitation.TwoPhotonAbsorptionBichromatic)
        sa = Base.string(process);               print(io, sa, "\n")
        println(io, "properties:                 $(process.properties)  ")
        println(io, "omegaLess:                  $(process.omegaLess)  ")
    end

    """
    `struct  MultiPhotonDeExcitation.TwoPhotonEmission  <:  MultiPhotonDeExcitation.AbstractMultiPhotonProcess`
        ... a struct to specify a two-photon emission process.

        + properties        ::Array{MultiPhotonDeExcitation.AbstractMultiPhotonProperty,1}
            ... List of the multi-photon processes that are to be calculated for this emission process.
        + noSharings  ::Int64   ... number of energy sharings for atomic transition, using the zeros of a Gauß-Legendre integration
    """
    struct   TwoPhotonEmission  <:  MultiPhotonDeExcitation.AbstractMultiPhotonProcess
        properties          ::Array{MultiPhotonDeExcitation.AbstractMultiPhotonProperty,1}
        noSharings          ::Int64
    end

    """
    `MultiPhotonDeExcitation.TwoPhotonEmission()`
        ... constructor for an 'default' instance of a MultiPhotonDeExcitation.TwoPhotonEmission() process.
    """
    function TwoPhotonEmission()
        TwoPhotonEmission( [EnergyDiffCs()], 4 )
    end

    function Base.string(process::MultiPhotonDeExcitation.TwoPhotonEmission)
        sa = "Two-photon emission process:"
        return( sa )
    end

    function Base.show(io::IO, process::MultiPhotonDeExcitation.TwoPhotonEmission)
        sa = Base.string(process);               print(io, sa, "\n")
        println(io, "properties:                 $(process.properties)  ")
    end

    """
    `struct  MultiPhotonDeExcitation.Settings  <:  AbstractProcessSettings`  ... defines a type for the settings in estimating multi-photon excitation and decay rates.

        + process           ::MultiPhotonDeExcitation.AbstractMultiPhotonProcess      ... considered multi-photon process.
        + multipoles        ::Array{EmMultipole}           ... Specifies the multipoles of the radiation field that are to be included.
        + gauges            ::Array{UseGauge}              ... Specifies the gauges to be included into the computations.
        + green             ::Array{AtomicState.GreenChannel,1}
                              ... Green function channels for these multi-photon calculations which need to be generated independently.
        + NoEnergySharings  ::Int64                        ... Number of energy sharings that are used for all emission lines in the
                                                               computations for each line.
        + printBefore       ::Bool                         ... True, if all energies and lines are printed before their evaluation.
        + lineSelection     ::LineSelection                ... Specifies the selected levels, if any.

    """
    struct Settings  <:  AbstractProcessSettings
        process             ::MultiPhotonDeExcitation.AbstractMultiPhotonProcess
        multipoles          ::Array{EmMultipole}
        gauges              ::Array{UseGauge}
        green               ::Array{AtomicState.GreenChannel,1}
        NoEnergySharings    ::Int64
        printBefore         ::Bool
        lineSelection       ::LineSelection
    end

    """
    `MultiPhotonDeExcitation.Settings()`  ... constructor for the default values of multi-photon excitation and decay rates.
    """
    function Settings()
        Settings( TwoPhotonEmission(), EmMultipole[], UseGauge[], AtomicState.GreenExpansion(), false, LineSelection() )
    end

    # `Base.show(io::IO, settings::MultiPhotonDeExcitation.Settings)`
    #		... prepares a proper printout of the variable settings::MultiPhotonDeExcitation.Settings.
     function Base.show(io::IO, settings::MultiPhotonDeExcitation.Settings)
        println(io, "process:                  $(settings.process)  ")
        println(io, "multipoles:               $(settings.multipoles)  ")
        println(io, "gauges:                   $(settings.gauges)  ")
        print(io,   "green:                     (settings.green)  ... with symmetries:   ")
        for channel in settings.green  print(io, "   $(channel.symmetry)")  end;   println(io, " ")
        println(io, "NoEnergySharings:         $(settings.NoEnergySharings)  ")
        println(io, "printBefore:              $(settings.printBefore)  ")
        println(io, "lineSelection:            $(settings.lineSelection)  ")
    end

    include("module-MultiPhotonDeExc-inc-2pAbsMono.jl")
    include("module-MultiPhotonDeExc-inc-2pEmission.jl")

end # module


using Printf

abstract type PerformanceCurve end

xmetric(::Type{<: PerformanceCurve}) = error("unimplemented")
xmetric(::T) where {T <: PerformanceCurve} = xmetric(T)

ymetric(::Type{<: PerformanceCurve}) = error("unimplemented")
ymetric(::T) where {T <: PerformanceCurve} = ymetric(T)

xlabel(::Type{<: PerformanceCurve}) = error("unimplemented")
xlabel(::T) where {T <: PerformanceCurve} = xlabel(T)

ylabel(::Type{<: PerformanceCurve}) = error("unimplemented")
ylabel(::T) where {T <: PerformanceCurve} = ylabel(T)

xvalues(::PerformanceCurve) = error("unimplemented")
yvalues(::PerformanceCurve) = error("unimplemented")

startpoint(::Type{<: PerformanceCurve}, xs, ys) = error("unimplemented")
startpoint(::T, args...) where {T <: PerformanceCurve} = startpoint(T, args...)

endpoint(::Type{<: PerformanceCurve}, xs, ys) = error("unimplemented")
endpoint(::T, args...) where {T <: PerformanceCurve} = endpoint(T, args...)

sortage!(::Type{<: PerformanceCurve}, xs, ys) = error("unimplemented")

function downsample(xs::AbstractVector{Float64}, ys::AbstractVector{Float64})
    keep = [1]
    for i in 2:length(xs)
        if !(xs[i] ≈ xs[i-1] && ys[i] ≈ ys[i-1])
            push!(keep, i)
        end
    end
    keep
end

function prepare!(::Type{T},
                 xs::AbstractVector{Float64},
                 ys::AbstractVector{Float64};
                 down=true) where {T <: PerformanceCurve}

    if size(xs) != size(ys)
        throw(DimensionMismatch("$(xlabel(T)) and $(ylabel(T)) vectors have different sizes"))
    elseif any(<(0.0), xs)
        throw(ArgumentError("$(xlabel(T)) less than 0.0 provided"))
    elseif any(<(0.0), ys)
        throw(ArgumentError("$(ylabel(T)) less than 0.0 provided"))
    elseif any(>(1.0), xs)
        throw(ArgumentError("$(xlabel(T)) greater than 1.0 provided"))
    elseif any(>(1.0), ys)
        throw(ArgumentError("$(ylabel(T)) greater than 1.0 provided"))
    end

    sortage!(T, xs, ys)

    if down
        d = downsample(xs, ys)
        xs = xs[d]
        ys = ys[d]
    end

    s = startpoint(T, xs, ys)
    if !all((xs[begin], ys[begin]) .≈ s)
        pushfirst!(xs, s[begin])
        pushfirst!(ys, s[end])
    end

    e = endpoint(T, xs, ys)
    if !all((xs[end], ys[end]) .≈ e)
        push!(xs, e[begin])
        push!(ys, e[end])
    end

    xs, ys
end

Base.length(c::PerformanceCurve) = length(xvalues(c))

function MLBase.roc(T::Type{<:PerformanceCurve},
                    arena::AbstractArena,
                    edges::Vector{EdgeEvidence},
                    args...)
    g = graph(arena)
    gt = [Int(has_edge(g, e.src, e.dst)) for e in edges]
    scores = [e.evidence for e in edges]
    if all(==(scores[1]), scores)
        T(gt)
    else
        x, y = xmetric(T), ymetric(T)
        r = roc(gt, scores, args...)
        T(map(x, r), map(y, r))
    end
end

auc(c::PerformanceCurve) = let xs = xvalues(c), ys = yvalues(c)
    @views dot(ys[2:end] + ys[1:end-1], xs[2:end] - xs[1:end-1]) / 2.0
end

struct Curve
    xs::Vector{Float64}
    ys::Vector{Float64}
    function Curve(xs::AbstractVector{Float64}, ys::AbstractVector{Float64})
        if length(xs) != length(ys)
            throw(DimensionMismatch("the x- and y-position vectors must have the same length"))
        end
        new(xs, ys)
    end
end
Curve(r::PerformanceCurve) = Curve(xvalues(r), yvalues(r))

Base.length(c::Curve) = length(c.xs)

function Base.:+(a::Curve, b::Curve)
    xs, ys = Float64[], Float64[]
    i, j = 1, 1
    @views while i <= length(a) && j <= length(b)
        if a.xs[i] ≈ b.xs[j]
            push!(xs, a.xs[i])
            push!(ys, a.ys[i] + b.ys[j])
            i += 1
            j += 1
        elseif a.xs[i] < b.xs[j]
            push!(xs, a.xs[i])
            push!(ys, a.ys[i] + interpolate(b.xs[j-1:j], b.ys[j-1:j], a.xs[i]))
            i += 1
        else
            push!(xs, b.xs[j])
            push!(ys, b.ys[j] + interpolate(a.xs[i-1:i], a.ys[i-1:i], b.xs[j]))
            j += 1
        end
    end
    Curve(xs, ys)
end

function Base.:-(a::Curve, b::Curve)
    xs, ys = Float64[], Float64[]
    i, j = 1, 1
    @views while i <= length(a) && j <= length(b)
        if a.xs[i] ≈ b.xs[j]
            push!(xs, a.xs[i])
            push!(ys, a.ys[i] - b.ys[j])
            i += 1
            j += 1
        elseif a.xs[i] < b.xs[j]
            push!(xs, a.xs[i])
            push!(ys, a.ys[i] - interpolate(b.xs[j-1:j], b.ys[j-1:j], a.xs[i]))
            i += 1
        else
            push!(xs, b.xs[j])
            push!(ys, b.ys[j] - interpolate(a.xs[i-1:i], a.ys[i-1:i], b.xs[j]))
            j += 1
        end
    end
    Curve(xs, ys)
end

function interpolate(xs::AbstractVector{Float64}, ys::AbstractVector{Float64}, x::Float64)
    m = (ys[2] - ys[1]) / (xs[2] - xs[1])
    m * (x - xs[1]) + ys[1]
end

Base.:/(c::Curve, k::Number) = Curve(c.xs, c.ys / k)

Statistics.mean(cs::AbstractVector{T}) where {T <: PerformanceCurve} = T(mean(Curve.(cs)))

struct ROC <: PerformanceCurve
    fpr::Vector{Float64}
    tpr::Vector{Float64}
    function ROC(fpr::Vector{Float64}, tpr::Vector{Float64}; kwargs...)
        fpr, tpr = prepare!(ROC, deepcopy(fpr), deepcopy(tpr); kwargs...)
        new(fpr, tpr)
    end
end
ROC() = ROC([0.0, 1.0], [0.0, 1.0])
ROC(::AbstractArray{Int}) = ROC()
ROC(c::Curve) = ROC(c.xs, c.ys)

fpr(r::ROC) = r.fpr

tpr(r::ROC) = r.tpr

xmetric(::Type{ROC}) = false_positive_rate

ymetric(::Type{ROC}) = true_positive_rate

xlabel(::Type{ROC}) = :FPR

ylabel(::Type{ROC}) = :TPR

startpoint(::Type{ROC}, fpr, tpr) = (0.0, 0.0)

endpoint(::Type{ROC}, fpr, tpr) = (1.0, 1.0)

xvalues(r::ROC) = fpr(r)

yvalues(r::ROC) = tpr(r)

function sortage!(::Type{ROC}, fpr, tpr)
    if !(issorted(fpr) && issorted(tpr))
        p = sortperm(collect(zip(fpr, tpr)))
        fpr[:] = fpr[p]
        tpr[:] = tpr[p]
    end

    if !issorted(tpr)
        throw(ArgumentError("the $T curve is not monotonic"))
    end

    fpr, tpr
end

struct PRC <: PerformanceCurve
    recall::Vector{Float64}
    precision::Vector{Float64}
    function PRC(recall::Vector{Float64}, precision::Vector{Float64}; kwargs...)
        recall, precision = prepare!(PRC, deepcopy(recall), deepcopy(precision); kwargs...)
        new(recall, precision)
    end
end
PRC(baseline::Float64=0.5) = PRC([0.0, 0.0, 1.0, 1.0], [1.0, baseline, baseline, 0.0])
PRC(gt::AbstractArray{Int}) = PRC(mean(gt))
PRC(c::Curve) = PRC(c.xs, c.ys)

recall(r::PRC) = r.recall

precision(r::PRC) = r.precision

xmetric(::Type{PRC}) = MLBase.recall

ymetric(::Type{PRC}) = MLBase.precision

xlabel(::Type{PRC}) = :Recall

ylabel(::Type{PRC}) = :Precision

startpoint(::Type{PRC}, recall, precision) = (0.0, 1.0)

endpoint(::Type{PRC}, recall, precision) = (1.0, 0.0)

xvalues(r::PRC) = recall(r)

yvalues(r::PRC) = precision(r)

function prcorder((r1, p1), (r2, p2))
    isequal(r1, r2) && return isless(p2, p1)
    isless(r1, r2)
end

function sortage!(::Type{PRC}, recall, precision)
    p = sortperm(collect(zip(recall, precision)); lt = prcorder)
    recall[:] = recall[p]
    precision[:] = precision[p]

    recall, precision
end

@recipe function plotroc(r::ROC)
    a = auc(r)
    if !get(plotattributes, :noauc, false)
        label := if haskey(plotattributes, :label) && plotattributes[:label] != ""
            @sprintf "%s (AUC = %0.3f)" plotattributes[:label] a
        elseif !haskey(plotattributes, :label)
            @sprintf "AUC = %0.3f" a
        end
    end

    xlims := (0, 1)
    ylims := (0, 1)

    grid -> false
    legend --> ifelse(a ≥ 0.5, :bottomright, :topleft)
    size --> (600, 600)
    xguide --> "FPR"
    yguide --> "TPR"
    linewidth --> 2

    @series begin
        r.fpr, r.tpr
    end

    @series begin
        seriestype := :path
        line := :dash
        linewidth := 2
        seriescolor := :dimgray
        label := nothing
        [0,1], [0,1]
    end
end

safelog(x::Float64) = iszero(x) ? zero(x) : log(x)
function rescale(scores::AbstractArray{Float64})
    a, b = extrema(scores)
    @. (scores - a) / (b - a)
end

function perf(arena::AbstractArena, edges::Vector{EdgeEvidence}, kernel::Function)
    g = graph(arena)
    gt = Float64[has_edge(g, e.src, e.dst) for e in edges]
    scores = rescale([e.evidence for e in edges])
    ll = 0.
    for (g, s) in zip(gt, scores)
        ll += kernel(g, s)
    end
    ll / length(edges)
end

logloss(arena, edges) = perf(arena, edges, (g,s) -> -(g * safelog(s) + (1. - g) * safelog(1. - s)))
brier(arena, edges) = perf(arena, edges, (g,s) -> (g - s)^2)


#-------------------------------------------------------------------------------------------------------#
#											Input Types
#-------------------------------------------------------------------------------------------------------#
abstract type StandardArray{T,N} <: AbstractArray{T,N} end
abstract type AbstractInput{N} <: StandardArray{Float64,N} end
struct RandInput{N} <: AbstractInput{N}
	data::Array{<:Any,N}
	stages::Vector{Int}
	input_bool::Vector{Bool}
	reward_bool::Vector{Bool}
	punishment_bool::Vector{Bool}
	cum_stages::Vector{Int}
end
struct SparseInput{N} <: AbstractInput{N}
	data::Array{<:Any,N}
	stages::Vector{Int}
	input_bool::Vector{Bool}
	reward_bool::Vector{Bool}
	punishment_bool::Vector{Bool}
	cum_stages::Vector{Int}
end
struct RestInput{N} <: AbstractInput{N}
	data::Array{<:Any,N}
	stages::Vector{Int}
	input_bool::Vector{Bool}
	reward_bool::Vector{Bool}
	punishment_bool::Vector{Bool}
	cum_stages::Vector{Int}
end
struct SparseRandInput{N} <: AbstractInput{N}
	data::Array{<:Any,N}
	stages::Vector{Int}
	input_bool::Vector{Bool}
	reward_bool::Vector{Bool}
	punishment_bool::Vector{Bool}
	cum_stages::Vector{Int}
end
struct ColorInput{N} <: AbstractInput{N}
	data::Array{<:Any,N}
	stages::Vector{Int}
	input_bool::Vector{Bool}
	reward_bool::Vector{Bool}
	punishment_bool::Vector{Bool}
	cum_stages::Vector{Int}
	image::Image
end

Inputs = Union{RandInput, SparseInput, RestInput, SparseRandInput, ColorInput}

#-------------------------------------------------------------------------------------------------------#
#											Network Types
#-------------------------------------------------------------------------------------------------------#

struct NeuronLayer{T,N} <: AbstractArray{T,N}
	data::Array{T,1}
	dims::NTuple{1,Int}
end

struct NeuronLayers{T,N}
	layers::Array{NeuronLayer{T,N},1}
end

struct SynapseLayer{T,N} <: AbstractArray{T,N}
	data::Array{T,N}
	dims::NTuple{N,Int}
end

struct MBONLayer{T,N} <: AbstractArray{T,N}
	data::Array{T,N}
	dims::NTuple{N,Int}
end

struct SynapseLayers{T,N}
	layers::Array{SynapseLayer{T,N},1}
end

abstract type AbstractSparseConnection{T,Int,N} <: AbstractSparseArray{T,Int,N}
end

abstract type AbstractDenseConnection{T,N} <: AbstractArray{T,N}
end

struct SparseConnection{T,N} <: AbstractSparseConnection{T,Int,N}
	data::AbstractSparseArray{T,Int,N}
	dims::NTuple{N,Int}
end

struct DenseConnection{T,N} <: AbstractDenseConnection{T,N}
	data::Array{T,N}
	dims::NTuple{N,Int}
end

Connection = Union{SparseConnection,DenseConnection}
NormProduct = Union{Number, AbstractSparseArray}
ConnectionLayer = Union{NeuronLayer, SynapseLayer, MBONLayer}

struct DiverseLayers{T}
	layers::Array{T,1}
end

Iterable = Union{Array, Tuple}
BrainTypes = Union{NeuronLayers, SynapseLayers, DiverseLayers}

mutable struct Dopamine <: AbstractArray{Float64, 1}
	reward::Float64
	punishment::Float64
end

size(da::Dopamine) = (2,)
function getindex(da::Dopamine, i::Int)
	if i == 1
		return da.reward
	elseif i == 2
		return da.punishment
	else
		error("Tried to access a dopamine type with index != 1 or 2")
	end
end
function setindex!(da::Dopamine, f, i::Int)
	if i == 1
		da.reward = f
	elseif i == 2
		da.punishment = f
	else
		error("Tried to access a dopamine type with index != 1 or 2")
	end
end
Dopamine(num::Number) = Dopamine(num,num)