#===============================================================================

    Defines:
        - DiffusionGlobalWorkspace      : global workspace for diffusions
        - DiffusionGlobalSubworkspace   : convenience object with containers
        - DiffusionLocalWorkspace       : local workspace for diffusions
        - DiffusionLocalSubworkspace    : convenience object with containers
        - LocalUpdtParamNames           : convenience object for setting params
    and methods for them.

===============================================================================#

#===============================================================================
                        Global workspace for diffusions
===============================================================================#

"""
    struct DiffusionGlobalSubworkspace{T,TGP,TW,TWn,TX} <: eMCMC.GlobalWorkspace{T}
        P::TGP
        WW::TW
        Wnr::TWn
        XX::TX
    end

Convenience sub-workspace that holds
- `P::Vector{<:Vector{<:GuidProp}}`: the `GuidProp` laws [TODO change to PP or
                                     PPP]
- `WW::Vector{<:Vector{<:Trajectory}}`: the containers for sampling the Wiener
                                        noise
- `Wnr::Vector{<:Wiener}`: the flags for the Wiener noise
- `XX::Vector{<:Vector{<:Trajectory}}`: the containers for sampling the
                                        diffusion paths
The outer array corresponds to successive `recording`s, the inner arrays
correspond to successive observations within each recording. Each
`DiffusionSubworkspace` has two `DiffusionGlobalSubworkspace`s, one for the
accepted law/path another for the proposal.

    DiffusionGlobalSubworkspace{T}(
        aux_laws, data::AllObservations, tts, args,
    ) where T

A standard constructor. `aux_laws` is a list of auxiliary laws (one for each
recording), `data` holds the observations, `tts` are the time-grids on which
to do imputation when sampling and `args` are additional positional arguments
passed to initializers of `GuidProp`.
"""
struct DiffusionGlobalSubworkspace{T,TGP,TW,TWn,TX} <: eMCMC.GlobalWorkspace{T}
    P::TGP
    WW::TW
    Wnr::TWn
    XX::TX

    function DiffusionGlobalSubworkspace{T}(
            aux_laws, data::AllObservations, tts, args,
        ) where T
        N = num_recordings(data)
        P = map(1:N) do i
            build_guid_prop(aux_laws[i], data.recordings[i], tts[i], args[i])
        end

        traj = map(1:N) do i
            trajectory(tts[i], data.recordings[i].P)
        end
        XX = map(x->x.process, traj)
        WW = map(x->x.wiener, traj)
        Wnr = [Wiener(data.recordings[i].P) for i=1:N]
        init_paths!(P, WW, Wnr, XX, data)
        new{T,typeof(P),typeof(WW),typeof(Wnr),typeof(XX)}(P, WW, Wnr, XX)
    end
end

"""
    init_paths!(P, WW, Wnr, XX, data)

Sample paths of guided proposals without using the preconditioned Crank–Nicolson
scheme.
"""
function init_paths!(P, WW, Wnr, XX, data)
    num_recordings = length(P)

    # this loop should be entirely parallelizable, as recordings
    # are—by design—conditionally independent
    for i in 1:num_recordings
        success = false
        while !success
            y1 = rand(data.recordings[i].x0_prior)
            success = forward_guide!(P[i], XX[i], WW[i], y1; Wnr=Wnr[i])
        end
    end
end

"""
    struct DiffusionGlobalWorkspace{T,SW,SWD,TD,TX} <: eMCMC.GlobalWorkspace{T}
        sub_ws::SW
        sub_ws_diff::SWD
        sub_ws_diff°::SWD
        data::TD
        stats::GenericChainStats{T}
        XX_buffer::TX
        pnames::Vector{Symbol}
    end

Global workspace for MCMC problems concerning diffusions.

# Arguments
- `sub_ws`: `GenericGlobalSubworkspace` from `ExtensibleMCMC` that stores info
            about the parameter chain
- `sub_ws_diff`: `DiffusionGlobalSubworkspace` containers for the accepted
                 diffusion law/paths
- `sub_ws_diff°`: `DiffusionGlobalSubworkspace` containers for the proposed
                  diffusion law/paths
- `data`: observations
- `stats::GenericChainStats`: generic online statistics for the param chain
- `XX_buffer`: cache for saving a chain of paths
- `pnames`: a list of parameter names that are being updated
"""
struct DiffusionGlobalWorkspace{T,SW,SWD,TD,TX} <: eMCMC.GlobalWorkspace{T}
    sub_ws::SW
    sub_ws_diff::SWD
    sub_ws_diff°::SWD
    data::TD
    stats::GenericChainStats{T}
    XX_buffer::TX
    pnames::Vector{Symbol}
end

function eMCMC.init_global_workspace(
        backend::DiffusionMCMCBackend,
        num_mcmc_steps,
        updates::Vector{<:eMCMC.MCMCUpdate},
        data,
        θinit::Nothing;
        kwargs...
    ) where T
    dummy = OrderedDict(
        :none => -Inf, # no parameters needed for smoothing
    )
    eMCMC.init_global_workspace(
        backend, num_mcmc_steps, updates, data, dummy; kwargs...
    )
end

function eMCMC.init_global_workspace(
        ::DiffusionMCMCBackend,
        num_mcmc_steps,
        updates::Vector{<:eMCMC.MCMCUpdate},
        data,
        θinit::OrderedDict{Symbol,T};
        kwargs...
    ) where T
    dkwargs = Dict(kwargs)
    M, NU = num_mcmc_steps, length(updates)

    _θ_init = collect(values(θinit))
    _pnames = collect(keys(θinit))

    sub_ws = eMCMC.StandardGlobalSubworkspace(M, NU, data, _θ_init)

    sub_ws_diff = DiffusionGlobalSubworkspace{T}(
        extract_aux_laws(updates, data),
        data,
        OBS.setup_time_grids(
            data,
            get(dkwargs, :dt, 0.01),
            get(dkwargs, :τ, _DEFAULT_TIME_TRANSFORM),
            get(dkwargs, :Ttime, Float64),
            get(dkwargs, :tt, nothing),
        ),
        package(get(dkwargs, :guid_prop_args, tuple()), data), # from ObservationSchemes
    )

    path_saving_buffer = PathSavingBuffer(
        sub_ws_diff.XX,
        get(dkwargs, :path_buffer_size, 1)
    )

    SW, SWD, TD = typeof(sub_ws), typeof(sub_ws_diff), typeof(data)
    TX = typeof(path_saving_buffer)

    DiffusionGlobalWorkspace{T,SW,SWD,TD,TX}(
        sub_ws,
        sub_ws_diff,
        deepcopy(sub_ws_diff),
        data,
        GenericChainStats(_θ_init, NU, M),
        path_saving_buffer,
        _pnames,
    )
end

function extract_aux_laws(updates::Vector{<:eMCMC.MCMCUpdate}, data)
    imps = filter(x->typeof(x)<:PathImputation, updates)
    N = length(imps)
    @assert N > 0
    aux_laws = map(x->OBS.package(x.aux_laws, data), imps)
    @assert N == 1 || all([aux_laws[1]==aux_laws[i] for i=2:N])
    aux_laws[1]
end

#===============================================================================
                        Local workspace for diffusions
===============================================================================#
"""
    struct DiffusionLocalSubworkspace{T,TP,TW,TWn,TX,Tz0} <: eMCMC.LocalWorkspace{T}
        ll::Vector{Float64}
        ll_history::Vector{Vector{Float64}}
        P::TP
        WW::TW
        Wnr::TWn
        XX::TX
        z0s::Tz0
    end

Convenience sub-workspace that holds
- `ll::Vector{Float64}`: computational buffer for log-likelihoods (one for each
                         recording)
- `ll_history::Vector{Vector{Float64}}`:  history of all log-likelihoods
- `P`: appropriately shaped views to containers with `GuidProp`s
- `WW`: appropriately shaped views to containers for sampled Wiener process
- `Wnr`: list of Wiener flags for each recording
- `XX`: appropriately shaped views to containers for sampled diffusion paths
- `z0s`: list of white noise for starting points (one for each recording)
"""
struct DiffusionLocalSubworkspace{T,TP,TW,TWn,TX,Tz0} <: eMCMC.LocalWorkspace{T}
    ll::Vector{Float64}
    ll_history::Vector{Vector{Float64}}
    P::TP
    WW::TW
    Wnr::TWn
    XX::TX
    z0s::Tz0

    function DiffusionLocalSubworkspace{T}(
            ll, ll_hist, P::TP, WW::TW, Wnr::TWn, XX::TX, z0s::Tz0
        ) where {T,TP,TW,TWn,TX,Tz0}
        new{T,TP,TW,TWn,TX,Tz0}(ll, ll_hist, P, WW, Wnr, XX, z0s)
    end
end

function DiffusionLocalSubworkspace{T}(
        sws::DiffusionGlobalSubworkspace,
        block_layout,
        num_mcmc_steps
    ) where T
    N, M = num_mcmc_steps, length(block_layout)
    DiffusionLocalSubworkspace{T}(
        Float64[-Inf for _=1:M],
        [zeros(Float64, M) for _=1:N],
        block_view(sws.P, block_layout),
        block_view(sws.WW, block_layout),
        block_wiener(sws.Wnr, block_layout),
        block_view(sws.XX, block_layout),
        block_start_pts(sws.XX, block_layout)
    )
end

function block_view(v, layout)
    [ view(v[bl_i], bl_i_range) for (bl_i, bl_i_range) in layout ]
end

function block_view(v, layout::Nothing)
    [ view(v[i], :) for i in 1:length(v) ]
end

function block_start_pts(XX, layout::Nothing)
    [ deepcopy(X[1].x[1]) for X in XX ]
end

function block_start_pts(XX, layout)
    [ deepcopy(XX[bl_i][bl_i_range[1]].x[1]) for (bl_i, bl_i_range) in layout ]
end

block_wiener(Wnr, layout::Nothing) = Wnr

block_wiener(Wnr, layout) = [ deepcopy(Wnr[bl_i]) for (bl_i, _) in layout ]

const _SYMS = Tuple{Vararg{Pair{Int64,Symbol},N} where N}
const _OBSIDX = Tuple{Vararg{Pair{Int64,Int64},N} where N}

struct LocalUpdtParamNames
    vp_names::Vector{Tuple{Vararg{Symbol,N} where N}}
    vp_names_aux::Vector{Vector{Tuple{Vararg{Symbol,N} where N}}}
    θ_local_names::Vector{_SYMS}
    θ_local_aux_names::Vector{Vector{_SYMS}}
    θ_local_obs::Vector{Vector{_OBSIDX}}

    function LocalUpdtParamNames(
            PPP::AbstractArray{<:AbstractArray{<:GuidProp}},
            data::AllObservations,
            θinit::Vector{Symbol},
        )
        psym, oind = OBS.local_symbols(data, PPP, x->x.P_target, θinit)
        pauxsym, _ = OBS.local_symbols(data, PPP, x->x.P_aux, θinit)

        new(
            [DD.var_parameter_names(PP) for PP in PPP],
            [[DD.var_parameter_names(P.P_aux) for P in PP] for PP in PPP],
            first.(psym), # target laws are the same on each obs of a single rec
            pauxsym,
            oind,
        )
    end

    LocalUpdtParamNames() = new()
end

"""
    struct DiffusionLocalWorkspace{T,SW,SWD,Tpr,AH} <: eMCMC.LocalWorkspace{T}
        sub_ws::SW
        sub_ws°::SW
        sub_ws_diff::SWD
        sub_ws_diff°::SWD
        xx0_priors::Tpr
        acceptance_history::Vector{AH}
        loc_pnames::LocalUpdtParamNames
        critical_param_change::Vector{Bool}
    end

Local workspace for MCMC problems concerning diffusions.

# Arguments
- `sub_ws`: `GenericLocalSubworkspace` from `ExtensibleMCMC` that acts as a
            local cache for some light parameter computations
- `sub_ws°`: same as `sub_ws`, but concerns the proposed parameter
- `sub_ws_diff`: `DiffusionLocalSubworkspace`, views to containers for the
                 accepted diffusion law/paths
- `sub_ws_diff°`: `DiffusionLocalSubworkspace`, views to containers for the
                  proposed diffusion law/paths
- `xx0_priors`: priors over starting points
- `acceptance_history`: history of results from the accept/reject
                        Metropolis–Hastings step
- `loc_pnames`: various lists of parameter names helping in performing the
                update in a clean way
- `critical_param_change`: a list of flags for whether a given update by itself
                           prompts for recomputation of the guiding term
"""
struct DiffusionLocalWorkspace{T,SW,SWD,Tpr,CGH,AH} <: eMCMC.LocalWorkspace{T}
    sub_ws::SW
    sub_ws°::SW
    sub_ws_diff::SWD
    sub_ws_diff°::SWD
    xx0_priors::Tpr
    acceptance_history::Vector{AH}
    loc_pnames::LocalUpdtParamNames
    cu_helper::CGH
    critical_param_change::Vector{Bool}
end

