# Functions to run pedigree models
function run_pedigree_model_multirecord(tbl, markerInfo, pedigree, Rvar, Gvar, Pvar; chainlength = 10000, outputfilename = "MCMC_samples", covariates = [])
    namesG11 = markerInfo[:,3] .* "_" .* markerInfo[:,1] .* "_" .* markerInfo[:,2] .* "_G11"
    namesG12 = markerInfo[:,3] .* "_" .* markerInfo[:,1] .* "_" .* markerInfo[:,2] .* "_G12"
    namesG22 = markerInfo[:,3] .* "_" .* markerInfo[:,1] .* "_" .* markerInfo[:,2] .* "_G22"
    covars = [namesG11[1]; namesG12[1]; namesG22[1]; covariates]
    printcovars = [covars; "anml_keyx"]
    eqcovars = [covars; "anml_key"; "anml_keyx"]
    model_equation = "yd_trait ="
    model_terms = join(eqcovars, " + ")
    model_equation = join([model_equation model_terms], " ")
    print("MODEL EQUATION:  ")
    println(model_equation)
    model = build_model(model_equation, Rvar)
    for term in covars
        set_covariate(model, term)
    end
    set_random(model,"anml_key",pedigree,Gvar)
    set_random(model, "anml_keyx", Pvar)
    for term in printcovars
        outputMCMCsamples(model, term)
    end
    out = runMCMC(model, tbl, output_samples_frequency = 10, chain_length = chainlength, output_folder = outputfilename)
end

function run_pedigree_model(tbl, markerInfo, pedigree, Rvar, Gvar; chainlength = 10000, outputfilename = "MCMC_samples", covariates = [])
    namesG11 = markerInfo[:,3] .* "_" .* markerInfo[:,1] .* "_" .* markerInfo[:,2] .* "_G11"
    namesG12 = markerInfo[:,3] .* "_" .* markerInfo[:,1] .* "_" .* markerInfo[:,2] .* "_G12"
    namesG22 = markerInfo[:,3] .* "_" .* markerInfo[:,1] .* "_" .* markerInfo[:,2] .* "_G22"
    covars = [namesG11[1]; namesG12[1]; namesG22[1]; covariates]
    printcovars = covars
    eqcovars = [covars; "anml_key"]
    model_equation = "yd_trait ="
    model_terms = join(eqcovars, " + ")
    model_equation = join([model_equation model_terms], " ")
    print("MODEL EQUATION:  ")
    println(model_equation)
    model = build_model(model_equation, Rvar)
    for term in covars
        set_covariate(model, term)
    end
    set_random(model,"anml_key",pedigree,Gvar)
    for term in printcovars
        outputMCMCsamples(model, term)
    end
    out = runMCMC(model, tbl, output_samples_frequency = 10, chain_length = chainlength, output_folder = outputfilename)
end

function run_pedigree_model_eQTL(tbl, pedigree, Rvar, Gvar; chainlength = 10000, freq = 10, outputfilename = "eQTL.MCMC_samples")
    pheno = String(names(tbl)[2])
    snp = String(names(tbl)[3])
    model_equation = pheno .* " = intercept + " .* snp .* " + anml_key"
    model = build_model(model_equation, Rvar)
    set_covariate(model, snp)
    set_random(model,"anml_key",pedigree, Gvar)
    outputMCMCsamples(model, snp)
    out = runMCMC(model, tbl, output_samples_frequency = freq, chain_length = chainlength, output_folder = outputfilename)
end

function run_spliceQTL(snp_data, pheno_data, pedigree, Rvar, Gvar; chainlength = 10000, freq = 10, outdir = "MCMC_samples")
    k = size(pheno_data.header, 1)
    m = size(snp_data.header, 1)
    for i in 1:k
        pheno_symbol = Symbol(pheno_data.header[i])
        phenoi = table(pheno_data.iid, pheno_data.data[:,i], names = [:anml_key, pheno_symbol], pkey=:anml_key)
        outputdir = outdir .* "/" .* pheno_data.header[i]
        run(`mkdir -p $outputdir`)
        for j in 1:m
            outputfilename = outputdir .* "/" .* snp_data.header[j]
            snp_symbol = Symbol(snp_data.header[j])
            if any(snp_data.data[:,j] .== "NA") # change NAs to mean(genotype)
                na_index = findall(snp_data.data[:,j] .== "NA")
                snp_data.data[na_index,j] = repeat([missing], size(na_index,1))
                snp_data.data[na_index,j] = repeat([mean(skipmissing(snp_data.data[:,j]))], size(na_index,1))
            end
            snpj = table(snp_data.iid, snp_data.data[:,j], names = [:anml_key, snp_symbol], pkey = :anml_key)
            tbl = join_tables(phenoi, snpj)
            meanA = mean(IndexedTables.select(tbl, snp_symbol))
            tbl = IndexedTables.transform(tbl, snp_symbol => snp_symbol => x -> x - meanA)
            df = DataFrame(tbl)
            run_pedigree_model_eQTL(df, pedigree, Rvar, Gvar; chainlength = chainlength, freq = freq, outputfilename = outputfilename)
        end
    end
end


module ShapeReconstructionPaLS
using jInv
using ParamLevelSet
include("Utils/Utils.jl")
include("DipTransform/DipTransform.jl")
include("ShapeFromSilhouette/ShapeFromSilhouette.jl")
include("Direct/Direct.jl")
include("PointCloud/PointCloud.jl")
end # module


# Read the spectrum data file format associated with handheld Bruker XRF units.

function readbrukerpdz(fn::AbstractString)
    open(fn, "r") do f
        res = readbrukerpdz(f)
        res[:Name] = splitpath(fn)[end]
        return res
    end
end

"""
    readbrukerpdz(io::IO|fn::AbstractString)

Read the default data file type for the Bruker handheld XRF unit. (Reverse engineered.)
"""
function readbrukerpdz(io::IO)
    mm = Mmap.mmap(io)
    nCh = reinterpret(UInt16, mm[7:8])[1]
    eVpCh = reinterpret(Float64, mm[51:58])[1]
    yr, mon, _, day, hour, min, sec = reinterpret(UInt16, mm[147:160])
    e0, pc = reinterpret(Float32, mm[163:170])
    rt, dt, _, lt = reinterpret(Float32, mm[343:358])
    filt = reinterpret(UInt16, mm[115:130]) # Filter / thickness
    filters = [
        (PeriodicTable.elements[filt[2*i-1]], filt[2*i] * 1.0e-6) for
        i in filter(i -> filt[2*i-1] > 0, 1:4)
    ]
    dt = DateTime(yr, mon, day, hour, min, sec)
    props = Dict{Symbol,Any}(
        :RealTime => convert(Float64, rt),
        :LiveTime => convert(Float64, lt),
        :AcquisitionTime => dt,
        :BeamEnergy => convert(Float64, e0) * 1000.0,
        :ProbeCurrent => convert(Float64, pc) * 1000.0,
        :XRFFilter => filters,
    )
    counts = reinterpret(UInt32, mm[length(mm)-4*nCh+1:length(mm)])
    return Spectrum(LinearEnergyScale(0.0, eVpCh), collect(counts), props)
end

function detectBrukerPDZ(ios::IO)
    try
        return read(io, UInt32) == 0x00170101
    catch
        return false
    end
end


using ArgParse
using BSON
using MarkovModels
using Distributed

# BUT specific package to work with the SGE cluster
# See https://github.com/BUTSpeechFIT/BUTSGEManager
using BUTSGEManager

function getargs()
    s = ArgParseSettings()
    @add_arg_table! s begin
        "--add-sil", "-s"
            help = "add silence add the beginning/end of each utterance"
            action = :store_true
        "--jobs", "-j"
            help = "number of jobs to run in parallel"
            arg_type = Int
            default = 4
        "--jobs-args", "-a"
            help = "arguments to pass to the parallel engine"
            arg_type = String
            default = ""
        "trans"
            help = "list of word level transcription"
            arg_type = String
            required = true
        "lexicon"
            help = "mapping word -> pronunciation"
            arg_type = String
            required = true
        "hmms"
            help = "unit hmms"
            arg_type = String
            required = true
        "outdir"
            help = "output directory where will be stored the alignment graphs"
            arg_type = String
            required = true
    end
    args = parse_args(s)
end

const args = getargs()
const trans = args["trans"]
const lexiconfile = args["lexicon"]
const hmmsfile = args["hmms"]
const outdir = args["outdir"]

@info "loading the pronunciations"
const pronuns = Dict{String, Vector{FSM}}()
open(lexiconfile, "r") do f
    for line in eachline(f)
        tokens = split(line)
        word = tokens[1]
        pronun = tokens[2:end]
        ps = get(pronuns, word, FSM[])
        push!(ps, LinearFSM(pronun))
        pronuns[word] = ps
    end
end

@info "compiling pronunciation graphs"
const lexicon = Dict{String, FSM}()
for (k, v) in pronuns
    lexicon[k] = union(v...) |> minimize!
end

# Make sure the directory exists
run(`mkdir -p $outdir`)

@info "initializing the workers"
addprocs(SGEManager(args["jobs"]), args = args["jobs-args"],
         exeflags = "--project=$(Base.active_project())")
@everywhere using BSON
@everywhere using MarkovModels
@everywhere const outdir = $outdir
@everywhere const lexicon = $lexicon
@everywhere const hmmsfile = $hmmsfile
@everywhere hmms = BSON.load(hmmsfile)[:hmms]

@info "compiling alignment graphs"
@sync @distributed for line in readlines(trans)
    tokens = split(line)
    uttid = tokens[1]
    words = tokens[2:end]
    if args["add-sil"]
        insert!(words, 1, "sil")
        push!(words, "sil")
    end

    ali = compose!(compose!(LinearFSM(words), lexicon), hmms) |> removenilstates!
    bson(joinpath(outdir, uttid * ".bson"), Dict(:ali => ali))
end


"""
    Calculators

This module exists to bridge the gap between the `Models` and the `Dynamics`.

Here we provide functions and types for evaluating and storing quantities obtained from the
`Models`.
In addition any further manipulation of those quantities, such as computing eigenvalues,
is included here.

This module is largely needed to facilitate seemless integration of both ring polymer and
classical dynamics, using the same models and functions for both.
Specific ring polymer types are provided that have the extra fields and methods needed
to evaluate the quantities for each bead.
"""
module Calculators

using LinearAlgebra: LinearAlgebra, Hermitian, I, Eigen, tr
using StaticArrays: SMatrix, SVector

using NQCModels: NQCModels, Model, nstates
using NQCModels.AdiabaticModels: AdiabaticModel
using NQCModels.DiabaticModels: DiabaticModel, DiabaticFrictionModel
using NQCModels.FrictionModels: AdiabaticFrictionModel

using NQCDynamics: RingPolymers, ndofs

"""
    AbstractCalculator{M<:Model}

Top-level type for all calculators.

Each concrete calculator contains the `Model` and the fields to store the quantities
obtained from the model.
"""
abstract type AbstractCalculator{T,M<:Model} end
abstract type AbstractAdiabaticCalculator{T,M<:AdiabaticModel} <: AbstractCalculator{T,M} end
abstract type AbstractDiabaticCalculator{T,M<:Union{DiabaticFrictionModel,DiabaticModel}} <: AbstractCalculator{T,M} end
abstract type AbstractStaticDiabaticCalculator{T,M} <: AbstractDiabaticCalculator{T,M} end
abstract type AbstractFrictionCalculator{T,M<:AdiabaticFrictionModel} <: AbstractCalculator{T,M} end

Base.eltype(::AbstractCalculator{T}) where {T} = T

mutable struct AdiabaticCalculator{T,M} <: AbstractAdiabaticCalculator{T,M}
    model::M
    potential::T
    derivative::Matrix{T}
    function AdiabaticCalculator{T}(model::M, atoms::Integer) where {T,M<:Model}
        new{T,M}(model, 0, zeros(ndofs(model), atoms))
    end
end

struct RingPolymerAdiabaticCalculator{T,M} <: AbstractAdiabaticCalculator{T,M}
    model::M
    potential::Vector{T}
    derivative::Array{T,3}
    function RingPolymerAdiabaticCalculator{T}(model::M, atoms::Integer, beads::Integer) where {T,M<:Model}
        new{T,M}(model, zeros(beads), zeros(ndofs(model), atoms, beads))
    end
end

mutable struct DiabaticCalculator{T,M,S,L} <: AbstractDiabaticCalculator{T,M}
    model::M
    potential::Hermitian{T,SMatrix{S,S,T,L}}
    derivative::Matrix{Hermitian{T,SMatrix{S,S,T,L}}}
    eigen::LinearAlgebra.Eigen{T,T,SMatrix{S,S,T,L},SVector{S,T}}
    adiabatic_derivative::Matrix{SMatrix{S,S,T,L}}
    nonadiabatic_coupling::Matrix{SMatrix{S,S,T,L}}
    tmp_position::Matrix{T}
    tmp_mat::Matrix{T}
    tmp_mat_complex1::Matrix{Complex{T}}
    tmp_mat_complex2::Matrix{Complex{T}}
    function DiabaticCalculator{T}(model::M, atoms::Integer) where {T,M<:Model}
        n = nstates(model)
        matrix_template = NQCModels.DiabaticModels.matrix_template(model, T)
        vector_template = NQCModels.DiabaticModels.vector_template(model, T)

        potential = Hermitian(matrix_template)
        derivative = [Hermitian(matrix_template) for _=1:ndofs(model), _=1:atoms]
        eigen = Eigen(vector_template, matrix_template + I)
        adiabatic_derivative = [matrix_template for _ in CartesianIndices(derivative)]
        nonadiabatic_coupling = [matrix_template for _ in CartesianIndices(derivative)]
        tmp_position = zeros(T, ndofs(model), atoms)
        tmp_mat = zeros(T, n, n)
        tmp_mat_complex1 = zeros(Complex{T}, n, n)
        tmp_mat_complex2 = zeros(Complex{T}, n, n)

        new{T,M,n,n^2}(model,
            potential, derivative, eigen,
            adiabatic_derivative, nonadiabatic_coupling,
            tmp_position, tmp_mat, tmp_mat_complex1, tmp_mat_complex2)
    end
