using Bingomatic
using Test

@testset "Bingomatic.jl" begin
    @test_throws ErrorException sample_words(["a","b"], 2)
    @test_throws ErrorException sample_words(["a", "b", "c"], 4)
    n_rows = 3
    word_list = ["a","b","c","d","e","f","g","h","i"]
    words = sample_words(word_list, n_rows, n_rows)
    @test size(words) == (n_rows,n_rows)
    @test words[5] == "Free Space"
    card = make_card(words)
end


@testset "Conjugated orbitals" begin
    o = o"1s"
    co = conj(o)
    @test co == Conjugate(o)
    @test string(co) == "1s†"
end


platform = load_platform("$HetaSimulatorDir/test/examples/single_comp", rm_out = false);
model = platform.models[:nameless];
@test test_show(platform)
@test test_show(model)

@test isfile("$HetaSimulatorDir/test/examples/single_comp/_julia/model.jl")

# Wrong input tests
## no saveat or tspan
@test_throws ErrorException("Please, provide either `saveat` or `tspan` value.") Scenario(model)
## wrong observables
@test_throws ErrorException("The following observables have not been found in the model: [:GG]") Scenario(model; observables=[:r1, :GG], tspan=(0,120))

# Simulate default scenario
s = Scenario(model, tspan = (0., 200.), saveat = [0.0, 150., 250.], observables=[:r1]) |> sim
@test test_show(s)
@test typeof(s) <: HetaSimulator.SimResult
@test status(s) == :Success
@test times(s)[end] == 250.
@test vals(s)[1][:r1] == 0.1

# tests related to saveat, tspan behavior
s1 = Scenario(model, tspan = (0., 200.)) |> sim
s2 = Scenario(model, saveat = [150., 200.]) |> sim
@test times(s1)[1] == 0.0
@test times(s1)[end] == 200.0
@test times(s2) == [150., 200.]

# Simulate scenario
scn1 = Scenario(model; parameters = [:k1=>0.02], tspan = (0., 200.), saveat = [0.0, 150., 250.], observables=[:r1])
scn2 = Scenario(model; parameters = [:k1=>0.015], tspan = (0., 200.), observables=[:r1])
@test typeof(scn1) <: Scenario
@test test_show(scn1)
@test saveat(scn1) == [0.0, 150., 250.]
@test saveat(scn2) == Float64[]
@test tspan(scn1) == (0., 250.)
@test tspan(scn2) == (0., 200.)
@test parameters(scn1)[:k1] == 0.02
@test parameters(scn2)[:k1] == 0.015

cs1 = sim(scn1)
cs2 = sim([:one=>scn1,:two=>scn2], parameters_upd=[:k1=>0.03])
@test typeof(cs1) <: HetaSimulator.SimResult
@test typeof(cs2) <: Vector{Pair{Symbol, HetaSimulator.SimResult}}
@test length(parameters(cs1))==0
@test parameters(last(cs2[:one]))[:k1] == 0.03

# Monte-Carlo simulation tests
mciter = 100
mc1 =  mc(model, [:k1=>Normal(0.02,1e-3)], mciter; tspan = (0., 200.), observables=[:r1])
mc2 = mc([:one=>scn1,:two=>scn2], [:k1=>Normal(0.02,1e-3)], mciter)
gsar = gsa(mc1,200)
@test typeof(mc1) <: HetaSimulator.MCResult
@test test_show(mc1)
@test typeof(mc2) <: Vector{Pair{Symbol, HetaSimulator.MCResult}}
@test typeof(mc1[1]) <: HetaSimulator.Simulation
@test typeof(mc2[1]) <: Pair{Symbol, HetaSimulator.MCResult}
@test length(mc1) == mciter
@test typeof(parameters(mc1)) <: Vector
@test times(mc1[1])[end] == 200.
@test keys((parameters(mc1[1]))) == (:k1,)
@test size(pearson(gsar)) == (1,1)
@test size(partial(gsar)) == (1,1)
@test size(standard(gsar)) == (1,1)

# Fitting tests
fscn1 = Scenario(model; parameters = [:k1=>0.02], tspan = (0., 200.), observables=[:A, :B, :r1])
fscn2 = Scenario(model; parameters = [:k1=>0.015], tspan = (0., 200.), observables=[:A, :B, :r1])
data = read_measurements("$HetaSimulatorDir/test/examples/single_comp/single_comp_data.csv")
add_measurements!(fscn1, data; subset = [:scenario => :one])
add_measurements!(fscn2, data; subset = [:scenario => :two])
fres = fit([:one=>fscn1, :two=>fscn2], [:k1=>0.01])

@test typeof(fres) <: HetaSimulator.FitResult
@test test_show(fres)
@test length(measurements(fscn1)) == 24
@test length(measurements(fscn2)) == 24
@test status(fres) == :FTOL_REACHED
@test obj(fres) ≈ 146.056244
@test typeof(optim(fres)) == Vector{Pair{Symbol, Float64}}
@test length(optim(fres)) == 1


#=
Copyright (c) 2018-2022 Chris Coey, Lea Kapelevich, and contributors

This Julia package Hypatia.jl is released under the MIT license; see LICENSE
file in the root directory or at https://github.com/chriscoey/Hypatia.jl

common code for native examples
=#

abstract type ExampleInstanceNative{T <: Real} <: ExampleInstance{T} end

# fallback: just check optimal status
function test_extra(
    inst::ExampleInstanceNative,
    solve_stats::NamedTuple,
    ::NamedTuple,
    )
    @test solve_stats.status == Solvers.Optimal
end

function run_instance(
    ex_type::Type{<:ExampleInstanceNative{T}}, # an instance of a native example
    inst_data::Tuple,
    inst_options::NamedTuple = NamedTuple(),
    solver_type::Type{<:Solvers.Solver} = Solvers.Solver{T};
    default_options::NamedTuple = NamedTuple(),
    test::Bool = true,
    rseed::Int = 1,
    verbose::Bool = true,
    ) where {T <: Real}
    new_options = merge(default_options, inst_options)

    verbose && println("setup model")
    setup_time = @elapsed begin
        Random.seed!(rseed)
        inst = ex_type(inst_data...)
        model = build(inst)
        model_stats = get_model_stats(model)
        solver = Solvers.Solver{T}(; default_options..., inst_options...)
    end
    flush(stdout); flush(stderr)

    verbose && println("solve and check")
    check_time = @elapsed begin
        Solvers.load(solver, model)
        Solvers.solve(solver)
        flush(stdout); flush(stderr)

        (solve_stats, solution) = process_result(model, solver)

        test && test_extra(inst, solve_stats, solution)
    end
    flush(stdout); flush(stderr)

    return (; model_stats..., solve_stats..., setup_time,
        check_time, :script_status => "Success")
end


using Test

using BraidChains
using BraidChains: SynchronicBraid, Hash, MetaData

import PeaceCypher: sign, hash
using PeaceCypher

using SynchronicBallot

include("basics.jl")

include("services.jl")

include("keychain.jl")


#=
# Infernal.jl usage example
=#

import Infernal
import Rfam

# Retrieve the path to Infernal programs

isfile(Infernal.path("cmalign"))

# Execute a command

success(Infernal.exe("cmfetch", Rfam.cm(), "RF00162"))


function str2et( dates::Array{ASCIIString,1} )
        ets = Array{Float64}( length(dates) )
    for i = 1:length(dates)
         ets[i] = str2et( dates[i] )
    end
    return ets
end

function cnmfrm( cnames::Array{ASCIIString,1} )
    frcodes = Array{Int32,1}( length(cnames) )
    frnames = Array{String,1}( length(cnames) )
    founds = Array{Bool,1}( length(cnames) )
    for i = 1:length(cnames)
        (frcodes[i], frnames[i], founds[i]) = cnmfrm( cnames[i] )
    end
    return (frcodes, frnames, founds)
end

function bodc2n( codes::Array{Int64,1} )
    names = Array{String,2}( 1, length(codes) )
    founds = Array{Bool,2}( 1, length(codes) )
    for i = 1:length(codes)
        (names[1, i], founds[1, i]) = bodc2n( codes[i] )
    end
    return (names, founds )
end
#

function bodc2s( codes::Array{Int64,1} )
    names = Array{String,1}( length(codes) )
    for i = 1:length(codes)
        names[i] = bodc2s( codes[i] )
    end
    return names
end

function cidfrm( cents::Array{Int64,1} )
    frcodes = Array{Int32,1}( length(cents) )
    frnames = Array{String,1}( length(cents) )
    founds = Array{Bool,1}( length(cents) )
    for i = 1:length(cents)
        (frcodes[i], frnames[i], founds[i]) = cidfrm( cents[i] )
    end
    return (frcodes, frnames, founds)
end

function cnmfrm( cnames::Array{ASCIIString,1} )
    frcodes = Array{Int32,1}( length(cnames) )
    frnames = Array{String,1}( length(cnames) )
    founds = Array{Bool,1}( length(cnames) )
    for i = 1:length(cnames)
        (frcodes[i], frnames[i], founds[i]) = cnmfrm( cnames[i] )
    end
    return (frcodes, frnames, founds)
end

function et2lst(
    ets::Array{Float64}, # SpiceDouble
    body::Int64, # SpiceInt
    lon::Float64, # SpiceDouble
    _type::AbstractString) # Ptr{ConstSpiceChar}
    hrs = Array{Int64,1}( length(ets) )      # Ptr{SpiceInt}
    mns = Array{Int64,1}( length(ets) )   # Ptr{SpiceInt}
    scs = Array{Int64,1}( length(ets) )   # Ptr{SpiceInt}
    times = Array{String,1}( length(ets) )  # Ptr{SpiceChar}
    ampm = Array{String,1}( length(ets) )   # Ptr{SpiceChar}
    for i = 1;length(ets)
        (hrs[i], mns[i], scs[i], times[i], ampm[i]) = et2lst( ets[i], body, lon, _type )
    end
    return (hrs, mns, scs, times, ampms)
end


module TestDatasets

using FileIO
using ColorTypes: Gray
using Pkg.Artifacts
using StringDistances

const artifacts_toml = abspath(joinpath(@__DIR__, "..", "Artifacts.toml"))

export testdataset

const remotedatasets = [
    "cameraman512_10",
    "cameraman256_10",
    "cameraman256_5",
    "cameraman128_5",
    "faces_5",
    "faces_10",
    "matches_5",
    "matches_10",
    "mandrill",
    "peppers",
    "playing_cards_5"
]

function testdataset(datasetname; download_only::Bool = false)

    datasetfiles = dataset_path(full_datasetname(datasetname))
    download_only && return datasetfiles
    dataset = load_dataset(datasetfiles)
    return dataset

end

function full_datasetname(datasetname)
    idx = findfirst(remotedatasets) do x
        startswith(x, datasetname)
    end
    if idx === nothing
        warn_msg = "\"$datasetname\" not found in `TestDatasets.remotedatasets`."

        best_match = _findnearest(datasetname)
        if isnothing(best_match[2])
            similar_matches = remotedatasets[_findall(datasetname)]
            if !isempty(similar_matches)
                similar_matches_msg = "  * \"" * join(similar_matches, "\"\n  * \"") * "\""
                warn_msg = "$(warn_msg) Do you mean one of the following?\n$(similar_matches_msg)"
            end
            throw(ArgumentError(warn_msg))
        else
            idx = best_match[2]
            @warn "$(warn_msg) Load \"$(remotedatasets[idx])\" instead."
        end
    end
    return remotedatasets[idx]
end

function dataset_path(datasetname)
    ensure_artifact_installed("datasets", artifacts_toml)

    dataset_dir = artifact_path(artifact_hash("datasets", artifacts_toml))
    return joinpath(dataset_dir, "datasets" ,datasetname)
end

function load_dataset(datasetfiles)
    image_pairs = readlines(joinpath(datasetfiles,"filelist.txt"))
    M,N = size(load(joinpath(datasetfiles,split(image_pairs[1],",")[1])))
    true_images = zeros(M,N,length(image_pairs))
    data_images = zeros(M,N,length(image_pairs))
    for i = 1:length(image_pairs)
        pair = split(image_pairs[i],",")
        true_images[:,:,i] = load(joinpath(datasetfiles,pair[1]))
        data_images[:,:,i] = load(joinpath(datasetfiles,pair[2]))
    end
    return true_images,data_images
end

_findall(name; min_score=0.6) = findall(name, remotefiles,JaroWinkler(), min_score=min_score)
_findnearest(name; min_score=0.8) = findnearest(name, remotefiles, JaroWinkler(), min_score=min_score)

end


# ideally this should be in the GenerativeAD pkg
function compute_ranks(rt)
	mask_nan_max = (x) -> (isnan(x) ? -Inf : x)
	rs = []
	for row in eachrow(rt)
		push!(rs, StatsBase.competerank(mask_nan_max.(Vector(row)), rev = true))
	end
	# each row represents ranks of a method
	reduce(hcat, rs)
end

const MODEL_RENAME = Dict(
    "aae_full" => "aae",
    "wae_full" => "wae",
    "vae_full" => "vae")

