# SparseArrays.dropstored!
size(env::NatureEnv) = env.world_size[1:2]

function NatureEnv12(;
        num_starting_players=2,
        world_size=(32, 32),
        window=3,
        episode_len=100,
        num_frames=4,
        player_starting_food=5f0,
        food_generators=[
            FoodGen([5,5],[1,1]),
            FoodGen([25,25],[1,1]),
           ],
        vocab_size=0) #width, heigh, channels, batch

    num_food_types = length(food_generators)
    # self, other, x pos, y pos, food_poses, food_counts
    num_channels = num_frames * (2+2+num_food_types+num_food_types+vocab_size)

    observation_space = Space(
        ClosedInterval.(
            fill(typemin(Int), 2window+1, 2window+1, num_channels),
            fill(typemax(Int), 2window+1, 2window+1, num_channels)
        )
    )
    place_records = bigger_dd_builder()

    NatureEnv12{num_food_types}(
        1,
        episode_len,
        num_starting_players,
        [],
        num_food_types,
        Vector{Matrix{Float32}}(),
        food_generators,
        player_starting_food,
        world_size,
        window,
        (2*window+1, 2*window+1, num_channels),
        observation_space,
        place_records,
        zeros(Float32, num_food_types),
        vocab_size,
        [[] for _ in 1:episode_len],
        num_frames,
        [[] for _ in 1:num_starting_players]
       )
end

RLBase.state_space(env::NatureEnv,::Observation{Any}, players::Dict{Int,Player}) = Dict(p=>state_space(env, p) for p in keys(players))
RLBase.state(env::NatureEnv,::Observation{Any}, players::Vector{Int}) = Dict(p=>state(env, p) for p in players)
RLBase.action_space(env::NatureEnv, players::Dict{Int,Player}) = Dict(p=>action_space(env, p) for p in keys(players))
RLBase.reward(env::NatureEnv, players::Dict{Int,Player}) = Dict(p=>reward(env, p) for p in keys(players))
RLBase.reward(env::NatureEnv, players::Vector{Int}) = Dict(p=>reward(env, p) for p in players)

RLBase.state_space(env::NatureEnv,::Observation{Any}, player::Int) = env.observation_space
RLBase.action_space(env::NatureEnv, player::Int) = Base.OneTo(4 + 2env.food_types + env.vocab_size)

function RLBase.reward(env::NatureEnv, player::Int)
    if all(env.players[player].food_counts .> 0f0) && !env.players[player].dead
        return 1f0
    else
        return 0f0
    end
end

RLBase.is_terminated(env::NatureEnv, player::Int) = env.players[player].dead
RLBase.is_terminated(env::NatureEnv, players::Vector{Int}) = Dict(p=>is_terminated(env, p) for p in players)
function RLBase.is_terminated(env::NatureEnv)
    if all(is_terminated(env, p) for p in keys(env.players)) ||
        env.step >= env.episode_len
        return true
    else
        return false
    end
end
function RLBase.reset!(env::NatureEnv)
    # Add Food according to each generator
    env.step=1
    env.food_frames = []
    for type in 1:env.food_types
        push!(env.food_frames, zeros(Float32, size(env)...))
        for nf::Tuple in new_food(env.food_generators[type], 10000)
            env.food_frames[type][nf...] += 0.1f0
        end
    end

    env.players = [Player(rand_pos(env)...,env.food_types, env.player_starting_food) for _ in 1:env.num_starting_players]
    env.comms = [[] for _ in 1:env.episode_len]
    frame_channels = Int(env.obs_size[3] / env.num_frames)
    env.past_states = [
        [[zeros(Float32, env.obs_size[1:2]..., frame_channels) for _ in 1:env.num_frames-1]...] for _ in 1:env.num_starting_players]

    empty!(env.place_record)
    env.exchanges = zeros(Float32, env.food_types)
    nothing
end

function RLBase.state(env::NatureEnv, player::Int)
    if env.players[player].dead
        return zeros(Float32, env.obs_size)
    end
    # Don't create a new frame if episode is over
    if env.step > env.episode_len
        return cat(env.past_states[player][end-env.num_frames+1:end]..., dims=3)
    end
    w = env.window
    px, py = env.players[player].pos
    self_frame = make_frame(size(env)..., w)
    other_frame = make_frame(size(env)..., w)
    if env.players[player].dead
        return zeros(Float32, env.obs_size...)
    end
    # Player frames
    for i in 1:length(env.players)
        if i == player
            self_frame[(w .+ env.players[i].pos)...] = 1
        else
            other_frame[(w .+ env.players[i].pos)...] = 1
        end
    end

    # positions
    xpos = hcat([i*ones(Float32, env.world_size[1]+(2*w)) for i in 1-w:env.world_size[1]+w]...) ./ (env.world_size[1]+w)
    ypos = vcat([i*ones(Float32, 1, env.world_size[2]+(2*w)) for i in 1-w:env.world_size[2]+w]...) ./ (env.world_size[2]+w)

    # Food frames and food counts
    food_frames = []
    food_counts = []
    for (f, env_food_frame) in enumerate(env.food_frames)
        ff = make_frame(size(env)..., w)
        fc = make_frame(size(env)..., w) .* 0f0
        ff[w+1:end-w,w+1:end-w] = env_food_frame ./ 10f0
        for p in env.players
            p.dead && continue
            fc[w+p.pos[1],w+p.pos[2]] = p.food_counts[f] / 10f0
        end
        push!(food_frames, ff)
        push!(food_counts, fc)
    end
    comm_frames = [make_frame(size(env)...,w) for _ in 1:env.vocab_size]
    for (x, y, _, symbol) in env.comms[env.step]
        comm_frames[symbol][w+x,w+y] = 1
    end

    frames = [self_frame, other_frame, xpos, ypos, food_frames..., food_counts..., comm_frames...]
    frames = cat(frames..., dims=3)
    frames = frames[px:px+(2*w), py:py+(2*w), :]
    if length(env.past_states[player]) == env.num_frames + env.step - 2
        push!(env.past_states[player], frames)
    elseif length(env.past_states[player]) == env.num_frames + env.step - 1
        env.past_states[player][end] = frames
    end
    frame_history = cat(env.past_states[player][end-env.num_frames+1:end]..., dims=3)
    frame_history
end

function (env::NatureEnv)(actions::Dict)
    if is_terminated(env)
        return nothing
    end
    env.step > env.episode_len && begin
        println("exceeded")
        return nothing
    end
    env.step += 1
    Dict(p=>env(a.action[1], p) for (p, a) in actions)
end

function (env::NatureEnv)(action::Int, player::Int)

    if is_terminated(env, player)
        return nothing
    end

    MOVE_RANGE = 1:4
    FOOD_RANGE = (MOVE_RANGE.stop+1)$(2*env.food_types)
    VOCAB_RANGE = (FOOD_RANGE.stop+1)$(env.vocab_size+1)

    # Players lose 0.1 food per tick, floored at 0
    if action in MOVE_RANGE
        env.players[player].food_counts = env.players[player].food_counts .- 0.15f0 .|> x->max(x, 0f0)
    else
        env.players[player].food_counts = env.players[player].food_counts .- 0.1f0 .|> x->max(x, 0f0)
    end

    (action in MOVE_RANGE) && move(env, player, action)
    (action in FOOD_RANGE) && food(env, player, action - FOOD_RANGE.start + 1)
    (action in VOCAB_RANGE) && comm(env, player, action - VOCAB_RANGE.start + 1)
    if min(env.players[player].food_counts...) <= 0f0
        env.players[player].dead = true
    end
    nothing
end

RLBase.legal_action_space_mask(::NatureEnv) = nothing
RLBase.current_player(env::NatureEnv) =  Dict{Int, Player}(i=>p for (i,p) in enumerate(env.players) if !p.dead)
# https://juliareinforcementlearning.org/docs/How_to_write_a_customized_environment/#More-Complicated-Environments
RLBase.NumAgentStyle(env::NatureEnv) = MultiAgent(length(env.players))
RLBase.DynamicStyle(::NatureEnv) = SIMULTANEOUS
RLBase.ActionStyle(::NatureEnv) = FULL_ACTION_SET
RLBase.InformationStyle(::NatureEnv) = IMPERFECT_INFORMATION
# RLBase.StateStyle(::TicTacToeEnv) =
#     (Observation{Any}(), Observation{Int}(), Observation{BitArray{3}}())
# RLBase.RewardStyle(::TicTacToeEnv) = TERMINAL_REWARD
# RLBase.UtilityStyle(::TicTacToeEnv) = GENERAL_SUM
RLBase.ChanceStyle(::NatureEnv) = DETERMINISTIC


# Copyright (c) 2021 Idiap Research Institute, http://www.idiap.ch/
#  Niccolò Antonello <nantonel@idiap.ch>

using ArgParse

s = ArgParseSettings()
@add_arg_table! s begin
    "--conf"
        help = "configuration setup"
        arg_type = String
        default = "2a"
end
parsed_args = parse_args(ARGS, s)
setup = parsed_args["conf"]

using HMMGradients, Flux, Zygote
using Random, Statistics, LinearAlgebra
using FiniteStateTransducers
using DSP, MFCC
using BSON, JLD2, LibSndFile, FileIO, UUIDs
include("WFSTs.jl")
include("Models.jl")
include("Utils.jl")

include("conf/$(setup)/feat_conf.jl")

# get transition matrix
lexicon, ilexicon = get_lexicon()
H, L = get_HL(lexicon)
a, A = get_aA(H)
ippsym = get_iisym(H)

BSON.@load "models/$setup/best_modely_final.bson" best_modely
Flux.testmode!(best_modely)

feat_dir = get_feat_dir(setup)
data = load(joinpath(feat_dir,"train.jld2"))
uttID2feats, uttID2text, uttID2phones = data["uttID2feats"], data["uttID2text"], data["uttID2phones"]

uttID2text_dec   = Dict()
uttID2phones_dec = Dict()
min_dist=2

for uttID in keys(uttID2feats)
  x = uttID2feats[uttID]
  x = feats_post(x)
  y = best_modely(Flux.unsqueeze(x,3))
  gamma = logposterior(size(y,1),a,A,y[:,:])
  ps  = posterior2phones(ippsym,gamma)
  ws  = phones2words_greedy(ilexicon,ps; min_dist=min_dist)
  uttID2phones_dec[uttID] = ps
  uttID2text_dec[uttID] = strip(prod(ws.*" "))
end

uttID2err_textdec = Dict{String,Tuple{String,String}}()
for uttID in keys(uttID2text)
  text, dec = uttID2text[uttID], uttID2text_dec[uttID]
  if text != dec
    uttID2err_textdec[uttID] = (text,dec)
  end
end

accuracy = 1-length(uttID2err_textdec) / length(uttID2text)
wer = get_error_rate(uttID2text, uttID2text_dec)
per = get_error_rate(uttID2phones, uttID2phones_dec; is_split=true)
println("# Setup $setup")
println("* Phone Error Rate (PER): $(round(per * 100, digits=3)) %")
println("* Word Error Rate (WER) : $(round(wer * 100, digits=3)) %")
println("* Accuracy: $(round(accuracy,digits=3))")