function create_workspace(
        backend::DiffusionMCMCBackend,
        mcmcupdate,
        global_ws,
        block_layout,
        num_mcmc_steps,
    )
    _state = (
        typeof(mcmcupdate) <: eMCMC.MCMCParamUpdate ?
        copy(state(global_ws, mcmcupdate)) :
        Float64[]
    )
    sub_ws = eMCMC.StandardLocalSubworkspace(_state, num_mcmc_steps)
    xx0_priors = block_setup_x0_priors(
        global_ws.data, block_layout, global_ws.sub_ws_diff
    )
    T = eltype(state(global_ws))
    sub_ws_diff = DiffusionLocalSubworkspace{T}(
        global_ws.sub_ws_diff, block_layout, num_mcmc_steps
    )
    init_ll!(sub_ws_diff)

    sub_ws_diff° = DiffusionLocalSubworkspace{T}(
        global_ws.sub_ws_diff°, block_layout, num_mcmc_steps
    )
    a_h = (
        typeof(mcmcupdate) <: eMCMC.MCMCParamUpdate ?
        fill(false, num_mcmc_steps) :
        [similar(v, Bool) for v in sub_ws_diff.ll_history]
    )

    init_update!(mcmcupdate, block_layout)

    lpn = (
        typeof(mcmcupdate) <: eMCMC.MCMCParamUpdate ?
        LocalUpdtParamNames(
            sub_ws_diff.P,
            global_ws.data,
            global_ws.pnames[eMCMC.coords(mcmcupdate)]
        ) :
        LocalUpdtParamNames()
    )

    cgu_helper = (
        typeof(mcmcupdate) <: DiffusionConjugGsnUpdate ?
        ConjugGsnHelper(
            mcmcupdate,
            global_ws,
            sub_ws_diff.P,
            lpn,
            #NOTE finally, computation type, stay with outofplace for now
        ) :
        ConjugGsnHelper()
    )

    crit_updt = (
        typeof(mcmcupdate) <: eMCMC.MCMCParamUpdate ?
        [
            GP.is_critical_update(
                PP, lpn.θ_local_aux_names[i], lpn.θ_local_obs[i]
            ) for (i,PP) in enumerate(sub_ws_diff.P)
        ] :
        [false for _ in eachindex(sub_ws_diff.P)]
    )

    DiffusionLocalWorkspace{
        T, typeof(sub_ws), typeof(sub_ws_diff), typeof(xx0_priors),
        typeof(cgu_helper), eltype(a_h)
    }(
        sub_ws, deepcopy(sub_ws), sub_ws_diff, sub_ws_diff°, xx0_priors, a_h,
        lpn, cgu_helper, crit_updt,
    )
end

function init_ll!(sub_ws_diff)
    XXX, PPP = sub_ws_diff.XX, sub_ws_diff.P
    for i in 1:length(sub_ws_diff.ll)
        recompute_guiding_term!(PPP[i])
        sub_ws_diff.ll[i] = loglikhd(PPP[i], XXX[i])
    end
end

function block_setup_x0_priors(data, block_layout, sub_ws_diff)
    #println(block_layout)
    map(block_layout) do b_l
        bl_i, bl_i_range = b_l
        #println(bl_i, bl_i_range)
        i1 = bl_i_range[1]
        pr = (
            i1 == 1 ?
            data.recordings[bl_i].x0_prior :
            KnownStartingPt(sub_ws_diff.XX[bl_i][i1].x[1])
        )
        pr
    end
end

function eMCMC.create_workspaces(backend::DiffusionMCMCBackend, mcmc::MCMC)
    map(1:length(mcmc.updates)) do i
        create_workspace(
            backend,
            mcmc.updates[i],
            mcmc.workspace,
            mcmc.schedule.extra_info.blocking_layout[i],
            mcmc.schedule.num_mcmc_steps
        )
    end
end

function eMCMC.update_workspaces!(
        local_updt::MCMCDiffusionImputation,
        global_ws::eMCMC.GlobalWorkspace,
        local_ws::eMCMC.LocalWorkspace,
        step,
        prev_ws,
    )
    for i in 1:length(step.same_layout)
        if !step.same_layout[i]
            change_laws_for_blocking!(local_ws, i) # TODO
            recompute_loglikhd!(local_ws, i)
        else
            prev_ws===nothing || ( ll(local_ws)[i] = ll(prev_ws)[i] )
        end
    end
    local_ws, local_updt
end

accepted(ws::DiffusionLocalWorkspace, i::Int) = ws.acceptance_history[i]
#set_accepted!(ws::DiffusionLocalWorkspace, i::Int, accepted, j::Int=1) = (ws.acceptance_history[i][j] = accepted)
ll(ws::DiffusionLocalWorkspace) = ws.sub_ws_diff.ll
ll°(ws::DiffusionLocalWorkspace) = ws.sub_ws_diff°.ll
ll(ws::DiffusionLocalWorkspace, i::Int) = ws.sub_ws_diff.ll_history[i]
ll°(ws::DiffusionLocalWorkspace, i::Int) = ws.sub_ws_diff°.ll_history[i]


#!/usr/bin/env julia

# Compute probabilities that feature-level splicing changes are occuring
# monotonically across time points.
#
# As input it it takes a YAML formatted file listing timepoints. Each timepoint
# is a just a list of condition names, which should match the condition names in
# the experiment file fed to isolator. E.g.
#
#     - [h7-day20, empty-vector]
#     - [adult, h7-1yr, let-7g-oe, fetal-ventricle, fetal-atrium, IMR90]
#
# Any condition that was present in the experiment but not present here will be
# ignored when considering monotonicity.
#
# This will print a table containing any feature with > 0.5 probability of
# changing across timepoints in a specific pattern.

using DataStructures, DataFrames, YAML

if length(ARGS) != 1
    println(STDERR,
    """
        Usage: isolator summarize condition-feature-splicing-samples isolator-output.h5 | ./monotonicity.jl timepoints.yml
    """)
    exit(1)
end

function read_permutations(timepoints_filename)
    equiv_classes = YAML.load_file(timepoints_filename)

    header = split(strip(readline(STDIN)), '\t')
    condition_names = [replace(name, r"_splice_rate_samples$", "") for name in header[7:end]]

    equiv = Dict{Int, Int}()
    for (i, equiv_class) in enumerate(equiv_classes)
        for cond in equiv_class
            equiv[findfirst(condition_names, cond)] = i
        end
    end

    k = maximum(values(equiv))
    permutations = Vector{Uint64}[]

    for i in 1:length(condition_names)
        if !haskey(equiv, i)
            equiv[i] = 0
        end
    end

    features = Any[]
    for line in eachline(STDIN)
        entry = split(strip(line), '\t')
        push!(features, entry[1:6])
        rows = [split(row, ',') for row in entry[7:end]]

        M = [parse(Float32, rows[i][j])
             for i in 1:length(rows), j in 1:length(rows[1])]

        permutations_row = Array(Uint64, length(rows[1]))

        for j in 1:size(M, 2)
            ps = sortperm(M[:,j])
            for i in 1:length(ps)
                ps[i] = equiv[ps[i]]
            end

            u = Uint64(0)
            for p in ps
                if p == 0
                    continue
                end
                u *= k
                u += p
            end

            permutations_row[j] = u
        end
        push!(permutations, permutations_row)
    end

    return permutations, features, k
end

permutations, features, k = read_permutations(ARGS[1])

function decode_permutation(p, k, n)
    decoded = Array(Uint64, n)
    for i in 1:n
        decoded[n-i+1] = p % k
        p = div(p, k)
    end
    return decoded
end

function interesting_transcripts(permutations, features)
    num_samples = length(permutations[1])
    interesting = Any[]
    cutoff = 0.5
    tally = DefaultDict(Uint64, Uint64, () -> 0)
    for i in 1:size(permutations, 1)
        for key in keys(tally)
            tally[key] = 0
        end

        for j in 1:num_samples
            tally[permutations[i][j]] += 1
        end

        for (k, v) in tally
            freq = v / num_samples
            if freq > cutoff
                push!(interesting, (features[i], k, freq))
            end
        end
    end

    return interesting
end

interesting = interesting_transcripts(permutations, features)

println("gene_names\tgene_ids\tincluding_transcript_ids\texcluding_transcript_ids\tlocus\tfeature_type\tprobability")
for feature in interesting
    println(join(feature[1], '\t'), '\t', feature[3])
end


using Test
using ManifoldsBase
@testset "ManifoldsBase" begin

    if VERSION >= v"1.1"
        num_ambiguities = length(Test.detect_ambiguities(ManifoldsBase))
        #num_ambiguities > 0 && @warn "The number of ambiguities in ManifoldsBase is $(num_ambiguities)."
        if VERSION >= v"1.6"
            @test num_ambiguities <= 1
        else
            @test num_ambiguities == 0
        end
    end
    include("allocation.jl")
    include("numbers.jl")
    include("bases.jl")
    include("decorator_manifold.jl")
    include("empty_manifold.jl")
    include("errors.jl")
    include("default_manifold.jl")
    include("complex_manifold.jl")
    include("validation_manifold.jl")
    include("embedded_manifold.jl")
    include("power.jl")
    include("domain_errors.jl")
    include("vector_transport_along.jl")
end


using JEMSS

println("Loading sim.")
sim = initSim(joinpath(@__DIR__, "input/sim_config.xml"))

println("Simulating.")
simulate!(sim)

println("Writing stats files to \"$(abspath(sim.outputPath))\"")
writeStatsFiles(sim)

# read some stats from file
println("Some statistics from stats_dict.csv output file:")
statsDict = readStatsDictFile(sim.outputFiles["statsDict"].path)
@show statsDict["ambs_avgDailyBusyDurationHours"]
@show statsDict["calls_all_fracResponsesInTime"]

# tables = readTablesFromFile(sim.outputFiles["ambulancesStats"].path)
# sim.outputFiles["callsStats"].path

# # Animate, if you like:
# animate!(sim)

nothing


using ArgParse

function parse()
    s = ArgParseSettings()

    @add_arg_table s begin
        "graph"
            help = "A NetworkX-format adjacency graph with population and district metadata."
            required = true
        "outfile"
            help = "The name of the JSONL-format log "
        "--pop-col"
            default = "population"
            help = "The population column in the graph metadata."
        "--n-plans"
            default = 1000
            help = "The number of plans to generate."
        "--init-plan-col"
            default = "district"
            help = "The column specifying the initial district assignment in the graph metadata."
    end
end


## Convenient tools
# all-zero dynamics (for test)
function ẋ(env::Fym, x, t)
    println("all-zero dynamics for test; see `LazyFym.ẋ`")
    env_names = names(env)
    if env_names == []
        return zero(x)
    else
        zero_values = env_names |> Map(name -> ẋ(getfield(env, name), x[name], t))
        return (; zip(env_names, zero_values)...)  # NamedTuple
    end
end

# default update (for test)
"""
    update(env::Fym, ẋ, x, t, Δt)

Update law of a simulator.
You may have to customise this function to realise high-level agent
such as zero-order-hold (ZOH) input.
# Examples
function update(rlenv::RLEnv, ẋ, x, t, Δt)
    action = command(rlenv.policy, x.env, t)
    _datum = Dict(:x => x, :t => t, :action => action)
    x_next = ∫(rlenv, ẋ, x, t, Δt; action=action)  # `action` will be a kwarg of ẋ
    _datum[:x_next] = x_next
    datum = (; zip(keys(_datum), values(_datum))...)
    return datum, x_next
end
"""
function update(env::Fym, ẋ, x, t, Δt)  # provided
    _datum = Dict(:x => x, :t => t)
    x_next = ∫(env, ẋ, x, t, Δt)
    _datum[:x_next] = x_next
    datum = (; zip(keys(_datum), values(_datum))...)
    return datum, x_next
end

# automatic completion of initial condition
function initial_condition(env::Fym)
    env_names = names(env)
    values = env_names |> Map(name -> initial_condition(getfield(env, name)))
    return (; zip(env_names, values)...)  # NamedTuple
end


### A Pluto.jl notebook ###
# v0.14.2

using Markdown
using InteractiveUtils

# ╔═╡ caa2cdba-a76b-11eb-3461-bb8695eb2e3b
using Revise, Curiosity, Plots, PlutoUI, Statistics, Random

# ╔═╡ 0af438b0-ffe6-4f8e-9077-2f7d37b3bf3f
using ProgressMeter

# ╔═╡ cfb65788-a6be-4d5e-a404-3ebd38a1c539
include("../experiments/1d-tmaze.jl")

# ╔═╡ dfd01378-0c1d-4385-b84c-d06fb049364e
ODTMU = Curiosity.OneDTMazeUtils