end

mutable struct RingPolymerDiabaticCalculator{T,M,S,L} <: AbstractDiabaticCalculator{T,M}
    model::M
    potential::Vector{Hermitian{T,SMatrix{S,S,T,L}}}
    traceless_potential::Vector{Hermitian{T,SMatrix{S,S,T,L}}}
    V̄::Vector{T}
    derivative::Array{Hermitian{T,SMatrix{S,S,T,L}},3}
    traceless_derivative::Array{Hermitian{T,SMatrix{S,S,T,L}},3}
    D̄::Array{T,3}
    eigen::Vector{LinearAlgebra.Eigen{T,T,SMatrix{S,S,T,L},SVector{S,T}}}
    adiabatic_derivative::Array{SMatrix{S,S,T,L},3}
    traceless_adiabatic_derivative::Array{SMatrix{S,S,T,L},3}
    nonadiabatic_coupling::Array{SMatrix{S,S,T,L},3}

    centroid_potential::Hermitian{T,SMatrix{S,S,T,L}}
    centroid_derivative::Matrix{Hermitian{T,SMatrix{S,S,T,L}}}
    centroid_eigen::LinearAlgebra.Eigen{T,T,SMatrix{S,S,T,L},SVector{S,T}}
    centroid_adiabatic_derivative::Matrix{SMatrix{S,S,T,L}}
    centroid_nonadiabatic_coupling::Matrix{SMatrix{S,S,T,L}}
    tmp_centroid::Matrix{T}

    tmp_mat::Matrix{T}
    tmp_mat_complex1::Matrix{Complex{T}}
    tmp_mat_complex2::Matrix{Complex{T}}
    function RingPolymerDiabaticCalculator{T}(model::M, atoms::Integer, beads::Integer) where {T,M<:Model}
        n = nstates(model)
        matrix_template = NQCModels.DiabaticModels.matrix_template(model, T)
        vector_template = NQCModels.DiabaticModels.vector_template(model, T)

        potential = [Hermitian(matrix_template) for _=1:beads]
        traceless_potential = [Hermitian(matrix_template) for _=1:beads]
        V̄ = zeros(beads)

        derivative = [Hermitian(matrix_template) for _=1:ndofs(model), _=1:atoms, _=1:beads]
        traceless_derivative = [Hermitian(matrix_template) for _=1:ndofs(model), _=1:atoms, _=1:beads]
        D̄ = zeros(ndofs(model), atoms, beads)

        adiabatic_derivative = [matrix_template for _=1:ndofs(model), _=1:atoms, _=1:beads]
        traceless_adiabatic_derivative = [Hermitian(matrix_template) for _=1:ndofs(model), _=1:atoms, _=1:beads]

        eigen = [Eigen(vector_template, matrix_template + I) for _=1:beads]
        nonadiabatic_coupling = [matrix_template for _=1:ndofs(model), _=1:atoms, _=1:beads]

        centroid_potential = Hermitian(matrix_template)
        centroid_derivative = [Hermitian(matrix_template) for _=1:ndofs(model), _=1:atoms]
        centroid_eigen = Eigen(vector_template, matrix_template + I)
        centroid_adiabatic_derivative = [matrix_template for _ in CartesianIndices(centroid_derivative)]
        centroid_nonadiabatic_coupling = [matrix_template for _ in CartesianIndices(centroid_derivative)]
        tmp_centroid = zeros(T, ndofs(model), atoms)

        tmp_mat = zeros(T, n, n)
        tmp_mat_complex1 = zeros(Complex{T}, n, n)
        tmp_mat_complex2 = zeros(Complex{T}, n, n)
        new{T,M,n,n^2}(model,
            potential, traceless_potential, V̄,
            derivative, traceless_derivative, D̄,
            eigen, adiabatic_derivative, traceless_adiabatic_derivative, nonadiabatic_coupling,
            centroid_potential, centroid_derivative, centroid_eigen,
            centroid_adiabatic_derivative, centroid_nonadiabatic_coupling, tmp_centroid,
            tmp_mat, tmp_mat_complex1, tmp_mat_complex2)
    end
end

function Calculator(model::DiabaticModel, atoms::Integer, t::Type{T}) where {T}
    DiabaticCalculator{t}(model, atoms)
end
function Calculator(model::AdiabaticModel, atoms::Integer, t::Type{T}) where {T}
    AdiabaticCalculator{t}(model, atoms)
end
function Calculator(model::DiabaticModel, atoms::Integer, beads::Integer, t::Type{T}) where {T}
    RingPolymerDiabaticCalculator{t}(model, atoms, beads)
end
function Calculator(model::AdiabaticModel, atoms::Integer, beads::Integer, t::Type{T}) where {T}
    RingPolymerAdiabaticCalculator{t}(model, atoms, beads)
end

function evaluate_potential!(calc::AbstractCalculator, R)
    calc.potential = NQCModels.potential(calc.model, R)
end

function evaluate_potential!(calc::AbstractCalculator, R::AbstractArray{T,3}) where {T}
    @views for i in axes(R, 3)
        calc.potential[i] = NQCModels.potential(calc.model, R[:,:,i])
    end
end

function evaluate_V̄!(calc::RingPolymerDiabaticCalculator)
    for i in 1:length(calc.V̄)
        calc.V̄[i] = tr(calc.potential[i]) / nstates(calc.model)
    end
end

function evaluate_traceless_potential!(calc::RingPolymerDiabaticCalculator)
    n = nstates(calc.model)
    for I in 1:length(calc.traceless_potential)
        calc.traceless_potential[I] = Hermitian(SMatrix{n,n}(
            i != j ? calc.potential[I][j,i] : calc.potential[I][j,i] - calc.V̄[I] for j=1:n, i=1:n
        ))
    end
end

function evaluate_centroid_potential!(calc::AbstractCalculator, R::AbstractMatrix)
    calc.centroid_potential = NQCModels.potential(calc.model, R)
end

function evaluate_derivative!(calc::AbstractCalculator, R)
    NQCModels.derivative!(calc.model, calc.derivative, R)
end

function evaluate_derivative!(calc::AbstractCalculator, R::AbstractArray{T,3}) where {T}
    @views for i in axes(R, 3)
        NQCModels.derivative!(calc.model, calc.derivative[:,:,i], R[:,:,i])
    end
end

function evaluate_D̄!(calc::RingPolymerDiabaticCalculator)
    for I in eachindex(calc.D̄)
        calc.D̄[I] = tr(calc.derivative[I]) / nstates(calc.model)
    end
end

function evaluate_traceless_derivative!(calc::RingPolymerDiabaticCalculator)
    n = nstates(calc.model)
    for I in eachindex(calc.derivative)
        calc.traceless_derivative[I] = Hermitian(SMatrix{n,n}(
            i != j ? calc.derivative[I][j,i] : calc.derivative[I][j,i] - calc.D̄[I] for j=1:n, i=1:n
        ))
    end
end

function evaluate_traceless_adiabatic_derivative!(calc::RingPolymerDiabaticCalculator)
    n = nstates(calc.model)
    for I in eachindex(calc.adiabatic_derivative)
        calc.traceless_adiabatic_derivative[I] = SMatrix{n,n}(
            i != j ? calc.adiabatic_derivative[I][j,i] : calc.adiabatic_derivative[I][j,i] - calc.D̄[I] for j=1:n, i=1:n
        )
    end
end

function evaluate_centroid_derivative!(calc::AbstractCalculator, R::AbstractMatrix)
    NQCModels.derivative!(calc.model, calc.centroid_derivative, R)
end

function eigen!(calc::DiabaticCalculator)
    eig = LinearAlgebra.eigen(calc.potential)
    corrected_vectors = correct_phase(eig.vectors, calc.eigen.vectors)
    calc.eigen = Eigen(eig.values, corrected_vectors)
    return nothing
end

function centroid_eigen!(calc::RingPolymerDiabaticCalculator)
    eig = LinearAlgebra.eigen(calc.centroid_potential)
    corrected_vectors = correct_phase(eig.vectors, calc.centroid_eigen.vectors)
    calc.centroid_eigen = Eigen(eig.values, corrected_vectors)
    return nothing
end

function eigen!(calc::RingPolymerDiabaticCalculator)
    for i=1:length(calc.potential)
        eig = LinearAlgebra.eigen(calc.potential[i])
        corrected_vectors = correct_phase(eig.vectors, calc.eigen[i].vectors)
        calc.eigen[i] = Eigen(eig.values, corrected_vectors)
    end
    return nothing
end

function correct_phase(new_vectors::SMatrix, old_vectors::SMatrix)
    n = size(new_vectors, 1)
    vect = SVector{n}(sign(LinearAlgebra.dot(new_vectors[:,i], old_vectors[:,i])) for i=1:n)
    return new_vectors .* vect'
end

function transform_derivative!(calc::AbstractDiabaticCalculator)
    for I in eachindex(calc.derivative)
        calc.adiabatic_derivative[I] = calc.eigen.vectors' * calc.derivative[I] * calc.eigen.vectors
    end
end

function transform_centroid_derivative!(calc::RingPolymerDiabaticCalculator)
    for I in eachindex(calc.centroid_derivative)
        calc.centroid_adiabatic_derivative[I] = calc.centroid_eigen.vectors' * calc.centroid_derivative[I] * calc.centroid_eigen.vectors
    end
end

function transform_derivative!(calc::RingPolymerDiabaticCalculator)
    for i in axes(calc.derivative, 3) # Beads
        for j in axes(calc.derivative, 2) # Atoms
            for k in axes(calc.derivative, 1) # DoFs
                calc.adiabatic_derivative[k,j,i] = calc.eigen[i].vectors' * calc.derivative[k,j,i] * calc.eigen[i].vectors
            end
        end
    end
end

function evaluate_nonadiabatic_coupling!(calc::AbstractDiabaticCalculator)
    for I in eachindex(calc.adiabatic_derivative)
        calc.nonadiabatic_coupling[I] = evaluate_nonadiabatic_coupling(calc.adiabatic_derivative[I], calc.eigen.values)
    end
end

function evaluate_centroid_nonadiabatic_coupling!(calc::RingPolymerDiabaticCalculator)
    for I in eachindex(calc.centroid_adiabatic_derivative)
        calc.centroid_nonadiabatic_coupling[I] = evaluate_nonadiabatic_coupling(calc.centroid_adiabatic_derivative[I], calc.centroid_eigen.values)
    end
end

function evaluate_nonadiabatic_coupling!(calc::RingPolymerDiabaticCalculator)
    for i in axes(calc.nonadiabatic_coupling, 3) # Beads
        for I in CartesianIndices(size(calc.adiabatic_derivative)[1:2])
            calc.nonadiabatic_coupling[I,i] = evaluate_nonadiabatic_coupling(calc.adiabatic_derivative[I,i], calc.eigen[i].values)
        end
    end
end

"""
# References

- HammesSchifferTully_JChemPhys_101_4657_1994 Eq. (32)
- SubotnikBellonzi_AnnuRevPhyschem_67_387_2016, section 2.3
"""
function evaluate_nonadiabatic_coupling(adiabatic_derivative::SMatrix, eigenvalues::SVector)
    n = length(eigenvalues)
    SMatrix{n,n}(
        (i != j ? adiabatic_derivative[j,i] / (eigenvalues[i] - eigenvalues[j]) : 0
        for j=1:n, i=1:n))
end

"""
Evaluates all electronic properties for the current position `r`.

# Properties evaluated:
- Diabatic potential
- Diabatic derivative
- Eigenvalues and eigenvectors
- Adiabatic derivative
- Nonadiabatic coupling
"""
function update_electronics!(calculator::AbstractDiabaticCalculator, r::AbstractArray)
    evaluate_potential!(calculator, r)
    evaluate_derivative!(calculator, r)
    eigen!(calculator)
    transform_derivative!(calculator)
    evaluate_nonadiabatic_coupling!(calculator)
end

function update_electronics!(calculator::RingPolymerDiabaticCalculator, r::AbstractArray{T,3}) where {T}
    evaluate_potential!(calculator, r)
    evaluate_derivative!(calculator, r)
    eigen!(calculator)
    transform_derivative!(calculator)
    evaluate_nonadiabatic_coupling!(calculator)

    update_centroid_electronics!(calculator, r)
end

function update_centroid_electronics!(calculator::RingPolymerDiabaticCalculator, r::AbstractArray{T,3}) where {T}
    RingPolymers.get_centroid!(calculator.tmp_centroid, r)
    evaluate_centroid_potential!(calculator, calculator.tmp_centroid)
    evaluate_centroid_derivative!(calculator, calculator.tmp_centroid)
    centroid_eigen!(calculator)
    transform_centroid_derivative!(calculator)
    evaluate_centroid_nonadiabatic_coupling!(calculator)
end

include("large_diabatic.jl")
include("friction.jl")

position(calc) = calc.tmp_position

for field ∈ (:potential, :derivative)
    get_field = Symbol(:get, :_, field)
    evaluate_field = Symbol(:evaluate, :_, field, :!)

    eval(quote
        function $get_field(calc, r)
            if position(calc) ≈ r
                return calc.$field
            else
                copyto!(position(calc), r)
                return $evaluate_field(calc, r)
            end
        end
    end)
end

end # module


struct Operations end

Primitives() = Multi(Operations(), Basic())

exprtype(ir, x) = IRTools.exprtype(ir, x, typeof = Const)

xtypeof(x::XScalar) = Const(x)
xtypeof(x::Tuple) = ptuple(xtypeof.(x)...)
xtypeof(x::Array{<:XScalar}) = Mjolnir.Shape{typeof(x)}(size(x))
xtypeof(x::XArray{T,N}) where {T,N} = Mjolnir.Shape{Array{T,N}}(size(x))
xtypeof(x) = isbits(x) && nfields(x) == 0 ? Const(x) : Partial{typeof(x)}((; map(f -> f=>xtypeof(getfield(x, f)), fieldnames(typeof(x)))...))

