using JsonGrinder, JSON, Test, SparseArrays, Flux, Random, HierarchicalUtils
using JsonGrinder: ExtractScalar, ExtractCategorical, ExtractArray, ExtractDict, ExtractVector
using JsonGrinder: extractempty
using Mill
using Mill: catobs, nobs, MaybeHotMatrix
using LinearAlgebra

function misequals(a, b)
	for i in 1:length(a)
		ismissing(b[i]) && !ismissing(a[i]) && return(false)
		ismissing(a[i]) && !ismissing(b[i]) && return(false)
		!ismissing(a[i]) && a[i] != b[i] && return(false)
	end
	return(true)
end

@testset "Testing ExtractScalar" begin
	sc = ExtractScalar(Float64,2,3)
	@test all(sc("5").data .== [9])
	@test all(sc(5).data .== [9])
	@test all(sc(nothing).data .=== [missing])
	@test all(sc(missing).data .=== [missing])
	@test nobs(sc(missing)) == 1
	@test nobs(sc(nothing)) == 1
	@test sc(extractempty).data isa Matrix{Float64}
	@test nobs(sc(extractempty)) == 0
	@test nobs(sc(5)) == 1

	sc = ExtractScalar(Float32, 0.5, 4.0)
	@test sc(1).data isa Matrix{Float32}
	@test sc(extractempty).data isa Matrix{Float32}
end