"""
	sorted_rank(d, agg; val_metric=:val_auc, tst_metric=:tst_auc, downsample=Dict{String, Int}())

Computes rank table from `d` and aggregation `agg`. Applies model and dataset aliases.
Sorts models based on model_type. Returns the rank table and also the aggregated dataframe.
"""
function sorted_rank(d, agg, val_metric=:val_auc, tst_metric=:tst_auc, downsample=Dict{String, Int}(); verbose=true)
	df_agg = agg(d, val_metric; downsample=downsample, verbose=verbose)

	df_agg["model_type"] = copy(df_agg["modelname"])
	apply_aliases!(df_agg, col="modelname", d=MODEL_RENAME)
	apply_aliases!(df_agg, col="modelname", d=MODEL_ALIAS)
	apply_aliases!(df_agg, col="dataset", d=DATASET_ALIAS)
	apply_aliases!(df_agg, col="model_type", d=MODEL_TYPE)

	sort!(df_agg, (:dataset, :model_type, :modelname))
	rt = rank_table(df_agg, tst_metric)
	names(rt)[2:end], rt
end

"""
automatic aggregation of combined dataframes for single/multi class datasets

single class datasets are marked with anomaly_class = -1 and will be aggregated usin aggregate_stats_mean_max
multi class datasets are marked with anomaly_class > 0 and will be aggregated usin aggregate_stats_max_mean

"""
function aggregate_stats_auto(df::DataFrame, criterion_col=:val_auc; kwargs...)
	df_single = select(filter(x -> x.anomaly_class == -1, df), Not(:anomaly_class))
	df_multi = filter(x -> x.anomaly_class > 0, df)

	if nrow(df_single) > 0 && nrow(df_multi) > 0
		df_agg_single = aggregate_stats_mean_max(df_single, criterion_col; kwargs...)
		df_agg_multi = aggregate_stats_max_mean(df_multi, criterion_col; kwargs...)
		return	vcat(df_agg_single, df_agg_multi, cols=:intersect)
	elseif nrow(df_single) > 0
		return aggregate_stats_mean_max(df_single, criterion_col; kwargs...)
	else
		return aggregate_stats_max_mean(df_multi, criterion_col; kwargs...)
	end
end

# top 10 mean-max
function aggregate_stats_mean_max_top_10(df::DataFrame,
										criterion_col=:val_auc;
										min_samples=("anomaly_class" in names(df)) ? 10 : 3,
										add_col=nothing, verbose=true)
	agg_cols = vcat(_prefix_symbol.("val", BASE_METRICS), _prefix_symbol.("tst", BASE_METRICS))
	agg_cols = vcat(agg_cols, _prefix_symbol.("val", PAT_METRICS), _prefix_symbol.("tst", PAT_METRICS))
	agg_cols = vcat(agg_cols, _prefix_symbol.("val", PATN_METRICS), _prefix_symbol.("tst", PATN_METRICS))
	agg_cols = vcat(agg_cols, _prefix_symbol.("val", PAC_METRICS), _prefix_symbol.("tst", PAC_METRICS))
	agg_cols = vcat(agg_cols, Symbol.(TRAIN_EVAL_TIMES))
	agg_cols = (add_col !== nothing) ? vcat(agg_cols, add_col) : agg_cols
	top10_std_cols = _prefix_symbol.(agg_cols, "top_10_std")
	std_cols = _prefix_symbol.(agg_cols, "std")

	# agregate by seed over given hyperparameter and then choose best
	results = []
	for (dkey, dg) in pairs(groupby(df, :dataset))
		for (mkey, mg) in pairs(groupby(dg, :modelname))
			pg = groupby(mg, :phash)
			mg_suff = reduce(vcat, [g for g in pg if nrow(g) >= min_samples])
			pg_agg = combine(groupby(mg_suff, :phash),
						nrow => :psamples,
						agg_cols .=> mean .=> agg_cols,
						agg_cols .=> std,
						:parameters => unique => :parameters)

			# sort by criterion_col
			sort!(pg_agg, order(criterion_col, rev=true))

			# take best 10 models
			best_10 = first(pg_agg, 10)
			best = combine(best_10,
							agg_cols .=> mean .=> agg_cols,
							std_cols .=> mean .=> std_cols,
							agg_cols .=> std  .=> top10_std_cols)

			# add grouping keys
			best[:dataset] = dkey.dataset
			best[:modelname] = mkey.modelname

			push!(results, best)
		end
	end
	vcat(results...)
end

# top 10 max-mean
function aggregate_stats_max_mean_top_10(df::DataFrame,
										criterion_col=:val_auc; add_col=nothing, verbose=true)
	agg_cols = vcat(_prefix_symbol.("val", BASE_METRICS), _prefix_symbol.("tst", BASE_METRICS))
	agg_cols = vcat(agg_cols, _prefix_symbol.("val", PAT_METRICS), _prefix_symbol.("tst", PAT_METRICS))
	agg_cols = vcat(agg_cols, _prefix_symbol.("val", PATN_METRICS), _prefix_symbol.("tst", PATN_METRICS))
	agg_cols = vcat(agg_cols, _prefix_symbol.("val", PAC_METRICS), _prefix_symbol.("tst", PAC_METRICS))
	agg_cols = vcat(agg_cols, Symbol.(TRAIN_EVAL_TIMES))
	agg_cols = (add_col !== nothing) ? vcat(agg_cols, add_col) : agg_cols
	top10_std_cols = _prefix_symbol.(agg_cols, "top_10_std")
	std_cols = _prefix_symbol.(agg_cols, "std")

	agg_keys = ("anomaly_class" in names(df)) ? [:seed, :anomaly_class] : [:seed]

	# choose best for each seed/anomaly_class cobination and average over them
	results = []
	for (dkey, dg) in pairs(groupby(df, :dataset))
		for (mkey, mg) in pairs(groupby(dg, :modelname))
			partial_results = []

			if ("anomaly_class" in names(df))
				classes = unique(mg.anomaly_class)
				dif = setdiff(collect(1:10), classes)
				if (length(classes) < 10) && verbose
					@warn "$(mkey.modelname) - $(dkey.dataset): missing runs on anomaly_class $(dif)."
				end
			else
				seeds = unique(mg.seed)
				dif = setdiff(collect(1:5), seeds)
				if (length(seeds) < 3) && verbose
					@warn "$(mkey.modelname) - $(dkey.dataset): missing runs on seed $(dif)."
				end
			end

			# iterate over seed-anomaly_class groups
			for (skey, sg) in pairs(groupby(mg, agg_keys))
				ssg = sort(sg, order(criterion_col, rev=true))
				# best hyperparameter after sorting by criterion_col
				best_10 = first(ssg, 10)
				best = combine(best_10,
							agg_cols .=> mean .=> agg_cols,
							agg_cols .=> std  .=> top10_std_cols)

				push!(partial_results, best)
			end

			best_per_seed = reduce(vcat, partial_results)
			# average over seed-anomaly_class groups
			best = combine(best_per_seed,
						agg_cols .=> mean .=> agg_cols,
						top10_std_cols .=> mean .=> top10_std_cols,
						agg_cols .=> std)

			# add grouping keys
			best[:dataset] = dkey.dataset
			best[:modelname] = mkey.modelname

			push!(results, best)
		end
	end
	vcat(results...)
end


# ----------------------------------------------------------------------------------- #
# Copyright (c) 2016 Varnerlab
# Robert Frederick Smith School of Chemical and Biomolecular Engineering
# Cornell University, Ithaca NY 14850
#
# 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.
# ----------------------------------------------------------------------------------- #
#
# ----------------------------------------------------------------------------------- #
# Function: calculate_input_array
# Description: Calculate the simulation inputs at time t
# Generated on: 2016-11-02T07:00:00.258
#
# Input arguments:
# t::Float64 => Current time value (scalar)
# x::Array{Float64,1} => State array (number_of_species x 1)
# data_dictionary::Dict{AbstractString,Any} => Dictionary holding model parameters
#
# Output arguments:
# u::Array{Float64,1} => Input array (number_of_species x 1) at time t
# ----------------------------------------------------------------------------------- #
function calculate_input_array(t::Float64,x::Array{Float64,1},data_dictionary::Dict{AbstractString,Any})

	# Initialize default -
	u_array = zeros(length(x))

	# return -
	return u_array
end


module DataWrangling

export
       # Coarse graining
       coarse_grain,

       # Feature scaling
       AbstractFeatureScaling, ZeroMeanUnitVarianceScaling, MinMaxScaling, scale, unscale,

       # Convective adjustment
       convective_adjust!,

       # Animations
       animate_gif

using Statistics
using OrderedCollections
using Plots

include("coarse_graining.jl")
include("feature_scaling.jl")
include("convective_adjust.jl")
include("animate_gif.jl")

end


module DifferencesBase

using SparseArrays

export AbstractDifference,
    VectorDifference,
    MatrixDifference,
    DictDifference,
    NamedTupleDifference,
    SetDifference,
    added_indices,
    modified_indices,
    removed_indices,
    added,
    common,
    modified,
    removed

include("types.jl")
include("differences.jl")
include("math.jl")
include("show.jl")

end # module


function solve(filename)
    values = readlines(filename*".txt")
    values = split.(values, " | ")

    input = String[]
    output = String[]
    for line in values
        input = push!(input, line[1])
        output = push!(output, line[2])
    end
    input = split.(input, ' ')
    output = split.(output, ' ')

    count = 0
    for line in output
        for el in line
            if length(el) == 2 || length(el) == 4 || length(el) == 3 || length(el) == 7
                count += 1
            end
        end
    end
    println("Count of 1, 4, 7, 8: "*string(count))
end

function find_patterns(input)
    codesperline = []
    zero = one = two = three = four = five = six = seven = eight = nine = ""
    upper_line = mid_line = lower_line = right_upper = right_lower = left_upper = left_lower = ""
    for line in input
        #one
        for el in line
            if length(el) == 2
                one = join(sort(collect(el)))
                break
            end
        end
        #seven
        for el in line
            if length(el) == 3
                seven = join(sort(collect(el)))
                upper_line = replace(el, one => "")
                break
            end
        end
        #four
        for el in line
            if length(el) == 4
                four = join(sort(collect(el)))
                break
            end
        end
        #eight
        for el in line
            if length(el) == 7
                eight = join(sort(collect(el)))
                break
            end
        end
        #three is length five and contains one
        length_five_candidates = []
        length_six_candidates = []
        for el in line
            if length(el) == 5
                length_five_candidates = push!(length_five_candidates, join(sort(collect(el))))
            elseif length(el) == 6
                length_six_candidates = push!(length_six_candidates, join(sort(collect(el))))
            end
        end
        #three
        for el in length_five_candidates
            if stringcontains(one, el)
                three = el
                setdiff!(length_five_candidates, three)
                break
            end
        end
        #six
        for el in length_six_candidates
            if stringcontains(one, el) == false
                six = el
                setdiff!(length_six_candidates, six)
                break
            end
        end
        #nine
        for el in length_six_candidates
            if stringcontains(three, el)
                nine = el
                setdiff!(length_six_candidates, nine)
                break
            end
        end
        #zero
        for el in length_six_candidates
            if stringcontains(nine, el) == false && stringcontains(six, el) == false
                zero = el
                break
            end
        end

        #two
        for char in nine
            if occursin(char, three) == false
                left_upper = char
                break
            end
        end

        for el in length_five_candidates
            if stringcontains(left_upper, el) == false && el != three
                two = el
                setdiff!(length_five_candidates, two)
                break
            end
        end
        #five
        for el in length_five_candidates
            if stringcontains(three, el) == false && stringcontains(two, el) == false
                five = el
                break
            end
        end

        numbers = [zero, one, two, three, four, five, six, seven, eight, nine]
        for (i,n) in enumerate(numbers)
            numbers[i] = join(sort(collect(n)))
        end
        codesperline = push!(codesperline, numbers)
    end
    return codesperline
end

function stringcontains(ptrn, str)
    res = true
    for char in ptrn
        if occursin(char, str) == false
            res = false
            break
        end
    end
    return res
end

function decode(output, codesperline)
    output_numbers = []
    for (i,line) in enumerate(output)
        numbers = codesperline[i]
        decoded = ""
        for el in line
            el = join(sort(collect(el)))
            for i in 0:9
                if el == numbers[i+1]
                    decoded = decoded * string(i)
                    break
                end
            end
        end
        decoded_int = parse(Int, decoded)
        output_numbers = push!(output_numbers, decoded_int)
    end

    sum = 0
    for number in output_numbers
        sum += number
    end
    return sum
end

function solve2(filename)
    values = readlines(filename*".txt")
    values = split.(values, " | ")

    input = String[]
    output = String[]
    for line in values
        input = push!(input, line[1])
        output = push!(output, line[2])
    end
    input = split.(input, ' ')
    output = split.(output, ' ')

    codes = find_patterns(input)
    res = decode(output, codes)
    println("Sum of all numbers: "*string(res))
end

solve("smalltest")
solve("input")

solve2("smalltest")
solve2("input")