layout(x::Type{<:XScalar}) = [x]
layout(::Const) = []
layout(x::XShape) = [x]
layout(x::Mjolnir.Shape) = [XShape(eltype(x), size(x))]
layout(x::Type{<:Array{<:XScalar}}) = [x]
layout(x::Partial) = vcat(map(i -> layout(x.value[i]), 1:fieldcount(widen(x)))...)
layout(x::Type) = vcat(map(f -> layout(fieldtype(x, f)), fieldnames(x))...)
layout(x::Mjolnir.Node) = layout(widen(x))

subtype(T::Partial, i) = T.value[i]
subtype(T::Type, i) = fieldtype(T, i)

function subrange(T, i)
  offset = sum(Int[length(layout(subtype(T, j))) for j = 1:i-1])
  len = length(layout(subtype(T, i)))
  return offset .+ (1:len)
end

struct Print
  data
end

@abstract Operations repr(x) = x
@abstract Operations println(xs...) = Partial{Print}(Any[Mjolnir.ptuple(xs...)])

@abstract Operations Base.vect(xs::Const...) = Const([x.value for x in xs])

@abstract Operations NNlib.DenseConvDims(xs::Const...; kw...) =
  Const(DenseConvDims(map(x -> x.value, xs)...; kw...))

@abstract Operations NNlib.PoolDims(xs::Const...; kw...) =
  Const(PoolDims(map(x -> x.value, xs)...; kw...))

@abstract Operations conv(x::Array, w::Array, dims::Const{<:DenseConvDims}) =
  Shape{widen(x)}((NNlib.output_size(dims.value)..., size(w)[end], size(x)[end]))

@abstract Operations ∇conv_data(dy::Array, w::Array, dims::Const{<:DenseConvDims}) =
  Shape{widen(dy)}((NNlib.input_size(dims.value)..., NNlib.channels_in(dims.value), size(dy)[end]))

@abstract Operations ∇conv_filter(dy::Array, w::Array, dims::Const{<:DenseConvDims}) =
  Shape{widen(dy)}((NNlib.kernel_size(dims.value)..., NNlib.channels_in(dims.value), NNlib.channels_out(dims.value)))

function transpose!(ir, v, sz = size(exprtype(ir, v)))
  insertafter!(ir, v, xcall(Reshape(length(sz):-1:1, reverse(sz)), v))
end

xpadding() = ()
xpadding(a, b, ps...) = ((a, b), xpadding(ps...)...)

function xlaop!(ir, v, ::AType{typeof(conv)}, X, W, ::AType{DenseConvDims{N,K,C_in,C_out,S,P,D,F}}) where {N,K,C_in,C_out,S,P,D,F}
  @assert length(size(X)) == 4 "Only 2D conv currently supported"
  _, x, w = ir[v].expr.args
  F || (w = insertafter!(ir, w, xcall(Rev([1, 2]), w)))
  ir[v] = xcall(Conv(S, xpadding(P...), [1, 1], [1, 1]), x, w)
end

function xlaop!(ir, v, ::AType{typeof(∇conv_data)}, Ȳ, W, ::AType{DenseConvDims{N,K,C_in,C_out,S,P,D,F}}) where {N,K,C_in,C_out,S,P,D,F}
  @assert length(size(Ȳ)) == 4 "Only 2D conv currently supported"
  _, dy, w = ir[v].expr.args
  F && (w = insertafter!(ir, w, xcall(Rev([1, 2]), w)))
  w = insertafter!(ir, w, xcall(Reshape([1, 2, 4, 3], getindex.((size(W),), [1, 2, 4, 3])), w))
  p = xpadding(P...)
  p = ntuple(i -> size(W)[i]-1 .- p[i], length(p))
  ir[v] = xcall(Conv(S, p, [1, 1], [1, 1]), dy, w)
end

function xlaop!(ir, v, ::AType{typeof(∇conv_filter)}, X, Ȳ, dims::Const{DenseConvDims{N,K,C_in,C_out,S,P,D,F}}) where {N,K,C_in,C_out,S,P,D,F}
  @assert length(size(X)) == 4 "Only 2D conv currently supported"
  _, x, dy = ir[v].expr.args
  W = exprtype(ir, v)
  x = insertafter!(ir, x, xcall(Reshape([1, 2, 4, 3], getindex.((size(X),), [1, 2, 4, 3])), x))
  dy = insertafter!(ir, dy, xcall(Reshape([1, 2, 4, 3], getindex.((size(Ȳ),), [1, 2, 4, 3])), dy))
  p = xpadding(P...)
  dw = insert!(ir, v, xcall(Conv(S, p, [1, 1], [1, 1]), x, dy))
  F || (dw = insertafter!(ir, dw, xcall(Rev([1, 2]), dw)))
  ir[v] = xcall(Reshape([1, 2, 4, 3], size(W)), dw)
end

@abstract Operations function maxpool(x::Array, dims::Const{<:PoolDims})
  Shape{widen(x)}((NNlib.output_size(dims.value)..., size(x)[end-1:end]...))
end

@abstract Operations function ∇maxpool(dy::Array, y::Array, x::Array, dims::Const{<:PoolDims})
  Shape{widen(x)}(size(x))
end

poolwindow(x, k) = ntuple(i -> i > length(k) ? 1 : k[i], length(x))

function xlaop!(ir, v, ::AType{typeof(maxpool)}, X, dims::Const{<:PoolDims})
  _, x = ir[v].expr.args
  reducer = insert!(ir, v, code_xla(max, eltype(X), eltype(X)))
  init = zero(eltype(X))
  ir[v] = xcall(ReduceWindow(poolwindow(size(X), NNlib.kernel_size(dims.value))), reducer, x, init)
end

function xlaop!(ir, v, ::AType{typeof(∇maxpool)}, Ȳ, Y, X, dims::Const{<:PoolDims})
  _, dy, y, x = ir[v].expr.args
  select = insert!(ir, v, code_xla(>, eltype(X), eltype(X)))
  scatter = insert!(ir, v, code_xla((x, y) -> y, eltype(X), eltype(X)))
  ir[v] = xcall(SelectAndScatter(poolwindow(size(X), NNlib.kernel_size(dims.value))), select, scatter, x, dy, zero(eltype(X)))
end

# Base's `<` does something complicated.
simplelt(a, b) = <(promote(a, b)...)

instead(::Operations, args, ::AType{typeof(<)}, a::AType{<:XScalar}, b::AType{<:XScalar}) =
  widen(a) == widen(b) ? nothing : ([simplelt, args[2], args[3]], (Const(simplelt), a, b))