"""
    iht(filename, k, d, phenotypes=6, covariates="", summaryfile="iht.summary.txt",
        betafile="iht.beta.txt", kwargs...)

Runs IHT with sparsity level `k`.

# Arguments
- `filename`: A `String` for VCF, binary PLINK, or BGEN file. VCF files should end
    in `.vcf` or `.vcf.gz`. Binary PLINK files should exclude `.bim/.bed/.fam`
    trailings but the trio must all be present in the same directory. BGEN files
    should end in `.bgen`.
- `k`: An `Int` for sparsity parameter = number of none-zero coefficients
- `d`: Distribution of phenotypes. Specify `Normal` for quantitative traits,
    `Bernoulli` for binary traits, `Poisson` or `NegativeBinomial` for
    count traits, and `MvNormal` for multiple quantitative traits.

# Optional Arguments
- `phenotypes`: Phenotype file name (`String`), an integer, or vector of integer. Integer(s)
    coresponds to the column(s) of PLINK's `.fam` file that stores phenotypes (default `phenotypes=6`).
    Enter multiple integers for multivariate analysis (e.g. `phenotypes=[6, 7]`).
    We recognize missing phenotypes as `NA` or `-9`. For quantitative traits
    (univariate or multivariate), missing phenotypes are imputed with the mean. Binary
    and count phenotypes cannot be imputed. Phenotype files are read using `readdlm` function
    in Julia base. We require each subject's phenotype to occupy a different row, and for
    multivariate analysis, each phenotype is comma separated. The phenotype file
    should not include a header line. Each row should be listed in the same order as in
    the PLINK and (for multivariate analysis) be comma separated.
- `covariates`: Covariate file name. Default `covariates=""` (in which case an intercept
    term will be automatically included). If `covariates` file specified, it will be
    read using `readdlm` function in Julia base. We require the covariate file to be
    comma separated, and not include a header line. Each row should be listed in the
    same order as in the PLINK. The first column should be all 1s to indicate an
    intercept. All other columns not specified in `exclude_std_idx` will be standardized
    to mean 0 variance 1.
- `summaryfile`: Output file name for saving IHT's summary statistics. Default
    `summaryfile="iht.summary.txt"`.
- `betafile`: Output file name for saving IHT's estimated genotype effect sizes.
    Default `betafile="iht.beta.txt"`.
- `covariancefile`: Output file name for saving IHT's estimated trait covariance
    matrix for multivariate analysis. Default `covariancefile="iht.cov.txt"`.
- `exclude_std_idx`: Indices of non-genetic covariates that should be excluded from
    standardization.
- `dosage`: Currently only guaranteed to work for VCF files. If `true`, will read
    genotypes dosages (i.e. `X[i, j] ∈ [0, 2]` before standardizing)
- All optional arguments available in [`fit_iht`](@ref)
"""
function iht(
    filename::AbstractString,
    k::Int,
    d::UnionAll;
    phenotypes::Union{AbstractString, Int, AbstractVector{Int}} = 6,
    covariates::AbstractString = "",
    summaryfile::AbstractString = "iht.summary.txt",
    betafile::AbstractString = "iht.beta.txt",
    covariancefile::AbstractString = "iht.cov.txt",
    exclude_std_idx::AbstractVector{<:Integer} = Int[],
    dosage::Bool = false,
    kwargs...
    )
    # read genotypes
    X, X_sampleID, X_chr, X_pos, X_ids, X_ref, X_alt = parse_genotypes(filename, dosage)
    if typeof(X) <: SnpData
        xla = SnpLinAlg{Float64}(X.snparray, model=ADDITIVE_MODEL,
            center=true, scale=true, impute=true)
    else
        xla = X # numeric matrix from VCF or BGEN files
    end

    # read phenotypes
    y = parse_phenotypes(X, phenotypes, d())

    # read and standardize covariates
    z = covariates == "" ? ones(size(xla, 1)) :
        parse_covariates(covariates, exclude_std_idx, standardize=true)
    is_multivariate(y) && (z = convert(Matrix{Float64}, Transpose(z)))

    # run IHT
    io = open(summaryfile, "w")
    if is_multivariate(y)
        result = fit_iht(y, Transpose(xla), z, k=k, io=io; kwargs...)
    else
        l = d == NegativeBinomial ? LogLink() : canonicallink(d()) # link function
        result = fit_iht(y, xla, z, k=k, d=d(), l=l, io=io; kwargs...)
    end

    show(io, result)

    if is_multivariate(y)
        open(betafile, "w") do io
            print(io, "SNPid")
            for i in 1:size(y, 1)
                print(io, '\t', "beta_$i")
            end
            print(io, '\n')
            writedlm(io, [X_ids result.beta'])
        end
        # writedlm(betafile, DataFrame([X_ids result.beta'], ["SNPid", "Estimated_beta"]))
        writedlm(covariancefile, result.Σ)
    else
        open(betafile, "w") do io
            println(io, "SNPid", '\t', "Estimated_beta")
            writedlm(io, [X_ids result.beta])
        end
    end

    close(io)
    flush(io)

    return result
end

# adhoc constructor for empty MvNormal distribution
MvNormal() = MvNormal(Float64[])

"""
    parse_phenotypes(x, col::Union{Int, AbstractVector{Int}}, ::Distribution)

Reads phenotypes to numeric array. If `x` is a `SnpData`, columns `col` of the `.fam`
file will be parsed as phenotypes. Otherwise, will read `x` as comma-separated text
file where each sample occupies a row. We recognize missing phenotypes as `NA` or
`-9`. For quantitative traits (univariate or multivariate), missing phenotypes are
imputed with the mean. Binary and count phenotypes cannot be imputed.
"""
function parse_phenotypes end

function parse_phenotypes(x::SnpData, col::AbstractVector{Int}, ::MvNormal)
    n = x.people
    r = length(col) # number of traits
    y = Matrix{Float64}(undef, r, n)
    offset = 5

    # impute missing phenotypes "-9" by mean of observed phenotypes
    missing_idx = Int[]
    for c in col
        empty!(missing_idx)
        s = 0.0
        for i in 1:n
            if phenotype_is_missing(x.person_info[i, c])
                y[c - offset, i] = 0.0
                push!(missing_idx, i)
            else
                y[c - offset, i] = parse(Float64, x.person_info[i, c])
                s += y[c - offset, i]
            end
        end
        avg = s / (n - length(missing_idx))
        for i in missing_idx
            y[c - offset, i] = avg
        end
    end
    return y
end

parse_phenotypes(::SnpData, ::Int, ::MvNormal) =
    throw(ArgumentError("Multivariate analysis requires multiple phenotypes! Please specify " *
        "e.g. phenotypes=[6, 7] or save each sample's phenotypes in a comma-" *
        "separated file where each sample occupies a different row and each" *
        " phenotype is separated by a single comma."))

function parse_phenotypes(x::SnpData, col::Int, ::Normal)
    n = x.people
    y = Vector{Float64}(undef, n)

    # impute missing phenotypes by mean of observed phenotypes
    missing_idx = Int[]
    s = 0.0
    for i in 1:n
        if phenotype_is_missing(x.person_info[i, col])
            y[i] = 0.0
            push!(missing_idx, i)
        else
            y[i] = parse(Float64, x.person_info[i, col])
            s += y[i]
        end
    end
    avg = s / (n - length(missing_idx))
    for i in missing_idx
        y[i] = avg
    end
    return y
end

function parse_phenotypes(x::SnpData, col::Int, ::UnivariateDistribution)
    n = x.people
    y = Vector{Float64}(undef, n)

    # missing phenotypes NOT allowed for binary/count phenotypes
    for i in 1:n
        if phenotype_is_missing(x.person_info[i, col])
            throw(MissingException("Missing phenotype detected for sample $i. Automatic " *
                "phenotype imputation are only possible for quantitative traits. " *
                "Please exclude missing phenotypes or impute them first."))
        else
            y[i] = parse(Float64, x.person_info[i, col])
        end
    end
    return y
end

function parse_phenotypes(::Any, pheno_filename::AbstractString, d)
    y = readdlm(pheno_filename, ',', Float64)
    if is_multivariate(y)
        y = convert(Matrix{Float64}, Transpose(y))
    else
        y = dropdims(y, dims=2)
    end
    return y
end

"""
    parse_covariates(filename, exclude_std_idx; standardize::Bool=true)

Reads a comma separated text file `filename`. Each row should be a sample ordered the
same as in the plink file. The first column should be array of 1 (representing
intercept). Each covariate should be comma separated. If `standardize=true`,
all columns except those in `exclude_std_idx` will be standardized.
"""
function parse_covariates(filename::AbstractString, exclude_std_idx::AbstractVector{<:Integer};
    standardize::Bool=true)
    z = readdlm(filename, ',', Float64)

    if eltype(exclude_std_idx) == Bool
        std_idx = .!exclude_std_idx
    else
        std_idx = trues(size(z, 2))
        std_idx[exclude_std_idx] .= false
    end

    if all(x == 1 for x in @view(z[:, 1]))
        std_idx[1] = false # don't standardize intercept
    else
        @warn("Covariate file provided but did not detect an intercept. An intercept will NOT be included in IHT!")
    end

    standardize && standardize!(@view(z[:, std_idx]))
    return z
end

function phenotype_is_missing(s::AbstractString)
    return s == "-9" || s == "NA"
end

"""
    cross_validate(filename, d, path=1:20, phenotypes=6, covariates="",
        cv_summaryfile="cviht.summary.txt", q=5, kwargs...)

Runs cross-validation to determinal optimal sparsity level `k`. Different
sparsity levels are specified in `path`.

# Arguments
- `filename`: A `String` for input PLINK file name (without `.bim/.bed/.fam` suffixes)
- `d`: Distribution of phenotypes. Specify `Normal` for quantitative traits,
    `Bernoulli` for binary traits, `Poisson` or `NegativeBinomial` for
    count traits, and `MvNormal` for multiple quantitative traits.

# Optional Arguments
- `path`: Different values of `k` that should be tested. One can input a vector of
    `Int` (e.g. `path=[5, 10, 15, 20]`) or a range (default `path=1:20`).
- `phenotypes`: Phenotype file name (`String`), an integer, or vector of integer. Integer(s)
    coresponds to the column(s) of `.fam` file that stores phenotypes (default 6).
    We recognize missing phenotypes as `NA` or `-9`. For quantitative traits
    (univariate or multivariate), missing phenotypes are imputed with the mean. Binary
    and count phenotypes cannot be imputed. Phenotype files are read using `readdlm` function
    in Julia base. We require each subject's phenotype to occupy a different row, and for
    multivariate analysis, each phenotype is comma separated. The phenotype file
    should not include a header line. Each row should be listed in the same order as in
    the PLINK.
- `covariates`: Covariate file name. Default `covariates=""` (in which case an intercept
    term will be automatically included). If `covariates` file specified, it will be
    read using `readdlm` function in Julia base. We require the covariate file to be
    comma separated, and not include a header line. Each row should be listed in the
    same order as in the PLINK. The first column should be all 1s to indicate an
    intercept. All other columns not specified in `exclude_std_idx` will be standardized
    to mean 0 variance 1
- `cv_summaryfile`: Output file name for saving IHT's cross validation summary statistics.
    Default `cv_summaryfile="cviht.summary.txt"`.
- `q`: Number of cross validation folds. Larger means more accurate and more computationally
    intensive. Should be larger 2 and smaller than 10. Default `q=5`.
- `dosage`: Currently only guaranteed to work for VCF files. If `true`, will read
    genotypes dosages (i.e. `X[i, j] ∈ [0, 2]` before standardizing)
- All optional arguments available in [`cv_iht`](@ref)
"""
function cross_validate(
    filename::AbstractString,
    d::UnionAll;
    path::AbstractVector{<:Integer} = 1:20,
    phenotypes::Union{AbstractString, Int, AbstractVector{Int}} = 6,
    covariates::AbstractString = "",
    cv_summaryfile::AbstractString = "cviht.summary.txt",
    q::Int = 5,
    exclude_std_idx::AbstractVector{<:Integer} = Int[],
    dosage::Bool = false,
    kwargs...
    )
    start_time = time()

    # read genotypes
    X, X_sampleID, X_chr, X_pos, X_ids, X_ref, X_alt = parse_genotypes(filename, dosage)
    if typeof(X) <: SnpData
        x = SnpLinAlg{Float64}(X.snparray, model=ADDITIVE_MODEL,
            center=true, scale=true, impute=true)
    else
        x = X # numeric matrix from VCF or BGEN files
    end

    # read phenotypes
    y = parse_phenotypes(X, phenotypes, d())

    # read and standardize covariates
    z = covariates == "" ? ones(size(x, 1)) :
        parse_covariates(covariates, exclude_std_idx, standardize=true)
    is_multivariate(y) && (z = convert(Matrix{Float64}, Transpose(z)))

    # run cross validation
    if is_multivariate(y)
        mse = cv_iht(y, Transpose(x), z, path=path, q=q; kwargs...)
    else
        l = d == NegativeBinomial ? LogLink() : canonicallink(d()) # link function
        mse = cv_iht(y, x, z, path=path, q=q, d=d(), l=l; kwargs...)
    end

    # save results
    open(cv_summaryfile, "w") do io
        k = path[argmin(mse)]
        print_cv_results(io, mse, path, k)
        end_time = time() - start_time
        println(io, "Total cross validation time = $end_time seconds")
    end

    return mse
end

"""
    convert_bgen_gt(t::Type{T}, b::Bgen)

Imports BGEN genotypes and chr/sampleID/pos/snpID/ref/alt into numeric arrays.
Genotypes are centered and scaled to mean 0 variance 1. Missing genotypes will
be replaced with the mean. Assumes every variant is biallelic (ie only 1 alt allele).

# Input
- `b`: a `Bgen` object
- `T`: Type for genotype array

# Output
- `G`: matrix of genotypes with type `T`.
"""
function convert_bgen_gt(t::Type{T}, b::Bgen) where T <: Real
    n = n_samples(b)
    p = n_variants(b)

    # return arrays
    G = Matrix{t}(undef, n, p)
    Gchr = Vector{String}(undef, p)
    Gpos = Vector{Int}(undef, p)
    GsnpID = [String[] for _ in 1:p] # each variant can have >1 rsid, although we don't presently allow this
    Gref = Vector{String}(undef, p)
    Galt = [String[] for _ in 1:p] # each variant can have >1 alt allele, although we don't presently allow this

    # import each variant
    i = 1
    for v in iterator(b; from_bgen_start=true)
        dose = ref_allele_dosage!(b, v; T=t) # this reads REF allele as 1
        BGEN.alt_dosage!(dose, v.genotypes.preamble) # switch 2 and 0 (ie treat ALT as 1)
        copyto!(@view(G[:, i]), dose)
        # store chr/pos/snpID/ref/alt info
        Gchr[i], Gpos[i] = chrom(v), pos(v)
        push!(GsnpID[i], rsid(v))
        ref_alt_alleles = alleles(v)
        length(ref_alt_alleles) > 2 && error("Marker $i of BGEN is not biallelic!")
        Gref[i] = ref_alt_alleles[1]
        push!(Galt[i], ref_alt_alleles[2])
        i += 1
        clear!(v)
    end

    # center/scale/impute
    standardize_genotypes!(G)

    return G, b.samples, Gchr, Gpos, GsnpID, Gref, Galt
end

"""
    standardize_genotypes!(G::AbstractMatrix)

Centers and scales each column (SNP) of `G` to mean 0 variance 1. Also each
missing entry will be imputed as mean.
"""
function standardize_genotypes!(G::AbstractMatrix)
    T = eltype(G)
    @inbounds for snp in eachcol(G)
        μi, mi = zero(T), 0
        @simd for i in eachindex(snp)
            μi += isnan(snp[i]) ? zero(t) : snp[i]
            mi += isnan(snp[i]) ? 0 : 1
        end
        μi /= mi
        σi = sqrt(μi * (1 - μi / 2))
        @simd for i in eachindex(snp)
            isnan(snp[i]) && (snp[i] = μi) # impute
            snp[i] -= μi # center
            σi > 0 && (snp[i] /= σi) # scale
        end
    end
    return nothing
end

"""
    parse_genotypes(tgtfile::AbstractString, dosage=false)

Imports genotype data from `tgtfile`. If binary PLINK files are supplied, genotypes
will not be decompressed to numeric matrices (~2 bit per entry). VCF or BGEN genotypes
will be stored in double precision matrices (64 bit per entry).

# Inputs
- `tgtfile`: VCF, binary PLINK, or BGEN file. VCF files should end in `.vcf` or
    `.vcf.gz`. Binary PLINK files should exclude `.bim/.bed/.fam` trailings but
    the trio must all be present in the same directory. BGEN files should end in
    `.bgen`

# Optional Inputs
- `dosage`: Currently only guaranteed to work for VCF files. If `true`, will read
    genotypes dosages (i.e. `X[i, j] ∈ [0, 2]` before standardizing)

# Output
- `X`: a `n × p` genotype matrix of type `Float64` (VCF or BGEN inputs) or
    `SnpArrays.SnpData` (binary PLINK inputs)
- `Gchr`: Vector of `String`s holding chromosome number for each variant
- `Gpos`: Vector of `Int` holding each variant's position
- `GsnpID`: Vector of `String`s holding variant ID for each variant
- `Gref`: Vector of `String`s holding reference allele for each variant
- `Galt`: Vector of `String`s holding alterante allele for each variant
"""
function parse_genotypes(tgtfile::AbstractString, dosage=false)
    if (endswith(tgtfile, ".vcf") || endswith(tgtfile, ".vcf.gz"))
        f = dosage ? convert_ds : convert_gt
        X, X_sampleID, X_chr, X_pos, X_ids, X_ref, X_alt =
            f(Float64, tgtfile, trans=false,
            save_snp_info=true, msg = "Importing from VCF file...")
        # convert missing to NaN
        replace!(X, missing => NaN)
        X = convert(Matrix{Float64}, X) # drop Missing from Matrix type
        # center/scale/impute
        standardize_genotypes!(X)
    elseif endswith(tgtfile, ".bgen")
        # dosage && error("Currently don't support reading dosage data from BGEN files")
        samplefile = isfile(tgtfile[1:end-5] * ".sample") ?
            tgtfile[1:end-5] * ".sample" : nothing
        indexfile = isfile(tgtfile * ".bgi") ? tgtfile * ".bgi" : nothing
        bgen = Bgen(tgtfile; sample_path=samplefile, idx_path=indexfile)
        X, X_sampleID, X_chr, X_pos, X_ids, X_ref, X_alt = MendelIHT.convert_bgen_gt(Float64, bgen)
    elseif isplink(tgtfile)
        dosage && error("PLINK files detected but dosage = true!")
        X = SnpArrays.SnpData(tgtfile)
        # get other relevant information
        X_sampleID = X.person_info[!, :iid]
        X_chr = X.snp_info[!, :chromosome]
        X_pos = X.snp_info[!, :position]
        X_ids = X.snp_info[!, :snpid]
        X_ref = X.snp_info[!, :allele1]
        X_alt = X.snp_info[!, :allele2]
    else
        error("Unrecognized target file format: target file can only be VCF" *
            " files (ends in .vcf or .vcf.gz), BGEN (ends in .bgen) or PLINK" *
            " (do not include.bim/bed/fam) and all trio must exist in 1 directory)")
    end
    return X, X_sampleID, X_chr, X_pos, X_ids, X_ref, X_alt
end

isplink(tgtfile::AbstractString) = isfile(tgtfile * ".bed") &&
                                   isfile(tgtfile * ".fam") &&
                                   isfile(tgtfile * ".bim")


using Revise
using Plots
using Pkg
# Pkg.activate(".")
using RadiativeTransfer
using RadiativeTransfer.Architectures
using RadiativeTransfer.Absorption
using RadiativeTransfer.Scattering
using RadiativeTransfer.vSmartMOM
using RadiativeTransfer.SolarModel
using InstrumentOperator
using Interpolations
using Polynomials
using ForwardDiff
using Distributions
using NCDatasets

## Atmospheric Radiative Transfer

# Load parameters from file
parameters = vSmartMOM.parameters_from_yaml("test/test_parameters/O2Parameters.yaml")
#parameters.architecture = CPU()
FT = Float64

# Load OCO Data:
# File names:
L1File = "/net/fluo/data1/group/oco2/L1bSc/oco2_L1bScND_26780a_190715_B10003r_200429212407.h5"
metFile = "/net/fluo/data1/group/oco2/L2Met/oco2_L2MetND_26780a_190715_B10003r_200429212406.h5"
dictFile = "/home/cfranken/code/gitHub/InstrumentOperator.jl/json/oco2.yaml"

# Load L1 file (could just use filenames here as well)
oco = InstrumentOperator.load_L1(dictFile,L1File, metFile);

# oco2_L1bScND_18688a_180105_B8100r_180206190633

oco_file = "/net/fluo/data1/group/oco2/L1bSc/oco2_L1bScGL_15258a_170515_B10003r_200214061601.h5"
oco_met_file = "/net/fluo/data1/group/oco2/L2Met/oco2_L2MetGL_26777a_190715_B10003r_200429213029.h5"
ils_file = "test/prototyping/ils_oco2.json"

# Pick some bands as tuple (or just one)
bands = (1,);
# Indices within that band:
indices = (92:885,);
# Geo Index (footprint,sounding):
GeoInd = [5,5000];

# Get data for that sounding:
oco_sounding = InstrumentOperator.getMeasurement(oco, bands, indices, GeoInd)

fp = 5
iOrbit = 5000
band = 1

sza_ = 32.4436
vza_ = [0.072]

function conv_spectra_local(m::VariableKernelInstrument, ν, spectrum; stride=1)
    # FT = eltype(m.ν_out)
    # Define grid where to perform convolution:

    # Padding at both sides required:
    off = ceil(Int, size(m.kernel, 1) / 2)
    ind = off:stride:(length(ν) - off)

    # knots where convolution will be applied to
    knots = view(ν, ind)
    te = LinearInterpolation(m.ν_out, Float32.(m.ind_out))
    spec_out = zeros(Real, length(knots));
    for i in eachindex(knots)
        # Simple first, nearest neighbor ILS
        ind_fraction = round(Int, te(knots[i]));
        kernel = view(m.kernel, :, ind_fraction)
        for j in eachindex(kernel)
            spec_out[i] += kernel[j] * spectrum[ind[i] + j]
        end
    end
    # Change this later to only perform conv around output grid!
    fin = LinearInterpolation(ν[ind], spec_out; extrapolation_bc=Interpolations.Flat())
    return fin(m.ν_out)
end;

# Runner is used to set AD fields as duals
function runner!(y, x, parameters=parameters, oco_file=oco_file,
                                              oco_met_file=oco_met_file,
                                              ils_file=ils_file, fp=fp, iOrbit=iOrbit,
                                              sza_=sza_, vza_=vza_)

    # Set parameters fields as the dual numbers
    parameters.brdf = [vSmartMOM.LambertianSurfaceLegendre([x[1],x[3],x[4]])]

    parameters.scattering_params.rt_aerosols[1].τ_ref = exp(x[2]);
    parameters.scattering_params.rt_aerosols[1].p₀    = 70000.0; #x[4]
    # parameters.scattering_params.rt_aerosols[1].aerosol.size_distribution = LogNormal(log(x[3]), log(x[4]), check_args=false)

    # parameters.scattering_params.rt_aerosols[1].aerosol.nᵣ = x[5];
    # parameters.scattering_params.rt_aerosols[1].aerosol.nᵢ = x[6];

    # parameters.scattering_params.rt_aerosols[1].p₀ = x[7];
    # parameters.scattering_params.rt_aerosols[1].σp = x[8];
    ocoData = Dataset(oco_file)
    SoundingGeometry = ocoData.group["SoundingGeometry"]
    ϕ = SoundingGeometry["sounding_polarization_angle"][fp,iOrbit]
    m₁ = 0.5
    m₂ = cosd(2ϕ)/2
    m₃ = sind(2ϕ)/2
    # Stokes:

    # Set profiles properly
    met = Dataset(oco_met_file);
    T_met = met.group["Meteorology"]["temperature_profile_met"][:,fp,iOrbit];
    ak = met.group["MeteorologyDiagnostics"]["ak"][:,fp,iOrbit];
    bk = met.group["MeteorologyDiagnostics"]["bk"][:,fp,iOrbit];
    ak = [0.000000e+00, 4.804826e-02, 6.593752e+00, 1.313480e+01, 1.961311e+01, 2.609201e+01,
3.257081e+01, 3.898201e+01, 4.533901e+01, 5.169611e+01, 5.805321e+01, 6.436264e+01,
7.062198e+01, 7.883422e+01, 8.909992e+01, 9.936521e+01, 1.091817e+02, 1.189586e+02,
1.286959e+02, 1.429100e+02, 1.562600e+02, 1.696090e+02, 1.816190e+02, 1.930970e+02,
2.032590e+02, 2.121500e+02, 2.187760e+02, 2.238980e+02, 2.243630e+02, 2.168650e+02,
2.011920e+02, 1.769300e+02, 1.503930e+02, 1.278370e+02, 1.086630e+02, 9.236572e+01,
7.851231e+01, 6.660341e+01, 5.638791e+01, 4.764391e+01, 4.017541e+01, 3.381001e+01,
2.836781e+01, 2.373041e+01, 1.979160e+01, 1.645710e+01, 1.364340e+01, 1.127690e+01,
9.292942e+00, 7.619842e+00, 6.216801e+00, 5.046801e+00, 4.076571e+00, 3.276431e+00,
2.620211e+00, 2.084970e+00, 1.650790e+00, 1.300510e+00, 1.019440e+00, 7.951341e-01,
6.167791e-01, 4.758061e-01, 3.650411e-01, 2.785261e-01, 2.113490e-01, 1.594950e-01,
1.197030e-01, 8.934502e-02, 6.600001e-02, 4.758501e-02, 3.270000e-02, 2.000000e-02,
1.000000e-02][end:-1:1]

bk = [1.000000e+00, 9.849520e-01, 9.634060e-01, 9.418650e-01, 9.203870e-01, 8.989080e-01,
8.774290e-01, 8.560180e-01, 8.346609e-01, 8.133039e-01, 7.919469e-01, 7.706375e-01,
7.493782e-01, 7.211660e-01, 6.858999e-01, 6.506349e-01, 6.158184e-01, 5.810415e-01,
5.463042e-01, 4.945902e-01, 4.437402e-01, 3.928911e-01, 3.433811e-01, 2.944031e-01,
2.467411e-01, 2.003501e-01, 1.562241e-01, 1.136021e-01, 6.372006e-02, 2.801004e-02,
6.960025e-03, 8.175413e-09, 0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00,
0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00,
0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00,
0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00,
0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00,
0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00,
0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00,
0.000000e+00][end:-1:1]

    p_surf = met.group["Meteorology"]["surface_pressure_met"][fp,iOrbit];
    p_half = (ak + bk * p_surf);
    #p_half = vcat(p_half, p_surf);
    q = met.group["Meteorology"]["specific_humidity_profile_met"][:,fp,iOrbit];
    parameters.p = p_half / 100
    parameters.q = q # zeros(size(q))
    parameters.T = T_met
    #parameters.T[end-30:end] .-= 10
    parameters.sza = sza_
    parameters.vza = vza_

    @show parameters.sza, parameters.vza

    # @show parameters.sza
    # plot!(parameters.p)
    # return

    model = model_from_parameters(parameters);
    @show sum(model.τ_rayl[1] )
    # Run the model to obtain reflectance matrix
    R = vSmartMOM.rt_run(model, i_band=1);
    RR = m₁*R[1,1,:] + m₂*R[1,2,:] + m₃*R[1,3,:]
    @show size(R)
    # Produce black-body in wavenumber range
    T = 5777
    λ_grid = reverse(1e4 ./ parameters.spec_bands[1]) #collect(757:0.01:777)
    black_body = planck_spectrum_wl(T, λ_grid) * 2.1629e-05 * pi
    black_body = SolarModel.watts_to_photons(λ_grid, black_body)

    # Get solar transmittance spectrum
    #solar_transmission = solar_transmission_from_file("/home/rjeyaram/RadiativeTransfer/src/solar_merged_20160127_600_26316_100.out", parameters.spec_bands[1])
    solar_transmission = solar_transmission_from_file("/home/rjeyaram/RadiativeTransfer/src/SolarModel/solar.out", parameters.spec_bands[1])
    # Get outgoing solar radiation
    sun_out = reverse(solar_transmission) .* black_body

    # Apply Earth reflectance matrix
    earth_out = sun_out .* reverse(RR[:])

    # Set up InstrumentOperator
    ils_Δ, ils_in, dispersion = InstrumentOperator.read_ils_table(oco_file, ils_file);

    # Define model grid:
    res = 0.001

    # Just consider the ILS within ± 0.35nm
    grid_x = -0.35e-3:res*1e-3:0.35e-3

    extended_dims = [fp,band] # Footprint, band

    # Re-interpolate I from ν_grid to new grid/resolution
    interp_I = LinearInterpolation(λ_grid, earth_out);
    wl = 757.5:res:771.0
    I_wl = interp_I(wl/1000)

    # Pixels to be used
    ind_out = collect(0:1015);

    # Eventual grid of OCO-2 for Band 1, FP 5:
    dispPoly = Polynomial(view(dispersion, :, extended_dims...))
    ν = Float32.(dispPoly.(1:1016))

    # Prepare ILS table:
    ils_pixel   = InstrumentOperator.prepare_ils_table(grid_x, ils_in, ils_Δ,extended_dims)
    oco2_kernel = VariableKernelInstrument(ils_pixel, ν, ind_out)

    # Convolve input spectrum with variable kernel
    I_conv = conv_spectra_local(oco2_kernel, wl./1e3, I_wl)

    y[:] = I_conv[:]

end

directory = "/net/fluo/data1/group/oco2/L1bSc/"
files_list = filter(x->endswith(x, ".h5") && startswith(x, "oco2_L1bScND_26"), readdir(directory))

met_files_list = filter(x->endswith(x, ".h5"), readdir("/net/fluo/data1/group/oco2/L2Met/"))

I_convs_all = zeros(1016, length(files_list))
oco2_Abands_all = zeros(1016, length(files_list))

albedos = zeros(length(files_list))
cloud_levels = zeros(length(files_list))

optimal_x = zeros(8, length(files_list))

#for i in 1:length(files_list)
i = 1
ind = 92:885
@show i

oco_file = files_list[i]

    # try
ocoData = Dataset(directory * oco_file)
oco_met_file = "/net/fluo/data1/group/oco2/L2Met/" * filter(x->contains(x, oco_file[14:19]), met_files_list)[1]
oco_file = directory * oco_file
met = Dataset(oco_met_file)
cloud_levels[i] = met.group["Meteorology"]["cloud_liquid_water_path_met"][fp,iOrbit]

## Getting an OCO spectrum for fun to fit:
o2AbandSpectra = ocoData.group["SoundingMeasurements"]["radiance_o2"]
SoundingGeometry = ocoData.group["SoundingGeometry"]
ϕ = SoundingGeometry["sounding_polarization_angle"][fp,iOrbit]
m₁ = 0.5
m₂ = cosd(2ϕ)/2
m₃ = sind(2ϕ)/2
vza = SoundingGeometry["sounding_zenith"]
sza = SoundingGeometry["sounding_solar_zenith"]
lat = SoundingGeometry["sounding_latitude"]
lon = SoundingGeometry["sounding_longitude"]
o2AbandSpectra = ocoData.group["SoundingMeasurements"]["radiance_o2"]
oco2_Aband = o2AbandSpectra[:,fp,iOrbit]

println("$(round(Float64(lat[fp,iOrbit]), digits=2)) $(round(Float64(lon[fp,iOrbit]), digits=4))")
# continue
sza_ = sza[fp,iOrbit]
vza_ = [vza[fp,iOrbit]]

# State vector
x = FT[0.2377,
        -1, -0.00244, 0]#,

# Run FW model:
# @time runner(x);
I_conv = zeros(1016)
y = oco2_Aband[ind];
for i=1:3
    dfdx = ForwardDiff.jacobian(runner!, I_conv, x);
    Fx = I_conv[ind];
    K = dfdx[ind,:];
    dx = K \ (y-Fx)
    x += dx
    @show x
end
runner!(I_conv,x)
Fx = I_conv[ind];

# plot fits
ils_Δ, ils_in, dispersion = InstrumentOperator.read_ils_table(oco_file, ils_file);
extended_dims = [fp,band]; # Footprint, band
# Eventual grid of OCO-2 for Band 1, FP 5:
dispPoly = Polynomial(view(dispersion, :, extended_dims...));
ν = Float32.(dispPoly.(1:1016));
plot(ν[ind]*1e3,  y/1e20, label="Meas")
plot!(ν[ind]*1e3, Fx/1e20, label="Mod")
plot!(ν[ind]*1e3, (y-Fx)/1e20, label="Meas-mod")

function sampleOCO_Orbit(fp,iOrbit,l1,met)

end


# This file is adapted from test/abstractarray.jl from Julia.

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

A_abs = rand(5,4,3)
A = OpenCV.Mat(A_abs)
@testset "Bounds checking" begin
    @test checkbounds(Bool, A, 1, 1, 1) == true
    @test checkbounds(Bool, A, 5, 4, 3) == true
    @test checkbounds(Bool, A, 0, 1, 1) == false
    @test checkbounds(Bool, A, 1, 0, 1) == false
    @test checkbounds(Bool, A, 1, 1, 0) == false
    @test checkbounds(Bool, A, 6, 4, 3) == false
    @test checkbounds(Bool, A, 5, 5, 3) == false
    @test checkbounds(Bool, A, 5, 4, 4) == false
    @test checkbounds(Bool, A, 1) == true           # linear indexing
    @test checkbounds(Bool, A, 60) == true
    @test checkbounds(Bool, A, 61) == false
    @test checkbounds(Bool, A, 2, 2, 2, 1) == true  # extra indices
    @test checkbounds(Bool, A, 2, 2, 2, 2) == false
    @test checkbounds(Bool, A, 1, 1)  == false
    @test checkbounds(Bool, A, 1, 12) == false
    @test checkbounds(Bool, A, 5, 12) == false
    @test checkbounds(Bool, A, 1, 13) == false
    @test checkbounds(Bool, A, 6, 12) == false
end

@testset "single CartesianIndex" begin
    @test checkbounds(Bool, A, CartesianIndex((1, 1, 1))) == true
    @test checkbounds(Bool, A, CartesianIndex((5, 4, 3))) == true
    @test checkbounds(Bool, A, CartesianIndex((0, 1, 1))) == false
    @test checkbounds(Bool, A, CartesianIndex((1, 0, 1))) == false
    @test checkbounds(Bool, A, CartesianIndex((1, 1, 0))) == false
    @test checkbounds(Bool, A, CartesianIndex((6, 4, 3))) == false
    @test checkbounds(Bool, A, CartesianIndex((5, 5, 3))) == false
    @test checkbounds(Bool, A, CartesianIndex((5, 4, 4))) == false
    @test checkbounds(Bool, A, CartesianIndex((1,))) == false
    @test checkbounds(Bool, A, CartesianIndex((60,))) == false
    @test checkbounds(Bool, A, CartesianIndex((61,))) == false
    @test checkbounds(Bool, A, CartesianIndex((2, 2, 2, 1,))) == true
    @test checkbounds(Bool, A, CartesianIndex((2, 2, 2, 2,))) == false
    @test checkbounds(Bool, A, CartesianIndex((1, 1,)))  == false
    @test checkbounds(Bool, A, CartesianIndex((1, 12,))) == false
    @test checkbounds(Bool, A, CartesianIndex((5, 12,))) == false
    @test checkbounds(Bool, A, CartesianIndex((1, 13,))) == false
    @test checkbounds(Bool, A, CartesianIndex((6, 12,))) == false
end

@testset "mix of CartesianIndex and Int" begin
    @test checkbounds(Bool, A, CartesianIndex((1,)), 1, CartesianIndex((1,))) == true
    @test checkbounds(Bool, A, CartesianIndex((5, 4)), 3)  == true
    @test checkbounds(Bool, A, CartesianIndex((0, 1)), 1)  == false
    @test checkbounds(Bool, A, 1, CartesianIndex((0, 1)))  == false
    @test checkbounds(Bool, A, 1, 1, CartesianIndex((0,))) == false
    @test checkbounds(Bool, A, 6, CartesianIndex((4, 3)))  == false
    @test checkbounds(Bool, A, 5, CartesianIndex((5,)), 3) == false
    @test checkbounds(Bool, A, CartesianIndex((5,)), CartesianIndex((4,)), CartesianIndex((4,)))  == false
end

@testset "vector indices" begin
    @test checkbounds(Bool, A, 1:5, 1:4, 1:3) == true
    @test checkbounds(Bool, A, 0:5, 1:4, 1:3) == false
    @test checkbounds(Bool, A, 1:5, 0:4, 1:3) == false
    @test checkbounds(Bool, A, 1:5, 1:4, 0:3) == false
    @test checkbounds(Bool, A, 1:6, 1:4, 1:3) == false
    @test checkbounds(Bool, A, 1:5, 1:5, 1:3) == false
    @test checkbounds(Bool, A, 1:5, 1:4, 1:4) == false
    @test checkbounds(Bool, A, 1:60) == true
    @test checkbounds(Bool, A, 1:61) == false
    @test checkbounds(Bool, A, 2, 2, 2, 1:1) == true  # extra indices
    @test checkbounds(Bool, A, 2, 2, 2, 1:2) == false
    @test checkbounds(Bool, A, 1:5, 1:4) == false
    @test checkbounds(Bool, A, 1:5, 1:12) == false
    @test checkbounds(Bool, A, 1:5, 1:13) == false
    @test checkbounds(Bool, A, 1:6, 1:12) == false
end

@testset "logical" begin
    @test checkbounds(Bool, A, trues(5), trues(4), trues(3)) == true
    @test checkbounds(Bool, A, trues(6), trues(4), trues(3)) == false
    @test checkbounds(Bool, A, trues(5), trues(5), trues(3)) == false
    @test checkbounds(Bool, A, trues(5), trues(4), trues(4)) == false
    @test checkbounds(Bool, A, trues(60)) == true
    @test checkbounds(Bool, A, trues(61)) == false
    @test checkbounds(Bool, A, 2, 2, 2, trues(1)) == true  # extra indices
    @test checkbounds(Bool, A, 2, 2, 2, trues(2)) == false
    @test checkbounds(Bool, A, trues(5), trues(12)) == false
    @test checkbounds(Bool, A, trues(5), trues(13)) == false
    @test checkbounds(Bool, A, trues(6), trues(12)) == false
    @test checkbounds(Bool, A, trues(5, 4, 3)) == true
    @test checkbounds(Bool, A, trues(5, 4, 2)) == false
    @test checkbounds(Bool, A, trues(5, 12)) == false
    @test checkbounds(Bool, A, trues(1, 5), trues(1, 4, 1), trues(1, 1, 3)) == false
    @test checkbounds(Bool, A, trues(1, 5), trues(1, 4, 1), trues(1, 1, 2)) == false
    @test checkbounds(Bool, A, trues(1, 5), trues(1, 5, 1), trues(1, 1, 3)) == false
    @test checkbounds(Bool, A, trues(1, 5), :, 2) == false
end

@testset "array of CartesianIndex" begin
    @test checkbounds(Bool, A, [CartesianIndex((1, 1, 1))]) == true
    @test checkbounds(Bool, A, [CartesianIndex((5, 4, 3))]) == true
    @test checkbounds(Bool, A, [CartesianIndex((0, 1, 1))]) == false
    @test checkbounds(Bool, A, [CartesianIndex((1, 0, 1))]) == false
    @test checkbounds(Bool, A, [CartesianIndex((1, 1, 0))]) == false
    @test checkbounds(Bool, A, [CartesianIndex((6, 4, 3))]) == false
    @test checkbounds(Bool, A, [CartesianIndex((5, 5, 3))]) == false
    @test checkbounds(Bool, A, [CartesianIndex((5, 4, 4))]) == false
    @test checkbounds(Bool, A, [CartesianIndex((1, 1))], 1) == true
    @test checkbounds(Bool, A, [CartesianIndex((5, 4))], 3) == true
    @test checkbounds(Bool, A, [CartesianIndex((0, 1))], 1) == false
    @test checkbounds(Bool, A, [CartesianIndex((1, 0))], 1) == false
    @test checkbounds(Bool, A, [CartesianIndex((1, 1))], 0) == false
    @test checkbounds(Bool, A, [CartesianIndex((6, 4))], 3) == false
    @test checkbounds(Bool, A, [CartesianIndex((5, 5))], 3) == false
    @test checkbounds(Bool, A, [CartesianIndex((5, 4))], 4) == false
end

println("OpenCV.Mat tests passed")


function FieldVector3D(arg0::Field, arg1::Vector3D)
    return FieldVector3D((Field, Vector3D), arg0, arg1)
end

function FieldVector3D(arg0::RealFieldElement, arg1::FieldVector3D)
    return FieldVector3D((RealFieldElement, FieldVector3D), arg0, arg1)
end

function FieldVector3D(arg0::RealFieldElement, arg1::FieldVector3D, arg2::RealFieldElement, arg3::FieldVector3D)
    return FieldVector3D((RealFieldElement, FieldVector3D, RealFieldElement, FieldVector3D), arg0, arg1, arg2, arg3)
end

function FieldVector3D(arg0::RealFieldElement, arg1::FieldVector3D, arg2::RealFieldElement, arg3::FieldVector3D, arg4::RealFieldElement, arg5::FieldVector3D)
    return FieldVector3D((RealFieldElement, FieldVector3D, RealFieldElement, FieldVector3D, RealFieldElement, FieldVector3D), arg0, arg1, arg2, arg3, arg4, arg5)
end

function FieldVector3D(arg0::RealFieldElement, arg1::FieldVector3D, arg2::RealFieldElement, arg3::FieldVector3D, arg4::RealFieldElement, arg5::FieldVector3D, arg6::RealFieldElement, arg7::FieldVector3D)
    return FieldVector3D((RealFieldElement, FieldVector3D, RealFieldElement, FieldVector3D, RealFieldElement, FieldVector3D, RealFieldElement, FieldVector3D), arg0, arg1, arg2, arg3, arg4, arg5, arg6, arg7)
end

function FieldVector3D(arg0::RealFieldElement, arg1::RealFieldElement)
    return FieldVector3D((RealFieldElement, RealFieldElement), arg0, arg1)
end

function FieldVector3D(arg0::RealFieldElement, arg1::RealFieldElement, arg2::RealFieldElement)
    return FieldVector3D((RealFieldElement, RealFieldElement, RealFieldElement), arg0, arg1, arg2)
end

function FieldVector3D(arg0::RealFieldElement, arg1::Vector3D)
    return FieldVector3D((RealFieldElement, Vector3D), arg0, arg1)
end

function FieldVector3D(arg0::RealFieldElement, arg1::Vector3D, arg2::RealFieldElement, arg3::Vector3D)
    return FieldVector3D((RealFieldElement, Vector3D, RealFieldElement, Vector3D), arg0, arg1, arg2, arg3)
end

function FieldVector3D(arg0::RealFieldElement, arg1::Vector3D, arg2::RealFieldElement, arg3::Vector3D, arg4::RealFieldElement, arg5::Vector3D)
    return FieldVector3D((RealFieldElement, Vector3D, RealFieldElement, Vector3D, RealFieldElement, Vector3D), arg0, arg1, arg2, arg3, arg4, arg5)
end

function FieldVector3D(arg0::RealFieldElement, arg1::Vector3D, arg2::RealFieldElement, arg3::Vector3D, arg4::RealFieldElement, arg5::Vector3D, arg6::RealFieldElement, arg7::Vector3D)
    return FieldVector3D((RealFieldElement, Vector3D, RealFieldElement, Vector3D, RealFieldElement, Vector3D, RealFieldElement, Vector3D), arg0, arg1, arg2, arg3, arg4, arg5, arg6, arg7)
end

function FieldVector3D(arg0::Vector{RealFieldElement})
    return FieldVector3D((Vector{RealFieldElement},), arg0)
end

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

function FieldVector3D(arg0::jdouble, arg1::FieldVector3D, arg2::jdouble, arg3::FieldVector3D)
    return FieldVector3D((jdouble, FieldVector3D, jdouble, FieldVector3D), arg0, arg1, arg2, arg3)
end

function FieldVector3D(arg0::jdouble, arg1::FieldVector3D, arg2::jdouble, arg3::FieldVector3D, arg4::jdouble, arg5::FieldVector3D)
    return FieldVector3D((jdouble, FieldVector3D, jdouble, FieldVector3D, jdouble, FieldVector3D), arg0, arg1, arg2, arg3, arg4, arg5)
end

function FieldVector3D(arg0::jdouble, arg1::FieldVector3D, arg2::jdouble, arg3::FieldVector3D, arg4::jdouble, arg5::FieldVector3D, arg6::jdouble, arg7::FieldVector3D)
    return FieldVector3D((jdouble, FieldVector3D, jdouble, FieldVector3D, jdouble, FieldVector3D, jdouble, FieldVector3D), arg0, arg1, arg2, arg3, arg4, arg5, arg6, arg7)
end

function add(obj::FieldVector3D, arg0::FieldVector3D)
    return jcall(obj, "add", FieldVector3D, (FieldVector3D,), arg0)
end

function add(obj::FieldVector3D, arg0::RealFieldElement, arg1::FieldVector3D)
    return jcall(obj, "add", FieldVector3D, (RealFieldElement, FieldVector3D), arg0, arg1)
end

function add(obj::FieldVector3D, arg0::RealFieldElement, arg1::Vector3D)
    return jcall(obj, "add", FieldVector3D, (RealFieldElement, Vector3D), arg0, arg1)
end

function add(obj::FieldVector3D, arg0::Vector3D)
    return jcall(obj, "add", FieldVector3D, (Vector3D,), arg0)
end

function add(obj::FieldVector3D, arg0::jdouble, arg1::FieldVector3D)
    return jcall(obj, "add", FieldVector3D, (jdouble, FieldVector3D), arg0, arg1)
end

function add(obj::FieldVector3D, arg0::jdouble, arg1::Vector3D)
    return jcall(obj, "add", FieldVector3D, (jdouble, Vector3D), arg0, arg1)
end

function angle(::Type{FieldVector3D}, arg0::FieldVector3D, arg1::FieldVector3D)
    return jcall(FieldVector3D, "angle", RealFieldElement, (FieldVector3D, FieldVector3D), arg0, arg1)
end

function angle(::Type{FieldVector3D}, arg0::FieldVector3D, arg1::Vector3D)
    return jcall(FieldVector3D, "angle", RealFieldElement, (FieldVector3D, Vector3D), arg0, arg1)
end

function angle(::Type{FieldVector3D}, arg0::Vector3D, arg1::FieldVector3D)
    return jcall(FieldVector3D, "angle", RealFieldElement, (Vector3D, FieldVector3D), arg0, arg1)
end

function cross_product(::Type{FieldVector3D}, arg0::FieldVector3D, arg1::FieldVector3D)
    return jcall(FieldVector3D, "crossProduct", FieldVector3D, (FieldVector3D, FieldVector3D), arg0, arg1)
end

function cross_product(::Type{FieldVector3D}, arg0::FieldVector3D, arg1::Vector3D)
    return jcall(FieldVector3D, "crossProduct", FieldVector3D, (FieldVector3D, Vector3D), arg0, arg1)
end

function cross_product(::Type{FieldVector3D}, arg0::Vector3D, arg1::FieldVector3D)
    return jcall(FieldVector3D, "crossProduct", FieldVector3D, (Vector3D, FieldVector3D), arg0, arg1)
end

function cross_product(obj::FieldVector3D, arg0::FieldVector3D)
    return jcall(obj, "crossProduct", FieldVector3D, (FieldVector3D,), arg0)
end

function cross_product(obj::FieldVector3D, arg0::Vector3D)
    return jcall(obj, "crossProduct", FieldVector3D, (Vector3D,), arg0)
end

function distance(::Type{FieldVector3D}, arg0::FieldVector3D, arg1::FieldVector3D)
    return jcall(FieldVector3D, "distance", RealFieldElement, (FieldVector3D, FieldVector3D), arg0, arg1)
end

function distance(::Type{FieldVector3D}, arg0::FieldVector3D, arg1::Vector3D)
    return jcall(FieldVector3D, "distance", RealFieldElement, (FieldVector3D, Vector3D), arg0, arg1)
end

function distance(::Type{FieldVector3D}, arg0::Vector3D, arg1::FieldVector3D)
    return jcall(FieldVector3D, "distance", RealFieldElement, (Vector3D, FieldVector3D), arg0, arg1)
end

function distance(obj::FieldVector3D, arg0::FieldVector3D)
    return jcall(obj, "distance", RealFieldElement, (FieldVector3D,), arg0)
end

function distance(obj::FieldVector3D, arg0::Vector3D)
    return jcall(obj, "distance", RealFieldElement, (Vector3D,), arg0)
end

function distance1(::Type{FieldVector3D}, arg0::FieldVector3D, arg1::FieldVector3D)
    return jcall(FieldVector3D, "distance1", RealFieldElement, (FieldVector3D, FieldVector3D), arg0, arg1)
end

function distance1(::Type{FieldVector3D}, arg0::FieldVector3D, arg1::Vector3D)
    return jcall(FieldVector3D, "distance1", RealFieldElement, (FieldVector3D, Vector3D), arg0, arg1)
end

function distance1(::Type{FieldVector3D}, arg0::Vector3D, arg1::FieldVector3D)
    return jcall(FieldVector3D, "distance1", RealFieldElement, (Vector3D, FieldVector3D), arg0, arg1)
end

function distance1(obj::FieldVector3D, arg0::FieldVector3D)
    return jcall(obj, "distance1", RealFieldElement, (FieldVector3D,), arg0)
end

function distance1(obj::FieldVector3D, arg0::Vector3D)
    return jcall(obj, "distance1", RealFieldElement, (Vector3D,), arg0)
end

function distance_inf(::Type{FieldVector3D}, arg0::FieldVector3D, arg1::FieldVector3D)
    return jcall(FieldVector3D, "distanceInf", RealFieldElement, (FieldVector3D, FieldVector3D), arg0, arg1)
end

function distance_inf(::Type{FieldVector3D}, arg0::FieldVector3D, arg1::Vector3D)
    return jcall(FieldVector3D, "distanceInf", RealFieldElement, (FieldVector3D, Vector3D), arg0, arg1)
end

function distance_inf(::Type{FieldVector3D}, arg0::Vector3D, arg1::FieldVector3D)
    return jcall(FieldVector3D, "distanceInf", RealFieldElement, (Vector3D, FieldVector3D), arg0, arg1)
end

function distance_inf(obj::FieldVector3D, arg0::FieldVector3D)
    return jcall(obj, "distanceInf", RealFieldElement, (FieldVector3D,), arg0)
end

function distance_inf(obj::FieldVector3D, arg0::Vector3D)
    return jcall(obj, "distanceInf", RealFieldElement, (Vector3D,), arg0)
end

function distance_sq(::Type{FieldVector3D}, arg0::FieldVector3D, arg1::FieldVector3D)
    return jcall(FieldVector3D, "distanceSq", RealFieldElement, (FieldVector3D, FieldVector3D), arg0, arg1)
end

function distance_sq(::Type{FieldVector3D}, arg0::FieldVector3D, arg1::Vector3D)
    return jcall(FieldVector3D, "distanceSq", RealFieldElement, (FieldVector3D, Vector3D), arg0, arg1)
end

function distance_sq(::Type{FieldVector3D}, arg0::Vector3D, arg1::FieldVector3D)
    return jcall(FieldVector3D, "distanceSq", RealFieldElement, (Vector3D, FieldVector3D), arg0, arg1)
end

function distance_sq(obj::FieldVector3D, arg0::FieldVector3D)
    return jcall(obj, "distanceSq", RealFieldElement, (FieldVector3D,), arg0)
end

function distance_sq(obj::FieldVector3D, arg0::Vector3D)
    return jcall(obj, "distanceSq", RealFieldElement, (Vector3D,), arg0)
end

function dot_product(::Type{FieldVector3D}, arg0::FieldVector3D, arg1::FieldVector3D)
    return jcall(FieldVector3D, "dotProduct", RealFieldElement, (FieldVector3D, FieldVector3D), arg0, arg1)
end

function dot_product(::Type{FieldVector3D}, arg0::FieldVector3D, arg1::Vector3D)
    return jcall(FieldVector3D, "dotProduct", RealFieldElement, (FieldVector3D, Vector3D), arg0, arg1)
end

function dot_product(::Type{FieldVector3D}, arg0::Vector3D, arg1::FieldVector3D)
    return jcall(FieldVector3D, "dotProduct", RealFieldElement, (Vector3D, FieldVector3D), arg0, arg1)
end

function dot_product(obj::FieldVector3D, arg0::FieldVector3D)
    return jcall(obj, "dotProduct", RealFieldElement, (FieldVector3D,), arg0)
end

function dot_product(obj::FieldVector3D, arg0::Vector3D)
    return jcall(obj, "dotProduct", RealFieldElement, (Vector3D,), arg0)
end

function equals(obj::FieldVector3D, arg0::Object)
    return jcall(obj, "equals", jboolean, (Object,), arg0)
end

function get_alpha(obj::FieldVector3D)
    return jcall(obj, "getAlpha", RealFieldElement, ())
end

function get_delta(obj::FieldVector3D)
    return jcall(obj, "getDelta", RealFieldElement, ())
end

function get_minus_i(::Type{FieldVector3D}, arg0::Field)
    return jcall(FieldVector3D, "getMinusI", FieldVector3D, (Field,), arg0)
end

function get_minus_j(::Type{FieldVector3D}, arg0::Field)
    return jcall(FieldVector3D, "getMinusJ", FieldVector3D, (Field,), arg0)
end

function get_minus_k(::Type{FieldVector3D}, arg0::Field)
    return jcall(FieldVector3D, "getMinusK", FieldVector3D, (Field,), arg0)
end

function get_na_n(::Type{FieldVector3D}, arg0::Field)
    return jcall(FieldVector3D, "getNaN", FieldVector3D, (Field,), arg0)
end

function get_negative_infinity(::Type{FieldVector3D}, arg0::Field)
    return jcall(FieldVector3D, "getNegativeInfinity", FieldVector3D, (Field,), arg0)
end

function get_norm(obj::FieldVector3D)
    return jcall(obj, "getNorm", RealFieldElement, ())
end

function get_norm1(obj::FieldVector3D)
    return jcall(obj, "getNorm1", RealFieldElement, ())
end

function get_norm_inf(obj::FieldVector3D)
    return jcall(obj, "getNormInf", RealFieldElement, ())
end

function get_norm_sq(obj::FieldVector3D)
    return jcall(obj, "getNormSq", RealFieldElement, ())
end

function get_plus_i(::Type{FieldVector3D}, arg0::Field)
    return jcall(FieldVector3D, "getPlusI", FieldVector3D, (Field,), arg0)
end

function get_plus_j(::Type{FieldVector3D}, arg0::Field)
    return jcall(FieldVector3D, "getPlusJ", FieldVector3D, (Field,), arg0)
end

function get_plus_k(::Type{FieldVector3D}, arg0::Field)
    return jcall(FieldVector3D, "getPlusK", FieldVector3D, (Field,), arg0)
end

function get_positive_infinity(::Type{FieldVector3D}, arg0::Field)
    return jcall(FieldVector3D, "getPositiveInfinity", FieldVector3D, (Field,), arg0)
end

function get_x(obj::FieldVector3D)
    return jcall(obj, "getX", RealFieldElement, ())
end

function get_y(obj::FieldVector3D)
    return jcall(obj, "getY", RealFieldElement, ())
end

function get_z(obj::FieldVector3D)
    return jcall(obj, "getZ", RealFieldElement, ())
end

function get_zero(::Type{FieldVector3D}, arg0::Field)
    return jcall(FieldVector3D, "getZero", FieldVector3D, (Field,), arg0)
end

function hash_code(obj::FieldVector3D)
    return jcall(obj, "hashCode", jint, ())
end

function is_infinite(obj::FieldVector3D)
    return jcall(obj, "isInfinite", jboolean, ())
end

function is_na_n(obj::FieldVector3D)
    return jcall(obj, "isNaN", jboolean, ())
end

function negate(obj::FieldVector3D)
    return jcall(obj, "negate", FieldVector3D, ())
end

function normalize(obj::FieldVector3D)
    return jcall(obj, "normalize", FieldVector3D, ())
end

function orthogonal(obj::FieldVector3D)
    return jcall(obj, "orthogonal", FieldVector3D, ())
end

function scalar_multiply(obj::FieldVector3D, arg0::RealFieldElement)
    return jcall(obj, "scalarMultiply", FieldVector3D, (RealFieldElement,), arg0)
end

function scalar_multiply(obj::FieldVector3D, arg0::jdouble)
    return jcall(obj, "scalarMultiply", FieldVector3D, (jdouble,), arg0)
end

function subtract(obj::FieldVector3D, arg0::FieldVector3D)
    return jcall(obj, "subtract", FieldVector3D, (FieldVector3D,), arg0)
end

function subtract(obj::FieldVector3D, arg0::RealFieldElement, arg1::FieldVector3D)
    return jcall(obj, "subtract", FieldVector3D, (RealFieldElement, FieldVector3D), arg0, arg1)
end

function subtract(obj::FieldVector3D, arg0::RealFieldElement, arg1::Vector3D)
    return jcall(obj, "subtract", FieldVector3D, (RealFieldElement, Vector3D), arg0, arg1)
end

function subtract(obj::FieldVector3D, arg0::Vector3D)
    return jcall(obj, "subtract", FieldVector3D, (Vector3D,), arg0)
end

function subtract(obj::FieldVector3D, arg0::jdouble, arg1::FieldVector3D)
    return jcall(obj, "subtract", FieldVector3D, (jdouble, FieldVector3D), arg0, arg1)
end

function subtract(obj::FieldVector3D, arg0::jdouble, arg1::Vector3D)
    return jcall(obj, "subtract", FieldVector3D, (jdouble, Vector3D), arg0, arg1)
end

function to_array(obj::FieldVector3D)
    return jcall(obj, "toArray", Vector{RealFieldElement}, ())
end

function to_string(obj::FieldVector3D)
    return jcall(obj, "toString", JString, ())
end

function to_string(obj::FieldVector3D, arg0::NumberFormat)
    return jcall(obj, "toString", JString, (NumberFormat,), arg0)
end

function to_vector3d(obj::FieldVector3D)
    return jcall(obj, "toVector3D", Vector3D, ())
end


function debugger(arena::Honeycomb,gstack::Array{UInt64,1},list::Array{Pointer,1},deb::Debugstate,ticks::Time)
    if deb.d==1
        println("\nticks:",ticks.t)
        @inbounds for ind=1:length(list)
            #println(ind,"  instr: ",arena.a[list[ind].loc[1],list[ind].loc[2]],"  dir: ",showdir(list[ind].dir),"  loc: ",list[ind].loc,"  ",transpose(list[ind].lstack),"•")
            println(ind,"  instr: ",arena.a[list[ind].loc[1],list[ind].loc[2]],"  dir: ",showdir(list[ind].dir),"  loc: ",list[ind].loc,"  ",transpose(convert(Array{BigInt},list[ind].lstack)),"•")
        end
        print("\n[")
        @inbounds for n=1:length(gstack)
            print(Int128(gstack[n]))
            n<length(gstack) ? print(" ") : nothing
        end
        println("]•")
        print("\n")
        println("———————")
    elseif deb.d >= 2
        println("\nticks: ",ticks.t)
        gdebug(arena,gstack,list,deb);print("\n")
        println("———————")
    end
end

function gdebug(arena::Honeycomb,gstack::Array{UInt64,1},list::Array{Pointer,1},deb::Debugstate)
    rows=maximum(length(arena.a[:,1]))
    cols=maximum(length(arena.a[1,:]))
    comb=deepcopy(arena.a)
    @inbounds for ind=1:length(list)
        comb[list[ind].loc[1],list[ind].loc[2]]=Char(ind+944)#translates IP numbers to greek letters
    end

    @inbounds for ind=1:length(list)
        #println(Char(ind+944),"  ",ind,"  instr: ",arena.a[list[ind].loc[1],list[ind].loc[2]]," wait: ",showwait(list[ind].wait)," dir: ",showdir(list[ind].dir),"  loc: ",list[ind].loc,"  ",transpose(list[ind].lstack),"•")
        println(Char(ind+944),"  ",ind,"  instr: ",arena.a[list[ind].loc[1],list[ind].loc[2]]," wait: ",showwait(list[ind].wait)," dir: ",showdir(list[ind].dir),"  loc: ",list[ind].loc,"  ",transpose(convert(Array{BigInt},list[ind].lstack)),"•")
    end

    println("\ngstack: ",transpose(convert(Array{BigInt},gstack)),"•\n\n")

    @inbounds for r=1:rows
        @inbounds for c=1:cols
            if Int(comb[r,c])>944 && Int(comb[r,c])<976 #greek letters unicode range
                if deb.d==3
                    c<cols ? print_with_color(:red,"$(comb[r,c])") : print_with_color(:red,"$(comb[r,c])\n")
                elseif deb.d==2
                    c<cols ? print(comb[r,c]) : println(comb[r,c])
                end
            else c<cols ? print(comb[r,c]) : println(comb[r,c])
            end
        end
    end
end

function showwait(wait::Int)
    wait==0 ? (return "∅") :
    wait==1 ? (return "∫") :
    wait==2 ? (return "∬") : nothing
end

function showdir(dir::Int)
    dir==0 ? (return "→") :
    dir==1 ? (return "↗") :
    dir==2 ? (return "↖") :
    dir==3 ? (return "←") :
    dir==4 ? (return "↙") :
    dir==5 ? (return "↘") : nothing
end


########################
# at_dynamic generator #
########################

struct AtDynamicTrace{T,U,K}
    call::CallRecord{T}
    kernel_trace::U
    key::K
end

struct AtDynamicAssignment{T,U,K} <: Assignment
    trace::AtDynamicTrace{T,U,K}
end

get_call_record(trace::AtDynamicTrace) = trace.call
has_choices(trace::AtDynamicTrace) = has_choices(trace.kernel_trace)
get_assignment(trace::AtDynamicTrace) = AtDynamicAssignment(trace)
Base.isempty(assignment::AtDynamicAssignment) = !has_choices(assignment.trace.kernel_trace)
get_address_schema(::Type{AtDynamicAssignment}) = SingleDynamicKeyAddressSchema()
function has_internal_node(assignment::AtDynamicAssignment{T,U,K}, addr::K) where {T,U,K}
    assignment.trace.key == addr
end
function has_internal_node(assignment::AtDynamicAssignment{T,U,K}, addr::Pair{K,W}) where {T,U,K,W}
    (first, rest) = addr
    if assignment.trace.key == first
        sub_assignment = get_assignment(assignment.trace.kernel_trace)
        has_internal_node(sub_assignment, rest)
    else
        throw(KeyError(first))
    end
end
function get_internal_node(assignment::AtDynamicAssignment{T,U,K}, addr::K) where {T,U,K}
    if assignment.trace.key == addr
        get_assignment(assignment.trace.kernel_trace)
    else
        throw(KeyError(addr))
    end
end
function get_internal_node(assignment::AtDynamicAssignment{T,U,K}, addr::Pair{K,W}) where {T,U,K,W}
    (first, rest) = addr
    if assignment.trace.key == first
        sub_assignment = get_assignment(assignment.trace.kernel_trace)
        get_internal_node(sub_assignment, rest)
    else
        throw(KeyError(first))
    end
end

function get_internal_nodes(assignment::AtDynamicAssignment)
    node = (assignment.trace.key, get_assignment(assignment.trace.kernel_trace))
    (node,)
end

function get_leaf_nodes(assignment::AtDynamicAssignment)
    # TODO what about wrapping a distribuion, instead of a a generator (cf Plate)
    ()
end

"""
Generator that applies another generator at a specific, dynamically determined, address.
"""
struct AtDynamic{T,U,K} <: Generator{T, AtDynamicTrace{T,U,K}}
    kernel::Generator{T,U}
end

function at_dynamic(kernel::Generator{T,U}, ::Type{K}) where {T,U,K}
    AtDynamic{T,U,K}(kernel)
end

function get_static_argument_types(gen::AtDynamic{T,U,K}) where {T,U,K}
    [K, Tuple{get_static_argument_types(gen.kernel)...}]
end

include("simulate.jl")
include("assess.jl")

export at_dynamic


# ============================================================================ #
mutable struct Board
    hf::PyPlot.Figure
    ax::PyCall.PyObject
    squares::PatchDict
    coins::PatchDict
    player::Player
    counter::Int
end

function Board(hf, ax)
    if !isassigned(BOARD, 1)
        p = Player(1, 1, East)
        b = Board(hf, ax, PatchDict(), PatchDict(), p, 0)
        BOARD[1] = b
    else
        b = BOARD[1]
        b.player = Player(1, 1, East)
        b.hf = hf
        b.ax = ax
    end
    return b
end

increment!(b::Board) = begin b.counter += 1 end

function set_color(b::Board, x::Integer, y::Integer, col::AbstractString)
    if haskey(b.squares, (x, y))
        b.squares[(x,y)].set_facecolor(col)
    end
    return nothing
end

function get_color(b::Board, x::Integer, y::Integer)
    !haskey(b.squares, (x, y)) && error("Invalid patch corrdinates!")
    return colorname(b.squares[(x,y)].get_facecolor())
end

get_color(b::Board) = get_color(b, location(b.player)...)

function set_player(b::Board, x::Integer, y::Integer, facing::Bearing)
    set_to_face!(b.player, facing)
    moveto!(b.player, x, y)
    b.ax.add_patch(b.player.h)
    return nothing
end

function set_square(b::Board, x::Integer, y::Integer, color::AbstractString="gray")
    p = matplotlib.patches.Rectangle((x,y), 1, 1, color=color)
    b.squares[(x,y)] = p
    b.ax.add_patch(p)
    return nothing
end

function set_coin(b::Board, x::Integer, y::Integer)
    p = matplotlib.patches.CirclePolygon((x+.5,y+.5), .2, color="gold")
    b.coins[(x,y)] = p
    b.ax.add_patch(p)
    return nothing
end

function turn(b::Board, d::Direction)
    game_over() && return nothing
    increment!(b)
    turn!(b.player, d)
    sleep(1/SPEED)
    return nothing
end

@inline left(b::Board) = turn(b, Left)
@inline right(b::Board) = turn(b, Right)

function forward(b::Board)
    game_over() && return nothing
    increment!(b)
    dst = Tuple{Int,Int}(destination(b.player))
    if haskey(b.squares, dst)
        moveto!(b.player, dst...)
        b.hf.canvas.draw()

        if haskey(b.coins, dst)
            sleep(1/(SPEED*2))
            b.coins[dst].remove()
            delete!(b.coins, dst)
            if isempty(b.coins)
                message_box(b, "Final score: $(b.counter)!")
            end
        end
    end
    sleep(1/SPEED)
    return nothing
end

coins_remaining(b::Board) = length(b.coins)

@inline left() = isassigned(BOARD, 1) && left(BOARD[1])
@inline right() = isassigned(BOARD, 1) && right(BOARD[1])
@inline forward() = isassigned(BOARD, 1) && forward(BOARD[1])
@inline get_color() = isassigned(BOARD, 1) && get_color(BOARD[1])
@inline coins_remaining() = isassigned(BOARD, 1) && coins_remaining(BOARD[1])

@inline red() = get_color() == "red"
@inline green() = get_color() == "green"
@inline blue() = get_color() == "blue"
@inline gray() = get_color() == "gray"

function game_over()
    if isassigned(BOARD, 1)
        return isempty(BOARD[1].coins) || BOARD[1].counter > ceil(Int, length(BOARD[1].squares)* 1.5)
    else
        return true
    end
end

function refresh()
    if isassigned(RAW_STR, 1)
        make_board(RAW_STR[1])
    end
    return nothing
end

function clear!(b::Board)
    if plt.fignum_exists(b.hf.number)
        clear!(b.squares)
        clear!(b.coins)
        clear!(b.player)
    end
    b.counter = 0
    return b
end

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


function make_ner(;
    n_frames,
    height,
    width,
    color="black",
    words_in,
    words_out,
    margin=width/10
)

marg_width = width ÷ 2 - margin
marg_height = height ÷ 2 - margin
time_steps = floor.(Int, range(1, n_frames, length = 5))

Object(
    1:time_steps[end],
    (args...) -> JUtils.draw_text_doc(
        pos=O,
        sizex=100,
        sizey=100,
        color=color,
        n_lines=10
    )
)
n_in = length(words_in)
start_points = [(min(marg_width, marg_height) ÷ 1.5) * Point(cos(θ), -sin(θ)) for θ in (n_in > 1 ? range(π/2, -π/2, length=n_in) : [π - π/3])]
inwords = [Object(
    1:time_steps[3],
    @JShape begin
        sethue(color)
        fontsize(30)
        text(
            inword,
            start_point,
            halign=:center,
            valign=:middle
        )
    end
) for (start_point, inword) in zip(start_points, words_in)]
for (start_point, inword) in zip(start_points, inwords)
    act!(inword, Action(time_steps[2]:time_steps[3], anim_scale(0.0)))
    act!(
        inword,
        Action(time_steps[2]:time_steps[3], anim_translate(O - start_point))
    )
end

n_out = length(words_out)
outwords = [
    Object(
        time_steps[3]:time_steps[5],
        @JShape begin
            sethue(color)
            fontsize(30)
            text(
                outword,
                O,
                halign=:center,
                valign=:middle
            )
    end
) for outword in words_out
]
end_points = [
    (min(marg_width, marg_height) ÷ 1.5) * Point(cos(θ), -sin(θ))
    for θ in range(-π/2, π/2, length = n_out)
    ]
for (end_point, outword) in zip(end_points, outwords)
    act!(outword, Action(1:1, anim_scale(0.1)))
    act!(outword, Action(time_steps[1]:time_steps[2], anim_scale(1.0)))
    act!(outword, Action(time_steps[1]:time_steps[2], anim_translate(end_point)))
end
end


function add_constrained_variables(
    dual_problem,
    primal_model,
    variable_parameters,
)
    single_or_vector_variables_types =
        MOIU.sorted_variable_sets_by_cost(dual_problem.dual_model, primal_model)
    params = Set(variable_parameters)
    for (F, S) in single_or_vector_variables_types
        if F === MOI.VectorOfVariables
            _add_constrained_variables(
                dual_problem.primal_dual_map,
                primal_model,
                S,
                params,
            )
        elseif F === MOI.SingleVariable
            _add_constrained_variable(
                dual_problem.primal_dual_map,
                primal_model,
                S,
                params,
            )
        end
    end
end
const NO_CONSTRAINT = CI{Nothing,Nothing}(0)
# Function barrier for the type unstability of `F` and `S`.
function _add_constrained_variables(
    m::PrimalDualMap,
    primal_model,
    ::Type{S},
    params,
) where {S<:MOI.AbstractVectorSet}
    cis = MOI.get(
        primal_model,
        MOI.ListOfConstraintIndices{MOI.VectorOfVariables,S}(),
    )
    for ci in cis
        f = MOI.get(primal_model, MOI.ConstraintFunction(), ci)
        if all(
            vi -> !haskey(m.constrained_var_idx, vi) && !(vi in params),
            f.variables,
        )
            for (i, vi) in enumerate(f.variables)
                m.constrained_var_idx[vi] = (ci, i)
            end
            # Placeholder to indicate this constraint is part of constrained variables,
            # it will be replaced later with a dual constraints
            m.constrained_var_dual[ci] = NO_CONSTRAINT
        end
    end
end
function _add_constrained_variable(
    m::PrimalDualMap,
    primal_model,
    ::Type{S},
    params,
) where {S<:MOI.AbstractScalarSet}
    cis = MOI.get(
        primal_model,
        MOI.ListOfConstraintIndices{MOI.SingleVariable,S}(),
    )
    for ci in cis
        f = MOI.get(primal_model, MOI.ConstraintFunction(), ci)
        if !haskey(m.constrained_var_idx, f.variable) && !(f.variable in params)
            set = MOI.get(primal_model, MOI.ConstraintSet(), ci)
            if !iszero(MOI.constant(set))
                continue
            end
            m.constrained_var_idx[f.variable] = (ci, 1)
            # Placeholder to indicate this constraint is part of constrained variables,
            # it will be replaced later with a dual constraints
            m.constrained_var_dual[ci] = CI{Nothing,Nothing}(0)
        end
    end
end


include("boundary_conditions_adapted.jl")
"""
    calc_boundary_state!(::NumericalFluxSecondOrder, ::Impenetrable{FreeSlip}, ::ModelSetup)
apply free slip boundary condition for velocity
apply zero numerical flux in the normal direction
"""
function calc_boundary_state!(
    ::NumericalFluxSecondOrder,
    ::Impenetrable{FreeSlip},
    ::Union{ModelSetup},
    ::Nothing,
    state⁺,
    gradflux⁺,
    hyperflux⁺,
    aux⁺,
    n⁻,
    state⁻,
    gradflux⁻,
    hyperflux⁻,
    aux⁻,
    t,
    _...,
)
    state⁺.ρu = state⁻.ρu

    return nothing
end

function calc_boundary_state!(
    ::NumericalFluxSecondOrder,
    ::Impenetrable{FreeSlip},
    ::Union{ModelSetup},
    ::ConstantViscosity,
    state⁺,
    gradflux⁺,
    hyperflux⁺,
    aux⁺,
    n⁻,
    state⁻,
    gradflux⁻,
    hyperflux⁻,
    aux⁻,
    t,
    _...,
)
    state⁺.ρu = state⁻.ρu
    gradflux⁺.ν∇u = n⁻ * (@SVector [-0, -0, -0])'

    return nothing
end

"""
    calc_boundary_state!(::NumericalFluxSecondOrder, ::Impenetrable{NoSlip}, ::ModelSetup)
apply no slip boundary condition for velocity
sets ghost point to have no numerical flux on the boundary for U
"""
@inline function calc_boundary_state!(
    ::NumericalFluxSecondOrder,
    ::Impenetrable{NoSlip},
    ::Union{ModelSetup},
    ::Nothing,
    state⁺,
    gradflux⁺,
    hyperflux⁺,
    aux⁺,
    n⁻,
    state⁻,
    gradflux⁻,
    hyperflux⁻,
    aux⁻,
    t,
    _...,
)
    state⁺.ρu = -state⁻.ρu

    return nothing
end

@inline function calc_boundary_state!(
    ::NumericalFluxSecondOrder,
    ::Impenetrable{NoSlip},
    ::Union{ModelSetup},
    ::ConstantViscosity,
    state⁺,
    gradflux⁺,
    hyperflux⁺,
    aux⁺,
    n⁻,
    state⁻,
    gradflux⁻,
    hyperflux⁻,
    aux⁻,
    t,
    _...,
)
    state⁺.ρu = -state⁻.ρu
    gradflux⁺.ν∇u = gradflux⁻.ν∇u

    return nothing
end

"""
    calc_boundary_state!(::NumericalFluxSecondOrder, ::Insulating, ::HBModel)

apply insulating boundary condition for velocity
sets ghost point to have no numerical flux on the boundary for κ∇θ
"""
@inline function calc_boundary_state!(
    ::NumericalFluxSecondOrder,
    ::Union{Insulating, CoupledSecondaryBoundary, CoupledPrimaryBoundary},
    ::Union{ModelSetup},
    ::Nothing,
    state⁺,
    gradflux⁺,
    hyperflux⁺,
    aux⁺,
    n⁻,
    state⁻,
    gradflux⁻,
    hyperflux⁻,
    aux⁻,
    t,
    _...,
)
    state⁺.ρθ = state⁻.ρθ

    return nothing
end

@inline function calc_boundary_state!(
    ::NumericalFluxSecondOrder,
    ::Union{Insulating, CoupledSecondaryBoundary, CoupledPrimaryBoundary},
    ::Union{ModelSetup},
    ::ConstantViscosity,
    state⁺,
    gradflux⁺,
    hyperflux⁺,
    aux⁺,
    n⁻,
    state⁻,
    gradflux⁻,
    hyperflux⁻,
    aux⁻,
    t,
    _...,
)
    state⁺.ρθ = state⁻.ρθ
    gradflux⁺.κ∇θ = n⁻ * -0

    return nothing
end


module Elsa

using Requires

# function __init__()
#     # @require Pardiso = "46dd5b70-b6fb-5a00-ae2d-e8fea33afaf2" include("diagonalizers/pardiso.jl")
#     # @require KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77" include("diagonalizers/krylovkit.jl")
#     @require Arpack = "7d9fca2a-8960-54d3-9f78-7d1dccf2cb97" include("diagonalizers/arpack.jl")
#     # @require ArnoldiMethod = "ec485272-7323-5ecc-a04f-4719b315124d" include("diagonalizers/arnoldimethod.jl")
#     # @require IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153" include("diagonalizers/iterativesolvers.jl")
# end

using StaticArrays, NearestNeighbors, SparseArrays, LinearAlgebra, OffsetArrays,
      ProgressMeter, LinearMaps, Random

using SparseArrays: getcolptr

export sublat, bravais, lattice, dims, hopping, onsite, hamiltonian, randomstate,
       mul!, supercell, unitcell, semibounded, bloch, bloch!, optimize!, similarmatrix,
       sites, bandstructure, marchingmesh, defaultmethod, bands, vertices, states,
       flatten, wrap, transform!, combine,
       momentaKPM, dosKPM, averageKPM, densityKPM

export LatticePresets, RegionPresets, HamiltonianPresets

export LinearAlgebraPackage, ArpackPackage, KrylovKitPackage

export @SMatrix, @SVector, SMatrix, SVector

export ishermitian, I

const NameType = Symbol
const nametype = Symbol

const TOOMANYITERS = 10^8

include("iterators.jl")
include("presets.jl")
include("lattice.jl")
include("model.jl")
include("field.jl")
include("hamiltonian.jl")
include("state.jl")
include("mesh.jl")
include("diagonalizer.jl")
include("bandstructure.jl")
include("KPM.jl")
include("convert.jl")
include("tools.jl")

end


## arquivo modelo

using IL_simulations, PDBTools

top_dir     =  "/home/viniciusp/Documents/doutorado/ANALYSE/Repository/IL_simulations.jl/VSIL/ITP"
pdb_dir     =  "/home/viniciusp/Documents/doutorado/ANALYSE/Repository/IL_simulations.jl/VSIL/PDB"
input_dir   =  "/home/viniciusp/Documents/doutorado/ANALYSE/Repository/IL_simulations.jl/ubiquitin_files"
top_ion_dir =  "/home/viniciusp/Documents/doutorado/ANALYSE/Repository/IL_simulations.jl/ion_files"

atoms   = readPDB("$input_dir/ubq.pdb")
MMP     = mass(atoms)

cation1 = "BMIM"
anion1  = "NO3"
anion2  = "DCA"

# function to calculate the mean MDDF giving a set of MD trajectories
#function molarmass(anion, cation, pdb_dir)
#  atcation = readPDB("$pdb_dir/$(cation)_VSIL.pdb")
#  atanion  = readPDB("$pdb_dir/$(anion)_VSIL.pdb")
#  MMIL = mass(atcation) + mass(atanion)
#  println("molar mass of the ânion is = ", mass(atanion))
#  println("molar mass of the cation is = ", mass(atcation))
#  return MMIL
#end

MMIL  = molarmass(anion1, cation1, pdb_dir)
#MMIL2 = molarmass(anion2, cation1, pdb_dir)
c       =  0.10

## test with one IL/salt and a protein/polymer
#data = IL_simulations.Data_il(protein="$input_dir/ubq.pdb", MMP = MMP, cation = cation1, anion = anion1,MM = MMIL,c=c)
#nions, nwater, sides = prot_elec(data)
#IL_simulations.top(dict, top_dir, input_dir, data, nions, nwater)
#pack_input(data,pdb_dir,nions,nwater,sides)

função para escrever o arquivo

# test with one IL/salt
data = IL_simulations.Data_elec(cation = cation1, anion = anion1, MM = MMIL,c=c)
nions, nwater, sides = IL_simulations.sol_elec(data)
IL_simulations.topelec(dict, top_dir, top_ion_dir, data, nions, nwater)
IL_simulations.pack_input_sol(data,pdb_dir,nions,nwater,sides)
IL_simulations.analyzeIN(pdb_dir,data)

## test with two ILs/salts and a protein/polymer
#data = IL_simulations.Data_il(protein="$input_dir/ubq.pdb", MMP = MMP, cation = cation1, anion = anion1,MM = MM1,c=c)
#nions, nwater, sides = prot_melec(data)
#IL_simulations.top(dict, top_dir, input_dir, data, nions, nwater)
#pack_input(data,pdb_dir,nions,nwater,sides)
#
## test with two ILs/salts
#data = IL_simulations.Data_il(protein="$input_dir/ubq.pdb", MMP = MMP, cation = cation1, anion = anion1,MM = MM1,c=c)
#nions, nwater, sides = sol_melec(data)
#IL_simulations.top(dict, top_dir, input_dir, data, nions, nwater)
#pack_input(data,pdb_dir,nions,nwater,sides)


using OrdinaryDiffEq, Test
using DiffEqProblemLibrary.ODEProblemLibrary: importodeproblems; importodeproblems()
import DiffEqProblemLibrary.ODEProblemLibrary: van

@testset "Stiffness Detection Tests" begin
  prob1 = ODEProblem(van,[0,2.],(0.0,6),inv(0.003))
  # out-of-place test
  function _van(u, p, t)
    μ = p[1]
    [μ*((1-u[2]^2)*u[1] - u[2]),
     1*u[1]]
  end
  prob2 = ODEProblem(_van,[0,2.],(0.0,6),inv(0.003))
  probArr = [prob1, prob2]
  # Test if switching back and forth
  is_switching_fb(sol) = maximum(diff(findall(x->x==2, sol.alg_choice))) > 5
  for prob in probArr
    alg = AutoTsit5(Rodas5(); maxstiffstep=5, maxnonstiffstep=5, stiffalgfirst=true)
    sol = solve(prob, alg)
    @test length(sol.t) < 280
    @test typeof(alg.algs[sol.alg_choice[1]]) <: Rodas5
    @test is_switching_fb(sol)
    sol = solve(prob, AutoDP5(Rodas5(); maxstiffstep=2, maxnonstiffstep=2,
                              stifftol=11//10, nonstifftol=9//10),
                              reltol=1e-5, abstol=1e-5)
    @test length(sol.t) < 625
    @test is_switching_fb(sol)

    sol = solve(prob,AutoVern6(Kvaerno3(); maxstiffstep=4, maxnonstiffstep=4))
    @test length(sol.t) < 690
    @test is_switching_fb(sol)
    sol = solve(prob,AutoVern7(Hairer42(); maxstiffstep=4, maxnonstiffstep=4))
    @test length(sol.t) < 540
    @test is_switching_fb(sol)
    sol = solve(prob,AutoVern8(Rosenbrock23(); maxstiffstep=4, maxnonstiffstep=4))
    @test length(sol.t) < 910
    @test is_switching_fb(sol)
    sol = solve(prob,AutoVern9(KenCarp3(); maxstiffstep=4, maxnonstiffstep=4))
    @test length(sol.t) < 470
    @test is_switching_fb(sol)
  end
end


# This file is a part of TextSearch.jl

export TextModel, VectorModel, trainsize, vocsize,
    TfWeighting, IdfWeighting, TpWeighting,
    FreqWeighting, BinaryLocalWeighting, BinaryGlobalWeighting,
    LocalWeighting, GlobalWeighting,
    fit, prune, prune_select_top

#####
##
## LocalWeighting
##
#####
"""
    LocalWeighting

Abstract type for local weighting
"""
abstract type LocalWeighting end

"""
    TfWeighting()

Term frequency weighting
"""
struct TfWeighting <: LocalWeighting end

"""
    TpWeighting()

Term probability weighting
"""
struct TpWeighting <: LocalWeighting end

"""
    FreqWeighting()

Frequency weighting
"""
struct FreqWeighting <: LocalWeighting end

"""
    BinaryLocalWeighting()

The weight is 1 for known tokens, 0 for out of vocabulary tokens
"""
struct BinaryLocalWeighting <: LocalWeighting end

#####
##
## GlobalWeighting
##
#####
"""
    GlobalWeighting

Abstract type for global weighting
"""
abstract type GlobalWeighting end

"""
    IdfWeighting()

Inverse document frequency weighting
"""
struct IdfWeighting <: GlobalWeighting end

"""
    BinaryGlobalWeighting()

The weight is 1 for known tokens, 0 for out of vocabulary tokens
"""
struct BinaryGlobalWeighting <: GlobalWeighting end

#####
##
## TextModels
##
#####
"""
    Model

An abstract type that represents a weighting model
"""
abstract type TextModel end

mutable struct VectorModel{_G<:GlobalWeighting, _L<:LocalWeighting} <: TextModel
    global_weighting::_G
    local_weighting::_L
    voc::Vocabulary
    maxoccs::Int32
    mindocs::Int32
end

function Base.copy(
        e::VectorModel;
        local_weighting=e.local_weighting,
        global_weighting=e.global_weighting,
        voc=e.voc,
        maxoccs=e.maxoccs,
        mindocs=e.mindocs,
    )
    VectorModel(global_weighting, local_weighting, voc, maxoccs, mindocs)
end

@inline trainsize(model::VectorModel) = trainsize(model.voc)
@inline vocsize(model::VectorModel) = vocsize(model.voc)
@inline Base.length(model::VectorModel) = length(model.voc)
@inline occs(model::VectorModel, tokenID::Integer) = occs(model.voc, tokenID)
@inline ndocs(model::VectorModel, tokenID::Integer) = ndocs(model.voc, tokenID)
@inline weight(model::VectorModel, tokenID::Integer) = weight(model.voc, tokenID)
@inline token(model::VectorModel, tokenID::Integer) = token(model.voc, tokenID)
@inline Base.eachindex(model::VectorModel) = eachindex(model.voc)

"""
    VectorModel(global_weighting::GlobalWeighting, local_weighting::LocalWeighting, voc::Vocabulary; mindocs=1)

Creates a vector model using the input corpus.
"""
function VectorModel(global_weighting::GlobalWeighting, local_weighting::LocalWeighting, voc::Vocabulary; mindocs=1)
    maxocc = convert(Int32, maximum(voc.occs))
    mindocs = convert(Int32, mindocs)

    VectorModel(global_weighting, local_weighting, voc, maxocc, mindocs)
end

function VectorModel(global_weighting::GlobalWeighting, local_weighting::LocalWeighting, textconfig::TextConfig, corpus::AbstractVector; mindocs=1, minbatch=0)
    voc = Vocabulary(textconfig, corpus; minbatch)
    maxocc = convert(Int32, maximum(voc.occs))
    mindocs = convert(Int32, mindocs)
    VectorModel(global_weighting, local_weighting, voc, maxocc, mindocs)
end

Base.show(io::IO, model::VectorModel) = print(io, "{VectorModel global_weighting=$(model.global_weighting), local_weighting=$(model.local_weighting), train-voc=$(vocsize(model)), train-n=$(trainsize(model)), maxoccs=$(model.maxoccs)}")

"""
    prune(model::VectorModel, lowerweight::AbstractFloat)

Creates a new vector model without terms with smaller global weights than `lowerweight`.
"""
function prune(model::VectorModel, lowerweight::AbstractFloat)
    voc = Vocabulary(trainsize(model))
    old = model.voc
    for tokenID in eachindex(old)
        if prune_global_weighting(model, tokenID) >= lowerweight
            push!(voc, old.token[tokenID], old.occs[tokenID], old.ndocs[tokenID], old.weight[tokenID])
        end
    end

    VectorModel(model.global_weighting, model.local_weighting, voc, maximum(voc.occs), model.mindocs)
end

"""
    prune_select_top(model::VectorModel, k::Integer)
    prune_select_top(model::VectorModel, ratio::AbstractFloat)

Creates a new vector model with the best `k` tokens from `model` based on global weighting.
`ratio` is a floating point between 0 and 1 indicating the ratio of the vocabulary to be kept
"""
function prune_select_top(model::VectorModel, k::Integer)
    voc = model.voc
    w = [prune_global_weighting(model, tokenID) for tokenID in eachindex(voc)]
    sort!(w, rev=true)
    prune(model, float(w[k]))
end

prune_select_top(model::VectorModel, ratio::AbstractFloat) =
    prune_select_top(model, floor(Int, length(model.voc) * ratio))

"""
    vectorize(model::VectorModel, bow::BOW; normalize=true, mindocs=model.mindocs, minweight=1e-9) where Tv<:Real

Computes a weighted vector using the given bag of words and the specified weighting scheme.
"""
function vectorize(model::VectorModel{_G,_L}, bow::BOW; normalize=true, mindocs=model.mindocs, minweight=1e-9) where {_G,_L}
    numtokens::Int = 0
    if _L === TpWeighting
        for freq in values(bow)
            numtokens += freq
        end
    end

    maxoccs::Int = 0
    if _L === TfWeighting
        maxoccs = length(bow) == 0 ? 0 : maximum(bow)
    end

    vec = SVEC()
    voc = model.voc
    for (tokenID, freq) in bow
        voc.ndocs[tokenID] < mindocs && continue
        w = local_weighting(model.local_weighting, freq, maxoccs, numtokens) * global_weighting(model, tokenID)
        w >= minweight && (vec[tokenID] = w)
    end

    if length(vec) == 0
        vec[0] = 1f0
    else
        normalize && normalize!(vec)
    end

    vec
end

function vectorize(model::VectorModel, textconfig::TextConfig, text, buff::TextSearchBuffer; normalize=true, mindocs=model.mindocs, minweight=1e-9)
    empty!(buff)
    bow = vectorize(model.voc, textconfig, text, buff)
    vectorize(model, bow; normalize, mindocs, minweight)
end

function vectorize(model::VectorModel, textconfig::TextConfig, text; normalize=true, mindocs=model.mindocs, minweight=1e-9)
    buff = take!(CACHES)
    try
        copy(vectorize(model, textconfig, text, buff; normalize, mindocs, minweight))
    finally
        put!(CACHES, buff)
    end
end

function vectorize_corpus(model::VectorModel, textconfig::TextConfig, corpus::AbstractVector; normalize=true, minbatch=0)
    n = length(corpus)
    V = Vector{SVEC}(undef, n)
    minbatch = getminbatch(minbatch, n)

    #@batch minbatch=minbatch per=thread
    for i in 1:n
        text = corpus[i]
        buff = take!(CACHES)
        try
            V[i] = copy(vectorize(model, textconfig, text, buff; normalize))
        finally
            put!(CACHES, buff)
        end
    end

    V
end

function broadcastable(model::VectorModel)
    (model,)
end

# local weightings: TfWeighting, TpWeighting, FreqWeighting, BinaryLocalWeighting
# global weightings: IdfWeighting, BinaryGlobalWeighting

@inline local_weighting(::TfWeighting, occs, maxoccs, numtokens) = occs / maxoccs
@inline local_weighting(::FreqWeighting, occs, maxoccs, numtokens) = occs
@inline local_weighting(::TpWeighting, occs, maxoccs, numtokens) = occs / numtokens
@inline local_weighting(::BinaryLocalWeighting, occs, maxoccs, numtokens) = 1.0
@inline global_weighting(model::VectorModel{IdfWeighting}, tokenID) = @inbounds log2(trainsize(model) / (1 + ndocs(model, tokenID)))
@inline global_weighting(model::VectorModel{BinaryGlobalWeighting}, tokenID) = 1.0
@inline prune_global_weighting(model::VectorModel, tokenID) = global_weighting(model, tokenID)
@inline prune_global_weighting(model::VectorModel{BinaryGlobalWeighting}, tokenID) = @inbounds 1 / (1 + ndocs(model, tokenID))


verbose = true
 if verbose && !isdefined(Main,:RvSpectML)  println("# Loading RvSpecML")    end
 using RvSpectMLBase
 using EchelleInstruments, EchelleInstruments.NEID
 using RvSpectML
 if verbose println("# Loading NeidSolarScripts")    end
 using NeidSolarScripts
 using NeidSolarScripts.SolarRotation
 using NeidSolarScripts.DifferentialExtinction
 if verbose   println("# Loading other packages")    end
 using CSV, DataFrames, Query, Dates
 #using StatsBase, Statistics, Dates

fits_target_str = "Sun"
 if !@isdefined(target_subdir)
      target_subdir = ""
 end
 if !@isdefined(create_missing_continuum_files)
     create_missing_continuum_files = false
 end
 paths_to_search_for_param = [pwd(),pkgdir(NeidSolarScripts)]

 if verbose println("# Finding what data files are avaliable.")  end
 eval(read_data_paths(paths_to_search=paths_to_search_for_param))
 @assert isdefined(Main,:neid_data_path)
 @assert isdefined(Main,:output_dir)
 output_path = joinpath(neid_data_path, target_subdir,output_dir)
 manifest_filename = joinpath(output_path,"manifest.csv")
 manifest_calib_filename = joinpath(output_path,"manifest_calib.csv")

can_skip_generating_manifest = false
if isfile(manifest_filename) || islink(manifest_filename)
    df_files  = CSV.read(manifest_filename, DataFrame)
    warn_link =  islink(manifest_filename) ? " (a link)" : ""
    println("# Read ", manifest_filename, warn_link, ".")
    @assert size(df_files,1) >= 1
    @assert hasproperty(df_files,:Filename)
    @assert hasproperty(df_files,:target)
    @assert hasproperty(df_files,:bjd)
    @assert hasproperty(df_files,:ssbz)
    @assert hasproperty(df_files,:exptime)
    @assert hasproperty(df_files,:alt_sun)
    @assert hasproperty(df_files,:Δv_diff_ext)
    @assert hasproperty(df_files,:Δfwhm²)
    @assert hasproperty(df_files,:order_snrs)
    if eltype(df_files[!,:order_snrs]) == String
        df_files[!,:order_snrs] = map(i->parse.(Float64,split(df_files[i,:order_snrs][2:end-1],',')),1:size(df_files,1))
    end
    @assert eltype(df_files[!,:order_snrs]) == Vector{Float64}
    println("# Required fields present ", size(df_files), ".  No need to regenerate manifest.")
    can_skip_generating_manifest = true
    if any(isnan.(df_files[!,:order_snrs])) ||
       any(isnan.(df_files[!,:Δfwhm²]))    ||
       any(isnan.(df_files[!,:Δv_diff_ext]))
          can_skip_generating_manifest = false
    end
else
    println("# Will need to generate manifest at ", manifest_filename, ".")
end

#can_skip_generating_manifest = false
if can_skip_generating_manifest && !create_missing_continuum_files
    exit()
end

df_files = NEID.make_manifest(joinpath(neid_data_path, target_subdir))

calibration_target_substrings = ["Etalon","LFC","ThArS","LDLS","FlatBB"]
df_files_calib = df_files |>
        @filter( any(map(ss->contains(_.target,ss),calibration_target_substrings)) ) |>
        DataFrame

df_files_calib[!,:airmass] = fill(NaN,size(df_files_calib,1))
CSV.write(manifest_calib_filename, df_files_calib)

df_files_obs = df_files |>
    @filter( !any(map(ss->contains(_.target,ss),calibration_target_substrings)) ) |>
    DataFrame

df_files_use = df_files_obs |>
      @filter( _.target == fits_target_str ) |>
      #@filter(bjd_first_good <= _.bjd < bjd_last_good) |>
      #@filter( is_good_day(_.Filename) ) |>
      #@take(max_spectra_to_use) |>
      DataFrame

if fits_target_str == "Sun" || fits_target_str == "Solar"
    df_files_use[!,:alt_sun] = calc_solar_alt.(df_files_use.bjd,obs=:WIYN)
    df_files_use[!,:airmass] = DifferentialExtinction.f_airmass.(DifferentialExtinction.f_z.(deg2rad.(df_files_use.alt_sun)))
    println("# Computing differential extinction")
    try
        df_files_use[!,:Δv_diff_ext] = calc_Δv_diff_extinction.(df_files_use.bjd, obs=:WIYN)
    catch
       try
          println("# Encountered error with call to JplHorizons, sleeping and will try again.")
          sleep(60)
          df_files_use[!,:Δv_diff_ext] = calc_Δv_diff_extinction.(df_files_use.bjd, obs=:WIYN)
       catch
          println("# Encountered error with call to JplHorizons, using NaNs to proceed with rest of pipeline.")
          df_files_use[!,:Δv_diff_ext] = fill(NaN,length(df_files_use.bjd))
       end
    end
    println("# FWHM effect")
    df_files_use[!,:Δfwhm²] = CalculateFWHMDifference_SolarRotation_from_obs.(df_files_use.bjd,obs=:WIYN)
end

if create_missing_continuum_files
   df_files_use.continuum_filename = map(fn->joinpath(output_path,"continuum",match(r"(neidL1_\d+[T_]\d+)\.fits$", fn)[1] * "_continuum=afs.jld2"),df_files_use.Filename)
end

#=
if verbose println("# Reading in customized parameters from param.jl.")  end
   if !@isdefined(idx_day_to_use)
       idx_day_to_use = 1
   end
   eval(code_to_include_param_jl(paths_to_search=paths_to_search_for_param))
   if match(r"neidL1_(\d+)[T_](\d+)\.fits$", first(df_files_use.Filename))[1] ==  match(r"neidL1_(\d+)[T_](\d+)\.fits$", last(df_files_use.Filename))[1]
      date_str = match(r"neidL1_(\d+)[T_](\d+)\.fits$", first(df_files_use.Filename))[1]
    else
      date_str = string(match(r"neidL1_(\d+)[T_](\d+)\.fits$", first(df_files_use.Filename))[1]) * "-" * string(match(r"neidL1_(\d+)[T_](\d+)\.fits$", last(df_files_use.Filename))[1])
   end
   outputs["df_files_use"] = df_files_use

   outputs_filename = joinpath(output_dir,"solar_" * date_str * "_new.jld2")
   if isfile(outputs_filename) && false
     times_already_processed = load(outputs_filename, "times")
     files_in_day_to_process = size(df_files_solar_by_day.data[idx_day_to_use],1)
      if files_in_day_to_process == length(times_already_processed)
         println("# Already processed all ", length(times_already_processed), " files for ", date_str)
         exit()
      end
   end
=#

#using NaNMath
println("# Computing order SNRs")
df_files_use[!,:order_snrs] = fill(zeros(0),size(df_files_use,1))
   for (i,row) in enumerate(eachrow(df_files_use))
       #if row.target != "Sun" continue   end
       if any(map(ss->contains(row.target,ss),calibration_target_substrings))
           println("# Warning: target = ", row.target)
           continue
       end
       spec = NEID.read_data(row)
       all_orders = min_order(NEID2D()):max_order(NEID2D())
       pixels = min_pixel_in_order(NEID2D()):max_pixel_in_order(NEID2D())
       order_snr = map(ord->RvSpectMLBase.calc_snr(spec, pixels, ord), all_orders)
       df_files_use[i,:order_snrs] = order_snr
   end

println("# Writing manifest file.")
CSV.write(manifest_filename, df_files_use)
      # Can extract vector of Reals turned into a string via:
      # df2 = CSV.read("test.csv", DataFrame)
      # map(i->parse.(Float64,split(df2[i,:order_snrs][2:end-1],',')),1:size(df2,1))

#snr_matrix = hcat(df_files_use[!,:order_snrs]...)

#=
using JLD2, FileIO
if create_missing_continuum_files
   println("# Creating missing continuum files.")
   for row in eachrow(df_files_use)
     if isfile(row.continuum_filename) continue end
     println("# Need to make ", row.continuum_filename )

     output_filename = row.continuum_filename
     tmpdir = basename(output_filename)
     #=
     if !isdir(tmpdir)
        println("# Making ", tmpdir)
        mkdir(tmpdir)
     end
     =#
     #continue
     spec = NEID.read_data(row.Filename)
     try
        continuum = Continuum.calc_continuum_model(spec)
        @save output_filename continuum
     catch
        println("# ERROR computing continuum for ", output_filename)
     end
   end
end

=#