module Graficos

using Plots
import ..FEA: Estrutura, forcas_internas
import ..Otm: Otimizador

export plotar_estrutura_inicial,
       plotar_estrutura_otimizada

"""
    maior_dimensao(estrutura::Estrutura)::Float64

Retorna o tamanho da maior dimensão (horizontal ou vertical) da estrutura.
"""
function maior_dimensao(estrutura::Estrutura)::Float64
    min_x = Inf
    max_x = -Inf
    min_y = Inf
    max_y = -Inf

    for no in estrutura.nos
        min_x = min(min_x, no.x)
        max_x = max(max_x, no.x)
        min_y = min(min_y, no.y)
        max_y = max(max_y, no.y)
    end

    return max(max_x - min_x, max_y - min_y)
end

"""
    plotar_forcas_apoios!(p, estrutura::Estrutura, estrutura_inicial=true, ponderador=0.05)

Plota no gráfico `p` as forças e os apoios da estrutura.

# Argumentos
- `p`: Gráfico (plot).
- `estrutura::Estrutura`: Estrutura analisada.
- `estrutura_inicial::Bool=true`: Parâmetro que altera as cores dos indicadores de forças e apoios
                                  caso a estrutura seja a inicial ou a otimizada.
- `ponderador::Real`: Parâmetro que indica o tamanho dos indicadores de forças e apoios. É proporcional
                      à maior dimensão da estrutura.
"""
function plotar_forcas_apoios!(p, estrutura::Estrutura, estrutura_inicial::Bool=true, ponderador::Real=0.05)
    lmax = maior_dimensao(estrutura)
    cor_forcas = estrutura_inicial ? :red : :green
    cor_apoios = estrutura_inicial ? :blue : :black
    # Plotagem dos apoios e das forças.
    for no in estrutura.nos
        # Apoios
        if no.rx
            x = [NaN, no.x - ponderador/2*lmax, no.x]
            y = [NaN, no.y, no.y]
            plot!(p, x, y, lw=ponderador*50, label="", color=cor_apoios)
        end
        if no.ry
            x = [NaN, no.x, no.x]
            y = [NaN, no.y - ponderador/2*lmax, no.y]
            plot!(p, x, y, lw=ponderador*50, label="", color=cor_apoios)
        end

        # Forças
        if no.fx > 0
            x = [NaN, no.x, no.x + ponderador*lmax]
            y = [NaN, no.y, no.y]
            plot!(p, x, y, lw=ponderador*50, label="", color=cor_forcas, arrow=true, arrowsize=ponderador*10)
        elseif no.fx < 0
            x = [NaN, no.x, no.x - ponderador*lmax]
            y = [NaN, no.y, no.y]
            plot!(p, x, y, lw=ponderador*50, label="", color=cor_forcas, arrow=true, arrowsize=ponderador*10)
        end
        if no.fy > 0
            x = [NaN, no.x, no.x]
            y = [NaN, no.y, no.y + ponderador*lmax]
            plot!(p, x, y, lw=ponderador*50, label="", color=cor_forcas, arrow=true, arrowsize=ponderador*10)
        elseif no.fy < 0
            x = [NaN, no.x, no.x]
            y = [NaN, no.y, no.y - ponderador*lmax]
            plot!(p, x, y, lw=ponderador*50, label="", color=cor_forcas, arrow=true, arrowsize=ponderador*10)
        end
    end
end

"""
    plotar_estrutura_inicial(estrutura::Estrutura)

Plota a estrutura antes da otimização.
"""
function plotar_estrutura_inicial(estrutura::Estrutura)
    p = plot()
    # Plotagem dos elementos.
    x = []
    y = []
    for el in estrutura.elementos
        push!(x, [NaN, el.no₁.x, el.no₂.x]...)
        push!(y, [NaN, el.no₁.y, el.no₂.y]...)
    end
    plot!(p, x, y, label="", color=:black, aspect_ratio=:equal)

    plotar_forcas_apoios!(p, estrutura, true)

    display(p)
    sleep(10000)
end

"""
    plotar_estrutura_otimizada(otimizador::Otimizador)

Plota a estrutura após a otimização.
"""
function plotar_estrutura_otimizada(otimizador::Otimizador)
    x_norm = otimizador.x / maximum(otimizador.x)
    forcas = forcas_internas(otimizador.estrutura)

    p = plot()
    for i in eachindex(otimizador.estrutura.elementos)
        area = x_norm[i]
        if isapprox(area, 0; atol=eps())
            continue
        end

        forca = forcas[i]
        el = otimizador.estrutura.elementos[i]

        x = [NaN, el.no₁.x, el.no₂.x]
        y = [NaN, el.no₁.y, el.no₂.y]

        plot!(p, x, y, label="", lw=area*5, color= forca ≤ 0 ? :blue : :red, aspect_ratio=:equal)
    end

    plotar_forcas_apoios!(p, otimizador.estrutura, false)

    display(p)
    sleep(10000)
end

end


import MLJBase
import DataFrames
import CategoricalArrays

mutable struct XGBClassifier <: MLJBase.Probabilistic
    maxdepth::Integer
    nrounds::Integer
 end

function XGBClassifier(; maxdepth = 6, nrounds = 2)
    model = XGBClassifier(maxdepth, nrounds)
    return model
end

function MLJBase.scitypes(X::Vector{<:AbstractFactor})
    [MLJBase.Multiclass for x in X]
end

function MLJBase.scitype_union(y::AbstractFactor)
    MLJBase.Multiclass
end

function MLJBase.scitype_union(y::AbstractCovariate)
    MLJBase.Multiclass
end

function MLJBase.target_scitype_union(model::XGBClassifier)
    MLJBase.Multiclass
end

function MLJBase.selectrows(X::Vector{<:AbstractFactor}, r)
    (X, r)
end

function MLJBase.selectrows(X::JuML.DataFrame, r)
    (X, r)
end

function MLJBase.selectrows(y::AbstractFactor, r)
    (y, r)
end

function MLJBase.selectrows(y::AbstractCovariate, r)
    (y, r)
end

function MLJBase.fit(model::XGBClassifier, verbosity::Integer, X::DataFrames.DataFrame,
                     y::CategoricalArrays.CategoricalArray)

    X = JuML.DataFrame(X)
    levels = [string(level) for level in CategoricalArrays.index(y.pool)]
    T = eltype(y.refs)
    y = Factor{T}("y", levels, y.refs)
    label = covariate(y, (level -> level == levels[1] ? 1.0 : 0.0))
    trainsel = BoolVariate("", BitArray{1}(undef, 0))
    fitres = xgblogit(label, X.factors; trainselector = trainsel, maxdepth = model.maxdepth,
                      nrounds = model.nrounds)
    return fitres, nothing, nothing
end

function MLJBase.fit(model::XGBClassifier, verbosity::Integer, X::Tuple{Vector{<:AbstractFactor}, Any},
                     y::Tuple{AbstractCovariate, Any})

    trainsel = X[2] <: Colon ?
        begin
            BoolVariate("", BitArray{1}(undef, 0))
        end :
        begin
            n = length(y[1])
            b = falses(n)
            for i in X[2]
                b[i] = true
            end
            BoolVariate("", b)
        end

    factors = X[1]
    label = y[1]
    fitres = xgblogit(label, factors; trainselector = trainsel, maxdepth = model.maxdepth,
                      nrounds = model.nrounds)
    return fitres, nothing, nothing

end

function MLJBase.fit(model::XGBClassifier, verbosity::Integer, X::Tuple{JuML.DataFrame, Any},
                     y::Tuple{AbstractCovariate, Any})

    factors = X[1].factors
    MLJBase.fit(model, verbosity, (factors, X[2]), y)
end

function MLJBase.fit(model::XGBClassifier, verbosity::Integer, X::Tuple{JuML.DataFrame, Any},
                     y::Tuple{AbstractFactor, Any})

    factors = X[1].factors
    ylevels = getlevels(y[1])
    label = covariate(getname(y), (level -> level == ylevels[1] ? 1.0 : 0.0))
    MLJBase.fit(model, verbosity, (factors, X[2]), (label, y[2]))

end

function MLJBase.predict(model::XGBClassifier, fitresult, Xnew::JuML.DataFrame)

    predict(fitresult, Xnew)

end

function MLJBase.predict(model::XGBClassifier, fitresult, Xnew)

    predict(fitresult, JuML.DataFrame(Xnew))

end


# This example shows a simple use for a channel and a feeder

using Distributed
using Logging
using Base.CoreLogging

# Set Debug for Logging
global_logger(ConsoleLogger(stdout, Logging.Debug))

function read_and_print(channel::Channel{Tuple{Vector{Int}, Int}})

    @debug("Running updater.")

    while true

        x, i = try
            take!(channel);
        catch e

            if isa(e, InvalidStateException)
                break
            end

            @warn("Something wrong when reading from channel. Will skip.", e)

            continue

        end

        @debug("Updater has read values from channel. x = $(x) and i = $(i).")

    end

    @debug("Update channel closed. Exiting thread.")

end

c = Channel{Tuple{Vector{Int}, Int}}(0)

# Task waits for the channel to be fed
task1 = @async read_and_print(c)

# Feeding the channel
task2 = @async for i = 1:10
    put!(c, (ones(Int, 3), i))
end

sleep(1)

# Wrong feeding
put!(c, 1)

# This should close `task1`
close(c)

# Just to be sure
sleep(1)

println("Task1 done = $(istaskdone(task1))")

println("Task2 done = $(istaskdone(task2))")


import Base.show
using Base64: base64encode

struct WAVArray{T,N}
    Fs::Number
    data::AbstractArray{T,N}
    caption::AbstractString
end

WAVArray(Fs::Number, data::AbstractArray) = WAVArray(Fs, data, "")

function WAVArray(io::IO, caption::AbstractString; subrange=(:), format="double")
    data, Fs = wavread(io, subrange=subrange, format=format)
    return WAVArray(Fs, data, caption)
end

WAVArray(io::IO; subrange=(:), format="double") = WAVArray(io::IO, ""; subrange=subrange, format=format)

function WAVArray(filename::AbstractString, caption::AbstractString; subrange=(:), format="double")
    data, Fs = wavread(filename, subrange=subrange, format=format)
    return WAVArray(Fs, data, caption)
end

WAVArray(filename::AbstractString; subrange=(:), format="double") = WAVArray(filename, ""; subrange=subrange, format=format)

wavwrite(x::WAVArray, io::IO) = wavwrite(x.data, io; Fs=x.Fs)

function show(io::IO, ::MIME"text/html", x::WAVArray)
    buf = IOBuffer()
    wavwrite(x, buf)
    data = base64encode(String(take!(copy(buf))))
    markup = """<figure>
                $(ifelse(x.caption !== "", "<figcaption>$(x.caption)</figcaption>", ""))
                    <audio controls="controls" {autoplay}>
                    <source src="data:audio/wav;base64,$data" type="audio/wav" />
                    Your browser does not support the audio element.
                    </audio>
                </figure>"""
    print(io, markup)
end


using DNSS, Test

libraries = ["MinkowskiSphericalSymmetry"]

for file in libraries
    @info "Testing $file"
    include("$(file)Test.jl")