"""
$(TYPEDEF)

Statistics from the differential equation solver about the solution process.

## Fields

- nf: Number of function evaluations. If the differential equation is a split function,
  such as a `SplitFunction` for implicit-explicit (IMEX) integration, then `nf` is the
  number of function evaluations for the first function (the implicit function)
- nf2: If the differential equation is a split function, such as a `SplitFunction`
  for implicit-explicit (IMEX) integration, then `nf2` is the number of function
  evaluations for the second function, i.e. the function treated explicitly. Otherwise
  it is zero.
- nw: The number of W=I-gamma*J (or W=I/gamma-J) matrices constructed during the solving
  process.
- nsolve: The number of linear solves `W\b` required for the integration.
- njacs: Number of Jacobians calculated during the integration.
- nnonliniter: Total number of iterations for the nonlinear solvers.
- nnonlinconvfail: Number of nonlinear solver convergence failures.
- ncondition: Number of calls to the condition function for callbacks.
- naccept: Number of accepted steps.
- nreject: Number of rejected steps.
- maxeig: Maximum eigenvalue over the solution. This is only computed if the
  method is an auto-switching algorithm.
"""
mutable struct DEStats
  nf::Int
  nf2::Int
  nw::Int
  nsolve::Int
  njacs::Int
  nnonliniter::Int
  nnonlinconvfail::Int
  ncondition::Int
  naccept::Int
  nreject::Int
  maxeig::Float64
end

DEStats(x::Int = -1) = DEStats(x, x, x, x, x, x, x, x, x, x, 0.0)

function Base.show(io::IO, s::DEStats)
  println(io, summary(s))
  @printf io "%-50s %-d\n" "Number of function 1 evaluations:" s.nf
  @printf io "%-50s %-d\n" "Number of function 2 evaluations:" s.nf2
  @printf io "%-50s %-d\n" "Number of W matrix evaluations:" s.nw
  @printf io "%-50s %-d\n" "Number of linear solves:" s.nsolve
  @printf io "%-50s %-d\n" "Number of Jacobians created:" s.njacs
  @printf io "%-50s %-d\n" "Number of nonlinear solver iterations:" s.nnonliniter
  @printf io "%-50s %-d\n" "Number of nonlinear solver convergence failures:" s.nnonlinconvfail
  @printf io "%-50s %-d\n" "Number of rootfind condition calls:" s.ncondition
  @printf io "%-50s %-d\n" "Number of accepted steps:"  s.naccept
  @printf io "%-50s %-d" "Number of rejected steps:" s.nreject
  iszero(s.maxeig) || @printf io "\n%-50s %-d" "Maximum eigenvalue recorded:" s.maxeig
end


using DataFrames

"""
This function takes in a file path/name and exports a dataframe consisting of the circuit parameters
Specifically the file should be from LTSPICE (or it wont work)

If no input is given, the user will be prompted for a FileName
"""
function LTSpiceLoad(FileName=nothing)
        # if FileName===nothing
        #         FileName = open_dialog("Pick a file")
        # end
        f =open(FileName)
         Lines =  readlines(f)

         close(f)

         SPICE_DF = DataFrame((Type = Any[], Value=Any[], Node1=Any[], Node2=Any[], Name=Any[],ESR = Any[]))
        for i in 1:length(Lines)

                Spaces = vcat(findall(" ",Lines[i])...)

                if ~(Lines[i][1]=='*')
                        # println(Lines[i])
                end
                if (Lines[i][1]=='R')|(Lines[i][1]=='L')|(Lines[i][1]=='C')
                        if length(Spaces)==3
                                Value = Lines[i][Spaces[3]+1:end]
                        else
                                Value = Lines[i][Spaces[3]+1:Spaces[4]-1]
                        end
                        Node1 = Lines[i][Spaces[1]+1:Spaces[2]-1]
                        Node2 = Lines[i][Spaces[2]+1:Spaces[3]-1]
                        Name = Lines[i][1:Spaces[1]-1]
                        RSerInds = findall("Rser=", Lines[i])
                        if length(RSerInds)>0
                                RSerIndEnd =RSerInds[1][end]+1
                                if RSerInds[1][end]>maximum(Spaces)
                                        RSerVal = MakeNumericalVals(Lines[i][RSerIndEnd:end])
                                else
                                        Tmp = findfirst(x-> x>RSerIndEnd,Spaces)
                                        RSerVal = MakeNumericalVals(Lines[i][RSerIndEnd:Spaces[Tmp]])
                                end
                        else
                                RSerVal=0
                        end
                        push!(SPICE_DF,[Lines[i][1] MakeNumericalVals(Value) Node1 Node2 Name RSerVal])

                elseif (Lines[i][1]=='V')
                        Node1 = Lines[i][Spaces[1]+1:Spaces[2]-1]
                        Node2 = Lines[i][Spaces[2]+1:Spaces[3]-1]
                        Name = Lines[i][1:Spaces[1]-1]
                        push!(SPICE_DF,[Lines[i][1] 1 Node1 Node2 Name 0])

                end

        end

        return SPICE_DF
end

"""
A simple string to numeric conversion for metric prefixes
Input: String
Output Float64

e.g.
        MakeNumericalVals("1n")
         > 1.0e-9
        MakeNumericalVals.(["1n" , "2m"])
                2-element Array{Float64,1}:
                1.0e-9
                0.002
"""
function MakeNumericalVals(ValString::String)
        NewString = replace(ValString,"Meg"=>"e6")
        NewString = replace(NewString,"k"=>"e3")
        NewString = replace(NewString,"m"=>"e-3")
        NewString = replace(NewString,"u"=>"e-6")
        NewString = replace(NewString, "n"=>"e-9")
        NewString = replace(NewString, "p"=>"e-12")
        return parse(Float64,NewString)
end


# see also the non-macro versions in datasets_requires.jl

"""Load a well-known public regression dataset with `Continuous` features."""
macro load_boston()
    quote
        import CSV
        y, X = unpack(load_boston(), ==(:MedV), x->x != :Chas)
        (X, y)
    end
end

"""Load a reduced version of the well-known Ames Housing task"""
macro load_reduced_ames()
    quote
        import CSV
        y, X = unpack(load_reduced_ames(), ==(:target), x-> true)
        (X, y)
    end
end

"""Load the full version of the well-known Ames Housing task."""
macro load_ames()
    quote
        import CSV
        y, X = unpack(load_ames(), ==(:target), x->x != :Id)
        (X, y)
    end
end

"""Load a well-known public classification task with nominal features."""
macro load_iris()
    quote
        import CSV
        y, X = unpack(load_iris(), ==(:target), x-> true)
        (X, y)
    end
end

"""Load a well-known crab classification dataset with nominal features."""
macro load_crabs()
    quote
        import CSV
        y, X = unpack(load_crabs(), ==(:sp), x-> !(x in [:sex, :index]))
        (X, y)
    end
end


@testset "72.edit-distance.jl" begin
    @test min_distance("horse", "ros") == 3
    @test min_distance("intention", "execution") == 5
    @test min_distance("", "") == 0
end


using MarketData
export simulate

# run model without closing price updates

# convert TimeArray to Timestamps array
clt = Timestamp[]

for i in 1:length(cl)
    push!(clt, Timestamp(cl[i].timestamp[1], cl[i].values[1]))
end

cr = vcat(res, clt) # combine both arrays with txn array first
sd = sort(Date[c.timestamp for c in cr]) # get dates and sort
ud = unique(sd) # remove duplicates

mixed_array = cr[ud[1]] # initialize array to be pushed

for i in 2:length(ud)
    mixed_array = vcat(mixed_array, cr[ud[i]])
end

# this takes mixed_array and updates the following
# txn & mark_price explicity
# A, Pval, Aval by association

function simulate(tsv::Vector{Timestamp})
    for t in 1:length(tsv)
        if typeof(tsv[t].value) <: FinancialAsset
            push!(txn, tsv[t].value)
        elseif typeof(tsv[t].value) == Float64
            push!(mark_price, tsv[t].value)
        else
            error("rogue timestamp in your midst!")
        end
    end
end


#
# Quasi Submodule for Inference
#

include("exact.jl")
include("likelihood.jl")
include("gibbs.jl")
include("loopy_belief.jl")


#
println("Ag) Test of the QED model corrections to the level structure of atoms and ions.")
#
if   true
wa = Atomic.Computation(Atomic.Computation(), name="QED estimates for xenon-like Sn", grid=JAC.Radial.Grid(true), nuclearModel=Nuclear.Model(50.),
                        configs=[Configuration("1s^2 2s^2 2p^6")],
                        asfSettings=AsfSettings(true, false, Basics.DFSField(), "hydrogenic", Dict{Subshell, Orbital}(), [1],    40, 1.0e-6, JAC.Subshell[], JAC.Subshell[],
                                                true, false, QedPetersburg(), LSjjSettings(false),
                                                false, [1,2,3,4], false, JAC.LevelSymmetry[] )  )

elseif  false
wa = Atomic.Computation(Atomic.Computation(), name="QED estimates for hydrogen-like Sn", grid=JAC.Radial.Grid(true), nuclearModel=Nuclear.Model(50.),
                        configs=[Configuration("1s"), Configuration("2s"), Configuration("2p")],
                        asfSettings=AsfSettings(true, false, Basics.DFSField(), "hydrogenic", Dict{Subshell, Orbital}(), [1],    0, 1.0e-6, JAC.Subshell[], JAC.Subshell[],
                                                true, false, QedPetersburg(),  LSjjSettings(false),
                                                false, [1,2,3,4], false, JAC.LevelSymmetry[] )  )

end

wb = perform(wa)

# NoneQed() QedPetersburg()
#


function HansenThirdBodyLinear(arg0::jint, arg1::jint)
    return HansenThirdBodyLinear((jint, jint), arg0, arg1)
end

function compute_init_values(obj::HansenThirdBodyLinear, arg0::jdouble, arg1::jdouble, arg2::jdouble)
    return jcall(obj, "computeInitValues", void, (jdouble, jdouble, jdouble), arg0, arg1, arg2)
end

function get_derivative(obj::HansenThirdBodyLinear, arg0::jint, arg1::jdouble)
    return jcall(obj, "getDerivative", jdouble, (jint, jdouble), arg0, arg1)
end

function get_value(obj::HansenThirdBodyLinear, arg0::jint, arg1::jdouble)
    return jcall(obj, "getValue", jdouble, (jint, jdouble), arg0, arg1)
end


# load propagator bubbles
include("bubbles.jl")

# load parquet equations for different symmetries
include("parquet_lib/parquet_su2/parquet_su2.jl")
include("parquet_lib/parquet_u1_dm/parquet_u1_dm.jl")

# load flow equations for different symmetries
include("flow_lib/flow_su2/flow_su2.jl")
include("flow_lib/flow_u1_dm/flow_u1_dm.jl")

# load quadrature rule for vertex integration
include("quadrature.jl")

# load SDE computation
include("SDE.jl")

# load BSE computation of channels for frequency tuples (w1, w2, w3) and kernels k
include("BSE_s.jl")
include("BSE_t.jl")
include("BSE_u.jl")

# load self energy derivative calculation
include("flow_self.jl")

# load flow computation of channels for frequency tuples (w1, w2, w3) and kernels k
include("flow_s.jl")
include("flow_t.jl")
include("flow_u.jl")

# load full BSE and flow computation
include("compute.jl")

# load tests and timers
include("test.jl")
include("timers.jl")


for p in ("Knet","ArgParse","Images")
    Pkg.installed(p) == nothing && Pkg.add(p)
end
include(Pkg.dir("Knet","data","mnist.jl"))
include(Pkg.dir("Knet","data","imagenet.jl"))

module CGAN
using Knet
using Images
using ArgParse
using JLD2, FileIO

function main(args)
    o = parse_options(args)
    o[:seed] > 0 && Knet.setseed(o[:seed])

    # load models, data, optimizers
    wd, wg, md, mg = load_weights(o[:atype], o[:zdim], o[:loadfile])
    xtrn,ytrn,xtst,ytst = Main.mnist()
    dtrn = minibatch(xtrn, ytrn, o[:batchsize]; shuffle=true, xtype=o[:atype])
    optd = map(wi->eval(parse(o[:optim])), wd)
    optg = map(wi->eval(parse(o[:optim])), wg)
    z = sample_noise(o[:atype],o[:zdim],prod(o[:gridsize]))
    l0 = map(i->reshape(collect(1:9), 1, 9), 1:9)
    l1 = vec(vcat(l0...))

    if o[:outdir] != nothing && !isdir(o[:outdir])
        mkpath(o[:outdir])
        mkpath(joinpath(o[:outdir],"models"))
        mkpath(joinpath(o[:outdir],"generations"))
    end

    # training
    println("training started..."); flush(STDOUT)
    for epoch = 1:o[:epochs]
        dlossval = glossval = 0
        @time for (x,y) in dtrn
            x0 = 2x-1; x0 = reshape(x0, 784, div(length(x0),784))
            noise = sample_noise(o[:atype],o[:zdim],length(y))
            dlossval += train_discriminator!(wd,wg,md,mg,x0,y,noise,optd,o)
            noise = sample_noise(o[:atype],o[:zdim],length(y))
            glossval += train_generator!(wg,wd,mg,md,noise,y,optg,o)
        end
        dlossval /= length(dtrn); glossval /= length(dtrn)
        println((:epoch,epoch,:dloss,dlossval,:gloss,glossval))
        flush(STDOUT)

        # save models and generations
        if o[:outdir] != nothing
            filename = @sprintf("%04d.png",epoch)
            filepath = joinpath(o[:outdir],"generations",filename)
            plot_generations(
                wg, mg, l1; z=z, savefile=filepath,
                scale=o[:gridscale], gridsize=o[:gridsize])

            filename = @sprintf("%04d.jld2",epoch)
            filepath = joinpath(o[:outdir],"models",filename)
            save_weights(filepath,wd,wg,md,mg)
        end
    end

    return wd,wg,md,mg