instead(::Operations, args, ::AType{typeof(//)}, a::AType{<:XScalar}, b::AType{<:XScalar}) =
  [(/), args[2], args[3]], [Const(/), a, b]

@abstract Operations convert(::Type{T}, x::Const{<:XScalar}) where T<:XScalar =
  Const(convert(T, x.value))

@abstract Operations convert(::Type{T}, S::XScalar) where T<:XScalar =
  widen(S) <: T ? S : T

@abstract Operations (::Type{T})(S::XScalar) where T<:XScalar =
  abstract(Operations(), Const(convert), Const(T), S)

@abstract Operations fill(x::Const, sz::Const) = Const(fill(x.value, sz.value))
@abstract Operations ones(x::Const{<:Tuple{Vararg{Integer}}}) = Const(ones(x.value))

xlaop(args, ::AType{typeof(convert)}, ::AType{Type{T}}, _) where T<:XScalar =
  xcall(ConvertElementType(T), args[3])

xlaop(args, ::AType{Type{T}}, ::AType{<:XScalar}) where T<:XScalar =
  xcall(ConvertElementType(T), args[2])

for (op, xop) in [(+, :Add), (*, :Mul), (-, :Sub), (/, :Div), (^, :Pow), (>, :Gt),
                  (<, :Lt), (>=, :Ge), (<=, :Le), (max, :Max)]
  @eval @abstract Operations $op(a::Const{T}, b::Const{T}) where T<:XScalar =
    Const($op(a.value, b.value))
  @eval @abstract Operations $op(a::AType{T}, b::AType{T}) where T<:XScalar =
    Core.Compiler.return_type($op, Tuple{T,T})
  @eval xlaop(args, ::AType{typeof($op)}, a::AType{T}, b::AType{T}) where T<:XScalar =
          xcall($xop(), args[2:end]...)
end

@abstract Operations -(a::Const{T}) where T<:XScalar = Const(-a.value)
@abstract Operations -(a::AType{T}) where T<:XScalar = T
xlaop(args, ::AType{typeof(-)}, a::AType{T}) where T<:XScalar =
        xcall(Neg(), args[2:end]...)

for (op, xop) in [(exp, :Exp), (sin, :Sin), (cos, :Cos), (log, :Log)]
  @eval @abstract Operations $op(a::XFloat) = widen(a)
  @eval xlaop(args, ::AType{typeof($op)}, _) = xcall($xop(), args[2])
end

for (op, xop) in [(+, :Add), (-, :Sub)]
  @eval @abstract Operations $op(a::T, b::T) where T<:Array{<:XScalar} =
    Shape{Array{eltype(a),ndims(a)}}(size(a))
  @eval xlaop(args, ::AType{typeof($op)}, a::AType{T}, b::AType{T}) where T<:Array{<:XScalar} =
          xcall($xop(), args[2:end]...)
end

@abstract Operations reshape(A::Array{T}, i::NTuple{N,Int}) where {T<:XScalar,N} =
  Shape{Array{T,length(i.value)}}(i.value)

@abstract Operations function reshape(A::Array{T}, i::NTuple{N,Union{Int,Colon}}) where {T<:XScalar,N}
  dims = Base._reshape_uncolon(Fill(0, size(A)), i.value)
  Shape{Array{T,length(i.value)}}(dims)
end

xlaop(args, ::AType{typeof(reshape)}, x, sh::Const{NTuple{N,Int}}) where N =
  xcall(Reshape(1:ndims(x), sh.value), args[2])

xlaop(args, ::AType{typeof(reshape)}, x, sh::Const{<:Tuple{Vararg{Union{Int,Colon}}}}) =
  xcall(Reshape(1:ndims(x), Base._reshape_uncolon(Fill(0, size(x)), sh.value)), args[2])

@abstract Operations getindex(A::Vector{T}, i::Integer) where T<:XScalar = T

function xlaop!(ir, v, ::AType{typeof(getindex)}, A, i)
  args = ir[v].expr.args
  x = insert!(ir, v, xcall(DynamicSlice([1]), args[2], args[3:end]...))
  ir[v] = xcall(Reshape([1], []), x)
end

@abstract Operations function (a::Matrix{T} * b::Vector{T}) where T<:XScalar
  a isa Const && b isa Const && return Const(a.value * b.value)
  n, m = size(a)
  m == size(b)[1] || error("Dimension mismatch")
  Mjolnir.Shape{Vector{T}}((n,))
end

@abstract Operations function (a::Matrix{T} * b::Matrix{T}) where T<:XScalar
  a isa Const && b isa Const && return Const(a.value * b.value)
  n, m = size(a)
  m == size(b)[1] || error("Dimension mismatch")
  Mjolnir.Shape{Matrix{T}}((n,size(b)[2]))
end

@abstract Operations function (a::Vector{T} * b::Matrix{T}) where T<:XScalar
  a isa Const && b isa Const && return Const(a.value * b.value)
  n, = size(a)
  size(b)[1] == 1 || error("Dimension mismatch")
  Mjolnir.Shape{Vector{T}}((n,size(b)[2]))
end

function xlaop!(ir, v, ::AType{typeof(*)}, A::AType{<:Array{T}}, B::AType{<:Array{T}}) where T<:XScalar
  _, a, b = ir[v].expr.args
  A isa AType{<:Vector} && (a = insert!(ir, v, xcall(Reshape([1], [size(A)[1], 1]), a)))
  ir[v] = xcall(Dot(), a, b)
end

@abstract Operations adjoint(x::Shape{Vector{T}}) where T<:XScalar = Mjolnir.Shape{Matrix{T}}((1, size(x)[1]))
@abstract Operations adjoint(x::Shape{Matrix{T}}) where T<:XScalar = Mjolnir.Shape{Matrix{T}}(reverse(size(x)))
@abstract Operations adjoint(x::Const{<:Array{<:XScalar}}) = Const(collect(adjoint(x.value)))

xlaop(args, ::AType{typeof(adjoint)}, x::AType{<:Vector}) = xcall(Reshape([1], [1, size(x)[1]]), args[2])
xlaop(args, ::AType{typeof(adjoint)}, x::AType{<:Matrix}) = xcall(Reshape([2, 1], [reverse(size(x))...]), args[2])

xlaop(args, ::AType{typeof(repr)}, x) = args[2]

fieldnum(T, f) = findfirst(==(f), fieldnames(T))

function strip_self_arg!(ir)
  @assert exprtype(ir, arguments(ir)[1]) isa Const
  deletearg!(ir, 1)
  return ir
end

function tuplerange!(ir, v, xs, is)
  part(i) = length(layout(exprtype(ir, xs))) == 1 ? xs : insert!(ir, v, xcall(GetTupleElement(i-1), xs))
  ir[v] = length(is) == 1 ? part(is[]) : xcall(XTuple(), part.(is)...)
end

function tuplecat!(ir, v, xs)
  layouts = layout.(exprtype.((ir,), xs))
  if sum(length.(layouts)) == 1
    ir[v] = xs[findfirst(x -> !isempty(x), layouts)]
  else
    parts(i) = length(layouts[i]) == 1 ? [xs[i]] : [insert!(ir, v, xcall(GetTupleElement(j-1), xs[i])) for j = 1:length(layouts[i])]
    ir[v] = xcall(XTuple(), vcat([parts(i) for i = 1:length(xs)]...)...)
  end
end

function xlaop!(ir, v, ::AType{typeof(getindex)}, T::AType{<:Tuple}, i::Const{<:Integer})
  xs = ir[v].expr.args[2]
  is = subrange(T, i.value)
  tuplerange!(ir, v, xs, is)
end

function xlaop!(ir, v, ::AType{typeof(getfield)}, T::AType, i::Const{Symbol})
  xs = ir[v].expr.args[2]
  is = subrange(T, fieldnum(widen(T), i.value))
  tuplerange!(ir, v, xs, is)
end

function xlaop!(ir, v, ::AType{typeof(Mjolnir.__new__)}, ::AType{Type{T}}, xs...) where T
  args = ir[v].expr.args[3:end]
  tuplecat!(ir, v, args)
end

function xlaop!(ir, v, ::AType{typeof(Mjolnir.__splatnew__)}, ::AType{Type{T}}, xs) where T
  ir[v] = ir[v].expr.args[3]
end

function xlaop!(ir, v, ::AType{typeof(tuple)}, xs...)
  args = ir[v].expr.args[2:end]
  tuplecat!(ir, v, args)
end

function xlaop!(ir, v, ::AType{typeof(println)}, xs...)
  tuplecat!(ir, v, ir[v].expr.args[2:end])
end

function expand(ir, x, old, new)
  if x isa Variable
    x = insertafter!(ir, x, xcall(Reshape(1:length(old), ntuple(i -> i > length(old) ? 1 : old[i], length(new))), x))
    x = insertafter!(ir, x, xcall(BroadcastInDim(new, 1:length(new)), x))
  end
  return x
end

function xlaop!(ir, v, ::AType{typeof(broadcast)}, f, As...)
  args = ir[v].expr.args
  f = ir[args[2]].expr
  strip_self_arg!(f)
  xs = args[3:end]
  # TODO fold this into `expand`, do a generic type promotion
  for i = 1:length(xs)
    if As[i] isa Const{<:Number}
      xs[i] = fill(convert(eltype(As[1]), xs[i]), size(As[1]))
    end
  end
  new = Broadcast.broadcast_shape(size.(As)...)
  xs = expand.((ir,), xs, size.(As), (new,))
  ir[v] = Expr(:call, Map(), args[2], xs...)
end

function xlaop!(ir, v, ::AType{typeof(mapreduce)}, ::AType{typeof(identity)}, op, xs; dims = :)
  args = ir[v].expr.args
  args[1] isa KwFunc && (popfirst!(args); popfirst!(args))
  f = ir[args[3]].expr
  strip_self_arg!(f)
  if dims == (:)
    ir[v] = Expr(:call, Reduce(dims), args[3], args[4], zero(eltype(widen(xs))))
  else
    dims = filter(d -> d <= ndims(xs), collect(dims))
    out = insert!(ir, v, Expr(:call, Reduce(dims), args[3], args[4], zero(eltype(widen(xs)))))
    sz = size(ir[v].type)
    ir[v] = Expr(:call,
                 Reshape(1:length(size(xs))-length(dims),
                         ntuple(i -> i in dims ? 1 : size(xs)[i], length(sz))),
                 out)
  end
end

function xla_identity(T)
  ir = IR()
  return!(ir, argument!(ir, T))
  return ir
end

function xlaop!(ir, v, ::AType{typeof(ifelse)}, cond, T, F)
  _, cond, t, f = ir[v].expr.args
  tfunc = insert!(ir, v, xla_identity(T))
  ffunc = insert!(ir, v, xla_identity(F))
  ir[v] = xcall(Conditional(), cond, t, tfunc, f, ffunc)
end

xlaop!(ir, v, ::AType{<:KwFunc}, kw::Const, f, args...) =
  xlaop!(ir, v, f, args...; kw.value...)

function xlaop!(ir, v, args...; kw...)
  ir[v] = xlaop(ir[v].expr.args, args...; kw...)
end

function xlaops!(ir)
  for (v, st) in ir
    if st.expr isa Expr
      xlaop!(ir, v, exprtype.((ir,), st.expr.args)...)
    elseif st.expr isa IR
      ir[v] = xlaops!(st.expr)
    end
  end
  # TODO: more fine grained way of ensuring the right return value.
  t = push!(ir, xcall(XTuple(), returnvalue(blocks(ir)[end])))
  push!(ir, xcall(GetTupleElement(0), t))
  return ir
end

function flattenargs!(ir, T)
  arglayout = (layout(T)...,)
  args = copy(arguments(ir))
  deletearg!(ir, 1:length(args))
  components = [argument!(ir, arglayout[i]) for i = 1:length(arglayout)]
  env = Dict()
  for i = 1:length(args)
    cs = components[subrange(T,i)]
    env[args[i]] = length(cs) == 1 ? cs[1] : pushfirst!(ir, xcall(XTuple(), cs...))
  end
  prewalk!(x -> get(env, x, x), ir)
  for (v, st) in ir
    ir[v] = stmt(st, type = Any)
  end
  return ir
end

function convert_xla!(ir, T)
  xlaops!(ir)
  flattenargs!(ir, T)
  return ir
end


Materials = (
    tissueplastic = Material(
        "A-150 Tissue-Equivalent Plastic",
        0.54903,
        65.1eV,
        1.127g / cm^3,
        Dict(
            1 => 0.10133,
            6 => 0.775498,
            7 => 0.035057,
            8 => 0.052315,
            9 => 0.017423,
            20 => 0.018377,
        ),
    ),
    adipose = Material(
        "Adipose Tissue (ICRU-44)",
        0.55579,
        64.8eV,
        0.95g / cm^3,
        Dict(
            1 => 0.114,
            6 => 0.598,
            7 => 0.007,
            8 => 0.278,
            11 => 0.001,
            16 => 0.001,
            17 => 0.001,
        ),
    ),
    air = Material(
        "Air, Dry (near sea level)",
        0.49919,
        85.7eV,
        0.001205g / cm^3,
        Dict(6 => 0.000124, 7 => 0.755268, 8 => 0.231781, 18 => 0.012827),
    ),
    alanine = Material(
        "Alanine",
        0.53876,
        71.9eV,
        1.424g / cm^3,
        Dict(1 => 0.079192, 6 => 0.404437, 7 => 0.157213, 8 => 0.359157),
    ),
    boneplastic = Material(
        "B-100 Bone-Equivalent Plastic",
        0.5274,
        85.9eV,
        1.45g / cm^3,
        Dict(
            1 => 0.065473,
            6 => 0.536942,
            7 => 0.0215,
            8 => 0.032084,
            9 => 0.167415,
            20 => 0.176585,
        ),
    ),
    bakelite = Material(
        "Bakelite",
        0.52792,
        72.4eV,
        1.25g / cm^3,
        Dict(1 => 0.057444, 6 => 0.774589, 8 => 0.167968),
    ),
    wholeblood = Material(
        "Blood, Whole (ICRU-44)",
        0.54999,
        75.2eV,
        1.06g / cm^3,
        Dict(
            1 => 0.102,
            6 => 0.11,
            7 => 0.033,
            8 => 0.745,
            11 => 0.001,
            15 => 0.001,
            16 => 0.002,
            17 => 0.003,
            19 => 0.002,
            26 => 0.001,
        ),
    ),
    corticalbone = Material(
        "Bone, Cortical (ICRU-44)",
        0.51478,
        112.0eV,
        1.92g / cm^3,
        Dict(
            1 => 0.034,
            6 => 0.155,
            7 => 0.042,
            8 => 0.435,
            11 => 0.001,
            12 => 0.002,
            15 => 0.103,
            16 => 0.003,
            20 => 0.225,
        ),
    ),
    brain = Material(
        "Brain, Grey/White Matter (ICRU-44)",
        0.55239,
        73.9eV,
        1.04g / cm^3,
        Dict(
            1 => 0.107,
            6 => 0.145,
            7 => 0.022,
            8 => 0.712,
            11 => 0.002,
            15 => 0.004,
            16 => 0.002,
            17 => 0.003,
            19 => 0.003,
        ),
    ),
    breasttissue = Material(
        "Breast Tissue (ICRU-44)",
        0.55196,
        70.3eV,
        1.02g / cm^3,
        Dict(
            1 => 0.106,
            6 => 0.332,
            7 => 0.03,
            8 => 0.527,
            11 => 0.001,
            15 => 0.001,
            16 => 0.002,
            17 => 0.001,
        ),
    ),
    airplastic = Material(
        "C-552 Air-equivalent Plastic",
        0.49969,
        86.8eV,
        1.76g / cm^3,
        Dict(
            1 => 0.024681,
            6 => 0.50161,
            8 => 0.004527,
            9 => 0.465209,
            14 => 0.003973,
        ),
    ),
    CdTe = Material(
        "Cadmium Telluride",
        0.41665,
        539.3eV,
        6.2g / cm^3,
        Dict(48 => 0.468358, 52 => 0.531642),
    ),
    CaF = Material(
        "Calcium Fluoride",
        0.48671,
        166.0eV,
        3.18g / cm^3,
        Dict(9 => 0.486672, 20 => 0.513328),
    ),
    CaSO4 = Material(
        "Calcium Sulfate",
        0.49948,
        152.3eV,
        2.96g / cm^3,
        Dict(8 => 0.470081, 16 => 0.235534, 20 => 0.294385),
    ),
    ammoniumsulfate = Material(
        "15 mmol L-1 Ceric Ammonium Sulfate Solution",
        0.55282,
        76.7eV,
        1.030E+00g / cm^3,
        Dict(
            1 => 0.107694,
            7 => 0.000816,
            8 => 0.875172,
            16 => 0.01427958,
            58 => 0.00204,
        ),
    ),
    CsI = Material(
        "Cesium Iodide",
        0.41569,
        553.1eV,
        4.51g / cm^3,
        Dict(53 => 0.488451, 55 => 0.511549),
    ),
    concrete = Material(
        "Concrete, Ordinary",
        0.50932,
        124.5eV,
        2.3g / cm^3,
        Dict(
            1 => 0.0221,
            6 => 0.002484,
            8 => 0.57493,
            11 => 0.015208,
            12 => 0.001266,
            13 => 0.019953,
            14 => 0.304627,
            19 => 0.010045,
            20 => 0.042951,
            26 => 0.006435,
        ),
    ),
    concretebarite = Material(
        "Concrete, Barite (TYPE BA)",
        0.45714,
        248.2eV,
        3.35g / cm^3,
        Dict(
            1 => 0.003585,
            8 => 0.311622,
            12 => 0.001195,
            13 => 0.004183,
            14 => 0.010457,
            16 => 0.107858,
            20 => 0.050194,
            26 => 0.047505,
            56 => 0.4634,
        ),
    ),
    eyelens = Material(
        "Eye Lens (ICRU-44)",
        0.54709,
        74.3eV,
        1.07g / cm^3,
        Dict(
            1 => 0.096,
            6 => 0.195,
            7 => 0.057,
            8 => 0.646,
            11 => 0.001,
            15 => 0.001,
            16 => 0.003,
            17 => 0.001,
        ),
    ),
    ferroussulfate = Material(
        "Ferrous Sulfate Standard Fricke",
        0.55334,
        76.3eV,
        1.024g / cm^3,
        Dict(
            1 => 0.108376,
            8 => 0.878959,
            11 => 2.2e-05,
            16 => 0.012553,
            17 => 3.5e-05,
            26 => 5.5e-05,
        ),
    ),
    GOS = Material(
        "Gadolinium Oxysulfide",
        0.42265,
        493.3eV,
        7.44g / cm^3,
        Dict(8 => 0.084527, 16 => 0.084704, 64 => 0.830769),
    ),
    gafchromicsensor = Material(
        "Gafchromic Sensor",
        0.54384,
        67.2eV,
        1.3g / cm^3,
        Dict(1 => 0.0897, 6 => 0.6058, 7 => 0.1122, 8 => 0.1923),
    ),
    GaAs = Material(
        "Gallium Arsenide",
        0.44246,
        384.9eV,
        5.31g / cm^3,
        Dict(31 => 0.48203, 33 => 0.51797),
    ),
    pyrex = Material(
        "Glass, Borosilicate (Pyrex)",
        0.49707,
        134.0eV,
        2.23g / cm^3,
        Dict(
            5 => 0.040066,
            8 => 0.539559,
            11 => 0.028191,
            13 => 0.011644,
            14 => 0.37722,
            19 => 0.003321,
        ),
    ),
    leadglass = Material(
        "Glass, Lead",
        0.42101,
        526.4eV,
        6.22g / cm^3,
        Dict(
            8 => 0.156453,
            14 => 0.080866,
            22 => 0.008092,
            33 => 0.002651,
            82 => 0.751938,
        ),
    ),
    LiF = Material(
        "Lithium Fluride",
        0.46262,
        94.0eV,
        2.635g / cm^3,
        Dict(3 => 0.267585, 9 => 0.732415),
    ),
    LiB4O7 = Material(
        "Lithium Tetraborate",
        0.48485,
        94.6eV,
        2.44g / cm^3,
        Dict(3 => 0.082081, 5 => 0.255715, 8 => 0.662204),
    ),
    lung = Material(
        "Lung Tissue (ICRU-44)",
        0.55048,
        75.2eV,
        1.05g / cm^3,
        Dict(
            1 => 0.103,
            6 => 0.105,
            7 => 0.031,
            8 => 0.749,
            11 => 0.002,
            15 => 0.002,
            16 => 0.003,
            17 => 0.003,
            19 => 0.002,
        ),
    ),
    MgB4O7 = Material(
        "Magnesium Tetroborate",
        0.49012,
        108.3eV,
        2.53g / cm^3,
        Dict(5 => 0.24087, 8 => 0.623762, 12 => 0.135367),
    ),
    MgI2 = Material(
        "Mercuric Iodide",
        0.40933,
        684.5eV,
        6.36g / cm^3,
        Dict(53 => 0.55856, 80 => 0.44144),
    ),
    muscle = Material(
        "Muscle, Skeletal (ICRU-44)",
        0.55,
        74.6eV,
        1.05g / cm^3,
        Dict(
            1 => 0.102,
            6 => 0.143,
            7 => 0.034,
            8 => 0.71,
            11 => 0.001,
            15 => 0.002,
            16 => 0.003,
            17 => 0.001,
            19 => 0.004,
        ),
    ),
    ovary = Material(
        "Ovary (ICRU-44)",
        0.55149,
        75.0eV,
        1.05g / cm^3,
        Dict(
            1 => 0.105,
            6 => 0.093,
            7 => 0.024,
            8 => 0.768,
            11 => 0.002,
            15 => 0.002,
            16 => 0.002,
            17 => 0.002,
            19 => 0.002,
        ),
    ),
    emulsionkodak = Material(
        "Photographic Emulsion (Kodak Type AA)",
        0.48176,
        179.0eV,
        2.2g / cm^3,
        Dict(
            1 => 0.0305,
            6 => 0.2107,
            7 => 0.0721,
            8 => 0.1632,
            35 => 0.2228,
            47 => 0.3007,
        ),
    ),
    emulsionstandard = Material(
        "Photographic Emulsion (Standard Nuclear)",
        0.45453,
        331.0eV,
        3.815g / cm^3,
        Dict(
            1 => 0.0141,
            6 => 0.072261,
            7 => 0.01932,
            8 => 0.066101,
            16 => 0.00189,
            35 => 0.349104,
            47 => 0.474105,
            53 => 0.00312,
        ),
    ),
    vinyltoluene = Material(
        "Plastic Scintillator, Vinyltoluene",
        0.54141,
        64.7eV,
        1.032g / cm^3,
        Dict(1 => 0.085, 6 => 0.915),
    ),
    polyethylene = Material(
        "Polyethylene",
        0.57033,
        57.4eV,
        0.93g / cm^3,
        Dict(1 => 0.143716, 6 => 0.856284),
    ),
    mylar = Material(
        "Polyethylene Terephthalate, (Mylar)",
        0.52037,
        78.7eV,
        1.38g / cm^3,
        Dict(1 => 0.04196, 6 => 0.625016, 8 => 0.333024),
    ),
    PMMA = Material(
        "Polymethyl Methacrylate",
        0.53937,
        74.0eV,
        1.19g / cm^3,
        Dict(1 => 0.080541, 6 => 0.599846, 8 => 0.319613),
    ),
    polystyrene = Material(
        "Polystyrene",
        0.53768,
        68.7eV,
        1.06g / cm^3,
        Dict(1 => 0.077421, 6 => 0.922579),
    ),
    teflon = Material(
        "Polytetrafluoroethylene, (Teflon)",
        0.47993,
        99.1eV,
        2.25g / cm^3,
        Dict(6 => 0.240183, 9 => 0.759818),
    ),
    PVC = Material(
        "Polyvinyl Chloride",
        0.51201,
        108.2eV,
        1.406g / cm^3,
        Dict(1 => 0.048382, 6 => 0.384361, 17 => 0.567257),
    ),
    nylonfilm = Material(
        "Radiochromic Dye Film, Nylon Base",
        0.54987,
        64.5eV,
        1.08g / cm^3,
        Dict(1 => 0.101996, 6 => 0.654396, 7 => 0.098915, 8 => 0.144693),
    ),
    testis = Material(
        "Testis (ICRU-44)",
        0.552,
        74.7eV,
        1.04g / cm^3,
        Dict(
            1 => 0.106,
            6 => 0.099,
            7 => 0.02,
            8 => 0.766,
            11 => 0.002,
            15 => 0.001,
            16 => 0.002,
            17 => 0.002,
            19 => 0.002,
        ),
    ),
    softtissue = Material(
        "Tissue, Soft (ICRU-44)",
        0.54996,
        74.7eV,
        1.06g / cm^3,
        Dict(
            1 => 0.102,
            6 => 0.143,
            7 => 0.034,
            8 => 0.708,
            11 => 0.002,
            15 => 0.003,
            16 => 0.003,
            17 => 0.002,
            19 => 0.003,
        ),
    ),
    softtissue4 = Material(
        "Tissue, Soft (ICRU Four-Component)",
        0.54975,
        74.9eV,
        1.0g / cm^3,
        Dict(1 => 0.101174, 6 => 0.111, 7 => 0.026, 8 => 0.761826),
    ),
    methane = Material(
        "Tissue-Equivalent Gas, Methane Based",
        0.54992,
        61.2eV,
        0.001064g / cm^3,
        Dict(1 => 0.101873, 6 => 0.456177, 7 => 0.035172, 8 => 0.406778),
    ),
    propane = Material(
        "Tissue-Equivalent Gas, Propane Based",
        0.55027,
        59.5eV,
        0.001826g / cm^3,
        Dict(1 => 0.102676, 6 => 0.568937, 7 => 0.035022, 8 => 0.293365),
    ),
    water = Material(
        "Water, Liquid",
        0.55508,
        75.0eV,
        1.0g / cm^3,
        Dict(1 => 0.111898, 8 => 0.888102),
    ),
)


module ReduceNoise

using Base.Iterators
using ImageCore
using ImageCore.MappedArrays
using ImageCore.OffsetArrays
using ImageCore.PaddedViews
using ImageCore: NumberLike, GenericGrayImage, GenericImage
import ..NoiseAPI: AbstractImageDenoiseAlgorithm, reduce_noise, reduce_noise!

include("compat.jl")
include("BM3DDenoise.jl")
include("NonlocalMean.jl")

export
    reduce_noise, reduce_noise!,

    # BM3D
    BM3D,
    # Non-local mean filter for gaussian noise
    NonlocalMean, get_NonlocalMean_rp

end # module


abstract type Wishart3x3{T} <: AbstractMatrix{T} end

struct InverseWishart{T} <: Wishart3x3{T}
    data::NTuple{8,VE{T}}
end
struct CholInvWishart{T} <: Wishart3x3{T}
    data::NTuple{8,T}
end
struct RevCholWishart{T} <: Wishart3x3{T}
    data::NTuple{8,T}
end

const WishartFactor{T} = Union{CholInvWishart{T},RevCholWishart{T}}

@inline InverseWishart(x::Vararg{T,8}   ) where T = InverseWishart{T}(VE.(x))
@inline InverseWishart{T}(x::Vararg{T,8}) where T = InverseWishart{T}(VE.(x))
# @generated function CholInvWishart(x::Vararg{T,8}) where T
#     quote
#         $(Expr(:meta,:inline))
#         CholInvWishart{$T}($(Expr(:tuple,[:(@inbounds x[$i]) for i ∈ 1:8]...)))
#     end
# end

@inline function LinearAlgebra.logdet(l::Union{CholInvWishart{T},RevCholWishart{T}}) where {T}
    @inbounds SLEEFPirates.log(l[1] * l[4] * l[6])
end
@inline function logdetw(l::Union{CholInvWishart{T},RevCholWishart{T}}) where {T}
    @inbounds SLEEFwrap.log(l[1] * l[4] * l[6])
end

@inline function CholInvWishart{T}(x_1::T,x_2::T,x_3::T,x_4::T,x_5::T,x_6::T) where T
    x_7 = x_8 = one(T)
    CholInvWishart{T}((Base.Cartesian.@ntuple 8 x))
end
@inline function CholInvWishart{T}(x_1::T,x_2::T,x_3::T,x_4::T,x_5::T,x_6::T,x_7::T,x_8::T) where T
    CholInvWishart{T}((Base.Cartesian.@ntuple 8 x))
end
@inline function CholInvWishart(x_1::T,x_2::T,x_3::T,x_4::T,x_5::T,x_6::T) where T
    x_7 = x_8 = one(T)
    CholInvWishart{T}((Base.Cartesian.@ntuple 8 x))
end
@inline function CholInvWishart(x_1::T,x_2::T,x_3::T,x_4::T,x_5::T,x_6::T,x_7::T,x_8::T) where T
    CholInvWishart{T}((Base.Cartesian.@ntuple 8 x))
end
@inline function RevCholWishart{T}(x_1::T,x_2::T,x_3::T,x_4::T,x_5::T,x_6::T) where T
    x_7 = x_8 = one(T)
    RevCholWishart{T}((Base.Cartesian.@ntuple 8 x))
end
@inline function RevCholWishart{T}(x_1::T,x_2::T,x_3::T,x_4::T,x_5::T,x_6::T,x_7::T,x_8::T) where T
    RevCholWishart{T}((Base.Cartesian.@ntuple 8 x))
end
@inline function RevCholWishart{T1}(
            x_1::SVec{W,T},x_2::SVec{W,T},x_3::SVec{W,T},x_4::SVec{W,T},x_5::SVec{W,T},x_6::SVec{W,T},x_7::SVec{W,T},x_8::SVec{W,T}
        ) where {W,T,T1}
    RevCholWishart{SVec{W,T}}((Base.Cartesian.@ntuple 8 x))
end
@inline function RevCholWishart(x_1::T,x_2::T,x_3::T,x_4::T,x_5::T,x_6::T) where T
    x_7 = x_8 = one(T)
    RevCholWishart{T}((Base.Cartesian.@ntuple 8 x))
end
@inline function RevCholWishart(x_1::T,x_2::T,x_3::T,x_4::T,x_5::T,x_6::T,x_7::T,x_8::T) where T
    RevCholWishart{T}((Base.Cartesian.@ntuple 8 x))
end
# @inline RevCholWishart{T}(x::Vararg{T,8}) where T = RevCholWishart{T}(x)
# @inline RevCholWishart(x::Vararg{T,6}   ) where T = RevCholWishart{T}(x,T(1),T(1))
# @inline RevCholWishart{T}(x::Vararg{T,6}) where T = RevCholWishart{T}(x,T(1),T(1))
@inline CholInvWishart{T}(x::Vararg{SVec{W,T},8}) where {W,T} = CholInvWishart{SVec{W,T}}(x)

inv_type(::Type{CholInvWishart{T}}) where T = RevCholWishart{T}
inv_type(::Type{RevCholWishart{T}}) where T = CholInvWishart{T}
# @inline function SArray{Tuple{S},T,N,L}(vs::Vararg{SIMDPirates.SVec{W,T},L}) where {S,T,N,L,W}
#     SArray{S}(vs...)
# end
# @inline Base.size(::Wishart2x2) = (2,2)
# @inline Base.length(::Wishart2x2) = 4
@inline Base.size(::Wishart3x3) = (3,3)
@inline Base.length(::Wishart3x3) = 9
# ScatteredArrays.type_length(::Type{Wishart2x2}) = 3
# ScatteredArrays.type_length(::Type{<:Wishart3x3}) = 6
# ScatteredArrays.type_length(::Type{CholInvWishart{T}}) where T = 6

@inline extractval(x::Core.VecElement{T}) where T = x.value
@inline extractval(x) = x
@inline Base.getindex(w::Wishart3x3, i::Integer) = extractval(@inbounds w.data[i])

@inline Base.getindex(A::Wishart3x3, ::LinearStorage, i::Integer) = @inbounds A.data[i]

@inline function Base.getindex(w::Wishart3x3, i::CartesianIndex{2})
    w[i[1], i[2]]
end
@inline function Base.getindex(iw::InverseWishart, i::Integer, j::Integer)
    i, j = minmax(i, j)
    @boundscheck (j > 3 || i < 1) && throwboundserror()
    im1 = i - 1
    @inbounds iw.data[3im1 - ( (im1*i) >> 1) + j].value
end
@inline function Base.getindex(w::WishartFactor{T}, j::Integer, i::Integer) where T
    j < i && return zero(T)
    @boundscheck (j > 3 || i < 1) && throwboundserror()
    im1 = i - 1
    @inbounds extractval(w.data[3im1 - ( (im1*i) >> 1) + j])
end
function Base.show(io::IO, ::MIME"text/plain", x::InverseWishart{T}) where T
    denom = x[7] > 2 ? x[7]-2 : one(T)
    println(io, "InverseWishart{$T} with ν = $(x[7]+2)")
    Base.print_matrix(io, x / denom)
end
function Base.show(io::IO, ::MIME"text/plain", x::CholInvWishart{T}) where T
    denom = x[7] > 2 ? sqrt(x[7]-2) : one(T)
    println(io, "Cholesky factor of InverseWishart{$T} with ν = $(x[7]+2)")
    Base.print_matrix(io, x / denom)
end

@inline function InverseWishart(x::T, y::T, z::T) where T
    InverseWishart(x * x, x * y, x * z, y * y, y * z, z * z, one(T), one(T))
end

@inline function Base.:+(a::W, b::W) where W <: Wishart3x3
    W(SIMDPirates.vadd(a.data, b.data))
end
@inline function Base.:*(L::WishartFactor{T}, x::SVector{3,T}) where {T}
    @fastmath @inbounds SVector(
        L[1]*x[1],
        L[2]*x[1] + L[4]*x[2],
        L[3]*x[1] + L[5]*x[2] + L[6]*x[3]
    )
end
@inline function Base.:*(L::WishartFactor{SVec{W,T}}, x::SVector{3,T}) where {W,T}
    @fastmath @inbounds SVector(
        L[1]*x[1],
        L[2]*x[1] + L[4]*x[2],
        L[3]*x[1] + L[5]*x[2] + L[6]*x[3]
    )
end
# L2 is treated as a lower triangular matrix
@inline function Base.:*(L1::W, L2::Union{W,NTuple{6,T}}) where {T, W <: WishartFactor{T}}
    @fastmath @inbounds W(
        L1[1]*L2[1],
        L1[2]*L2[1] + L1[4]*L2[2],
        L1[3]*L2[1] + L1[5]*L2[2] + L1[6]*L2[3],
        L1[4]*L2[4],
        L1[5]*L2[4] + L1[6]*L2[5],
        L1[6]*L2[6],
        L1[7],L1[8]
    )
end
@inline function Base.:*(L1::NTuple{6,T}, L2::W) where {T, W <: WishartFactor{T}}
    @fastmath @inbounds W(
        L1[1]*L2[1],
        L1[2]*L2[1] + L1[4]*L2[2],
        L1[3]*L2[1] + L1[5]*L2[2] + L1[6]*L2[3],
        L1[4]*L2[4],
        L1[5]*L2[4] + L1[6]*L2[5],
        L1[6]*L2[6],
        L2[7],L2[8]
    )
end
@inline function Base.:/(a::W, x::T) where {T <: Number, W <: Wishart3x3{T}}
    @fastmath @inbounds begin
        xinv = 1 / x
        W(
            a[1] * xinv, a[2] * xinv, a[3] * xinv, a[4] * xinv, a[5] * xinv, a[6] * xinv, a[7], a[8]
        )
    end
end

@inline extract_ν(iw::Wishart3x3) = iw[7]
@generated function extract_ν(iw::AbstractVector{<:Wishart3x3{T}}, ::Val{NG}) where {NG,T}
    quote
        $(Expr(:meta, :inline))
        ptr_iw = Base.unsafe_convert(Ptr{$T}, pointer(iw))
        $(Expr(:tuple, [:(unsafe_load(ptr_iw, $((ng-1)*8) + 7 ) ) for ng ∈ 1:NG]...))
    end
end
@inline extract_α(iw::Union{InverseWishart,CholInvWishart}) = iw[8]
@generated function extract_α(iw::AbstractVector{<:Union{InverseWishart{T},CholInvWishart{T}}}, ::Val{NG}) where {NG,T}
    quote
        $(Expr(:meta, :inline))
        ptr_iw = Base.unsafe_convert(Ptr{$T}, pointer(iw))
        $(Expr(:tuple, [:(unsafe_load(ptr_iw, $((ng-1)*8) + 8 ) ) for ng ∈ 1:NG]...))
    end
end

@inline function finite_sqrt(x::T, min_value = eps(T)) where {T}
    rx = @fastmath sqrt(x)
    # min_value = eps(T)
    # min_value = T(0.005) # 5
    isfinite(rx) ? rx : min_value
end

# @inline
function inv_and_cholesky!(riwv::AbstractVector{RevCholWishart{T}},
                                    ciwv::AbstractVector{CholInvWishart{T}},
                                    iwv::AbstractVector{InverseWishart{T}}) where T

    @inbounds for i ∈ eachindex(iwv)
        iw = iwv[i]
        @fastmath begin
            min_value = T(0.05) * sqrt(iw[7] - 2)
            L11 = finite_sqrt(iw[1], min_value)
            R11 = 1 / L11
            L21 = R11 * iw[2]
            L31 = R11 * iw[3]
            L22 = finite_sqrt(iw[4] - L21^2, min_value)
            R22 = 1 / L22
            L32 = R22 * (iw[5] - L21 * L31)
            L33 = finite_sqrt(iw[6] - L31^2 - L32^2, min_value)
            R33 = 1 / L33

            R21 = - R22 * L21 * R11
            R31 = - R33 * ( L31*R11 + L32*R21 )
            R32 = - R33 * L32 * R22
        end
        ciwv[i] = CholInvWishart(
            L11, L21, L31, L22, L32, L33, iw[7], iw[8]
        )
        riwv[i] = RevCholWishart(
            R11, R21, R31, R22, R32, R33, iw[7], iw[8]
        )
    end
end
@inline function triangle_inv(L::NTuple{8,T}) where T
    @inbounds @fastmath begin
        L11           = L[1]
        L21, L22      = L[2], L[4]
        L31, L32, L33 = L[3], L[5], L[6]
        R11 = 1 / L11
        R22 = 1 / L22
        R33 = 1 / L33
        R21 = - R22 * L21 * R11
        R31 = - R33 * ( L31*R11 + L32*R21 )
        R32 = - R33 * L32 * R22
    end
    @inbounds (R11, R21, R31, R22, R32, R33, L[7], L[8])
end
@inline function triangle_inv(L::NTuple{6,T}) where T
    @inbounds @fastmath begin
        L11           = L[1]
        L21, L22      = L[2], L[4]
        L31, L32, L33 = L[3], L[5], L[6]
        R11 = 1 / L11
        R22 = 1 / L22
        R33 = 1 / L33
        R21 = - R22 * L21 * R11
        R31 = - R33 * ( L31*R11 + L32*R21 )
        R32 = - R33 * L32 * R22
    end
    R11, R21, R31, R22, R32, R33
end
@inline function Base.inv(w::W) where {T,W<:WishartFactor{T}}
    inv_type(W)(triangle_inv(w.data))
end

@inline function UpperTriangle3(U::WishartFactor{T}) where T
    @inbounds UpperTriangle3(SVector{6,T}(
        U[1], U[2], U[4], U[3], U[5], U[6]
    ))
end

@inline function randinvwishartfactor(rng::AbstractRNG, rcw::RevCholWishart{T}) where T
    rcwU = UpperTriangle3(rcw)
    ν = extract_ν(rcw)
    U11 = randchi(rng, ν - 2)
    U12 = randn(rng)
    U22 = randchi(rng, ν - 1)
    U13 = randn(rng)
    U23 = randn(rng)
    U33 = randchi(rng, ν)
    U =  UpperTriangle3(
        U11, U12, U22, U13, U23, U33
    )
#    @show rcw
#    @show rcwU
#    @show U
    inv(rcwU * U)
end
function randinvwishart(rng::AbstractRNG, rcw::RevCholWishart{T}) where T
    xtx( randinvwishartfactor(rng, rcw) )
end


# This analysis generates several different common plots used to
# diagnose and explore screen data including: (1) volcano plot of gene
# pvalues and effect sizes (2) precision-recall curve for the hits
# (3) raw count plots
using Gadfly
using ColorBrewer

function gen_plots(filepath; debug=false, quiet=false)
    colors = palette("Accent", 4)

    for screentype in [FacsScreen(), GrowthScreen()]
        for crisprtype in [CRISPRi(), CRISPRn()]
            if typeof(screentype) == GrowthScreen
                max_phenotype_dists = Dict{Symbol, Tuple{Float64, Sampleable}}(
                    :inactive => (0.83, Delta(0.0)),
                    :negcontrol => (0.05, Delta(0.0)),
                    :increasing => (0.02, TruncatedNormal(0.1, 0.1, 0.025, 1)),
                    :decreasing => (0.1, TruncatedNormal(-0.55, 0.2, -1, -0.1))
                )
                lib = Library(max_phenotype_dists, crisprtype)
            else
                lib = Library(crisprtype)
            end

            # redefining function each time and abusing lexical scoping
            gen_plots = (bc_counts, genes) -> begin
                curr_dir = dirname(filepath)
                new_filename = joinpath(curr_dir, "$(basename(filepath))_$(typeof(crisprtype))_$(typeof(screentype))_counts.svg")
                nopseudo = bc_counts[(bc_counts[:counts_bin1] .> 0.5) .& (bc_counts[:counts_bin2] .> 0.5), :]
                draw(SVG(new_filename, 10cm, 10cm), plot(nopseudo,
                x=:counts_bin1, y=:counts_bin2, color=:class, Scale.x_log10,
                Scale.y_log10, Scale.color_discrete_manual(colors..., levels=sort(unique(nopseudo[:class]))),
                Theme(highlight_width=0pt),
                Coord.cartesian(fixed=true),
                Guide.title("$(typeof(crisprtype)) $(typeof(screentype))"),
                Guide.xlabel("log counts t0"), Guide.ylabel("log counts t1")))
                new_filename = joinpath(curr_dir, "$(basename(filepath))_$(typeof(crisprtype))_$(typeof(screentype))_volcano.svg")
                draw(SVG(new_filename, 10cm, 10cm), plot(genes,
                x=:mean_bin2_div_bin1, y=:pvalue_bin2_div_bin1, color=:class,
                Scale.color_discrete_manual(colors..., levels=sort(unique(genes[:class]))),
                Theme(highlight_width=0pt),
                Guide.title("$(typeof(crisprtype)) $(typeof(screentype))"),
                Guide.xlabel("mean log2 fold change"), Guide.ylabel("-log10 pvalue")))
                (1)
            end

            run_exp(screentype, lib, gen_plots; run_idx=1)
        end
    end
end


export DoubleLearner

using Random

"""
    DoubleLearner(;L1, L2, rng=MersenneTwister())

This is a meta-learner, it will randomly select one learner and update another learner.
The estimation of an observation is the sum of result from two learners.
"""
Base.@kwdef struct DoubleLearner{T1<:AbstractLearner,T2<:AbstractLearner,R<:AbstractRNG} <:
                   AbstractLearner
    L1::T1
    L2::T2
    rng::R = MersenneTwister()
end

"""
    DoubleLearner(l1, l2; seed = nothing)
"""
DoubleLearner(l1, l2; seed = nothing) = DoubleLearner(l1, l2, MersenneTwister(seed))

(learner::DoubleLearner)(obs) = learner.L1(obs) .+ learner.L2(obs)

RLBase.extract_experience(t::AbstractTrajectory, learner::DoubleLearner) =
    extract_experience(t, learner.L1)

update!(learner::DoubleLearner, experience) =
    rand(learner.rng, Bool) ? update!(learner.L1, experience) :
    update!(learner.L2, experience)


# This file defines some struct types used in the FrameConnection package.

using DataFrames

"""
    SMD(connectID::Vector{Int}, x::Vector{Float32}, y::Vector{Float32},
        x_se::Vector{Float32}, y_se::Vector{Float32}, framenum::Vector{Int},
        nframes::Int, datasetnum::Vector{Int}, ndatasets::Int)

Single Molecule Data structure with information about localizations.
"""
mutable struct SMD
    connectID::Vector{Int}
    x::Vector{Float32}
    y::Vector{Float32}
    x_se::Vector{Float32}
    y_se::Vector{Float32}
    framenum::Vector{Int}
    nframes::Int
    datasetnum::Vector{Int}
    ndatasets::Int
    SMD() = new()
end
function SMD(nlocalizations::Int)
    smd = SMD()
    smd.x = Vector{Float32}(undef, nlocalizations)
    smd.y = Vector{Float32}(undef, nlocalizations)
    smd.x_se = Vector{Float32}(undef, nlocalizations)
    smd.y_se = Vector{Float32}(undef, nlocalizations)
    smd.framenum = Vector{Int}(undef, nlocalizations)
    smd.datasetnum = Vector{Int}(undef, nlocalizations)

    return smd
end
function SMD(data::DataFrames.DataFrame)
    smd = SMD()
    smd.datasetnum = Int.(data[:, 1])
    smd.framenum = Int.(data[:, 2])
    smd.x = Float32.(data[:, 3])
    smd.y = Float32.(data[:, 4])
    smd.x_se = Float32.(data[:, 5])
    smd.y_se = Float32.(data[:, 6])
    smd.nframes = Int(maximum(smd.framenum))
    smd.ndatasets = Int(length(unique(smd.datasetnum)))

    return smd
end

"""
    ParamStruct(initialdensity::Vector{Float32},
        k_on::Float32, k_off::Float32, k_bleach::Float32, p_miss::Float32,
        nsigmadev::Float32, maxframegap::Int, nnearestclusters::Int)

Structure of parameters needed for frame-connection.

# Fields
-`initialdensity`: Density of emitters at the start of the experiment.
                   (see estimatedensities()) (emitters/pixel^2)
-`nnearestclusters`: Number of nearest preclusters used for local density
                     estimates. (default = 2)(see estimatedensities())
-`k_on`: Rate at which dark emitters convert to the visible state.
         (see estimateparams())(1/frame)
-`k_off`: Rate at which visible emitters are converted to the reversible dark
          state. (see estimateparams())(1/frame)
-`k_bleach`: Rate at which visible emitters are irreversibly photobleached.
             (see estimateparams())(1/frame)
-`p_miss`: Probability of missing a localization of a visible emitter.
-`nsigmadev`: Multiplier of localization errors that defines a pre-clustering
              distance threshold. (default = 5)(see precluster())(pixels)
-`maxframegap`: Maximum frame gap between temporally adjacent localizations in
                a precluster. (default = 5)(see precluster())(frames)
-`nmaxnn`: Maximum number of nearest-neighbors inspected for precluster
           membership.  Ideally, this would be set to inf, but that's not
           feasible for most data. (default = 2)(see precluster())

"""
mutable struct ParamStruct
    initialdensity::Vector{Float32}
    nnearestclusters::Int
    k_on::Float32
    k_off::Float32
    k_bleach::Float32
    p_miss::Float32
    nsigmadev::Float32
    maxframegap::Int
    nmaxnn::Int
end
ParamStruct() = ParamStruct([], 2, 0, 0, 0, 0, 5, 5, 2)


# MIT License
#
# Copyright (c) 2018 Martin Biel
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""
    Optimizer(; <keyword arguments>)

Return a progressive-hedging optimizer.

...
# Arguments
- `subproblem_optimizer::AbstractOptimizer`: MathOptInterface solver capable of solving quadratic programs.
- `penalty::AbstractPenalizer = Fixed()`: Specify penalty update procedure (Fixed, Adaptive)
- `execution::AbstractExecuter = Serial`: Specify how algorithm should be executed (Serial, Synchronous, Asynchronous). Distributed variants requires worker cores.
- `penaltyterm::PenaltyTerm = Quadratic`: Specify penaltyterm variant ([`Quadratic`](@ref), [`Linearized`](@ref), [`InfNorm`](@ref), [`ManhattanNorm`][@ref])
- <keyword arguments>: Algorithm specific parameters, consult individual docstrings (see above list) for list of possible arguments and default values.
...
"""
mutable struct Optimizer <: AbstractStructuredOptimizer
    subproblem_optimizer
    sub_params::Dict{MOI.AbstractOptimizerAttribute, Any}
    execution::AbstractExecution
    penalizer::AbstractPenalizer
    penaltyterm::AbstractPenaltyTerm
    parameters::ProgressiveHedgingParameters{Float64}

    status::MOI.TerminationStatusCode
    primal_status::MOI.ResultStatusCode
    dual_status::MOI.ResultStatusCode
    raw_status::String
    solve_time::Float64
    progressivehedging::Union{AbstractProgressiveHedging, Nothing}

    function Optimizer(; subproblem_optimizer = nothing,
                       execution::AbstractExecution = nworkers() == 1 ? Serial() : Synchronous(),
                       penalty::AbstractPenalizer = Fixed(),
                       penaltyterm::AbstractPenaltyTerm = Quadratic(),
                       kw...)
        return new(subproblem_optimizer,
                   Dict{MOI.AbstractOptimizerAttribute, Any}(),
                   execution,
                   penalty,
                   penaltyterm,
                   ProgressiveHedgingParameters{Float64}(; kw...),
                   MOI.OPTIMIZE_NOT_CALLED,
                   MOI.NO_SOLUTION,
                   MOI.NO_SOLUTION,
                   "Progressive-hedging optimizer has not been run.",
                   NaN,
                   nothing)
    end
end

# Interface #
# ========================== #
function supports_structure(optimizer::Optimizer, ::ScenarioDecompositionStructure)
    return true
end

function default_structure(::UnspecifiedInstantiation, optimizer::Optimizer)
    if optimizer.execution isa Serial && nworkers() == 1
        return ScenarioDecomposition()
    else
        return DistributedScenarioDecomposition()
    end
end

function check_loadable(optimizer::Optimizer, ::ScenarioDecompositionStructure)
    if optimizer.subproblem_optimizer === nothing
        msg = "Subproblem optimizer not set. Consider setting `SubProblemOptimizer` attribute."
        throw(UnloadableStructure{Optimizer, ScenarioDecompositionStructure}(msg))
    end
    return nothing
end

function ensure_compatible_execution!(optimizer::Optimizer, ::ScenarioDecompositionStructure{2, 1, <:Tuple{ScenarioProblems}})
    if !(optimizer.execution isa Serial)
        @warn "Distributed execution policies are not compatible with a single-core scenario-decomposition structure. Switching to `Serial` execution by default."
        MOI.set(optimizer, Execution(), Serial())
    end
    return nothing
end

function ensure_compatible_execution!(optimizer::Optimizer, ::ScenarioDecompositionStructure{2, 1, <:Tuple{DistributedScenarioProblems}})
    if optimizer.execution isa Serial
        @warn "Serial execution not compatible with distributed scenario-decomposition structure. Switching to `Synchronous` execution by default."
        MOI.set(optimizer, Execution(), Synchronous())
    end
    return nothing
end

function load_structure!(optimizer::Optimizer, structure::ScenarioDecompositionStructure, x₀::AbstractVector)
    # Sanity check
    check_loadable(optimizer, structure)
    # Restore structure if optimization has been run before
    restore_structure!(optimizer)
    # Ensure that execution policy is compatible
    ensure_compatible_execution!(optimizer, structure)
    # Set subproblem optimizers
    set_subproblem_optimizer!(structure, optimizer.subproblem_optimizer)
    # Set subproblem optimizer attributes
    for (attr, value) in optimizer.sub_params
        MOI.set(scenarioproblems(structure), attr, value)
    end
    # Create new progressive-hedging algorithm
    optimizer.progressivehedging = ProgressiveHedgingAlgorithm(structure,
                                                               x₀,
                                                               optimizer.execution,
                                                               optimizer.penalizer,
                                                               optimizer.penaltyterm;
                                                               type2dict(optimizer.parameters)...)
    return nothing
end

function restore_structure!(optimizer::Optimizer)
    if optimizer.progressivehedging !== nothing
        restore_subproblems!(optimizer.progressivehedging)
    end
    return nothing
end

function MOI.optimize!(optimizer::Optimizer)
    if optimizer.progressivehedging === nothing
        throw(StochasticProgram.UnloadedStructure{Optimizer}())
    end
    # Run progressive-hedging procedure
    optimizer.status = optimizer.progressivehedging()
    # Check if optimal
    if optimizer.status == MOI.OPTIMAL
        optimizer.primal_status = MOI.FEASIBLE_POINT
        optimizer.dual_status = MOI.FEASIBLE_POINT
        optimizer.raw_status = "Progressive-hedging procedure converged to optimal solution."
    end
    # Extract solve time
    optimizer.solve_time = optimizer.progressivehedging.progress.tlast - optimizer.progressivehedging.progress.tinit
    return nothing
end

function num_iterations(optimizer::Optimizer)
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return num_iterations(optimizer.progressivehedging)
end

function optimizer_name(optimizer::Optimizer)
    return "$(str(optimizer.execution))Progressive-hedging with $(str(optimizer.penalizer))"
end

# MOI #
# ========================== #
function MOI.get(optimizer::Optimizer, ::MOI.Silent)
    return !MOI.get(optimizer, MOI.RawParameter("log"))
end

function MOI.set(optimizer::Optimizer, attr::MOI.Silent, flag::Bool)
    MOI.set(optimizer, MOI.RawParameter("log"), !flag)
    optimizer.sub_params[attr] = flag
    return nothing
end

function MOI.get(optimizer::Optimizer, param::MOI.RawParameter)
    name = Symbol(param.name)
    if !(name in fieldnames(ProgressiveHedgingParameters))
        error("Unrecognized parameter name: $(name).")
    end
    return getfield(optimizer.parameters, name)
end

function MOI.set(optimizer::Optimizer, param::MOI.RawParameter, value)
    name = Symbol(param.name)
    if !(name in fieldnames(ProgressiveHedgingParameters))
        error("Unrecognized parameter name: $(name).")
    end
    setfield!(optimizer.parameters, name, value)
    return nothing
end

function MOI.get(optimizer::Optimizer, ::MOI.TimeLimitSec)
    limit = MOI.get(optimizer, MOI.RawParameter("time_limit"))
    return isinf(limit) ? nothing : limit
end

function MOI.set(optimizer::Optimizer, ::MOI.TimeLimitSec, limit::Union{Real, Nothing})
    limit = limit === nothing ? Inf : limit
    MOI.set(optimizer, MOI.RawParameter("time_limit"), limit)
    return
end

function MOI.get(optimizer::Optimizer, ::PrimalTolerance)
    return MOI.get(optimizer, MOI.RawParameter("ϵ₁"))
end

function MOI.set(optimizer::Optimizer, ::PrimalTolerance, limit::Real)
    MOI.set(optimizer, MOI.RawParameter("ϵ₁"), limit)
    return nothing
end

function MOI.get(optimizer::Optimizer, ::DualTolerance)
    return MOI.get(optimizer, MOI.RawParameter("ϵ₂"))
end

function MOI.set(optimizer::Optimizer, ::DualTolerance, limit::Real)
    MOI.set(optimizer, MOI.RawParameter("ϵ₂"), limit)
    return nothing
end

function MOI.get(optimizer::Optimizer, ::MasterOptimizer)
    return MOI.get(optimizer, SubProblemOptimizer())
end

function MOI.get(optimizer::Optimizer, ::SubProblemOptimizer)
    if optimizer.subproblem_optimizer === nothing
        return MOI.get(optimizer, MasterOptimizer())
    end
    return MOI.OptimizerWithAttributes(optimizer.subproblem_optimizer, collect(optimizer.sub_params))
end

function MOI.set(optimizer::Optimizer, ::SubProblemOptimizer, optimizer_constructor)
    optimizer.subproblem_optimizer = optimizer_constructor
    # Clear any old parameters
    empty!(optimizer.sub_params)
    return nothing
end

function MOI.get(optimizer::Optimizer, ::SubProblemOptimizerAttribute, attr::MOI.AbstractOptimizerAttribute)
    if !haskey(optimizer.sub_params, attr)
        error("Subproblem optimizer attribute $(attr) has not been set.")
    end
    return optimizer.sub_params[attr]
end

function MOI.set(optimizer::Optimizer, ::SubProblemOptimizerAttribute, attr::MOI.AbstractOptimizerAttribute, value)
    optimizer.sub_params[attr] = value
    return nothing
end

function MOI.get(optimizer::Optimizer, param::RawSubProblemOptimizerParameter)
    moi_param = MOI.RawParameter(param.name)
    if !haskey(optimizer.sub_params, moi_param)
        error("Subproblem optimizer attribute $(param.name) has not been set.")
    end
    return optimizer.sub_params[moi_param]
end

function MOI.set(optimizer::Optimizer, param::RawSubProblemOptimizerParameter, value)
    moi_param = MOI.RawParameter(param.name)
    optimizer.sub_params[moi_param] = value
    return nothing
end

function MOI.get(optimizer::Optimizer, ::Execution)
    return optimizer.execution
end

function MOI.set(optimizer::Optimizer, ::Execution, execution::AbstractExecution)
    optimizer.execution = execution
    return nothing
end

function MOI.get(optimizer::Optimizer, ::Penalizer)
    return optimizer.penalizer
end

function MOI.set(optimizer::Optimizer, ::Penalizer, penalizer::AbstractPenalizer)
    optimizer.penalizer = penalizer
    return nothing
end

function MOI.get(optimizer::Optimizer, ::PenaltyTerm)
    return optimizer.penaltyterm
end

function MOI.set(optimizer::Optimizer, ::PenaltyTerm, penaltyterm::AbstractPenaltyTerm)
    optimizer.penaltyterm = penaltyterm
    return nothing
end

function MOI.get(optimizer::Optimizer, param::ExecutionParameter)
    return MOI.get(optimizer.execution, param)
end

function MOI.set(optimizer::Optimizer, param::ExecutionParameter, value)
    MOI.set(optimizer.execution, param, value)
    return nothing
end

function MOI.get(optimizer::Optimizer, param::PenalizationParameter)
    return MOI.get(optimizer.penalizer, param)
end

function MOI.set(optimizer::Optimizer, param::PenalizationParameter, value)
    MOI.set(optimizer.penalizer, param, value)
    return nothing
end

function MOI.get(optimizer::Optimizer, ::MOI.TerminationStatus)
    return optimizer.status
end

function MOI.get(optimizer::Optimizer, ::MOI.PrimalStatus)
    return optimizer.primal_status
end

function MOI.get(optimizer::Optimizer, ::MOI.DualStatus)
    return optimizer.dual_status
end

function MOI.get(optimizer::Optimizer, ::MOI.RawStatusString)
    return optimizer.raw_status
end

function MOI.get(optimizer::Optimizer, ::MOI.VariablePrimal, index::MOI.VariableIndex)
    if optimizer.progressivehedging === nothing
        throw(StochasticProgram.UnloadedStructure{Optimizer}())
    end
    return decision(optimizer.progressivehedging, index)
end

function MOI.get(optimizer::Optimizer, ::MOI.ConstraintPrimal, ci::MOI.ConstraintIndex)
    if optimizer.progressivehedging === nothing
        throw(StochasticProgram.UnloadedStructure{Optimizer}())
    end
    return scalar_subproblem_reduction(optimizer.progressivehedging) do subproblem
        return MOI.get(subproblem.optimizer, MOI.ConstraintPrimal(), ci)
    end
end

function MOI.get(optimizer::Optimizer, ::MOI.ConstraintDual, ci::MOI.ConstraintIndex)
    if optimizer.progressivehedging === nothing
        throw(StochasticProgram.UnloadedStructure{Optimizer}())
    end
    return scalar_subproblem_reduction(optimizer.progressivehedging) do subproblem
        return MOI.get(subproblem.optimizer, MOI.ConstraintDual(), ci)
    end
end

function MOI.get(optimizer::Optimizer, ::MOI.ObjectiveValue)
    if optimizer.progressivehedging === nothing
        throw(StochasticProgram.UnloadedStructure{Optimizer}())
    end
    return objective_value(optimizer.progressivehedging)
end

function MOI.get(optimizer::Optimizer, ::MOI.DualObjectiveValue)
    if optimizer.progressivehedging === nothing
        throw(StochasticProgram.UnloadedStructure{Optimizer}())
    end
    if optimizer.status == MOI.OPTIMAL
        return objective_value(optimizer.progressivehedging)
    elseif optimizer.status == MOI.INFEASIBLE
        sense = MOI.get(optimizer.progressivehedging.structure, MOI.ObjectiveSense())
        return sense == MOI.MAX_SENSE ? Inf : -Inf
    elseif optimizer.status == MOI.DUAL_INFEASIBLE
        sense = MOI.get(optimizer.progressivehedging.structure, MOI.ObjectiveSense())
        return sense == MOI.MAX_SENSE ? -Inf : Inf
    else
        return NaN
    end
end

function MOI.get(optimizer::Optimizer, ::MOI.SolveTime)
    return optimizer.solve_time
end

function MOI.get(optimizer::Optimizer, attr::Union{MOI.AbstractOptimizerAttribute, MOI.AbstractModelAttribute})
    # Fallback to through structure
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.get(optimizer.progressivehedging.structure, attr)
end

function MOI.get(optimizer::Optimizer, attr::MOI.AbstractConstraintAttribute, ci::MOI.ConstraintIndex)
    # Fallback to first-stage optimizer through structure
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.get(optimizer.progressivehedging.structure, attr, ci)
end

function MOI.set(optimizer::Optimizer, attr::Union{MOI.AbstractOptimizerAttribute, MOI.AbstractModelAttribute}, value)
    # Fallback to first-stage optimizer through structure
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.set(optimizer.progressivehedging.structure, attr, value)
end

function MOI.set(optimizer::Optimizer, attr::MOI.AbstractVariableAttribute, index::MOI.VariableIndex, value)
    # Fallback to first-stage optimizer through structure
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.set(optimizer.progressivehedging.structure, attr, index, value)
end

function MOI.set(optimizer::Optimizer, attr::MOI.AbstractConstraintAttribute, ci::MOI.ConstraintIndex, value)
    # Fallback to first-stage optimizer through structure
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.set(optimizer.progressivehedging.structure, attr, ci, value)
end

function MOI.get(optimizer::Optimizer, attr::ScenarioDependentModelAttribute)
    # Fallback to subproblem optimizer through structure
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.get(optimizer.progressivehedging.structure, attr)
end

function MOI.get(optimizer::Optimizer, attr::ScenarioDependentVariableAttribute, index::MOI.VariableIndex)
    # Fallback to subproblem optimizer through structure
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.get(optimizer.progressivehedging.structure, attr, index)
end

function MOI.get(optimizer::Optimizer, attr::ScenarioDependentConstraintAttribute, ci::MOI.ConstraintIndex)
    # Fallback to subproblem optimizer through structure
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.get(optimizer.progressivehedging.structure, attr, ci)
end

function MOI.set(optimizer::Optimizer, attr::ScenarioDependentModelAttribute, value)
    # Fallback to subproblem optimizer through structure
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.set(optimizer.progressivehedging.structure, attr, value)
end

function MOI.set(optimizer::Optimizer, attr::ScenarioDependentVariableAttribute, index::MOI.VariableIndex, value)
    # Fallback to subproblem optimizer through structure
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.set(optimizer.progressivehedging.structure, attr, index, value)
end

function MOI.set(optimizer::Optimizer, attr::ScenarioDependentConstraintAttribute, ci::MOI.ConstraintIndex, value)
    # Fallback to subproblem optimizer through structure
    if optimizer.progressivehedging === nothing
        throw(UnloadedStructure{Optimizer}())
    end
    return MOI.set(optimizer.progressivehedging.structure, attr, ci, value)
end

function MOI.is_empty(optimizer::Optimizer)
    return optimizer.progressivehedging === nothing
end

MOI.supports(::Optimizer, ::MOI.Silent) = true
MOI.supports(::Optimizer, ::MOI.TimeLimitSec) = true
MOI.supports(::Optimizer, ::MOI.RawParameter) = true
MOI.supports(::Optimizer, ::AbstractStructuredOptimizerAttribute) = true
MOI.supports(::Optimizer, ::RelativeTolerance) = false
MOI.supports(::Optimizer, ::RawInstanceOptimizerParameter) = true
MOI.supports(::Optimizer, ::AbstractProgressiveHedgingAttribute) = true

# High-level attribute setting #
# ========================== #
"""
    get_penalization_attribute(stochasticprogram::StochasticProgram, name::String)

Return the value associated with the penalization-specific attribute named `name` in `stochasticprogram`.

See also: [`set_penalization_attribute`](@ref), [`set_penalization_attributes`](@ref).
"""
function get_penalization_attribute(stochasticprogram::StochasticProgram, name::String)
    return return MOI.get(optimizer(stochasticprogram), RawPenalizationParameter(name))
end
"""
    set_penalization_attribute(stochasticprogram::StochasticProgram, name::Union{Symbol, String}, value)

Sets the penalization-specific attribute identified by `name` to `value`.

"""
function set_penalization_attribute(stochasticprogram::StochasticProgram, name::Union{Symbol, String}, value)
    return set_optimizer_attribute(stochasticprogram, RawPenalizationParameter(String(name)), value)
end
"""
    set_penalization_attributes(stochasticprogram::StochasticProgram, pairs::Pair...)

Given a list of `attribute => value` pairs or a collection of keyword arguments, calls
`set_penalization_attribute(stochasticprogram, attribute, value)` for each pair.

"""
function set_penalization_attributes(stochasticprogram::StochasticProgram, pairs::Pair...)
    for (name, value) in pairs
        set_penalization_attribute(stochasticprogram, name, value)
    end
end
function set_penalization_attributes(stochasticprogram::StochasticProgram; kw...)
    for (name, value) in kw
        set_penalization_attribute(stochasticprogram, name, value)
    end
end


using Distributed, K8sClusterManagers
using DataFrames
using CSV

workers = parse(Int,ARGS[1])
cpu_fraction = parse(Int,ARGS[2])
addprocs(K8sClusterManager(workers,cpu=cpu_fraction); exeflags="--project")

#addprocs(K8sClusterManager(10, cpu=1, memory="300Mi", pending_timeout=300); exeflags="--project")
#nprocs() == 1 && addprocs(; exeflags = "--project")

@everywhere include("twoblocks.jl")
@everywhere using Main.TwoBlocksPipeline

dataset = getiris()
X = dataset[:,1:4]
Y = dataset[:,5] |> collect
X = dataset[:,2:end-1]

#dataset=CSV.read("environment_data.csv",DataFrame)
#X = select(dataset, Not([:state,:time]))
#Y = dataset.state

results=TwoBlocksPipeline.twoblockspipelinesearch(X,Y)
println(results)
println("best pipeline is: ", results[1,1],", with mean-sd ",results[1,2], " ± ",results[1,3])


# developed with Julia 1.1.1
#
# functions to simulate an EMS

function simulate_sites(controller::AbstractController,
                        path_to_save_folder::String,
                        path_to_price_csv_file::String,
                        path_to_metadata_csv_file::String,
                        path_to_test_data_folder::String,
                        path_to_train_data_folder::Union{String, Nothing}=nothing)

    mkpath(path_to_save_folder)
    prices = load_prices(path_to_price_csv_file)
    sites = load_sites(path_to_metadata_csv_file, path_to_test_data_folder,
        path_to_train_data_folder, path_to_save_folder)

    elapsed = 0.0

    for site in sites

        elapsed += @elapsed simulate_site(controller, site, prices)

    end

    group_all_simulations(sites)
    println("Terminating model simulation in $(elapsed) seconds")

    return nothing

end

function simulate_sites_parallel(controller::EMSx.AbstractController,
                                 path_to_save_folder::String,
                                 path_to_price_folder::String,
                                 path_to_metadata_csv_file::String,
                                 path_to_test_data_folder::String,
                                 path_to_train_data_folder::Union{String, Nothing}=nothing)

    mkpath(path_to_save_folder)
    prices = load_prices(path_to_price_folder)
    sites = load_sites(path_to_metadata_csv_file, path_to_test_data_folder,
        path_to_train_data_folder, path_to_save_folder)

    to_do = length(sites)

    @sync begin
        for p in workers()
            @async begin
                while true
                    idx = to_do
                    to_do -= 1
                    if idx <= 0
                        break
                    end
                    println("processing a new job - jobs left in queue : $(idx-1) / $(length(sites))")
                    _ = remotecall_fetch(simulate_site, p, controller, sites[idx], prices)
                end
            end
        end
    end

    group_all_simulations(sites)

end

function simulate_site(controller::AbstractController,
                       site::Site,
                       prices::Prices)

    test_data, site = load_site_data(site)
    controller = initialize_site_controller(controller, site, prices)
    periods = unique(test_data[!, :period_id])
    simulations = Simulation[]

    @showprogress for period_id in periods

        test_data_period = test_data[test_data.period_id .== period_id, :]
        period = Period(string(period_id), test_data_period, site)
        simulation = simulate_period(controller, period, prices)
        push!(simulations, simulation)

    end

    save_simulations(site, simulations)

    return nothing

end

function simulate_period(controller::AbstractController,
                           period::Period,
                           prices::Prices)

    horizon = size(period.data, 1) - 96 # test data: 24h of history lag + period data
    id = Id(period.site.id, period.id, prices.name, string(typeof(controller)))
    state_of_charge = 0.
    result = Result(horizon)
    timer = zeros(horizon)

    for t in 1:horizon

        information = Information(t, prices, period, state_of_charge)
        timing = @elapsed control = compute_control(controller, information)

        stage_cost, state_of_charge = apply_control(t, horizon, prices, period,
            state_of_charge, control)

        result.cost[t] = stage_cost
        result.soc[t] = state_of_charge
        timer[t] = timing

    end

    return Simulation(result, timer, id)

end

function apply_control(t::Int64, horizon::Int64,
                       prices::Prices, period::Period,
                       soc::Float64, control::Float64)
    """
    note on the load and pv values:
    at the end of the period, values at t+1 cannot be accessed: replaced by values at t=horizon
    with a minor impact since the optimal control is to empty the battery anyway
    """

    load = period.data[! ,:actual_consumption][min(t+96+1, horizon+96)]
    pv = period.data[!, :actual_pv][min(t+96+1, horizon+96)]
    net_energy_demand = load-pv

    stage_cost = compute_stage_cost(period.site.battery, prices, t, control, net_energy_demand)
    new_state_of_charge = compute_stage_dynamics(period.site.battery, soc, control)

    return stage_cost, new_state_of_charge

end


abstract type AbstractPlanetEnvironment end

struct PlanetEnvironment <: AbstractPlanetEnvironment
    GravityModel::Function
    PlanetModel::Function
    AtmosphereFlag::Bool #Boolean to determine whether atmoshpere model should be looked at
    AtmosphereModel::Function
    WindFlag::Bool #Boolean to determine whether wind model should be looked at
    WindModel::Function
end


struct PlotlinesModel{T₁ <: AbstractVector} <: AbstractModel
    data::T₁
end

mutable struct PlotlinesControl <: AbstractControl
    isenabled::Bool
end

struct PlotContext{T₁ <: AbstractControl,   T₂ <: AbstractModel,  T₃ <: AbstractDisplayProperties} <: AbstractPlotContext
    control::T₁
    model::T₂
    display_properties::T₃
end

function (context::PlotContext{<: PlotlinesControl,   <: PlotlinesModel,  <: PlotlinesDisplayProperties})()
        control = context.control
        model = context.model
        display_properties = context.display_properties
        id = get_id(display_properties)
        captioner = get_captioner(display_properties)
        caption = captioner(nothing)
        is_new_window(display_properties) ? CImGui.Begin(caption*id,C_NULL, CImGui.ImGuiWindowFlags_NoBringToFrontOnFocus) : nothing
        isenabled(control) ? control(model, display_properties) : nothing
        is_new_window(display_properties) ? CImGui.End() : nothing
end

function get_data(model::PlotlinesModel)
    model.data
end

function (control::PlotlinesControl)(model::PlotlinesModel, properties::PlotlinesDisplayProperties)
    id = get_id(properties)
    captioner = get_captioner(properties)
    caption = captioner(nothing)
    col = get_col(properties)
    rectangle = get_layout(properties)
    xaxis = get_xaxis(properties)
    yaxis = get_yaxis(properties)
    pos = CImGui.GetCursorScreenPos()
    totalwidth = get_width(rectangle)
    totalheight = get_height(rectangle)
    padding = get_padding(properties)
    data = get_data(model)
    CImGui.Dummy(ImVec2(0, padding[1]))
    CImGui.Indent(padding[2])
    width = totalwidth - padding[2]
    height = totalheight
    CImGui.PlotLines(id, data , length(data), 0 , caption, minimum(data), maximum(data), (width, height))
    isenabled(xaxis) ? draw_horizontal_ticks(data, width, height, get_ytick(properties)) : nothing
    isenabled(yaxis) ? draw_vertical_ticks(data, width, height, get_xtick(properties)) : nothing
    CImGui.Unindent(padding[2])
end

function draw_vertical_ticks(data, width, height, tick::Tickmark)
    if isenabled(tick)
        # Draw the x-axis
        draw_list = CImGui.GetWindowDrawList()
        black = Base.convert(ImU32, ImVec4(0,0, 0, 1))
        pos = CImGui.GetCursorScreenPos()
        CImGui.AddLine(draw_list, pos, ImVec2(pos.x + width, pos.y), black, get_thickness(tick));
        # Draw the concomitant tick marks.
        x = pos.x
        y = pos.y
        interpret = get_interpreter(tick)
        # for xₙ in range(x, stop = x + width; length = floor(Int, width / get_spacing(tick)) )
        #     # This line represents the tick mark.
        #     CImGui.AddLine(draw_list, ImVec2(xₙ, y), ImVec2(xₙ, y + get_length(tick)), black, get_thickness(tick)) #TODO add tick direction
        #     # Display value associated with that tickmark
        #     index = round(Int,stretch_linearly(xₙ, x,  x + width, 1, length(data)))
        #     # Convert the raw data to a text description after applying the
        #     # transformation associated with the value2text function.
        #     str = string(interpret(index))
        #     xoffset = div(length(str) * CImGui.GetFontSize(), 4)
        #     yoffset = get_length(tick) + 2
        #     CImGui.AddText(draw_list, ImVec2(xₙ - xoffset, y + yoffset), black, "$str",);
        # end
        gap = floor(Int, width / get_spacing(tick))
        increment = round(Int, stretch_linearly(gap , 1, width, 1, length(data)))
        for index in 1:increment:length(data)
            xₙ = Cfloat(stretch_linearly(index , 1, length(data), x,  x + width))
            # This line represents the tick mark.
            CImGui.AddLine(draw_list, ImVec2(xₙ, y), ImVec2(xₙ, y + get_length(tick)), black, get_thickness(tick)) #TODO add tick direction
            # Display value associated with that tickmark
            # Convert the raw data to a text description after applying the
            # transformation associated with the value2text function.
            str = string(interpret(index))
            xoffset = div(length(str) * CImGui.GetFontSize(), 4)
            yoffset = get_length(tick) + 2
            CImGui.AddText(draw_list, ImVec2(xₙ - xoffset, y + yoffset), black, "$str",);
        end
    end
end

function draw_horizontal_ticks(data, width, height, tick::Tickmark)

    if isenabled(tick)
        draw_list = CImGui.GetWindowDrawList()
        black = Base.convert(ImU32, ImVec4(0,0, 0, 1))
        col = Base.convert(ImU32, ImVec4(0.95,0.1, 0.1, 0.1))
        pos = CImGui.GetCursorScreenPos()
        # Draw the y-axis
        CImGui.AddLine(draw_list, pos, ImVec2(pos.x, pos.y - height), black, get_thickness(tick));
        minval = minimum(data)
        maxval = maximum(data)
        # Draw the concomitant tick marks.
        x = pos.x
        y = pos.y
        interpret = get_interpreter(tick)
        # TODO: Perhaps better to iterate over bonafide values and round pixel coordinates?
        span = range(y - height, stop = y; length = floor(Int, height / get_spacing(tick)))
        count = 1
        for yₙ in span
            # This line represents the tick mark.
            CImGui.AddLine(draw_list, ImVec2(x, yₙ), ImVec2(x - get_length(tick), yₙ), black, get_thickness(tick)); #TODO add tick direction
            # Display value associated with that tickmark
            val = round(stretch_linearly(yₙ, y,  y - height, minval, maxval); digits = 3)
            # Convert the raw data to a text description after applying the
            # transformation associated with the value2text function.
            str = string(interpret(val))
            xoffset = get_length(tick) + 3 + div(length(str) * CImGui.GetFontSize(), 2)
            yoffset = div(CImGui.GetFontSize(), 2)
            CImGui.AddText(draw_list, ImVec2(x - xoffset, yₙ - yoffset), black, "$str",);
            # Draw a faint filled rectangle as a visual guide for every second row
            if iseven(count)
                 CImGui.AddRectFilled(draw_list, ImVec2(x, yₙ), ImVec2(x + width, yₙ + step(span)), col);
            end
            count += 1
        end
    end
end