end


# Plot coordinates of the C alpha with best occupancy
@recipe function plot(residues::AbstractVector{PDBResidue})
    ca = CAmatrix(residues)
    chains = [r.id.chain for r in residues if r.id.group == "ATOM"]
    group --> chains
    ca[:,1], ca[:,2], ca[:,3]
end


module Mapper2

const SRCDIR = @__DIR__
const PKGDIR = dirname(SRCDIR)
using Reexport

using Logging

# Aux modules
include("Helper.jl")
using .Helper

include("MapperGraphs.jl")

# Central data types.
include("MapperCore.jl")

# Placement modules
include("Place/SA/SA.jl")

# Routing modules
include("Route/Route.jl")

# exports from Helper.
# TODO: DO we really need these exports?
export Address, dim_max, dim_min, place!

# Use submodules to make exports visible.
@reexport using .MapperCore

#############
# PLACEMENT #
#############
# Placement Algorithms
@reexport using .SA

# Default Placement Algorithm
place!(map::Map; kwargs...) = SA.place!(map; kwargs...)

###########
# ROUTING #
###########
@reexport using .Routing

end #module Mapper2


"""
    update!(sim::PlanktonSimulation; time_offset = (vels = true, PAFR = true, temp = true))

Update the `PlanktonSimulation` for `sim.iterations` time steps.
`time_offset` is used when velocities (or PARF or temperature) starts from timestep 1, but model.t is not.
It is usually used when velocity fields are too large and need to be broken down into several parts.
Only one part of the whole velocity fields can be constructed into a `PlanktonSimulation`, so in this PlanktonSimulation
`model.iteration` might no be 1, but the velocity fields need to start from 1 (same for PARF or temperature fields).

"""
function update!(sim::PlanktonSimulation; time_offset = (vels = false, PARF = false, temp = false))
    if sim.input.vels ≠ (;)
        for i in 1:sim.iterations
            model_t_vels = time_offset.vels ? (i-1)*sim.ΔT : sim.model.t
            model_t_PARF = time_offset.PARF ? (i-1)*sim.ΔT : sim.model.t
            model_t_temp = time_offset.temp ? (i-1)*sim.ΔT : sim.model.t

            t_vel = Int(model_t_vels÷sim.input.ΔT_vel)+1 # starting from 1
            vel_copy!(sim.model.timestepper.vel₁, sim.input.vels.u[:,:,:,t_vel],
                    sim.input.vels.v[:,:,:,t_vel], sim.input.vels.w[:,:,:,t_vel], sim.model.grid)

            t_par = Int(model_t_PARF÷sim.input.ΔT_PAR)+1 # starting from 1
            copyto!(sim.model.timestepper.PARF, sim.input.PARF[:,:,t_par])

            t_temp = Int(model_t_temp÷sim.input.ΔT_temp)+1 # starting from 1
            copy_interior!(sim.model.timestepper.temp, sim.input.temp[:,:,:,t_temp], sim.model.grid)

            TimeStep!(sim.model, sim.ΔT, sim.diags)

            write_output!(sim.output_writer, sim.model, sim.diags, sim.ΔT)
        end
    else
        for i in 1:sim.iterations
            model_t_PARF = time_offset.PARF ? sim.model.t : (i-1)*sim.ΔT
            model_t_temp = time_offset.temp ? sim.model.t : (i-1)*sim.ΔT

            t_par = Int(model_t_PARF÷sim.input.ΔT_PAR)+1 # starting from 1
            copyto!(sim.model.timestepper.PARF, sim.input.PARF[:,:,t_par])

            t_temp = Int(model_t_temp÷sim.input.ΔT_temp)+1 # starting from 1
            copy_interior!(sim.model.timestepper.temp, sim.input.temp[:,:,:,t_temp], sim.model.grid)

            TimeStep!(sim.model, sim.ΔT, sim.diags)

            write_output!(sim.output_writer, sim.model, sim.diags, sim.ΔT)
        end
    end
end


using wembed
using Base.Test

@test isfile("../res/GoogleNews-vectors-negative300.bin")

bin_path = "../res/GoogleNews-vectors-negative300.bin"
model_path = "../res/word2vec_model.hdf5"
@time @test wembed.bin2hdf5(bin_path, model_path)

@time @test wembed.init(model_path)
@time @test typeof(wembed.getvector("queen", false)) == Vector{Float32}
@time @test typeof(getsimilarity("king", "queen")) == Float32
@time @test typeof(wordnearest("king", 5)) == Array{Pair{AbstractString,Float32},1}

wordsAsString = "breakfast sushi dinner lunch"
@time @test typeof(odd1out(wordsAsString)) == String

wordset1 = ["king", "man"]
wordset2 = ["queen", "woman"]
@time @test typeof(nsimilarity(wordset1, wordset2)) == Float32

words = map(String, split("Julia is a developing language"))
getvector(words, fill(0.25,300))


module FitnessPredictionsScript

using SynapseClient
using DISSEQT
using DISSEQT.SynapseTools
using DISSEQT.Plots
using DISSEQT.AnnotatedArrays
using LinearAlgebra
using DataFrames
using CSV
using Statistics
using JLD
using Gadfly
using Colors

# put mutagen doses into low,medium,high categories
function addlevel!(metadata)
    levelDict = Dict(("5FU","50uM")=>"Low", ("5FU","100uM")=>"Medium", ("5FU","150uM")=>"Medium", ("5FU","200uM")=>"High", ("AML","50uM")=>"Low", ("AML","100uM")=>"Medium", ("AML","200uM")=>"High", ("AZC","50uM")=>"Low", ("AZC","100uM")=>"Medium", ("AZC","200uM")=>"Medium", ("AZC","300uM")=>"High", ("Mn","0.25mM")=>"Medium", ("Mn","0.33mM")=>"Medium", ("Mn","0.5mM")=>"High", ("Mock","0uM")=>"Low", ("RIBA","50uM")=>"Low", ("RIBA","100uM")=>"Medium", ("RIBA","200uM")=>"Medium", ("RIBA","300uM")=>"High")
    metadata[:Level] = map( (m,d)->get(levelDict, (m,d), missing), metadata[:Mutagen], metadata[:Dose] )
end