end

function parse_options(args)
    s = ArgParseSettings()
    s.description = "Conditional GANs on MNIST."

    @add_arg_table s begin
        ("--atype"; default=(gpu()>=0?"KnetArray{Float32}":"Array{Float32}");
         help="array and float type to use")
        ("--batchsize"; arg_type=Int; default=100; help="batch size")
        ("--zdim"; arg_type=Int; default=100; help="noise dimension")
        ("--epochs"; arg_type=Int; default=20; help="# of training epochs")
        ("--seed"; arg_type=Int; default=-1; help="random seed")
        ("--gridsize"; arg_type=Int; nargs=2; default=[9,9])
        ("--gridscale"; arg_type=Float64; default=2.0)
        ("--optim"; default="Adam(;lr=0.0002, beta1=0.5)")
        ("--loadfile"; default=nothing; help="file to load trained models")
        ("--outdir"; default=nothing; help="output dir for models/generations")
    end

    isa(args, AbstractString) && (args=split(args))
    o = parse_args(args, s; as_symbols=true)
    o[:atype] = eval(parse(o[:atype]))
    if o[:outdir] != nothing
        o[:outdir] = abspath(o[:outdir])
    end
    return o
end

function load_weights(atype,zdim,loadfile=nothing)
    if loadfile == nothing
        wd, md = initwd(atype)
        wg, mg = initwg(atype)
    else
        @load loadfile wd wg md mg
        wd = convert_weights(wd, atype)
        wg = convert_weights(wg, atype)
        md = convert_moments(md, atype)
        mg = convert_moments(mg, atype)
    end
    return wd, wg, md, mg
end

function save_weights(savefile,wd,wg,md,mg)
    save(savefile,
         "wd", convert_weights(wd),
         "wg", convert_weights(wg),
         "md", convert_moments(md),
         "mg", convert_moments(mg))
end

function convert_weights(w, atype=Array{Float32})
    w0 = map(wi->convert(atype, wi), w)
    w1 = convert(Array{Any}, w0)
end

function convert_moments(moments,atype=Array{Float32})
    clone = map(mi->bnmoments(), moments)
    for k = 1:length(clone)
        if moments[k].mean != nothing
            clone[k].mean = convert(atype, moments[k].mean)
        end

        if moments[k].var != nothing
            clone[k].var = convert(atype, moments[k].var)
        end
    end
    return convert(Array{Any,1}, clone)
end

function leaky_relu(x, alpha=0.2)
    pos = max(0,x)
    neg = min(0,x) * alpha
    return pos + neg
end

function sample_noise(atype,zdim,nsamples,mu=0.5,sigma=0.5)
    noise = convert(atype, randn(zdim,nsamples))
    normalized = (noise-mu)/sigma
end

function initwd(atype, winit=0.01)
    w = Any[]
    m = Any[]

    # fc layers for inputs (x,y)
    push!(w, winit*randn(1024,784))
    push!(w, zeros(1024,1))
    push!(w, winit*randn(1024,10))
    push!(w, zeros(1024,1))

    # fc layers for concatenated
    push!(w, winit*randn(512,2048))
    push!(w, bnparams(512))
    push!(m, bnmoments())
    push!(w, winit*randn(256,512))
    push!(w, bnparams(256))
    push!(m, bnmoments())

    push!(w, winit*randn(2,256))
    push!(w, zeros(2,1))
    return convert_weights(w,atype), m
end

function dnet(w,x,y,m; training=true, pdrop=0.5, alpha=0.2)
    x1 = leaky_relu.(w[1] * x .+ w[2], alpha)
    x2 = leaky_relu.(w[3][:,y] .+ w[4])
    x3 = vcat(x1,x2)
    x4 = leaky_relu.(batchnorm(w[5]*x3, m[1], w[6]), alpha)
    x5 = leaky_relu.(batchnorm(w[7]*x4, m[2], w[8]), alpha)
    x6 = w[9] * x5 .+ w[10]
end

function dloss(w,m,real_images,real_labels,fake_images,fake_labels,ygold)
    yreal = dnet(w,real_images,ygold,m)
    real_loss = nll(yreal, real_labels)
    yfake = dnet(w,fake_images,ygold,m)
    fake_loss = nll(yfake, fake_labels)
    return real_loss + fake_loss
end

dlossgradient = gradloss(dloss)

function train_discriminator!(wd,wg,md,mg,real_images,ygold,noise,optd,o)
    fake_images = gnet(wg,noise,ygold,mg; training=true)
    nsamples = div(length(real_images),784)
    real_labels = ones(Int64, 1, nsamples)
    fake_labels = 2ones(Int64, 1, nsamples)
    gradients, lossval = dlossgradient(
        wd,md,real_images,real_labels,fake_images,fake_labels,ygold)
    update!(wd, gradients, optd)
    return lossval
end

function initwg(atype, winit=0.01)
    w = Any[]
    m = Any[]

    push!(w, winit*randn(256,100))
    push!(w, bnparams(256))
    push!(m, bnmoments())

    push!(w, winit*randn(256,10))
    push!(w, bnparams(256))
    push!(m, bnmoments())

    push!(w, winit*randn(512,512))
    push!(w, bnparams(512))
    push!(m, bnmoments())

    push!(w, winit*randn(1024,512))
    push!(w, bnparams(1024))
    push!(m, bnmoments())

    push!(w, winit*randn(784,1024))
    push!(w, zeros(784,1))

    return convert_weights(w,atype), m
end

function gnet(w,z,y,m; training=true, pdrop=0.5)
    x0 = relu.(batchnorm(w[1]*z, m[1], w[2]))
    x1 = relu.(batchnorm(w[3][:,y], m[2], w[4]))
    x2 = vcat(x0,x1)
    x3 = relu.(batchnorm(w[5]*x2, m[3], w[6]))
    x4 = relu.(batchnorm(w[7]*x3, m[4], w[8]))
    x5 = tanh.(w[9] * x4 .+ w[10])
end

function gloss(wg,wd,mg,md,noise,ygold,labels)
    fake_images = gnet(wg,noise,labels,mg)
    ypred = dnet(wd,fake_images,labels,md)
    return nll(ypred, ygold)
end

glossgradient = gradloss(gloss)

function train_generator!(wg,wd,mg,md,noise,labels,optg,o)
    ygold = ones(Int64, 1, length(labels))
    gradients, lossval = glossgradient(wg,wd,mg,md,noise,ygold,labels)
    update!(wg,gradients,optg)
    return lossval
end

function plot_generations(
    wg, mg, labels; z=nothing, gridsize=(8,8), scale=1.0, savefile=nothing)
    if z == nothing
        nimg = prod(gridsize)
        zdim = size(wg[1],2)
        atype = wg[1] isa KnetArray ? KnetArray{Float32} : Array{Float32}
        z = sample_noise(atype,zdim,nimg)
    end
    output = Array(0.5*(1+gnet(wg,z,labels,mg; training=false)))
    images = map(i->reshape(output[:,i], 28, 28, 1), 1:size(output,2))
    grid = Main.make_image_grid(images; gridsize=gridsize, scale=scale)
    if savefile == nothing
        display(colorview(Gray, grid))
    else
        save(savefile, grid)
    end
end

splitdir(PROGRAM_FILE)[end] == "conditional-gan.jl" && main(ARGS)

end # module


using Crux
using POMDPModels
using Test
using DataStructures
using CUDA
using Flux
using Random

## data structures
@test MinHeap == MutableBinaryHeap{Float32, DataStructures.FasterForward}

## Fenwick Trees
t = FenwickTree(ones(Float32, 10))
@test t isa FenwickTree{Float32}
v = rand(16)
t2pow = FenwickTree(v)

# Inverse query on non power of 2 arrays
@test inverse_query(t, 0.5) == 1
@test inverse_query(t, 1.0) == 1
@test inverse_query(t, 0) == 1
@test inverse_query(t, -1) == 1
@test inverse_query(t, 3.0) == 3
@test inverse_query(t, 3.5) == 4
@test inverse_query(t, 4.0) == 4
@test inverse_query(t, 9.5) == 10
@test inverse_query(t, 10) == 10
@test inverse_query(t, 11) == 11

# Search up to a given range (used due to some numerical problems when searching the full range)
priorities = FenwickTree(zeros(Float32, 100000))
N = 13500
for i=1:N
    update!(priorities, i, rand(Float32))
end
ptot = prefixsum(priorities, N)

B = 320
Δp = ptot / B
inverse_query(priorities, B*Δp)
inverse_query(priorities, ptot)
@test inverse_query(priorities, B*Δp, N) <= N+1
inverse_query(priorities, ptot, N-1)

for i = 1:1000
    rns = [rand() for _=1:B]
    indices = [inverse_query(priorities, (j + rns[j] - 1) * Δp, N-1) for j=1:B]
    if !all(indices .<= N)
        i = findfirst(indices .> N)
        println("index: ", i, " rn: ", rns[i], " val: ", inverse_query(priorities, (i + rns[i] - 1) * Δp))
    end
    @test all(indices .<= N)
end

# power of 2 arrays
for i=1:1000
    rval = 5*rand()
    ind = inverse_query(t2pow, rval, 16)

    @test prefixsum(t2pow, ind) ≥ rval
    @test prefixsum(t2pow, ind-1) < rval
end
@test inverse_query(t2pow, prefixsum(t2pow, 16)) == 16

# value test
@test all(t[i] == 1 for i=1:10)
@test all(t2pow[i] ≈ v[i] for i=1:16)

# update! test
update!(t, 4, 2.5)
@test t[4] ≈ 2.5
@test prefixsum(t, 4) ≈ 5.5

## Binary heap test
v = rand(19)
b = MutableBinaryHeap{Float64, DataStructures.FasterForward}(v)
@test first(b) == minimum(v)

push!(b, 1.0)
@test b[20] == 1.

update!(b, 1, 0.)
@test first(b) == 0.00

vlarge = rand(100000)
t = FenwickTree(vlarge)
h = MutableBinaryHeap{Float64, DataStructures.FasterForward}(vlarge)

# @time first(h)
# @time push!(h, 0.5)
# @time h[999]
#
# @time t[999]
# @time inverse_query(t, 200)
# @time value(t, 999)

## mdp_data
d1 = mdp_data(ContinuousSpace(3), ContinuousSpace(4), 100)
d2 = mdp_data(ContinuousSpace(3), ContinuousSpace(4), 100, [:weight, :t, :advantage, :return, :logprob])
# @test_throws ErrorException  mdp_data(ContinuousSpace(3), ContinuousSpace(4), 100, [:bad_key]) # Now this just throws a warning
d3 = mdp_data(ContinuousSpace(3), ContinuousSpace(4), 100, ArrayType = CuArray)

@test d1[:s] == zeros(Float32, 3, 100) && d2[:s] == zeros(Float32, 3, 100) && size(d3[:s]) == (3, 100)
@test d1[:sp] == zeros(Float32, 3, 100) && d2[:sp] == zeros(Float32, 3, 100) && size(d3[:sp]) == (3, 100)
@test d1[:a] == zeros(Bool, 4, 100) && size(d2[:a]) == (4, 100) && size(d3[:a]) == (4, 100)
@test d1[:r] == zeros(Float32, 1, 100) && size(d2[:r]) == (1, 100) && size(d3[:r]) == (1, 100)
@test d1[:done] == zeros(Bool, 1, 100) && size(d2[:done]) == (1, 100) && size(d3[:done]) == (1, 100)
@test d2[:return] == zeros(1, 100) && d2[:advantage] == zeros(1, 100) && d2[:logprob] == zeros(1, 100) && d2[:t] == zeros(1, 100) && d2[:weight] == ones(1, 100)
@test !haskey(d1, :return) && !haskey(d1, :advantage) && !haskey(d1, :logprob) && !haskey(d1, :t) && !haskey(d1, :weight)
@test d3[:s] isa CuArray

## circular_indices
circ_inds(start, Nsteps, l) = mod1.(start:start + Nsteps - 1, l)
@test circ_inds(4, 60, 100) == [4:63 ...]
@test circ_inds(1, 100, 100) == [1:100 ...]
@test circ_inds(1, 101, 100) == [1:100 ..., 1]
@test circ_inds(1, 120, 100) == [1:100 ..., 1:20 ...]
@test circ_inds(90, 20, 100) == [90:100 ..., 1:9 ...]

## Get last N indices test
# A not full buffer
b = ExperienceBuffer(ContinuousSpace(2), ContinuousSpace(1), 100)
d = Dict(:s => 2*ones(2,50), :a => ones(Bool, 1,50), :sp => ones(2,50), :r => ones(1,50), :done => zeros(1,50), :weight=>zeros(1,50))
push!(b, d)

@test get_last_N_indices(b, 10) == collect(41:50)
@test get_last_N_indices(b, 1) == [50]
@test get_last_N_indices(b, 50) == collect(1:50)
@test get_last_N_indices(b, 51) == collect(1:50)
@test get_last_N_indices(b, 1000) == collect(1:50)

# full buffer
push!(b, d)
push!(b, d)
@test get_last_N_indices(b, 10) == collect(41:50)
@test get_last_N_indices(b, 1) == [50]
@test get_last_N_indices(b, 50) == collect(1:50)
@test get_last_N_indices(b, 51) == [100, collect(1:50)...]
@test get_last_N_indices(b, 100) == [collect(51:100)..., collect(1:50)...]
@test get_last_N_indices(b, 1000) == [collect(51:100)..., collect(1:50)...]