# ╔═╡ 578231af-8b90-4f2f-9745-989eae047a38
default_args(ϵ) =
Dict(
	"logger_interval" => 100,

	# Behaviour Items
	# "behaviour_eta" => 0.1/8,
	"behaviour_gamma" => 0.95,
	"behaviour_learner" => "Q",
	"behaviour_update" => "ESARSA",
	"behaviour_reward_projector" => "base",
	"behaviour_rp_tilings" => 1,
	"behaviour_rp_tiles" => 16,
	"behaviour_trace" => "ReplacingTraces",
	"behaviour_opt" => "Auto",
	"behaviour_lambda" => 0.9,
	"behaviour_alpha_init" => 0.1,
	"exploration_param" => ϵ,
	"exploration_strategy" => "epsilon_greedy",
	"ϵ_range" => (0.4,0.1),
	"decay_period" => 5000,
	"warmup_steps" => 1000,
	"behaviour_w_init" => 10,

	# Demon Attributes
	"demon_alpha_init" => 0.1,
	# "demon_eta" => 0.1/8,
	"demon_discounts" => 0.9,
	"demon_learner" => "SR",
	"demon_update" => "TB",
	"demon_policy_type" => "greedy_to_cumulant",
	"demon_opt" => "Auto",
	"demon_lambda" => 0.9,
	"demon_trace"=> "AccumulatingTraces",
	"demon_beta_m" => 0.99,
	"demon_beta_v" => 0.99,

	#shared
	"num_tiles" => 8,
	"num_tilings" =>2,
	"demon_rep" => "ideal_martha",
	# "demon_rep" => "tilecoding",
	"demon_num_tiles" => 6,
	"demon_num_tilings" => 1,
	"eta" => 0.2,

	# Environment Config
	"constant_target"=> [-10.0,10.0],
	"cumulant_schedule" => "DrifterDistractor",
	"distractor" => (1.0, 5.0),
	"drifter" => (1.0, sqrt(0.01)),
	"exploring_starts"=>"beg",
	"env_step_penalty" => -0.01,

	# Agent and Logger
	"horde_type" => "regular",
	"random_first_action" => false,
	"intrinsic_reward" => "weight_change",
	# "logger_keys" => [LoggerKey.TTMAZE_ERROR],
	"save_dir" => "OneDTMazeExperiment",
	"seed" => 1,
	"steps" => 10000,
	"use_external_reward" => true,

	"logger_keys" => [LoggerKey.ONEDTMAZEERROR, LoggerKey.ONED_GOAL_VISITATION, LoggerKey.EPISODE_LENGTH, LoggerKey.INTRINSIC_REWARD]
)

# ╔═╡ 6d51db0c-72b1-47ed-8ce2-8822f15473cc
env = OneDTMaze(Curiosity.TMazeCumulantSchedules.Constant(1), "beg")

# ╔═╡ db493356-572c-4475-93ac-f9c785d8cefd
let
	@show MinimalRLCore.step!(env, 1)
	plot(env)
end

# ╔═╡ 99ad1698-d4f7-449d-8865-18a5917bb47c
function main_experiment(parsed=default_args(); progress=false, working=false)

    num_steps = parsed["steps"]
    Random.seed!(parsed["seed"])

    cumulant_schedule = ODTMU.get_cumulant_schedule(parsed)

    # exploring_starts = parsed["exploring_starts"]
    env = OneDTMaze(cumulant_schedule, parsed["exploring_starts"], parsed["env_step_penalty"])

    agent = OneDTmazeExperiment.construct_agent(parsed)
	@show parsed["behaviour_learner"],parsed["demon_learner"]

    goal_visitations = zeros(4)

	anim = Animation()

    eps = 1
    max_num_steps = num_steps
    steps = Int[]

    prg_bar = ProgressMeter.Progress(num_steps, "Step: ")
    while sum(steps) < max_num_steps
        cur_step = 0
        max_episode_steps = max_num_steps - sum(steps)
        tr, stp =
            run_episode!(env, agent, max_episode_steps) do (s, a, s_next, r, t)
                #This is a callback for every timestep where logger can go
                # agent is accesible in this scope
				plot(env)
                if t == true && working==true
                    goals = s_next[3:end]
                    f = findfirst(!iszero, goals)
                    goal_visitations[f] += 1
                end

                if progress
                    next!(prg_bar)
                end
                cur_step+=1
				frame(anim)

            end
        push!(steps, stp)
        eps += 1
    end
    if working == true
        println("Goal Visitation: ", goal_visitations)
    end
	anim
end

# ╔═╡ 15a72300-bd1a-4476-a3b6-d108035d0f25
anim = main_experiment(default_args(0.1),progress=true)

# ╔═╡ bcea5162-66e6-4e1f-8bcc-b4d738277213
mp4(anim,"../plotting/plots/tst2.mp4")

# ╔═╡ Cell order:
# ╠═caa2cdba-a76b-11eb-3461-bb8695eb2e3b
# ╠═cfb65788-a6be-4d5e-a404-3ebd38a1c539
# ╠═dfd01378-0c1d-4385-b84c-d06fb049364e
# ╠═578231af-8b90-4f2f-9745-989eae047a38
# ╠═6d51db0c-72b1-47ed-8ce2-8822f15473cc
# ╠═db493356-572c-4475-93ac-f9c785d8cefd
# ╠═99ad1698-d4f7-449d-8865-18a5917bb47c
# ╠═0af438b0-ffe6-4f8e-9077-2f7d37b3bf3f
# ╠═15a72300-bd1a-4476-a3b6-d108035d0f25
# ╠═bcea5162-66e6-4e1f-8bcc-b4d738277213


# To be used instead of _PMACDC.variable_dc_converter_ne() - supports deduplication of variables
function variable_dc_converter_ne(pm::_PM.AbstractDCPModel; kwargs...)
    variable_ne_converter_indicator(pm; kwargs..., relax=true) # FlexPlan version: replaces _PMACDC.variable_converter_ne().
    variable_ne_converter_investment(pm; kwargs...)
    _PMACDC.variable_converter_active_power_ne(pm; kwargs...)
    _PMACDC.variable_dcside_power_ne(pm; kwargs...)
    _PMACDC.variable_converter_filter_voltage_ne(pm; kwargs...)
    _PMACDC.variable_converter_internal_voltage_ne(pm; kwargs...)
    _PMACDC.variable_converter_to_grid_active_power_ne(pm; kwargs...)

    _PMACDC.variable_conv_transformer_active_power_to_ne(pm; kwargs...)
    _PMACDC.variable_conv_reactor_active_power_from_ne(pm; kwargs...)
end


"""
    readData(path_to_data_folder, complete_trayek = true)
reads the data folder given into dataframes ready to be used.
the data folder must contain:
- khazanah (id, name, x_coor, y_coor, capacity, cjarak, cpeti, transit)
- trayek (u, v, moda)
- moda (name, capacity, cjarak, cpeti, transit)
- demand (location_id, period, pecahan, value)
- stock (location_id, pecahan, value)
by default completes the trayek by ensuring for every a -> b there exists b -> a.
"""
function readData(path::String, complete_trayek::Bool=true)
    khazanah = CSV.read(joinpath(path, "khazanah.csv"), DataFrame)
    demand = CSV.read(joinpath(path, "demand.csv"), DataFrame)
    stock = CSV.read(joinpath(path, "stock.csv"), DataFrame)
    trayek = CSV.read(joinpath(path, "trayek.csv"), DataFrame)
    moda = CSV.read(joinpath(path, "moda.csv"), DataFrame)

    if complete_trayek
        append!(trayek, DataFrame(u=trayek.v, v=trayek.u, moda=trayek.moda))
        unique!(trayek)
    end

    trayek = innerjoin(trayek, moda, on=:moda => :name)

    return khazanah, demand, stock, trayek
end


using FileIO, FixedPointNumbers, Colors, ImageCore, Test, Statistics

workdir = joinpath(tempdir(), "ImagesNRRD")
isdir(workdir) || mkdir(workdir)
const teembase = "http://teem.sourceforge.net/nrrd/files"
const base3D   = "http://www.tc18.org/code_data_set/3D_greyscale"

function getfile(name, base=teembase)
    file = joinpath(workdir, name)
    if !isfile(file)
        file = download(joinpath(base, name), file)
    end
    file
end

@testset "readremote" begin
    fool_file = getfile("fool.nrrd")
    imgu = load(fool_file)
    @test eltype(imgu) == UInt8
    @test size(imgu) == (128,128)
    @test imgu[120,5] < imgu[5,5] < imgu[120,120]  # a simple orientation test
    foolf_file = getfile("foolf.nrrd")
    imgf = load(foolf_file)
    @test eltype(imgf) == Float32
    @test size(imgf) == (128,128)
    @test imgf[120,5] < imgf[5,5] < imgf[120,120]
    @test maximum(abs.(imgf-imgu))/maximum(imgu) <= 1/200
    foolc_file = getfile("foolc.nrrd")
    imgc = load(foolc_file, RGB)
    @test eltype(imgc) == RGB{N0f8}
    @test size(imgc) == (128,128)
    imgg = convert(Array{Gray}, imgc)
    @test mean(abs.(Int.(rawview(channelview(imgg))) - Int.(imgu))) < 12
#     aneurism_rawfile = getfile("aneurism.raw.gz", base3D)
#     aneurism_file = joinpath(workdir, "aneurism.nhdr")
#     aneurism_rawfile_path = joinpath(workdir, "aneurism.raw.gz")
#     open(aneurism_file, "w") do io
#         write(io, """
# NRRD0001
# content: aneurism
# # Courtesy of Philips Research, Hamburg, Germany
# dimension: 3
# type: uchar
# sizes: 256 256 256
# spacings: 1 1 1
# data file: $aneurism_rawfile_path
# encoding: gzip
# """)
#     end
#     imga = load(aneurism_file)
#     @test eltype(imga) == UInt8
#     @test size(imga) == (256,256,256)
#     @test sdims(imga) == 3
#     @test pixelspacing(imga) == (1.0,1.0,1.0)
end

GC.gc()  # to close any mmapped files
rm(workdir, recursive=true)

nothing


using Base: func_for_method_checked, structdiff, Const
using GLMakie
using Statistics
using Colors
using Distributions
include("load_eeglab.jl")

##

# Ticks
"""
Create 10 x-ticks from 10% to 90% of the x-axis as time in seconds.
time_idx should be the lower bound of the plotted data.
"""
function set_time_ticks(axis, time_idx, plot_range)
    xtick_pos = LinRange(round(plot_range*0.1),round_int(plot_range*0.9), 10)
    # Convert x_pos to time
    xtick_pos_label = map(x->round((time_idx + x)/srate, digits=2), xtick_pos)
    axis.xticks = (xtick_pos, ["$(x)s" for x in xtick_pos_label])
end

# Events
"""
Check what events lines are in the current plot range and make them visible.
Make events outside of this range invisible.
"""
function redraw_event_markers(lines, lower_bound, upper_bound)
    # Check for visibility
    vis_events = lower_bound .< event_df[:,:latency] .< upper_bound

    for (event_idx, event) in enumerate(lines[vis_events])
        event.visible = true
    end
    for (event_idx, event) in enumerate(lines[.!vis_events])
        event.visible = false
    end
end

"""
Check what events tags are in the current plot range and make them visible.
Make events outside of this range invisible.
"""
function redraw_event_text(tags, event_df, lower_bound, upper_bound, tag_height)
    # Check for visibility
    vis_events = lower_bound .< event_df[:,:latency] .< upper_bound
    for tag in tags[vis_events]
        # Index of tag in unfiltered array
        idx_unfiltered = findfirst((x)-> x==tag, tags)
        # Observable position doesn't work for text -> manually update position
        corrected_lat = event_df[idx_unfiltered,:latency] - lower_bound
        tag.attributes.attributes[:position][] = (corrected_lat, tag_height)
        tag.visible = true
    end
    for tag in tags[.!vis_events]
        tag.visible = false
    end
end

# Reject regions mark time regions for rejection
"""
Create a new reject region at the passed region if it does not overlap with an exisiting one.
Stores the plot in the reject_plots array and its limits in reject_limits array.
"""
function create_new_reject_region!(axis, new_rect, reject_limits, reject_plots, slider_time, nsample_to_show_obs, lower_bound, upper_bound)
    start_idx = deepcopy(new_rect.origin[1]+max(1, slider_time-nsample_to_show_obs/2))
    end_idx = deepcopy(start_idx + new_rect.widths[1])
    overlaps = false
    for limit in reject_limits
        if limit[1] < start_idx < limit[2] || limit[1] < end_idx < limit[2]
            overlaps = true
        end
    end
    if !overlaps
        append!(reject_limits, [[start_idx, end_idx]])

        plot_start = @lift(max(0, start_idx-$lower_bound))
        plot_end = @lift(min(end_idx - $lower_bound, $upper_bound))
        y_lim = @lift([$(axis.finallimits).origin[2], $(axis.finallimits).widths[2]])
        y_height = @lift(abs($y_lim[1]-$y_lim[2]))
        draw_rect = @lift(FRect($plot_start, $y_lim[1], abs($plot_end-$plot_start), $y_height))
        append!(reject_plots, [draw_reject_region(axis, draw_rect)])
    end
end