function main()
    syn = SynapseClient.login()

    # projectFolder should point to MyProject folder (containing subfolders Metadata and Analysis)
    projectFolder  = "syn11639899" # VignuzziLabPublic/Projects/FitnessLandscapes
    analysisFolder = getchildbyname(syn, projectFolder, "Analysis")
    fitnessFolder  = getchildbyname(syn, projectFolder, "Fitness")

    isdir("MutantSwarms") || mkdir("MutantSwarms")
    isdir("plots") || mkdir("plots")

    # set to true to upload files
    doUpload = true
    uploadName = "fitness landscapes invitro"

    metadataID = getchildbyname(syn, projectFolder, "Metadata", "fitness landscapes invitro.csv")
    metadata = getsamplemetadata(syn, metadataID, :Rejected=>"No") # Always remove rejected samples
    metadata = metadata[ .~ismissing.(metadata[:Dose]), : ] # skip samples with unknown dose

    seqSpaceRepID = getchildbyname(syn, analysisFolder, "SequenceSpace", uploadName, "SequenceSpaceRepresentation", "seqspacerepresentation.jld")
    seqSpaceRepDict = load(localpath(syn, seqSpaceRepID))
    U = seqSpaceRepDict["U"]
    Σ = seqSpaceRepDict["Σ"]
    V = seqSpaceRepDict["V"]
    sampleID  = seqSpaceRepDict["SampleID"]
    positions = seqSpaceRepDict["positions"]
    codons    = seqSpaceRepDict["codons"]
    consensus = permutedims(seqSpaceRepDict["consensus"],(2,1)) # transpose to get Samples x Codons

    nbrDims = length(Σ)

    Y = V.*Σ' # sample representation scaled by singular values

    # also load the Consensus-based PCA representation
    seqSpaceRepConsensusID = getchildbyname(syn, analysisFolder, "SequenceSpace", uploadName, "SequenceSpaceRepresentation", "seqspacerepresentationConsensusPCA.jld")
    seqSpaceRepConsensusDict = load(localpath(syn, seqSpaceRepConsensusID))
    UConsensus = seqSpaceRepConsensusDict["U"]
    ΣConsensus = seqSpaceRepConsensusDict["Σ"]
    VConsensus = seqSpaceRepConsensusDict["V"]
    @assert sampleID == seqSpaceRepConsensusDict["SampleID"] # make sure the same samples were used
    nbrDimsConsensus = length(ΣConsensus)

    YConsensus = VConsensus.*ΣConsensus' # sample representation scaled by singular values

    # filter metadata by samples included in the Sequence Space Representation
    ind = indexin(sampleID, metadata[:SampleID])
    @assert all(ind.!=0) "Unknown SampleID"
    metadata = metadata[ind,:]

    addlevel!(metadata) # categorize dosages as low, medium, high

    # get fitness
    fitnessID = getchildbyname(syn,fitnessFolder,"fitness landscapes invitro fitness.csv")
    appendfitness!(metadata, getsamplefitness(syn, fitnessID))

    # setup dependencies
    dependencies = unique(vcat(seqSpaceRepID, seqSpaceRepConsensusID, metadata[:MetadataID], collect(skipmissing(metadata[:FitnessID]))))

    uploads = UploadList()

    virusColors = Dict("WT"  =>colorant"black",
                       "Less"=>colorant"blue",
                       "More"=>colorant"green",
                       "Stop"=>colorant"red")
    virusColorMap = groupcolors(metadata[:Reference],virusColors)

    mutagenColors = Dict("5FU" =>RGB(1,0,0),
                         "AML" =>RGB(0,1,0),
                         "AZC" =>RGB(0,0,1),
                         "Mn"  =>RGB(0.8,0.8,0),
                         "Mock"=>RGB(1,0,1),
                         "RIBA"=>RGB(0,1,1))
    mutagenColorMap = groupcolors(metadata[:Mutagen],mutagenColors)

    # Isomap representation
    K = Y*Y'
    D = sqrt.(diag(K) .+ diag(K)' .- 2*K) # sample distance matrix
    Z,stress,converged,dists = kruskalisomap(D, 2, 2, maximum(D)/20)
    converged || error("Kruskal's stress MDS did not converge.")

    # Isomap representation (Consensus PCA)
    KConsensus = YConsensus*YConsensus'
    DConsensus = sqrt.(max.(0., diag(KConsensus) .+ diag(KConsensus)' .- 2*KConsensus)) # sample distance matrix (make sure it's nonnegative to handle rounding errors for colocated points)
    ZConsensus,_,converged,_ = kruskalisomap(DConsensus, 2, 2, maximum(DConsensus)/20)
    converged || error("Kruskal's stress MDS did not converge for Consensus PCA.")

    # prediction
    fitnessMask = .~ismissing.(metadata[:Fitness])

    # make a table with all predictors
    predictionResults = DataFrame(Name=String[], Type=String[], VarianceExplained=Float64[], NbrSamples=Int[])

    # run twice, the second time excluding samples that constitute singleton groups in the "Lineage, Mutagen, Dose" group predictor
    mask = copy(fitnessMask)
    for i=1:2
        fitness = collect(skipmissing(metadata[mask,:Fitness]))

        # Isomap landscape
        σValues = 10.0.^range(-3,stop=1,length=1000) * sqrt(mean(sum(Z[mask,:].^2; dims=2))) # scale by point size cloud
        @time σPerSample,_,_ = leaveoneoutlandscapekernelwidth(Z[mask,:], fitness, σValues, nbrIter=1000)
        f = leaveoneoutpredict(LandscapeModel, Z[mask,:], fitness, perSampleModelArgs=σPerSample)
        push!(predictionResults, ("[GKS] Isomap (2d)", "Gaussian Kernel Smoother", 1.0.-varianceunexplained(f,fitness), count(mask)) )
        # println(predictionResults)

        # Reduced space landscape
        σValues = 10.0.^range(-3,stop=1,length=1000) * sqrt(mean(sum(Y[mask,:].^2; dims=2))) # scale by point size cloud
        @time σPerSample,_,_ = leaveoneoutlandscapekernelwidth(Y[mask,:], fitness, σValues, nbrIter=1000)
        f = leaveoneoutpredict(LandscapeModel, Y[mask,:], fitness, perSampleModelArgs=σPerSample)
        push!(predictionResults, ("[GKS] SMSSVD ($(nbrDims)d)", "Gaussian Kernel Smoother", 1.0.-varianceunexplained(f,fitness), count(mask)) )
        # println(predictionResults)

        # Isomap landscape (Consensus PCA)
        σValues = 10.0.^range(-3,stop=1,length=1000) * sqrt(mean(sum(ZConsensus[mask,:].^2; dims=2))) # scale by point size cloud
        @time σPerSample,_,_ = leaveoneoutlandscapekernelwidth(ZConsensus[mask,:], fitness, σValues, nbrIter=1000)
        f = leaveoneoutpredict(LandscapeModel, ZConsensus[mask,:], fitness, perSampleModelArgs=σPerSample)
        push!(predictionResults, ("[GKS] Consensus Isomap (2d)", "Gaussian Kernel Smoother", 1.0.-varianceunexplained(f,fitness), count(mask)) )
        # println(predictionResults)

        # Reduced space landscape  (Consensus PCA)
        σValues = 10.0.^range(-3,stop=1,length=1000) * sqrt(mean(sum(YConsensus[mask,:].^2; dims=2))) # scale by point size cloud
        @time σPerSample,_,_ = leaveoneoutlandscapekernelwidth(YConsensus[mask,:], fitness, σValues, nbrIter=1000)
        f = leaveoneoutpredict(LandscapeModel, YConsensus[mask,:], fitness, perSampleModelArgs=σPerSample)
        push!(predictionResults, ("[GKS] Consensus PCA ($(nbrDimsConsensus)d)", "Gaussian Kernel Smoother", 1.0.-varianceunexplained(f,fitness), count(mask)) )

        # Nearest neighbor predictor: Reduced dimension
        f = leaveoneoutpredict(NearestNeighborModel, Y[mask,:], fitness)
        push!(predictionResults, ("[NN] SMSSVD ($(nbrDims)d)", "Nearest Neighbor", 1.0.-varianceunexplained(f,fitness), count(mask)) )

        # Nearest neighbor predictor: Consensus
        f = leaveoneoutpredict(NearestNeighborModel, consensus[mask,:], fitness)
        push!(predictionResults, ("[NN] Consensus", "Nearest Neighbor", 1.0.-varianceunexplained(f,fitness), count(mask)) )

        # Group predictor: Lineage, Drug
        f = leaveoneoutpredict(GroupModel, metadata[mask,[:Reference,:Mutagen]], fitness)
        push!(predictionResults, ("[G] Lineage/Mutagen", "Group", 1.0.-varianceunexplained(f,fitness), count(mask)) )

        # Group predictor: Lineage, Level
        f = leaveoneoutpredict(GroupModel, metadata[mask,[:Reference,:Level]], fitness)
        push!(predictionResults, ("[G] Lineage/Dose", "Group", 1.0.-varianceunexplained(f,fitness), count(mask)) )

        # Group predictor: Lineage, Drug, Dose
        #f = leaveoneoutpredict(GroupModel, metadata[mask,[:Reference,:Mutagen,:Dose]], fitness)
        f = leaveoneoutpredict(GroupModel, metadata[mask,[:Reference,:Mutagen,:Level]], fitness)
        push!(predictionResults, ("[G] Lineage/Mutagen/Dose", "Group", 1.0.-varianceunexplained(f,fitness), count(mask)) )
        newMask = i==1 ? .~isnan.(f) : trues(size(f)) # For next iteration, exclude singleton groups
        i==1 && println(count(.~newMask), " singleton groups removed.")

        mask[mask] = newMask
    end

    println(predictionResults)
    CSV.write("plots/fitnesspredictions.csv", predictionResults; missingstring="NA")
    markforupload!(uploads, "FitnessPredictionPlots", "plots/fitnesspredictions.csv", dependencies)

    N = div(size(predictionResults,1),2)
    xMax = maximum(predictionResults[.~isnan.(predictionResults[:VarianceExplained]), :VarianceExplained])

    # with all samples
    df = predictionResults[1:N,:]
    df = df[.~isnan.(df[:VarianceExplained]),:] # remove predictors that couldn't be used ([G] Lineage/Mutagen/Dose)
    plotArgs = (layer(df, x=:VarianceExplained, color=:Name, Geom.bar(orientation=:horizontal)),
                Coord.Cartesian(xmin=0, xmax=xMax, yflip=true),
                Guide.yticks(ticks=nothing),
                Guide.xlabel("Variance Explained"), Guide.ylabel(""))
    themeArgs = (:key_title_font_size=>2*11pt, :key_label_font_size=>2*8pt, :major_label_font_size=>2*11pt, :minor_label_font_size=>2*8pt, :point_label_font_size=>2*8pt)

    pl = plot(plotArgs..., Theme(;themeArgs...))
    markforupload!(uploads, "FitnessPredictionPlots", saveplot(pl, [:png,:svg,:pdf], "plots/fitnesspredictions_allsamples"), dependencies)

    pl = plot(plotArgs..., Theme(key_position=:none; themeArgs...)) # no legend
    markforupload!(uploads, "FitnessPredictionPlots", saveplot(pl, [:png,:svg,:pdf], "plots/fitnesspredictions_allsamples_nolegend"), dependencies)

    # with singleton groups removed
    df = predictionResults[N+1:end,:]
    df = df[.~isnan.(df[:VarianceExplained]),:] # remove predictors that couldn't be used (none)
    plotArgs = (layer(df, x=:VarianceExplained, color=:Name, Geom.bar(orientation=:horizontal)),
                Coord.Cartesian(xmin=0, xmax=xMax, yflip=true),
                Guide.yticks(ticks=nothing),
                Guide.xlabel("Variance Explained"), Guide.ylabel(""))
    themeArgs = (:key_title_font_size=>2*11pt, :key_label_font_size=>2*8pt, :major_label_font_size=>2*11pt, :minor_label_font_size=>2*8pt, :point_label_font_size=>2*8pt)

    pl = plot(plotArgs..., Theme(;themeArgs...))
    markforupload!(uploads, "FitnessPredictionPlots", saveplot(pl, [:png,:svg,:pdf], "plots/fitnesspredictions"), dependencies)

    pl = plot(plotArgs..., Theme(key_position=:none; themeArgs...)) # no legend
    markforupload!(uploads, "FitnessPredictionPlots", saveplot(pl, [:png,:svg,:pdf], "plots/fitnesspredictions_nolegend"), dependencies)

    doUpload && uploadfiles(syn, createchildfolder(syn, analysisFolder, "SequenceSpace", uploadName), uploads, @__FILE__, "Sequence Space Plots")

    # seqSpaceRepDict, metadata, U, Σ, V, Y, Z
    nothing
end

end
FitnessPredictionsScript.main();
# seqSpaceRepDict, metadata, U, Σ, V, Y, Z = FitnessPredictionsScript.main();


using ProgressMeter

T = 50
N = 50

pomdp = BabyPOMDP()

# bds = IndependentBounds(PORollout(FeedWhenCrying(), PreviousObservationUpdater()), 0.0)
bds = IndependentBounds(SemiPORollout(FeedWhenCrying()), 0.0)

solver = AdaOPSSolver(epsilon_0=0.1,
                      m_min=200,
                      max_depth=50,
                      bounds=bds,
                      T_max=Inf,
                      max_trials=500,
                      rng=MersenneTwister(4)
                     )

rsum = 0.0
fwc_rsum = 0.0
@showprogress for i in 1:N
    planner = solve(solver, pomdp)
    sim = RolloutSimulator(max_steps=T, rng=MersenneTwister(i))
    fwc_sim = deepcopy(sim)
    global rsum += simulate(sim, pomdp, planner)
    global fwc_rsum += simulate(fwc_sim, pomdp, FeedWhenCrying())
end

@show rsum/N
@show fwc_rsum/N


macro checkpoint(alg, adtool, loop)
    ex = quote
        function tobedifferentiated(inputs)
            local F_H = similar(inputs)
            local F = inputs
            $(loop.args[2])
            outputs = F
            return outputs
        end
        if isa($alg, Revolve)
            storemap = Dict{Int32,Int32}()
            check = 0
            F_Check = Array{Any, 2}(undef, 3, $alg.acp)
            F_final = Array{Float64, 1}(undef, 2)
            while true
                next_action = next_action!($alg)
                if (next_action.actionflag == Checkpointing.store)
                    check = check+1
                    storemap[next_action.iteration-1]=check
                    $alg.fstore(F,F_Check,t,check)
                elseif (next_action.actionflag == Checkpointing.forward)
                    for j= next_action.startiteration:(next_action.iteration - 1)
                        $(loop.args[2])
                    end
                elseif (next_action.actionflag == Checkpointing.firstuturn)
                    $(loop.args[2])
                    F_final .= F
                    L .= [0, 1]
                    t = 1.0-h
                    L_H .= L
                    L = Checkpointing.jacobian(tobedifferentiated, F_H, $adtool)[2,:]
                elseif (next_action.actionflag == Checkpointing.uturn)
                    L_H .= L
                    F_H = F
                    res = Checkpointing.jacobian(tobedifferentiated, F_H, $adtool)
                    L =  transpose(res)*L
                    t = t - h
                    if haskey(storemap,next_action.iteration-1-1)
                        delete!(storemap,next_action.iteration-1-1)
                        check=check-1
                    end
                elseif (next_action.actionflag == Checkpointing.restore)
                    F, t = $alg.frestore(F_Check,storemap[next_action.iteration-1])
                elseif next_action.actionflag == Checkpointing.done
                    if haskey(storemap,next_action.iteration-1-1)
                        delete!(storemap,next_action.iteration-1-1)
                        check=check-1
                    end
                    break
                end
            end
            F .= F_final
        elseif isa($alg, Periodic)
            check = 0
            F_Check = Array{Any, 2}(undef, 3, $alg.acp)
            F_final = Array{Float64, 1}(undef, 2)
            F_Check_inner = Array{Any, 2}(undef, 3, $alg.period)
            for i = 1:$alg.acp
                $alg.fstore(F,F_Check,t,i)
                for j= (i-1)*$alg.period: (i)*$alg.period-1
                    $(loop.args[2])
                end
            end
            F_final .= F
            L .= [0, 1]
            t = 1.0-h
            L_H .= L
            for i = $alg.acp:-1:1
                F,t = $alg.frestore(F_Check,i)
                for j= 1:$alg.period
                    $alg.fstore(F,F_Check_inner,t,j)
                    $(loop.args[2])
                end
                for j= $alg.period:-1:1
                    F,t = $alg.frestore(F_Check_inner,j)
                    L_H .= L
                    F_H .= F
                    res = Checkpointing.jacobian(tobedifferentiated, F_H, $adtool)
                    L =  transpose(res)*L
                    t = t - h
                end
            end
            F .= F_final
        end
    end
    esc(ex)
end


# Interface

###########################################################################################################
##### Methods common to all: AbstractLinearQuadraticModel, AbstractConicModel, AbstractNonlinearModel #####
##### see: http://mathprogbasejl.readthedocs.io/en/latest/solverinterface.html                        #####
###########################################################################################################

getobjgap(m::SCIPMathProgModel) = error("Not implemented for SCIP.jl!")

getrawsolver(m::SCIPMathProgModel) = error("Not implemented for SCIP.jl!")

getsolvetime(m::SCIPMathProgModel) = ccall((:SCIPgetSolvingTime, libcsip), Cdouble, (Ptr{Void},), _getInternalSCIP(m))

getsense(m::SCIPMathProgModel) = error("Not implemented for SCIP.jl!")

numvar(m::SCIPMathProgModel) = _getNumVars(m)

"The number of proper constraints, excluding those from lazy callbacks."
numconstr(m::SCIPMathProgModel) = _getNumConss(m)

freemodel!(m::SCIPMathProgModel) = error("Not implemented for SCIP.jl!")

# TODO: mapping for :SemiCont, :SemiInt
const vartypemap = Dict{Symbol, Cint}(
  :Cont => 3,
  :Bin => 0,
  :Int => 1
)
const revvartypemap = Dict{Cint, Symbol}(
  3 => :Cont,
  0 => :Bin,
  1 => :Int
)

function getvartype(m::SCIPMathProgModel)
    nvars = _getNumVars(m)
    types = zeros(nvars)
    for idx = one(Cint):nvars
        types[idx] = _getVarType(m, idx - one(Cint))
    end
    return map(x -> revvartypemap[x], types)
end

function setvartype!(m::SCIPMathProgModel, vartype::Vector{Symbol})
    nvars = Cint(length(vartype))
    scipvartypes = map(vt -> vartypemap[vt], vartype)
    for idx = one(Cint):nvars
        _chgVarType(m, idx - one(Cint), scipvartypes[idx])
    end
end

optimize!(m::SCIPMathProgModel) = _solve(m)

function status(m::SCIPMathProgModel)
    statusmap = [:Optimal,
                 :Infeasible,
                 :Unbounded,
                 :InfeasibleOrUnbounded,
                 :UserLimit, # node limit
                 :UserLimit, # time limit
                 :UserLimit, # memory limit
                 :UserLimit, # user limit
                 :Unknown    # TODO: find good value
                 ]
    stat = _getStatus(m)
    return statusmap[stat + 1]
end

getobjbound(m::SCIPMathProgModel) = _getObjBound(m)

getobjval(m::SCIPMathProgModel) = _getObjValue(m)

function getsolution(m::SCIPMathProgModel)
    nvars = _getNumVars(m)
    values = zeros(nvars)
    _getVarValues(m, values)
    values
end

function setsense!(m::SCIPMathProgModel, sense)
    if sense == :Max
        _setSenseMaximize(m)
    else
        _setSenseMinimize(m)
    end
end

###########################################################################
##### Methods specific to AbstractLinearQuadraticModel                #####
##### see: http://mathprogbasejl.readthedocs.io/en/latest/lpqcqp.html #####
###########################################################################

loadproblem!(m::SCIPMathProgModel, filename::AbstractString) =
    error("Not implemented for SCIP.jl")

function loadproblem!(m::SCIPMathProgModel, A, varlb, varub, obj, rowlb, rowub, sense)
    # TODO: clean old model?

    nrows, ncols = size(A)
    nvars = Cint(ncols)
    varindices = collect(zero(Cint):nvars - one(Cint))

    for v in 1:ncols
        # TODO: define enum for vartype?
        _addVar(m, float(varlb[v]), float(varub[v]),
                Cint(3), Ptr{Cint}(C_NULL))
    end
    for c in 1:nrows
        # TODO: care about sparse matrices
        denserow = float(collect(A[c, :]))
        _addLinCons(m, nvars, varindices, denserow,
                    float(rowlb[c]), float(rowub[c]), Ptr{Cint}(C_NULL))
    end

    _setObj(m, nvars, varindices, float(obj))

    if sense == :Max
        _setSenseMaximize(m)
    else
        _setSenseMinimize(m)
    end
end

writeproblem(m::SCIPMathProgModel, filename::AbstractString) =
    error("Not implemented for SCIP.jl")

getvarLB(m::SCIPMathProgModel) = error("Not implemented for SCIP.jl")

function setvarLB!(m::SCIPMathProgModel, lb)
    indices = collect(zero(Cint):Cint(numvar(m) - 1))
    _chgVarLB(m, Cint(numvar(m)), indices, float(lb))
end

getvarUB(m::SCIPMathProgModel) = error("Not implemented for SCIP.jl")

function setvarUB!(m::SCIPMathProgModel, ub)
    indices = collect(zero(Cint):Cint(numvar(m) - 1))
    _chgVarUB(m, Cint(numvar(m)), indices, float(ub))
end

getconstrLB(m::SCIPMathProgModel) = error("Not implemented for SCIP.jl")

# setconstrLB!(m::SCIPMathProgModel, lb) = error("Not implemented for SCIP.jl")

getconstrUB(m::SCIPMathProgModel) = error("Not implemented for SCIP.jl")

# setconstrUB!(m::SCIPMathProgModel, ub) = error("Not implemented for SCIP.jl")

getobj(m::SCIPMathProgModel) = error("Not implemented for SCIP.jl")

function setobj!(m::SCIPMathProgModel, obj)
    indices = collect(zero(Cint):Cint(numvar(m) - 1))
    _setObj(m, Cint(numvar(m)), indices, float(obj))
end

getconstrmatrix(m::SCIPMathProgModel) = error("Not implemented for SCIP.jl")

addvar!(m::SCIPMathProgModel, constridx, constrcoef, l, u, objcoef) =
    error("Not implemented for SCIP.jl")

addvar!(m::SCIPMathProgModel, l, u, objcoef) = error("Not implemented for SCIP.jl")

function addconstr!(m::SCIPMathProgModel, varidx, coef, lb, ub)
    _addLinCons(m, Cint(length(varidx)), Vector{Cint}(varidx - 1), float(coef),
                float(lb), float(ub), Ptr{Cint}(C_NULL))
end

numlinconstr(m::SCIPMathProgModel) = error("Not implemented for SCIP.jl")

getconstrsolution(m::SCIPMathProgModel) = error("Not implemented for SCIP.jl")

getreducedcosts(m::SCIPMathProgModel) = error("Not implemented for SCIP.jl")

getconstrduals(m::SCIPMathProgModel) = error("Not implemented for SCIP.jl")

getinfeasibilityray(m::SCIPMathProgModel) = error("Not implemented for SCIP.jl")

getbasis(m::SCIPMathProgModel) = error("Not implemented for SCIP.jl")

getunboundedray(m::SCIPMathProgModel) = error("Not implemented for SCIP.jl")

getsimplexiter(m::SCIPMathProgModel) = error("Not implemented for SCIP.jl")

getbarrieriter(m::SCIPMathProgModel) = error("Not implemented for SCIP.jl")

function setwarmstart!(m::SCIPMathProgModel, v)
    # does not support incomplete solutions (with NaN)
    _setInitialSolution(m, float(v))
end

##########################################################################
##### Methods specific to Integer Programming                        #####
##########################################################################

getnodecount(m::SCIPMathProgModel) = error("Not implemented for SCIP.jl")

function addsos1!(m::SCIPMathProgModel, idx, weight)
    nidx = Cint(length(idx))
    cidx = convert(Vector{Cint}, idx - 1)
    _addSOS1(m, nidx, cidx, weight, Ptr{Cint}(C_NULL))
end

function addsos2!(m::SCIPMathProgModel, idx, weight)
    nidx = Cint(length(idx))
    cidx = convert(Vector{Cint}, idx - 1)
    _addSOS2(m, nidx, cidx, weight, Ptr{Cint}(C_NULL))
end

##########################################################################
##### Methods specific to Quadratic Programming                      #####
##########################################################################

numquadconstr(m::SCIPMathProgModel) = error("Not implemented for SCIP.jl")

setquadobj!{T<:Real}(m::SCIPMathProgModel, Q::Array{T, 2}) =
    error("Not implemented for SCIP.jl")

setquadobj!(m::SCIPMathProgModel, rowidx, colidx, quadval) =
    error("Not implemented for SCIP.jl")

setquadobjterms!(m::SCIPMathProgModel, rowidx, colidx, quadval) =
    _setQuadObj(m, Cint(0), Array{Cint}(0), Array{Cdouble}(0),
                 Cint(length(rowidx)), convert(Vector{Cint}, rowidx - 1),
                 convert(Vector{Cint}, colidx - 1), quadval)

function addquadconstr!(m::SCIPMathProgModel, linearidx, linearval, quadrowidx, quadcolidx, quadval, sense, rhs)
    clhs = -Inf
    crhs =  Inf
    if sense == '<'
        crhs = rhs
    elseif sense == '>'
        clhs = rhs
    else
        @assert sense == '='
        clhs = rhs
        crhs = rhs
    end
    _addQuadCons(m, Cint(length(linearidx)), convert(Vector{Cint}, linearidx - 1),
                 linearval, Cint(length(quadrowidx)), convert(Vector{Cint}, quadrowidx - 1),
                 convert(Vector{Cint}, quadcolidx - 1), quadval, clhs, crhs, Ptr{Cint}(C_NULL))
end

getquadconstrsolution(m::SCIPMathProgModel) = error("Not implemented for SCIP.jl")

getquadconstrduals(m::SCIPMathProgModel) = error("Not implemented for SCIP.jl")

getquadinfeasibilityray(m::SCIPMathProgModel) = error("Not implemented for SCIP.jl")

getquadconstrRHS(m::SCIPMathProgModel) = error("Not implemented for SCIP.jl")

setquadconstrRHS!(m::SCIPMathProgModel, lb) = error("Not implemented for SCIP.jl")

##########################################################################
##### Methods specific to MIP Callbacks                              #####
##########################################################################

# use a different type for heuristic callback and multiple dispatch to implements
# the methods that they share
abstract SCIPCallbackData <: MathProgCallbackData

type SCIPLazyCallbackData <: SCIPCallbackData
    model::SCIPMathProgModel
    csip_lazydata::Ptr{Void}
end

# this is the function that should fit the CSIP_LAZYCALLBACK signature
function lazycb_wrapper(csip_model::Ptr{Void}, csip_lazydata::Ptr{Void},
                        userdata::Ptr{Void})
    # m, f = unsafe_pointer_to_objref(userdata)::(SCIPMathProgModel, Function)
    # WTF: TypeError: typeassert: expected Type{T}, got Tuple{DataType,DataType}
    m, f = unsafe_pointer_to_objref(userdata)
    d = SCIPLazyCallbackData(m, csip_lazydata)
    ret = f(d)
    ret == :Exit && _interrupt(m)

    return convert(Cint, 0) # CSIP_RETCODE_OK
end

function setlazycallback!(m::SCIPMathProgModel, f)
    # f is function(d::SCIPLazyCallbackData)

    cbfunction = cfunction(lazycb_wrapper, Cint, (Ptr{Void}, Ptr{Void}, Ptr{Void}))
    m.lazy_userdata = (m, f)

    _addLazyCallback(m, cbfunction, m.lazy_userdata)
end

# if we are called from a lazy callback, we check whether the LP relaxation is integral
function cbgetstate(d::SCIPLazyCallbackData)
    context = _lazyGetContext(d.csip_lazydata)
    mapping = [:MIPNode, :MIPSol, :Other]
    mapping[context + 1]
end

function cbgetlpsolution(d::SCIPLazyCallbackData, output)
    _lazyGetVarValues(d.csip_lazydata, output)
end

#TODO: what should we do if there is no best solution?
cbgetmipsolution(d::SCIPLazyCallbackData, output) = cbgetlpsolution(d, output)

function cbgetlpsolution(d::SCIPCallbackData)
    output = Array(Float64, _getNumVars(m))
    cbgetlpsolution(d, output)
end

#TODO: what should we do if there is no best solution?
cbgetmipsolution(d::SCIPLazyCallbackData) = cbgetlpsolution(d)

function cbaddlazy!(d::SCIPLazyCallbackData, varidx, varcoef, sense, rhs)
    clhs = -Inf
    crhs =  Inf
    if sense == '<'
        crhs = rhs
    elseif sense == '>'
        clhs = rhs
    else
        @assert sense == '='
        clhs = rhs
        crhs = rhs
    end
    _lazyAddLinCons(d.csip_lazydata, convert(Cint, length(varidx)),
                  convert(Vector{Cint}, varidx - 1), varcoef, clhs, crhs,
                  convert(Cint, 0))
end

# function cbaddlazylocal!(d::SCIPLazyCallbackData, varidx, varcoef, sense, rhs)
# end

type SCIPHeurCallbackData <: SCIPCallbackData
    model::SCIPMathProgModel
    csip_heurdata::Ptr{Void}
    sol::Vector{Float64}
end

# this is the function that should fit the CSIP_HEURCALLBACK signature
function heurcb_wrapper(csip_model::Ptr{Void}, csip_heurdata::Ptr{Void},
                        userdata::Ptr{Void})
    # m, f = unsafe_pointer_to_objref(userdata)::(SCIPMathProgModel, Function)
    # WTF: TypeError: typeassert: expected Type{T}, got Tuple{DataType,DataType}
    m, f = unsafe_pointer_to_objref(userdata)
    d = SCIPHeurCallbackData(m, csip_heurdata, fill(NaN, numvar(m)))

    ret = f(d)
    ret == :Exit && _interrupt(m)

    return convert(Cint, 0) # CSIP_RETCODE_OK
end

function setheuristiccallback!(m::SCIPMathProgModel, f)
    # f is function(d::SCIPHeurCallbackData)

    cbfunction = cfunction(heurcb_wrapper, Cint, (Ptr{Void}, Ptr{Void}, Ptr{Void}))
    m.heur_userdata = (m, f)

    _addHeuristicCallback(m, cbfunction, m.heur_userdata)
end

# TODO: detect :MIPSol like with lazy constraints?
cbgetstate(d::SCIPHeurCallbackData) = :MIPNode

function cbgetlpsolution(d::SCIPHeurCallbackData, output)
    _heurGetVarValues(d.csip_heurdata, output)
end

function cbaddsolution!(d::SCIPHeurCallbackData)
    # check for unspecified values (NaN)
    complete = (findfirst(v -> v === NaN, d.sol) == 0)

    if complete
        # add solution that was filled from cbsetsolutionvalue
        _heurAddSolution(d.csip_heurdata, d.sol)
    else
        warn("Incomplete solutions not supported, skipping candidate.")
    end

    # reset solution vector for the next solution
    d.sol = fill(NaN, numvar(d.model))
end

function cbsetsolutionvalue!(d::SCIPHeurCallbackData, varidx, value)
    d.sol[varidx] = value
end

##########################################################################
##### Methods specific to AbstractConicModel                         #####
##### see: http://mathprogbasejl.readthedocs.io/en/latest/conic.html #####
##########################################################################

########################################################################
##### Methods specific to AbstractNonlinear                        #####
##### see: http://mathprogbasejl.readthedocs.io/en/latest/nlp.html #####
########################################################################
const julia_op_to_csip_op = Dict{Symbol, Int}(
   :- => 9,
   :/ => 11,
   :sqrt => 13,
   :^ => 14,
   :exp => 17,
   :log => 18,
   :+ => 64,
   :* => 65
)

#inspired in ReverseDiffSparse's conversion.jl
immutable CSIPNodeData
    nodetype::Int
    childids::Array{Int}
end

# for backward compatibility with 0.4 (from JuMP/macros.jl)
function comparison_to_call(ex)
    if Meta.isexpr(ex,:comparison) && length(ex.args) == 3
        return Expr(:call, ex.args[2], ex.args[1], ex.args[3])
    else
        return ex
    end
end

# parse constraint expression into something easier for CSIP
function constr_expr_to_nodedata(ex::Expr)
    values = Float64[]
    csipex = CSIPNodeData[]

    # for compatibility with 0.4
    ex = comparison_to_call(ex)

    # a ex.args is [side, comp, expr, comp, side] or [expr, comp, side]
    if length(ex.args) == 5 # two sided constraint
        expr_to_csip(ex.args[3], values, csipex)
    else
        @assert length(ex.args) == 3
        expr_to_csip(ex.args[2], values, csipex)
    end
    return csipex, values
end

# parse objective expression into something easier for CSIP
function obj_expr_to_nodedata(ex::Expr)
    values = Float64[]
    csipex = CSIPNodeData[]
    expr_to_csip(ex, values, csipex)
    return csipex, values
end

# store expression tree as an array of CSIPNodeData
# Each element of a CSIPNodeData consists of the operator identifier
# and an array with the position of the operator's children.
# For each node in the expression tree, first process its children
# (creating an array with their position) and afterwards add create a new
# node with the node's operator and the position of its children.
function expr_to_csip(ex::Expr, values, csipex)
    if Meta.isexpr(ex,:call) # functional operator
        # create node
        childids = []
        op = ex.args[1]
        nchildren = length(ex.args)
        for c in 2:nchildren  # the first actual children is in positin 2
            pos = expr_to_csip(ex.args[c], values, csipex)
            push!(childids, pos)
        end
        push!(csipex, CSIPNodeData(julia_op_to_csip_op[op], childids))

    elseif Meta.isexpr(ex, :ref) # variable
        @assert ex.args[1] == :x
        push!(csipex, CSIPNodeData(1, [ex.args[2]]))

    else # not supported I guess
        error("Unrecognized expression $ex: $(ex.head), probably not supported")
    end
    return length(csipex)
end

# parsing values
function expr_to_csip(ex::Number, values, csipex)
    valueidx = length(values)+1
    push!(values,ex)
    push!(csipex, CSIPNodeData(2, [valueidx]))
    return length(csipex)
end

# the AbstractNLPEvaluator contains the constraints and objective.
# One can get different data from it (with initialize)
# SCIP doesn't need most of the data, just the epression
function loadproblem!(m::SCIPMathProgModel, numVars, numConstr,
                      l, u, lb, ub, sense, d::AbstractNLPEvaluator)
    # add variables
    for v in 1:numVars
        _addVar(m, float(l[v]), float(u[v]), Cint(3), Ptr{Cint}(C_NULL))
    end

    # we want to get the expressions
    initialize(d, [:ExprGraph])

    # add constraints
    for c in 1:numConstr
        csipex, values = constr_expr_to_nodedata(constr_expr(d, c))
        children = Cint[]
        beg = Cint[1]
        ops = Cint[]
        for node in csipex
            ops = [ops; node.nodetype]
            children = [children; node.childids]
            push!(beg, beg[end]+length(node.childids))
        end
        _addNonLinCons(m, Cint(length(ops)),
                       convert(Vector{Cint},ops),
                       convert(Vector{Cint},children - 1),
                       convert(Vector{Cint},beg - 1),
                       values, float(lb[c]), float(ub[c]),
                       Ptr{Cint}(C_NULL))
    end

    # add objective
    csipex, values = obj_expr_to_nodedata(obj_expr(d))
    children = Cint[]
    beg = Cint[1]
    ops = Cint[]
    for node in csipex
        ops = [ops; node.nodetype]
        children = [children; node.childids]
        push!(beg, beg[end]+length(node.childids))
    end
    _setNonlinearObj(m, Cint(length(ops)),
                     convert(Vector{Cint},ops),
                     convert(Vector{Cint},children - 1),
                     convert(Vector{Cint},beg - 1),
                     values)

    # set sense
    if sense == :Max
        _setSenseMaximize(m)
    else
        _setSenseMinimize(m)
    end
end


module Utils_Model
using JuMP
using DataFrames

include(dirname(dirname(@__FILE__))*"\\Model\\R_sep_Model.jl")
include(dirname(dirname(@__FILE__))*"\\Model\\C_sep_Model.jl")
using .R_sep_Model
using .C_sep_Model

export export_model,
       batch_model,
       read_model

""" Exmports the JuMP model to SDPA format """
export_model(model,file_name::String = pwd()*"\\Defualt_sep_mod.dat-s") =
    JuMP.write_to_file(model, file_name, format=MOI.FileFormats.FORMAT_SDPA)

""" Reads a JuMP model from SDPA format text document """
function read_model(file_name::String)
    @assert splitext(file_name)[end] == ".dat-s" "Can only load models from .dat-s (SDPA) files!"
    return JuMP.read_from_file(file_name, format=MOI.FileFormats.FORMAT_SDPA)
end

end


struct AudioLocation
  azimuth::Int # ∈ -179:180
  elevation::Int # ∈ -90:90
end


using Test
using QuantumLattices.Prerequisites.VectorSpaces
using QuantumLattices.Interfaces: dimension, rank

import QuantumLattices.Prerequisites.VectorSpaces: shape, ndimshape
import QuantumLattices.Prerequisites.Traits: contentnames, getcontent

struct SimpleVectorSpace{B, N} <: EnumerativeVectorSpace{B}
    sorted::Bool
    table::NTuple{N, B}
    SimpleVectorSpace(sorted::Bool, table::NTuple{N, B}) where {B, N} = new{B, N}(sorted, table)
end
@inline contentnames(::Type{<:SimpleVectorSpace}) = (:sorted, :table)
@inline Base.issorted(vs::SimpleVectorSpace) = vs.sorted

@testset "EnumerativeVectorSpace" begin
    @test contentnames(EnumerativeVectorSpace) == (:table,)

    id0, id4 = (1, 0), (1, 4)
    id1, id2, id3 = (1, 1), (1, 2), (1, 3)
    vs = SimpleVectorSpace(true, (id1, id2, id3))
    @test vs == deepcopy(vs)
    @test isequal(vs, deepcopy(vs))
    @test vs|>size == (3,)
    @test vs|>dimension == 3
    @test vs|>collect == [id1, id2, id3]
    @test vs[1]==id1 && vs[2]==id2 && vs[3]==id3
    @test searchsortedfirst(vs, id0)==1 && searchsortedfirst(vs, id4)==4
    @test searchsortedfirst(vs, id1)==1 && searchsortedfirst(vs, id2)==2 && searchsortedfirst(vs, id3)==3
    @test isnothing(findfirst(id0, vs)) && isnothing(findfirst(id4, vs))
    @test findfirst(id1, vs)==1 && findfirst(id2, vs)==2 && findfirst(id3, vs)==3
    @test (id0∉vs) && (id4 ∉ vs)
    @test (id1∈vs) && (id2∈vs) && (id3 ∈ vs)

    vs = SimpleVectorSpace(false, (id1, id2, id3))
    @test searchsortedfirst(vs, id0)==4 && searchsortedfirst(vs, id4)==4
    @test searchsortedfirst(vs, id1)==1 && searchsortedfirst(vs, id2)==2 && searchsortedfirst(vs, id3)==3
    @test isnothing(findfirst(id0, vs)) && isnothing(findfirst(id4, vs))
    @test findfirst(id1, vs)==1 && findfirst(id2, vs)==2 && findfirst(id3, vs)==3
end

struct SimpleIndices{N} <: CartesianVectorSpace{CartesianIndex{N}}
    shape::NTuple{N, UnitRange{Int}}
    SimpleIndices(shape::NTuple{N, UnitRange{Int}}) where N = new{N}(shape)
end
SimpleIndices(shape::UnitRange{Int}...) = SimpleIndices(shape)
@inline shape(vs::SimpleIndices) = vs.shape
@inline ndimshape(::Type{<:SimpleIndices{N}}) where N = N
@inline Base.CartesianIndex(basis::CartesianIndex{N}, ::SimpleIndices{N}) where N = basis

@testset "CartesianVectorSpace" begin
    foi = SimpleIndices(2:3, 2:3, 2:3)
    @test dimension(foi) == 8
    @test ndimshape(foi) == ndimshape(typeof(foi)) == 3
    @test issorted(foi) == true
    @test foi|>collect == CartesianIndex.([(2, 2, 2), (3, 2, 2), (2, 3, 2), (3, 3, 2), (2, 2, 3), (3, 2, 3), (2, 3, 3), (3, 3, 3)])
    for (i, index) in enumerate(CartesianIndices((2:3, 2:3, 2:3)))
        @test foi[i] == foi[index] == index
        @test findfirst(index, foi) == i
        @test searchsortedfirst(foi, index) == i
        @test index∈foi
    end
    i1, i2 = CartesianIndex(1, 1, 1), CartesianIndex(4, 4, 4)
    @test isnothing(findfirst(i1, foi)) && isnothing(findfirst(i2, foi))
    @test searchsortedfirst(foi, i1) == 1 && searchsortedfirst(foi, i2) == 9
    @test i1∉foi && i2∉foi
end

@testset "NamedVectorSpace" begin
    nvs = ZipNamedVectorSpace{(:t, :U)}([1, 2], [8.0, 9.0])
    @test nvs==deepcopy(nvs) && isequal(nvs, deepcopy(nvs))
    @test nvs≠ZipNamedVectorSpace{(:t, :V)}([1, 2], [8.0, 9.0])
    @test !isequal(nvs, ZipNamedVectorSpace{(:t, :V)}([1, 2], [8.0, 9.0]))
    @test nvs|>keys == nvs|>typeof|>keys == (:t, :U)
    @test nvs|>values == ([1, 2], [8.0, 9.0])
    @test nvs|>pairs|>collect == [:t=>[1, 2], :U=>[8.0, 9.0]]
    @test eltype(nvs, 1) == eltype(nvs|>typeof, 1) == Int
    @test eltype(nvs, 2) == eltype(nvs|>typeof, 2) == Float64
    @test nvs|>rank == nvs|>typeof|>rank == 1
    @test shape(nvs) == (1:2,)
    @test ndimshape(nvs) == ndimshape(typeof(nvs)) == 1
    elements = [(t = 1, U = 8.0), (t = 2, U = 9.0)]
    for i = 1:dimension(nvs)
        @test NamedTuple(CartesianIndex(elements[i], nvs), nvs) == elements[i]
    end
    @test nvs|>collect == elements

    nvs = ProductNamedVectorSpace{(:t, :U)}([1.0, 2.0], [8.0, 9.0])
    @test nvs|>rank == nvs|>typeof|>rank == 2
    @test shape(nvs) == (1:2, 1:2)
    @test ndimshape(nvs) == ndimshape(typeof(nvs)) == 2
    elements = [(t = 1.0, U = 8.0), (t = 2.0, U = 8.0), (t = 1.0, U = 9.0), (t = 2.0, U = 9.0)]
    for i = 1:dimension(nvs)
        @test NamedTuple(CartesianIndex(elements[i], nvs), nvs) == elements[i]
    end
    @test nvs|>collect == elements
end


# Samplers for Categorical distribution

abstract AbstractCategoricalSampler

include("drawtable.jl")
include("alias.jl")
include("huffman.jl")


using RecipesBase
using KernelDensity
using Random
using PlotUtils
using Colors

#=
!!!! IMPORTANT !!!!
Each code path of the @series macro must return the data for the series.
data = ... should be the last statement
=#
@recipe function abc_output_recipe(abco::ABCOutput;
      runs=nothing, params=nothing, population_colors=nothing)
      if params === nothing
            params = [i for i in 1:abco.n_params]
      end
      if runs === nothing
            runs = [i for i in 1:length(abco.population)]
      elseif isa(runs, Int)
            runs = [i for i in max(length(abco.population) - runs + 1, 1):length(abco.population)]
      end
      legend --> false
      layout := length(params) ^ 2
      for (i, par1) in enumerate(params)
            for (j, par2) in enumerate(params)
                  subplot := (i - 1) * length(params) + j
                  if i == j
                        @series begin
                              seriestype := :histogram
                              ylabel := "Accepted"
                              xlabel := "Parameter $(par1)"
                              bins --> 50
                              if isa(abco.population, Vector)
                                    data = abco.population[end][:, par1]
                              else
                                    data = abco.population[:, par1]
                              end # if Vector
                        end  # @series
                  elseif j < i
                        seriestype := :scatter
                        markerstrokecolor --> false
                        if isa(abco.population, Vector)
                              for r in runs
                                    @series begin
                                          xlabel --> "Parameter $(par2)"
                                          ylabel --> "Parameter $(par1)"
                                          if population_colors !== nothing
                                                idx = r % length(population_colors)
                                                if idx == 0
                                                      idx = length(population_colors)
                                                end
                                                seriescolor := population_colors[idx]
                                          end
                                          pop = abco.population[r]
                                          data = (pop[:, par2], pop[:, par1])
                                    end # @series
                              end # for r
                        else
                              @series begin
                                    data = (abco.population[:, par2], abco.population[:, par1])
                              end # @series
                        end # if Vector
                  else
                        @series begin
                              ylabel := ""
                              xlabel := ""
                              grid := false
                              xaxis := false
                              yaxis := false
                              data = []
                        end # @series
                  end # if/elseif
            end # for j, par2
      end # for i, par1
end # @recipe

function scale_outer_intervals(intervals, scale=1.0)
    int_min = min([int[1] for int in intervals]...)
    int_max = max([int[2] for int in intervals]...)
    int_scaled_half = (int_max - int_min) * scale / 2.0
    int_mid = (int_min + int_max) / 2.0
    int_mid - int_scaled_half, int_mid + int_scaled_half
end

my_linspace(bounds::Tuple{Float64, Float64}, length::Int64) = range(bounds[1], stop=bounds[2], length=length)

@recipe function emulation_vs_simulation(emu_out::EmulatedABCSMCOutput, sim_out::SimulatedABCSMCOutput, true_params::AbstractVector{<:Real}=Float64[];
            pop_idx=0,
            pop_color = "#007731",
            simulation_color = "#08519c",
            emulation_color = "#800013",
            contour_levels=8)
      if pop_idx == 0
            pop_idx = min(length(sim_out.population), length(emu_out.population))
      end
      kernel_bandwidth_scale = 0.09
      bounds_scale = 1.2
      sim_count = 30
      params = 1:emu_out.n_params
      legend --> false
      layout := length(params) ^ 2
      for i in params
            for j in params
                  subplot := (i - 1) * length(params) + j
                  if j < i
                        x_data_emu = emu_out.population[pop_idx][:,j]
                        y_data_emu = emu_out.population[pop_idx][:,i]
                        x_data_sim = sim_out.population[pop_idx][:,j]
                        y_data_sim = sim_out.population[pop_idx][:,i]
                        sim_size = size(sim_out.population[pop_idx], 1)
                        if sim_size > sim_count
                            idx = randperm(sim_size)[1:sim_count]
                            x_data_sim = x_data_sim[idx]
                            y_data_sim = y_data_sim[idx]
                        end
                        x_extr_emu = extrema(x_data_emu)
                        y_extr_emu = extrema(y_data_emu)
                        x_extr_sim = extrema(x_data_sim)
                        y_extr_sim = extrema(y_data_sim)

                        bandwidth = (-kernel_bandwidth_scale * -(x_extr_emu...), -kernel_bandwidth_scale * -(y_extr_emu...))
                        kde_joint = kde((x_data_emu, y_data_emu), bandwidth=bandwidth)
                        contour_x = my_linspace(scale_outer_intervals([x_extr_emu], bounds_scale), 100)
                        contour_y = my_linspace(scale_outer_intervals([y_extr_emu], bounds_scale), 100)
                        contour_z = pdf(kde_joint, contour_x, contour_y)
                        @series begin
                              seriestype := :contour
                              color := ColorGradient([:white, emulation_color])
                              levels := contour_levels
                              fill := true
                              (contour_x, contour_y, contour_z)
                        end
                        @series begin
                              seriestype := :scatter
                              markersize := 4
                              markershape := :x
                              color := simulation_color
                              (x_data_sim, y_data_sim)
                        end
                        if length(true_params) > 0
                              seriestype := :path
                              linestyle := :dash
                              color := :black
                              markershape := :none
                              @series begin
                                    ([true_params[j], true_params[j], minimum(contour_x)],
                                    [minimum(contour_y), true_params[i], true_params[i]])
                              end
                        end
                  elseif j > i
                        pop_colors = range(parse(Colorant, "white"), parse(Colorant, pop_color), length=pop_idx+1)
                        seriestype := :scatter
                        markershape := :circle
                        markersize := 2
                        for k in 1:pop_idx
                              @series begin
                                    markercolor := pop_colors[k+1]
                                    markerstrokecolor := pop_colors[k+1]
                                    (emu_out.population[k][:,j], emu_out.population[k][:,i])
                              end
                        end
                        if length(true_params) > 0
                              seriestype := :path
                              linestyle := :dash
                              color := :black
                              markershape := :none
                              @series begin
                                    ([true_params[j], true_params[j], minimum(emu_out.population[1][:,j])],
                                    [minimum(emu_out.population[1][:,i]), true_params[i], true_params[i]])
                              end
                        end
                  else # i == j
                        seriestype := :path
                        markershape := :none
                        emu_data = emu_out.population[pop_idx][:,i]
                        sim_data = sim_out.population[pop_idx][:,i]
                        extr_emu = extrema(emu_data)
                        extr_sim = extrema(sim_data)
                        kde_emu = kde(emu_data, bandwidth=-kernel_bandwidth_scale * -(extr_emu...))
                        kde_sim = kde(sim_data, bandwidth=-kernel_bandwidth_scale * -(extr_sim...))
                        x_bounds = scale_outer_intervals([extr_emu, extr_sim], bounds_scale)
                        x_plot = my_linspace(x_bounds, 100)
                        y_emu_plot = pdf(kde_emu,x_plot)
                        y_sim_plot = pdf(kde_sim,x_plot)
                        @series begin
                              linestyle := :solid
                              color := simulation_color
                              (x_plot, y_sim_plot)
                        end
                        @series begin
                              linestyle := :solid
                              color := emulation_color
                              (x_plot, y_emu_plot)
                        end
                        if length(true_params) > 0
                              true_param = true_params[i]
                              max_pdf = max(pdf(kde_emu, true_param), pdf(kde_sim, true_param))
                              @series begin
                                    linestyle := :dash
                                    color := :black
                                    ([true_param, true_param], [0, max_pdf])
                              end
                        end
                  end # if/else/elseif i/j
            end # for j
      end # for i
end

# Plot recipe for mdoel selection output
@recipe function modelselection_plotrecipe(::Type{ModelSelectionOutput}, mso::ModelSelectionOutput)
    seriestype := :line
    xlabel --> "Population"
    ylabel --> "Number of accepted particles"
    labels --> [string("Model ", m) for m in 1:mso.M]
    data = [[mso.n_accepted[i][j] for i in 1:size(mso.n_accepted,1)] for j in 1:mso.M]
end


Pkg.update()
Pkg.add("JuMP")
Pkg.add("Ipopt")
Pkg.add("JSON")
Pkg.add("ArgParse")


@testset "Split Sets" begin
    @testset "split_cols" begin
        cols = ["A", "B", "C", "D"]
        target = "A"
        exclude_target = exclude_elem(cols, target)
        @test cols == ["A", "B", "C", "D"]
        @test exclude_target == ["B", "C", "D"]

        target = "C"
        exclude_target2 = exclude_elem(cols, target)
        @test exclude_target2 == ["A", "B", "D"]
    end
end

@testset "Date Validation" begin
    @testset "Invalid date range" begin
        # one month
        _year = 2018
        _month = 1
        from_date = ZonedDateTime(_year, _month, 1, 0, tz"Asia/Seoul")
        to_date = ZonedDateTime(_year, _month, Dates.daysinmonth(DateTime(_year, _month)), 23, tz"Asia/Seoul")
        dates = collect(from_date:Dates.Hour(1):to_date)
        df = DataFrame(date = dates)

        @test_throws ArgumentError validate_dates(to_date, from_date, 1, df)
    end

    @testset "Invalid window size" begin
        # one month
        _year = 2018
        _month = 1
        from_date = ZonedDateTime(_year, _month, 1, 0, tz"Asia/Seoul")
        to_date = ZonedDateTime(_year, _month, 1, 23, tz"Asia/Seoul")
        dates = collect(from_date:Dates.Hour(1):to_date)
        df = DataFrame(date = dates)

        @test_throws BoundsError validate_dates(from_date, to_date, 48, df)
    end
end

@testset "Window" begin
    @testset "Basic window" begin
        df = DataFrame(
            A = 1:12,
            B = 13:24,
            C = 25:36)
        sample_size = 4
        output_size = 2
        dfs = window_df(df, sample_size, output_size)
        @test length(dfs) == 7
        @test dfs[1] == DataFrame(
            A = 1:6,
            B = 13:18,
            C = 25:30)
        @test dfs[7] == DataFrame(
            A = 7:12,
            B = 19:24,
            C = 31:36)
    end

    @testset "Window + Date Range" begin
        # one month
        _year = 2018
        _month = 1
        from_date = ZonedDateTime(_year, _month, 1, 0, tz"Asia/Seoul")
        to_date = ZonedDateTime(_year, _month, Dates.daysinmonth(DateTime(_year, _month)), 23, tz"Asia/Seoul")
        dates = collect(from_date:Dates.Hour(1):to_date)
        n = length(dates)

        df = DataFrame(
            A = 1:n,
            B = n+1:2n,
            C = 2n+1:3n,
            date = dates
        )
        sample_size = 12
        output_size = 2
        dfs = window_df(df, sample_size, output_size,
            ZonedDateTime(_year, _month, 10, 0, 1, tz"Asia/Seoul"),
            ZonedDateTime(_year, _month, 14, 23, 59, tz"Asia/Seoul"))

        @test length(dfs) == 106
        @test dfs[1] == DataFrame(
            A = 218:231,
            B = 962:975,
            C = 1706:1719,
            date = collect(ZonedDateTime(_year, _month, 10, 1, tz"Asia/Seoul"):Hour(1):ZonedDateTime(_year, _month, 10, 14, tz"Asia/Seoul"))
        )
        @test dfs[end] == DataFrame(
            A = 323:336,
            B = 1067:1080,
            C = 1811:1824,
            date = collect(ZonedDateTime(_year, _month, 14, 10, tz"Asia/Seoul"):Hour(1):ZonedDateTime(_year, _month, 14, 23, tz"Asia/Seoul"))
        )
    end

    @testset "Window + Date Range (Last day)" begin
        # one month
        _year = 2018
        _month = 1
        from_date = ZonedDateTime(_year, _month, 1, 0, tz"Asia/Seoul")
        to_date = ZonedDateTime(_year, _month, Dates.daysinmonth(DateTime(_year, _month)), 23, tz"Asia/Seoul")
        dates = collect(from_date:Dates.Hour(1):to_date)
        n = length(dates)

        df = DataFrame(
            A = 1:n,
            B = n+1:2n,
            C = 2n+1:3n,
            date = dates
        )

        sample_size = 12
        output_size = 2
        window_size = sample_size + output_size
        dfs = window_df(df, sample_size, output_size,
            ZonedDateTime(_year, _month, 30, 0, tz"Asia/Seoul"),
            ZonedDateTime(_year, _month, 31, 23, tz"Asia/Seoul"))

        @test length(dfs) == 35
        @test dfs[35] == DataFrame(
            A = 731:744,
            B = 1475:1488,
            C = 2219:2232,
            date = collect(ZonedDateTime(_year, _month, 31, 10, tz"Asia/Seoul"):Hour(1):ZonedDateTime(_year, _month, 31, 23, tz"Asia/Seoul"))
        )
    end

    @testset "Window + Date Range + StationCode" begin
        # one month
        _year = 2018
        _month = 1
        from_date = ZonedDateTime(_year, _month, 1, 0, tz"Asia/Seoul")
        to_date = ZonedDateTime(_year, _month, Dates.daysinmonth(DateTime(_year, _month)), 23, tz"Asia/Seoul")
        dates = collect(from_date:Dates.Hour(1):to_date)
        n = length(dates)
        baseCode = 1000
        tot_stn = 10

        listA = 1:n
        listB = n+1:2n
        listC = 2n+1:3n

        # stationCode 1 1 1 1 2 2 2 2 3 3 3 3 ...
        # A, B, C, date 1 2 3 4 1 2 3 4 1 2 3 4 ...
        df = DataFrame(
            stationCode = repeat(range(baseCode, step=1, length=tot_stn), inner=n),
            A = repeat(listA, outer=tot_stn),
            B = repeat(listB, outer=tot_stn),
            C = repeat(listC, outer=tot_stn),
            date = repeat(dates, outer=tot_stn))

        sample_size = 12
        output_size = 2
        window_size = sample_size + output_size

        target_stncode = baseCode + abs(rand(Int)) % tot_stn
        dfs = window_df(df, sample_size, output_size,
            ZonedDateTime(_year, _month, 30, 0, tz"Asia/Seoul"),
            ZonedDateTime(_year, _month, 31, 23, tz"Asia/Seoul"),
            target_stncode)

        df = dfs[1]
        @test DataFrames.nrow(df) == window_size
        @test df[!, :stationCode] == ones(window_size) * target_stncode
    end
end

@testset "Splitting Sizes" begin
    @testset "split_sizes3_mod0" begin
        total_size = 100
        batch_size = 10
        train_size, valid_size, test_size = split_sizes3(total_size, batch_size)
        @test train_size == 64
        @test valid_size == 16
        @test test_size == 20
    end

    @testset "split_sizes3_prime" begin
        total_size = 97
        batch_size = 10
        train_size, valid_size, test_size = split_sizes3(total_size, batch_size)
        @test train_size == 62
        @test valid_size == 16
        @test test_size == 19
    end

    @testset "split_sizes2_mod0" begin
        total_size = 100
        batch_size = 10
        train_size, valid_size = split_sizes2(total_size, batch_size)
        @test train_size == 80
        @test valid_size == 20
    end

    @testset "split_sizes2_prime" begin
        total_size = 97
        batch_size = 10
        train_size, valid_size = split_sizes2(total_size, batch_size)
        @test train_size == 78
        @test valid_size == 19
    end
end

@testset "Input" begin
    @testset "getX & getY" begin
        sample_size = 4
        output_size = 3
        window_size = sample_size + output_size

        start = 10
        lists = collect(Iterators.partition((start+1):(start + 3*window_size), window_size))

        df = DataFrame(
            A = lists[1],
            B = lists[2],
            C = lists[3])
        features = [:A, :B]

        X = getX(df, features, sample_size)
        @test X == Matrix(hcat([11, 12, 13, 14], [18, 19, 20, 21]))

        Y = getY(df, :C, sample_size)
        @test Y == Vector([29, 30, 31])
    end
end

@testset "DNN Input" begin
    @testset "DNN Pair construction" begin
        sample_size = 24
        output_size = 12

        _year = 2015
        _month = 1
        from_date = ZonedDateTime(_year, _month, 1, 0, tz"Asia/Seoul")
        to_date = ZonedDateTime(_year, _month, Dates.daysinmonth(DateTime(_year, _month)), 23, tz"Asia/Seoul")
        dates = from_date:Dates.Hour(1):to_date
        len_dates = length(dates)
        lists = collect(Iterators.partition(1:3*len_dates, len_dates))
        eltype = Integer

        df = DataFrame(
            date = dates,
            A = lists[1],
            B = lists[2],
            C = lists[3])

        features = [:A, :B]
        ycol = :C

        wd = window_df(df, sample_size, output_size, from_date, to_date)
        # pair = ([1,...,24,len_df,...,len_df+24],[2*len_df,...,2*len_df+12])
        pair = make_pair_DNN(wd[1], ycol, features, sample_size, output_size, eltype)

        # X length
        @test length(pair[1]) == sample_size * length(features)
        # Y length
        @test length(pair[2]) == output_size

        @test pair[1] == reduce(vcat, [collect(1:sample_size), collect(len_dates + 1:len_dates + sample_size)])
        @test pair[2] == collect((sample_size + 2*len_dates + 1):(sample_size + 2*len_dates + output_size))
    end

    @testset "DNN Batch(size 2) construction" begin
        sample_size = 24
        output_size = 12
        batch_size = 2

        _year = 2015
        _month = 1
        from_date = ZonedDateTime(_year, _month, 1, 0, tz"Asia/Seoul")
        to_date = ZonedDateTime(_year, _month, Dates.daysinmonth(DateTime(_year, _month)), 23, tz"Asia/Seoul")
        dates = from_date:Dates.Hour(1):to_date
        len_dates = length(dates)
        lists = collect(Iterators.partition(1:3*len_dates, len_dates))
        eltype = Integer

        df = DataFrame(
            date = dates,
            A = lists[1],
            B = lists[2],
            C = lists[3])

        features = [:A, :B]
        ycol = :C

        # type contraint to Floating point
        for col in [:A, :B, :C]
            df[!, col] = Float32.(df[!, col])
        end

        wd = window_df(df, sample_size, output_size, from_date, to_date)

        batch = make_batch_DNN(wd[1:batch_size], :C, [:A, :B], sample_size, output_size, batch_size,
            0.5, eltype)

        # indices for comparison
        in_chnks = [1:sample_size,
                    2:(sample_size+1)]
        out_chnks = [(sample_size+1):(sample_size+output_size),
                     (sample_size+2):(sample_size+output_size+1)]

        # Test X
        @test batch[1] == hcat(
            vcat(df[in_chnks[1], :A], df[in_chnks[1], :B]),
            vcat(df[in_chnks[2], :A], df[in_chnks[2], :B]))
        # Test Y
        @test batch[2] == hcat(
            vcat(df[out_chnks[1], :C]),
            vcat(df[out_chnks[2], :C]))

        @test size(batch[1]) == (sample_size * 2, batch_size)
        @test size(batch[2]) == (output_size, batch_size)
    end

    @testset "DNN Batch(size 3) construction" begin
        sample_size = 24
        output_size = 12
        batch_size = 3

        _year = 2015
        _month = 1
        from_date = ZonedDateTime(_year, _month, 1, 0, tz"Asia/Seoul")
        to_date = ZonedDateTime(_year, _month, Dates.daysinmonth(DateTime(_year, _month)), 23, tz"Asia/Seoul")
        dates = from_date:Dates.Hour(1):to_date
        len_dates = length(dates)
        lists = collect(Iterators.partition(1:3*len_dates, len_dates))
        eltype = Integer

        df = DataFrame(
            date = dates,
            A = lists[1],
            B = lists[2],
            C = lists[3])

        features = [:A, :B]
        ycol = :C

        # type contraint to Floating point
        for col in [:A, :B, :C]
            df[!, col] = Float32.(df[!, col])
        end

        wd = window_df(df, sample_size, output_size, from_date, to_date)

        batch = make_batch_DNN(wd[1:batch_size], :C, [:A, :B], sample_size, output_size, batch_size,
            0.5, eltype)

        # indices for comparison
        in_chnks = [1:sample_size,
                    2:(sample_size+1),
                    3:(sample_size+2)]
        out_chnks = [(sample_size+1):(sample_size+output_size),
                     (sample_size+2):(sample_size+output_size+1),
                     (sample_size+3):(sample_size+output_size+2)]
        # Test X
        @test batch[1] == hcat(
            vcat(df[in_chnks[1], :A], df[in_chnks[1], :B]),
            vcat(df[in_chnks[2], :A], df[in_chnks[2], :B]),
            vcat(df[in_chnks[3], :A], df[in_chnks[3], :B]))
        @test batch[2] == hcat(
            vcat(df[out_chnks[1], :C]),
            vcat(df[out_chnks[2], :C]),
            vcat(df[out_chnks[3], :C]))
        @test size(batch[1]) == (sample_size * 2, 3)
        @test size(batch[2]) == (output_size, 3)
    end
end

@testset "Validate Pairs" begin
    @testset "validate_pairs_normal" begin
        X = Matrix([ 1  2  3  4;
                     7  8  9 10;
                    13 14 15 16;
                    19 20 21 22;
                    25 26 27 28])
        Y = Matrix([5  6  7  8  9 10])

        remove_sparse_input!(X, Y, 0.5)

        @test X == Matrix([ 1  2  3  4;
                     7  8  9 10;
                    13 14 15 16;
                    19 20 21 22;
                    25 26 27 28])
        @test Y == Matrix([5  6  7  8  9 10])
    end

    @testset "validate_pairs_missings" begin
        X = Matrix([ 1  2  3  4;
                     7  8  9 10;
                    13 14 15 16;
                    19 20 21 22;
                    25 26 27 28])
        Y = Matrix([23 24 25 missing missing missing])

        remove_sparse_input!(X, Y, 0.5)

        @test X == Matrix([ 0.0 0.0 0.0 0.0;
                            0.0 0.0 0.0 0.0;
                            0.0 0.0 0.0 0.0;
                            0.0 0.0 0.0 0.0;
                            0.0 0.0 0.0 0.0;])
        @test Y == Matrix([0.0 0.0 0.0 0.0 0.0 0.0])
    end

    @testset "validate_pairs_missings_0.3" begin
        X = Matrix([ 1  2  3  4;
                     7  8  9 10;
                    13 14 15 16;
                    19 20 21 22;
                    25 26 27 28])
        Y = Matrix([23 24 25 26 missing missing])

        remove_sparse_input!(X, Y, 0.3)

        @test X == Matrix([ 0.0 0.0 0.0 0.0;
                            0.0 0.0 0.0 0.0;
                            0.0 0.0 0.0 0.0;
                            0.0 0.0 0.0 0.0;
                            0.0 0.0 0.0 0.0;])
        @test Y == Matrix([0.0 0.0 0.0 0.0 0.0 0.0])
    end

    @testset "validate_pairs_mx" begin
        X = Matrix([ 1  2  3  4;
                     7  8  9 10;
                    13 14 15 16;
                    19 20 21 22;
                    25 26 27 28])
        Y = Matrix([23 24 25 0.0 missing missing])

        remove_sparse_input!(X, Y, 0.5)

        @test X == Matrix([ 0.0 0.0 0.0 0.0;
                            0.0 0.0 0.0 0.0;
                            0.0 0.0 0.0 0.0;
                            0.0 0.0 0.0 0.0;
                            0.0 0.0 0.0 0.0;])
        @test Y == Matrix([0.0 0.0 0.0 0.0 0.0 0.0])
    end
end

@testset "CNN to RNN Batch Input Conversion" begin
    @testset "Batch Serialization" begin
        m = 3
        n = 5
        a = reshape(1:m*n, m, n)
        @test size(a) == (m, n)

        # 2D -> Arrays of 1D
        b = serializeBatch(a)
        @test size(b) == (n,)
        @test size(collect(b[1])) == (m,)
    end

    @testset "Batch Serialization" begin
        m1 = 3
        m2 = 4
        n = 5
        X = reshape(1:m1*n, m1, n)
        Y = reshape(101:(100+m2*n), m2, n)
        Z = serializeBatch(X, Y)

        @test size(Z) == (n,)
        @test size(collect(Z[1])) == (2,)
        # row size should be maintained
        @test size(Z[1][1]) == (m1,)
        @test size(Z[1][2]) == (m2,)
    end

    @testset "Unpack Sequence 3" begin
        x = reshape(vec(1:24), (2, 3, 4))

        @test unpack_seq(x) == reshape([
            [1 7 13 19; 3 9 15 21; 5 11 17 23],
            [2 8 14 20; 4 10 16 22; 6 12 18 24]], 1, 1, 2)
    end

    @testset "Unpack Sequence 4" begin
        x = reshape(vec(1:24), (2, 1, 3, 4))

        @test unpack_seq(x) == reshape([
            [1 7 13 19; 3 9 15 21; 5 11 17 23],
            [2 8 14 20; 4 10 16 22; 6 12 18 24]], 1, 1, 2, 1)
    end
end


libtorch = pyimport_conda("torch", "pytorch")
py"""
import torch
import random

# test
# 0 - 4 array-dim
# strides
# reshape
# offset
# slice
# mutation
#
# combine

def random_slice(x):
  shape = x.shape
  def randrange(s):
    s1 = random.randint(0, s-1)
    s2 = random.randint(s1+1, s)
    return (s1, s2)
  slices = list(map(randrange, shape))
  if len(shape) == 1:
    s1, s2 = slices[0]
    y = x[s1:s2]
  elif len(shape) == 2:
    s11, s12 = slices[0]
    s21, s22 = slices[1]
    y = x[s11:s12, s21:s22]
  elif len(shape) == 3:
    s11, s12 = slices[0]
    s21, s22 = slices[1]
    s31, s32 = slices[2]
    y = x[s11:s12, s21:s22, s31:s32]
  return y

def random_stride(x):
  shape = x.shape
  def randrange(s):
    s1 = random.randint(0, s//2-1)
    s2 = random.randint(s//2+2, s-1)
    si = random.randint(2, 3)
    return (s1, si, s2)
  strides = list(map(randrange, shape))
  if len(shape) == 1:
    s1, si, s2 = strides[0]
    y = x[s1:s2:si]
  elif len(shape) == 2:
    s11, si1, s12 = strides[0]
    s21, si2, s22 = strides[1]
    y = x[s11:s12:si1, s21:s22:si2]
  elif len(shape) == 3:
    s11, si1, s12 = strides[0]
    s21, si2, s22 = strides[1]
    s31, si3, s32 = strides[2]
    y = x[s11:s12:si1, s21:s22:si2, s31:s32:si3]
  return y

def pyslice(x, i):
  return x[i:]

"""
thload = libtorch.load
thsave = libtorch.save
random_slice = py"random_slice"
random_stride = py"random_stride"
pyslice = py"pyslice"


@testset "get_vertexval(g, v, ::Symbol)" begin

    g1 = ValGraph{Int8}(2, vertexval_types=(a=Int, b=Bool), vertexval_init=v -> (v, iseven(v)))
    g2 = ValOutDiGraph{Int16}(2, vertexval_types=(a=Int, b=Bool), vertexval_init=v -> (v, iseven(v)))
    g3 = ValDiGraph{UInt8}(2, vertexval_types=(a=Int, b=Bool), vertexval_init=v -> (v, iseven(v)))

    @test get_vertexval(g1, 1, :a) == 1
    @test get_vertexval(g1, Int8(2), :b) == true

    @test get_vertexval(g2, UInt8(1), :b) == false
    @test get_vertexval(g2, 2, :a) == 2

    @test get_vertexval(g3, 1, :b) == false
    @test get_vertexval(g3, 2, :a) == 2
end