## Priority Param Construction
p = PriorityParams(100, α=0.7)
@test length(p.minsort_priorities) == 100
@test length(p.priorities) == 100
@test p.α == 0.7f0
@test p.max_priority == 1.0f0

p2 = PriorityParams(1000, p)
@test length(p2.priorities)==1000
@test p2.α == p.α
@test p2.max_priority == p.max_priority

## Construction
b = ExperienceBuffer(ContinuousSpace(2), DiscreteSpace(4), 100,)
bpriority = ExperienceBuffer(ContinuousSpace(2), DiscreteSpace(4), 50, prioritized=true)
b_gpu = b |> gpu

bcopy = ExperienceBuffer(deepcopy(b.data), elements=0)
@test b.data == bcopy.data
@test b.priority_params == bcopy.priority_params
bcopy = ExperienceBuffer(deepcopy(b.data), prioritized=true)
@test isprioritized(bcopy)

@test b isa ExperienceBuffer{Array}
@test b_gpu isa ExperienceBuffer{CuArray}

@test length(keys(b.data)) == length(keys(b_gpu.data))
@test !haskey(b, :weight)
@test :s in keys(b_gpu.data)
@test :a in keys(b_gpu.data)
@test :sp in keys(b_gpu.data)
@test :r in keys(b_gpu.data)
@test :done in keys(b_gpu.data)
@test haskey(bpriority, :weight)
@test bpriority.data[:weight] == ones(Float32, 1, 50)
@test size(b[:a]) == (4,0)

# Buffer_like function
bsmall = buffer_like(b, capacity=3, device=cpu)
@test Crux.device(bsmall) == cpu
@test keys(bsmall) == keys(b)
@test capacity(bsmall) == 3
@test length(bsmall) == 0

## Base functions
@test keys(b) == keys(b.data)

@test size(b[:s]) == (2,0)
@test length(b) == 0
@test capacity(b) == 100

@test size(b_gpu[:s]) == (2,0)
@test length(b_gpu) == 0
@test capacity(b_gpu) == 100

@test isprioritized(bpriority)
@test Crux.device(b) == cpu
@test Crux.device(bpriority) == cpu
@test Crux.device(b_gpu) == gpu

# Convert to and from GPU
@test Crux.device(gpu(b)) == gpu
@test Crux.device(cpu(b_gpu)) == cpu

## push!
#push dictionary with one element
d = Dict(:s => 2*ones(2,1), :a => ones(Bool, 4,1), :sp => ones(2,1), :r => ones(1,1), :done => zeros(1,1), :weight=>zeros(1,1))
push!(b, d)
@test length(b) == 1
@test b[:s] == 2*ones(2,1)
@test b[:a] == ones(Int, 4,1)
@test b[:sp] == ones(2,1)
@test b[:r] == ones(1,1)
@test b[:done] == zeros(1,1)

# push dictionary with more than one element
d = Dict(:s => 3*ones(2,3), :a => rand(4,3) .< 0.5, :sp => 5*ones(2,3), :r => 6*ones(1,3), :done => ones(1,3), :weight=>zeros(1,3))
push!(b, d)
@test length(b) == 4
@test b[:s][:,2:end] ==  3*ones(2,3)
@test b[:a][:,2:end] == d[:a]
@test b[:sp][:,2:end] == 5*ones(2,3)
@test b[:r][:,2:end] == 6*ones(1,3)
@test b[:done][:,2:end] == ones(1,3)

# push a buffer
push!(b, b)
@test length(b) == 8
for k in keys(b)
    @test b[k][:, 1:4] == b[k][:, 5:8]
end

## clear!
bcopy = deepcopy(b)
clear!(bcopy)
@test length(bcopy) == 0
@test bcopy.next_ind == 1

## minibatch
I = [1,2,4]
d2 = minibatch(b, I)
for k in keys(d2)
    @test all(d2[k] .== b[k][:, I])
end

## Split
@test Crux.split_batches(100, [0.5, 0.5]) == [50,50]
@test Crux.split_batches(100, [1.0]) == [100]
@test Crux.split_batches(100, [1/3, 1/3, 1/3]) == [34,33,33]
@test_throws AssertionError Crux.split_batches(100, 0.4)

v = rand(3,100)
bsplit = ExperienceBuffer(Dict(:a => v))

b1, b2 = split(bsplit, [0.5, 0.5])
@test length(b1) == 50
@test b1[:a] == v[:,1:50]
@test length(b2) == 50
@test b2[:a] == v[:,51:100]

## update_priorities!
update_priorities!(bpriority, [1,2,3], [1., 2., 3.])
@test bpriority.priority_params.max_priority == 3.0
@test bpriority.priority_params.priorities[1] ≈ 1f0^bpriority.priority_params.α
@test bpriority.priority_params.priorities[2] ≈ 2f0^bpriority.priority_params.α
@test bpriority.priority_params.priorities[3] ≈ 3f0^bpriority.priority_params.α
@test bpriority.priority_params.minsort_priorities[1] ≈ 1f0^bpriority.priority_params.α
@test bpriority.priority_params.minsort_priorities[2] ≈ 2f0^bpriority.priority_params.α
@test bpriority.priority_params.minsort_priorities[3] ≈ 3f0^bpriority.priority_params.α

push!(bpriority, d)
d
push!(bpriority, d)
@test bpriority.priority_params.max_priority == 3.0
for i=1:6
    @test bpriority.priority_params.priorities[i] ≈ 3f0^bpriority.priority_params.α
    @test bpriority.priority_params.minsort_priorities[i] ≈ 3f0^bpriority.priority_params.α
end

## sampling

# uniform sample
t = ExperienceBuffer(ContinuousSpace(2), DiscreteSpace(4), 10)
rand!(t, b)

t = ExperienceBuffer(ContinuousSpace(2), DiscreteSpace(4), 3)
Random.seed!(0)
ids = rand(1:length(b), 3)
Random.seed!(0)
rand!(t, b)

for k in keys(t)
    @test t[k] == b[k][:,ids]
end

# Test the multi-buffer sampling
t1 = ExperienceBuffer(ContinuousSpace(2), DiscreteSpace(4), 10)
d = Dict(:s => ones(2,1), :a => ones(Bool, 4,1), :sp => ones(2,1), :r => ones(1,1), :done => zeros(1,1))
push!(t1, d)

t2 = ExperienceBuffer(ContinuousSpace(2), DiscreteSpace(4), 10)
d = Dict(:s => 2*ones(2,1), :a => ones(Bool, 4,1), :sp => ones(2,1), :r => ones(1,1), :done => zeros(1,1))
push!(t2, d)

t3 = ExperienceBuffer(ContinuousSpace(2), DiscreteSpace(4), 10)
d = Dict(:s => 3*ones(2,1), :a => ones(4,1), :sp => ones(2,1), :r => ones(1,1), :done => zeros(1,1))
push!(t3, d)

t = ExperienceBuffer(ContinuousSpace(2), DiscreteSpace(4), 10)
rand!(t, t1, t2, t3)

@test all(t[:s][:, 1:4] .== 1.0)
@test all(t[:s][:, 5:7] .== 2.0)
@test all(t[:s][:, 8:10] .== 3.0)

# Priority samples
bpriority[:s] .= rand(Float32, 2, 6)
update_priorities!(bpriority, [1:6...], [1.:6. ...])
t = ExperienceBuffer(ContinuousSpace(2), DiscreteSpace(4), 1000, [:weight])
priorities = [bpriority.priority_params.priorities[i] for i=1:length(bpriority)]

rand!(t, bpriority)

# Test the the frequency of samples is proportional to their priority
length(t)
freqs = [sum(t.indices .== i) for i=1:6] ./ length(t)
probs = priorities ./ sum(priorities)
relerr = abs.(freqs .- probs) ./ probs
@test all(relerr .< 0.01)

@test all(t[:s] .== bpriority[:s][:,t.indices])
@test all(t[:weight] .<= 1.)

## merge
# bmerge = merge(b, b, capacity = 300)
# @test length(bmerge) == 200
# @test capacity(bmerge) == 300
# @test bmerge[:s][:, 1:100] == bmerge[:s][:, 101:200]

## Multi-D states
b2d = ExperienceBuffer(ContinuousSpace((2,2), UInt8), DiscreteSpace(4), 100;)
@test ndims(b2d[:s]) == 3
b2d[:s]

d = Dict(:s => 3*ones(2,2,1), :a => ones(4,1), :sp => ones(2,2,1), :r => ones(1,1), :done => zeros(1,1))
push!(b2d, d)

@test all(b2d[:s] .== 3)


# # Indexing
containers = Dict("<:AbstractRange" => 1:1000,
                  "StaticRange" => srange(1,1000,step=1),
                  "Type{StaticRange}" => OneToSRange{1000},
                  "SVector" => SVector(1:1000...),
                  "Vector" => [1:1000...],
                  "Matrix" => rand(Int, 20, 50),
                  "SMatrix" => SMatrix{20,50}(1:1000...));

by_r(x) = x[range(1, 500, step=1)];
by_sr(x) = x[srange(1, 500, step=1)];
# TODO: should this be a thing?
#by_srtype(x) = x[OneToSRange{500}];

by_step_r(x) = x[range(1, 500, step=2)];
by_step_sr(x) = x[srange(1, 500, step=2)];

fxns = Dict("range" => by_r,
            "srange" => by_sr,
            "step_range" => by_step_r,
            "step_srange" => by_step_sr)
by_nd_r(x::AbstractMatrix) = x[range(1, 20, step=1), range(1, 20, step=1)];
by_nd_sr(x::AbstractMatrix) = x[srange(1, 20, step=1), srange(1, 20, step=1)];

SRBench["indexing"] = BenchmarkGroup();
for  (ck,cv) in containers
    for (fk, fv) in fxns
        if !haskey(SRBench["indexing"], ck)
            SRBench["indexing"][ck] = BenchmarkGroup([ck])
        end
        SRBench["indexing"][ck][fk] = @benchmarkable $(fv)($cv)
    end
end

SRBench["indexing"]["Matrix"]["by_nd_r"] = @benchmarkable by_nd_r($(containers["Matrix"]))
SRBench["indexing"]["Matrix"]["by_nd_sr"] = @benchmarkable by_nd_sr($(containers["Matrix"]))
SRBench["indexing"]["SMatrix"]["by_nd_r"] = @benchmarkable by_nd_r($(containers["SMatrix"]))
SRBench["indexing"]["SMatrix"]["by_nd_sr"] = @benchmarkable by_nd_sr($(containers["SMatrix"]))


### A Pluto.jl notebook ###
# v0.14.0

using Markdown
using InteractiveUtils

# ╔═╡ 030f5512-2a66-11ec-1a7a-21059e6facac
begin

# Import packages used

using Random
using DataFrames
using CSV
using StatFiles
using Statistics

end

# ╔═╡ af9b5960-2a65-11ec-3015-53185d92610c
# Replication of DellaVigna and Pope, 2018, "What Motivates Effort? Evidence and Expert Forecasts"

# The code in this notebook performs the estimates to replicate columns (1) (3) panel A and columns (1) (4) panel B of Table 5 using minimum distance.

# ╔═╡ 03372860-2a66-11ec-2b31-298b2c9afc9d
# Authors:
# Massimiliano Pozzi (Bocconi University, pozzi.massimiliano@studbocconi.it)
# Salvatore Nunnari (Bocconi University, https://snunnari.github.io, salvatore.nunnari@unibocconi.it)

# ╔═╡ 7c88c030-55f8-11ec-0159-99a5b81bf850
# This notebook works with both Neptune and Pluto; the "begin" and "end" tags in each cell are necessary for Pluto notebooks but do not play any role in Neptune

# ╔═╡ 7dcdea60-55f8-11ec-2fdc-1d683f29a01c
# This notebook was tested with the following packages versions:

# Pozzi: julia 1.6.1, Pluto 0.14.7, Neptune 0.14.0, DataFrames 1.2.2, CSV 0.9.5

# ╔═╡ 03008800-2a66-11ec-0c80-1767685eb4fa
begin

# 1. Data Cleaning and Data Preparation

# We import the relevant dataset containing data on the number of buttonpresses in the different treatments and for different piece rates wage that the participants received when completing the task. We then compute the means for the treatments that are used to compute the minimum distance estimates.

dt = DataFrame(load(raw"../input/mturk_clean_data_short.dta"))

emp_moments = round.(combine(groupby(dt, :treatment), :buttonpresses => mean)[2], digits=0)

end;

# ╔═╡ 944bb54e-55f9-11ec-32b3-c953374ecd8f
# We set up the Bootstrap procedure for the standard errors. The following cell creates 2000 new samples and computes the mean for the relevant treatments.

# ╔═╡ 030039e0-2a66-11ec-3766-ad9593bffbd3
begin

# Treatments are as follows:
# 1.1: benchmark specification with piece rate of 0.01
# 1.2: benchmark specification with piece rate of 0.10
# 1.3: benchmark specification with piece rate of 0.00
# 3.1: social preferences (charity) with piece rate of 0.01
# 3.2: social preferences (charity) with piece rate of 0.10
# 10 : social preferences (gift exchange) bonus of 40 cents (independently of nr buttonpresses)
# 4.1: time discounting with extra 0.01 paid two weeks later
# 4.2: time discounting with extra 0.01 paid four weeks later