"""
Check if reject region is under passed x-position.
Delete region from axis, plot_array, and region limits.
"""
function delete_reject_region(axis, position, reject_limits, reject_plots, slider_time, nsample_to_show_obs)
    for (region_idx, limit) in enumerate(reject_limits)
        if limit[1] < position+max(1, slider_time-nsample_to_show_obs/2) < limit[2]
            deleteat!(reject_limits, region_idx)
            delete!(axis, reject_plots[region_idx])
            deleteat!(reject_plots, region_idx)
        end
    end
end

"""
Draw a transparent rectangle that represents a region to be rejected
"""
function draw_reject_region(axis, rect)
    poly!(axis, rect, raw = true, visible = true, color = RGBAf0(1.0, 0.08, 0.58, 0.3), strokewidth = 0)
end

"""
Check what reject regions are in the current plot range and make them visible.
Make reject regions outside of this range invisible.
"""
function redraw_reject_regions(limits, plots, lower_bound, upper_bound)
    # Create dataframe with lower and upper bounds of the regions
    limits_df = DataFrame( lower=[], upper=[])
    for region in limits
        push!(limits_df, region)
    end
    # Check what regions are visible
    vis_regions = ((lower_bound .< limits_df[:,:lower] .< upper_bound)
                .| (lower_bound .< limits_df[:,:upper] .< upper_bound))

    for region in plots[vis_regions]
        region.visible = true
    end
    for region in plots[.!vis_regions]
        region.visible = false
    end
end

# Channel rejection
"""
Find channel that is closest to mouseposition at click and color it red.
If channel is marked already, unmark it.
"""
function mark_channel(axis, mouse_pos, channel_plots)
    mark_plot = Makie.pick(axis.scene, mouse_pos, 500)[1]
    # Lineplot is an actual EEG plot

    if mark_plot in channel_plots
        # If marked, unmark it
        if mark_plot.attributes.color[] == :red
            mark_plot.attributes.color = :grey0
        else
            mark_plot.attributes.color = :red
        end
    end

    findall(x -> x==mark_plot, channel_plots)
end

# Helper functions
"""
Check if mouse is inside scene.
"""
function is_mouseinside(scene)
    return Vec(scene.events.mouseposition[]) in pixelarea(scene)[]
end

"""
Take a rectangle with origin and widths and return the same rectangle but with
the origin in the lower left corner.
This ensures that widths are positive.
"""
function absrect(rectangle)
    origin = rectangle.origin
    widths = rectangle.widths
    opposite_corner = origin + widths
    new_origin = min.(origin, opposite_corner)
    new_widths = abs.(widths)
    new_rectangle = FRect(new_origin[1], new_origin[2], new_widths[1], new_widths[2])
end

"""
Select a time region in the plot when Left-Shift is hold.
Return the selected region.
"""
function select_rectangle(axis; blocking = false, priority = 2, strokewidth = 3.0, kwargs...)
    scene = axis.scene
    key = Keyboard.left_shift
    waspressed = Node(false)
    rect = Node(FRect(0, 0, 1, 1)) # plotted rectangle
    rect_ret = Node(FRect(0, 0, 1, 1)) # returned rectangle
    y_lim = @lift([$(axis.finallimits).origin[2], $(axis.finallimits).widths[2]])
    y_height = @lift(abs($y_lim[1]-$y_lim[2]))
    # Create an initially hidden rectangle
    plotted_rect = poly!(
        scene, rect, raw = true, visible = false, color = RGBAf0(1.0, 0.08, 0.58, 0.3), strokewidth = 0, transparency = false, kwargs...,
    )

    # Control key press and hold
    on(events(scene).keyboardbutton , priority=priority) do event
        # Start tracking on key press
        if event.key == key
            if event.action == Keyboard.press && is_mouseinside(scene)
                mp = mouseposition(scene)
                waspressed[] = true
                plotted_rect[:visible] = true # start displaying
                rect[] = FRect(mp[1], y_lim[][1], 0.0, y_height[])
                return Consume(blocking)
            end
        end
        # End if key is released
        if !(event.key == key && event.action == Keyboard.repeat)
            if waspressed[] # User has selected the rectangle
                waspressed[] = false
                r = absrect(rect[])
                w, h = widths(r)
                if w > 0.0 && h > 0.0 # Ensure that the rectangle has non0 size.
                    rect_ret[] = r
                end
            end
            # always hide if not the right key is pressed
            plotted_rect[:visible] = false # make the plotted rectangle invisible
            return Consume(blocking)
        end

        return Consume(false)
    end
    # Adapt rectangle depending on mouse position; only if currently tracking.
    on(events(scene).mouseposition, priority=priority) do event
        if waspressed[]
            mp = mouseposition(scene)
            mini = minimum(rect[])
            rect[] = FRect(mini[1], y_lim[][1], mp[1] - mini[1], y_height[])
            return Consume(blocking)
        end
        return Consume(false)
    end

    return rect_ret
end

"""
Disable Makie-native scroll zoom, pan scroll and rectangle zoom.
"""
function disable_native_zoom(axis)
    axis.yzoomlock = true
    axis.xzoomlock = true
    axis.xpanlock = true
    axis.ypanlock = true
    axis.xrectzoom = false
    axis.yrectzoom = false
end

"""
Make channels in mask visible and channels not in it invisible.
"""
function toggle_channel_visibility(channel_plots, mask)
    for (plot_idx, plot) in enumerate(channel_plots)
        if plot_idx in mask
            plot.visible = true
        else
            plot.visible = false
        end
    end
end

"""
Find index of element in array. If element is not in array, return 0.
"""
function findfirst_zero(element, array)
    idx = findfirst(x-> x==element, array)
    if isnothing(idx)
        idx = 0
    end
    idx
end

"""
Round number to neares Integer.
"""
function round_int(x)
    round(Int, x)
end

"""
Check what event regions are in the current plot range and make them visible.
Make event regions outside of this range invisible.
"""
function redraw_event_regions(events, plots, lower_bound, upper_bound)
    # Create dataframe with lower and upper bounds of the regions
    # limits_df = DataFrame( lower=[], upper=[])
    # for event in events
    #     push!(limits_df, region)
    # end
    # Check what regions are visible
    vis_regions = ((lower_bound .< events[:,:latency] .< upper_bound)
                .| (lower_bound .< events[:,:latency] + events[:,:duration] .< upper_bound))

    for region in plots[vis_regions]
        region.visible = true
    end
    for region in plots[.!vis_regions]
        region.visible = false
    end
end

"""
Control a reasonable scale bar size.
"""
function get_scale_factor(scale)
    1/round(scale; digits=1)
end

"""
Plot the data and set up the various interactions.
"""
function data_plot(data,srate,event_df,chanlabels,title, path_to_output, output_file_name)

    # Output file
    io = open(path_to_output * output_file_name, "w")
    reject_chans = []

    # Data size
    nchan = size(data,1)
    nsamples = size(data,2)

    # Parameters to set
    offset = 5.0
    y_estate = offset*nchan
    nsample_to_show = 1000 # only for initial plotting

    tag_height = (nchan+5)*offset
    plot_high = (nchan+10)*offset
    plot_low = -offset
    event_upper_lim = tag_height/plot_high-0.01

    # Set up figure
    fig = Figure(resolution = (2560, 1440))
    axis = fig[1, 1] = Axis(fig, title = title)

    # disable native zoom
    disable_native_zoom(axis)

    # Set up time slider
    time_slider = Slider(fig[2, 1], range = 1:1:nsamples-nsample_to_show, startvalue = 1)
    slider_time = time_slider.value

    # Channel scroll
    channel_slider = IntervalSlider(fig[1, 2], range = 1:1:nchan, startvalues = (1,nchan), horizontal = false)
    channel_interval = channel_slider.interval

    # Time range to plot
    nsample_to_show_obs = Node(nsample_to_show)
    lower_bound = @lift(max(1, round_int($slider_time)-round_int($nsample_to_show_obs/2)))
    upper_bound = @lift(min(nsamples, $lower_bound+round_int($nsample_to_show_obs)))
    range = @lift([$lower_bound,$upper_bound])

    # Scale Observables
    scale_obs = Node(0.4)
    amp_scalable = Node(false)

    # Channel Observables
    offset_obs = Node(offset)

    # Plot EEG
    channel_plots = []
    for chan in 1:nchan
        # Get data chunk currently to be plottet
        data_chunk = @lift(data[chan,$range[1]:$range[2]])
        chunk_mean = @lift(mean($data_chunk))
        chan_pos = @lift(findfirst_zero(chan, ($channel_interval[1]:$channel_interval[2])))
        # Mean center, scale and shift to plot all channels above each other in y-direction
        data_chunk_plot = @lift((($data_chunk .-$chunk_mean)  .* $scale_obs) .+ $chan_pos*$offset_obs)
        append!(channel_plots, [lines!(axis, data_chunk_plot, color = :grey0)])
    end

    # Plot events as vertical lines
    event_lines = map(eachrow(event_df)) do event
        plot_pos = @lift(event.latency-$lower_bound)
        vlines!(axis, plot_pos, ymax=event_upper_lim, visible = false, color = :steelblue1)
    end

    # event_regions = map(eachrow(event_df)) do event
    #     plot_start = @lift(event.latency-$lower_bound)
    #     plot_end = @lift(min(event.latency + event.duration - $lower_bound, $upper_bound))
    #     y_lim = @lift([$(axis.finallimits).origin[2], $(axis.finallimits).widths[2]])
    #     y_height = @lift(abs($y_lim[1]-$y_lim[2]))
    #     draw_rect = @lift(FRect($plot_start, $y_lim[1], abs($plot_end-$plot_start), $y_height))

    #     poly!(axis, draw_rect, raw = true, visible = false, color = :steelblue1, strokewidth = 0)
    # end

    #TW SPECIFIC
    level_1 = ["experimentStart", "longBaselineStart",   "longBaselineWaypoint", "longBaselineTargetpoint", "intertrialBaselineStart", "intertrialBaselineTransitpoint", "intertrialBaselineTargetpoint", "intertrialBaselineStop", "triangleStart", "triangleTaskFirstpoint", "triangleTaskSecondpoint", "triangleTaskHomingtaskStart", "triangleTaskHomingtaskLocationConfirmed", "triangleTaskHomingtaskDistanceerror", "triangleStop", "longBaselineStop", "experimentStop"]
    level_2 = ["rotationStart", "rotationStop", "translationStart", "translationStop","triangleTaskPoleVisible", "triangleTaskPoleSpotted", "longBaselineSphereSpotted", "intertrialBaselineSphereSpotted"]
    level_3 = ["miniBaselineStart", "miniBaselineStop", "subtaskFinished"]
    # TW SPECIFIC

    # Plot event tag above event line
    event_tags = map(eachrow(event_df)) do event
        plot_pos = @lift(event.latency-$lower_bound)
        # TW SPECIFIC
        if event.event ∈ level_1
            text!(axis.scene, event.event, position=(plot_pos,tag_height), offset=(0,30), align=(:center, :center), visible=false, color = :green)
        elseif event.event ∈ level_2
            text!(axis.scene, event.event, position=(plot_pos,tag_height), offset=(0,10), align=(:center, :center), visible=false, color = :orange)
        else
            text!(axis.scene, event.event, position=(plot_pos,tag_height), offset=(0,-10), align=(:center, :center), visible=false, color = :red)
        end
        # TW SPECIFIC
        #text!(axis.scene, event.event, position=(plot_pos,tag_height), align=(:center, :center), visible=false, color = :black)
    end

    # Plot range changes (i.e. zoom or scroll)
    on(range) do value
        redraw_event_markers(event_lines, lower_bound[], upper_bound[])
        redraw_event_text(event_tags, event_df, lower_bound[], upper_bound[], tag_height)
        redraw_reject_regions(reject_limits, reject_plots, lower_bound[], upper_bound[])
        # redraw_event_regions(event_df, event_regions)
        set_time_ticks(axis, lower_bound[], nsample_to_show_obs[])
    end

    # Set initial ticks
    axis.yticks = (offset:offset:(nchan)*offset, chanlabels)
    axis.yticklabelsize = 10.0
    set_time_ticks(axis, 0, nsample_to_show_obs[])

    # Scrolling time/amplitude-zoom
    on(events(fig).scroll, priority = 100) do scroll
        if amp_scalable[]
            scale_obs[] = max(0, scale_obs[] + 0.01*scroll[2])
        else
            nsample_to_show_obs[] = clamp(nsample_to_show_obs[]-floor(scroll[2]*30), 150, nsamples)
        end
    end

    # Enable amplitude-zoom on left-ctrl hold
    on(events(fig).keyboardbutton) do event
        if event.key == Keyboard.left_control && event.action in (Keyboard.repeat, Keyboard.press)
            amp_scalable[] = true
        else
            amp_scalable[] = false
        end
    end

    # Color channel on click
    on(events(fig).mousebutton) do event
        if event.action == Mouse.press && event.button == Mouse.left
            mouse_pos = events(fig).mouseposition[]
            channel_id = mark_channel(axis, mouse_pos, channel_plots)
            if !isempty(channel_id)
                if !(channel_id[1] in reject_chans)
                    append!(reject_chans, channel_id[1])
                else
                    deleteat!(reject_chans, findall(x->x==channel_id[1], reject_chans))
                end
            end
        end
        return Consume(false)
    end

    # Keyboard input
    on(events(fig).keyboardbutton) do event
        if event.action in (Keyboard.press, Keyboard.repeat)
            # Time-scroll on left/right arrow key
            if event.key == Keyboard.left
                set_close_to!(time_slider, max(1, slider_time[]-100))
            end
            if event.key == Keyboard.right
                set_close_to!(time_slider, min(slider_time[]+100, nsamples-nsample_to_show))
            end

            # Close file
            if event.key == Keyboard.p
                for chan_id in reject_chans
                    println(io, chan_id)
                end
                close(io)
            end
        end
        return Consume(false)
    end

    # Mark time windows
    rect = select_rectangle(axis)
    reject_limits = Array{Float32}[]
    reject_plots = []
    on(rect) do new_rect
        create_new_reject_region!(axis, new_rect, reject_limits, reject_plots, slider_time[], nsample_to_show_obs[], lower_bound, upper_bound)
    end

    # Delete time window on right click
    on(events(fig).mousebutton) do event
        if event.button == Mouse.right && event.action == Mouse.press
            mp_x = mouseposition(axis.scene)[1]
            delete_reject_region(axis, mp_x, reject_limits, reject_plots, slider_time[], nsample_to_show_obs[])
        end
    end

    # Time input
    time_input = Textbox(fig[2, 2], width = 100, validator = Float64, placeholder=string(nsample_to_show_obs[]/srate))
    time_value = time_input.stored_string

    on(time_value) do time_range
        new_range = min(parse(Float64, time_range), nsamples/srate)
        nsample_to_show_obs[] = round_int(new_range * srate)
    end

    # Update time zoom variables
    on(nsample_to_show_obs) do nsample
        xlims!(0,nsample)
        time_input.displayed_string = string(nsample/srate)
    end

    # Channel scroll
    on(channel_interval) do interval
        toggle_channel_visibility(channel_plots, (interval[1]:interval[2]))
        offset_obs[] = y_estate/(interval[2]-interval[1]+1)
        new_nchan = interval[2]-interval[1]+1
        new_chanlabels = chanlabels[interval[1]:interval[2]]
        axis.yticks = (offset_obs[]:offset_obs[]:(new_nchan)*offset_obs[], new_chanlabels)
    end

    # Inital event rendering
    redraw_event_markers(event_lines, lower_bound[], upper_bound[])
    redraw_event_text(event_tags, event_df, lower_bound[], upper_bound[], tag_height)

    # Set visible limits (doesn't affect plotted data)
    tightlimits!(axis, Left(), Right()) # No left/right margin
    ylims!(plot_low, plot_high)

    # Scalebar
    scale_x_pos = @lift([round_int($nsample_to_show_obs*0.98)])
    scale_y_pos = [round_int((plot_low + plot_high)/2)]
    scale_size = @lift($scale_obs*100/2) # Divide by 2 because it extends in both directions
    errorbars!(scale_x_pos, scale_y_pos, scale_size, color = :red, linewidth = 3, whiskerwidth = 10)

    # Scalelabel
    label_pos = @lift(($scale_x_pos[1], scale_y_pos[1]))
    label_offset = @lift((-round_int($nsample_to_show_obs*0.02), 0))
    label_text = string(100) * " μV"

    text!(axis.scene, label_text, position = label_pos, color = :red, align=(:right, :center), offset = (-30, 0))

    fig
