module BEASTXMLConstructor

export BEASTXMLElement,
       MCMCOptions,
       save_xml,
       make_residual_xml,
       make_pfa_xml,
       make_joint_xml,
       set_mbd_precision,
       set_residual_precision,
       set_loadings,
       set_factor_precisions,
       DiffusionModel,
       ResidualVarianceModel,
       IntegratedFactorModel,
       JointProcessModel,
       DataModel,
       tip_dimension,
       data_dimension,
       set_options!,
       set_chain_length!,
       get_operators,
       set_file_logEvery,
       add_factor_proportion

using LightXML, LinearAlgebra, UnPack, StatsBase

include("BeastNames.jl")
bn = BeastNames

XMLOrNothing = Union{XMLElement, Nothing}

abstract type MyXMLElement end
abstract type SimpleXMLElement <: MyXMLElement end
abstract type AbstractDataXMLElement <: MyXMLElement end

import Base: copy

import LightXML: name
function name(x::MyXMLElement)
    error("Not yet implemented for object of type " * string(typeof(x)))
end

mutable struct BEASTXMLElement
    el::XMLOrNothing
    components::Vector{MyXMLElement}
end

mutable struct MCMCOptions
    chain_length::Int
    file_log_every::Int
    screen_log_every::Int
    likelihood_check_count::Int
end

function MCMCOptions(;
                     chain_length::Int = 10_000,
                     file_log_every::Int = max(1, div(chain_length, 1_000)),
                     screen_log_every::Int = max(1, div(chain_length, 100)),
                     likelihood_check_count::Int = -1)

    if likelihood_check_count == -1
        check_count = div(chain_length, 10)
        check_count = max(check_count, 100)
        check_count = min(check_count, 1000)
    else
        check_count = likelihood_check_count
    end

    return MCMCOptions(chain_length, file_log_every, screen_log_every,
                       likelihood_check_count)
end

import Base.show

function Base.show(io::IO, bx::BEASTXMLElement)
    disp = "$(typeof(bx)) with components:"
    for comp in bx.components
        disp *= "\n\t$(typeof(comp))"
    end
    println(io, disp)
end

function BEASTXMLElement()
    return BEASTXMLElement(nothing, MyXMLElement[])
end

abstract type ModelExtensionXMLElement <: MyXMLElement end
abstract type OperatorXMLElement <: MyXMLElement end

const element_dir = joinpath(@__DIR__, "XMLConstructor", "XMLElements")
const dir = joinpath(@__DIR__, "XMLConstructor")

include(joinpath(dir, "constructor_helpers.jl"))
include(joinpath(dir, "integration.jl"))

include(joinpath(element_dir, "NothingXMLElement.jl"))
include(joinpath(element_dir, "CompoundXMLElement.jl"))
include(joinpath(element_dir, "MatrixParameter.jl"))
include(joinpath(element_dir, "DataXMLElement.jl"))
include(joinpath(element_dir, "NewickXMLElement.jl"))
include(joinpath(element_dir, "TreeModelXMLElement.jl"))
include(joinpath(element_dir, "WishartPriorXMLElement.jl"))
include(joinpath(element_dir, "MBDXMLElement.jl"))
include(joinpath(element_dir, "RepeatedMeasuresXMLElement.jl"))
include(joinpath(element_dir, "MatrixShrinkage.jl"))
include(joinpath(element_dir, "IntegratedFactorsXMLElement.jl"))
include(joinpath(element_dir, "TraitLikelihoodXMLElement.jl"))
include(joinpath(element_dir, "LatentFactorModelXMLElement.jl"))
include(joinpath(element_dir, "LogPredictiveDensity.jl"))
include(joinpath(element_dir, "NormalGammaPrecisionOperatorXMLElement.jl"))
include(joinpath(element_dir, "PrecisionGibbsOperatorXMLElement.jl"))
include(joinpath(element_dir, "OperatorsXMLElement.jl"))
include(joinpath(element_dir, "LoggablesXMLElement.jl"))
include(joinpath(element_dir, "MCMCXMLElement.jl"))
include(joinpath(element_dir, "TimerXMLElement.jl"))
include(joinpath(element_dir, "TraitValidationXMLElement.jl"))
include(joinpath(element_dir, "OldLoadingsGibbsOperatorXMLElement.jl"))
include(joinpath(element_dir, "HMCOperatorXMLElement.jl"))
include(joinpath(element_dir, "LoadingsGibbsOperatorXMLElement.jl"))
include(joinpath(element_dir, "LatentFactorPrecisionOperatorXMLElement.jl"))
include(joinpath(element_dir, "FactorTreeGibbsOperatorXMLElement.jl"))
include(joinpath(element_dir, "TraitLoggerXMLElement.jl"))
include(joinpath(element_dir, "LKJPrecisionXMLElement.jl"))
include(joinpath(element_dir, "NormalMatrixNormLikelihood.jl"))
include(joinpath(element_dir, "ScaledOrthogonalMatrix.jl"))
include(joinpath(element_dir, "MultivariateGammaLikelihood.jl"))
include(joinpath(element_dir, "NormalDistributionLikelihood.jl"))
include(joinpath(element_dir, "MultiplicativeScalePrior.jl"))
include(joinpath(element_dir, "ForceOrderedLikelihood.jl"))
include(joinpath(element_dir, "IntegratedFactorsGibbsOperator.jl"))
include(joinpath(element_dir, "IndependentNormalDistributionModel.jl"))
include(joinpath(element_dir, "LoadingsScaleGibbsOperator.jl"))
include(joinpath(element_dir, "LatentLiability.jl"))
include(joinpath(element_dir, "FactorProportionStatistic.jl"))

include(joinpath(dir, "utils.jl"))
include(joinpath(dir, "constructors.jl"))
include(joinpath(dir, "getters.jl"))
include(joinpath(dir, "setters.jl"))

function make_xml(bx::BEASTXMLElement)
    xdoc = XMLDocument()
    bx.el = new_element(bn.BEAST)
    set_root(xdoc, bx.el)

    for el in bx.components
    add_el(bx, el)
    end

    return xdoc
end

function save_xml(path::String, bx::BEASTXMLElement;
                  change_filename::Bool = true)

    nm = basename(path)
    s = split(nm, '.')

    if length(s) == 1
        path = path * ".xml"
    elseif length(s) == 2
        @assert s[2] == "xml"
    end

    if change_filename
        filename = s[1]
        set_filename(bx, filename)
    end

    xdoc = make_xml(bx)
    save_file(xdoc, path)
    free(xdoc)
    return nothing
end

import LightXML: add_child
function add_child(xel::XMLElement, no::Nothing)
    # do nothing
end

function add_child(bx::BEASTXMLElement, el::MyXMLElement)
    push!(bx.components, el)
end

function add_child(bx::BEASTXMLElement, el::MyXMLElement, ind::Int)
    insert!(bx.components, ind, el)
end

import LightXML: find_element

function find_element(bx::BEASTXMLElement, type::Type)
    ind = find_element_ind(bx, type)
    if isnothing(ind)
        return nothing
    end
    return bx.components[find_element_ind(bx, type)]
end

function find_element_ind(bx::BEASTXMLElement, type::Type)
    if !(type <: MyXMLElement)
        error("Can only find objects of type MyXMLElement.")
    end

    found = findall(x -> typeof(x) <: type, bx.components)

    if length(found) == 0
        return nothing
    elseif length(found) > 1
        error("Multiple elements of type $type. " *
              "Please use function 'find_elements'.")
    end

    return found[1]
end

function find_elements(bx::BEASTXMLElement, type::Type)
    found = findall(x -> typeof(x) <: type, bx.components)
    return bx.components[found]
end

function add_parameter_id(pel::XMLElement, id::String)
    el = new_element(bn.PARAMETER)
    set_attribute(el, bn.ID, id)
    add_child(pel, el)
    return el
end

function add_loggable(bx::BEASTXMLElement, el::MyXMLElement;
                        already_made::Bool = false)

    loggables = LoggablesXMLElement([el], [already_made])
    add_loggables(bx, loggables)
end

function add_loggables(bx::BEASTXMLElement, loggables::LoggablesXMLElement)
    old_loggables = find_elements(bx, LoggablesXMLElement)
    if length(old_loggables) == 0
        mcmc = get_mcmc(bx)
        ind = findfirst(x -> x === mcmc, bx.components)
        mcmc_loggables = mcmc.loggables
        push!(mcmc_loggables, loggables)
        insert!(bx.components, ind, mcmc_loggables)
    elseif length(old_loggables) == 1
        push!(old_loggables[1], loggables)
    else
        error("Cannot join with more than one LoggablesXMLElement")
    end
    return nothing
end

function make_xml(n::Nothing)
    # do nothing
end

function add_el(bx::BEASTXMLElement, no::Nothing)
    #do nothing
end

function add_el(bx::BEASTXMLElement, el::TextXMLElement)
    add_child(bx.el, el.el)
end

function add_el(bx::BEASTXMLElement, el::MyXMLElement)
    make_xml(el)
    add_child(bx.el, el.el)
end

function add_el(bx::BEASTXMLElement, mbd_el::MBDXMLElement)
    make_xml(mbd_el)
    add_child(bx.el, mbd_el.el)
    add_el(bx, mbd_el.precision_prior)
end

function add_el(bx::BEASTXMLElement, n_el::NothingXMLElement)
    # do nothing
end

function add_el(bx::BEASTXMLElement, cp_el::CompoundXMLElement)
    for el in cp_el.els
        add_el(bx, el)
    end
end

function add_el(bx::BEASTXMLElement, rm_el::RepeatedMeasuresXMLElement)
    make_xml(rm_el)
    add_child(bx.el, rm_el.el)
    add_child(bx.el, rm_el.precision_prior.el)
end

function add_el(bx::BEASTXMLElement, if_el::IntegratedFactorsXMLElement)
    make_xml(if_el)
    add_el(bx, if_el.loadings)
    add_el(bx, if_el.loadings_prior)
    add_child(bx.el, if_el.el)
    add_child(bx.el, if_el.precision_prior_el)
end

function add_el(bx::BEASTXMLElement, lfm_el::LatentFactorModelXMLElement)
    make_xml(lfm_el)
    if !lfm_el.parameters_already_made
        add_el(bx, lfm_el.loadings)
        add_el(bx, lfm_el.loadings_prior)
    end
    add_child(bx.el, lfm_el.el)
    if !lfm_el.parameters_already_made
        add_child(bx.el, lfm_el.precision_prior_el)
    end
end

function add_el(bx::BEASTXMLElement, el::FactorLogPredictiveDensity)
    make_xml(el)
    add_child(bx.el, el.fac_el)
    add_child(bx.el, el.like_el)
    add_child(bx.el, el.cl_el)
end

function add_el(bx::BEASTXMLElement, els::Array{MyXMLElement})
    for el in els
        add_el(bx, el)
    end
end

function add_el(bx::BEASTXMLElement, loggables::LoggablesXMLElement)
    for i = 1:length(loggables.els)
        if !loggables.already_made[i]
            l = loggables.els[i]
            add_el(bx, l)
        end
    end
end

function add_el(bx::BEASTXMLElement, el::ScaledOrthogonalMatrix)
    make_xml(el)
    add_child(bx.el, el.scale.el)
    add_child(bx.el, el.U.el)
    add_child(bx.el, el.el)
end

function add_el(bx::BEASTXMLElement, el::MatrixShrinkageLikelihoods)
    make_xml(el)
    for mult in el.mult_els
        add_child(bx.el, mult)
    end
    for gp in el.gp_els
        add_child(bx.el, gp)
    end
    for ls in el.local_scale_els
        add_child(bx.el, ls)
    end
    for gs in el.global_scale_els
        add_child(bx.el, gs)
    end
    for bb in el.bb_els
        add_child(bx.el, bb)
    end
    add_child(bx.el, el.ms_el)
    for gp in el.global_prior_els
        add_child(bx.el, gp)
    end
end

function add_el(bx::BEASTXMLElement, el::MultiplicativeScalePrior)
    add_el(bx, el.mults)
    add_el(bx, el.precs)
    add_el(bx, el.mult_prior)
    add_el(bx, el.scale_prior)
end

end


using Base: Float64
using AlphaZero
using ProgressLogging

include("../mu_game_wrapper.jl")
include("../network.jl")
include("../trace.jl")
include("../play.jl")
#include("memory.jl")
include("../training.jl")

# experiment = Examples.experiments["tictactoe"]
# session = Session(experiment)
# env = session.env

# oracle = Network.copy(env.bestnn, on_gpu=env.params.self_play.sim.use_gpu, test_mode=true)
# player = MctsPlayer(env.gspec, oracle, env.params.self_play.mcts)

# game = GI.init(env.gspec)

# function AlphaZero.think(p::MctsPlayer, game)
#     hidden_state = Representation()(GI.current_state(game))
#     mugame = MuGameEnvWrapper(game,
#         CachedOracle(Dynamics(gspec=env.gspec),Dict{Tuple{typeof(hidden_state), Int},Tuple{Float64, typeof(hidden_state)}}()),
#         Representation(),
#         hidden_state,
#         true,
#         GI.white_playing(game),
#         0.)

#     if isnothing(p.timeout) # Fixed number of MCTS simulations
#         MCTS.explore!(p.mcts, mugame, p.niters)
#     else # Run simulations until timeout
#         start = time()
#         while time() - start < p.timeout
#         MCTS.explore!(p.mcts, mugame, p.niters)
#         end
#     end
#     return MCTS.policy(p.mcts, mugame)
# end

# actions, π_target = think(player, game) #overload this function, perhaps MuMCTSplayer

gspec = Examples.games["tictactoe"]