@testset "ExtractCategorical" begin
	e = ExtractCategorical(["a","b"])
	@test e("a").data ≈ [1, 0, 0]
	@test e("b").data ≈ [0, 1, 0]
	@test e("z").data ≈ [0, 0, 1]
	@test isequal(e(nothing).data, [missing missing missing]')
	@test isequal(e(missing).data, [missing missing missing]')
	@test typeof(e("a").data) == MaybeHotMatrix{Int64,Array{Int64,1},Int64,Bool}
	@test typeof(e(nothing).data) == MaybeHotMatrix{Missing,Array{Missing,1},Int64,Missing}
	@test typeof(e(missing).data) == MaybeHotMatrix{Missing,Array{Missing,1},Int64,Missing}
	@test e(extractempty).data isa MaybeHotMatrix{Int64,Array{Int64,1},Int64,Bool}
	@test nobs(e(extractempty)) == 0

	@test e(["a", "b"]).data ≈ [1 0; 0 1; 0 0]
	@test isequal(e(["a", missing]).data, [true missing; false missing; false missing])
	@test isequal(e(["a", missing, "x"]).data, [true missing false; false missing false; false missing true])
	@test typeof(e(["a", "b"]).data) == MaybeHotMatrix{Int64,Array{Int64,1},Int64,Bool}
	@test typeof(e(["a", "b", nothing]).data) == MaybeHotMatrix{Union{Missing, Int64},Array{Union{Missing, Int64},1},Int64,Union{Missing, Bool}}

	@test isnothing(ExtractCategorical([]))
	e2 = ExtractCategorical(JsonGrinder.Entry(Dict("a"=>1,"c"=>1), 2))
	@test e2("a").data ≈ [1, 0, 0]
	@test e2("c").data ≈ [0, 1, 0]
	@test e2("b").data ≈ [0, 0, 1]
	@test isequal(e2(nothing).data, [missing missing missing]')
	@test isequal(e2(missing).data, [missing missing missing]')

	@test catobs(e("a"), e("b")).data ≈ [1 0; 0 1; 0 0]
	@test reduce(catobs, [e("a").data, e("b").data]) ≈ [1 0; 0 1; 0 0]
	# this is not working now, due to type stability issues
	# @test catobs(e("a").data, e("b").data) ≈ [1 0; 0 1; 0 0]
	@test isequal(e(Dict(1=>2)).data |> collect, [missing missing missing]')

	@test nobs(e("a")) == 1
	@test nobs(e("b")) == 1
	@test nobs(e("z")) == 1
	@test nobs(e(nothing)) == 1
	@test nobs(e(missing)) == 1
	@test nobs(e(missing).data) == 1
	@test nobs(e([missing, nothing])) == 2
	@test nobs(e([missing, nothing, "a"])) == 3

	e3 = ExtractCategorical(JsonGrinder.Entry(Dict(1=>1,2=>1), 2))
	@test e3(1).data ≈ [1, 0, 0]
	@test e3(2).data ≈ [0, 1, 0]
	@test e3(4).data ≈ [0, 0, 1]
	@test e3(1.).data ≈ [1, 0, 0]
	@test e3(2.).data ≈ [0, 1, 0]
	@test e3(4.).data ≈ [0, 0, 1]
	@test isequal(e3([]).data, [missing missing missing]')

	e4 = ExtractCategorical(JsonGrinder.Entry(Dict(1.0=>1,2.0=>1), 2))
	@test e4(1).data ≈ [1, 0, 0]
	@test e4(2).data ≈ [0, 1, 0]
	@test e4(4).data ≈ [0, 0, 1]
	@test e4(1.).data ≈ [1, 0, 0]
	@test e4(2.).data ≈ [0, 1, 0]
	@test e4(4.).data ≈ [0, 0, 1]
	@test isequal(e4([]).data, [missing missing missing]')
end

@testset "Testing array conversion" begin
	Mill.emptyismissing!(false)
	sc = ExtractArray(ExtractCategorical(2:4))
	@test all(sc([2,3,4]).data.data .== Matrix(1.0I, 4, 3))
	@test nobs(sc(nothing).data) == 0
	@test sc(nothing).data.data isa MaybeHotMatrix{Int64,Array{Int64,1},Int64,Bool}
	@test nobs(sc(nothing).data.data) == 0
	@test all(sc(nothing).bags.bags .== [0:-1])

	Mill.emptyismissing!(true)
	@test sc(nothing).data isa Missing
	@test all(sc(nothing).bags.bags .== [0:-1])
	Mill.emptyismissing!(false)

	@test nobs(sc(extractempty).data.data) == 0
	@test nobs(sc(extractempty).data) == 0
	@test isempty(sc(extractempty).bags.bags)
	@test sc(extractempty).data.data isa MaybeHotMatrix{Int64,Array{Int64,1},Int64,Bool}

	sc = ExtractArray(ExtractScalar(Float32))
	@test all(sc([2,3,4]).data.data .== [2 3 4])
	@test nobs(sc(nothing).data) == 0
	@test all(sc(nothing).bags.bags .== [0:-1])

	@test nobs(sc(extractempty).data.data) == 0
	@test nobs(sc(extractempty).data) == 0
	@test isempty(sc(extractempty).bags.bags)
	@test sc(extractempty).data.data isa Matrix{Float32}

	@test sc(nothing).data.data isa Matrix{Float32}
end

@testset "Testing feature vector conversion" begin
	sc = ExtractVector(5)
	@test sc([1, 2, 2, 3, 4]).data isa Matrix
	@test all(sc([1, 2, 2, 3, 4]).data .== [1, 2, 2, 3, 4])
	@test sc([1, 2, 2, 3, 4]).data isa Array{Float32, 2}
	@test sc(extractempty).data isa Matrix{Float32}
	@test nobs(sc(extractempty).data) == 0

	sc = ExtractVector{Int64}(5)
	@test all(sc([1, 2, 2, 3, 4]).data .== [1, 2, 2, 3, 4])
	@test sc([1, 2, 2, 3, 4]).data isa Array{Int64, 2}
	@test sc(nothing).data isa Matrix
	@test all(sc(nothing).data .=== missing)
	@test sc(extractempty).data isa Matrix{Int64}
	@test nobs(sc(extractempty).data) == 0

	# feature vector longer than expected
	sc = ExtractVector{Float32}(5)
	@test all(sc([1, 2, 2, 3, 4, 5]).data .== [1, 2, 2, 3, 4])
	@test typeof(sc([1, 2, 3, 4, 5]).data) == Array{Float32,2}
	@test sc([5, 6]).data[1:2] ≈ [5, 6]
	@test typeof(sc([1, 2]).data) == Array{Union{Missing,Float32},2}
	@test sc([1, 2]).data isa Matrix
	@test all(sc([5, 6]).data[3:5] .=== missing)
	@test all(sc(Dict(1=>2)).data .=== missing)
end

@testset "Testing ExtractDict" begin
	dict = Dict("a" => ExtractScalar(Float64,2,3),
				"b" => ExtractScalar(Float64),
				"c" => ExtractArray(ExtractScalar(Float64,2,3)))
	br = ExtractDict(dict)
	a1 = br(Dict("a" => 5, "b" => 7, "c" => [1,2,3,4]))
	a2 = br(Dict("a" => 5, "b" => 7))
	a3 = br(Dict("a" => 5, "c" => [1,2,3,4]))

	@test catobs(a1,a1)[:a].data ≈ [9 9]
	@test catobs(a1,a1)[:b].data ≈ [7 7]
	@test catobs(a1,a1)[:c].data.data ≈  [-3 0 3 6 -3 0 3 6]
	@test all(catobs(a1,a1)[:c].bags .== [1:4,5:8])

	@test catobs(a1,a2)[:a].data ≈ [9 9]
	@test catobs(a1,a2)[:b].data ≈ [7 7]
	@test catobs(a1,a2)[:c].data.data ≈ [-3 0 3 6]
	@test all(catobs(a1,a2)[:c].bags .== [1:4,0:-1])

	@test catobs(a2,a3)[:a].data ≈ [9 9]
	@test misequals(catobs(a2,a3)[:b].data, [7.0 missing])
	@test catobs(a2,a3)[:c].data.data ≈ [-3 0 3 6]
	@test all(catobs(a2,a3)[:c].bags .== [0:-1,1:4])

	@test catobs(a1,a3)[:a].data ≈ [9 9]
	@test misequals(catobs(a1,a3)[:b].data, [7.0 missing])
	@test catobs(a1,a3)[:c].data.data ≈ [-3 0 3 6 -3 0 3 6]
	@test all(catobs(a1,a3)[:c].bags .== [1:4,5:8])

	@test a1[:a].data ≈ [9]
	@test a1[:b].data ≈ [7]
	@test a2[:a].data ≈ [9]
	@test a2[:b].data ≈ [7]
	@test a3[:a].data ≈ [9]
	@test misequals(a3[:b].data, [missing])

	@test a1[:c].data.data ≈ [-3 0 3 6]
	@test all(a1[:c].bags .== [1:4])

	@test a3[:c].data.data ≈ [-3 0 3 6]
	@test all(a3[:c].bags .== [1:4])
	@test catobs(a3,a3)[:c].data.data ≈ [-3 0 3 6 -3 0 3 6]
	@test all(catobs(a3,a3)[:c].bags .== [1:4,5:8])

	a4 = br(extractempty)
	@test nobs(a4[:a]) == 0
	@test a4[:a].data isa Matrix{Float64}
	@test nobs(a4[:b]) == 0
	@test a4[:b].data isa Matrix{Float64}
	@test nobs(a4[:c]) == 0
	@test nobs(a4[:c].data) == 0
	@test a4[:c].data.data isa Matrix{Float64}
end

@testset "Testing Nested Missing Arrays" begin
	dict = Dict("a" => ExtractArray(ExtractScalar(Float32,2,3)),
		"b" => ExtractArray(ExtractScalar(Float32,2,3)))
	br = ExtractDict(dict)
	a1 = br(Dict("a" => [1,2,3], "b" => [1,2,3,4]))
	a2 = br(Dict("b" => [2,3,4]))
	a3 = br(Dict("a" => [2,3,4]))
	a4 = br(Dict{String,Any}())

	@test all(catobs(a1,a2).data[1].data.data .== [-3.0  0.0  3.0  6.0  0.0  3.0  6.0])
	@test all(catobs(a1,a2).data[1].bags .== [1:4, 5:7])
	@test all(catobs(a1,a2).data[2].data.data .== [-3.0  0.0  3.0])
	@test all(catobs(a1,a2).data[2].bags .== [1:3, 0:-1])

	@test all(catobs(a2,a3).data[1].data.data .== [0.0  3.0  6.0])
	@test all(catobs(a2,a3).data[1].bags .== [1:3, 0:-1])
	@test all(catobs(a2,a3).data[2].data.data .== [0 3 6])
	@test all(catobs(a2,a3).data[2].bags .== [0:-1, 1:3])

	@test all(catobs(a1,a4).data[1].data.data .== [-3.0  0.0  3.0  6.0])
	@test all(catobs(a1,a4).data[1].bags .== [1:4, 0:-1])
	@test all(catobs(a1,a4).data[2].data.data .== [-3.0  0.0  3.0])
	@test all(catobs(a1,a4).data[2].bags .== [1:3, 0:-1])

	@test all(a4.data[2].data.data isa Array{Float32,2})
	a4 = br(extractempty)
	@test nobs(a4[:a]) == 0
	@test nobs(a4[:a].data) == 0
	@test a4[:a].data.data isa Matrix{Float32}
	@test nobs(a4[:b]) == 0
	@test nobs(a4[:b].data) == 0
	@test a4[:b].data.data isa Matrix{Float32}
end

@testset "ExtractOneHot" begin
	samples = ["{\"name\": \"a\", \"count\" : 1}",
		"{\"name\": \"b\", \"count\" : 2}",]
	vs = JSON.parse.(samples)

	e = ExtractOneHot(["a","b"], "name", "count")
	@test e(vs).data[:] ≈ [1, 2, 0]
	@test e(nothing).data[:] ≈ [0, 0, 0]
	@test e(missing).data[:] ≈ [0, 0, 0]
	@test e(vs).data isa SparseMatrixCSC{Float32,Int64}
	@test e(nothing).data isa SparseMatrixCSC{Float32,Int64}
	@test e(missing).data isa SparseMatrixCSC{Float32,Int64}
	@test e(extractempty).data isa SparseMatrixCSC{Float32,Int64}
	@test nobs(e(extractempty)) == 0
	@test nobs(e(extractempty).data) == 0

	e = ExtractOneHot(["a","b"], "name", nothing)
	@test e(vs).data[:] ≈ [1, 1, 0]
	@test e(nothing).data[:] ≈ [0, 0, 0]
	@test e(missing).data[:] ≈ [0, 0, 0]
	@test typeof(e(vs).data) == SparseMatrixCSC{Float32,Int64}
	@test typeof(e(nothing).data) == SparseMatrixCSC{Float32,Int64}
	@test typeof(e(missing).data) == SparseMatrixCSC{Float32,Int64}
	vs = JSON.parse.(["{\"name\": \"c\", \"count\" : 1}"])
	@test e(vs).data[:] ≈ [0, 0, 1]
	@test typeof(e(vs).data) == SparseMatrixCSC{Float32,Int64}
	@test e(extractempty).data isa SparseMatrixCSC{Float32,Int64}
	@test nobs(e(extractempty)) == 0
	@test nobs(e(extractempty).data) == 0
end

@testset "ExtractString" begin
	e = ExtractString()
	@test e("Hello").data.s == ["Hello"]
	@test e(Symbol("Hello")).data.s == ["Hello"]
	@test e(["Hello", "world"]).data.s == ["Hello", "world"]
	@test e("Hello") isa ArrayNode{NGramMatrix{String,Vector{String},Int64},Nothing}
	@test e("Hello").data isa NGramMatrix{String,Vector{String},Int64}
	@test all(e(missing).data.s .=== [missing])
	@test e(missing).data isa NGramMatrix{Missing,Vector{Missing},Missing}
end

@testset "Extractor of keys as field" begin
	JsonGrinder.updatemaxkeys!(1000)
	js = [Dict(randstring(5) => rand()) for _ in 1:1000]
	sch = JsonGrinder.schema(js)
	ext = JsonGrinder.suggestextractor(sch, (;key_as_field = 500))

	b = ext(js[1])
	k = only(keys(js[1]))
	@test b.data[:item].data[1] ≈ (js[1][k] - ext.item.c) * ext.item.s
	@test b.data[:key].data.s[1] == k

	b = ext(nothing)
	@test nobs(b) == 1
	@test nobs(b.data) == 0
	b = ext(Dict())
	@test nobs(b) == 1
	@test nobs(b.data) == 0

	b = ext(extractempty)
	@test nobs(b) == 0
	@test nobs(b.data) == 0
	@test nobs(b.data[:item]) == 0
	@test b.data[:item].data isa Matrix{Float32}
	@test nobs(b.data[:key]) == 0
	@test b.data[:key].data isa Mill.NGramMatrix{String,Array{String,1},Int64}

	js = [Dict(randstring(5) => Dict(:a => rand(), :b => randstring(1))) for _ in 1:1000]
	sch = JsonGrinder.schema(js)
	ext = JsonGrinder.suggestextractor(sch, (;key_as_field = 500))

	b = ext(js[1])
	k = only(keys(js[1]))
	i = ext.item(js[1][k])
	@test b.data[:item][:a].data == i[:a].data
	@test b.data[:item][:b].data ==i[:b].data
	@test b.data[:key].data.s[1] == k

	b = ext(nothing)
	@test nobs(b) == 1
	@test nobs(b.data) == 0
	b = ext(Dict())
	@test nobs(b) == 1
	@test nobs(b.data) == 0
end

@testset "equals and hash test" begin
	other1 = Dict(
		"a" => ExtractArray(ExtractScalar(Float64,2,3)),
		"b" => ExtractArray(ExtractScalar(Float64,2,3)),
		"c" => ExtractCategorical(["a","b"]),
		"d" => ExtractVector(4),
	)
	br1 = ExtractDict(other1)
	other11 = Dict(
		"a" => ExtractArray(br1),
		"b" => ExtractScalar(Float64,2,3),
	)
	br11 = ExtractDict(other11)

	other2 = Dict(
		"a" => ExtractArray(ExtractScalar(Float64,2,3)),
		"b" => ExtractArray(ExtractScalar(Float64,2,3)),
		"c" => ExtractCategorical(["a","b"]),
		"d" => ExtractVector(4),
	)
	br2 = ExtractDict(other2)
	other22 = Dict("a" => ExtractArray(br2), "b" => ExtractScalar(Float64,2,3))
	br22 = ExtractDict(other22)

	@test hash(br11) === hash(br22)
	@test hash(br11) !== hash(br1)
	@test hash(br11) !== hash(br2)
	@test hash(br1) === hash(br2)

	@test br11 == br22
	@test br11 != br1
	@test br11 != br2
	@test br1 == br2
end
# todo: add tests for settings of suggestextractor!!!
# todo: add tests for single keyed dict, that model generated from such jsons behaves as we want
@testset "Extractor skip empty lists" begin
	j1 = JSON.parse("""{"a": [{"a":1},{"b":2}], "b": [], "c": [[], []], "d": 1, "e": {"a": 1}, "f": {"a": []}}""")
	j2 = JSON.parse("""{"a": [{"a":3},{"b":4}], "b": [], "c": [[]], "d": 1, "e": {"a": 2}, "f": {"a": []}}""")
	j3 = JSON.parse("""{"a": [{"a":1},{"b":3}], "b": [], "c": [[], [], [], []], "d": 2, "e": {"a": 3}, "f": {"a": []}}""")
	j4 = JSON.parse("""{"a": [{"a":2},{"b":4}], "b": [], "c": [], "d": 4, "e": {"a": 3}, "f": {"a": []}}""")

	sch = JsonGrinder.schema([j1, j2, j3, j4])
	ext = suggestextractor(sch)
	@test ext[:a] isa ExtractArray
	@test isnothing(ext[:b])
	@test isnothing(ext[:c])
	@test ext[:d] isa ExtractScalar
	@test ext[:e] isa ExtractDict
	@test isnothing(ext[:f])
end

@testset "Extractor of numbers as strings" begin
	j1 = JSON.parse("""{"a": "1", "b": "a", "c": "1.1", "d": 1.1, "e": "1.2", "f": 1}""")
	j2 = JSON.parse("""{"a": "2", "b": "b", "c": "2", "d": 2, "e": "1.3", "f": 2}""")
	j3 = JSON.parse("""{"a": "3", "b": "c", "c": "2.3", "d": 2.3, "e": "1.4", "f": 3}""")
	j4 = JSON.parse("""{"a": "4", "b": "c", "c": "5", "d": 5, "e": "1.4", "f": 3}""")

	sch = JsonGrinder.schema([j1, j2, j3, j4])
	ext = suggestextractor(sch)

	@test ext[:a] isa ExtractCategorical
	@test ext[:b] isa ExtractString
	@test ext[:c] isa ExtractCategorical
	@test ext[:d] isa ExtractScalar{Float32}
	@test ext[:e] isa ExtractCategorical
	@test ext[:f] isa ExtractScalar{Float32}

	ext_j1 = ext(j1)
	ext_j2 = ext(j2)
	ext_j3 = ext(j3)
	ext_j4 = ext(j4)

	@test eltype(ext_j1[:a].data) <: Bool
	@test eltype(ext_j1[:b].data) <: Int64
	@test eltype(ext_j1[:c].data) <: Bool
	@test eltype(ext_j1[:d].data) <: Float32
	@test eltype(ext_j1[:e].data) <: Bool
	@test eltype(ext_j1[:f].data) <: Float32

	@test eltype(ext_j2[:a].data) <: Bool
	@test eltype(ext_j2[:b].data) <: Int64
	@test eltype(ext_j2[:c].data) <: Bool
	@test eltype(ext_j2[:d].data) <: Float32
	@test eltype(ext_j2[:e].data) <: Bool
	@test eltype(ext_j2[:f].data) <: Float32

	@test eltype(ext_j3[:a].data) <: Bool
	@test eltype(ext_j3[:b].data) <: Int64
	@test eltype(ext_j3[:c].data) <: Bool
	@test eltype(ext_j3[:d].data) <: Float32
	@test eltype(ext_j3[:e].data) <: Bool
	@test eltype(ext_j3[:f].data) <: Float32

	@test eltype(ext_j4[:a].data) <: Bool
	@test eltype(ext_j4[:b].data) <: Int64
	@test eltype(ext_j4[:c].data) <: Bool
	@test eltype(ext_j4[:d].data) <: Float32
	@test eltype(ext_j4[:e].data) <: Bool
	@test eltype(ext_j4[:f].data) <: Float32

	@test ext_j1["U"].data ≈ [0]
	@test ext_j2["U"].data ≈ [3/13]
	@test ext_j3["U"].data ≈ [4/13]
	@test ext_j4["U"].data ≈ [1]

	m = reflectinmodel(sch, ext)
	# this is for stable version
	# @test buf_printtree(m) == """
	# ProductModel ↦ ArrayModel(Dense(42, 10))
	#   ├── a: ArrayModel(Dense(5, 10))
	#   ├── b: ArrayModel(Dense(2053, 10))
	#   ├── c: ArrayModel(Dense(5, 10))
	#   ├── d: ArrayModel(identity)
	#   ├── e: ArrayModel(Dense(4, 10))
	#   └── f: ArrayModel(identity)"""
	  # this is test for master
	@test buf_printtree(m) == """
	ProductModel … ↦ ArrayModel(Dense(42, 10))
	  ├── a: ArrayModel(Dense(5, 10))
	  ├── b: ArrayModel(Dense(2053, 10))
	  ├── c: ArrayModel(Dense(5, 10))
	  ├── d: ArrayModel(identity)
	  ├── e: ArrayModel(Dense(4, 10))
	  └── f: ArrayModel(identity)"""
end
# todo: add separate tests for reflectinmodel(sch, ext) to test behavior with various missing and non-missing stuff in schema
# e.g. empty string, missing keys, irregular schema and MultiRepresentation

@testset "Suggest feature vector extraction" begin
	j1 = JSON.parse("""{"a": "1", "b": [1, 2, 3], "c": [1, 2, 3]}""")
	j2 = JSON.parse("""{"a": "2", "b": [2, 2, 3], "c": [1, 2, 3, 4]}""")
	j3 = JSON.parse("""{"a": "3", "b": [3, 2, 3], "c": [1, 2, 3, 4, 5]}""")
	j4 = JSON.parse("""{"a": "4", "b": [2, 3, 4], "c": [1, 2, 3]}""")

	sch = JsonGrinder.schema([j1, j2, j3, j4])
	ext = suggestextractor(sch)

	@test ext[:a] isa ExtractCategorical
	@test ext[:b] isa ExtractVector
	@test ext[:b].n == 3
	@test ext[:c] isa ExtractArray{ExtractScalar{Float32}}

	ext_j1 = ext(j1)
	ext_j2 = ext(j2)
	ext_j3 = ext(j3)
	ext_j4 = ext(j4)

	@test ext_j1["U"].data ≈ [1, 2, 3]
	@test ext_j2["U"].data ≈ [2, 2, 3]
	@test ext_j3["U"].data ≈ [3, 2, 3]
	@test ext_j4["U"].data ≈ [2, 3, 4]

	@test ext_j1["s"].data ≈ [0 .25 .5]
	@test ext_j2["s"].data ≈ [0 .25 .5 .75]
	@test ext_j3["s"].data ≈ [0 .25 .5 .75 1.]
	@test ext_j4["s"].data ≈ [0 .25 .5]
end

@testset "Suggest complex" begin
	JsonGrinder.updatemaxkeys!(1000)
	js = [Dict("a" => rand(), "b" => Dict(randstring(5) => rand()), "c"=>[rand(), rand()], "d"=>[rand() for i in 1:rand(1:10)]) for _ in 1:1000]
	sch = JsonGrinder.schema(js)
	ext = JsonGrinder.suggestextractor(sch, (;key_as_field = 500))

	@test ext[:a] isa ExtractScalar{Float32}
	@test ext[:b] isa JsonGrinder.ExtractKeyAsField
	@test ext[:b].key isa ExtractString
	@test ext[:b].item isa ExtractScalar{Float32}
	@test ext[:c] isa ExtractVector
	@test ext[:c].n == 2
	@test ext[:d] isa ExtractArray
end

@testset "Extractor typed missing bags" begin
	j1 = JSON.parse("""{"a": []}""")
	j2 = JSON.parse("""{"a": [1, 2]}""")
	j3 = JSON.parse("""{"a": [1, 2, 3]}""")
	j4 = JSON.parse("""{"a": [1, 2, 3, 4]}""")
	j5 = JSON.parse("""{"a": [1, 2, 3, 4, 5]}""")

	sch = JsonGrinder.schema([j1, j2, j3, j4])
	ext = suggestextractor(sch)

	ext_j1 = ext(j1)
	ext_j2 = ext(j2)
	ext_j3 = ext(j3)
	ext_j4 = ext(j4)

	@test ext_j1[:a].data.data isa Array{Float32,2}
	@test ext_j2[:a].data.data isa Array{Float32,2}
	@test ext_j3[:a].data.data isa Array{Float32,2}
	@test ext_j4[:a].data.data isa Array{Float32,2}
end

@testset "testing irregular extractor" begin
	j1 = JSON.parse("""{"a": 4}""")
	j2 = JSON.parse("""{"a": { "a": "hello", "b":[5,6]}}""")
	j3 = JSON.parse("""{"a": [1, 2, 3, 4]}""")
	j4 = JSON.parse("""{"a": 5}""")
	j5 = JSON.parse("""{"a": 3}""")

	sch = schema([j1,j2,j3,j4,j5])
	ext = suggestextractor(sch)
	a = ext(j1)
	@test a[:a][:e1].data[1] == 0.5
	@test ismissing(a[:a][:e2].data[:a].data.s[1])
	@test nobs(a[:a][:e3]) == 1
	@test nobs(a[:a][:e3].data) == 1
end

@testset "Mixed scalar extraction" begin
	j1 = JSON.parse("""{"a": "1"}""")
	j2 = JSON.parse("""{"a": 4}""")
	j3 = JSON.parse("""{"a": "3.1"}""")
	j4 = JSON.parse("""{"a": 2.5}""")

	sch = JsonGrinder.schema([j1, j2, j3, j4])
	sch_hash = hash(sch)
	ext = suggestextractor(sch)
	@test sch_hash == hash(sch)	# testing that my entry parkour does not modify schema
	@test nchildren(ext[:a]) == 1

	e1 = ext(j1)
	e2 = ext(j2)
	e3 = ext(j3)
	e4 = ext(j4)
	@test e1["k"].data ≈ [0]
	@test e2["k"].data ≈ [1]
	@test e3["k"].data ≈ [0.7]
	@test e4["k"].data ≈ [0.5]
end

@testset "Mixed scalar extraction with other types" begin
	j1 = JSON.parse("""{"a": "1"}""")
	j2 = JSON.parse("""{"a": 2.5}""")
	j3 = JSON.parse("""{"a": "3.1"}""")
	j4 = JSON.parse("""{"a": 5}""")
	j5 = JSON.parse("""{"a": "4.5"}""")
	j6 = JSON.parse("""{"a": [1, 2, 3, 4, 5]}""")
	j7 = JSON.parse("""{"a": {"Sylvanas is the worst warchief ever": "yes"}}""")

	sch = JsonGrinder.schema([j1, j2, j3, j4, j5, j6, j7])
	ext = suggestextractor(sch)

	@test buf_printtree(sch) ==
    """
    [Dict] (updated = 7)
      └── a: [MultiEntry] (updated = 7)
               ├── 1: [Scalar - String], 3 unique values, updated = 3
               ├── 2: [Scalar - Float64,Int64], 2 unique values, updated = 2
               ├── 3: [List] (updated = 1)
               │        └── [Scalar - Int64], 5 unique values, updated = 5
               └── 4: [Dict] (updated = 1)
                        └── Sylvanas is the worst warchief ever: [Scalar - String], 1 unique values, updated = 1"""

	@test buf_printtree(ext) ==
    """
	Dict
	  └── a: MultiRepresentation
	           ├── e1: FeatureVector with 5 items
	           ├── e2: Dict
	           │         └── Sylvanas is the worst warchief ever: String
	           └── e3: Float32"""

	e1 = ext(j1)
	e2 = ext(j2)
	e3 = ext(j3)
	e4 = ext(j4)
	e5 = ext(j5)
	@test e1["s"].data ≈ [0]
	@test e2["s"].data ≈ [0.375]
	@test e3["s"].data ≈ [0.525]
	@test e4["s"].data ≈ [1.0]
	@test e5["s"].data ≈ [0.875]
	@test buf_printtree(e1) ==
	"""
	ProductNode with 1 obs
	  └── a: ProductNode with 1 obs
	           ├── e1: ArrayNode(5×1 Array with Missing elements) with 1 obs
	           ├── e2: ProductNode with 1 obs
	           │         └── Sylvanas is the worst warchief ever: ArrayNode(2053×1 NGramMatrix with Missing elements) with 1 obs
	           └── e3: ArrayNode(1×1 Array with Float32 elements) with 1 obs"""
end

@testset "mixing numeric and non-numeric strings" begin
	j1 = JSON.parse("""{"a": "hello"}""")
	j2 = JSON.parse("""{"a": "4"}""")
	j3 = JSON.parse("""{"a": 5}""")
	j4 = JSON.parse("""{"a": 3.0}""")
	j5 = JSON.parse("""{"a": [1, 2, 3, 4, 5]}""")

	sch = JsonGrinder.schema([j1, j2, j3, j4, j5])
	ext = suggestextractor(sch)

	@test buf_printtree(sch) ==
	"""
	[Dict] (updated = 5)
	  └── a: [MultiEntry] (updated = 5)
	           ├── 1: [Scalar - String], 2 unique values, updated = 2
	           ├── 2: [Scalar - Float64,Int64], 2 unique values, updated = 2
	           └── 3: [List] (updated = 1)
	                    └── [Scalar - Int64], 5 unique values, updated = 5"""

	@test buf_printtree(ext) ==
	"""
	Dict
	  └── a: MultiRepresentation
	           ├── e1: String
	           ├── e2: Float32
	           └── e3: FeatureVector with 5 items"""

	e1 = ext(j1)
	e2 = ext(j2)
	e3 = ext(j3)
	e4 = ext(j4)
	e5 = ext(j5)
	@test all(e1["k"].data .=== [missing])
	@test e2["k"].data ≈ [0.5]
	@test e3["k"].data ≈ [1.0]
	@test e4["k"].data ≈ [0.0]
	@test all(e5["k"].data .=== [missing])

	@test hash(ext) !== hash(suggestextractor(JsonGrinder.schema([j1, j2, j4, j5])))
end

@testset "empty string and substring" begin
	a, b, c = split("a b ", " ")
	d = "d"
	e = ""
	f = 5.2
	@test a isa SubString{String}
	@test b isa SubString{String}
	@test c isa SubString{String}
	@test d isa String
	@test e isa String
	@test f isa AbstractFloat
	@test c == ""
	ext = JsonGrinder.extractscalar(AbstractString)
	@test SparseMatrixCSC(ext(c).data) == SparseMatrixCSC(ext(e).data)
	@test all(ext(f).data.s .=== ext(nothing).data.s)
end

@testset "key as field" begin
	j1 = JSON.parse("""{"a": 1}""")
	j2 = JSON.parse("""{"b": 2.5}""")
	j3 = JSON.parse("""{"c": 3.1}""")
	j4 = JSON.parse("""{"d": 5}""")
	j5 = JSON.parse("""{"e": 4.5}""")
	j6 = JSON.parse("""{"f": 5}""")
	sch = JsonGrinder.schema([j1, j2, j3, j4, j5, j6])
	e = JsonGrinder.key_as_field(sch, NamedTuple(), path = "")
	@test e isa JsonGrinder.ExtractKeyAsField

	e2 = JsonGrinder.key_as_field(sch, NamedTuple(), path = "")
	@test hash(e) === hash(e2)
	@test e == e2
end

@testset "AuxiliaryExtractor HUtils" begin
	e2 = ExtractCategorical(["a","b"])
	e = AuxiliaryExtractor(e2, (ext, sample)->ext(String(sample)))

	@test e("b") == e(:b)
	@test e("b").data ≈ [0, 1, 0]
	@test e(:b).data ≈ [0, 1, 0]

    @test buf_printtree(e) ==
	"""
    Auxiliary extractor with
      └── Categorical d = 3"""
end

@testset "Skipping single dict key" begin
	j1 = JSON.parse("""{"a": []}""")
	j2 = JSON.parse("""{"a": [1, 2]}""")
	j3 = JSON.parse("""{"a": [1, 2, 3]}""")
	j4 = JSON.parse("""{"a": [1, 2, 3, 4]}""")
	j5 = JSON.parse("""{"a": [1, 2, 3, 4, 5]}""")

	sch = JsonGrinder.schema([j1, j2, j3, j4, j5])
	ext = suggestextractor(sch)

	@test buf_printtree(ext) ==
	"""
	Dict
	  └── a: Array of
	           └── Float32"""

	ext_j2 = ext(j2)
	@test buf_printtree(ext_j2) ==
    """
	ProductNode with 1 obs
	  └── a: BagNode with 1 obs
	           └── ArrayNode(1×2 Array with Float32 elements) with 2 obs"""
		   # todo: add more tests for integration with Mill to make sure it's propagated well
end


#=

tribe wrapper.

Ignacio Quintero Mächler

t(-_-t)

April 27 2017

=#

"""
    tribe(tree_file::String,
          data_file::String,
          out_file ::String;
          min_dt   ::Float64           = 0.01,
          niter    ::Int64             = 50_000,
          nburn    ::Int64             = 50_000,
          nthin    ::Int64             = 500,
          saveXY   ::Tuple{Bool,Int64} = (false, 1_000),
          saveDM   ::Tuple{Bool,Int64} = (false, 1_000),
          ωxprior  ::NTuple{2,Float64} = (0.,10.),
          ω1prior  ::NTuple{2,Float64} = (0.,10.),
          ω0prior  ::NTuple{2,Float64} = (0.,10.),
          σ²prior  ::Float64           = 1e-1,
          λprior   ::Float64           = 1e-1,
          weight   ::NTuple{5,Float64} = (0.15,0.05,0.02,0.02,5e-3),
          λ1i      ::Float64           = 1.0,
          λ0i      ::Float64           = 0.5,
          ωxi      ::Float64           = 0.0,
          ω1i      ::Float64           = 0.0,
          ω0i      ::Float64           = 0.0,
          fix_ωx   ::Bool              = false,
          fix_ω1   ::Bool              = false,
          fix_ω0   ::Bool              = false,
          delim    ::Char              = '\t',
          eol      ::Char              = '\r')

Run tribe. Wrapper for all functions.
"""
function tribe(tree_file::String,
               data_file::String,
               out_file ::String;
               min_dt   ::Float64           = 0.01,
               niter    ::Int64             = 50_000,
               nburn    ::Int64             = 50_000,
               nthin    ::Int64             = 500,
               saveXY   ::Tuple{Bool,Int64} = (false, 1_000),
               saveDM   ::Tuple{Bool,Int64} = (false, 1_000),
               ωxprior  ::NTuple{2,Float64} = (0.,10.),
               ω1prior  ::NTuple{2,Float64} = (0.,10.),
               ω0prior  ::NTuple{2,Float64} = (0.,10.),
               σ²prior  ::Float64           = 1e-1,
               λprior   ::Float64           = 1e-1,
               weight   ::NTuple{5,Float64} = (0.15,0.05,0.02,0.02,5e-3),
               λ1i      ::Float64           = 1.0,
               λ0i      ::Float64           = 0.5,
               ωxi      ::Float64           = 0.0,
               ω1i      ::Float64           = 0.0,
               ω0i      ::Float64           = 0.0,
               fix_ωx   ::Bool              = false,
               fix_ω1   ::Bool              = false,
               fix_ω0   ::Bool              = false,
               delim    ::Char              = '\t',
               eol      ::Char              = '\r')

  tip_values, tip_areas, tree, bts =
    read_data_tribe(tree_file, data_file)

  X, Y, B, ncoup, δt, tree, si =
    initialize_data(tip_values, tip_areas, min_dt, tree, bts)

  R = tribe_mcmc(X, Y, ncoup, δt,
                 deepcopy(tree.ed),
                 deepcopy(tree.el),
                 B,
                 niter    = niter,
                 nthin    = nthin,
                 nburn    = nburn,
                 saveXY   = saveXY,
                 saveDM   = saveDM,
                 ωxprior  = ωxprior,
                 ω1prior  = ω1prior,
                 ω0prior  = ω0prior,
                 σ²prior  = σ²prior,
                 λprior   = λprior,
                 out_file = out_file,
                 weight   = weight,
                 λ1i      = λ1i,
                 λ0i      = λ0i,
                 ωxi      = ωxi,
                 ω1i      = ω1i,
                 ω0i      = ω0i,
                 σ²i      = si,
                 fix_ωx   = fix_ωx,
                 fix_ω1   = fix_ω1,
                 fix_ω0   = fix_ω0)

  return R

end

"""
    tribe(tip_values::Dict{Int64,Float64},
          tip_areas ::Dict{Int64,Array{Int64,1}},
          tree      ::rtree,
          bts       ::Array{Float64,1},
          out_file  ::String;
          min_dt    ::Float64           = 0.01,
          niter     ::Int64             = 500_000,
          nburn     ::Int64             = 500_000,
          nthin     ::Int64             = 1_000,
          saveXY    ::Tuple{Bool,Int64} = (false, 1_000),
          saveDM    ::Tuple{Bool,Int64} = (false, 1_000),
          ωxprior   ::NTuple{2,Float64} = (0.,10.),
          ω1prior   ::NTuple{2,Float64} = (0.,10.),
          ω0prior   ::NTuple{2,Float64} = (0.,10.),
          σ²prior   ::Float64           = 1e-1,
          λprior    ::Float64           = 1e-1,
          weight    ::NTuple{5,Float64} = (0.15,0.05,0.02,0.02,5e-3),
          λ1i       ::Float64           = 1.0,
          λ0i       ::Float64           = 0.4,
          ωxi       ::Float64           = 0.,
          ω1i       ::Float64           = 0.,
          ω0i       ::Float64           = 0.,
          fix_ωx    ::Bool              = false,
          fix_ω1    ::Bool              = false,
          fix_ω0    ::Bool              = false,
          delim     ::Char              = '\t',
          eol       ::Char              = '\r')

Run tribe for simulations. Wrapper for all functions.
"""
function tribe(tip_values::Dict{Int64,Float64},
               tip_areas ::Dict{Int64,Array{Int64,1}},
               tree      ::rtree,
               bts       ::Array{Float64,1},
               out_file  ::String;
               min_dt    ::Float64           = 0.01,
               niter     ::Int64             = 500_000,
               nburn     ::Int64             = 500_000,
               nthin     ::Int64             = 1_000,
               saveXY    ::Tuple{Bool,Int64} = (false, 1_000),
               saveDM    ::Tuple{Bool,Int64} = (false, 1_000),
               ωxprior   ::NTuple{2,Float64} = (0.,10.),
               ω1prior   ::NTuple{2,Float64} = (0.,10.),
               ω0prior   ::NTuple{2,Float64} = (0.,10.),
               σ²prior   ::Float64           = 1e-1,
               λprior    ::Float64           = 1e-1,
               weight    ::NTuple{5,Float64} = (0.15,0.05,0.02,0.02,5e-3),
               λ1i       ::Float64           = 1.0,
               λ0i       ::Float64           = 0.4,
               ωxi       ::Float64           = 0.,
               ω1i       ::Float64           = 0.,
               ω0i       ::Float64           = 0.,
               fix_ωx    ::Bool              = false,
               fix_ω1    ::Bool              = false,
               fix_ω0    ::Bool              = false,
               delim     ::Char              = '\t',
               eol       ::Char              = '\r')

  X, Y, B, ncoup, δt, tree, si =
    initialize_data(tip_values, tip_areas, min_dt, tree, bts)

  R = tribe_mcmc(X, Y, ncoup, δt,
                 deepcopy(tree.ed),
                 deepcopy(tree.el),
                 B,
                 niter    = niter,
                 nthin    = nthin,
                 nburn    = nburn,
                 saveXY   = saveXY,
                 saveDM   = saveDM,
                 ωxprior  = ωxprior,
                 ω1prior  = ω1prior,
                 ω0prior  = ω0prior,
                 σ²prior  = σ²prior,
                 λprior   = λprior,
                 out_file = out_file,
                 weight   = weight,
                 λ1i      = λ1i,
                 λ0i      = λ0i,
                 ωxi      = ωxi,
                 ω1i      = ω1i,
                 ω0i      = ω0i,
                 σ²i      = si,
                 fix_ωx   = fix_ωx,
                 fix_ω1   = fix_ω1,
                 fix_ω0   = fix_ω0)

  return R

end

"""
    tribe(out_file::String;
          niter   ::Int64             = 500_000,
          nburn   ::Int64             = 500_000,
          nthin   ::Int64             = 1_000,
          ωxprior ::NTuple{2,Float64} = (0.,10.),
          ω1prior ::NTuple{2,Float64} = (0.,10.),
          ω0prior ::NTuple{2,Float64} = (0.,10.),
          σ²prior ::Float64           = 1e-1,
          λprior  ::Float64           = 1e-1,
          weight  ::NTuple{4,Float64} = (0.15,0.05,0.02,0.02),
          σ²i     ::Float64           = 1.,
          ωxi     ::Float64           = 0.,
          ω1i     ::Float64           = 0.01,
          ω0i     ::Float64           = 0.01,
          λ1i     ::Float64           = 1.0,
          λ0i     ::Float64           = 0.2,
          fix_ωx  ::Bool              = false,
          fix_ω1  ::Bool              = false,
          fix_ω0  ::Bool              = false)

Run tribe **under the prior**. Wrapper for all functions.
"""
function tribe(out_file::String;
               niter   ::Int64             = 500_000,
               nburn   ::Int64             = 500_000,
               nthin   ::Int64             = 1_000,
               ωxprior ::NTuple{2,Float64} = (0.,10.),
               ω1prior ::NTuple{2,Float64} = (0.,10.),
               ω0prior ::NTuple{2,Float64} = (0.,10.),
               σ²prior ::Float64           = 1e-1,
               λprior  ::Float64           = 1e-1,
               weight  ::NTuple{4,Float64} = (0.15,0.05,0.02,0.02),
               σ²i     ::Float64           = 1.,
               ωxi     ::Float64           = 0.,
               ω1i     ::Float64           = 0.01,
               ω0i     ::Float64           = 0.01,
               λ1i     ::Float64           = 1.0,
               λ0i     ::Float64           = 0.2,
               fix_ωx  ::Bool              = false,
               fix_ω1  ::Bool              = false,
               fix_ω0  ::Bool              = false)

  R = tribe_mcmc(out_file,
                 niter    = niter,
                 nthin    = nthin,
                 nburn    = nburn,
                 ωxprior  = ωxprior,
                 ω1prior  = ω1prior,
                 ω0prior  = ω0prior,
                 σ²prior  = σ²prior,
                 λprior   = λprior,
                 weight   = weight,
                 λ1i      = λ1i,
                 λ0i      = λ0i,
                 ωxi      = ωxi,
                 ω1i      = ω1i,
                 ω0i      = ω0i,
                 σ²i      = σ²i,
                 fix_ωx   = fix_ωx,
                 fix_ω1   = fix_ω1,
                 fix_ω0   = fix_ω0)

  return R

end


using DataFrames
using CategoricalArrays
using Random

import Base: hash, size, getindex, length, show

struct ModalInstance <: AbstractVector{SubDataFrame}
    # worlds::Set{<:AbstractOntology} # TODO add set of worlds, for the semantics (DT)
    rows::Vector{DataFrameRow}
end

ModalInstance() = ModalInstance(DataFrameRow[])
DataFrames.nrow(::Array{DataFrameRow,1}) = 1

function show(io::IO, ::MIME"text/plain", mi::ModalInstance)
    println(io, "Modal instance with $(length(mi.rows)) frames")
    for i in 1:length(mi.rows)
        println(io, "Frame #$i")
        println(io, mi.rows[i])
    end
end

size(mi::ModalInstance) = (nrow(mi.rows),)
getindex(mi::ModalInstance, i::Int) = mi.rows[i]
length(mi::ModalInstance, fid::Int) = max([length(mi[fid][i]) for i in 1:length(mi[fid])]...)

mutable struct ModalFrame <: AbstractVector{SubDataFrame}
    # name::String # TODO add name?
    dimension::Int  # (0=static, 1=timeseries, 2=images, ecc.)
    data::DataFrame
    views::Vector{DataFrameRow} # intraframe instances (i.e., first, second, .., series)
end

dimension(mf::ModalFrame) = mf.dimension

@inline function ModalFrame(data::DataFrame)
    colnames = names(data)

    @assert length(colnames) > 0 && nrow(data) > 0 "There must be at least one column and one instance in data to compute its dimension."

    dimension = ndims(data[1, Symbol(colnames[1])])

    for colname in colnames
        curr_dimension = ndims(data[1, Symbol(colname)])
        @assert dimension == curr_dimension "Each attribute must have the same dimension."
    end

    # if dimension > 0
    #     typ = eltype(data[1, Symbol(colnams[1])])
    #     for colname in colnames
    #         curr_typ = eltype(data[1, Symbol(colname)])
    #         @assert typ == cur_typ "Each modal frame with higher-than-zero dimension must have the same element types."
    #     end
    # end

    views = Vector{DataFrameRow}(undef, nrow(data))
    for i in 1:nrow(data)
        views[i] = view(data, i, :)
    end
    ModalFrame(dimension, data, views)
end

size(mf::ModalFrame) = (nrow(mf.data),)
getindex(mf::ModalFrame, i::Int) = mf.views[i]

struct ClassificationDataset
    ldim::Int
    hdim::Int
    frames::Vector{ModalFrame}
    instances::Vector{ModalInstance}
    domains::Dict{Int, Dict{Symbol, Vector{Number}}}
    classes::CategoricalArray
    unique_classes::Vector
end

function ClassificationDataset(frames::Vector{ModalFrame}, classes::CategoricalArray)
    instances = ModalInstance[]
    domains = Dict{Int, Dict{Symbol, Vector{Number}}}()

    n_instances = size(frames[1])

    for i in 1:length(frames)
        curr_n_instances = size(frames[i])
        @assert n_instances == curr_n_instances "Each frame must have the same number of (intraframe) instances!"
    end

    @assert n_instances[1] == length(classes) "Each instance must have one and only one class!"

    # Computing the instances.
    # for each instance
    for i in 1:n_instances[1]
        instance = ModalInstance()
        # for each frame
        for f in 1:length(frames)
            # push the instance
            push!(instance.rows, frames[f][i])
        end
        push!(instances, instance)
    end

    # Computing the domains.
    # for each frame
    for f in 1:length(frames)
        attrs = names(frames[f][1])
        # for each attribute
        attr_domain = Dict{Symbol, Vector{Number}}()
        for attr in attrs
            # convert to symbol
            A = Symbol(attr)
            # for each instance
            # for i in 1:n_instances[1]

            # end
            attr_domain[A] = collect(Set([v for i in 1:n_instances[1] for v in frames[f][i][A]]))
            domains[f] = attr_domain
        end
    end

    ldim, hdim = extrema([dimension(i) for i in frames])

    ClassificationDataset(ldim, hdim, frames, instances, domains, classes, unique(classes))
end

instance(cs::ClassificationDataset, i::Int) = cs.instances[i] # TODO cs -> ds (either sup or unsupervised)
instance(cs::ClassificationDataset, i::Int, f::Int) = cs.instances[i][f]
function attributes(ds::ClassificationDataset)
    d = Dict{Int,Array{String,1}}()
    for (fid,frame) in enumerate(ds.frames)
        d[fid] = names(frame.data)
    end
    return d
end
attributes(ds::ClassificationDataset, fid::Int) = names(ds.frames[fid].data)

length(ds::ClassificationDataset) = length(ds.instances)

function show(io::IO, ::MIME"text/plain", ds::ClassificationDataset)
    println(io, "Classification dataset with $(length(ds.instances)) instances")

    for (i, inst) in enumerate(ds.instances)
        println(io, "Instance #$(i):")
        println(io, instance(ds, i))
        println(io)
    end
end

# TODO
# possible test: all([auslan.instances[i].rows[1][attr] === auslan.frames[1].data[i,attr] for i in 1:length(auslan.instances) for attr in attributes(auslan, 1)])
function transform!(ds::ClassificationDataset, f::Function, fid::Int; kwargs...)
    for attr in attributes(ds, fid)
        for i in 1:nrow(ds.frames[fid].data)
            ds.frames[fid].data[i,attr] = f(ds.frames[fid].data[i,attr]; kwargs...)
        end
    end
end

# function Base.minimum(ds::ClassificationDataset)
#     d = Dict{Int,Array{Float64,1}}()

#     for (fid,frame) in enumerate(ds.frames)
#         frame = ds.frames[fid]
#         if frame.dimension == 0
#             d[fid] = missing
#         else
#             d[fid] = [Base.minimum(vcat(frame.data[!,attr]...)) for attr in attributes(ds, fid)]
#         end
#     end
#     return d
# end

# function Base.maximum(ds::ClassificationDataset)
#     d = Dict{Int,Array{Float64,1}}()

#     for (fid,frame) in enumerate(ds.frames)
#         frame = ds.frames[fid]
#         if frame.dimension == 0
#             d[fid] = missing
#         else
#             d[fid] = [Base.maximum(vcat(frame.data[!,attr]...)) for attr in attributes(ds, fid)]
#         end
#     end
#     return d
# end

# TODO fix problem with select_attributes!
# TODO make it with bang!
function select_attributes(ds::ClassificationDataset, frame_attributes::Dict{Int, Array{Int, 1}})
    frames = collect(keys(frame_attributes))

    # Frame level.
    for i in 1:length(frames)
        ds.frames[frames[i]].data = ds.frames[frames[i]].data[:, frame_attributes[frames[i]]]
    end

    modal_frames = ModalFrame[]
    for i in 1:length(frames)
        push!(modal_frames, ModalFrame(ds.frames[frames[i]].data))
    end

    return ClassificationDataset(modal_frames, ds.classes)
end

# function transform!(ds::ClassificationDataset, f::Function; kwargs...)

# using MultivariateTimeSeries
# X, y = read_data_labeled(joinpath(dirname(pathof(GBDTs)), "..", "data", "auslan_youtube8"));
#
# # adesso voglio passare da MTS a ClassificationDataset
#
# my_frame = DataFrame()
#
# for attr in names(X)
#     insertcols!(my_frame, attr => Array{Float64,1}[])
# end
#
# # for each instance
# for i in 1:length(X)
#     array_data = Array{Float64, 2}(X[i])
#     push!(my_frame, [array_data[:, j] for j in 1:length(names(X))])
# end
#
# auslan = ClassificationDataset([ModalFrame(my_frame)], CategoricalArray(y))

function sample(ds::ClassificationDataset,
        range::Union{UnitRange{Int},Nothing};
        invert::Bool=false)

    n_instances = length(ds.instances)

    @assert range ⊆ 1:n_instances "The range must be a subset of 1:$(n_instances)."

    if ismissing(range)
        return ds
    else

        output_classes = Int[]
        timeseries_df = DataFrame()

        attrs = attributes(ds, 1) # TODO up to now, only timeseries

        for attr in attrs
            insertcols!(timeseries_df, attr => Array{Float64, 1}[]) # TODO also here Array{Float64, 1}[], i.e., timeseries
        end

        if invert == false
            for idx in range
                push!(timeseries_df, ds.instances[idx][1]) # ds.instances[idx][1] <-- [1] means the first frame TODO in future
                push!(output_classes, ds.classes[idx])
            end
        else
            for idx in setdiff(1:n_instances,range)
                push!(timeseries_df, ds.instances[idx][1]) # ds.instances[idx][1] <-- [1] means the first frame TODO in future
                push!(output_classes, ds.classes[idx])
            end
        end
    end

    return ClassificationDataset([ModalFrame(timeseries_df)], CategoricalArray(output_classes))
end

function sample(ds::ClassificationDataset,
    percentage::Float64;
    seed::Int=1,
    bias::Float64=0.0, # bias to uniform
    replace::Bool=true,
    invert::Bool=false)

    @assert 0 ≤ percentage && percentage ≤ 1 "Percentage p must be 0 ≤ p ≤ 1."

    n_instances = length(ds.instances)
    n_instances_per_class = countmap(ds.classes)
    classes = keys(n_instances_per_class)
    n_classes = length(n_instances_per_class)

    instances_per_class = Dict{CategoricalArrays.CategoricalValue{Int64, UInt32}, Array{Int, 1}}()

    for class in classes
        instances_per_class[class] = ModalInstance[]
    end

    instances_to_sample = Dict{CategoricalArrays.CategoricalValue{Int64, UInt32}, Int64}()

    for class in classes
        sample_size = Int(round(percentage * ((1-bias)*n_instances_per_class[class] + bias*n_instances / n_classes)))
        instances_to_sample[class] = sample_size
    end

    for i in 1:n_instances
        push!(instances_per_class[ds.classes[i]], i)
    end

    r = Random.seed!(seed)

    output_classes = Int[]
    timeseries_df = DataFrame()

    attrs = attributes(ds, 1) # TODO up to now, only timeseries

    for attr in attrs
        insertcols!(timeseries_df, attr => Array{Float64, 1}[]) # TODO also here Array{Float64, 1}[], i.e., timeseries
    end

    if replace
        for class in classes
            for i in 1:instances_to_sample[class]
                idx = instances_per_class[class][rand(r, 1:n_instances_per_class[class])]
                push!(timeseries_df, ds.instances[idx][1]) # ds.instances[idx][1] <-- [1] means the first frame TODO in future
                push!(output_classes, ds.classes[idx])
                # @show idx
                # @show ds.classes[idx]
            end
        end
    else
        for class in classes
            n_eligible = length(instances_per_class[class])
            selected = instances_per_class[class]

            # the selected elements are those at the end of then selected array
            for i in 1:instances_to_sample[class]
                idx = rand(r, 1:n_eligible)
                selected[idx], selected[n_eligible] = selected[n_eligible], selected[idx]
                n_eligible -= 1
            end

            # take the last instances_to_sample[class] instances
            if invert == false
                for i in n_instances_per_class[class]-instances_to_sample[class]+1:n_instances_per_class[class]
                    idx = selected[i]
                    push!(timeseries_df, ds.instances[idx][1]) # ds.instances[idx][1] <-- [1] means the first frame TODO in future
                    push!(output_classes, ds.classes[idx])
                end
            # take the first n_instances_per_class[class]-instances_to_sample[class] instances
            else
                for i in 1:n_instances_per_class[class]-instances_to_sample[class]
                    idx = selected[i]
                    push!(timeseries_df, ds.instances[idx][1]) # ds.instances[idx][1] <-- [1] means the first frame TODO in future
                    push!(output_classes, ds.classes[idx])
                end
            end
        end
    end

    return ClassificationDataset([ModalFrame(timeseries_df)], CategoricalArray(output_classes))
end

ds = SymbolicML.trip_no_trip("/home/edu/.julia/dev/SymbolicML/data/Siemens/Features");
# ds = SymbolicML.trip_no_trip("/home/edu/.julia/dev/SymbolicML/data/Siemens"; n_machine=42740);
# d = Dict{Int,Array{Int,1}}()
# d[1] = [1,2,4,5,9,19,22,28,29,30]
# SymbolicML.select_attributes!(ds, d)
SymbolicML.transform!(ds, SymbolicML.window_slice, 1; from=1,to=120)

grammar = @grammar begin
    b = G(bvec) | F(bvec) | G(implies(bvec,bvec))
    bvec = and(bvec, bvec)
    bvec = or(bvec, bvec)
    bvec = not(bvec)
    bvec = lt(rvec, rvec)
    bvec = lte(rvec, rvec)
    bvec = gt(rvec, rvec)
    bvec = gte(rvec, rvec)
    bvec = f_lt(x, xid, v, vid)
    bvec = f_lte(x, xid, v, vid)
    bvec = f_gt(x, xid, v, vid)
    bvec = f_gte(x, xid, v, vid)
    rvec = x[xid]
    # xid = |(["x_1","y_1","z_1","roll_1","pitch_1","yaw_1","thumbbend_1","forebend_1","middlebend_1","ringbend_1","littlebend_1","x_2","y_2","z_2","roll_2","pitch_2","yaw_2","thumbbend_2","forebend_2","middlebend_2","ringbend_2","littlebend_2"])
    xid = |("Fce", "Fsfc", "G2", "MAD", "T3-2", "T13-12", "T16-15", "T13-5", "T14-6", "T15-7")
    # xid = |("Ambient_air_humidity","Compr_IGV_position","Gas_fuel_valve_position","Ambient_air_temperature","Compressor_outlet_temperature","Exhaust_temperature_1","Exhaust_temperature_2","Exhaust_temperature_3","Exhaust_temperature_4","Exhaust_temperature_5","Exhaust_temperature_6","Exhaust_temperature_7","Exhaust_temperature_8","Exhaust_temperature_9","Exhaust_temperature_10","Exhaust_temperature_11","Exhaust_temperature_12","Exhaust_temperature_13","Exhaust_temperature_14","Exhaust_temperature_15","Exhaust_temperature_16","Compressor_outlet_pressure","Gas_fuel_mass_flow_rate","Compressor_inlet_air_mass_flow_rate","Rotational_speed","Power_output")
    # xid = |("Ambient_air_humidity","Compr_IGV_position","Gas_fuel_valve_position","Ambient_air_temperature","Compressor_outlet_temperature","Compressor_outlet_pressure","Gas_fuel_mass_flow_rate","Compressor_inlet_air_mass_flow_rate","Rotational_speed","Power_output")
    vid = |(1:10)
end

G(v) = all(v)                                                #globally
F(v) = any(v)                                                #eventually
f_lt(x, xid, v, vid) = lt(x[xid], v[xid][vid])               #feature is less than a constant
f_lte(x, xid, v, vid) = lte(x[xid], v[xid][vid])             #feature is less than or equal to a constant
f_gt(x, xid, v, vid) = gt(x[xid], v[xid][vid])               #feature is greater than a constant
f_gte(x, xid, v, vid) = gte(x[xid], v[xid][vid])             #feature is greater than or equal to a constant

#workarounds for slow dot operators:
implies(v1, v2) = (a = similar(v1); a .= v2 .| .!v1)         #implies
not(v) = (a = similar(v); a .= .!v)                          #not
and(v1, v2) = (a = similar(v1); a .= v1 .& v2)               #and
or(v1, v2) = (a = similar(v1); a .= v1 .| v2)                #or
lt(x1, x2) = (a = Vector{Bool}(undef,length(x1)); a .= x1 .< x2)   #less than
lte(x1, x2) = (a = Vector{Bool}(undef,length(x1)); a .= x1 .≤ x2)  #less than or equal to
gt(x1, x2) = (a = Vector{Bool}(undef,length(x1)); a .= x1 .> x2)   #greater than
gte(x1, x2) = (a = Vector{Bool}(undef,length(x1)); a .= x1 .≥ x2)  #greater than or equal to

balanced_ds = SymbolicML.sample(ds, 0.727; seed=1, bias=1.0, replace=false);
train_ds = SymbolicML.sample(balanced_ds, 0.8, seed=1, bias=1.0, replace=false);
test_ds = SymbolicML.sample(balanced_ds, 0.8, seed=1, bias=1.0, replace=false, invert=true);

# all_instances = union(train_ds.instances, test_ds.instances)

# A = for (i, inst) in enumerate(all_instances)
#     A = Bool[]
#     @show i
#     if inst in balanced_ds.instances
#         push!(A, true)
#     else
#         push!(A, false)
#     end
#     A
# end

const v = Dict{String,Vector{Float64}}()
mins, maxes = minimum(train_ds)[1], maximum(train_ds)[1]
for (i,xid) in enumerate(attributes(train_ds,1))#enumerate(names(X))
    v[xid] = collect(range(mins[i],stop=maxes[i],length=10))
end;

p = MonteCarlo(1000, 5)
# genetic = GeneticProgram(100,20,5,0.5,0.4,0.1)
using Random; Random.seed!(1)
model = induce_tree(grammar, :b, p, train_ds.frames[1], train_ds.classes, 5, min_members_per_branch=1, prevent_same_label=true);

train_ind = collect(1:length(train_ds.instances));
train_y_pred = classify(model, train_ds.frames[1], train_ind);
train_accuracy = count(train_y_pred .== train_ds.classes[train_ind]) / length(train_ind)

test_ind = collect(1:length(test_ds.instances));
test_y_pred = classify(model, test_ds.frames[1], test_ind);
test_accuracy = count(test_y_pred .== test_ds.classes[test_ind]) / length(test_ind)

# TODO now it works for the specific problem
function confusion_matrix(ŷ, y)
    @assert length(ŷ) == length(y) "The number of predictions must be equal to the number of true classes."
    # @show y
    # @show length(unique(y))
    # mat = zeros(Int, length(unique(y)), length(unique(y)))
    mat = zeros(Int, 2, 2)

    for i in 1:length(y)
        mat[Int(ŷ[i].ref), Int(y.refs[i])] += 1
    end

    return mat
end

function performance_metrics(M)
    return tn, fn, fp, tp = M[1,1], M[1,2], M[2,1], M[2,2]
end

function sensitivity(M)
    tn, fn, fp, tp = performance_metrics(M)
    return tp/(tp+fn)
end

function specificity(M)
    tn, fn, fp, tp = performance_metrics(M)
    return tn/(tn+fp)
end

@show train_accuracy
@show test_accuracy

ds = SymbolicML.trip_no_trip("/home/edu/.julia/dev/SymbolicML/data/Siemens"; n_machine=42740);

const v = Dict{String,Vector{Float64}}()

for seed in 1:10
    for window_sz in 1:2
        for hour in 24:-1:19
            start_t = 1440 - hour*60 + 1
            end_t = 1440 - hour*60 + (window_sz*60)

            balanced_ds = SymbolicML.sample(ds; seed=seed, percentage=0.727, bias=1.0, replace=false);
            SymbolicML.transform!(balanced_ds, SymbolicML.window_slice, 1; from=start_t,to=end_t)
            # @show countmap(balanced_ds.classes)
            train_ds = SymbolicML.sample(balanced_ds, seed=seed, percentage=0.8, bias=1.0, replace=false);
            # @show countmap(train_ds.classes)
            test_ds = SymbolicML.sample(balanced_ds, seed=seed, percentage=0.8, bias=1.0, replace=false, invert=true);
            # @show countmap(test_ds.classes)

            # local v = Dict{String,Vector{Float64}}()
            mins, maxes = minimum(train_ds)[1], maximum(train_ds)[1]
            for (i,xid) in enumerate(attributes(train_ds,1))#enumerate(names(X))
                v[xid] = collect(range(mins[i],stop=maxes[i],length=10))
            end

            p = MonteCarlo(500, 5)

            Random.seed!(seed)

            model = induce_tree(grammar, :b, p, train_ds.frames[1], train_ds.classes, 7, min_members_per_branch=1);

            # train_ind = collect(1:length(train_ds.instances));
            # train_y_pred = classify(model, train_ds.frames[1], train_ind);
            # train_accuracy = count(train_y_pred .== train_ds.classes[train_ind]) / length(train_ind)

            test_ind = collect(1:length(test_ds.instances));
            test_y_pred = classify(model, test_ds.frames[1], test_ind);
            test_accuracy = round(count(test_y_pred .== test_ds.classes[test_ind]) / length(test_ind), digits=4)

            test_M = confusion_matrix(test_y_pred, test_ds.classes[test_ind])
            test_sens = round(sensitivity(test_M),digits=4)
            test_spec = round(specificity(test_M),digits=4)

            @show (seed,window_sz,hour,test_accuracy,test_sens,test_spec)
        end
    end
end

for seed in 1:10
    # @show seed
    for window_sz in 1:2
        # @show window_sz
        for hour in 24:-1:19
            # @show hour
            start_t = 1440 - hour*60 + 1
            end_t = 1440 - hour*60 + (window_sz*60)

            balanced_ds = SymbolicML.sample(ds, 0.727; seed=seed, bias=1.0, replace=false);
            SymbolicML.transform!(balanced_ds, SymbolicML.window_slice, 1; from=start_t,to=end_t)
            test_M = zeros(Int, 2, 2)
            for r in 1:length(balanced_ds.instances)
                # @show r
                train_ds = SymbolicML.sample(balanced_ds, r:r, invert=true);
                test_ds = SymbolicML.sample(balanced_ds, r:r, invert=false);

                mins, maxes = minimum(train_ds)[1], maximum(train_ds)[1]
                for (i,xid) in enumerate(attributes(train_ds,1))#enumerate(names(X))
                    v[xid] = collect(range(mins[i],stop=maxes[i],length=10))
                end

                p = MonteCarlo(500, 5)

                Random.seed!(seed)

                model = induce_tree(grammar, :b, p, train_ds.frames[1], train_ds.classes, 7, min_members_per_branch=1);

                test_ind = collect(1:length(test_ds.instances));
                test_y_pred = classify(model, test_ds.frames[1], test_ind);
                # test_accuracy = round(count(test_y_pred .== test_ds.classes[test_ind]) / length(test_ind), digits=4)

                # @show test_y_pred
                # @show convert(Int, test_y_pred[1])
                # # @show test_ds.classes[r]
                # @show test_ds.classes[test_ind][1]
                # @show Int(test_y_pred[1].ref)
                # @show test_ds.classes[test_ind]
                # @show test_ds.classes[r]
                # tn
                if convert(Int, test_y_pred[1]) == 0 && convert(Int, test_ds.classes[test_ind][1]) == 0
                    test_M[1,1] += 1
                # tp
                elseif convert(Int, test_y_pred[1]) == 1 && convert(Int, test_ds.classes[test_ind][1]) == 1
                    test_M[2,2] +=1
                # fn
                elseif convert(Int, test_y_pred[1]) == 0 && convert(Int, test_ds.classes[test_ind][1]) == 1
                    test_M[1,2] += 1
                # fp
                elseif convert(Int, test_y_pred[1]) == 1 && convert(Int, test_ds.classes[test_ind][1]) == 0
                    test_M[2,1] += 1
                end

                # @show test_M
            end

            # @show test_M

            test_accuracy = (test_M[1,1] + test_M[2,2])/length(balanced_ds.instances)
            test_sens = round(sensitivity(test_M),digits=4)
            test_spec = round(specificity(test_M),digits=4)
            @show (seed,window_sz,hour,test_accuracy,test_sens,test_spec)
        end
    end
end


"Infer fluents that are never modified by some action in a domain."
function infer_static_fluents(domain::Domain)
    affected = Set(infer_affected_fluents(domain))
    static = setdiff(keys(get_fluents(domain)), affected)
    return collect(static)
end

"Infer fluents that are modified by some action in a domain."
function infer_affected_fluents(domain::Domain)
    affected = Symbol[]
    for action in values(get_actions(domain))
        append!(affected, get_affected(action))
    end
    return unique!(affected)
end

const AFFECTED_FUNCS = Dict{Symbol, Function}()

"Return the names of all fluents affected by an action."
function get_affected(action::Action)
    return get_affected(get_effect(action))
end

function get_affected(effect::Term)
    affected_fn! = get(AFFECTED_FUNCS, effect.name, builtin_affected!)
    return affected_fn!(Symbol[], effect)
end

function builtin_affected!(affected::Vector{Symbol}, effect::Term)
    if effect.name == :and
        for eff in effect.args append!(affected, get_affected(eff)) end
    elseif effect.name == :when
        append!(affected, get_affected(effect.args[2]))
    elseif effect.name == :forall
        append!(affected, get_affected(effect.args[2]))
    elseif effect.name == :assign
        push!(affected, effect.args[1].name)
    elseif effect.name in keys(GLOBAL_MODIFIERS)
        push!(affected, effect.args[1].name)
    elseif effect.name == :not
        push!(affected, effect.args[1].name)
    else
        push!(affected, effect.name)
    end
    return affected
end

"Check if term is static or composed of static subterms."
function is_static(term::Term, domain::Domain,
                   statics=infer_static_fluents(domain))
    if is_external_func(term, domain)
        return all(is_static(a, domain, statics) for a in term.args)
    else
        return term.name in statics
    end
end

function is_static(term::Const, domain::Domain,
                   statics=infer_static_fluents(domain))
    !is_fluent(term, domain) || term.name in statics
end

"Simplify away static fluents within a `term`."
function simplify_statics(term::Term, domain::Domain, state::State,
                          statics=infer_static_fluents(domain))
    # Simplify predicates if they are static and ground
    if is_static(term, domain, statics) && is_ground(term)
        return Const(evaluate(domain, state, term))
    elseif !(term.name in (:and, :or, :imply, :not))
        return term
    end
    # Simplify logical compounds
    args = Term[simplify_statics(a, domain, state, statics) for a in term.args]
    if term.name == :and
        true_idxs = Int[]
        for (i, a) in enumerate(args)
            a.name == false && return Const(false)
            a.name == true && push!(true_idxs, i)
        end
        length(true_idxs) == length(args) && return Const(true)
        deleteat!(args, true_idxs)
        return length(args) == 1 ? args[1] : Compound(:and, args)
    elseif term.name == :or
        false_idxs = Int[]
        for (i, a) in enumerate(args)
            a.name == true && return Const(true)
            a.name == false && push!(false_idxs, i)
        end
        length(false_idxs) == length(args) && return Const(false)
        deleteat!(args, false_idxs)
        return length(args) == 1 ? args[1] : Compound(:or, args)
    elseif term.name == :imply
        cond, query = args
        cond.name == true && return query
        cond.name == false && return Const(true)
        return Compound(:imply, args)
    elseif term.name == :not
        val = args[1].name
        return val isa Bool ? Const(!val) : Compound(:not, args)
    else
        error("Unrecognized logical operator: $(term.name)")
    end
end


module VariationalAE

export VAE, fit!, predict, generate, reconstruct

include("core.jl")

end # module


function get_inverse(obj::DecompositionSolver)
    return jcall(obj, "getInverse", RealMatrix, ())
end

function is_non_singular(obj::DecompositionSolver)
    return jcall(obj, "isNonSingular", jboolean, ())
end

function solve(obj::DecompositionSolver, arg0::RealMatrix)
    return jcall(obj, "solve", RealMatrix, (RealMatrix,), arg0)
end

function solve(obj::DecompositionSolver, arg0::RealVector)
    return jcall(obj, "solve", RealVector, (RealVector,), arg0)
end


"""
	vav2tsoft(file_results,file_output;sitename)
Convert VAV06 results (Earth tides) to TSoft tidal group info (written
parameters can be used in "LOCAT.TSD")

**Input**
* file_results: file containing VAV results (="analysisXX.dat") of gravity tidal analysis
* file_output: output file with tidal groups in tsoft format. Paste these results to TSoft/Tidal parameter set
* site: name of the site (identical to TSoft database)
* name: name of the tidal parameters (unique for given site)
* filemode: write new file ("w", default) or append to existing ("a")

file_vav = "test/input/vav_analysis.dat";
file_tsoft = "test/output/vav2tsoft.txt"
vav2tsoft(file_vav,file_tsoft,site="Cantlay",name="test");
"""
function vav2tsoft(file_results::String,file_output::String;
					site::String="Site",name::String="Name",
					filemode::String="w")::Nothing
	# Get number of header lines
	hl = FileTools.headlines(file_results,"results at filter frequency")
	# open file for writting + add header
	fo = tide2tsoft(file_output,site,name,filemode);
	# write parameters
	if hl != 0
		open(file_results,"r") do fid
			for i in 1:hl+1;readline(fid); end
			row = readline(fid);
			while !occursin(">",row)
				if !occursin("---",row) && !occursin("results",row)
					vav2tsoft_write(fo,row);
				end
				row = readline(fid);
				eof(fid) ? break : nothing;
			end
		end
	end
	close(fo);
end
"""
Auxiliary function to write the output file of vav2tsoft function
"""
function vav2tsoft_write(f::IOStream,s::String)::Nothing
	@printf(f,"COMP: %s  %s  %s  %s  %s\n",
			s[1:10],s[13:20],s[50:57],s[65:73],replace(s[27:31]," "=>""));
end

"""
	eterna2tsoft(file_results,file_output;file_groups,sitename)
Convert ETERNA34 "PRN" results (Earth tides) to TSoft tidal group info (written
parameters can be used in "LOCAT.TSD")
> Make sense only when the analyze parameter TIDALPOTEN= 4

**Input**
* file_results: file containing tidal analysis results (="projectname.prn")
* file_output: output file with tidal groups in tsoft format. Paste these results to TSoft/Tidal parameter set
* site: name of the site (identical to TSoft database)
* name: name of the tidal parameters (unique for given site)
* filemode: write new file ("w", default) or append to existing ("a")

file_eterna = "test/input/eterna_analysis.prn";
file_tsoft = "test/output/eterna2tsoft.txt"
eterna2tsoft(file_eterna,file_tsoft,site="Cantlay",name="test");
"""
function eterna2tsoft(file_results::String,file_output::String;
					site::String="Site",name::String="Name",
					filemode::String="w")::Nothing
	# Get number of header lines
	hl = FileTools.headlines(file_results,"Adjusted tidal parameters :")
	# open file for writting + add header
	fo = tide2tsoft(file_output,site,name,filemode);
	# write parameters
	if hl != 0
		open(file_results,"r") do fid
			for i in 1:hl+5;readline(fid); end
			row = readline(fid);
			while !occursin("Adjusted",row)
				length(row)>73 ? eterna2tsoft_write(fo,row) : nothing
				row = readline(fid);
				eof(fid) ? break : nothing;
			end
		end
	end
	close(fo);
end
"""
Auxiliary function to write the output file of eterna2tsoft function
"""
function eterna2tsoft_write(f::IOStream,s::String)::Nothing
	@printf(f,"COMP: %s %s %s %s %s\n",
			s[7:15],s[16:24],s[39:48],s[57:66],replace(s[25:29]," "=>""));
end
"""
Auxiliary function to write TSD header
"""
function tide2tsoft(filein::String,site::String,name::String,filemode::String)::IOStream
	f = open(filein,filemode);
	@printf(f,"[TIDECOMP] %s\n",site)
	@printf(f,"NAME: %s\nTYPE: 1\nAZIM: 0.0000000\nGRAV: 0.0000000\n",name);
	return f
end


"""
A dense table stores the data as a vector of row-tuples.
"""
immutable DenseTable{Index,DataTypes} <: AbstractTable{Index,DefaultKey()}
    data::Vector{Row{Index,DataTypes}}

    function DenseTable(data_in)
        check_row(Index,DataTypes)
        new(data_in)
    end
end

DenseTable{Index}(::Index) = DenseTable{Index,eltypes(Index)}(Vector{Row{Index,eltypes(Index)}}())
DenseTable{Index,DataTypes}(::Index,data::Vector{DataTypes}) = DenseTable{Index,DataTypes}(Vector{Row{Index,DataTypes}}(data))
# Constructor that takes index and seperate values
# Conversion between Dense and normal Tables??
DenseTable{Index,StorageTypes}(table::Table{Index,StorageTypes}) = DenseTable(Index,collect(zip(table.data...)))

Base.length{Index,DataTypes}(table::DenseTable{Index,DataTypes}) = length(table.data)
Base.size{Index,DataTypes}(table::DenseTable{Index,DataTypes}) = (length(Index),length(table.data))
Base.size{Index,DataTypes}(table::DenseTable{Index,DataTypes},i::Int) = i == 1 ? length(Index) : (i == 2 ? length(table.data) : error("Tables are two-dimensional"))
Base.eltype{Index,DataTypes}(table::DenseTable{Index,DataTypes}) = eltype(table.data)

Base.push!{Index,DataTypes}(table::DenseTable{Index,DataTypes},row::Row{Index,DataTypes}) = push!(table.data,row)
Base.push!{Index,DataTypes}(table::DenseTable{Index,DataTypes},data_in::DataTypes) = push!(table.data,Row{Index,DataTypes}(data_in))
Base.push!{Index,DataTypes}(table::DenseTable{Index,DataTypes},data_in...) = push!(table,Row{Index,DataTypes}((data_in...))) # TODO probably slow
Base.append!{Index,DataTypes}(table::DenseTable{Index,DataTypes},table_in::DenseTable{Index,DataTypes}) = append!(table.data,table_in.data)
Base.append!{Index,DataTypes}(table::DenseTable{Index,DataTypes},rows::Vector{Row{Index,DataTypes}}) = append!(table.data,rows)

# Some output
function Base.summary{Index,DataTypes}(io::IO,table::DenseTable{Index,DataTypes})
    println(io,"$(length(table.data[1]))-row DenseTable with columns $Index")
end

function Base.show{Index,DataTypes}(io::IO,table::DenseTable{Index,DataTypes})
    summary(io,table)
    show(io,table.data)
end


@testset "Drawing boxes" begin
    # Testing that boxes are drawn as expected

    ### square model
    yolomod = YOLO.v3_tiny_COCO(w=416, h=416, silent=true)
    #Nonsquare low aspect ratio image
    img = fill(Gray(1), 200, 100)
    batch = emptybatch(yolomod)
    batch[:,:,:,1], padding = prepareImage(img, yolomod)
    res = collect([ padding[1] padding[2] 1.0-padding[3] 1.0-padding[4] 0.0 0.0;]') #note the transpose!
    imgboxes = drawBoxes(img, yolomod, padding, res)
    @test all(imgboxes[1,:] .== Gray(0))
    @test all(imgboxes[:,1] .== Gray(0))
    @test all(imgboxes[end,:] .== Gray(0))
    @test all(imgboxes[:,end] .== Gray(0))

    #Nonsquare high aspect ratio image
    img = fill(Gray(1), 100, 200)
    batch = emptybatch(yolomod)
    batch[:,:,:,1], padding = prepareImage(img, yolomod)
    res = collect([ padding[1] padding[2] 1.0-padding[3] 1.0-padding[4] 0.0 0.0;]') #note the transpose!
    imgboxes = drawBoxes(img, yolomod, padding, res)
    @test all(imgboxes[1,:] .== Gray(0))
    @test all(imgboxes[:,1] .== Gray(0))
    @test all(imgboxes[end,:] .== Gray(0))
    @test all(imgboxes[:,end] .== Gray(0))

    #Square image, square model
    img = fill(Gray(1), 100, 100)
    batch = emptybatch(yolomod)
    batch[:,:,:,1], padding = prepareImage(img, yolomod)
    res = collect([ padding[1] padding[2] 1.0-padding[3] 1.0-padding[4] 0.0 0.0;]') #note the transpose!
    imgboxes = drawBoxes(img, yolomod, padding, res)
    @test all(imgboxes[1,:] .== Gray(0))
    @test all(imgboxes[:,1] .== Gray(0))
    @test all(imgboxes[end,:] .== Gray(0))
    @test all(imgboxes[:,end] .== Gray(0))

    ### nonsquare low aspect ratio model
    yolomod = YOLO.v3_tiny_COCO(w=512, h=416, silent=true)
    #Square image
    img = fill(Gray(1), 100, 100)
    batch = emptybatch(yolomod)
    batch[:,:,:,1], padding = prepareImage(img, yolomod)
    res = collect([ padding[1] padding[2] 1.0-padding[3] 1.0-padding[4] 0.0 0.0;]') #note the transpose!
    imgboxes = drawBoxes(img, yolomod, padding, res)
    @test all(imgboxes[1,:] .== Gray(0))
    @test all(imgboxes[:,1] .== Gray(0))
    @test all(imgboxes[end,:] .== Gray(0))
    @test all(imgboxes[:,end] .== Gray(0))

    #nonsquare low aspect ratio image
    img = fill(Gray(1), 200, 100)
    batch = emptybatch(yolomod)
    batch[:,:,:,1], padding = prepareImage(img, yolomod)
    res = collect([ padding[1] padding[2] 1.0-padding[3] 1.0-padding[4] 0.0 0.0;]') #note the transpose!
    imgboxes = drawBoxes(img, yolomod, padding, res)
    @test all(imgboxes[1,:] .== Gray(0))
    @test all(imgboxes[:,1] .== Gray(0))
    @test all(imgboxes[end,:] .== Gray(0))
    @test all(imgboxes[:,end] .== Gray(0))

    ### nonsquare high aspect ratio model
    yolomod = YOLO.v3_tiny_COCO(w=416, h=512, silent=true)
    #Square image
    img = fill(Gray(1), 100, 100)
    batch = emptybatch(yolomod)
    batch[:,:,:,1], padding = prepareImage(img, yolomod)
    res = collect([ padding[1] padding[2] 1.0-padding[3] 1.0-padding[4] 0.0 0.0;]') #note the transpose!
    imgboxes = drawBoxes(img, yolomod, padding, res)
    @test all(imgboxes[1,:] .== Gray(0))
    @test all(imgboxes[:,1] .== Gray(0))
    @test all(imgboxes[end,:] .== Gray(0))
    @test all(imgboxes[:,end] .== Gray(0))

    #nonsquare low aspect ratio image
    img = fill(Gray(1), 100, 200)
    batch = emptybatch(yolomod)
    batch[:,:,:,1], padding = prepareImage(img, yolomod)
    res = collect([ padding[1] padding[2] 1.0-padding[3] 1.0-padding[4] 0.0 0.0;]') #note the transpose!
    imgboxes = drawBoxes(img, yolomod, padding, res)
    @test all(imgboxes[1,:] .== Gray(0))
    @test all(imgboxes[:,1] .== Gray(0))
    @test all(imgboxes[end,:] .== Gray(0))
    @test all(imgboxes[:,end] .== Gray(0))

    ## Non-transposed
    imgboxes = drawBoxes(collect(img'), yolomod, padding, res, transpose=false)
    @test all(imgboxes[1,:] .== Gray(0))
    @test all(imgboxes[:,1] .== Gray(0))
    @test all(imgboxes[end,:] .== Gray(0))
    @test all(imgboxes[:,end] .== Gray(0))

end


#=
using .Soss
using .Soss: ConditionalModel, getmoduletypencoding, argvals, observations,
             Assign, Return, LineNumber, Sample, testvalue, @q, type2model,
             loadvals, getntkeys, buildSource, isleaf, _unwrap_type
using NestedTuples
using MacroTools
ascube(m::ConditionalModel{A, B}, _data::NamedTuple) where {A,B} = ascube(m | _data)

function ascube(m::ConditionalModel{A, B}) where {A,B}
    return _ascube(getmoduletypencoding(m), Model(m), argvals(m), observations(m))
end

# function ascube(m::Model{EmptyNTtype, B}) where {B}
#     return ascube(m,NamedTuple())
# end

ascube(d, _data) = nothing

ascube(μ::MT.AbstractMeasure,  _data::NamedTuple=NamedTuple()) = ascube(μ)
ascube(μ::Dists.Distribution,  _data::NamedTuple=NamedTuple()) = ascube(μ)

ascube(d::Dists.AbstractMvNormal, _data::NamedTuple=NamedTuple()) = ascube(d)
#ascube(::NamedTuple{(), Tuple{}}) = nothing

export sourceascube

sourceascube(m::Model) = sourceascube()(m)

function sourceascube(_data=NamedTuple())
    function(_m::Model)

        _datakeys = getntkeys(_data)
        proc(_m, st::Assign)        = :($(st.x) = $(st.rhs))
        proc(_m, st::Return)     = nothing
        proc(_m, st::LineNumber) = nothing

        function proc(_m, st::Sample)
            x = st.x
            xname = QuoteNode(x)
            rhs = st.rhs

            thecode = @q begin
                _t = ascube($rhs, get(_data, $xname, NamedTuple()))
                if !isnothing(_t)
                    _result = merge(_result, ($x=_t,))
                end
            end

            # Non-leaves might be referenced later, so we need to be sure they
            # have a value
            isleaf(_m, st.x) || pushfirst!(thecode.args, :($x = Soss.testvalue($rhs)))

            return thecode
        end

        wrap(kernel) = @q begin
            _result = NamedTuple()
            $kernel
            $ascube(_result)
        end

        buildSource(_m, proc, wrap) |> MacroTools.flatten

    end
end

@gg function _ascube(M::Type{<:GeneralizedGenerated.TypeLevel}, _m::Model{Asub,B}, _args::A, _data) where {Asub,A,B}
    body = type2model(_m) |> sourceascube(_data) |> loadvals(_args, _data)
    @under_global from_type(_unwrap_type(M)) @q let M
        $body
    end
end
=#


#=
MIT License

Copyright (c) 2017 Joseph R. Peterson

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.
=#

# External Libraries
using DataTables
using ArgParse

# Create command line parser and parse arguments
function parse_cmd_line()
	argSettings = ArgParseSettings()
	@add_arg_table argSettings begin
		"--expressionData", "-i"
			help = "A tab/comma separated file containing gene expression data."
			arg_type = String
			required = true
		"--outputDir", "-o"
			help = "Name of directory into which results are stored."
			arg_type = String
			default = "out"
		"--threshold", "-t"
			help = "A threshold on the minimum coefficient of variation for genes to be considered by the algorithm (pre-algorithm filtering)."
			arg_type = Float64
			default = 0.0
		"--significant", "-s"
			help = "Significant interactions are those 's' standard deviations away from mean (post-algorithm filtering)."
			arg_type = Float64
			default = 2.0
		"--plot", "-p"
			help = "Plot intermediate results"
			default = false
			action = :store_true
	end

	# Actually pargse aguments
	return parse_args(argSettings)
end

# Main computes the network
function main()
	# Get command line arguments
	args = parse_cmd_line()

	# Read data files
	println("Reading inputs...")
	expressionData = readtable(args["expressionData"])

	# Run algorithm
	println("Computing couplings...")
	genes, couplingStrength = plogp(expressionData, args["threshold"])

	# Cull network
	println("Thresholding significant interactions...")

	# Plot Data (if requested)
	if args["plot"]
		println("Plotting...")
	end

	# Save resulting weight and graph files
	println("Saving results...")

	# Successfully completed
	println("Done!")
end
main()


#module HierarchicalBetaBinomial

import Distributions: Beta, Binomial, Normal
import LinearAlgebra

include("./src/types.jl")
include("./src/MH.jl")

#end


using XLSX
export deaths
# function Sterbestatistik()

MAX = 200
export deaths_data

deaths_data() =
    (deaths2016_2021 = XLSX.readxlsx(joinpath(datadep"Sterbestatistik 2016-2021", "sonderauswertung-sterbefaelle.xlsx")),
     deaths2000_2015 = XLSX.readxlsx(joinpath(datadep"Sterbestatistik 2000-2015", "sonderauswertung-sterbefaelle-endgueltige-daten.xlsx")))

gender_code = Dict("männlich" => "Männlich",
                   "Männlich" => "Männlich",
                   "weiblich" => "Weiblich",
                   "Weiblich" => "Weiblich",
                   "insgesamt" => "Ins",
                   "Insgesamt" => "Ins",
                   )

Bundesländer_index = Dict(r => i for (i,r) in enumerate(Bundesländer))

Altersgruppen_Monate       = [ 0:15, 15:30, [a:a+5 for a in 30:5:80]..., 85:MAX ];
Altersgruppen_KW           = [ 0:30, [a:a+5 for a in 30:5:80]..., 85:MAX ];
Altersgruppen_Bundesländer = [ 0:65, 65:75, 75:85, 85:MAX]

"""
    Altersgruppen_Monate = [ 0:15, 15:30, [a:a+5 for a in 30:5:80]..., 85:MAX ];
    Altersgruppen_KW = [ 0:30, [a:a+5 for a in 30:5:80]..., 85:MAX ];
    Altersgruppen_Bundesländer = [0:65, 65:75, 75:85, 85:MAX]

Note: Sonderauswertung is more fine-grained, population data is only providing 85+ group.
"""
Altersgruppen_Monate, Altersgruppen_KW, Altersgruppen_Bundesländer

geschlechter = ["Männlich", "Weiblich"] # , "Insgesamt"

alterzeile = Dict(missing=>1, [r => i+2 for (i,r) in enumerate(Altersgruppen_Bundesländer)]...)

"""
    deaths(data; jahr, kw=missing, monat=missing, alter=missing, geschlecht="Insgesamt", bundesland = missing)
"""
function deaths(data;
                jahr, kw=missing, monat=missing,
                alter::Union{UnitRange,Missing}, geschlecht,
                bundesland = missing)
    # @assert geschlecht in geschlechter
    geschlecht = get(gender_code, geschlecht, geschlecht)

    dset, dsheet, joffset = if 2016<=jahr
        data.deaths2016_2021, 2016, 2021
    else
        data.deaths2000_2015, 2000, 2015
    end

    r = if kw === missing
        if monat === missing
	    return sum([ deaths(data; bundesland = bundesland, alter=alter, geschlecht=geschlecht, jahr=jahr, monat=i)
                         for i in 1:12
                             if deaths(data; bundesland = bundesland, alter=alter, geschlecht=geschlecht, jahr=jahr, kw=i) isa Number])
        else # monat !== missing

            if bundesland !== missing
                sheet = dset["BL_$(dsheet)_$(joffset)_Monate_AG_$geschlecht"]

                zeile = 5 + # skipping lines,
                    (joffset-jahr)*80 + # lines blocks for years,
                    5 * get(Bundesländer_index,"$bundesland") do  # lines blocks for Bundesländer
                        error("bundesland must be in  in $(Bundesländer)")
                    end +
                        get(alterzeile, alter) do # year steps
                            error("alter must be in $(keys(alterzeile))")
                        end - 1

                col = monat + 4
                sheet[zeile, col]
            else
                # skipping 11 lines,
                zeile = 11 +
                    # 16 lines blocks for years,
                    (joffset-jahr)*17 +
                    # 5 year steps beginning age 30
                    if alter.start == 0
                        0
                    elseif alter.start == 15
                        1
                    else
                        1 + max(0, div(alter.start-25,5))
                    end
                ## look up gender tab
                tab = dset["D_$(dsheet)-$(joffset)_Monate_AG_$geschlecht"]
                # and check in cw column, skipping first 3
                # @info "deaths lookup" "$dsheet$geschlecht", zeile, monat
                if alter==Altersgruppen_Monate[end]
                    sum([tab[i, monat+3] for i in zeile:(zeile+2)])
                else
                    tab[zeile, monat+3]
                end
	    end
        end
    else  # kw !== missing
        if monat !== missing
            error("proveide either monat or kw, but not both")
        end

        if bundesland !== missing
            sheet = dset["BL_$(dsheet)_$(joffset)_KW_AG_$geschlecht"]
            zeile = 5 + # skipping lines,
                (joffset-jahr)*80 + # lines blocks for years,
                5 * get(Bundesländer_index,"$bundesland") do  # lines blocks for Bundesländer
                    error("bundesland must be in  in $(Bundesländer)")
                end +
                    get(alterzeile, alter) do # year steps
                        error("alter must be in $(keys(alterzeile))")
                    end - 2

            col = kw + 4
            sheet[zeile, col]
        else

            # @assert alter isa UnitRange && alter.stop-alter.start==5 && alter.start<=95
            if alter.start >= 95
	        #@warn "alter" alter
                alter = 95:MAX
            end
            if alter.stop <= 30
	        #@warn "alter" alter
                alter = 0:30
            end
            # skipping 11 lines,
            zeile = 11 +
                # 16 lines blocks for years,
                (joffset-jahr)*16 +
                # 5 year steps beginning age 30
                max(0, div(alter.start-25,5))
            ## look up gender tab
            tab = dset["D_$(dsheet)_$(joffset)_KW_AG_$geschlecht"]
            # @info "deaths lookup" "$dsheet$geschlecht", zeile, kw
            # and check in cw column, skipping first 3
            if alter==Altersgruppen_KW[end]
                sum([tab[i, kw+3] for i in zeile:(zeile+2)])
            else
                tab[zeile, kw+3]
            end
        end
    end
    r isa Number ? r : missing
end

# end


# Core functions for evaluating PDDL formulae and state transitions

"""
    satisfy(formulae, state, domain=nothing; mode=:any)

Returns whether `formulae` can be satisfied by facts in the given `state`,
and any axioms / derived predicates in `domain`. If `formulae` contains
unbound variables, `mode=:any` returns the first substitution found,
while `mode=:all` returns all substitutions.
"""
function satisfy(formulae::Vector{<:Term}, state::State,
                 domain::Union{Domain,Nothing}=nothing; mode::Symbol=:any)
    # Do quick check as to whether formulae are in the set of facts
    function in_facts(f::Term)
        if !isempty(state.fluents)
            f = eval_term(f, Subst(), state.fluents) end
        if f in state.facts || f in state.types
            return true end
        if !isnothing(domain) && f in get_const_facts(domain)
            return true end
        if f.name in Julog.comp_ops || f.name in keys(state.fluents)
            return eval_term(f, Subst(), state.fluents).name == true end
        return false
    end
    if all(f -> f.name == :not ? !in_facts(f.args[1]) : in_facts(f), formulae)
        return true, [Subst()] end
    # Initialize Julog knowledge base
    clauses = isnothing(domain) ? Clause[] : get_clauses(domain)
    clauses = Clause[clauses; collect(state.types); collect(state.facts)]
    # Pass in fluents and function definitions as a dictionary of functions
    funcs = isnothing(domain) ?
        state.fluents : merge(state.fluents, domain.funcdefs)
    return resolve(formulae, clauses; funcs=funcs, mode=mode)
end

satisfy(formula::Term, state::State, domain::Union{Domain,Nothing}=nothing;
        options...) = satisfy(Term[formula], state, domain; options...)

"""
    evaluate(formula, state, domain=nothing; as_const=true)

Evaluates `formula` as fully as possible with respect to the fluents defined
in `state`, along with any axioms / derived predicates defined in `domain`.
Returns a Julog `Const` if `as_const=true`, otherwise return an unwrapped
Julia value.
"""
function evaluate(formula::Term, state::State,
                  domain::Union{Domain,Nothing}=nothing; as_const::Bool=true)
    # Evaluate formula as fully as possible
    funcs = isnothing(domain) ?
        state.fluents : merge(state.fluents, domain.funcdefs)
    val = eval_term(formula, Subst(), funcs)
    # Return if formula evaluates to a Const (unwrapping if as_const=false)
    if isa(val, Const) && !isa(val.name, Symbol)
        return as_const ? val : val.name end
    # If val is not a Const, check if holds true in the state
    sat, _ = satisfy(Term[val], state, domain)
    return as_const ? Const(sat) : sat # Wrap in Const if as_const=true
end

"""
    find_matches(formula, state, domain=nothing)

Returns a list of all matching substitutions of `formula` with respect to
a given `state` and `domain`.
"""
function find_matches(formula::Term, state::State,
                      domain::Union{Domain,Nothing}=nothing)
    if formula.name in keys(state.fluents)
        clauses = Vector{Clause}(get_fluents(state))
        _, subst = resolve(formula, clauses; mode=:all)
    else
        clauses = isnothing(domain) ? Clause[] : get_clauses(domain)
        clauses = Clause[clauses; collect(state.types); collect(state.facts)]
        funcs = state.fluents
        _, subst = resolve(formula, clauses; funcs=funcs, mode=:all)
    end
    matches = Term[substitute(formula, s) for s in subst]
    return matches
end

"Construct initial state from problem definition."
function init_state(problem::Problem)
    types = Term[@julog($ty(:o)) for (o, ty) in problem.objtypes]
    state = State(problem.init, types)
    return state
end

"Construct goal state from problem definition."
function goal_state(problem::Problem)
    types = Term[@julog($ty(:o)) for (o, ty) in problem.objtypes]
    state = State(flatten_conjs(problem.goal), types)
    return state
end

"Construct initial state from problem definition."
initialize(problem::Problem) = init_state(problem)

"""
    transition(domain, state, action::Term; kwargs...)
    transition(domain, state, actions::Set{<:Term}; kwargs...)

Returns the successor to `state` in the given `domain` after applying a single
`action` or a set of `actions` in parallel, along with any events triggered
by the effects of those actions. Keyword arguments specify whether to `check`
if action preconditions hold, and the `fail_mode` (`:error` or `:no_op`)
if they do not.
"""
function transition(domain::Domain, state::State, action::Term;
                    check::Bool=true, fail_mode::Symbol=:error)
    state = execute(action, state, domain; check=check, fail_mode=fail_mode)
    if length(domain.events) > 0
        state = trigger(domain.events, state, domain)
    end
    return state
end

function transition(domain::Domain, state::State, actions::Set{<:Term};
                    check::Bool=true, fail_mode::Symbol=:error)
    # Execute all actions in parallel
    state = execpar(actions, state, domain, check=check, fail_mode=fail_mode)
    if length(domain.events) > 0
        state = trigger(domain.events, state, domain)
    end
    return state
end

"""
    simulate(domain, state, actions; kwargs...)

Returns the state trajectory that results from applying a sequence of `actions`
to an initial `state` in a given `domain`. Keyword arguments specify whether
to `check` if action preconditions hold, the `fail_mode` (`:error` or `:no_op`)
if they do not, and a `callback` function to apply after each step.
"""
function simulate(domain::Domain, state::State, actions::Vector{<:Term};
                  check::Bool=true, fail_mode::Symbol=:error,
                  callback::Function=(d,s,a)->nothing)
    trajectory = State[state]
    callback(domain, state, Const(:start))
    for act in actions
        state = transition(domain, state, act; check=check, fail_mode=fail_mode)
        push!(trajectory, state)
        callback(domain, state, act)
    end
    return trajectory
end

function simulate(domain::Domain, state::State, actions::Vector{Set{<:Term}};
                  check::Bool=true, fail_mode::Symbol=:error,
                  callback::Function=(d,s,a)->nothing)
    trajectory = State[state]
    callback(domain, state, Set([Const(:start)]))
    for acts in actions
        state = transition(domain, state, acts;
                           check=check, fail_mode=fail_mode)
        push!(trajectory, state)
        callback(domain, state, acts)
    end
    return trajectory
end


export apply_back!, apply_back

const Rotor{N,T} = Union{RotationGate{N,T},PutBlock{N,<:Any,<:RotationGate{<:Any,T}}}

as_scalar(arr::AbstractArray{T,0}) where {T} = arr[]
as_scalar(arr) = arr

"""
    generator(rot::Rotor) -> AbstractBlock

Return the generator of rotation block.
"""
generator(rot::RotationGate) = rot.block
generator(rot::PutBlock{N,C,GT}) where {N,C,GT<:RotationGate} =
    PutBlock{N}(generator(rot |> content), rot |> occupied_locs)

"""
    apply_back!((ψ, ∂L/∂ψ*), circuit::AbstractBlock, collector) -> AbstractRegister

back propagate and calculate the gradient ∂L/∂θ = 2*Re(∂L/∂ψ*⋅∂ψ*/∂θ), given ∂L/∂ψ*.
`ψ` is the output register, ∂L/∂ψ* should also be register type.

Note: gradients are stored in `Diff` blocks, it can be access by either `diffblock.grad` or `gradient(circuit)`.
Note2: now `apply_back!` returns the inversed gradient!
"""
function apply_back!(st, block::AbstractBlock, collector) #,AbstractContainer{<:PrimitiveBlock}
    out, outδ = st
    if nparameters(block) == 0
        adjblock = block'
        in = apply!(out, adjblock)
        inδ = apply!(outδ, adjblock)
        return (in, inδ)
    else
        throw(MethodError(apply_back, (st, block, collector)))
    end
end

function apply_back!(st, block::Concentrator{N}, collector) where {N}
    out, outδ = st
    focus!(out, block.locs)
    focus!(outδ, block.locs)
    apply_back!((out, outδ), content(block), collector)
    relax!(out, block.locs; to_nactive = N)
    relax!(outδ, block.locs; to_nactive = N)
    return (out, outδ)
end

function apply_back!(st, block::Rotor{N}, collector) where {N}
    out, outδ = st
    adjblock = block'
    backward_params!((out, outδ), block, collector)
    in = apply!(out, adjblock)
    inδ = apply!(outδ, adjblock)
    return (in, inδ)
end

function apply_back!(st, block::TimeEvolution{N}, collector) where {N}
    out, outδ = st
    adjblock = block'

    out, outδ = st
    input = apply!(out, adjblock)
    for o in outδ
        !all(x -> x ≈ 0.0im, o.state) && apply!(o, adjblock)
    end
    pushfirst!(collector, -sum(imag(input' * apply!(copy(outδ), block.H))))
    return (input, outδ)
end

function apply_back!(st, block::PutBlock{N}, collector) where {N}
    out, outδ = st
    adjblock = block'
    in = apply!(out, adjblock)
    adjmat = outerprod(outδ, in)
    mat_back!(datatype(in), block, adjmat, collector)
    inδ = apply!(outδ, adjblock)
    return (in, inδ)
end

function apply_back!(st, block::KronBlock{N}, collector) where {N}
    apply_back!(st, chain(N, [put(loc => block[loc]) for loc in block.locs]), collector)
end

function apply_back!(st, block::ControlBlock{N}, collector) where {N}
    out, outδ = st
    adjblock = block'
    in = apply!(out, adjblock)
    #adjm = adjcunmat(outerprod(in, outδ), N, block.ctrl_locs, block.ctrl_config, mat(content(block)), block.locs)
    adjmat = outerprod(outδ, in)
    mat_back!(datatype(in), block, adjmat, collector)
    inδ = apply!(outδ, adjblock)
    return (in, inδ)
end

function apply_back!(st, block::Daggered, collector)
    out, outδ = st
    adjblock = block'
    in = apply!(out, adjblock)
    adjmat = outerprod(in, outδ)
    mat_back!(datatype(in), content(block), adjmat, collector)
    inδ = apply!(outδ, adjblock)
    return (in, inδ)
end

function apply_back!(st, block::Scale, collector)
    out, outδ = st
    apply_back!((out, outδ), content(block), collector)
    outδ.state .= outδ.state .* conj(factor(block))
    out.state .= out.state ./ factor(block)
    return (out, outδ)
end

function apply_back!(st, circuit::ChainBlock, collector)
    for blk in Base.Iterators.reverse(subblocks(circuit))
        st = apply_back!(st, blk, collector)
    end
    return st
end

function apply_back!(st, circuit::Add, collector; in)
    out, outδ = st
    adjmat = outerprod(outδ, in)
    for blk in Base.Iterators.reverse(subblocks(circuit))
        mat_back!(datatype(in), blk, adjmat, collector)
    end
    inδ = apply!(outδ, adjblock)
    (in, inδ)
end

function apply_back!(st, block::RepeatedBlock{N,C}, collector) where {N,C}
    if nparameters(content(block)) == 0
        return apply!.(st, Ref(block'))
    end
    res = Any[]
    st = apply_back!(st, chain(N, [put(loc => content(block)) for loc in block.locs]), res)
    res = dropdims(sum(reshape(res, :, C), dims = 2), dims = 2) |> as_scalar
    prepend!(collector, res)
    return st
end

# TODO: concentrator, repeat, kron
apply_back!(st, block::Measure, collector) = throw(MethodError(apply_back!, (st, block, collector)))

function backward_params!(st, block::Rotor, collector)
    in, outδ = st
    Σ = generator(block)
    g = dropdims(sum(conj.(state(in |> Σ)) .* state(outδ), dims = (1, 2)), dims = (1, 2)) |> as_scalar
    pushfirst!(collector, -imag(g) / 2)
    in |> Σ
    nothing
end

"""
    apply_back(st::Tuple{<:ArrayReg, <:ArrayReg}, block::AbstractBlock; kwargs...) -> (out, outδ), paramsδ

The backward function of `apply!`. Returns a tuple of ((input register, gradient of input register), parameter gradients)
"""
function apply_back(st::Tuple{<:ArrayReg,<:ArrayReg}, block::AbstractBlock; kwargs...)
    col = []
    in, inδ = apply_back!(st, block, col; kwargs...)
    (in, inδ), col
end

Base.adjoint(::typeof(expect)) = Adjoint(expect)
Base.show(io::IO, ::Adjoint{Any,typeof(expect)}) = print(io, "expect'")
Base.show(io::IO, ::MIME"text/plain", ::Adjoint{Any,typeof(expect)}) = print(io, "expect'")
"""
expect')(op::AbstractBlock, circuit::Pair{<:ArrayReg, <:AbstractBlock})

"""
function (::Adjoint{Any,typeof(expect)})(op::AbstractBlock, circuit::Pair{<:ArrayReg,<:AbstractBlock})
    reg, c = circuit
    out = copy(reg) |> c
    outδ = copy(out) |> op
    (in, inδ), paramsδ = apply_back((out, outδ), c)
    return outδ => paramsδ .* 2
end


function construct_network!(
    psi_container::PSIContainer,
    sys::PSY.System,
    ::Type{CopperPlatePowerModel},
)
    buses = PSY.get_components(PSY.Bus, sys)
    bus_count = length(buses)

    get_slack_variables(psi_container.settings) &&
        add_slacks!(psi_container, CopperPlatePowerModel)
    copper_plate(psi_container, :nodal_balance_active, bus_count)

    return
end

function construct_network!(
    psi_container::PSIContainer,
    sys::PSY.System,
    ::Type{StandardPTDFModel},
)
    buses = PSY.get_components(PSY.Bus, sys)
    ac_branches = get_available_components(PSY.ACBranch, sys)
    ptdf = get_PTDF(psi_container)

    if isnothing(ptdf)
        throw(ArgumentError("no PTDF matrix supplied"))
    end

    get_slack_variables(psi_container.settings) &&
        add_slacks!(psi_container, StandardPTDFModel)

    ptdf_networkflow(psi_container, ac_branches, buses, :nodal_balance_active, ptdf)

    dc_branches = get_available_components(PSY.DCBranch, sys)
    dc_branch_types = typeof.(dc_branches)
    for btype in Set(dc_branch_types)
        typed_dc_branches = IS.FlattenIteratorWrapper(
            btype,
            Vector([[b for b in dc_branches if typeof(b) == btype]]),
        )
        flow_variables!(psi_container, StandardPTDFModel, typed_dc_branches)
    end
    return
end

function construct_network!(
    psi_container::PSIContainer,
    sys::PSY.System,
    ::Type{T},
) where {T <: PM.AbstractPowerModel}
    incompat_list = [
        PM.SDPWRMPowerModel,
        PM.SparseSDPWRMPowerModel,
        PM.SOCBFPowerModel,
        PM.SOCBFConicPowerModel,
    ]
    if T in incompat_list
        throw(ArgumentError("$(T) formulation is not currently supported in PowerSimulations"))
    end

    get_slack_variables(psi_container.settings) && add_slacks!(psi_container, T)

    powermodels_network!(psi_container, T, sys)
    add_pm_var_refs!(psi_container, T, sys)
    return
end