end
##

subject = 21

path = "/Users/timotheusberg/Documents/Work/BPN/triangle_walking/04_data/3_EEG-preprocessing/s$subject/"

input_file = "s$(subject)_preprocessed.set"

output_file = "s$(subject)_reject.txt"

data,srate,event_df,chanlocs_df,EEG = import_eeglab(path * input_file)

# Test data
time = 1:5000
sin_time = sin.(time./10).*50
test_data = repeat(transpose(sin_time), 128, 1)

chanlabels = chanlocs_df[:, :labels]

##

fig = data_plot(data,srate,event_df,chanlabels,input_file,path, output_file)

##
# PROBLEMS
# - make sure that nsample to display
#   doesn't exceed data bounds and scroll when end is reached

# TODO
# - show amplitude scale
# - read in reject stuff
# - export reject regions

#
# OPTIONAL
# - show mini plot under time-scroll that visualizes position and zoom level


"""
```
IntegrationParameters(;
    method::AbstractIntegrator,
    tmin = 0.0,
    tmax = 1e13,
    dtmin = 1e-3,
    dtmax = Inf,
    abstol = 1e-6,
    reltol = 1e-6,
    maxchange = 1.2,
    exponent = 20.0,
    maxiter = 10,
)
```
Creates [`IntegrationParameters`](@ref).
"""
function IntegrationParameters(;
    method::AbstractIntegrator,
    tmin = 0.0,
    tmax = 1e13,
    dtmin = 1e-3,
    dtmax = Inf,
    abstol = 1e-6,
    reltol = 1e-6,
    maxchange = 1.2,
    exponent = 20.0,
    maxiter = 10,
)
    return IntegrationParameters(
        method,
        tmin,
        tmax,
        dtmin,
        dtmax,
        abstol,
        reltol,
        maxchange,
        exponent,
        maxiter,
    )
end

"""
```
IntegrationTime(; dt = 0.0, time = 0.0, step = 0)
```
Creates [`IntegrationTime`](@ref).
"""
function IntegrationTime(; dt = 0.0, time = 0.0, step = 0)
    return IntegrationTime(dt, time, step)
end


using CompressedStacks

#import CompressedStacks.stack_test
import CompressedStacks.CompressedStack
#import CompressedStacks.access
#import CompressedStacks.run!

#stack_test(70,3)

#original benchmark function
#function benchmark()
    # Any setup code goes here.

    # Run once, to force compilation.
#    srand(666)
#    @time code_to_profile()

    # Run a second time, with profiling.
#    println("\n\n======================= Second run:")
#    srand(666)
#    Profile.init(delay=0.01)
#    clear_malloc_data()
#    @profile @time code_to_profile()

    # Write profile results to profile.bin.
#    r = Profile.retrieve()
#    f = open("profile.bin", "w")
#    serialize(f, r)
#    close(f)
#end

function printsum(a)
    # summary generates a summary of an object
    println(summary(a), ": ", repr(a))
end

function condition_push(stack::CompressedStack, elt::Int)

return true

end

function condition_pop(stack::CompressedStack, elt::Int)

  #println("La lio? ",isempty(stack))
  # This sort of reads the next one?
    if !isempty(stack)
      popRatio= parse(ARGS[3])
      min=parse(ARGS[4])
      max=parse(ARGS[5])
      range=max-min
      topElement=top(stack)

    #flush(STDOUT)
    #  println("should I POP ",aux[1]," ? : ",(parse(aux[1]) < (min+popRatio*range)))
    #flush(STDOUT)
        if (elt==(max+1))
          #println("POPPOP")
          return true
        else
          return ( (elt+topElement) < (min+2*popRatio*range))
        end
    else
        print(stack)
        println("IN this case the compressed stack was empty so I should NOT pop, unless I am comparing the run! of two stacks")
        return false
    end

end

function action_push(stack::CompressedStack{Int,Int}, elt::Int)
  stack.context = Nullable(stack.index)
  # print(stack)
end

function action_pop(stack::CompressedStack, elt::Int)
#  println("Pop element : $elt")
end

#now do the same thing again for the normal stack

function condition_push2(stack::NormalStack, elt::Int)
#  aux = elt

#  popRatio= parse(ARGS[3])
#  min=parse(ARGS[4])
#  max=parse(ARGS[5])
#  range=max-min
#  return (aux > (min+popRatio*range))
  return true
end

function condition_pop2(stack::NormalStack, elt::Int)
#println("La lio? ",isempty(stack))
# This sort of reads the next one?
  if !isempty(stack)
    popRatio= parse(ARGS[3])
    min=parse(ARGS[4])
    max=parse(ARGS[5])
    range=max-min
    topElement=top(stack)

  #flush(STDOUT)
  #  println("should I POP ",aux[1]," ? : ",(parse(aux[1]) < (min+popRatio*range)))
  #flush(STDOUT)
      if (elt==(max+1))
          #println("POPPOP")
          return true
      else
        return ( (elt+topElement) < (min+2*popRatio*range))
      end

  else
      print(stack)
      println("IN this case the normal stack was empty so I should NOT pop, unless I am comparing the run! of two stacks")
      return false
  end

end

function action_push2(stack::NormalStack{Int,Int}, elt::Int)
  stack.context = Nullable(stack.index)
  # print(stack)
end

function action_pop2(stack::NormalStack, elt::Int)
#  println("Pop element : $elt")
end

function randomPushPopTestCompressed(name,size,space)

  println("randomPushPopTest with compressed stack, with Parameters size:",size," space: ",space," name ",name)

  context_type = Int
  data_type = Int

  #ns = NormalStack(name, action_pop2, action_push2, condition_pop2,
  #condition_push2, context_type, data_type)

  #cs = CompressedStack(name, action_pop, action_push, condition_pop,
  #condition_push, context_type, data_type)

  stack = CompressedStack(name, action_pop, action_push, condition_pop, condition_push, context_type, data_type,size,space)
 normalstack = NormalStack(name, action_pop2, action_push2, condition_pop2, condition_push2, context_type, data_type)

#println("Running ")
#try run!(stack,normalstack)
#try  run!(stack)
#try run!(normalstack)
#catch Exception e
#  println("Yagos code caught this exception ",e)
#  println("The stack was ")
#  print(stack)
#end

# seekend(stack.input)
# seekend(normalstack.input)
# println(position(stack.input)==position(normalstack.input))

run!(stack,normalstack)
#run!(stack)

  stack = 0

end

#function randomPushPopTest(repetitions,popRatio,min,max)
function randomPushPopTest(name)

# First let us make a normal stack
#q=Int[]
#println("randomPushPopTest, Parameters ",repetitions," ",popRatio," ",min," ",max," ")
#range=max-min
# Careful, as the rand is of integers this is not really working as a random choice with popRatio parameter
#for i = 1:repetitions
#  (x=rand(min:max)) > (min+popRatio*range) ? push!(q,x) : isempty(q) ? true : pop!(q)
  #println(i)
#end

println("randomPushPopTest with normal stack, with Parameters name:",name)

context_type = Int
data_type = Int

#ns = NormalStack(name, action_pop2, action_push2, condition_pop2,
#condition_push2, context_type, data_type)

#cs = CompressedStack(name, action_pop, action_push, condition_pop,
#condition_push, context_type, data_type)

#stack = CompressedStack(name, action_pop, action_push, condition_pop, condition_push, context_type, data_type,size,space)
normalstack = NormalStack(name, action_pop2, action_push2, condition_pop2, condition_push2, context_type, data_type)

run!(normalstack)
#try run!(normalstack)
#catch Exception e
#  println("normal stack caught this exception ",e)
#  println("The stack was ")
#  print(normalstack)
#end

normalstack = 0

end

#printsum(q)
if (isempty(ARGS) || length(ARGS)<8) println("WRONG PARAMETERS! NEED 5 (repetitions popRatio min max n p testType(0 normal 1 compressed) )")
else

  IOFileName=ARGS[1]
  size=parse(ARGS[2])
  min=parse(ARGS[4])
  max=parse(ARGS[5])
  n=parse(ARGS[6])
  p=parse(ARGS[7])

   println("Testing : FileName ",IOFileName,"\n repetitions: ",size," popRatio  ",ARGS[3]," min ",min,"  max ",max,"  n ",n," p ",p," testType(0 normal 1 compressed) ",ARGS[8])

    # Run once, to force compilation.
  #println("======================= First run:")
  #srand(666)
  #@time parse(ARGS[8])==0 ? randomPushPopTest(10,0.1,1,5) : randomPushPopTestCompressed("/home/yago/.julia/v0.4/CompressedStacks/ioexample/input1")

  # Run a second time, with profiling.
  #println("\n\n======================= Second run:")
  srand(666)

  Profile.init(delay=0.01)
  Profile.clear()

  Profile.clear_malloc_data()
  gc()
  #@profile @time parse(ARGS[8])==0 ? randomPushPopTest(parse(ARGS[2]),parse(ARGS[3]),parse(ARGS[4]),parse(ARGS[5])) : randomPushPopTestCompressed(IOFileName)
#parse(ARGS[8])==0 ? randomPushPopTest(parse(ARGS[2]),parse(ARGS[3]),parse(ARGS[4]),parse(ARGS[5])) : randomPushPopTestCompressed(IOFileName,n,p)
parse(ARGS[8])==0 ? randomPushPopTest( IOFileName ) : randomPushPopTestCompressed(IOFileName,n,p)

  # Write profile results to profile.bin.