representationhp = RepresentationHP(width=200, depth=4, hiddenstate_shape=32)
representation_oracle = RepresentationNetwork(gspec, representationhp)

dynamicshp = DynamicsHP(hiddenstate_shape=32, width=200, depth_common=4)
dynamics_oracle = DynamicsNetwork(gspec, dynamicshp)

predictionhp = PredictionHP(hiddenstate_shape=32, width=200, depth_common=4)
prediction_oracle = PredictionNetwork(gspec, predictionhp)

nns = (f=prediction_oracle, g=dynamics_oracle, h=representation_oracle)

import LinearAlgebra
LinearAlgebra.BLAS.set_num_threads(1)

n=10
sim=SimParams(
    num_games=10,
    num_workers=n,
    batch_size=n,
    use_gpu=false,
    reset_every=4,
    flip_probability=0.,
    alternate_colors=false)

  mcts = MctsParams(
    num_iters_per_turn=400,
    cpuct=1.0,
    temperature=ConstSchedule(1.0),
    dirichlet_noise_ϵ=0.2,
    dirichlet_noise_α=1.0)

  self_play = (; sim, mcts)

  #memory = CircularBuffer{Trace{GI.state_type(gspec)}}(1024)

  μparams = MuParams(self_play, 1, 100)

  env = MuEnv(gspec, μparams, nns)

  t = @timed self_play_step!(env)
  # push!(res, (n, t))

n_workers = 1:16
res2 = []
for n in n_workers
  sim=SimParams(
    num_games=100,
    num_workers=n,
    batch_size=n,
    use_gpu=false,
    reset_every=4,
    flip_probability=0.,
    alternate_colors=false)

  mcts = MctsParams(
    num_iters_per_turn=400,
    cpuct=1.0,
    temperature=ConstSchedule(1.0),
    dirichlet_noise_ϵ=0.2,
    dirichlet_noise_α=1.0)

  self_play = (; sim, mcts)

  #memory = CircularBuffer{Trace{GI.state_type(gspec)}}(1024)

  μparams = MuParams(self_play, 1, 100)

  env = MuEnv(gspec, μparams, nns)

  t = @timed self_play_step!(env)
  push!(res2, (n,t))
  @info res2

end

using Plots
# bar(first.(res), last.(res), xlabel="num_workers", ylabel="time [s]", label="num_games=100")
bar(first.(res2), last.(res2), xlabel="num_workers", ylabel="time [s]", label="num_games=100")


"""
    gbset(A::GBArray, option, value)
    gbset(option, value)

Set an option either for a specific GBArray, or globally. The commonly used options are:
    - `:format = [:byrow | :bycol]`: The global default or array specific
    column major or row major ordering.
    - `:nthreads = [Integer]`: The global number of OpenMP threads to use.
    - `:burble = [Bool]`: Print diagnostic output.
    - `:sparsity_control = [:full | :bitmap | :sparse | :hypersparse]`: Set the sparsity of a
    single GBArray.
"""
gbset

function gbset(option, value)
    option = option_toconst(option)
    value = option_toconst(value)
    libgb.GxB_Global_Option_set(option, value)
    return nothing
end

function gbget(option)
    option = option_toconst(option)
    return libgb.GxB_Global_Option_get(option)
end

function gbset(A::GBMatrix, option, value)
    option = option_toconst(option)
    value = option_toconst(value)
    libgb.GxB_Matrix_Option_set(A, option, value)
    return nothing
end

function gbget(A::GBMatrix, option)
    option = option_toconst(option)
    return libgb.GxB_Matrix_Option_get(A, option)
end

function gbset(A::GBVector, option, value)
    option = option_toconst(option)
    value = option_toconst(value)
    libgb.GxB_Matrix_Option_set(A, option, value)
    return nothing
end

function gbget(A::GBVector, option)
    option = option_toconst(option)
    return libgb.GxB_Matrix_Option_get(A, option)
end

function format(A::GBVecOrMat)
    t = gbget(A, SPARSITY_STATUS)
    f = gbget(A, FORMAT)
    return (GBSparsity(t), GBFormat(f))
end

const HYPER_SWITCH = libgb.GxB_HYPER_SWITCH
const BITMAP_SWITCH = libgb.GxB_BITMAP_SWITCH
const FORMAT = libgb.GxB_FORMAT
const SPARSITY_STATUS = libgb.GxB_SPARSITY_STATUS
const SPARSITY_CONTROL = libgb.GxB_SPARSITY_CONTROL
const BASE1 = libgb.GxB_PRINT_1BASED
const NTHREADS = libgb.GxB_GLOBAL_NTHREADS
const BURBLE = libgb.GxB_BURBLE

const BYROW = libgb.GxB_BY_ROW
const BYCOL = libgb.GxB_BY_COL

#only translate if it's a symbol
option_toconst(option) = option
function option_toconst(sym::Symbol)
    sym === :format && return FORMAT
    sym === :nthreads && return NTHREADS
    sym === :burble && return BURBLE
    sym === :byrow && return BYROW
    sym === :bycol && return BYCOL
    sym === :sparsity_status && return SPARSITY_STATUS
    sym === :sparsity_control && return SPARSITY_CONTROL
    sym === :full && return GBDENSE
    sym === :bitmap && return GBBITMAP
    sym === :sparse && return GBSPARSE
    sym === :hypersparse && return GBHYPER
end

"""
Sparsity options for GraphBLAS. values can be summed to produce additional options.
"""
@enum GBSparsity::Int32 begin
    GBDENSE = 8 #libgb.GxB_FULL
    GBBITMAP = 4 #libgb.GxB_BITMAP
    GBSPARSE = 2 #libgb.GxB_SPARSE
    GBHYPER = 1 #libgb.GxB_HYPERSPARSE
    GBANYSPARSITY = 15 #libgb.GxB_ANY_SPARSITY
    GBDENSE_OR_BITMAP = 12 #libgb.GxB_FULL + libgb.GXB_BITMAP
    GBSPARSE_OR_HYPER = 3 #libgb.GxB_SPARSE + libgb.GXB_HYPERSPARSE
    #... Probably don't need others
end


module CIFAR10Reader

export

    readdata!,
    readdata

const NROW = 32
const NCOL = 32
const NCHAN = 3
const NBYTE = NROW * NCOL * NCHAN + 1 # "+ 1" for label
const CHUNK_SIZE = 10_000

function readnext!(buffer::Array{UInt8}, io::IO)
    y = Int(read(io, UInt8))
    read!(io, buffer)
    buffer, y
end

function readdata!(buffer::Array{UInt8}, io::IO, index::Integer)
    seek(io, (index - 1) * NBYTE)
    readnext!(buffer, io)
end

function readdata(io::IO, index::Integer)
    buffer = Array{UInt8}(undef, NROW, NCOL, NCHAN)
    readdata!(buffer, io, index)
end

function readdata(io::IO)
    X = Array{UInt8}(undef, NROW, NCOL, NCHAN, CHUNK_SIZE)
    Y = Array{Int}(undef, CHUNK_SIZE)
    buffer = Array{UInt8}(undef, NROW, NCOL, NCHAN)
    @inbounds for index in 1:CHUNK_SIZE
        _, ty = readnext!(buffer, io)
        copyto!(view(X,:,:,:,index), buffer)
        Y[index] = ty
    end
    X, Y
end

function readdata(file::AbstractString, index::Integer)
    open(file, "r") do io
        readdata(io, index)
    end::Tuple{Array{UInt8,3},Int}
end

function readdata(file::AbstractString)
    open(file, "r") do io
        readdata(io)
    end::Tuple{Array{UInt8,4},Vector{Int}}
end

end


module Plotting

using Reexport
@reexport using Plots
@reexport using StructArrays
@reexport using RecursiveArrayTools
@reexport using LaTeXStrings

using Flight.Modeling

export TimeHistory
export plots, save_plots, thplot, thplot!

########################## plots method ###############################

#our entry plotting entry point cannot be a recipe. a recipe is called within
#the plot() pipeline, which creates a single figure. however, a Vector of System
#outputs will typically need to generate multiple plots from its values. for
#this we define a new method plots(). maybe define different detail levels /
#modes for

plots(args...; kwargs...) = println("Not implemented")

function plots(mdl::Model; mode::Symbol = :basic,
    save_path::Union{String,Nothing} = nothing, kwargs...)
    #generate default path tmp/plots/current_date
    save_path = (save_path === nothing ?
        joinpath("tmp", Dates.format(now(), "yyyy_mm_dd_HHMMSS")) : save_path)
    mkpath(save_path)
    plots(mdl.log.t, mdl.log.saveval; mode, save_path, kwargs...)
end

#for Systems without outputs
plots(::AbstractVector{<:Real}, ::AbstractVector{Nothing}; kwargs...) = nothing

#takes care of SystemGroups
function plots(t, data::AbstractVector{<:NamedTuple}; mode, save_path, kwargs...)

    c = data |> StructArray |> StructArrays.components
    for (c_label, c_data) in zip(keys(c), values(c))
        println("Generating plots for $c_label")
        save_path_c = mkpath(joinpath(save_path, String(c_label)))
        plots(t, c_data; mode, save_path = save_path_c, kwargs...)
    end

end

function save_plots(d::Dict{String,Plots.Plot}; save_path, format = :png)
    for (id, p) in zip(keys(d), values(d))
        savefig(p, joinpath(save_path, id*"."*String(format)))
    end
end

############################ TimeHistory #################################

mutable struct TimeHistory{D}
    t::AbstractVector{<:Real}
    data::D
end

#for all the following TimeHistory subtypes, a single figure is enough, so they
#can be handled by recipes
@recipe function f(th::TimeHistory{<:AbstractVector{<:Real}})

    xguide --> L"$t \: (s)$"
    return th.t, th.data

end