# define the vector containing the treatments (used for the loop) and the vectors that will store the rounded mean buttonpresses in our 2000 new samples.
# E11 is a 1x1000 vector that will contain the average buttonpresses in each of our 1000 samples for treatment 1.1 etc.

treatment = ["1.1", "1.2", "1.3", "3.1", "3.2", "10", "4.1", "4.2"]
E11, E12, E13, E31, E32, E10, E41, E42 = Int32[], Int32[], Int32[], Int32[], Int32[], Int32[], Int32[], Int32[]

# We first get a smaller dataframe containing only observations for a specific treatment, we then resample the observations and compute the rounded mean of buttonpresses, save the result and then pass onto the next treatment. we do this 1000 times.

for i=1:1000                 # we want 1000 new samples
    for treat in treatment   # we want to compute the mean for all treatments
        db = copy(dt)        # copy the original dataset
        db = db[db[:,2] .== treat,:] # keep only observation for treatment 'treat'
        bootsample = db[vec(rand(1:size(db[!,1],1),1,size(db[!,1],1))),1] # resample the dataset by taking random rows from the first column
		if treat == "1.1" append!(E11, round(mean(bootsample))) end
        if treat == "1.2" append!(E12, round(mean(bootsample))) end
        if treat == "1.3" append!(E13, round(mean(bootsample))) end
        if treat == "3.1" append!(E31, round(mean(bootsample))) end
        if treat == "3.2" append!(E32, round(mean(bootsample))) end
        if treat == "10"  append!(E10, round(mean(bootsample))) end
        if treat == "4.1" append!(E41, round(mean(bootsample))) end
        if treat == "4.2" append!(E42, round(mean(bootsample))) end
	end
end

end

# ╔═╡ ced0482e-2a66-11ec-06ca-0b961f35a057
# For a more detailed explanation of the model we refer the reader to the python notebook

# Define the function to estimate the parameters using minimum distance:

# There are two specifications: exponential cost function and power cost function
# The parameters are found by imposing the theoretical means found from the agent's problem being equal to the empirical means found in the data.
# E11 up to E42 are the relevant empirical moments
# Specification can either be 'Exp' or 'Power'
# P is a vector containing the different piece-rates
# To simplify the computations the authors take the log of the variables to estimate gamma k and s.

function mindisest(E11,E12,E13,E31,E32,E10,E41,E42,specification)

    P = [0, 0.01, 0.1]

	# gamma, k and s are found by solving the system of N equations in N unknowns

    if specification == "Exp"
        log_k = (log(P[3]) - log(P[2])*(E12)/(E11))/(1 - (E12)/(E11))
        log_gamma = log((log(P[2]) - log_k)/(E11))
        log_s = exp(log_gamma)*(E13) + log_k
	end

    if specification == "Power"
        log_k = (log(P[3]) - log(P[2])*log(E12)/log(E11))/(1 - log(E12)/log(E11))
        log_gamma = log((log(P[2]) - log_k)/log(E11))
        log_s = exp(log_gamma)*log(E13) + log_k
	end

    k = exp(log_k)
    g = exp(log_gamma)
    s = exp(log_s)

    if specification == "Exp"
        EG31, EG32, EG10, EG41, EG42 = exp.(E31*g), exp.(E32*g), exp.(E10*g), exp.(E41*g), exp.(E42*g)

        alpha = 100/9*k*(EG32-EG31)
        a = 100*k*EG31-100*s-alpha
        s_ge = k*EG10 - s
        delta = sqrt((k*EG42-s)/(k*EG41-s))
        beta  = 100*(k*EG41-s)/(delta^2)
	end

    if specification == "Power"
        alpha = 100/9*k*(E32^g-E31^g)
        a = 100*k*E31^g-100*s-alpha
        s_ge = k*E10^g - s
        delta = sqrt((k*E42^g-s)/(k*E41^g-s))
        beta  = 100*(k*E41^g-s)/(delta^2)
	end

    return k, g, s, alpha, a, s_ge, beta, delta
end

# ╔═╡ 9ac87bc0-5608-11ec-3f39-3713d508fdfb
begin

# Table 5 minimum distance estimates: columns (1) (3) panel A and columns (1) (4) panel B

Table5Exp = mindisest(emp_moments[1],emp_moments[2],emp_moments[3],emp_moments[7],emp_moments[8],emp_moments[5],emp_moments[9],emp_moments[10],"Exp")

Table5Power = mindisest(emp_moments[1],emp_moments[2],emp_moments[3],emp_moments[7],emp_moments[8],emp_moments[5],emp_moments[9],emp_moments[10],"Power")

end;

# ╔═╡ 4b1369f0-2a70-11ec-2cf3-e514a4899e06
begin

# Store mean and standard error of estimates for the exponential cost function specification. Our function returns an array containing 1000 vectors of 8 elements each. The estimate for k is in position 1 in each 1000 vectors etc

estimatesExp = mindisest.(E11,E12,E13,E31,E32,E10,E41,E42,"Exp")
k_exp_mean = mean([item[1] for item in estimatesExp])
k_exp_sd = std([item[1] for item in estimatesExp])
g_exp_mean = mean([item[2] for item in estimatesExp])
g_exp_sd = std([item[2] for item in estimatesExp])
s_exp_mean = mean([item[3] for item in estimatesExp])
s_exp_sd = std([item[3] for item in estimatesExp])
alpha_exp_mean = mean([item[4] for item in estimatesExp])
alpha_exp_sd = std([item[4] for item in estimatesExp])
a_exp_mean = mean([item[5] for item in estimatesExp])
a_exp_sd = std([item[5] for item in estimatesExp])
s_ge_exp_mean = mean([item[6] for item in estimatesExp])
s_ge_exp_sd = std([item[6] for item in estimatesExp])
beta_exp_mean = mean([item[7] for item in estimatesExp])
beta_exp_sd = std([item[7] for item in estimatesExp])
delta_exp_mean = mean([item[8] for item in estimatesExp])
delta_exp_sd = std([item[8] for item in estimatesExp])

end;

# ╔═╡ 2e15f490-2a6f-11ec-3bec-0f631d9fbd70
begin

# Store mean and standard error of estimates for the power cost function specification

estimatesPower = mindisest.(E11,E12,E13,E31,E32,E10,E41,E42,"Power")
k_power_mean = mean([item[1] for item in estimatesPower])
k_power_sd = std([item[1] for item in estimatesPower])
g_power_mean = mean([item[2] for item in estimatesPower])
g_power_sd = std([item[2] for item in estimatesPower])
s_power_mean = mean([item[3] for item in estimatesPower])
s_power_sd = std([item[3] for item in estimatesPower])

# A couple of elements in the arrays take value "NaN". We need to remove them before computing the mean and std

alpha_power = [item[4] for item in estimatesPower]
alpha_power_mean = mean(alpha_power[(!isnan).(alpha_power)])
alpha_power_sd = std(alpha_power[(!isnan).(alpha_power)])
a_power = [item[5] for item in estimatesPower]
a_power_mean = mean(a_power[(!isnan).(a_power)])
a_power_sd = std(a_power[(!isnan).(a_power)])
s_ge_power = [item[6] for item in estimatesPower]
s_ge_power_mean = mean(s_ge_power[(!isnan).(s_ge_power)])
s_ge_power_sd = std(s_ge_power[(!isnan).(s_ge_power)])
beta_power = [item[7] for item in estimatesPower]
beta_power_mean = mean(beta_power[(!isnan).(beta_power)])
beta_power_sd = std(beta_power[(!isnan).(beta_power)])
delta_power = [item[8] for item in estimatesPower]
delta_power_mean = mean(delta_power[(!isnan).(delta_power)])
delta_power_sd = std(delta_power[(!isnan).(delta_power)])

end;

# ╔═╡ c2d506c0-2a6f-11ec-0871-fdb55e56de67
begin

# To obtain confidence intervals in panel B table 5. CI are the 2.5% and 97.5% quantiles of the distribution of our parameters vectors. We obtain CI only for alpha, a, s_ge, beta, delta as in the paper

CI_Exp, CI_Power = [], []
for ci=4:8
    a  = sort!([item[ci] for item in estimatesExp]) # sort from smallest to biggest number
    append!(CI_Exp, [quantile!(a,0.025) ,quantile!(a,0.975)])
end
vecest = [alpha_power, a_power, s_ge_power, beta_power, delta_power]
for ci in vecest
    a  = sort!(ci[(!isnan).(ci)])
    append!(CI_Power, [quantile!(a,0.025) ,quantile!(a,0.975)])
end

end

# ╔═╡ a9cedae0-2a7c-11ec-3b2f-a3577f3decf7
begin

# Print our results for Table 5 (GMM estimates)

# Standard errors are different since the seed we used for the bootstrap procedure is different from the one used by the authors since random generation across softwares/languages is not easily replicated (each software uses its own algorithm)

sd_exp = [k_exp_sd,g_exp_sd,s_exp_sd,alpha_exp_sd,a_exp_sd,s_ge_exp_sd,beta_exp_sd,delta_exp_sd]

sd_power = [k_power_sd,g_power_sd,s_power_sd,alpha_power_sd,a_power_sd,s_ge_power_sd,beta_power_sd,delta_power_sd]

params_name = ["Level k of cost of effort", "Curvature γ of cost function","Intrinsic motivation s","Social preferences α", "Warm glow coefficient a","Gift exchange Δs", "Present bias β","(Weekly) discount factor δ"]

Table5Results = DataFrame()
Table5Results.params = params_name

# Table5Power is saved as a NTuple{8,Float64} but DataFrame only accepts vectors as columns. Using Collect transforms the NTuple into a vector

Table5Results.estimatesPower = collect(Table5Power)
Table5Results.sdPower = sd_power
Table5Results.estimatesExp = collect(Table5Exp)
Table5Results.sdExp = sd_exp
Table5Results

end

# ╔═╡ 88d81640-5609-11ec-1234-23840357ca64
# Save Table5Results as a csv file

CSV.write(raw"../output/table5GMM_julia.csv", Table5Results)

# ╔═╡ Cell order:
# ╠═af9b5960-2a65-11ec-3015-53185d92610c
# ╠═03372860-2a66-11ec-2b31-298b2c9afc9d
# ╠═7c88c030-55f8-11ec-0159-99a5b81bf850
# ╠═7dcdea60-55f8-11ec-2fdc-1d683f29a01c
# ╠═030f5512-2a66-11ec-1a7a-21059e6facac
# ╠═03008800-2a66-11ec-0c80-1767685eb4fa
# ╠═944bb54e-55f9-11ec-32b3-c953374ecd8f
# ╠═030039e0-2a66-11ec-3766-ad9593bffbd3
# ╠═ced0482e-2a66-11ec-06ca-0b961f35a057
# ╠═9ac87bc0-5608-11ec-3f39-3713d508fdfb
# ╠═4b1369f0-2a70-11ec-2cf3-e514a4899e06
# ╠═2e15f490-2a6f-11ec-3bec-0f631d9fbd70
# ╠═c2d506c0-2a6f-11ec-0871-fdb55e56de67
# ╠═a9cedae0-2a7c-11ec-3b2f-a3577f3decf7
# ╠═88d81640-5609-11ec-1234-23840357ca64


module MathOptInterfaceMosek

import MathOptInterface
using Mosek

immutable MosekSolver <: MathOptInterface.AbstractSolver
  options
end

MosekSolver(;kwargs...) = MosekSolver(kwargs)
export MosekSolver

include("LinkedInts.jl")

const DEBUG = false

mosek_block_type_unallocated = 0
mosek_block_type_zero   = 1
mosek_block_type_nonneg = 2
mosek_block_type_nonpos = 3
mosek_block_type_range  = 4
mosek_block_type_qcone  = 5
mosek_block_type_rqcone = 6
mosek_block_type_psd    = 7
mosek_block_type_integer = 8

problemtype_linear    = 0
problemtype_conic     = 1
problemtype_quadratic = 2

import MathOptInterface

boundflag_lower = 0x1
boundflag_upper = 0x2
boundflag_cone  = 0x4
boundflag_int   = 0x8
boundflag_all   = 0x0f