#  r = Profile.retrieve()
#  f = open("profile.bin", "w")
#  serialize(f, r)
#  close(f)

end


const maxMantissa = (UInt64(1) << 53) - 1
todouble(sign, exp, mant) = Core.bitcast(Float64, (UInt64(sign) << 63) | (UInt64(exp) << 52) | (UInt64(mant)))

@testset "Parsers.writeshortest" begin

@testset "Float64" begin

@testset "Basic" begin
    @test Parsers.writeshortest(0.0) == "0.0"
    @test Parsers.writeshortest(-0.0) == "-0.0"
    @test Parsers.writeshortest(1.0) == "1.0"
    @test Parsers.writeshortest(-1.0) == "-1.0"
    @test Parsers.writeshortest(NaN) == "NaN"
    @test Parsers.writeshortest(Inf) == "Inf"
    @test Parsers.writeshortest(-Inf) == "-Inf"
end

@testset "SwitchToSubnormal" begin
    @test "2.2250738585072014e-308" == Parsers.writeshortest(2.2250738585072014e-308)
end

@testset "MinAndMax" begin
    @test "1.7976931348623157e308" == Parsers.writeshortest(Core.bitcast(Float64, 0x7fefffffffffffff))
    @test "5.0e-324" == Parsers.writeshortest(Core.bitcast(Float64, Int64(1)))
end

@testset "LotsOfTrailingZeros" begin
    @test "2.9802322387695312e-8" == Parsers.writeshortest(2.98023223876953125e-8)
end

@testset "Regression" begin
    @test "-2.109808898695963e16" == Parsers.writeshortest(-2.109808898695963e16)
    @test "4.940656e-318" == Parsers.writeshortest(4.940656e-318)
    @test "1.18575755e-316" == Parsers.writeshortest(1.18575755e-316)
    @test "2.989102097996e-312" == Parsers.writeshortest(2.989102097996e-312)
    @test "9.0608011534336e15" == Parsers.writeshortest(9.0608011534336e15)
    @test "4.708356024711512e18" == Parsers.writeshortest(4.708356024711512e18)
    @test "9.409340012568248e18" == Parsers.writeshortest(9.409340012568248e18)
    @test "1.2345678" == Parsers.writeshortest(1.2345678)
end

@testset "LooksLikePow5" begin
    # These numbers have a mantissa that is a multiple of the largest power of 5 that fits,
    # and an exponent that causes the computation for q to result in 22, which is a corner
    # case for Ryu.
    @test "5.764607523034235e39" == Parsers.writeshortest(Core.bitcast(Float64, 0x4830F0CF064DD592))
    @test "1.152921504606847e40" == Parsers.writeshortest(Core.bitcast(Float64, 0x4840F0CF064DD592))
    @test "2.305843009213694e40" == Parsers.writeshortest(Core.bitcast(Float64, 0x4850F0CF064DD592))
end

@testset "OutputLength" begin
    @test "1.0" == Parsers.writeshortest(1.0) # already tested in Basic
    @test "1.2" == Parsers.writeshortest(1.2)
    @test "1.23" == Parsers.writeshortest(1.23)
    @test "1.234" == Parsers.writeshortest(1.234)
    @test "1.2345" == Parsers.writeshortest(1.2345)
    @test "1.23456" == Parsers.writeshortest(1.23456)
    @test "1.234567" == Parsers.writeshortest(1.234567)
    @test "1.2345678" == Parsers.writeshortest(1.2345678) # already tested in Regressi
    @test "1.23456789" == Parsers.writeshortest(1.23456789)
    @test "1.234567895" == Parsers.writeshortest(1.234567895) # 1.234567890 would be trimm
    @test "1.2345678901" == Parsers.writeshortest(1.2345678901)
    @test "1.23456789012" == Parsers.writeshortest(1.23456789012)
    @test "1.234567890123" == Parsers.writeshortest(1.234567890123)
    @test "1.2345678901234" == Parsers.writeshortest(1.2345678901234)
    @test "1.23456789012345" == Parsers.writeshortest(1.23456789012345)
    @test "1.234567890123456" == Parsers.writeshortest(1.234567890123456)
    @test "1.2345678901234567" == Parsers.writeshortest(1.2345678901234567)

  # Test 32-bit chunking
    @test "4.294967294" == Parsers.writeshortest(4.294967294) # 2^32 -
    @test "4.294967295" == Parsers.writeshortest(4.294967295) # 2^32 -
    @test "4.294967296" == Parsers.writeshortest(4.294967296) # 2^
    @test "4.294967297" == Parsers.writeshortest(4.294967297) # 2^32 +
    @test "4.294967298" == Parsers.writeshortest(4.294967298) # 2^32 +
end

# Test min, max shift values in shiftright128
@testset "MinMaxShift" begin
    # 32-bit opt-size=0:  49 <= dist <= 50
    # 32-bit opt-size=1:  30 <= dist <= 50
    # 64-bit opt-size=0:  50 <= dist <= 50
    # 64-bit opt-size=1:  30 <= dist <= 50
    @test "1.7800590868057611e-307" == Parsers.writeshortest(todouble(false, 4, 0))
    # 32-bit opt-size=0:  49 <= dist <= 49
    # 32-bit opt-size=1:  28 <= dist <= 49
    # 64-bit opt-size=0:  50 <= dist <= 50
    # 64-bit opt-size=1:  28 <= dist <= 50
    @test "2.8480945388892175e-306" == Parsers.writeshortest(todouble(false, 6, maxMantissa))
    # 32-bit opt-size=0:  52 <= dist <= 53
    # 32-bit opt-size=1:   2 <= dist <= 53
    # 64-bit opt-size=0:  53 <= dist <= 53
    # 64-bit opt-size=1:   2 <= dist <= 53
    @test "2.446494580089078e-296" == Parsers.writeshortest(todouble(false, 41, 0))
    # 32-bit opt-size=0:  52 <= dist <= 52
    # 32-bit opt-size=1:   2 <= dist <= 52
    # 64-bit opt-size=0:  53 <= dist <= 53
    # 64-bit opt-size=1:   2 <= dist <= 53
    @test "4.8929891601781557e-296" == Parsers.writeshortest(todouble(false, 40, maxMantissa))

    # 32-bit opt-size=0:  57 <= dist <= 58
    # 32-bit opt-size=1:  57 <= dist <= 58
    # 64-bit opt-size=0:  58 <= dist <= 58
    # 64-bit opt-size=1:  58 <= dist <= 58
    @test "1.8014398509481984e16" == Parsers.writeshortest(todouble(false, 1077, 0))
    # 32-bit opt-size=0:  57 <= dist <= 57
    # 32-bit opt-size=1:  57 <= dist <= 57
    # 64-bit opt-size=0:  58 <= dist <= 58
    # 64-bit opt-size=1:  58 <= dist <= 58
    @test "3.6028797018963964e16" == Parsers.writeshortest(todouble(false, 1076, maxMantissa))
    # 32-bit opt-size=0:  51 <= dist <= 52
    # 32-bit opt-size=1:  51 <= dist <= 59
    # 64-bit opt-size=0:  52 <= dist <= 52
    # 64-bit opt-size=1:  52 <= dist <= 59
    @test "2.900835519859558e-216" == Parsers.writeshortest(todouble(false, 307, 0))
    # 32-bit opt-size=0:  51 <= dist <= 51
    # 32-bit opt-size=1:  51 <= dist <= 59
    # 64-bit opt-size=0:  52 <= dist <= 52
    # 64-bit opt-size=1:  52 <= dist <= 59
    @test "5.801671039719115e-216" == Parsers.writeshortest(todouble(false, 306, maxMantissa))

    # https:#github.com/ulfjack/ryu/commit/19e44d16d80236f5de25800f56d82606d1be00b9#commitcomment-30146483
    # 32-bit opt-size=0:  49 <= dist <= 49
    # 32-bit opt-size=1:  44 <= dist <= 49
    # 64-bit opt-size=0:  50 <= dist <= 50
    # 64-bit opt-size=1:  44 <= dist <= 50
    @test "3.196104012172126e-27" == Parsers.writeshortest(todouble(false, 934, 0x000FA7161A4D6E0C))
end

@testset "SmallIntegers" begin
    @test "9.007199254740991e15" == Parsers.writeshortest(9007199254740991.0)
    @test "9.007199254740992e15" == Parsers.writeshortest(9007199254740992.0)

    @test "1.0" == Parsers.writeshortest(1.0e+0)
    @test "12.0" == Parsers.writeshortest(1.2e+1)
    @test "123.0" == Parsers.writeshortest(1.23e+2)
    @test "1234.0" == Parsers.writeshortest(1.234e+3)
    @test "12345.0" == Parsers.writeshortest(1.2345e+4)
    @test "123456.0" == Parsers.writeshortest(1.23456e+5)
    @test "1.234567e6" == Parsers.writeshortest(1.234567e+6)
    @test "1.2345678e7" == Parsers.writeshortest(1.2345678e+7)
    @test "1.23456789e8" == Parsers.writeshortest(1.23456789e+8)
    @test "1.23456789e9" == Parsers.writeshortest(1.23456789e+9)
    @test "1.234567895e9" == Parsers.writeshortest(1.234567895e+9)
    @test "1.2345678901e10" == Parsers.writeshortest(1.2345678901e+10)
    @test "1.23456789012e11" == Parsers.writeshortest(1.23456789012e+11)
    @test "1.234567890123e12" == Parsers.writeshortest(1.234567890123e+12)
    @test "1.2345678901234e13" == Parsers.writeshortest(1.2345678901234e+13)
    @test "1.23456789012345e14" == Parsers.writeshortest(1.23456789012345e+14)
    @test "1.234567890123456e15" == Parsers.writeshortest(1.234567890123456e+15)

  # 10^i
    @test "1.0" == Parsers.writeshortest(1.0e+0)
    @test "10.0" == Parsers.writeshortest(1.0e+1)
    @test "100.0" == Parsers.writeshortest(1.0e+2)
    @test "1000.0" == Parsers.writeshortest(1.0e+3)
    @test "10000.0" == Parsers.writeshortest(1.0e+4)
    @test "100000.0" == Parsers.writeshortest(1.0e+5)
    @test "1.0e6" == Parsers.writeshortest(1.0e+6)
    @test "1.0e7" == Parsers.writeshortest(1.0e+7)
    @test "1.0e8" == Parsers.writeshortest(1.0e+8)
    @test "1.0e9" == Parsers.writeshortest(1.0e+9)
    @test "1.0e10" == Parsers.writeshortest(1.0e+10)
    @test "1.0e11" == Parsers.writeshortest(1.0e+11)
    @test "1.0e12" == Parsers.writeshortest(1.0e+12)
    @test "1.0e13" == Parsers.writeshortest(1.0e+13)
    @test "1.0e14" == Parsers.writeshortest(1.0e+14)
    @test "1.0e15" == Parsers.writeshortest(1.0e+15)

  # 10^15 + 10^i
    @test "1.000000000000001e15" == Parsers.writeshortest(1.0e+15 + 1.0e+0)
    @test "1.00000000000001e15" == Parsers.writeshortest(1.0e+15 + 1.0e+1)
    @test "1.0000000000001e15" == Parsers.writeshortest(1.0e+15 + 1.0e+2)
    @test "1.000000000001e15" == Parsers.writeshortest(1.0e+15 + 1.0e+3)
    @test "1.00000000001e15" == Parsers.writeshortest(1.0e+15 + 1.0e+4)
    @test "1.0000000001e15" == Parsers.writeshortest(1.0e+15 + 1.0e+5)
    @test "1.000000001e15" == Parsers.writeshortest(1.0e+15 + 1.0e+6)
    @test "1.00000001e15" == Parsers.writeshortest(1.0e+15 + 1.0e+7)
    @test "1.0000001e15" == Parsers.writeshortest(1.0e+15 + 1.0e+8)
    @test "1.000001e15" == Parsers.writeshortest(1.0e+15 + 1.0e+9)
    @test "1.00001e15" == Parsers.writeshortest(1.0e+15 + 1.0e+10)
    @test "1.0001e15" == Parsers.writeshortest(1.0e+15 + 1.0e+11)
    @test "1.001e15" == Parsers.writeshortest(1.0e+15 + 1.0e+12)
    @test "1.01e15" == Parsers.writeshortest(1.0e+15 + 1.0e+13)
    @test "1.1e15" == Parsers.writeshortest(1.0e+15 + 1.0e+14)

  # Largest power of 2 <= 10^(i+1)
    @test "8.0" == Parsers.writeshortest(8.0)
    @test "64.0" == Parsers.writeshortest(64.0)
    @test "512.0" == Parsers.writeshortest(512.0)
    @test "8192.0" == Parsers.writeshortest(8192.0)
    @test "65536.0" == Parsers.writeshortest(65536.0)
    @test "524288.0" == Parsers.writeshortest(524288.0)
    @test "8.388608e6" == Parsers.writeshortest(8388608.0)
    @test "6.7108864e7" == Parsers.writeshortest(67108864.0)
    @test "5.36870912e8" == Parsers.writeshortest(536870912.0)
    @test "8.589934592e9" == Parsers.writeshortest(8589934592.0)
    @test "6.8719476736e10" == Parsers.writeshortest(68719476736.0)
    @test "5.49755813888e11" == Parsers.writeshortest(549755813888.0)
    @test "8.796093022208e12" == Parsers.writeshortest(8796093022208.0)
    @test "7.0368744177664e13" == Parsers.writeshortest(70368744177664.0)
    @test "5.62949953421312e14" == Parsers.writeshortest(562949953421312.0)
    @test "9.007199254740992e15" == Parsers.writeshortest(9007199254740992.0)

  # 1000 * (Largest power of 2 <= 10^(i+1))
    @test "8000.0" == Parsers.writeshortest(8.0e+3)
    @test "64000.0" == Parsers.writeshortest(64.0e+3)
    @test "512000.0" == Parsers.writeshortest(512.0e+3)
    @test "8.192e6" == Parsers.writeshortest(8192.0e+3)
    @test "6.5536e7" == Parsers.writeshortest(65536.0e+3)
    @test "5.24288e8" == Parsers.writeshortest(524288.0e+3)
    @test "8.388608e9" == Parsers.writeshortest(8388608.0e+3)
    @test "6.7108864e10" == Parsers.writeshortest(67108864.0e+3)
    @test "5.36870912e11" == Parsers.writeshortest(536870912.0e+3)
    @test "8.589934592e12" == Parsers.writeshortest(8589934592.0e+3)
    @test "6.8719476736e13" == Parsers.writeshortest(68719476736.0e+3)
    @test "5.49755813888e14" == Parsers.writeshortest(549755813888.0e+3)
    @test "8.796093022208e15" == Parsers.writeshortest(8796093022208.0e+3)