@recipe function f(th::TimeHistory{<:AbstractMatrix{<:Real}}; th_split = :none)

    xguide --> L"$t \ (s)$"

    vlength = size(th.data)[2]
    label --> (vlength <= 3 ?  ["x" "y" "z"][:, 1:vlength] : (1:vlength)')
    if th_split === :h
        layout --> (1, vlength)
        link --> :y #alternative: :none
    elseif th_split === :v
        layout --> (vlength, 1)
    else
        layout --> 1
    end

    return th.t, th.data

end

#convert to TimeHistory{Matrix} and return it to the pipeline for dispatching
@recipe function f(th::TimeHistory{<:AbstractVector{<:AbstractVector{<:Real}}})
    return TimeHistory(th.t, Array(VectorOfArray(th.data))')
end

thplot(t, data; kwargs...) = plot(TimeHistory(t, data); kwargs...)
thplot!(t, data; kwargs...) = plot!(TimeHistory(t, data); kwargs...)

end


#MDDatasets: Support broadcast of DataRS types
#-------------------------------------------------------------------------------

#==result_type: Figure out result type for a function call with given arguments
===============================================================================#
#=NOTES:
 - result_type never returns concrete DataRS{T}.  Only abstract "DataRS".

 - The default return types specified here are not necessarily correct.
   Methods should be overwritten for special functions.
   Sadly: can't currently dispatch on function.  Need to think about this.
=#

result_type(fn::Function, ::Type{T1}, ::Type{T2}) where {T1<:DataRS, T2<:DataRS} = DataRS
result_type(fn::Function, ::Type{T1}, ::Type{T2}) where {T1<:DataRS, T2} = DataRS
result_type(fn::Function, ::Type{T1}, ::Type{T2}) where {T1, T2<:DataRS} = DataRS

#Reducting functions
#-------------------------------------------------------------------------------
result_type(CT::CastTypeRed, fn::Function, ::Type{DataRS{T1}}) where T1<:Number = T1
result_type(CT::CastTypeRed, fn::Function, ::Type{DataRS{T}}) where T = DataRS

#==eltype/valtype: Figure out result type for a function call with given arguments
===============================================================================#

#Element type of an operation with a DataRS (result: DataRS):
Base.eltype(fn::Function, d1::DataRS) = result_type(fn, eltype(d1))

Base.eltype(fn::Function, d1::DataRS, d2::DataRS) =
	result_type(fn, eltype(d1), eltype(d2))
Base.eltype(fn::Function, d1::DataRS, d2) =
	result_type(fn, eltype(d1), typeof(d2))
Base.eltype(fn::Function, d1, d2::DataRS) =
	result_type(fn, typeof(d1), eltype(d2))

#Reducting functions
#-------------------------------------------------------------------------------
#Element type afer reducing fn on Array{Number}:
Base.eltype(CT::CastTypeRed, fn::Function, a::Vector{T}) where T<:Number = result_type(CT, fn, typeof(a))

#Element type afer reducing fn on Array{DataF1}:
function Base.eltype(CT::CastTypeRed, fn::Function, a::Array{DataF1})
	result = Set{DataType}()
	for elem in a
		RT = result_type(CT, fn, typeof(elem))
		push!(result, RT)
	end
	return promote_type(result...)
end

#Element type afer reducing fn on Array{DataRS}:
function Base.eltype(CT::CastTypeRed, fn::Function, a::Vector{DataRS})
	result = Set{DataType}()
	RT = result_type(CT, fn, typeof(a[1]))
	if DataRS == RT; return DataRS; end
	for elem in a
		RT = result_type(CT, fn, typeof(elem))
		push!(result, RT)
	end
	return promote_type(result...)
end

#==Broadcast tools for fn(DataMD) - Can broadcast on 1 argument
===============================================================================#

#Broadcast functions capable of operating directly on 1 base type (Number):
#-------------------------------------------------------------------------------
#fn(DataRS) - core: fn(Number):
function broadcastMD(CT::CastType1{Number,1}, fn::Function, d::DataRS, args...; kwargs...)
	result = DataRS{eltype(fn, d)}(d.sweep)
	for i in 1:length(d.sweep)
		result.elem[i] = broadcastMD(CT, fn, d.elem[i], args...; kwargs...)
	end
	return result
end

#Data reducing fn(DataRS{Number}) - core: fn(Number):
broadcastMD(CT::CastTypeRed1{Number,1}, fn::Function, d::DataRS{T}, args...; kwargs...) where T<:Number =
	fn(d.elem, args...; kwargs...)

#Data reducing fn(DataRS) - core: fn(Number):
function broadcastMD(CT::CastTypeRed1{Number,1}, fn::Function, d::DataRS, args...; kwargs...)
	RT = eltype(CT, fn, d.elem)
	result = DataRS{RT}(d.sweep)
	for i in 1:length(d.sweep)
		result.elem[i] = broadcastMD(CT, fn, d.elem[i], args...; kwargs...)
	end
	return result
end

#Broadcast functions capable of operating only on a dataF1 value:
#-------------------------------------------------------------------------------
#fn(DataRS) - core: fn(DataF1):
function broadcastMD(CT::CastType1{DataF1,1}, fn::Function, d::DataRS{T}, args...; kwargs...) where T
	result = DataRS{T}(d.sweep)
	for i in 1:length(d.sweep)
		result.elem[i] = broadcastMD(CT, fn, d.elem[i], args...; kwargs...)
	end
	return result
end
#fn(???, DataRS) - core: fn(???, DataF1):
function broadcastMD(CT::CastType1{DataF1,2}, fn::Function, dany1, d::DataRS{T}, args...; kwargs...) where T
	result = DataRS{T}(d.sweep)
	for i in 1:length(d.sweep)
		result.elem[i] = broadcastMD(CT, fn, dany1, d.elem[i], args...; kwargs...)
	end
	return result
end

#Data reducing fn(DataRS) - core: fn(Number):
broadcastMD(CT::CastTypeRed1{DataF1,1}, fn::Function, d::DataRS{T}, args...; kwargs...) where T<:Number =
	fn(DataF1(d.sweep.v, d.elem), args...; kwargs...)

#Data reducing fn(DataRS) - core: fn(Number):
function broadcastMD(CT::CastTypeRed1{DataF1,1}, fn::Function, d::DataRS, args...; kwargs...)
	RT = eltype(CT, fn, d.elem)
	result = DataRS{RT}(d.sweep)
	for i in 1:length(d.sweep)
		result.elem[i] = broadcastMD(CT, fn, d.elem[i], args...; kwargs...)
	end
	return result
end

#==Broadcast tools for fn(DataMD, DataMD) - Can broadcast on 2 arguments
===============================================================================#

#Broadcast functions capable of operating directly on base types (Number, Number):
#-------------------------------------------------------------------------------
#fn(DataRS, DataRS) - core: fn(Number, Number):
function broadcastMD(CT::CastType2{Number,1,Number,2}, fn::Function,
	d1::DataRS, d2::DataRS, args...; kwargs...)
	if d1.sweep != d2.sweep
		msg = "Sweeps do not match (not yet supported):"
		throw(ArgumentError(string(msg, "\n", d1.sweep, "\n", d2.sweep)))
	end
	result = DataRS{eltype(fn, d1, d2)}(d1.sweep)
	for i in 1:length(d1.sweep)
		result.elem[i] = broadcastMD(CT, fn, d1.elem[i], d2.elem[i], args...; kwargs...)
	end
	return result
end

#fn(DataRS, DataF1/Number) - core: fn(Number, Number):
function broadcastMD(CT::CastType2{Number,1,Number,2}, fn::Function,
	d1::DataRS, d2::DF1_Num, args...; kwargs...)
	result = DataRS{eltype(fn, d1, d2)}(d1.sweep)
	for i in 1:length(d1.sweep)
		result.elem[i] = broadcastMD(CT, fn, d1.elem[i], d2, args...; kwargs...)
	end
	return result
end

#fn(DataF1/Number, DataRS) - core: fn(Number, Number):
function broadcastMD(CT::CastType2{Number,1,Number,2}, fn::Function,
	d1::DF1_Num, d2::DataRS, args...; kwargs...)
	result = DataRS{eltype(fn, d1, d2)}(d2.sweep)
	for i in 1:length(d2.sweep)
		result.elem[i] = broadcastMD(CT, fn, d1, d2.elem[i], args...; kwargs...)
	end
	return result
end

#Broadcast functions capable of operating on DataF1 values:
#-------------------------------------------------------------------------------
const DF1_DRS = Union{DataF1,DataRS}
#fn(DataRS, DataRS) - core: fn(DataF1, DataF1):
function broadcastMD(CT::CastType2{DataF1,1,DataF1,2}, fn::Function,
	d1::DataRS{T1}, d2::DataRS{T2}, args...; kwargs...) where {T1<:DF1_DRS,T2<:DF1_DRS}
	if d1.sweep != d2.sweep
		msg = "Sweeps do not match (not yet supported):"
		throw(ArgumentError(string(msg, "\n", d1.sweep, "\n", d2.sweep)))
	end
	result = DataRS{eltype(fn, d1, d2)}(d1.sweep)
	for i in 1:length(d1.sweep)
		result.elem[i] = broadcastMD(CT, fn, d1.elem[i], d2.elem[i], args...; kwargs...)
	end
	return result
end
function broadcastMD(CT::CastType2{DataF1,1,DataF1,2}, fn::Function,
	d1::DataRS{T1}, d2::DataRS{T2}, args...; kwargs...) where {T1<:Number,T2<:Number}

	if d1.sweep.id != d2.sweep.id
		msg = "Sweep ids do not match:"
		throw(ArgumentError(string(msg, "\n", d1.sweep, "\n", d2.sweep)))
	end
	return fn(DataF1(d1.sweep.v, d1.elem), DataF1(d2.sweep.v, d2.elem), args...; kwargs...)
end
function broadcastMD(CT::CastType2{DataF1,1,DataF1,2}, fn::Function,
	d1::DataRS{T1}, d2::DataMD, args...; kwargs...) where T1<:Number
	return broadcastMD(CT, fn, DataF1(d1.sweep.v, d1.elem), d2, args...; kwargs...)
end
function broadcastMD(CT::CastType2{DataF1,1,DataF1,2}, fn::Function,
	d1::DataMD, d2::DataRS{T2}, args...; kwargs...) where T2<:Number
	return broadcastMD(CT, fn, d1, DataF1(d2.sweep.v, d2.elem), args...; kwargs...)
end

#fn(DataRS, DataF1) - core: fn(DataF1, DataF1):
function broadcastMD(CT::CastType2{DataF1,1,DataF1,2}, fn::Function,
	d1::DataRS{TRS}, d2::DataF1, args...; kwargs...) where TRS<:DF1_DRS
	result = DataRS{eltype(fn, d1, d2)}(d1.sweep)
	for i in 1:length(d1.sweep)
		result.elem[i] = broadcastMD(CT, fn, d1.elem[i], d2, args...; kwargs...)
	end
	return result
end
#fn(DataF1, DataRS) - core: fn(DataF1, DataF1):
function broadcastMD(CT::CastType2{DataF1,1,DataF1,2}, fn::Function,
	d1::DataF1, d2::DataRS{TRS}, args...; kwargs...) where TRS<:DF1_DRS
	result = DataRS{eltype(fn, d1, d2)}(d2.sweep)
	for i in 1:length(d2.sweep)
		result.elem[i] = broadcastMD(CT, fn, d1, d2.elem[i], args...; kwargs...)
	end
	return result
end
#Last line


abstract AbstractTuner

###Factory function, currently implemented are :monitor and :scale
tuner(s::Symbol,p::Int,args...;keyargs...) = _tuner(Val{s},p,args...;keyargs...)

period(tuner::AbstractTuner) = tuner.period
verbose(tuner::AbstractTuner) = tuner.verbose
needstuning(tuner::AbstractTuner,i::Int) = mod(i,period(tuner))==0

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

abstract AbstractTunerState

###Factory function for tunerstates
tunerstate{T<:AbstractFloat}(t::AbstractTuner,nburnin::Int,::Type{T}) = _tunerstate(t,nburnin,T)

accepted(state::AbstractTunerState) = state.accepted
proposed(state::AbstractTunerState) = state.proposed
rate(state::AbstractTunerState) = accepted(state)./proposed(state)
total(state::AbstractTunerState) = sum(proposed(state))
index(state::AbstractTunerState) = state.index
numtunesteps(state::AbstractTunerState) = length(state.accepted)

current(state::AbstractTunerState) = (a = state.accepted[state.index] ; p = state.proposed[state.index] ; tuple(a,p))

tune(tuner::AbstractTuner,state::AbstractTunerState) = ()

@inline function accepted!(state::AbstractTunerState,indicator::AbstractIndicatorMatrix)
    if state.index <= numtunesteps(state)
        state.accepted[state.index] += accepted(indicator)
        state.proposed[state.index] += numsamples(indicator)
    else
        warn("TunerState is full. Additional tuning steps cannot be stored")
    end
end

@inline nextindex!(state::AbstractTunerState) = (state.index+=1)


"""
    test_chan_ext(file::String, fmt::String, id::String, fs::Float64, i::Int64, t0::Int64)

Test for Issue 34
"""

function test_chan_ext(file::String, fmt::String, id::String, fs::Float64, i::Int64, t0::Int64)

  # Comparison struct
  S = verified_read_data(fmt, file)
  j = findid(id, S)
  if j == 0
    println("Test failing; cannot find ID!")
    println(rpad("File: ", 12), file)
    println(rpad("Format: ", 12), fmt)
    println("Channel ID: ", id)
    println("Output IDs: ", S.id)
    throw("Test failed; ID not found!")
  end

  # Create SeisData struct with data present using correct headers
  ny = max(100, round(Int64, 10*fs))
  y = rand(Float32, ny)
  S1 = SeisData(
    SeisChannel(
      id = id,
      fs = fs,
      gain = S.gain[j],
      units = S.units[j],
      loc = deepcopy(S.loc[j]),
      resp = deepcopy(S.resp[j]),
      t = [1 rand(0:10000000); ny 0],
      x = copy(y)
      )
    )

  # Read into S1
  verified_read_data!(S1, fmt, file)
  t1 = t_expand(S1.t[i], S.fs[i])[ny+1]
  t2 = S.t[j][1,2]

  # Check that read preserves time, data
  @test S1.x[i][1:ny] == y
  @test S1.x[i][ny+1:end] == S.x[j]
  @test t0 == t1 == t2
  return nothing
end


module SumProductTest

using Test
using ForneyLab
using ForneyLab: generateId, addNode!, associate!, inferUpdateRule!, outboundType, isApplicable, setTargets!, messagePassingSchedule
using ForneyLab: SPClamp

# Integration helper
mutable struct MockNode <: FactorNode
    id::Symbol
    interfaces::Vector{Interface}
    i::Dict{Int,Interface}

    function MockNode(vars::Vector{Variable}; id=generateId(MockNode))
        n_interfaces = length(vars)
        self = new(id, Array{Interface}(undef, n_interfaces), Dict{Int,Interface}())
        addNode!(currentGraph(), self)

        for idx = 1:n_interfaces
            self.i[idx] = self.interfaces[idx] = associate!(Interface(self), vars[idx])
        end

        return self
    end
end

@sumProductRule(:node_type     => MockNode,
                :outbound_type => Message{PointMass},
                :inbound_types => (Nothing, Message{PointMass}, Message{PointMass}),
                :name          => SPMockOutPP)

@testset "@SumProductRule" begin
    @test SPMockOutPP <: SumProductRule{MockNode}
    @test isApplicable(SPMockOutPP, [Nothing, Message{PointMass}, Message{PointMass}])
    @test !isApplicable(SPMockOutPP, [Nothing, Message{PointMass}, Message{PointMass}, Message{PointMass}])
end

# Composite definition for inferUpdateRule! testset
@composite TestComposite (b,a) begin
    @RV z ~ GaussianMeanVariance(a, constant(1.0))
    b = constant(2.0) + z
end

@testset "inferUpdateRule!" begin
    FactorGraph()
    nd = MockNode([Variable(), constant(0.0), constant(0.0)])
    inferred_outbound_types = Dict(nd.i[2].partner => Message{PointMass}, nd.i[3].partner => Message{PointMass})

    entry = ScheduleEntry(nd.i[1], SumProductRule{MockNode})
    inferUpdateRule!(entry, entry.message_update_rule, inferred_outbound_types)

    @test entry.message_update_rule == SPMockOutPP

    # Internal msg passing tests
    FactorGraph()
    a = constant(0.0, id=:a)
    b = Variable(id=:b)
    tc = TestComposite(b, a)
    inferred_outbound_types = Dict(tc.i[:a].partner => Message{PointMass})
    entry = ScheduleEntry(tc.i[:b], SumProductRule{Nothing})

    inferUpdateRule!(entry, entry.message_update_rule, inferred_outbound_types)

    @test isdefined(entry, :internal_schedule)
    @test length(entry.internal_schedule) == 4
    @test entry.message_update_rule == entry.internal_schedule[end].message_update_rule
end

@testset "messagePassingSchedule" begin
    fg = FactorGraph()
    x = Variable()
    nd = MockNode([x, constant(0.0), constant(0.0)])

    pfz = PosteriorFactorization()
    pf = PosteriorFactor(fg)
    setTargets!(pf, pfz, target_variables=Set{Variable}([x]))
    schedule = messagePassingSchedule(pf)

    @test length(schedule) == 3
    @test ScheduleEntry(nd.i[2].partner, SPClamp{Univariate}) in schedule
    @test ScheduleEntry(nd.i[3].partner, SPClamp{Univariate}) in schedule
    @test ScheduleEntry(nd.i[1], SPMockOutPP) in schedule
end

@testset "messagePassingAlgorithm" begin
    fg = FactorGraph()
    x = Variable()
    nd = MockNode([x, constant(0.0), constant(0.0)])

    pfz = PosteriorFactorization(fg)
    algo = messagePassingAlgorithm(x)

    @test isa(algo, InferenceAlgorithm)
end

end # module


#=
countDfs:
- Julia version:
- Author: Dan
- Date: 2021-11-28
=#

push!(LOAD_PATH, "/mnt/g/Python Programme/BALL.jl/src" )
using BALL
using BenchmarkTools

pdb_filepath = "1AKE.pdb"

comp = System(pdb_filepath)
dfs = DataFrameSystem(comp)

function distanceDfs()
    return distance(@view(dfs.residues[50,:]), @view(dfs.residues[60,:]))
end

distanceDfs()

elapsed = @elapsed distanceDfs()
println(elapsed)
@show @benchmark distanceDfs()


using AugmentorAddons
using Test
using ImageIO
using ImageFiltering
using Augmentor

@testset "AugmentorAddons.jl" begin
    ref = testpattern() |> size
    lo = (50, 50)
    hi = (10, 10)

    sz = augment(RCutOut(100, 100)) |> size
    @test (@. sz == ref) |> all

    sz = augment(PadBorders(RGB(0, 0, 0), lo)) |> size
    @test (@. sz == (ref + 2 * lo)) |> all

    sz = augment(PadBorders(:reflect, lo...)) |> size
    @test (@. sz == (ref + 2 * lo)) |> all

    sz = augment(PadBorders(RGB(0, 0, 0), lo, hi)) |> size
    @test (@. sz == (ref + lo + hi)) |> all

    x = rand(32, 32, 3, 4)
    ref = size(x)
    lo = (50, 50, 0, 0)
    hi = (10, 10, 0, 0)
    sz = augment(x, PadBorders(0, lo, hi)) |> size
    @test (@. sz == (ref + lo + hi)) |> all
end


import .Tracker
using .Tracker: Tracker,
                TrackedReal,
                TrackedVector,
                TrackedMatrix,
                TrackedArray,
                TrackedVecOrMat,
                @grad,
                track,
                data,
                param
using LinearAlgebra

# Different from Tracker.istracked
_istracked(x::AbstractArray{<:TrackedReal}) = true
_istracked(::TrackedArray) = false

Tracker.dual(x::Bool, p) = x
Base.prevfloat(r::TrackedReal) = track(prevfloat, r)
@grad function prevfloat(r::Real)
    prevfloat(data(r)), Δ -> Δ
end
Base.nextfloat(r::TrackedReal) = track(nextfloat, r)
@grad function nextfloat(r::Real)
    nextfloat(data(r)), Δ -> Δ
end
for i = 0:2, c = Tracker.combinations([:AbstractArray, :TrackedArray, :TrackedReal, :Number], i), f = [:hcat, :vcat]
    if :TrackedReal in c
        cnames = map(_ -> gensym(), c)
        @eval Base.$f($([:($x::$c) for (x, c) in zip(cnames, c)]...), x::Union{TrackedArray,TrackedReal}, xs::Union{AbstractArray,Number}...) =
            track($f, $(cnames...), x, xs...)
    end
end
@grad function vcat(x::Real)
    vcat(data(x)), (Δ) -> (Δ[1],)
end
@grad function vcat(x1::Real, x2::Real)
    vcat(data(x1), data(x2)), (Δ) -> (Δ[1], Δ[2])
end
@grad function vcat(x1::AbstractVector, x2::Real)
    vcat(data(x1), data(x2)), (Δ) -> (Δ[1:length(x1)], Δ[length(x1)+1])
end

function Base.copy(
    A::TrackedArray{T, 2, <:Adjoint{T, <:AbstractTriangular{T, <:AbstractMatrix{T}}}},
) where {T <: Real}
    return track(copy, A)
end
@grad function Base.copy(
    A::TrackedArray{T, 2, <:Adjoint{T, <:AbstractTriangular{T, <:AbstractMatrix{T}}}},
) where {T <: Real}
    return copy(data(A)), ∇ -> (copy(∇),)
end

Base.:*(A::TrackedMatrix, B::AbstractTriangular) = track(*, A, B)
Base.:*(A::AbstractTriangular{T}, B::TrackedVector) where {T} = track(*, A, B)
Base.:*(A::AbstractTriangular{T}, B::TrackedMatrix) where {T} = track(*, A, B)
Base.:*(A::Adjoint{T, <:AbstractTriangular{T}}, B::TrackedMatrix) where {T} = track(*, A, B)
Base.:*(A::Adjoint{T, <:AbstractTriangular{T}}, B::TrackedVector) where {T} = track(*, A, B)

_eps(::Type{<:TrackedReal{T}}) where {T} = _eps(T)

# AD implementations
function jacobian(
    b::Union{<:ADBijector{<:TrackerAD}, Inverse{<:ADBijector{<:TrackerAD}}},
    x::Real
)
    return data(Tracker.gradient(b, x)[1])
end
function jacobian(
    b::Union{<:ADBijector{<:TrackerAD}, Inverse{<:ADBijector{<:TrackerAD}}},
    x::AbstractVector{<:Real}
)
    # We extract `data` so that we don't return a `Tracked` type
    return data(Tracker.jacobian(b, x))
end

# implementations for Shift bijector
function _logabsdetjac_shift(a::TrackedReal, x::Real, ::Val{0})
    return param(_logabsdetjac_shift(data(a), data(x), Val(0)))
end
function _logabsdetjac_shift(a::TrackedReal, x::AbstractVector{<:Real}, ::Val{0})
    return param(_logabsdetjac_shift(data(a), data(x), Val(0)))
end
function _logabsdetjac_shift(
    a::Union{TrackedReal, TrackedVector{<:Real}},
    x::AbstractVector{<:Real},
    ::Val{1}
)
    return param(_logabsdetjac_shift(data(a), data(x), Val(1)))
end
function _logabsdetjac_shift(
    a::Union{TrackedReal, TrackedVector{<:Real}},
    x::AbstractMatrix{<:Real},
    ::Val{1}
)
    return param(_logabsdetjac_shift(data(a), data(x), Val(1)))
end

# Log bijector

@grad function logabsdetjac(b::Log{1}, x::AbstractVector)
    return -sum(log, data(x)), Δ -> (nothing, -Δ ./ data(x))
end
@grad function logabsdetjac(b::Log{1}, x::AbstractMatrix)
    return -vec(sum(log, data(x); dims = 1)), Δ -> (nothing, .- Δ' ./ data(x))
end

# implementations for Scale bijector
# Adjoints for 0-dim and 1-dim `Scale` using `Real`
function _logabsdetjac_scale(a::TrackedReal, x::Real, ::Val{0})
    return track(_logabsdetjac_scale, a, data(x), Val(0))
end
@grad function _logabsdetjac_scale(a::Real, x::Real, ::Val{0})
    return _logabsdetjac_scale(data(a), data(x), Val(0)), Δ -> (inv(data(a)) .* Δ, nothing, nothing)
end
# Need to treat `AbstractVector` and `AbstractMatrix` separately due to ambiguity errors
function _logabsdetjac_scale(a::TrackedReal, x::AbstractVector, ::Val{0})
    return track(_logabsdetjac_scale, a, data(x), Val(0))
end
@grad function _logabsdetjac_scale(a::Real, x::AbstractVector, ::Val{0})
    da = data(a)
    J = fill(inv.(da), length(x))
    return _logabsdetjac_scale(da, data(x), Val(0)), Δ -> (transpose(J) * Δ, nothing, nothing)
end
function _logabsdetjac_scale(a::TrackedReal, x::AbstractMatrix, ::Val{0})
    return track(_logabsdetjac_scale, a, data(x), Val(0))
end
@grad function _logabsdetjac_scale(a::Real, x::AbstractMatrix, ::Val{0})
    da = data(a)
    J = fill(size(x, 1) / da, size(x, 2))
    return _logabsdetjac_scale(da, data(x), Val(0)), Δ -> (transpose(J) * Δ, nothing, nothing)
end
# adjoints for 1-dim and 2-dim `Scale` using `AbstractVector`
function _logabsdetjac_scale(a::TrackedVector, x::AbstractVector, ::Val{1})
    return track(_logabsdetjac_scale, a, data(x), Val(1))
end
@grad function _logabsdetjac_scale(a::TrackedVector, x::AbstractVector, ::Val{1})
    # ∂ᵢ (∑ⱼ log|aⱼ|) = ∑ⱼ δᵢⱼ ∂ᵢ log|aⱼ|
    #                 = ∂ᵢ log |aᵢ|
    #                 = (1 / aᵢ) ∂ᵢ aᵢ
    #                 = (1 / aᵢ)
    da = data(a)
    J = inv.(da)
    return _logabsdetjac_scale(da, data(x), Val(1)), Δ -> (J .* Δ, nothing, nothing)
end
function _logabsdetjac_scale(a::TrackedVector, x::AbstractMatrix, ::Val{1})
    return track(_logabsdetjac_scale, a, data(x), Val(1))
end
@grad function _logabsdetjac_scale(a::TrackedVector, x::AbstractMatrix, ::Val{1})
    da = data(a)
    Jᵀ = repeat(inv.(da), 1, size(x, 2))
    return _logabsdetjac_scale(da, data(x), Val(1)), Δ -> (Jᵀ * Δ, nothing, nothing)
end
# TODO: implement analytical gradient for scaling a vector using a matrix
# function _logabsdetjac_scale(a::TrackedMatrix, x::AbstractVector, ::Val{1})
#     track(_logabsdetjac_scale, a, data(x), Val{1})
# end
# @grad function _logabsdetjac_scale(a::TrackedMatrix, x::AbstractVector, ::Val{1})
#     throw
# end

# implementations for Stacked bijector
function logabsdetjac(b::Stacked, x::TrackedMatrix{<:Real})
    return mapvcat(eachcol(x)) do c
        logabsdetjac(b, c)
    end
end
# TODO: implement custom adjoint since we can exploit block-diagonal nature of `Stacked`
function (sb::Stacked)(x::TrackedMatrix{<:Real})
    return eachcolmaphcat(sb, x)
end

# Simplex adjoints

function _simplex_bijector(X::TrackedVecOrMat, b::SimplexBijector)
    return track(_simplex_bijector, X, b)
end
function _simplex_inv_bijector(Y::TrackedVecOrMat, b::SimplexBijector)
    return track(_simplex_inv_bijector, Y, b)
end

@grad function _simplex_bijector(X::AbstractVector, b::SimplexBijector)
    Xd = data(X)
    return _simplex_bijector(Xd, b), Δ -> (simplex_link_jacobian(Xd)' * Δ, nothing)
end
@grad function _simplex_inv_bijector(Y::AbstractVector, b::SimplexBijector)
    Yd = data(Y)
    return _simplex_inv_bijector(Yd, b), Δ -> (simplex_invlink_jacobian(Yd)' * Δ, nothing)
end

@grad function _simplex_bijector(X::AbstractMatrix, b::SimplexBijector)
    Xd = data(X)
    return _simplex_bijector(Xd, b), Δ -> begin
        maphcat(eachcol(Xd), eachcol(Δ)) do c1, c2
            simplex_link_jacobian(c1)' * c2
        end, nothing
    end
end
@grad function _simplex_inv_bijector(Y::AbstractMatrix, b::SimplexBijector)
    Yd = data(Y)
    return _simplex_inv_bijector(Yd, b), Δ -> begin
        maphcat(eachcol(Yd), eachcol(Δ)) do c1, c2
            simplex_invlink_jacobian(c1)' * c2
        end, nothing
    end
end

replace_diag(::typeof(log), X::TrackedMatrix) = track(replace_diag, log, X)
@grad function replace_diag(::typeof(log), X)
    Xd = data(X)
    f(i, j) = i == j ? log(Xd[i, j]) : Xd[i, j]
    out = f.(1:size(Xd, 1), (1:size(Xd, 2))')
    out, ∇ -> begin
        g(i, j) = i == j ? ∇[i, j]/Xd[i, j] : ∇[i, j]
        return (nothing, g.(1:size(Xd, 1), (1:size(Xd, 2))'))
    end
end

replace_diag(::typeof(exp), X::TrackedMatrix) = track(replace_diag, exp, X)
@grad function replace_diag(::typeof(exp), X)
    Xd = data(X)
    f(i, j) = ifelse(i == j, exp(Xd[i, j]), Xd[i, j])
    out = f.(1:size(Xd, 1), (1:size(Xd, 2))')
    out, ∇ -> begin
        g(i, j) = ifelse(i == j, ∇[i, j]*exp(Xd[i, j]), ∇[i, j])
        return (nothing, g.(1:size(Xd, 1), (1:size(Xd, 2))'))
    end
end

logabsdetjac(b::SimplexBijector, x::TrackedVecOrMat) = track(logabsdetjac, b, x)
@grad function logabsdetjac(b::SimplexBijector, x::AbstractVector)
    xd = data(x)
    return logabsdetjac(b, xd), Δ -> begin
        (nothing, simplex_logabsdetjac_gradient(xd) * Δ)
    end
end
@grad function logabsdetjac(b::SimplexBijector, x::AbstractMatrix)
    xd = data(x)
    return logabsdetjac(b, xd), Δ -> begin
        (nothing, maphcat(eachcol(xd), Δ) do c, g
            simplex_logabsdetjac_gradient(c) * g
        end)
    end
end

for header in [
    (:(u::TrackedArray), :w),
    (:u, :(w::TrackedArray)),
    (:(u::TrackedArray), :(w::TrackedArray)),
]
    @eval begin
        function get_u_hat($(header...))
            if u isa TrackedArray
                T = typeof(u)
            else
                T = typeof(w)
            end
            x = w' * u
            return (u .+ (planar_flow_m(x) - x) .* w ./ sum(abs2, w))::T
        end
    end
end

for header in [
    (:(z::TrackedArray), :w, :b),
    (:z, :(w::TrackedArray), :b),
    (:z, :w, :(b::TrackedReal)),
    (:(z::TrackedArray), :(w::TrackedArray), :b),
    (:(z::TrackedArray), :w, :(b::TrackedReal)),
    (:z, :(w::TrackedArray), :(b::TrackedReal)),
    (:(z::TrackedArray), :(w::TrackedArray), :(b::TrackedReal)),
]
    @eval begin
        function ψ($(header...))
            if z isa AbstractMatrix
                if z isa TrackedMatrix
                    T = typeof(z)
                elseif w isa TrackedVector
                    T = matrixof(typeof(w))
                else
                    T = matrixof(typeof(b))
                end
            else
                if z isa TrackedVector
                    T = typeof(z)
                elseif w isa TrackedVector
                    T = typeof(w)
                else
                    T = vectorof(typeof(b))
                end
            end
            return ((1 .- tanh.(w' * z .+ b).^2) .* w)::T    # for planar flow from eq(11)
        end
    end
end

for header in [
    (:(u::TrackedArray), :w, :b, :(z::AbstractVecOrMat)),
    (:u, :(w::TrackedArray), :b, :(z::AbstractVecOrMat)),
    (:u, :w, :(b::TrackedReal), :(z::AbstractVecOrMat)),
    (:u, :w, :b, :(z::TrackedVecOrMat)),
    (:(u::TrackedArray), :(w::TrackedArray), :b, :(z::AbstractVecOrMat)),
    (:(u::TrackedArray), :w, :(b::TrackedReal), :(z::AbstractVecOrMat)),
    (:(u::TrackedArray), :w, :b, :(z::TrackedVecOrMat)),
    (:u, :(w::TrackedArray), :(b::TrackedReal), :(z::AbstractVecOrMat)),
    (:u, :(w::TrackedArray), :b, :(z::TrackedVecOrMat)),
    (:u, :w, :(b::TrackedArray), :(z::TrackedVecOrMat)),
    (:(u::TrackedArray), :(w::TrackedArray), :(b::TrackedReal), :(z::AbstractVecOrMat)),
    (:(u::TrackedArray), :(w::TrackedArray), :b, :(z::TrackedVecOrMat)),
    (:(u::TrackedArray), :w, :(b::TrackedReal), :(z::TrackedVecOrMat)),
    (:u, :(w::TrackedArray), :(b::TrackedReal), :(z::TrackedVecOrMat)),
    (:(u::TrackedArray), :(w::TrackedArray), :(b::TrackedReal), :(z::TrackedVecOrMat)),
]
    @eval begin
        function _planar_transform($(header...))
            u_hat = get_u_hat(u, w)
            if z isa AbstractVector
                temp = w' * z + b + zero(eltype(u_hat))
                if z isa TrackedVector
                    T = typeof(z)
                elseif u_hat isa TrackedVector
                    T = typeof(u_hat)
                else
                    T = vectorof(typeof(temp))
                end
            else
                temp = w' * z .+ (b + zero(eltype(u_hat)))
                if z isa TrackedMatrix
                    T = typeof(z)
                elseif u_hat isa TrackedVector
                    T = matrixof(typeof(u_hat))
                else
                    T = matrixof(typeof(temp'))
                end
            end
            transformed::T = z .+ u_hat .* tanh.(temp) # from eq(10)
            return (transformed = transformed, u_hat = u_hat)
        end
    end
end

for header in [
    (:(α_::TrackedReal), :β, :z_0, :(z::AbstractVector)),
    (:α_, :(β::TrackedReal), :z_0, :(z::AbstractVector)),
    (:α_, :β, :(z_0::TrackedVector), :(z::AbstractVector)),
    (:α_, :β, :z_0, :(z::TrackedVector)),
    (:(α_::TrackedReal), :(β::TrackedReal), :z_0, :(z::AbstractVector)),
    (:(α_::TrackedReal), :β, :(z_0::TrackedVector), :(z::AbstractVector)),
    (:(α_::TrackedReal), :β, :z_0, :(z::TrackedVecOrMat)),
    (:(α_::TrackedReal), :(β::TrackedReal), :(z_0::TrackedVector), :(z::AbstractVector)),
    (:(α_::TrackedReal), :(β::TrackedReal), :z_0, :(z::TrackedVector)),
    (:(α_::TrackedReal), :β, :(z_0::TrackedVector), :(z::TrackedVector)),
    (:α_, :(β::TrackedReal), :(z_0::TrackedVector), :(z::TrackedVector)),
    (:(α_::TrackedReal), :(β::TrackedReal), :(z_0::TrackedVector), :(z::TrackedVector)),
]
    @eval begin
        function _radial_transform($(header...))
            α = softplus(α_)            # from A.2
            β_hat = -α + softplus(β)    # from A.2
            if β_hat isa TrackedReal
                TV = vectorof(typeof(β_hat))
                T = vectorof(typeof(β_hat))
            elseif z_0 isa TrackedVector
                TV = typeof(z_0)
                T = typeof(z_0)
            else
                T = TV = typeof(z)
            end
            Tr = promote_type(eltype(z), eltype(z_0))
            r::Tr = norm((z .- z_0)::TV)
            transformed::T = z .+ β_hat ./ (α .+ r') .* (z .- z_0)   # from eq(14)
            return (transformed = transformed, α = α, β_hat = β_hat, r = r)
        end
    end
end

for header in [
    (:(α_::TrackedReal), :β, :z_0, :(z::AbstractMatrix)),
    (:α_, :(β::TrackedReal), :z_0, :(z::AbstractMatrix)),
    (:α_, :β, :(z_0::TrackedVector), :(z::AbstractMatrix)),
    (:α_, :β, :z_0, :(z::TrackedMatrix)),
    (:(α_::TrackedReal), :(β::TrackedReal), :z_0, :(z::AbstractMatrix)),
    (:(α_::TrackedReal), :β, :(z_0::TrackedVector), :(z::AbstractMatrix)),
    (:(α_::TrackedReal), :β, :z_0, :(z::TrackedMatrix)),
    (:(α_::TrackedReal), :(β::TrackedReal), :(z_0::TrackedVector), :(z::AbstractMatrix)),
    (:(α_::TrackedReal), :(β::TrackedReal), :z_0, :(z::TrackedMatrix)),
    (:(α_::TrackedReal), :β, :(z_0::TrackedVector), :(z::TrackedMatrix)),
    (:α_, :(β::TrackedReal), :(z_0::TrackedVector), :(z::TrackedMatrix)),
    (:(α_::TrackedReal), :(β::TrackedReal), :(z_0::TrackedVector), :(z::TrackedMatrix)),
]
    @eval begin
        function _radial_transform($(header...))
            α = softplus(α_)            # from A.2
            β_hat = -α + softplus(β)    # from A.2
            if β_hat isa TrackedReal
                TV = vectorof(typeof(β_hat))
                T = matrixof(TV)
            elseif z_0 isa TrackedVector
                TV = typeof(z_0)
                T = matrixof(TV)
            else
                T = typeof(z)
                TV = vectorof(T)
            end
            r::TV = mapvcat(norm, eachcol(z .- z_0))
            transformed::T = z .+ β_hat ./ (α .+ r') .* (z .- z_0)   # from eq(14)
            return (transformed = transformed, α = α, β_hat = β_hat, r = r)
        end
    end
end

function matrixof(::Type{TrackedArray{T, 1, Vector{T}}}) where {T <: Real}
    return TrackedArray{T, 2, Matrix{T}}
end
function matrixof(::Type{TrackedReal{T}}) where {T <: Real}
    return TrackedArray{T, 2, Matrix{T}}
end
function vectorof(::Type{TrackedArray{T, 2, Matrix{T}}}) where {T <: Real}
    return TrackedArray{T, 1, Vector{T}}
end
function vectorof(::Type{TrackedReal{T}}) where {T <: Real}
    return TrackedArray{T, 1, Vector{T}}
end

(b::Exp)(x::TrackedVector) = exp.(x)::vectorof(float(eltype(x)))
(b::Log)(x::TrackedVector) = log.(x)::vectorof(float(eltype(x)))
logabsdetjac(b::Log{0}, x::TrackedVector) = .-log.(x)::vectorof(float(eltype(x)))
logabsdetjac(b::Log{1}, x::TrackedMatrix) = - vec(sum(log.(x); dims = 1))

_logabsdetjac_shift(a::Real, x::TrackedVector{T}, ::Val{0}) where {T<:Real} = zeros(T, length(x))
_logabsdetjac_shift(a::T1, x::TrackedVector{T2}, ::Val{1}) where {T1<:Union{Real, TrackedVector}, T2<:Real} = zero(T2)
_logabsdetjac_shift(a::T1, x::TrackedMatrix{T2}, ::Val{1}) where {T1<:Union{Real, TrackedVector}, T2<:Real} = zeros(T2, size(x, 2))


@info "Create EOS table"

@info "Creating file \"eos.quantities\"..."

# regular_quantities = [1, 2, 12, 20, 21, 22, 23, 24]
# additional_quantities = []
# derivative_quantities = []
# pair_quantities = []
# quadruple_quantities = []
# microscopic_quantities = []
# error_quantities = []

# These are all the quantities (why not?)
regular_quantities = collect(1:24)
additional_quantities = collect(1:1)
derivative_quantities = collect(1:10)
pair_quantities = [10, 11, 4002, 0]
quadruple_quantities = [1]
microscopic_quantities = [10040, 11040, 10050, 11050]
error_quantities = collect(1:8)
output_format = 0               # 1=ASCII, 0=HDF5

open("eos.quantities", "w") do fh
    println(fh,
            "# number of regular, additional and derivative quantities (see table 7.1)")
    println(fh,
            join([length(regular_quantities), length(additional_quantities),
                  length(derivative_quantities)], " "))
    println(fh, "# indices of regular, additional and derivative quantities")
    println(fh,
            join([regular_quantities; additional_quantities;
                  derivative_quantities], " "))
    println(fh,
            "# number of pairs and quadruples for composition data (see table 3.2/3.3)")
    println(fh,
            join([length(pair_quantities), length(quadruple_quantities)], " "))
    println(fh, "# indices for pairs and quadruples for composition data")
    println(fh, join([pair_quantities; quadruple_quantities], " "))
    println(fh, "# number of microscopic quantities (see section 4.2.4)")
    println(fh, length(microscopic_quantities))
    println(fh, "# indices of microscopic quantities")
    println(fh, join(microscopic_quantities, " "))
    println(fh, "# number of error quantities")
    println(fh, length(error_quantities))
    println(fh, "# indices of error quantities")
    println(fh, join(error_quantities, " "))
    println(fh, "# format of output file, ASCII (1) or HDF5 (else)")
    println(fh, output_format)
    return nothing
end

@info "Creating file \"eos.parameters\"..."

# This selects all the points in the original table, presumably without interpolation (why not?)

file_t = split(read("eos.t", String))
log_t = parse(Int, file_t[1])
num_t = parse(Int, file_t[2])
min_t = parse(Float64, file_t[3])
max_t = parse(Float64, file_t[end])
@assert length(file_t) == num_t + 2

file_nb = split(read("eos.nb", String))
log_nb = parse(Int, file_nb[1])
num_nb = parse(Int, file_nb[2])
min_nb = parse(Float64, file_nb[3])
max_nb = parse(Float64, file_nb[end])
@assert length(file_nb) == num_nb + 2

file_yq = split(read("eos.yq", String))
log_yq = parse(Int, file_yq[1])
num_yq = parse(Int, file_yq[2])
min_yq = parse(Float64, file_yq[3])
max_yq = parse(Float64, file_yq[end])
@assert length(file_yq) == num_yq + 2

open("eos.parameters", "w") do fh
    println(fh, "# order of interpolation in first, second and third index")
    println(fh, "1 1 1")
    println(fh,
            "# calculation of beta-equilibrium (1: yes, else: no) and for given entropy (1: yes, else: no)")
    println(fh, "0 0")
    println(fh,
            "# tabulation scheme (0 = explicit listing, 1 = loops, see manual)")
    println(fh, "1")
    println(fh,
            "# parameter values (first, second and third index) depending on tabulation scheme")
    println(fh, join([min_t, min_nb, min_yq], " "))
    println(fh, join([max_t, max_nb, max_yq], " "))
    println(fh, join([num_t, num_nb, num_yq], " "))
    println(fh, join([log_t, log_nb, log_yq], " "))
    return nothing
end

@info "Generate EOS file..."
run(pipeline(`echo 3`, `/Users/eschnett/src/codehdf5/compose`))

# h5repack -v -f SHUF -f GZIP=9 eoscompose.h5 eoscompose_compressed.h5

@info "Done."


push!(LOAD_PATH, "$(homedir())/Development/")
using JSort

path = "$(homedir())/master/sortering/"*"28si.yaml"
parameters = Parameters(path)
#sortfile(parameters)
#savecoefficients_eΔe(parameters, coefficients...)
#@time events = loadlabr(parameters);
#@time makeeΔebin(events, parameters)
@time coeff_e, coeff_Δe = featurealign2d(parameters.readpath, plot=true)
@show coeff_e
p_e = ParticleCalibrator(coeff_e, var=:e)
p_Δe = ParticleCalibrator(coeff_e, var=:Δe)
write(joinpath(parameters.savepath, "coefficients_e.txt"), p_e)
write(joinpath(parameters.savepath, "coefficients_de.txt"), p_Δe)
#@time savecoefficients_eΔe(parameters, coefficients...)


## Pattern matching needs a new implementation. The ad hoc approach can't go further.
## This is a start ...

### Types for patterns

abstract type AbstractBlanks end

abstract type ABlank  <: AbstractBlanks end
abstract type ABlankSequence <: AbstractBlanks end
abstract type ABlankNullSequence <: AbstractBlanks end

mutable struct BlankNoHead <: ABlank end
mutable struct BlankWithHead{T} <: ABlank
    head::T
end

mutable struct BlankSequenceNoHead <: ABlankSequence end
mutable struct BlankSequenceWithHead{T} <: ABlankSequence
    head::T
end

mutable struct BlankNullSequenceNoHead <: ABlankNullSequence end
mutable struct BlankNullSequenceWithHead{T} <: ABlankNullSequence
    head::T
end

const BlanksNoHead = Union{BlankNoHead,BlankSequenceNoHead,BlankNullSequenceNoHead}
const BlanksWithHead = Union{BlankWithHead,BlankSequenceWithHead,BlankNullSequenceWithHead}

topattern(mx::Mxpr) = mxpra(topattern(mhead(mx)), mapmargs(x -> topattern(x), margs(mx)))
topattern(x) = x

## replace this later with BlankT. after removing the old BlankT
topattern(b::Mxpr{:Blank}) = isempty(b) ? BlankNoHead() : BlankWithHead(_make_blank_head(b,margs(b)...))
topattern(b::Mxpr{:BlankSequence}) = isempty(b) ? BlankSequenceNoHead() : BlankSequenceWithHead(_make_blank_head(b,margs(b)...))
topattern(b::Mxpr{:BlankNullSequence}) = isempty(b) ? BlankNullSequenceNoHead() :
          BlankNullSequenceWithHead(_make_blank_head(b,margs(b)...))

_make_blank_head(b,head,args...) = error("More than one argument in Blank.")
_make_blank_head(b,head::Symbol) = @isdefined(head) ? eval(head) : head
_make_blank_head(b,head) = head

#### Matching1

pmatch(b::BlanksNoHead,ex) = true
pmatch(b::BlankWithHead{T},ex) where {T<:DataType} = isa(ex,b.head)
pmatch(b::BlankSequenceWithHead{T},ex) where {T<:DataType} = isa(ex,b.head)
pmatch(b::BlankNullSequenceWithHead{T},ex) where {T<:DataType} = isa(ex,b.head)
pmatch(b::BlanksWithHead,ex) = isa(ex,Mxpr{b.head})
pmatch(b::AbstractBlanks,ex) = symerror("match: Can't match Blank of type ", typeof(b))

#### Matching

mutable struct Capt
    c::Dict{SJSym,Any}
end

Capt() = Capt(Dict{SJSym,Any}())

abstract type  AbstractMatchData end

mutable struct SearchRange
    start::Int
    stop::Int
end

SearchRange(r::UnitRange) = SearchRange(r.start, r.stop)

mutable struct MatchDataIn
    _expr::Mxpr
    _indr::SearchRange
    _pat ## pattern in general sense
    _capt::Capt
end

## Neither flat nor orderless
mutable struct MatchData <: AbstractMatchData
    d::MatchDataIn
end

## flat
mutable struct MatchDataF <: AbstractMatchData
    d::MatchDataIn
end

## orderless
mutable struct MatchDataO <: AbstractMatchData
    d::MatchDataIn
end

## flat and orderless
mutable struct MatchDataFO <: AbstractMatchData
    d::MatchDataIn
end

getexpr(d::AbstractMatchData) = d.d._expr
getcap(d::AbstractMatchData) = d.d._capt
getstartind(d::AbstractMatchData) = d.d._indr.start
getstopind(d::AbstractMatchData) = d.d._indr.stop
setind(d::AbstractMatchData,i::Integer) = d.d._indr.start = i
incrind(d::AbstractMatchData) = d.d._indr.start += 1
getpat(d::AbstractMatchData) = d.d._pat

function setsearchrange(d::AbstractMatchData, r::UnitRange)
    d.d._indr.start = r.start
    d.d._indr.stop = r.stop
end

startsearchind(d::AbstractMatchData) = getstartind(d) + 1
stopsearchind(d::AbstractMatchData) = getstopind(d) + 1
firstel(d::AbstractMatchData) = d[startsearchind(d)]

Base.getindex(d::AbstractMatchData,i::Integer) =  getexpr(d)[i]
Base.length(d::AbstractMatchData) = length(getexpr(d))

matchdata(expr::Mxpr, pat) = matchdata(expr,pat,Capt())

function matchdata(expr::Mxpr, pat, capt::Capt)
    if isFlat(expr)
        if isOrderless(expr)
            MatchDataFO(expr,pat,capt)
        else
            MatchDataF(expr,pat,capt)
        end
    elseif isOrderless(expr)
        MatchDataO(expr,pat,capt)
    else
        MatchData(expr,pat,capt)
    end
end

_matchdatain(expr,pat,capt) = MatchDataIn(expr,SearchRange(0:0),pat,capt)

MatchData(expr,pat,capt) = MatchData(_matchdatain(expr,pat,capt))
MatchDataF(expr,pat,capt) = MatchDataF(_matchdatain(expr,pat,capt))
MatchDataO(expr,pat,capt) = MatchDataO(_matchdatain(expr,pat,capt))
MatchDataFO(expr,pat,capt) = MatchDataFO(_matchdatain(expr,pat,capt))

####

abstract type AbstractMatchIndex end

#abstract MatchIndexCompound <: AbstractMatchIndex

mutable struct MatchIndexSingle <: AbstractMatchIndex
    i::Int
end

mutable struct MatchIndexAbstractArray{T<:AbstractArray} <: AbstractMatchIndex
    a::T
end

MatchIndex(a::AbstractArray) = MatchIndexAbstractArray(a)
MatchIndex(a::Int) = MatchIndexSingle(a)

getstart(mi::MatchIndexAbstractArray) = start(mi.a)
getstop(mi::MatchIndexAbstractArray) = lastindex(mi.a)

getstart(mi::MatchIndexSingle) = mi.i
getstop(mi::MatchIndexSingle) = mi.i

####

abstract type AbstractMatchType end

mutable struct MatchOne{T<:AbstractMatchIndex} <: AbstractMatchType
    a::T
end

MatchOne(a) = MatchOne(MatchIndex(a))

mutable struct MatchNone <: AbstractMatchType
end

mutable struct MatchFail <: AbstractMatchType
end

struct MatchSequence{T<:AbstractMatchIndex} <: AbstractMatchType
    a::T
end

MatchSequence(a) = MatchSequence(MatchIndex(a))

struct MatchNullSequence{T<:AbstractMatchIndex} <: AbstractMatchType
    a::T
end

MatchNullSequence(a) = MatchNullSequence(MatchIndex(a))

getstart(mr::AbstractMatchType) = getstart(mr.a)
getstop(mr::AbstractMatchType) = getstop(mr.a)

function setsearchrange(d::AbstractMatchData, mr::AbstractMatchType)
    d.d._indr.start = getstart(mr)
    d.d._indr.stop = getstop(mr)
end

function setsearchstop(d::AbstractMatchData, mr::AbstractMatchType)
    d.d._indr.stop = getstop(mr)
end

#### Matching 2

function matchrange(data::MatchData, pat::ABlank)
#    pmatch(pat,data[startsearchind(data)]) ?
    #       MatchOne(startsearchind(data)) : MatchFail()
    (j1,j) = _matchrange_blank(data,pat)
    return j1 == 0 ? MatchFail() : MatchOne(j1:j)
end

function _matchrange_blank(data,pat)
    local j1
    matchedone::Bool = false
    for i in startsearchind(data):stopsearchind(data)
        j1 = i
        if pmatch(pat,data[i])
            matchedone = true
            break
        end
    end
    matchedone || return (0,0)
    j = j1
    #    for i in j1+1:length(data)
    maxi = min(length(data),stopsearchind(data))
#    for i in j1+1:stopsearchind(data)
    for i in j1+1:maxi
        pmatch(pat,data[i]) || break
        j = i
    end
    return (j1,j)
end

function matchrange(data::MatchData, pat::ABlankSequence)
    (j1,j) = _matchrange_blanksequence(data,pat)
    return j1 == 0 ? MatchFail() : MatchSequence(j1:j)
end

function matchrange(data::MatchData, pat::ABlankNullSequence)
    (j1,j) = _matchrange_blanksequence(data,pat)
    return j1 == 0 ? MatchNone() : MatchNullSequence(j1:j)
end

function _matchrange_blanksequence(data,pat)
    local j1
    matchedone::Bool = false
    for i in startsearchind(data):stopsearchind(data)
        j1 = i
        if pmatch(pat,data[i])
            matchedone = true
            break
        end
    end
    matchedone || return (0,0)
    j = j1
    for i in j1+1:length(data)
#    for i in j1+1:stopsearchind(data)
        pmatch(pat,data[i]) || break
        j = i
    end
    return (j1,j)
end

function advanceind(d::AbstractMatchData, mr::AbstractMatchType)
    nothing
end

function advanceind(d::AbstractMatchData, mr::MatchOne)
    setsearchrange(d,mr)
end

function advanceind(d::AbstractMatchData, mr::MatchSequence)
    setsearchrange(d,mr)
end

function advanceind(d::AbstractMatchData, mr::MatchNullSequence)
    setsearchstop(d,mr)
end

function matchranges(expr::Mxpr, pat)
    pat = topattern(pat)
    matchranges(matchdata(expr,pat))
end

function matchranges(d::AbstractMatchData)
    pat = getpat(d)
    a = Any[]
    for p in pat
        mr = matchrange(d,p)
        advanceind(d,mr)
        push!(a,mr)
    end
    a
end


# QuEST_jl/src/base/unitaries.jl
#

function compactUnitary(qureg       ::QuEST_Types.Qureg,
                        targetQubit ::Integer,
                        α           ::Complex{Qreal},
                        β           ::Complex{Qreal})   ::Nothing

    alpha = QuEST_Types.Complex(real(α),imag(α))
    beta  = QuEST_Types.Complex(real(β),imag(β))

    ccall(:compactUnitary, Cvoid,
          (QuEST_Types.Qureg, Cint,        QuEST_Types.Complex, QuEST_Types.Complex),
          qureg,              targetQubit, alpha,               beta)
end

function controlledCompactUnitary(qureg          ::QuEST_Types.Qureg,
                                  controlQubit   ::Integer,
                                  targetQubit    ::Integer,
                                  α              ::Complex{Qreal},
                                  β              ::Complex{Qreal} ) ::Nothing

    alpha = QuEST_Types.Complex(real(α),imag(α))
    beta  = QuEST_Types.Complex(real(β),imag(β))

    ccall(:controlledCompactUnitary,
          Cvoid,
          (QuEST_Types.Qureg, Cint, Cint, QuEST_Types.Complex, QuEST_Types.Complex),
          qureg,
          Cint(controlQubit),
          Cint(targetQubit),
          alpha,
          beta)
    return nothing
end

function controlledMultiQubitUnitary(qureg   ::QuEST_Types.Qureg,
                                     ctrl    ::Integer,
                                     targs   ::Vector{QubitIdx},
                                     u       ::QuEST_Types.ComplexMatrixN) ::Nothing
    @assert 0 ≤ ctrl < qureg.numQubitsRepresented
    @assert all( tidx -> 0 ≤ targs[tidx] < qureg.numQubitsRepresented,
                 1:length(targs) )
    @assert ctrl ∉ targs

    @assert u.numQubits < qureg.numQubitsRepresented

    ccall(:controlledMultiQubitUnitary, Cvoid,
          (QuEST_Types.Qureg, Cint, Ptr{Cint}, Cint,          QuEST_Types.ComplexMatrixN),
          qureg,              ctrl, targs,     length(targs), u)
end

function controlledNot(qureg         ::QuEST_Types.Qureg,
                       controlQubit  ::Integer,
                       targetQubit   ::Integer)         ::Nothing

    @assert 0 ≤ controlQubit < qureg.numQubitsRepresented
    @assert 0 ≤ targetQubit  < qureg.numQubitsRepresented
    @assert controlQubit != targetQubit

    ccall(:controlledNot, Cvoid,
          (QuEST_Types.Qureg, QubitIdx,      QubitIdx),
          qureg,              controlQubit,  targetQubit)
end

function controlledPauliY(qureg         ::QuEST_Types.Qureg,
                          controlQubit  ::Integer,
                          targetQubit   ::Integer)         ::Nothing

    @assert 0 ≤ controlQubit < qureg.numQubitsRepresented
    @assert 0 ≤ targetQubit  < qureg.numQubitsRepresented
    @assert controlQubit != targetQubit

    ccall(:controlledPauliY,  Cvoid,
          (QuEST_Types.Qureg, Cint,          Cint),
          qureg,              controlQubit,  targetQubit)
end

function controlledPhaseFlip(qureg         ::QuEST_Types.Qureg,
                             controlQubit  ::Integer,
                             targetQubit   ::Integer)   ::Nothing

    @assert 0 ≤ controlQubit < qureg.numQubitsRepresented
    @assert 0 ≤ targetQubit  < qureg.numQubitsRepresented
    @assert controlQubit != targetQubit

    ccall(:controlledPhaseFlip, Cvoid,
          (QuEST_Types.Qureg, Cint,         Cint),
          qureg,              controlQubit, targetQubit)
end

function controlledPhaseShift(qureg    ::QuEST_Types.Qureg,
                              idQubit1 ::Integer,
                              idQubit2 ::Integer,
                              angle    ::Qreal)   ::Nothing

    @assert 0 ≤ idQubit1 < qureg.numQubitsRepresented
    @assert 0 ≤ idQubit2  < qureg.numQubitsRepresented
    @assert idQubit1 != idQubit2

    ccall(:controlledPhaseShift,  Cvoid,
          (QuEST_Types.Qureg, Cint,     Cint,      Qreal),
          qureg,              idQubit1, idQubit2,  angle)
end

function controlledRotateAroundAxis(qureg        ::QuEST_Types.Qureg,
                                    controlQubit ::Integer,
                                    targetQubit  ::Integer,
                                    angle        ::Qreal,
                                    axis         ::Union{NTuple{3,Qreal},Vector{Qreal}})  ::Nothing

    @assert 0 ≤ controlQubit < qureg.numQubitsRepresented
    @assert 0 ≤ targetQubit  < qureg.numQubitsRepresented
    @assert controlQubit != targetQubit

    @assert length(axis) == 3

    q_axis = QuEST_Types.Vector(axis[1], axis[2], axis[3])

    ccall(:controlledRotateAroundAxis,  Cvoid,
          (QuEST_Types.Qureg, Cint,           Cint,         Qreal, QuEST_Types.Vector),
          qureg,              controlQubit,   targetQubit,  angle, q_axis)
end

function controlledRotateX(qureg         ::QuEST_Types.Qureg,
                           controlQubit  ::Integer,
                           targetQubit   ::Integer,
                           angle         ::Qreal)  ::Nothing

    @assert 0 ≤ controlQubit < qureg.numQubitsRepresented
    @assert 0 ≤ targetQubit  < qureg.numQubitsRepresented
    @assert controlQubit != targetQubit

    ccall(:controlledRotateX, Cvoid,
          (QuEST_Types.Qureg, Cint,          Cint,          Qreal),
          qureg,              controlQubit,  targetQubit,   angle)
end

function controlledRotateY(qureg         ::QuEST_Types.Qureg,
                           controlQubit  ::Integer,
                           targetQubit   ::Integer,
                           angle         ::Qreal)  ::Nothing

    @assert 0 ≤ controlQubit < qureg.numQubitsRepresented
    @assert 0 ≤ targetQubit  < qureg.numQubitsRepresented
    @assert controlQubit != targetQubit

    ccall(:controlledRotateY,    Cvoid,
          (QuEST_Types.Qureg, Cint,          Cint,        Qreal),
          qureg,              controlQubit,  targetQubit, angle)
end

function controlledRotateZ(qureg         ::QuEST_Types.Qureg,
    controlQubit  ::Integer,
    targetQubit   ::Integer,
    angle         ::Qreal)  ::Nothing

    ccall(:controlledRotateZ,
          Cvoid,
          (QuEST_Types.Qureg, Cint, Cint, Qreal),
          qureg,
          Cint(controlQubit),
          Cint(targetQubit),
          angle)
return nothing
end

function _quest_mtx_2(U ::Matrix{Complex{Qreal}}) ::QuEST_Types.ComplexMatrix2
    @assert size(U) == (2,2)
    u = QuEST_Types.ComplexMatrix2(
        ( (real(U[1,1]), real(U[1,2])), (real(U[2,1]), real(U[2,2])) ),
        ( (imag(U[1,1]), imag(U[1,2])), (imag(U[2,1]), imag(U[2,2])) )  )
    return u
end

function _quest_mtx_4(U ::Matrix{Complex{Qreal}}) ::QuEST_Types.ComplexMatrix4
    @assert size(U) == (4,4)
    u = QuEST_Types.ComplexMatrix4( ( ( real(U[1,1]), real(U[1,2]), real(U[1,3]), real(U[1,4]) ),
                          ( real(U[2,1]), real(U[2,2]), real(U[2,3]), real(U[2,4]) ),
                          ( real(U[3,1]), real(U[3,2]), real(U[3,3]), real(U[3,4]) ),
                          ( real(U[4,1]), real(U[4,2]), real(U[4,3]), real(U[4,4]) ) ),
                        ( ( imag(U[1,1]), imag(U[1,2]), imag(U[1,3]), imag(U[1,4]) ),
                          ( imag(U[2,1]), imag(U[2,2]), imag(U[2,3]), imag(U[2,4]) ),
                          ( imag(U[3,1]), imag(U[3,2]), imag(U[3,3]), imag(U[3,4]) ),
                          ( imag(U[4,1]), imag(U[4,2]), imag(U[4,3]), imag(U[4,4]) ) )  )
end

function controlledTwoQubitUnitary(qureg         ::QuEST_Types.Qureg,
                                   controlQubit  ::Integer,
                                   targetQubit1  ::Integer,
                                   targetQubit2  ::Integer,
                                   U             ::Matrix{Complex{Qreal}}) ::Nothing

    @assert 0 ≤ controlQubit < qureg.numQubitsRepresented
    @assert 0 ≤ targetQubit1 < qureg.numQubitsRepresented
    @assert 0 ≤ targetQubit2 < qureg.numQubitsRepresented
    @assert targetQubit1 != targetQubit2
    @assert controlQubit ∉ [targetQubit1,targetQubit2]

    @assert size(U) == (4, 4)
    u = _quest_mtx_4(U)
    ccall(:controlledTwoQubitUnitary,
          Cvoid,
          (QuEST_Types.Qureg, Cint, Cint, Cint, QuEST_Types.ComplexMatrix4),
          qureg,
          Cint(controlQubit),
          Cint(targetQubit1),
          Cint(targetQubit2),
          u)
return nothing
end

function controlledUnitary(qureg        ::QuEST_Types.Qureg,
                           controlQubit ::Integer,
                           targetQubit  ::Integer,
                           U            ::Matrix{Complex{Qreal}}) ::Nothing

    @assert 0 ≤ controlQubit < qureg.numQubitsRepresented
    @assert 0 ≤ targetQubit  < qureg.numQubitsRepresented
    @assert controlQubit != targetQubit

    u = _quest_mtx_2(U)

    ccall(:controlledUnitary,  Cvoid,
          (QuEST_Types.Qureg, Cint,          Cint,          QuEST_Types.ComplexMatrix2),
          qureg,              controlQubit,  targetQubit,   u)
end

function hadamard(qureg       ::QuEST_Types.Qureg,
                  targetQubit ::Integer) ::Nothing

    ccall(:hadamard, Cvoid,
          (QuEST_Types.Qureg, Cint),
          qureg,              targetQubit)
end

function multiControlledMultiQubitUnitary(qureg  ::QuEST_Types.Qureg,
                                          ctrls  ::Vector{QubitIdx},
                                          targs  ::Vector{QubitIdx},
                                          u      ::QuEST_Types.ComplexMatrixN) ::Nothing

    ccall(:multiControlledMultiQubitUnitary, Cvoid,
          (QuEST_Types.Qureg, Ptr{Cint}, Cint,          Ptr{Cint}, Cint,          QuEST_Types.ComplexMatrixN),
          qureg,              ctrls,     length(ctrls), targs,     length(targs), u)
end

function multiControlledPhaseFlip(qureg         ::QuEST_Types.Qureg,
                                  controlQubits ::Vector{QubitIdx}) ::Nothing

    ccall(:multiControlledPhaseFlip, Cvoid,
          (QuEST_Types.Qureg,     Ptr{Cint},     Cint),
          qureg,                  controlQubits, length(controlQubits))
end

function multiControlledPhaseShift(qureg         ::QuEST_Types.Qureg,
                                   controlQubits ::Vector{QubitIdx},
                                   angle         ::Qreal)         ::Nothing

    ccall(:multiControlledPhaseShift, Cvoid,
          (QuEST_Types.Qureg, Ptr{Cint},     Cint,                  Qreal),
          qureg,              controlQubits, length(controlQubits), angle)
end

function multiControlledTwoQubitUnitary(qureg           ::QuEST_Types.Qureg,
                                        controlQubits   ::Vector{QubitIdx},
                                        targetQubit1    ::Integer,
                                        targetQubit2    ::Integer,
                                        U               ::Matrix{Complex{Qreal}}) ::Nothing

    @assert 0 ≤ length(controlQubits) ≤ qureg.numQubitsRepresented -2
    @assert 0 ≤ targetQubit1 < qureg.numQubitsRepresented
    @assert 0 ≤ targetQubit2 < qureg.numQubitsRepresented
    @assert targetQubit1 != targetQubit2
    @assert targetQubit1 ∉ controlQubits
    @assert targetQubit2 ∉ controlQubits

    u = _quest_mtx_4(U)

    ccall(:multiControlledTwoQubitUnitary, Cvoid,
          (QuEST_Types.Qureg, Ptr{Cint},     Cint,                  Cint,          Cint,           QuEST_Types.ComplexMatrix4),
          qureg,              controlQubits, length(controlQubits), targetQubit1,  targetQubit2,   u)
end

function multiControlledUnitary(qureg         ::QuEST_Types.Qureg,
                                controlQubits ::Vector{QubitIdx},
                                targetQubit   ::Integer,
                                U             ::Matrix{Complex{Qreal}}) ::Nothing

    u = _quest_mtx_2(U)

    ccall(:multiControlledUnitary,  Cvoid,
          (QuEST_Types.Qureg, Ptr{Cint},     Cint,                  Cint,          QuEST_Types.ComplexMatrix2),
          qureg,              controlQubits, length(controlQubits), targetQubit,   u)
end

function multiQubitUnitary(qureg ::QuEST_Types.Qureg,
                           targs ::Vector{QubitIdx},
                           u     ::QuEST_Types.ComplexMatrixN)  ::Nothing

    @assert u.numQubits ≤ qureg.numQubitsRepresented

    ccall(:multiQubitUnitary, Cvoid,
          (QuEST_Types.Qureg, Ptr{Cint}, Cint,           QuEST_Types.ComplexMatrixN),
          qureg,              targs,     length(targs),  u)
end

function multiRotatePauli(qureg         ::QuEST_Types.Qureg,
                          targetQubits  ::Vector{QubitIdx},
                          targetPaulis  ::Vector{QuEST_Types.pauliOpType},
                          angle         ::Qreal)          ::Nothing

    @assert length(targetQubits) == length(targetPaulis)
    #@assert all( σ -> 0 ≤ σ ≤ 3,   targetPaulis )

    ccall(:multiRotatePauli, Cvoid,
          (QuEST_Types.Qureg, Ptr{Cint},     Ptr{QuEST_Types.pauliOpType}, Cint,                 Qreal),
          qureg,              targetQubits,  targetPaulis,                 length(targetPaulis), angle)
end

function multiRotateZ(qureg      ::QuEST_Types.Qureg,
                      qubits     ::Vector{QubitIdx},
                      angle      ::Qreal)         ::Nothing
    ccall(:multiRotateZ, Cvoid,
          (QuEST_Types.Qureg, Ptr{Cint}, Cint,           Qreal),
          qureg,              qubits,    length(qubits), angle)
end

function multiStateControlledUnitary(qureg           ::QuEST_Types.Qureg,
                                     controlQubits   ::Vector{QubitIdx},
                                     controlState    ::Vector{Cint},
                                     targetQubit     ::Integer,
                                     U               ::Matrix{Complex{Qreal}}) ::Nothing

    @assert length(controlQubits) == length(controlState)

    u = _quest_mtx_2(U)

    ccall(:multiStateControlledUnitary, Cvoid,
          (QuEST_Types.Qureg, Ptr{Cint},      Ptr{Cint},    Cint,                  Cint,        QuEST_Types.ComplexMatrix2),
          qureg,              controlQubits,  controlState, length(controlQubits), targetQubit, u)
end

function pauliX(qureg       ::QuEST_Types.Qureg,
                targetQubit ::Integer)         :: Nothing

    ccall(:pauliX, Cvoid,
          (QuEST_Types.Qureg, Cint),
          qureg,              targetQubit)
end

function pauliY(qureg       ::QuEST_Types.Qureg,
                targetQubit ::Integer)         :: Nothing

    ccall(:pauliY, Cvoid,
          (QuEST_Types.Qureg, Cint),
          qureg,              targetQubit)
end
function pauliZ(qureg ::QuEST_Types.Qureg, targetQubit ::Integer):: Nothing

    ccall(:pauliZ,  Cvoid,
          (QuEST_Types.Qureg, Cint),
          qureg,              targetQubit)
end

function phaseShift(qureg       ::QuEST_Types.Qureg,
                    targetQubit ::Integer,
                    angle       ::Qreal) ::Nothing

    ccall(:phaseShift, Cvoid,
          (QuEST_Types.Qureg, Cint,        Qreal),
          qureg,              targetQubit, angle)
end

function rotateAroundAxis(qureg         ::QuEST_Types.Qureg,
                          rotQubit      ::Integer,
                          angle         ::Qreal,
                          axis          ::Union{NTuple{3,Qreal},Vector{Qreal}}) ::Nothing

    @assert length(axis) == 3

    q_axis = QuEST_Types.Vector(axis[1], axis[2], axis[3])

    ccall(:rotateAroundAxis,  Cvoid,
          (QuEST_Types.Qureg, Cint,     Qreal, QuEST_Types.Vector),
          qureg,              rotQubit, angle, q_axis)
end

function rotateX(qureg    ::QuEST_Types.Qureg,
                 rotQubit ::Integer,
                 angle    ::Qreal)           ::Nothing

    ccall(:rotateX,  Cvoid,
          (QuEST_Types.Qureg, Cint,      Qreal),
          qureg,              rotQubit,  angle)
end

function rotateY(qureg    ::QuEST_Types.Qureg,
                 rotQubit ::Integer,
                 angle    ::Qreal)           ::Nothing

    ccall(:rotateY, Cvoid,
          (QuEST_Types.Qureg, Cint,      Qreal),
          qureg,              rotQubit,  angle)
end

function rotateZ(qureg    ::QuEST_Types.Qureg,
                 rotQubit ::Integer,
                 angle    ::Qreal)           ::Nothing

    ccall(:rotateZ, Cvoid,
          (QuEST_Types.Qureg, Cint,     Qreal),
          qureg,              rotQubit, angle)
end

function sGate(qureg       ::QuEST_Types.Qureg,
               targetQubit ::Integer)         ::Nothing

    ccall(:sGate, Cvoid,
          (QuEST_Types.Qureg, Cint),
          qureg,              targetQubit)
end

function sqrtSwapGate(qureg  ::QuEST_Types.Qureg,
                      qubit1 ::Integer,
                      qubit2 ::Integer)         ::Nothing

    ccall(:sqrtSwapGate, Cvoid,
          (QuEST_Types.Qureg, Cint,   Cint),
          qureg,              qubit1, qubit2)
end

function swapGate(qureg  ::QuEST_Types.Qureg,
                  qubit1 ::Integer,
                  qubit2 ::Integer)         ::Nothing

    ccall(:swapGate,  Cvoid,
          (QuEST_Types.Qureg, Cint,   Cint),
          qureg,              qubit1, qubit2)
end

function tGate(qureg       ::QuEST_Types.Qureg,
               targetQubit ::Integer)         ::Nothing

    ccall(:tGate, Cvoid,
          (QuEST_Types.Qureg, Cint),
          qureg,              targetQubit)
end

function twoQubitUnitary(qureg           ::QuEST_Types.Qureg,
                         targetQubit1    ::Integer,
                         targetQubit2    ::Integer,
                         U               ::Matrix{Complex{Qreal}}) ::Nothing

    u = _quest_mtx_4(U)

    ccall(:twoQubitUnitary, Cvoid,
          (QuEST_Types.Qureg, Cint,          Cint,          QuEST_Types.ComplexMatrix4),
          qureg,              targetQubit1,  targetQubit2,  u)
end

function unitary(qureg           ::QuEST_Types.Qureg,
                 targetQubit     ::Integer,
                 U               ::Matrix{Complex{Qreal}}) ::Nothing

    u = _quest_mtx_2(U)

    ccall(:unitary, Cvoid,
          (QuEST_Types.Qureg, Cint,           QuEST_Types.ComplexMatrix2),
          qureg,              targetQubit,    u)
end

#EOF


struct DiscreteModelMock <: DiscreteModel{2,2}
  grid::GridMock
  topo::GridTopologyMock
  labels::FaceLabeling
  function DiscreteModelMock()
    grid = GridMock()
    topo = GridTopologyMock()
    labels = _init_labeling(topo)
    new(grid,topo,labels)
  end
end

function _init_labeling(topo)
  d_to_num_dfaces = [ num_faces(topo,d) for d in 0:num_dims(topo)]
  labels = FaceLabeling(d_to_num_dfaces)
  get_face_entity(labels,0) .= get_isboundary_face(topo,0) .+ 1
  get_face_entity(labels,1) .= get_isboundary_face(topo,1) .+ 1
  get_face_entity(labels,2) .= get_isboundary_face(topo,2) .+ 1
  add_tag!(labels,"interior",[1,])
  add_tag!(labels,"boundary",[2,])
  add_tag_from_tags!(labels,"all",["interior","boundary"])
  labels
end

get_grid(model::DiscreteModelMock) = model.grid

get_grid_topology(model::DiscreteModelMock) = model.topo

get_face_labeling(model::DiscreteModelMock) = model.labels


#=
methods for handling backtest results of strategy objects
=#

mutable struct Results
    trades::Dict{String,TS}
    backtest::Dict{String,TS{Float64}}
    optimization::Matrix{Float64}
    function Results(trades::Dict{String,TS}=Dict{String,TS}(),
                     backtest::Dict{String,TS{Float64}}=Dict{String,TS{Float64}}(),
                     optimization::Matrix{Float64}=Matrix{Float64}(0,0))
        return new(trades, backtest, optimization)
    end
end


function assign_scenario(s::BasePriceScenario, config::VanillaOptions, params::VanillaParameters)
    params.cpricebase
end

function assign_scenario(s::OptimalPriceScenario, config::VanillaOptions, params::VanillaParameters)
    cprice_ubound = fill(Inf, config.N);
    cprice_ubound[1] = params.cpricebase[1];
    # Warning: If parameters are changed, the next equation might make base case infeasible.
    # If so, reduce tnopol so that we don't run out of resources.
    for i in 2:config.N
        if i > config.tnopol
            cprice_ubound[i] = 1000.0;
        end
    end
    cprice_ubound
end

function assign_scenario(s::Scenario, config::VanillaOptions, params::VanillaParameters)
    error("$(s) is not a valid scenario for v2013R (Vanilla flavour)");
end


"""
Convenience code for plotting result of DCPCA

Author: Eric Ford
Created: August 2020
"""

function plot_basis_vectors(λ::VR1, f_mean::V1, deriv::V2, proj::A3;
                              num_basis::Integer = 4, idx_plt::AR = 1:length(f_mean),
                              label = 0 ) where {
                              R1<:AbstractRange, VR1<:AbstractVector{R1}, T1<:Real, V1<:AbstractVector{T1},
                              T2<:Real, V2<:AbstractVector{T2}, T3<:Real, A3<:AbstractArray{T3,2}, AR<:AbstractUnitRange }
  @assert num_basis >=1
  @assert sum(length.(λ)) == length(f_mean)
  @assert length(f_mean) == length(deriv)
  @assert size(proj,1) == length(f_mean)
  @assert minimum(idx_plt) >= 1
  @assert maximum(idx_plt) <= length(f_mean)
  if num_basis > 4   @warn "plot_basis_vectors hasn't been generalized for num_basis!=4."   end
  λ_plt = get_λs(λ,idx_plt)
  λc = mean(λ_plt)
  plt0 = plot(λ_plt, (f_mean[idx_plt].-1.0)./maximum(abs.(f_mean[idx_plt].-1.0)),linecolor=:black, label=:none) #"Std Mean")
  plt0 = plot!(plt0,λ_plt, deriv[idx_plt]./maximum(abs.(deriv[idx_plt])),linecolor=:red,label=:none) #"Std Deriv")
  plt0 = plot!(plt0,[λc,λc],[-1.1,1.0], linecolor=:black, label=:none)
  ylims!(-1.1,1.1)
  pltall = [plt0]
  if num_basis >= 1
      plt0 = plot!(plt0,λ_plt, proj[idx_plt,1]./maximum(abs.(proj[idx_plt,1])),linecolor=:blue,label=:none)
      plt1 = plot(λ_plt, proj[idx_plt,1]./maximum(abs.(proj[idx_plt,1])),linecolor=:blue,label=:none) #"PC 1")
      plt1 = plot!(plt1,[λc,λc],[-1.1,1.0], linecolor=:black, label=:none)
      ylabel!(plt1,"PC1")
      ylims!(-1.1,1.1)
      push!(pltall,plt1)
  end
  if num_basis >= 2
      plt0 = plot!(plt0,λ_plt, proj[idx_plt,2]./maximum(abs.(proj[idx_plt,2])),linecolor=:green,label=:none)
      plt2 = plot(λ_plt, proj[idx_plt,2]./maximum(abs.(proj[idx_plt,2])),linecolor=:green,label=:none) #"PC 2")
      plt2 = plot!(plt2,[λc,λc],[-1.1,1.0], linecolor=:black, label=:none)
      ylabel!(plt2,"PC2")
      ylims!(-1.1,1.1)
      push!(pltall,plt2)
  end
  if num_basis >= 3
      plt3 = plot(λ_plt, proj[idx_plt,3]./maximum(abs.(proj[idx_plt,3])),linecolor=:cyan,label=:none) #"PC 3")
      plt3 = plot!(plt3,[λc,λc],[-1.1,1.0], linecolor=:black, label=:none)
      ylabel!(plt3,"PC3")
      ylims!(-1.1,1.1)
      push!(pltall,plt3)
  end
  if num_basis >= 4
      plt4 = plot(λ_plt, proj[idx_plt,4]./maximum(abs.(proj[idx_plt,4])),linecolor=:magenta,label=:none) #"PC 4")
      plt4 = plot!(plt4,[λc,λc],[-1.1,1.0], linecolor=:black, label=:none)
      ylabel!(plt4,"PC4")
      ylims!(-1.1,1.1)
      push!(pltall,plt4)
  end
  if num_basis == 1
      xlabel!(plt1,"λ (Å)  line# "  * string(label))
      pltall = plot(plt0,plt1, layout = (2,1) )
  elseif num_basis == 2
      xlabel!(plt2,"λ (Å)  line# "  * string(label))
      pltall = plot(plt0,plt1,plt2, layout = (3,1) )
  elseif num_basis == 3
      xlabel!(plt3,"λ (Å)  line# "  * string(label))
      pltall = plot(plt0,plt1,plt2,plt3, layout = (4,1) )
  elseif num_basis == 4
      xlabel!(plt4,"λ (Å)  line# "  * string(label))
      pltall = plot(plt0,plt1,plt2,plt3,plt4, layout = (5,1) )
  elseif num_basis == 5
      xlabel!(plt5,"λ (Å)  line# "  * string(label))
      pltall = plot(plt0,plt1,plt2,plt3,plt4,plt5, layout = (6,1) )
  end

  pltall
end

function plot_basis_scores(times::V1, rvs::V2, scores::A3;
                              num_basis::Integer = 4, idx_plt::AR = 1:length(times) ) where {
                              T1<:Real, V1<:AbstractVector{T1}, T2<:Real, V2<:AbstractVector{T2}, T3<:Real, A3<:AbstractArray{T3,2}, AR<:AbstractUnitRange }
  @assert length(times) >=1
  @assert length(times) == length(rvs)
  @assert length(times) == size(scores,2)
  @assert size(scores,1) >= num_basis
  @assert minimum(idx_plt) >= 1
  @assert maximum(idx_plt) <= length(times)
  if num_basis > 4   @warn "plot_basis_scores hasn't been generalized for num_basis>4."   end
  local plt0 = scatter(times[idx_plt],rvs[idx_plt],label=:none,color=:black)
  ylabel!(plt0,"RV (m/s)")
  if num_basis >= 1
      plt1 = scatter(times[idx_plt],scores[1,idx_plt],label=:none,color=:blue)
      pltall = plot(plt0,plt1, layout = (2,1) )
  end
  if num_basis >= 2
      plt2 = scatter(times[idx_plt],scores[2,idx_plt],label=:none,color=:green)
      pltall = plot(plt0,plt1,plt2, layout = (3,1) )
  end
  if num_basis >= 3
      plt3 = scatter(times[idx_plt],scores[3,idx_plt],label=:none,color=:cyan)
      pltall = plot(plt0,plt1,plt2,plt3, layout = (4,1) )
  end
  if num_basis >= 4
      plt4 = scatter(times[idx_plt],scores[4,idx_plt],label=:none,color=:magenta)
      pltall = plot(plt0,plt1,plt2,plt3,plt4, layout = (5,1) )
  end
  xlabel!(pltall,"Time")
  pltall
end

function plot_basis_scores_cor(rvs::V2, scores::A3;
                              num_basis::Integer = 4, idx_plt::AR = 1:length(rvs) ) where {
                              T2<:Real, V2<:AbstractVector{T2}, T3<:Real, A3<:AbstractArray{T3,2}, AR<:AbstractUnitRange }
  @assert length(rvs) >=1
  @assert length(rvs) == size(scores,2)
  @assert size(scores,1) >= num_basis
  @assert minimum(idx_plt) >= 1
  @assert maximum(idx_plt) <= length(rvs)
  if num_basis > 4   @warn "plot_basis_scores_cor hasn't been generalized for num_basis>4."   end
  if num_basis >= 1
      plt1 = scatter(rvs[idx_plt],vec(scores[1,idx_plt]),xlabel="RV",ylabel="PC1",legend=:none)
      pltall = plt1
  end
  if num_basis >= 2
      plt2 = scatter(rvs[idx_plt],vec(scores[2,idx_plt]),xlabel="RV",ylabel="PC2",legend=:none)
      plt4 = scatter(vec(scores[1,idx_plt]),vec(scores[2,idx_plt]),xlabel="PC1",ylabel="PC2",legend=:none)
      pltall = plot(plt1,plt2,plt4,layout=3)
  end
  if num_basis >= 3
      plt3 = scatter(rvs[idx_plt],vec(scores[3,idx_plt]),xlabel="RV",ylabel="PC3",legend=:none)
      plt5 = scatter(vec(scores[1,idx_plt]),vec(scores[3,idx_plt]),xlabel="PC1",ylabel="PC3",legend=:none)
      plt6 = scatter(vec(scores[2,idx_plt]),vec(scores[3,idx_plt]),xlabel="PC2",ylabel="PC3",legend=:none)
      pltall = plot(plt1,plt2,plt3,plt4,plt5,plt6,layout=(2,3))
  end
  if num_basis >= 4
      plt7 = scatter(vec(scores[1,idx_plt]),vec(scores[4,idx_plt]),xlabel="PC2",ylabel="PC4",legend=:none)
      plt8 = scatter(vec(scores[2,idx_plt]),vec(scores[4,idx_plt]),xlabel="PC2",ylabel="PC4",legend=:none)
      plt9 = scatter(vec(scores[3,idx_plt]),vec(scores[4,idx_plt]),xlabel="PC3",ylabel="PC4",legend=:none)
      pltall = plot(plt1,plt2,plt3,plt4,plt5,plt6,plt7,plt8,plt9,layout=(3,3))
  end
  pltall
end