# Mapping of all constraint types to its index
struct ConstraintMap
    x_lessthan           :: Dict{UInt64,Int} # (SingleVariable,LessThan) -> constraint number
    x_greaterthan        :: Dict{UInt64,Int} # (SingleVariable,GreaterThan) -> constraint number
    x_equalto            :: Dict{UInt64,Int} # (SingleVariable,EqualTo) -> constraint number
    x_interval           :: Dict{UInt64,Int} # (SingleVariable,Interval) -> constraint number
    x_nonpositives       :: Dict{UInt64,Int} # (SingleVariable,Nonpositives) -> constraint number
    x_nonnegatives       :: Dict{UInt64,Int} # (SingleVariable,Nonnegatives) -> constraint number
    x_binary             :: Dict{UInt64,Int} # (SingleVariable,ZeroOne) -> constraint number
    x_integer            :: Dict{UInt64,Int} # (SingleVariable,Integer) -> constraint number

    xs_nonpositives      :: Dict{UInt64,Int} # (VectorOfVariables,Nonpositives) -> constraint number
    xs_nonnegatives      :: Dict{UInt64,Int} # (VectorOfVariables,Nonnegatives) -> constraint number
    xs_zeros             :: Dict{UInt64,Int} # (VectorOfVariables,Zeros) -> constraint number
    xs_reals             :: Dict{UInt64,Int} # (VectorOfVariables,Reals) -> constraint number
    xs_qcone             :: Dict{UInt64,Int} # (VectorOfVariables,SecondOrderCone) -> constraint number
    xs_rqcone            :: Dict{UInt64,Int} # (VectorOfVariables,RotatedSecondOrderCone) -> constraint number
    xs_psdconetriangle   :: Dict{UInt64,Int} # (VectorOfVariables,PositiveSemidefiniteConeTriangle) -> constraint number

    axb_lessthan         :: Dict{UInt64,Int} # (ScalarAffineFunction,LessThan) -> constraint number
    axb_greaterthan      :: Dict{UInt64,Int} # (ScalarAffineFunction,GreaterThan) -> constraint number
    axb_equalto          :: Dict{UInt64,Int} # (ScalarAffineFunction,EqualTo) -> constraint number
    axb_interval         :: Dict{UInt64,Int} # (ScalarAffineFunction,Interval) -> constraint number
    axb_nonpositives     :: Dict{UInt64,Int} # (ScalarAffineFunction,Nonpositives) -> constraint number
    axb_nonnegatives     :: Dict{UInt64,Int} # (ScalarAffineFunction,Nonnegatives) -> constraint number
    axb_binary           :: Dict{UInt64,Int} # (ScalarAffineFunction,ZeroOne) -> constraint number
    axb_integer          :: Dict{UInt64,Int} # (ScalarAffineFunction,Integer) -> constraint number

    axbs_nonpositives    :: Dict{UInt64,Int} # (VectorAffineFunction,Nonpositives) -> constraint number
    axbs_nonnegatives    :: Dict{UInt64,Int} # (VectorAffineFunction,Nonnegatives) -> constraint number
    axbs_zeros           :: Dict{UInt64,Int} # (VectorAffineFunction,Zeros) -> constraint number
    axbs_reals           :: Dict{UInt64,Int} # (VectorAffineFunction,Reals) -> constraint number
    axbs_qcone           :: Dict{UInt64,Int} # (VectorAffineFunction,SecondOrderCone) -> constraint number
    axbs_rqcone          :: Dict{UInt64,Int} # (VectorAffineFunction,RotatedSecondOrderCone) -> constraint number
    axbs_psdconetriangle :: Dict{UInt64,Int} # (VectorAffineFunction,PositiveSemidefiniteConeTriangle) -> constraint number
end

ConstraintMap() = ConstraintMap([Dict{UInt64,Int}() for i in 1:30]...)
select(cm::ConstraintMap,::Type{MathOptInterface.SingleVariable},               ::Type{MathOptInterface.LessThan{Float64}}) =                cm.x_lessthan
select(cm::ConstraintMap,::Type{MathOptInterface.SingleVariable},               ::Type{MathOptInterface.GreaterThan{Float64}}) =             cm.x_greaterthan
select(cm::ConstraintMap,::Type{MathOptInterface.SingleVariable},               ::Type{MathOptInterface.EqualTo{Float64}}) =                 cm.x_equalto
select(cm::ConstraintMap,::Type{MathOptInterface.SingleVariable},               ::Type{MathOptInterface.Interval{Float64}}) =                cm.x_interval
select(cm::ConstraintMap,::Type{MathOptInterface.SingleVariable},               ::Type{MathOptInterface.Nonpositives}) =                     cm.x_nonpositives
select(cm::ConstraintMap,::Type{MathOptInterface.SingleVariable},               ::Type{MathOptInterface.Nonnegatives}) =                     cm.x_nonnegatives
select(cm::ConstraintMap,::Type{MathOptInterface.SingleVariable},               ::Type{MathOptInterface.ZeroOne}) =                          cm.x_binary
select(cm::ConstraintMap,::Type{MathOptInterface.SingleVariable},               ::Type{MathOptInterface.Integer}) =                          cm.x_integer
select(cm::ConstraintMap,::Type{MathOptInterface.VectorOfVariables},            ::Type{MathOptInterface.Nonpositives}) =                     cm.xs_nonpositives
select(cm::ConstraintMap,::Type{MathOptInterface.VectorOfVariables},            ::Type{MathOptInterface.Nonnegatives}) =                     cm.xs_nonnegatives
select(cm::ConstraintMap,::Type{MathOptInterface.VectorOfVariables},            ::Type{MathOptInterface.Zeros}) =                            cm.xs_zeros
select(cm::ConstraintMap,::Type{MathOptInterface.VectorOfVariables},            ::Type{MathOptInterface.Reals}) =                            cm.xs_reals
select(cm::ConstraintMap,::Type{MathOptInterface.VectorOfVariables},            ::Type{MathOptInterface.SecondOrderCone}) =                  cm.xs_qcone
select(cm::ConstraintMap,::Type{MathOptInterface.VectorOfVariables},            ::Type{MathOptInterface.RotatedSecondOrderCone}) =           cm.xs_rqcone
select(cm::ConstraintMap,::Type{MathOptInterface.VectorOfVariables},            ::Type{MathOptInterface.PositiveSemidefiniteConeTriangle}) = cm.xs_psdconetriangle
select(cm::ConstraintMap,::Type{MathOptInterface.ScalarAffineFunction{Float64}},::Type{MathOptInterface.LessThan{Float64}}) =                cm.axb_lessthan
select(cm::ConstraintMap,::Type{MathOptInterface.ScalarAffineFunction{Float64}},::Type{MathOptInterface.GreaterThan{Float64}}) =             cm.axb_greaterthan
select(cm::ConstraintMap,::Type{MathOptInterface.ScalarAffineFunction{Float64}},::Type{MathOptInterface.EqualTo{Float64}}) =                 cm.axb_equalto
select(cm::ConstraintMap,::Type{MathOptInterface.ScalarAffineFunction{Float64}},::Type{MathOptInterface.Interval{Float64}}) =                cm.axb_interval
select(cm::ConstraintMap,::Type{MathOptInterface.ScalarAffineFunction{Float64}},::Type{MathOptInterface.ZeroOne}) =                          cm.axb_binary
select(cm::ConstraintMap,::Type{MathOptInterface.ScalarAffineFunction{Float64}},::Type{MathOptInterface.Integer}) =                          cm.axb_integer
select(cm::ConstraintMap,::Type{MathOptInterface.ScalarAffineFunction{Float64}},::Type{MathOptInterface.Nonpositives}) =                     cm.axb_nonpositives
select(cm::ConstraintMap,::Type{MathOptInterface.ScalarAffineFunction{Float64}},::Type{MathOptInterface.Nonnegatives}) =                     cm.axb_nonnegatives
select(cm::ConstraintMap,::Type{MathOptInterface.VectorAffineFunction{Float64}},::Type{MathOptInterface.Nonpositives}) =                     cm.axbs_nonpositives
select(cm::ConstraintMap,::Type{MathOptInterface.VectorAffineFunction{Float64}},::Type{MathOptInterface.Nonnegatives}) =                     cm.axbs_nonnegatives
select(cm::ConstraintMap,::Type{MathOptInterface.VectorAffineFunction{Float64}},::Type{MathOptInterface.Zeros}) =                            cm.axbs_zeros
select(cm::ConstraintMap,::Type{MathOptInterface.VectorAffineFunction{Float64}},::Type{MathOptInterface.Reals}) =                            cm.axbs_reals
select(cm::ConstraintMap,::Type{MathOptInterface.VectorAffineFunction{Float64}},::Type{MathOptInterface.SecondOrderCone}) =                  cm.axbs_qcone
select(cm::ConstraintMap,::Type{MathOptInterface.VectorAffineFunction{Float64}},::Type{MathOptInterface.RotatedSecondOrderCone}) =           cm.axbs_rqcone
select(cm::ConstraintMap,::Type{MathOptInterface.VectorAffineFunction{Float64}},::Type{MathOptInterface.PositiveSemidefiniteConeTriangle}) = cm.axbs_psdconetriangle

Base.getindex{F,D}(cm::ConstraintMap,r :: MathOptInterface.ConstraintReference{F,D}) = select(cm,F,D)[r.value]

struct MosekSolution
    whichsol :: Soltype
    solsta   :: Solsta
    prosta   :: Prosta

    xxstatus :: Vector{Stakey}
    xx       :: Vector{Float64}
    slx      :: Vector{Float64}
    sux      :: Vector{Float64}
    snx      :: Vector{Float64}

    cstatus  :: Vector{Stakey}
    xc       :: Vector{Float64}
    slc      :: Vector{Float64}
    suc      :: Vector{Float64}
    y        :: Vector{Float64}
end

"""
    MosekModel <: MathOptInterface.AbstractModel

Linear variables and constraint can be deleted. For some reason MOSEK
does not support deleting PSD variables.

Note also that adding variables and constraints will permanently add
some (currently between 1 and 3) Int64s that a `delete!` will not
remove. This ensures that References (Variable and constraint) that
are deleted are thereafter invalid.
"""
mutable struct MosekModel  <: MathOptInterface.AbstractSolverInstance
    task :: Mosek.MSKtask

    """
    Number of variables explicitly created by user
    """
    publicnumvar :: Int

    """
    """
    constrmap :: ConstraintMap

    """
    The total length of `x_block` matches the number of variables in
    the underlying task, and the number of blocks corresponds to the
    number variables allocated in the Model.
    """
    x_block      :: LinkedInts

    """
    One entry per scalar variable in the task indicating which bound
    types are defined
    """
    x_boundflags :: Vector{Int}

    """
    One entry per scalar variable in the task defining the number of
    variable constraints imposed on that variable. It is not allowed
    to delete a variable without deleting the constraints on it first
    (to get around the problem of deleting roots in conic constraints).
    """
    x_numxc :: Vector{Int}

    """
    One entry per variable-constraint
    """
    xc_block     :: LinkedInts

    """
    One entry per variable-constraint block indicating which bound
    types it defines. The values are binary ORed `boundflag_...` values.
    """
    xc_bounds    :: Vector{UInt8}
    """
    One entry per variable-constraint block indicating which cone it belongs to, 0 if none.
    """
    xc_coneid    :: Vector{Int}

    """
    One entry per scalar variable-constraint, defining which native
    variables the bound block covers.
    """
    xc_idxs      :: Vector{Int}

    ###########################
    """
    One scalar entry per constraint in the underlying task. One block
    per constraint allocated in the Model.
    """
    c_block :: LinkedInts

    """
    One entry per allocated scalar constraint. Defines the fixed term
    on the left-hand for each scalar constraint.
    """
    c_constant :: Vector{Float64}

    """
    One entry per allocated scalar constraint.
    Each element is either
    - 0: meaning: no slack, when domain is defined directly as a bound,
    - positive: a `x_block` reference, e.g. for qcones, or
    - negative: Negated index of a PSD variable in the underlying task
    """
    c_block_slack   :: Vector{Int}
    """
    One entry per variable-constraint block indicating which cone it belongs to, 0 if none.
    """
    c_coneid   :: Vector{Int}

    ###########################
    conecounter :: Int

    ###########################
    trm :: Rescode
    solutions :: Vector{MosekSolution}
end

function MathOptInterface.SolverInstance(solver::MosekSolver)
    t = maketask()
    try
        be_quiet = false
        for (option,val) in solver.options
            parname = string(option)
            if parname == "QUIET"
                be_quiet = be_quiet || convert(Bool,val)
            elseif startswith(parname, "MSK_IPAR_")
                Mosek.putnaintparam(t, parname, convert(Integer, val))
            elseif startswith(parname, "MSK_DPAR_")
                Mosek.putnadouparam(t, parname, convert(AbstractFloat, val))
            elseif startswith(parname, "MSK_SPAR_")
                Mosek.putnastrparam(t, parname, convert(AbstractString, val))
            elseif isa(val, Integer)
                parname = "MSK_IPAR_$parname"
                Mosek.putnaintparam(t, parname, val)
            elseif isa(val, AbstractFloat)
                parname = "MSK_DPAR_$parname"
                Mosek.putnadouparam(t, parname, val)
            elseif isa(val, AbstractString)
                parname = "MSK_SPAR_$parname"
                Mosek.putnastrparam(t, parname, val)
            else
                error("Value $val for parameter $option has unrecognized type")
            end
        end
        if ! be_quiet
            Mosek.putstreamfunc(t,Mosek.MSK_STREAM_LOG,printstream)
        end
        MosekModel(t,# task
                   0, # public numvar
                   ConstraintMap(), # public constraints
                   LinkedInts(),# c_block
                   Int[], # x_boundflags
                   Int[], # x_boundflags
                   LinkedInts(), # xc_block
                   UInt8[], # xc_bounds
                   Int[], # xc_coneid
                   Int[], # xc_idxs
                   LinkedInts(), # c_block
                   Float64[], # c_constant
                   Int[], # c_block_slack
                   Int[], # c_coneid
                   0, # cone counter
                   Mosek.MSK_RES_OK,
                   MosekSolution[]) # trm
    catch
        deletetask(m.task)
        rethrow()
    end
end

function MathOptInterface.free!(m::MosekModel)
    deletetask(m.task)
end