end

end # Float64

@testset "Float32" begin

@testset "Basic" begin
    @test "0.0" == Parsers.writeshortest(Float32(0.0))
    @test "-0.0" == Parsers.writeshortest(Float32(-0.0))
    @test "1.0" == Parsers.writeshortest(Float32(1.0))
    @test "-1.0" == Parsers.writeshortest(Float32(-1.0))
    @test "NaN" == Parsers.writeshortest(Float32(NaN))
    @test "Inf" == Parsers.writeshortest(Float32(Inf))
    @test "-Inf" == Parsers.writeshortest(Float32(-Inf))
end

@testset "SwitchToSubnormal" begin
    @test "1.1754944e-38" == Parsers.writeshortest(1.1754944f-38)
end

@testset "MinAndMax" begin
    @test "3.4028235e38" == Parsers.writeshortest(Core.bitcast(Float32, 0x7f7fffff))
    @test "1.0e-45" == Parsers.writeshortest(Core.bitcast(Float32, Int32(1)))
end

# Check that we return the exact boundary if it is the shortest
# representation, but only if the original floating point number is even.
@testset "BoundaryRoundeven" begin
    @test "3.355445e7" == Parsers.writeshortest(3.355445f7)
    @test "9.0e9" == Parsers.writeshortest(8.999999f9)
    @test "3.436672e10" == Parsers.writeshortest(3.4366717f10)
end

# If the exact value is exactly halfway between two shortest representations,
# then we round to even. It seems like this only makes a difference if the
# last two digits are ...2|5 or ...7|5, and we cut off the 5.
@testset "exactValueRoundeven" begin
    @test "305404.12" == Parsers.writeshortest(3.0540412f5)
    @test "8099.0312" == Parsers.writeshortest(8.0990312f3)
end

@testset "LotsOfTrailingZeros" begin
    # Pattern for the first test: 00111001100000000000000000000000
    @test "0.00024414062" == Parsers.writeshortest(2.4414062f-4)
    @test "0.0024414062" == Parsers.writeshortest(2.4414062f-3)
    @test "0.0043945312" == Parsers.writeshortest(4.3945312f-3)
    @test "0.0063476562" == Parsers.writeshortest(6.3476562f-3)
end

@testset "Regression" begin
    @test "4.7223665e21" == Parsers.writeshortest(4.7223665f21)
    @test "8.388608e6" == Parsers.writeshortest(8388608f0)
    @test "1.6777216e7" == Parsers.writeshortest(1.6777216f7)
    @test "3.3554436e7" == Parsers.writeshortest(3.3554436f7)
    @test "6.7131496e7" == Parsers.writeshortest(6.7131496f7)
    @test "1.9310392e-38" == Parsers.writeshortest(1.9310392f-38)
    @test "-2.47e-43" == Parsers.writeshortest(-2.47f-43)
    @test "1.993244e-38" == Parsers.writeshortest(1.993244f-38)
    @test "4103.9004" == Parsers.writeshortest(4103.9003f0)
    @test "5.3399997e9" == Parsers.writeshortest(5.3399997f9)
    @test "6.0898e-39" == Parsers.writeshortest(6.0898f-39)
    @test "0.0010310042" == Parsers.writeshortest(0.0010310042f0)
    @test "2.882326e17" == Parsers.writeshortest(2.8823261f17)
    @test "7.038531e-26" == Parsers.writeshortest(7.0385309f-26)
    @test "9.223404e17" == Parsers.writeshortest(9.2234038f17)
    @test "6.710887e7" == Parsers.writeshortest(6.7108872f7)
    @test "1.0e-44" == Parsers.writeshortest(1.0f-44)
    @test "2.816025e14" == Parsers.writeshortest(2.816025f14)
    @test "9.223372e18" == Parsers.writeshortest(9.223372f18)
    @test "1.5846086e29" == Parsers.writeshortest(1.5846085f29)
    @test "1.1811161e19" == Parsers.writeshortest(1.1811161f19)
    @test "5.368709e18" == Parsers.writeshortest(5.368709f18)
    @test "4.6143166e18" == Parsers.writeshortest(4.6143165f18)
    @test "0.007812537" == Parsers.writeshortest(0.007812537f0)
    @test "1.0e-45" == Parsers.writeshortest(1.4f-45)
    @test "1.18697725e20" == Parsers.writeshortest(1.18697724f20)
    @test "1.00014165e-36" == Parsers.writeshortest(1.00014165f-36)
    @test "200.0" == Parsers.writeshortest(200f0)
    @test "3.3554432e7" == Parsers.writeshortest(3.3554432f7)
end

@testset "LooksLikePow5" begin
    # These numbers have a mantissa that is the largest power of 5 that fits,
    # and an exponent that causes the computation for q to result in 10, which is a corner
    # case for Ryu.
    @test "6.7108864e17" == Parsers.writeshortest(Core.bitcast(Float32, 0x5D1502F9))
    @test "1.3421773e18" == Parsers.writeshortest(Core.bitcast(Float32, 0x5D9502F9))
    @test "2.6843546e18" == Parsers.writeshortest(Core.bitcast(Float32, 0x5E1502F9))
end

@testset "OutputLength" begin
    @test "1.0" == Parsers.writeshortest(Float32(1.0))
    @test "1.2" == Parsers.writeshortest(Float32(1.2))
    @test "1.23" == Parsers.writeshortest(Float32(1.23))
    @test "1.234" == Parsers.writeshortest(Float32(1.234))
    @test "1.2345" == Parsers.writeshortest(Float32(1.2345))
    @test "1.23456" == Parsers.writeshortest(Float32(1.23456))
    @test "1.234567" == Parsers.writeshortest(Float32(1.234567))
    @test "1.2345678" == Parsers.writeshortest(Float32(1.2345678))
    @test "1.23456735e-36" == Parsers.writeshortest(Float32(1.23456735e-36))
end

end # Float32

@testset "Float16" begin

@testset "Basic" begin
    @test "0.0" == Parsers.writeshortest(Float16(0.0))
    @test "-0.0" == Parsers.writeshortest(Float16(-0.0))
    @test "1.0" == Parsers.writeshortest(Float16(1.0))
    @test "-1.0" == Parsers.writeshortest(Float16(-1.0))
    @test "NaN" == Parsers.writeshortest(Float16(NaN))
    @test "Inf" == Parsers.writeshortest(Float16(Inf))
    @test "-Inf" == Parsers.writeshortest(Float16(-Inf))
end

let x=floatmin(Float16)
    while x <= floatmax(Float16)
        @test parse(Float16, Parsers.writeshortest(x)) == x
        x = nextfloat(x)
    end
end

# function testfloats(T)
#     x = floatmin(T)
#     i = 0
#     fails = 0
#     success = 0
#     while x < floatmax(T)
#         test = parse(T, Parsers.writeshortest(x)) == x
#         if !test

#             fails += 1
#         else
#             success += 1
#         end
#         x = nextfloat(x)
#         i += 1

#     end
#     return fails / (fails + success)
# end

end # Float16

end


using CSV, DataFrames, Plots, StatsPlots, Missings, CategoricalArrays
using Flux
using MLDataUtils
using PyCall

data = CSV.read("julia/cleaned.csv")

### Analyze the participants by career stage
career_group = groupby(data, :Position);
career_breakdown = sort(combine(career_group, nrow), :x1);
career_breakdown[!, :Index] = [1, 3, 4, 2];
sort!(career_breakdown, :Index);

career_breakdown.Position

# plot the data
gr(size=(600, 500))
chart = bar(
    career_breakdown[!, :x1],
    color = "#b85231",
    leg = false,
    lw = 1.0,
    bar_width = 0.3,
    background_color = "#f9efde",
    titlefont = font(14, "Open Sans"),
)
title!("Participant breakdown")
ylabel!("Counts", font = font(12, "Open Sans"))
xlabel!("Career stage", font = font(12, "Open Sans"))
xticks!([1, 2, 3, 4], career_breakdown[!, :Position])
savefig("julia/career_plot.png")

# Generate a corner plot for everything
mat_rep = Matrix(data[[:Lat, :Long, :HotDogEncoding, :PosEncoding]]);

gr(size = (600, 500))
corrplot(
    mat_rep,
    label = ["Latitude", "Longitude", "Hot Dog?", "Career Step"],
    grid = false,
    compact = true,
)

# Get raw stats out; how likely each career stage is to be affirmative
true_stats = combine(career_group, :HotDogEncoding => sum);
true_stats[!,:nentries] = combine(career_group, nrow)[:x1]
true_stats[!,:percent] = true_stats[:HotDogEncoding_sum] ./ true_stats[:nentries];
true_stats[!,:Index] = [2,4,3,1];
sort!(true_stats, :Index)

gr(size=(600, 500))
chart = bar(
    true_stats[!, :percent],
    color = "#b85231",
    leg = false,
    lw = 1.0,
    bar_width = 0.3,
    background_color = "#f9efde",
    titlefont = font(14, "Open Sans"),
)
hline!([1.], lw=1.5)
title!("Sandwich vs. Career stage")
ylabel!("Fraction of Yes", font = font(12, "Open Sans"))
xlabel!("Career stage", font = font(12, "Open Sans"))
xticks!([1, 2, 3, 4], career_breakdown[!, :Position])
annotate!(1.3, 0.95, "Correct answer", font=font(1, "Open Sans"))
savefig("julia/career_sandwich.png")

for col in [:Position, :Hotdog]
    categorical!(data, col);
end

gr(size=(600, 500))
@df data violin(:Position, :HotDogEncoding, alpha=0.7, background_color = "#f9efde", color="#b85231", marker=(0.2,"#b85231",stroke(0)), leg=false)
@df data dotplot!(:Position, :HotDogEncoding, marker=(:black,stroke(0)), leg=false, alpha=0.4)
yticks!([0., 1.], ["No", "Yes"])
ylabel!("Response")
savefig("julia/career_sandwich.png")

yaynay = combine(groupby(data, :HotDogEncoding), nrow);

gr(size=(600, 500))
chart = bar(
    yaynay[!, :x1],
    color = "#b85231",
    leg = false,
    lw = 1.0,
    bar_width = 0.3,
    background_color = "#f9efde",
    titlefont = font(14, "Open Sans"),
)
title!("Are hotdogs sandwiches?")
ylabel!("Counts", font = font(12, "Open Sans"))
xlabel!("Response", font = font(12, "Open Sans"))
xticks!([1, 2], ["Yes", "No"])
annotate!(1., 70., "Correct answer", font=font(1, "Open Sans"))
savefig("julia/hotdog_responses.png")

US = filter(row -> row.Long < -20., data)
EU = filter(row -> row.Long > -20., data)

US_yay = sum(US.HotDogEncoding) / nrow(US)
EU_yay = sum(EU.HotDogEncoding) / nrow(EU)

df = DataFrame([[US_yay, EU_yay]])

gr(size=(600, 500))
chart = bar(
    df[:,:x1],
    color = "#b85231",
    leg = false,
    lw = 1.0,
    bar_width = 0.3,
    background_color = "#f9efde",
    titlefont = font(14, "Open Sans"),
)
# title!("Europeans more likely to be right")
ylabel!("Fraction of Yes", font = font(12, "Open Sans"))
xlabel!("Side of the World", font = font(12, "Open Sans"))
xticks!([1, 2], ["Negative longitude", "Positive longitude"])
annotate!(1.3, 0.95, "Correct answer", font=font(1, "Open Sans"))
hline!([1.], lw=1.5)
savefig("julia/split_world.png")