function MathOptInterface.optimize!(m::MosekModel)
    m.trm = optimize(m.task)
    m.solutions = MosekSolution[]
    if solutiondef(m.task,MSK_SOL_ITG)

        push!(m.solutions,
              MosekSolution(MSK_SOL_ITG,
                            getsolsta(m.task,MSK_SOL_ITG),
                            getprosta(m.task,MSK_SOL_ITG),
                            getskx(m.task,MSK_SOL_ITG),
                            getxx(m.task,MSK_SOL_ITG),
                            Float64[],
                            Float64[],
                            Float64[],
                            getskc(m.task,MSK_SOL_ITG),
                            getxc(m.task,MSK_SOL_ITG),
                            Float64[],
                            Float64[],
                            Float64[]))
    end
    if solutiondef(m.task,MSK_SOL_BAS)
        push!(m.solutions,
              MosekSolution(MSK_SOL_BAS,
                            getsolsta(m.task,MSK_SOL_BAS),
                            getprosta(m.task,MSK_SOL_BAS),
                            getskx(m.task,MSK_SOL_BAS),
                            getxx(m.task,MSK_SOL_BAS),
                            getslx(m.task,MSK_SOL_BAS),
                            getsux(m.task,MSK_SOL_BAS),
                            Float64[],
                            getskc(m.task,MSK_SOL_BAS),
                            getxc(m.task,MSK_SOL_BAS),
                            getslc(m.task,MSK_SOL_BAS),
                            getsuc(m.task,MSK_SOL_BAS),
                            gety(m.task,MSK_SOL_BAS)))
    end
    if solutiondef(m.task,MSK_SOL_ITR)
        push!(m.solutions,
              MosekSolution(MSK_SOL_ITR,
                            getsolsta(m.task,MSK_SOL_ITR),
                            getprosta(m.task,MSK_SOL_ITR),
                            getskx(m.task,MSK_SOL_ITR),
                            getxx(m.task,MSK_SOL_ITR),
                            getslx(m.task,MSK_SOL_ITR),
                            getsux(m.task,MSK_SOL_ITR),
                            getsnx(m.task,MSK_SOL_ITR),
                            getskc(m.task,MSK_SOL_ITR),
                            getxc(m.task,MSK_SOL_ITR),
                            getslc(m.task,MSK_SOL_ITR),
                            getsuc(m.task,MSK_SOL_ITR),
                            gety(m.task,MSK_SOL_ITR)))
    end
end

function MathOptInterface.writeproblem(m::MosekModel, filename :: String)
    putintparam(m.task,MSK_IPAR_OPF_WRITE_SOLUTIONS, MSK_ON)
    writedata(m.task,filename)
end

# For linear objectives we accept:
# EITER affine left-hand side and ranged, unbounded, half-open, fixed (equality), PSD or SOC domains
# OR affine and quadratic left-hand side, and ranged, unbounded, half-open, fixed (equality) domains (quadratic constraints must be unbounded or half-open)
#
# For non-quadratic problems we allow binary and integer variables (but not constraints)
function supportsconstraints(m::MosekSolver, constraint_types) :: Bool
    for (fun,dom) in constraint_types
        if  fun in [MathOptInterface.ScalarAffineFunction{Float64},
                    MathOptInterface.SingleVariable,
                    MathOptInterface.VectorAffineFunction{Float64}] &&
            dom in [MathOptInterface.Zeros,
                    MathOptInterface.Reals,
                    MathOptInterface.Nonnegatives,
                    MathOptInterface.Nonpositives,
                    MathOptInterface.GreaterThan{Float64},
                    MathOptInterface.LessThan{Float64},
                    MathOptInterface.EqualTo{Float64},
                    MathOptInterface.Interval{Float64},
                    MathOptInterface.SecondOrderCone,
                    MathOptInterface.RotatedSecondOrderCone,
                    MathOptInterface.PositiveSemidefiniteConeTriangle,
                    MathOptInterface.PositiveSemidefiniteConeScaled ]
            # ok
        elseif dom in [MathOptInterface.ZeroOne,
                       MathOptInterface.Integer] &&
                           fun in [MathOptInterface.SingleVariable,
                                   MathOptInterface.VectorOfVariables]
            # ok
        else
            return false
        end
    end
    true
end

MathOptInterface.supportsproblem(m::MosekSolver, ::Type{MathOptInterface.SingleVariable},                constraint_types) :: Bool = supportsconstraints(m,constraint_types)
MathOptInterface.supportsproblem(m::MosekSolver, ::Type{MathOptInterface.ScalarAffineFunction{Float64}}, constraint_types) :: Bool = supportsconstraints(m,constraint_types)
MathOptInterface.supportsproblem{F}(m::MosekSolver, ::Type{F}, constraint_types) :: Bool = false

ref2id(ref :: MathOptInterface.VariableReference) :: Int =
    if ref.value & 1 == 0
        Int(ref.value >> 1)
    else
        - Int(ref.value >> 1)
    end

ref2id(ref :: MathOptInterface.ConstraintReference) :: Int =
    if ref.value & 1 == 0
        Int(ref.value >> 1)
    else
        - Int(ref.value >> 1)
    end

id2vref(id :: Int) :: MathOptInterface.VariableReference =
    if id < 0
        MathOptInterface.VariableReference((UInt64(-id) << 1) | 1)
    else
        MathOptInterface.VariableReference(UInt64(id) << 1)
    end

include("objective.jl")
include("variable.jl")
include("constraint.jl")
include("attributes.jl")

export MosekModel

end # module


#!/usr/bin/env julia

"""
    nanremove(arr::Array{T,2}) where T<:Real

For a real 2D array with columns of e.g. time, stress and strain, or frequency, G' and G'',
find any row which contains NaN values and remove that row completely.
"""
function nanremove(arr::Array{T,2}) where T<:Real
    # get 1D array which is NaN for any corresponding row with NaNs in
    rowsums = sum(arr, dims=2)

    # get list of NaNs / not NaNs (true for not NaN, false for NaN)
    notnanlist = vec(.!isnan.(rowsums))

    if false in notnanlist
        # raise warning
        @warn "Please note that NaN data rows are not included in resultant data struct."
    end

    # extract only those rows whose sum did not evaluate to NaN
    arr[notnanlist, :]
end

"""
    importcsv(filepath::String; t_col::IntOrNone = nothing, σ_col::IntOrNone = nothing, ϵ_col::IntOrNone = nothing, ω_col::IntOrNone = nothing, Gp_col::IntOrNone = nothing, Gpp_col::IntOrNone = nothing, delimiter=',')

Load data from a CSV file (two/three columns, comma seperated by default but
delimiter can be specified in the `delimiter` keyword argument). Arguments must be
identified by providing the number of the column in which they are contained.

Can be used to construct either a RheoTimeData instance or a RheoFreqData
instance. Function detects whether time or frequency has been included
and proceeds accordingly. For oscillatory data, all three columns (Gp, Gpp, Frequency)
must be provided. For regular viscoelastic data only time, or time-stress, or time-strain or
time-stress-strain data can be provided.
"""
function importcsv(filepath::String; t_col::IntOrNone = nothing, σ_col::IntOrNone = nothing, ϵ_col::IntOrNone = nothing, ω_col::IntOrNone = nothing, Gp_col::IntOrNone = nothing, Gpp_col::IntOrNone = nothing, delimiter = ',', comment = "Imported from csv file", savelog = true)

    @assert (!isnothing(t_col) && isnothing(ω_col)) || (isnothing(t_col) && !isnothing(ω_col)) "Data must contain either \"time\" or \"frequency\" "
    @assert endswith(lowercase(filepath), ".csv") "filepath must point to a .csv file."

    if !isnothing(t_col)

        @assert isnothing(Gp_col) && isnothing(Gpp_col) "Loss and storage modulus not allowed for time data"
        # read data from file
        dataraw = readdlm(filepath, delimiter)

        # remove and rows with NaN values in any of the columns
        data = nanremove(dataraw)

        #info=(comment=comment, folder=pwd(), stats=(t_min=data[1,t_col],t_max=data[end,t_col], n_sample=size(data[:,t_col])))
        #log = RheoLogItem( (type=:source, funct=:importcsv, params=(filepath=filepath,), keywords=(t_col=t_col, σ_col=σ_col, ϵ_col=ϵ_col)), info )

        log = if savelog
                info = (comment=comment, folder=pwd(), stats=(t_min=data[1,t_col],t_max=data[end,t_col], n_sample=size(data[:,t_col])))
                RheoLogItem( (type=:source, funct=:importcsv, params=(filepath=filepath,), keywords=(t_col=t_col, σ_col=σ_col, ϵ_col=ϵ_col)), info )
              else
                nothing
              end

        # generate RheologyData struct and output
        if !isnothing(ϵ_col) && !isnothing(σ_col)
            return RheoTimeData(t = data[:, t_col], ϵ = data[:, ϵ_col], σ = data[:, σ_col], log = log)
        elseif !isnothing(ϵ_col) && isnothing(σ_col)
            return RheoTimeData(t = data[:, t_col], ϵ = data[:, ϵ_col], log = log)
        elseif !isnothing(σ_col) && isnothing(ϵ_col)
            return RheoTimeData(t = data[:, t_col], σ = data[:, σ_col], log = log)
        end

    elseif !isnothing(ω_col)

        # check colnames length is correct
        @assert isnothing(σ_col) || isnothing(ϵ_col) "Stress and strain not allowed for frequency data"
        @assert !isnothing(Gp_col) && !isnothing(Gpp_col) "\"Gp\" and \"Gpp\" are required."

        # read data from file
        dataraw = readdlm(filepath, delimiter)
        # remove and rows with NaN values in any of the columns
        data = nanremove(dataraw)

        log = if savelog
                info = (comment=comment, folder=pwd(), stats=(ω_min=data[1,ω_col],ω_max=data[end,ω_col], n_sample=size(data[:,ω_col])))
                RheoLogItem( (type=:source, funct=:importcsv, params=(filepath=filepath,), keywords=(ω_col=ω_col, Gp_col=Gp_col, Gpp_col=Gpp_col)), info )
              else
                nothing
              end

        return RheoFreqData(ω = data[:,ω_col], Gp = data[:,Gp_col], Gpp = data[:,Gpp_col], log = log)

    end
end

"""
    exportcsv(self::Union{RheoTimeData, RheoFreqData}, filedir::String; delimiter=',', colorder=nothing)

Export `RheoTimeData` or `RheoFreqData` type to csv format. May be useful for plotting/analysis in other software.
By default, full time data will be exported with columns ordered as (t, σ, ϵ). Partial time data will be ordered
as either (t, σ) or (t, ϵ). Full frequency data will be ordered as (ω, Gp, Gpp). The order of columns can be customised
by passing a NamedTuple to the `colorder` arguments. For example (σ = 1, t = 3, ϵ = 2) would export the columns in the
order (σ, ϵ, t). As with `importcsv`, the delimiter can be set by keyword argument.
"""
function exportcsv(self::Union{RheoTimeData, RheoFreqData}, filedir::String; delimiter=',', colorder=nothing)

    @assert endswith(lowercase(filedir), ".csv") "filedir must be a .csv file."

    # if ordering of columns not provided, get default ordering depending on data contained in struct
    if isnothing(colorder)
        if typeof(self)==RheoTimeData
            datacontained = RheoTimeDataType(self)
            if datacontained == stress_only
                colorder = (t=1, σ=2)
            elseif datacontained == strain_only
                colorder = (t=1, ϵ=2)
            elseif datacontained == strain_and_stress
                colorder = (t=1, σ=2, ϵ=3)
            end

        elseif typeof(self)==RheoFreqData
            # invalid_freq_data=-1 freq_only=0 with_modulus=1
            datacontained = RheoFreqDataType(self)
            colorder = (ω = 1, Gp = 2, Gpp = 3)
        end
    end

    cols = collect(keys(colorder))
    sortedcols = sort(cols, by=i->getfield(colorder, i))
    dataraw = getfield.((self,), sortedcols)
    dataout = hcat(dataraw...)
    open(filedir, "w") do io
        writedlm(filedir, dataout, delimiter)
    end
end

# """
#     savedata(self::Union{RheoTimeData, RheoFreqData}, filedir::String)

# Save RheoTimeData or RheoFreqData object using JLD2 format for re-use in
# a later Julia session.
# """
# function savedata(self::Union{RheoTimeData, RheoFreqData}, filedir::String)

#     @assert endswith(lowercase(filedir), ".jld2") "filedir must be a .jld2 file."

#     @save filedir self

# end

# """
#     loaddata(filedir::String)

# Loads RheoTimeData or RheoFreqData object from a jld2 file.
# """
# function loaddata(filedir::String)

#     @assert endswith(lowercase(filedir), ".jld2") "filedir must be a .jld2 file."

#     @load filedir self

#     return self

# end

# """
#     savemodel(self::RheologyModel, filedir::String; ext = ".jld2")

# Save RheologyModel object using JLD2 format reuse in a later Julia session.
# """
# function savemodel(self::RheologyModel, filedir::String; ext = ".jld2")

#     fulldir = string(filedir, ext)

#     @save fulldir self

# end

# """
#     loaddata(filedir::String)

# Loads RheologyModel from a JLD2 file.
# """
# function loadmodel(filedir::String)

#     @load filedir self

#     return self

# end


module MathOptCartesians

import MathOptInterface
const MOI = MathOptInterface

include("scalar_set.jl")
include("vector_set.jl")

end # module