length(unique(data.Org))


export data

abstract type TensorStorage end

#
# Generic ITensor storage functions
#

storage_randn!(S::TensorStorage) = randn!(data(S))
storage_norm(S::TensorStorage) = norm(data(S))
storage_conj(S::T) where {T<:TensorStorage}= T(conj(data(S)))


module MinHash

using DataStructures

include("hashset.jl")
include("minhash.jl")
include("distance.jl")

export MinHasher,
update!,
MinHashSketch,
minhash,
intersectionlength

end


@testset "solve_burgers.jl" begin
    @test_broken solve()
end


using Test

using FlightMechanicsSimulator

@testset "ISA1976" begin
    # Check constructor
    atmosphere = ISA1976(0.0)

    # Check return type is the expected one.
    @test isa(atmosphere, ISA1976)
    # Check inheritance.
    @test isa(atmosphere, FlightMechanicsSimulator.Atmosphere)

    # Check getters
    @test isapprox(get_temperature(atmosphere), 288.15)
    @test isapprox(get_pressure(atmosphere), 101325)
    @test isapprox(get_density(atmosphere), 1.225)
    @test isapprox(get_sound_velocity(atmosphere), 340.293988026089)
    @test isapprox(get_height(atmosphere), 0.0)

    # Check against other constructor
    atmosphere2 = ISA1976(
        atmosphere.temperature,
        atmosphere.pressure,
        atmosphere.density,
        atmosphere.sound_velocity,
        atmosphere.height,
    )
    @test atmosphere2 === atmosphere
end

@testset "F16StevensAtmosphere" begin
    # Check constructor
    atmosphere = F16StevensAtmosphere(0.0)

    # Check return type is the expected one.
    @test isa(atmosphere, F16StevensAtmosphere)
    # Check inheritance.
    @test isa(atmosphere, FlightMechanicsSimulator.Atmosphere)

    # Check getters
    @test isapprox(get_temperature(atmosphere), 288, atol=1)
    @test isapprox(get_pressure(atmosphere), 101325, atol=100)
    @test isapprox(get_density(atmosphere), 1.225, atol=0.001)
    @test isapprox(get_sound_velocity(atmosphere), 340, atol=0.5)
    @test isapprox(get_height(atmosphere), 0.0)

    # Check against other constructor
    atmosphere2 = F16StevensAtmosphere(
        atmosphere.temperature,
        atmosphere.pressure,
        atmosphere.density,
        atmosphere.sound_velocity,
        atmosphere.height,
    )
    @test atmosphere2 === atmosphere
end


export Chain

"""
    struct Chain{I, J, K, O} <: Conditional{I, J, K, O, Nothing, Q, SFunc{J, O, Nothing}}

A `Conditional` that chains its input `I` through a given function that returns an
`SFunc{J,O}`.
"""
struct Chain{I, J, K, O} <: Conditional{I, J, K, O, SFunc{J, O}}
    fn::Function
    """
    function Chain(I, J, O, fn)

    Chain an input through `fn`.

    The chain is an `SFunc{K,O}`, where `K` is the concatenation of tuples `I` and `J`.
    `fn` is a function that takes an argument of type `I` and returns an `SFunc{J,O}`.
    The `Chain` defines a generative conditional distribution as follows:
    - Given inputs `i` and `j`
    - Let `s::SFunc{J,O} = fn(i)`
    - Use `j` to generate a value from `s`

    For the common case, `Chain` has a special constructors where J is empty.
    """
    function Chain(I, J, O, fn)
        K = extend_tuple_type(I,J)
        new{I, J, K, O}(fn)
    end
    Chain(I, O, fn) = Chain(I, Tuple{}, O, fn)
end

function gensf(ch::Chain{I, J, K, O}, i::I) where {I, J, K, O}
    ch.fn(i)
end


"""
    Submodel{iip, F, NS, NI, NP, JR, JS, JI, JP, GT, IS, II, IP, OS, OI, OP}

Base type for submodels in `Aeroelasticity.jl`.
"""
struct Submodel{iip, F, NS, NI, NP, C, JR, JS, JI, JP, GT, IS, II, IP, OS, OI, OP}
    # residual function
    f::F
    # dimensions
    nx::NS
    ny::NI
    np::NP
    # constants
    constants::C
    # jacobians/gradients
    ratejac::JR
    statejac::JS
    inputjac::JI
    paramjac::JP
    tgrad::GT
    # input/output functions
    setstate::IS
    setinput::II
    setparam::IP
    sepstate::OS
    sepinput::OI
    sepparam::OP
end

"""
    Submodel{iip}(f, nx, ny, np; kwargs...)

Define a custom submodel for use with `Aeroelasticity.jl`.  Note that `iip` should be set to
`true` if `f` is defined in-place and `false` otherwise.

# Arguments
 - `f`: Residual function (`resid = f(dx, x, y, p, t)` or `f(resid, dx, x, y, p, t)`)
 - `nx`: Number of states.
 - `ny`: Number of inputs.
 - `np`: Number of parameters.

# Keyword Arguments
 - `ratejac = Nonlinear()`:  Rate jacobian description (of type [`AbstractJacobian`](@ref))
 - `statejac = Nonlinear()`: State jacobian description (of type [`AbstractJacobian`](@ref))
 - `inputjac = Nonlinear()`: Input jacobian description (of type [`AbstractJacobian`](@ref))
 - `paramjac = Nonlinear()`: Parameter jacobian description (of type [`AbstractJacobian`](@ref))
 - `tgrad = Nonlinear()`: Time gradient description (of type [`AbstractGradient`](@ref))
 - `setstate = (states; x) -> states .= x`: Function for setting rates/states using a
    custom set of keyword arguments.
 - `setinput = (inputs; y) -> inputs .= y`: Function for setting inputs using a custom set
    of keyword arguments.
 - `setparam = (params; p) -> params .= p`: Function for setting parameters using a custom
    set of keyword arguments.
 - `sepstate = (x) -> (x = x,)`: Function for separating rates/states into named variables.
 - `sepinput = (y) -> (y = y,)`: Function for separating inputs into named variables.
 - `sepparam = (p) -> (p = p,)`: Function for separating parameters into named variables.
"""
function Submodel{iip}(f, nx, ny, np;
    constants = (),
    ratejac = Nonlinear(),
    statejac = Nonlinear(),
    inputjac = Nonlinear(),
    paramjac = Nonlinear(),
    tgrad = Nonlinear(),
    setstate = (states; x) -> states .= x,
    setinput = (inputs; y) -> inputs .= y,
    setparam = (params; p) -> params .= p,
    sepstate = (x) -> (x = x,),
    sepinput = (y) -> (y = y,),
    sepparam = (p) -> (p = p,)) where iip

    # add jacobian definitions, if necessary
    ratejac = add_rate_jacobian(ratejac, f, iip, nx, ny, np)
    statejac = add_state_jacobian(statejac, f, iip, nx, ny, np)
    inputjac = add_input_jacobian(inputjac, f, iip, nx, ny, np)
    paramjac = add_parameter_jacobian(paramjac, f, iip, nx, ny, np)
    tgrad = add_time_gradient(tgrad, f, iip, nx, ny, np)

    return Submodel{iip,
        typeof(f), typeof(nx), typeof(ny), typeof(np), typeof(constants),
        typeof(ratejac), typeof(statejac), typeof(inputjac), typeof(paramjac), typeof(tgrad),
        typeof(setstate), typeof(setinput), typeof(setparam),
        typeof(sepstate), typeof(sepinput), typeof(sepparam)}(
            f, nx, ny, np, constants,
            ratejac, statejac, inputjac, paramjac, tgrad,
            setstate, setinput, setparam,
            sepstate, sepinput, sepparam)
end

"""
    Submodel(np; kwargs...)

Define a submodel with no state variables (see [`Submodel`](@ref))

# Arguments
 - `np`: Number of parameters.

# Keyword Arguments
 - `constants = ()`: Constants associated with this model
 - `setparam = (params; p) -> params .= p`: Function for setting parameters using a custom
    set of keyword arguments.
 - `sepparam = (p) -> (p = p,)`: Function for separating parameters into named variables.
"""
function Submodel(np;
    constants = (),
    setparam = (params; p) -> params .= p,
    sepparam = (p) -> (p = p,))

    return Submodel{false}((dx, x, y, p, t) -> Float64[], Val(0), Val(0), np;
        constants = constants,
        ratejac = Empty(),
        statejac = Empty(),
        inputjac = Empty(),
        paramjac = Empty(),
        tgrad = Empty(),
        setstate = (x) -> x,
        setinput = (y) -> y,
        setparam = setparam,
        sepstate = (x) -> (),
        sepinput = (y) -> (),
        sepparam = sepparam)
end

"""
    Coupling{iip, G, NX, NY, NP, NPC, JR, JS, JP, GT, IP, OP}

Base type for coupling models in Aeroelasticity.jl.
"""
struct Coupling{iip, G, NX, NY, NP, NPC, JR, JS, JP, GT, IP, OP}
    # coupling function
    g::G
    # dimensions
    nx::NX
    ny::NY
    np::NP
    npc::NPC
    # jacobians/gradients
    ratejac::JR
    statejac::JS
    paramjac::JP
    tgrad::GT
    # input/output functions
    setparam::IP
    sepparam::OP
end

"""
    Coupling{iip}(g, nx, ny, np, npc; kwargs...)

Define a custom coupling model for use with Aeroelasticity.jl.  Note that `iip` should
be set to `true` if `g` is defined in-place and `false` otherwise.

# Arguments
 - `g`: Coupling function (`y = g(dx, x, p, t)` or `g(y, dx, x, p, t)`)
 - `nx`: Number of states.
 - `ny`: Number of inputs.
 - `np`: Number of parameters (including parameters introduced by the coupling).
 - `npc`: Number of additional parameters introduced by the coupling.

# Keyword Arguments
 - `ratejac = Nonlinear()`:  Rate jacobian description (of type [`AbstractJacobian`](@ref))
 - `statejac = Nonlinear()`: State jacobian description (of type [`AbstractJacobian`](@ref))
 - `paramjac = Nonlinear()`: Parameter jacobian description (of type [`AbstractJacobian`](@ref))
 - `tgrad = Nonlinear()`: Time gradient description (of type [`AbstractGradient`](@ref))
 - `setparam = (params; p) -> params .= p`: Function for setting parameters introduced by
    the coupling using a custom set of keyword arguments.
 - `sepparam = (p) -> (p = p,)`: Function for separating parameters introduced by the
    coupling into named variables.
"""
function Coupling{iip}(g, nx, ny, np, npc;
    ratejac = Nonlinear(),
    statejac = Nonlinear(),
    paramjac = Nonlinear(),
    tgrad = Nonlinear(),
    setparam = (params; p) -> params .= p,
    sepparam = (p) -> (p = p,)) where {iip}

    # add jacobian definitions, if necessary
    ratejac = add_coupling_rate_jacobian(ratejac, g, iip, nx, ny, np)
    statejac = add_coupling_state_jacobian(statejac, g, iip, nx, ny, np)
    paramjac = add_coupling_parameter_jacobian(paramjac, g, iip, nx, ny, np)
    tgrad = add_coupling_time_gradient(tgrad, g, iip, nx, ny, np)

    return Coupling{iip,
        typeof(g), typeof(nx), typeof(ny), typeof(np), typeof(npc),
        typeof(ratejac), typeof(statejac), typeof(paramjac), typeof(tgrad),
        typeof(setparam), typeof(sepparam)}(
            g, nx, ny, np, npc,
            ratejac, statejac, paramjac, tgrad,
            setparam, sepparam)
end

"""
    CoupledModel{M,C}

Base type for coupled models in Aeroelasticity.jl.
"""
struct CoupledModel{M,C}
    submodels::M
    coupling::C
end

"""
    CoupledModel(submodels, coupling)

Define a coupled model given a tuple of submodels (of type [`Submodel`](@ref)) and a
coupling model (of type [`Coupling`](@ref)).
"""
CoupledModel(submodels, coupling)

# --- Inplaceness Trait --- #

abstract type InPlaceness end
struct InPlace <: InPlaceness end
struct OutOfPlace <: InPlaceness end

inplaceness(::Submodel{true}) = InPlace()
inplaceness(::Submodel{false}) = OutOfPlace()

inplaceness(::Coupling{true}) = InPlace()
inplaceness(::Coupling{false}) = OutOfPlace()

function inplaceness(model::CoupledModel)
    if any(isinplace.(model.submodels)) || isinplace(model.coupling)
        return InPlace()
    else
        return OutOfPlace()
    end
end

isinplace(x) = isinplace(inplaceness(x))
isinplace(::InPlace) = true
isinplace(::OutOfPlace) = false