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 Flux: OneHotMatrix
using LinearAlgebra

function less_categorical_scalar_extractor()
	[
	(e -> (keys_len = length(keys(e)); keys_len / e.updated < 0.5 && length(keys(e)) <= 10 && JsonGrinder.is_numeric_or_numeric_string(e)),
		(e, uniontypes) -> ExtractCategorical(keys(e), uniontypes)),
	(e -> JsonGrinder.is_intable(e),
		(e, uniontypes) -> JsonGrinder.extractscalar(Int32, e, uniontypes)),
	(e -> JsonGrinder.is_floatable(e),
	 	(e, uniontypes) -> JsonGrinder.extractscalar(JsonGrinder.FloatType, e, uniontypes)),
	# it's important that condition here would be lower than maxkeys
	(e -> (keys_len = length(keys(e)); keys_len / e.updated < 0.1 && keys_len < 10000 && !JsonGrinder.is_numeric_or_numeric_string(e)),
		(e, uniontypes) -> ExtractCategorical(keys(e), uniontypes)),
	(e -> true,
		(e, uniontypes) -> JsonGrinder.extractscalar(JsonGrinder.unify_types(e), e, uniontypes)),]
end

testing_settings = (; scalar_extractors = less_categorical_scalar_extractor())

function with_emptyismissing(f::Function, a)
    orig_val = Mill.emptyismissing()
    Mill.emptyismissing!(a)
    f()
    Mill.emptyismissing!(orig_val)
end

@testset "ExtractScalar" begin
    @testset "with uniontypes" begin
	    sc = ExtractScalar(Float64,2,3,true)
	    @test sc.uniontypes == true
	    @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{Union{Missing, Float64}}
	    @test nobs(sc(extractempty)) == 0
	    @test nobs(sc(5)) == 1

	    sc = ExtractScalar(Float32, 0.5, 4.0, true)
	    @test sc.uniontypes == true
	    @test sc(1).data isa Matrix{Union{Missing, Float32}}
	    @test sc(extractempty).data isa Matrix{Union{Missing, Float32}}

	    sc = JsonGrinder.extractscalar(Float32, true)
	    @test sc.uniontypes == true
	    @test sc(1).data isa Matrix{Union{Missing, Float32}}
	    @test sc(Dict(1=>1)).data isa Matrix{Union{Missing, Float32}}
	    @test length(sc) == 1

		@testset "store_input" begin
    		@test sc("5", store_input=true).data == sc("5", store_input=false).data
			@test sc("5", store_input=true).metadata == fill("5",1,1)
			@test isnothing(sc("5", store_input=false).metadata)
			@test sc(5, store_input=true).data == sc(5, store_input=false).data
			@test sc(5, store_input=true).metadata == fill(5,1,1)
			@test isnothing(sc(5, store_input=false).metadata)
			@test sc(nothing, store_input=true).data ≃ sc(nothing, store_input=false).data
			@test sc(nothing, store_input=true).metadata == fill(nothing,1,1)
			@test isnothing(sc(nothing, store_input=false).metadata)
		end
    end

    @testset "without uniontypes" begin
    	sc1 = ExtractScalar(Float64,2,3,false)
    	@test sc1.uniontypes == false
    	@test all(sc1("5").data .== [9])
    	@test all(sc1(5).data .== [9])
    	@test_throws ErrorException sc1(nothing)
    	@test_throws ErrorException sc1(missing)
		@test nobs(sc1(5)) == 1

    	sc2 = ExtractScalar(Float32, 0.5, 4.0, false)
    	@test sc2.uniontypes == false
    	@test sc2(1).data isa Matrix{Float32}
    	sc3 = JsonGrinder.extractscalar(Float32,false)
    	@test sc3.uniontypes == false
    	@test sc3(1).data isa Matrix{Float32}
    	@test_throws ErrorException sc3(Dict(1=>1))
    	@test length(sc3) == 1

		@testset "extractempty" begin
    		@test sc1(extractempty).data isa Matrix{Float64}
    		@test nobs(sc1(extractempty)) == 0
			@test sc2(extractempty).data isa Matrix{Float32}
	    	@test isnothing(sc2(extractempty, store_input=false).metadata)
	    	@test sc2(extractempty, store_input=true).metadata isa Matrix{UndefInitializer}
	    	@test size(sc2(extractempty, store_input=true).metadata) == (1,0)
		end

		@testset "store_input" begin
    		@test sc3("5", store_input=true).data == sc3("5", store_input=false).data
       		@test sc3("5", store_input=true).metadata == fill("5",1,1)
       		@test isnothing(sc3("5", store_input=false).metadata)
       		@test sc3(5, store_input=true).data == sc3(5, store_input=false).data
       		@test sc3(5, store_input=true).metadata == fill(5,1,1)
       		@test isnothing(sc3(5, store_input=false).metadata)
       		@test_throws ErrorException sc3(nothing, store_input=true)
       		@test_throws ErrorException sc3(nothing, store_input=true)
		end
    end
end
# todo: add for each extractor tests with and without uniontypes, and with and without store_input
# todo: categorize tests w.r.t. individual extractors, make a checklist of what each extractor needs to test
@testset "ExtractArray" begin
	sc = ExtractArray(ExtractCategorical(2:4))
	with_emptyismissing(false) do
		e234 = sc([2,3,4], store_input=false)
		e234 = sc([2,3,4], store_input=false)
		en = sc(nothing, store_input=false)
		e234s = sc([2,3,4], store_input=true)
		ens = sc(nothing, store_input=true)
		@test all(e234.data.data .== Matrix(1.0I, 4, 3))
		@test nobs(en.data) == 0
		@test en.data.data isa MaybeHotMatrix{Union{Missing, UInt32}, UInt32, Union{Missing, Bool}}
		@test nobs(en.data.data) == 0
		@test all(en.bags.bags .== [0:-1])

		@testset "store_input" begin
			@test e234s.data.data == e234.data.data
			@test e234s.data.metadata == [2,3,4]
			@test e234s.data[1].metadata == [2]
			@test e234s.data[2].metadata == [3]
			@test e234s.data[3].metadata == [4]
			@test ens.data.data == en.data.data
			@test ens.metadata == [nothing]
			@test isnothing(en.metadata)
			@test isnothing(ens.data.metadata)
			@test isnothing(e234.data.metadata)
			@test isnothing(en.data.metadata)
		end
	end
	with_emptyismissing(true) do
		en = sc(nothing, store_input=false)
		ens = sc(nothing, store_input=true)
		em = sc(missing, store_input=false)
		ems = sc(missing, store_input=true)
		@test en.data isa Missing
		@test en.data ≃ em.data
		@test ens.data isa Missing
		@test ens.data ≃ ems.data
		@test isnothing(en.metadata)
		@test isnothing(em.metadata)
		@test ens.metadata == [nothing]
		@test ems.metadata ≃ [missing]
		@test all(en.bags.bags .== [0:-1])
		@test all(em.bags.bags .== [0:-1])
	end

	@testset "extractempty" begin
		@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{Union{Missing, UInt32}, UInt32, Union{Missing, Bool}}
		with_emptyismissing(true) do
			@test nobs(sc(extractempty)) == 0
		end
		with_emptyismissing(false) do
			@test nobs(sc(extractempty)) == 0
		end
	end
	sc = ExtractArray(ExtractScalar(Float32))
	e234 = sc([2,3,4], store_input=false)
	en = sc(nothing, store_input=false)
	e234s = sc([2,3,4], store_input=true)
	ens = sc(nothing, store_input=true)
	@test all(e234.data.data .== [2 3 4])
	@test nobs(en.data) == 0
	@test all(en.bags.bags .== [0:-1])
	@test nobs(sc(Dict(1=>1)).data) == 0

	@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{Union{Missing, Float32}}

	@test en.data.data isa Matrix{Union{Missing, Float32}}
	@test e234s.data.data == e234.data.data
	@test e234s.data.metadata == [2 3 4]
	@test e234s.data[1].metadata == fill(2,1,1)
	@test e234s.data[2].metadata == fill(3,1,1)
	@test e234s.data[3].metadata == fill(4,1,1)
	@test ens.data.data == en.data.data
	@test isnothing(en.data.metadata)
	@test ens.metadata == [nothing]
	@test isnothing(e234.data.metadata)
	@test isnothing(en.data.metadata)

	# testing https://github.com/CTUAvastLab/JsonGrinder.jl/issues/76
	e = ExtractArray(ExtractScalar(Float32))
	@test e([1,2,3]) == BagNode(
		ArrayNode([1. 2. 3.]),
		[1:3]
	)
	@test_throws MethodError e((1,2,3))
	@test e(nothing) == BagNode(
		ArrayNode(Matrix(fill(zero(Float32),1,0))),
		[0:-1]
	)
	@test_throws MethodError e(e)
end

@testset "ExtractVector" begin
	@testset "with uniontypes" begin
		sc1 = ExtractVector(5, true)
		e1 = sc1([1, 2, 2, 3, 4], store_input=false)
		e1s = sc1([1, 2, 2, 3, 4], store_input=true)
		e2s = sc1([1, 2, 2, 3], store_input=true)
		n1 = sc1(missing, store_input=false)
		n1s = sc1(missing, store_input=true)
		@test sc1.uniontypes == true
		@test e1.data == e1s.data
		@test e1.data isa Matrix{Union{Missing, Float32}}

		sc2 = ExtractVector{Int64}(5, true)
		e12234 = sc2([1, 2, 2, 3, 4], store_input=false)
		e12234s = sc2([1, 2, 2, 3, 4], store_input=true)
		en = sc2(nothing, store_input=false)
		ens = sc2(nothing, store_input=true)
		@test sc2.uniontypes == true
		@test e12234.data ≈ [1, 2, 2, 3, 4]
		@test e12234.data isa Matrix{Union{Missing, Int64}}
		@test e12234s.data ≈ e12234.data
		@test ens.data ≃ en.data
		@test en.data isa Matrix{Union{Missing, Int64}}
		@test en.data ≃ [missing missing missing missing missing]'

		sc3 = ExtractVector{Float32}(5, true)
		@test sc3.uniontypes == true
		# feature vector longer than expected
		sc122345 = sc3([1, 2, 2, 3, 4, 5], store_input=false)
		sc122345s = sc3([1, 2, 2, 3, 4, 5], store_input=true)
		sc12345 = sc3([1, 2, 3, 4, 5])
		@test sc122345.data == [1f0 2f0 2f0 3f0 4f0]'
		@test sc12345.data isa Matrix{Union{Missing, Float32}}
		e56 = sc3([5, 6], store_input=false)
		e56s = sc3([5, 6], store_input=true)
		@test e56.data ≃ [5f0 6f0 missing missing missing]'
		@test e56s.data ≃ e56.data
		@test sc122345s.data ≈ sc122345.data
		@test all(sc3(Dict(1=>2)).data .=== missing)

		@testset "extractempty" begin
			@test sc1(extractempty).data isa Matrix{Union{Missing, Float32}}
			@test nobs(sc1(extractempty).data) == 0
			@test nobs(sc1(extractempty)) == 0
			@test sc2(extractempty).data isa Matrix{Union{Missing, Int64}}
			@test nobs(sc2(extractempty).data) == 0
			@test nobs(sc2(extractempty)) == 0
		end

		@testset "store_input" begin
			@test e1s.metadata == [[1, 2, 2, 3, 4]]
			@test isnothing(e1.metadata)
			@test catobs(e1s, e1s).data == [1 1; 2 2; 2 2; 3 3; 4 4]
			@test catobs(e1s, e1s).metadata == [[1, 2, 2, 3, 4], [1, 2, 2, 3, 4]]
			@test catobs(e1s, e1s)[1].data == [1 2 2 3 4]'
			@test catobs(e1s, e1s)[1].metadata == [[1, 2, 2, 3, 4]]
			@test n1s.metadata ≃ [missing]
			@test catobs(e1s, n1s).metadata ≃ [[1, 2, 2, 3, 4], missing]
			@test e2s.data ≃ [1 2 2 3 missing]'
			@test e2s.metadata == [[1, 2, 2, 3]]
			@test e12234s.metadata == [[1, 2, 2, 3, 4]]
			@test ens.metadata == [nothing]
			@test sc122345s.metadata == [[1, 2, 2, 3, 4, 5]]
			@test isnothing(e56.metadata)
			@test e56s.metadata == [[5, 6]]
		end
	end
	@testset "without uniontypes" begin
		sc1 = ExtractVector(5, false)
		@test sc1.uniontypes == false
		e1 = sc1([1, 2, 2, 3, 4], store_input=false)
		e1s = sc1([1, 2, 2, 3, 4], store_input=true)
		@test e1.data == [1f0 2f0 2f0 3f0 4f0]'
		@test e1.data isa Matrix{Float32}
		@test e1.data == e1s.data

		sc2 = ExtractVector{Int64}(5, false)
		@test sc2.uniontypes == false
		@test all(sc2([1, 2, 2, 3, 4]).data .== [1, 2, 2, 3, 4])
		@test sc2([1, 2, 2, 3, 4]).data isa Array{Int64, 2}

		sc3 = ExtractVector{Float32}(5, false)
		@test sc3.uniontypes == false

		# feature vector longer than expected
		@test sc3([1, 2, 2, 3, 4, 5]).data == [1f0 2f0 2f0 3f0 4f0]'
		@test sc3([1, 2, 3, 4, 5]).data isa Matrix{Float32}

		@testset "extractempty" begin
			@test sc1(extractempty).data isa Matrix{Float32}
			@test nobs(sc1(extractempty).data) == 0
			@test sc2(extractempty).data isa Matrix{Int64}
			@test nobs(sc2(extractempty).data) == 0
			@test sc3(extractempty).data isa Matrix{Float32}
			@test nobs(sc3(extractempty).data) == 0
		end

		@testset "store_input" begin
			@test e1s.metadata == [[1, 2, 2, 3, 4]]
			@test catobs(e1s, e1s).data == [1 1; 2 2; 2 2; 3 3; 4 4]
			@test catobs(e1s, e1s).metadata == [[1, 2, 2, 3, 4], [1, 2, 2, 3, 4]]
			@test catobs(e1s, e1s)[1].data == [1 2 2 3 4]'
			@test catobs(e1s, e1s)[1].metadata == [[1, 2, 2, 3, 4]]
			@test isnothing(e1.metadata)
		end

		@testset "missing" begin
			@test_throws ErrorException sc1(nothing, store_input=false)
			@test_throws ErrorException sc1(nothing, store_input=true)
			@test_throws ErrorException sc2(nothing, store_input=false)
			@test_throws ErrorException sc2(nothing, store_input=true)
			@test_throws ErrorException sc1(missing)
			@test_throws ErrorException sc2(missing)
			@test_throws ErrorException sc1([5, 6])
			@test_throws ErrorException sc3([5, 6])
			@test_throws ErrorException sc3(Dict(1=>2))
		end
	end
end

@testset "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 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 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 a3[:b].data ≃ hcat(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])

	@testset "extractempty" begin
		a4 = br(extractempty)
		@test nobs(a4) == 0
		@test nobs(a4[:a]) == 0
		@test a4[:a].data isa Matrix{Union{Missing, Float64}}
		@test nobs(a4[:b]) == 0
		@test a4[:b].data isa Matrix{Union{Missing, Float64}}
		@test nobs(a4[:c]) == 0
		@test nobs(a4[:c].data) == 0
		@test a4[:c].data.data isa Matrix{Union{Missing, Float64}}
	end
	@testset "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]), store_input=false)
		a2 = br(Dict("b" => [2,3,4]), store_input=false)
		a3 = br(Dict("a" => [2,3,4]), store_input=false)
		a4 = br(Dict{String, Any}(), store_input=false)
		a1s = br(Dict("a" => [1,2,3], "b" => [1,2,3,4]), store_input=true)
		a2s = br(Dict("b" => [2,3,4]), store_input=true)
		a3s = br(Dict("a" => [2,3,4]), store_input=true)
		a4s = br(Dict{String, Any}(), store_input=true)

		@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 Matrix{Union{Missing, Float32}})
		# todo: ověřit catobsování empty bagů s plnými bagy s metadaty jestli to funguje u všech možných extraktorů uvnitř bagu
		@test catobs(a1,a4).data[1].data.data == catobs(a1s,a4s).data[1].data.data
		@test catobs(a1,a4).data[1].bags == catobs(a1s,a4s).data[1].bags
		@test catobs(a1,a4).data[2].data.data == catobs(a1s,a4s).data[2].data.data
		@test catobs(a1,a4).data[2].bags == catobs(a1,a4).data[2].bags

		@test catobs(a1s,a4s).data[1].data.metadata == [1 2 3 4]
		@test catobs(a1s,a4s).data[2].data.metadata == [1 2 3]
		@test a1s.data[1].data.metadata == [1 2 3 4]
		@test a1s.data[2].data.metadata == [1 2 3]
		@test a4s.data[1].data.metadata == zeros(1,0)
		@test a4s.data[2].data.metadata == zeros(1,0)

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

@testset "ExtractCategorical" begin
	e = ExtractCategorical(["a","b"], true)
	ea = e("a", store_input=false)
	eb = e("b", store_input=false)
	ez = e("z", store_input=false)
	en = e(nothing, store_input=false)
	em = e(missing, store_input=false)
	eas = e("a", store_input=true)
	ebs = e("b", store_input=true)
	ezs = e("z", store_input=true)
	ens = e(nothing, store_input=true)
	@test e.uniontypes == true
	@test ea.data ≈ [1, 0, 0]
	@test eb.data ≈ [0, 1, 0]
	@test ez.data ≈ [0, 0, 1]
	@test en.data ≃ [missing missing missing]'
	@test em.data ≃ [missing missing missing]'
	@test typeof(ea.data) == MaybeHotMatrix{Union{Missing, UInt32}, UInt32, Union{Missing, Bool}}
	@test typeof(en.data) == MaybeHotMatrix{Union{Missing, UInt32}, UInt32, Union{Missing, Bool}}
	@test typeof(em.data) == MaybeHotMatrix{Union{Missing, UInt32}, UInt32, Union{Missing, Bool}}
	@test e(extractempty).data isa MaybeHotMatrix{Union{Missing, UInt32}, UInt32, Union{Missing, Bool}}
	@test nobs(e(extractempty)) == 0

	@test e(["a", "b"]).data ≃ [missing missing missing]'
	@test mapreduce(e, catobs, ["a", "b"]).data ≈ [1 0; 0 1; 0 0]
	@test e(["a", missing]).data ≃ [missing missing missing]'
	@test mapreduce(e, catobs, ["a", missing]).data ≃ [true missing; false missing; false missing]
	@test e(["a", missing, "x"]).data ≃ [missing missing missing]'
	@test mapreduce(e, catobs, ["a", missing, "x"]).data ≃ [true missing false; false missing false; false missing true]
	@test typeof(e(["a", "b"]).data) == MaybeHotMatrix{Union{Missing, UInt32}, UInt32, Union{Missing, Bool}}
	@test typeof(mapreduce(e, catobs, ["a", "b"]).data) == MaybeHotMatrix{Union{Missing, UInt32}, UInt32, Union{Missing, Bool}}
	@test typeof(e(["a", "b", nothing]).data) == MaybeHotMatrix{Union{Missing, UInt32}, UInt32, Union{Missing, Bool}}
	@test typeof(mapreduce(e, catobs, ["a", "b", nothing]).data) == MaybeHotMatrix{Union{Missing, UInt32}, UInt32, Union{Missing, Bool}}

	@test catobs(ea, eb).data ≈ [1 0; 0 1; 0 0]
	@test reduce(catobs, [ea.data, eb.data]) ≈ [1 0; 0 1; 0 0]
	@test hcat(ea.data, eb.data) ≈ [1 0; 0 1; 0 0]
	@test e(Dict(1=>2)).data |> collect ≃ [missing missing missing]'

	@test catobs(eas, ebs).data == catobs(ea, eb).data
	@test reduce(catobs, [eas.data, ebs.data]) == reduce(catobs, [ea.data, eb.data])
	@test e(Dict(1=>2)).data ≃ [missing missing missing]'

	@test catobs(eas, ebs).metadata == ["a", "b"]
	@test e(Dict(1=>2), store_input=true).metadata == [Dict(1=>2)]

	@test nobs(ea) == 1
	@test nobs(eb) == 1
	@test nobs(ez) == 1
	@test nobs(en) == 1
	@test nobs(em) == 1
	@test nobs(em.data) == 1
	@test nobs(e([missing, nothing])) == 1
	@test nobs(mapreduce(e, catobs, [missing, nothing])) == 2
	@test nobs(e([missing, nothing, "a"])) == 1
	@test nobs(mapreduce(e, catobs, [missing, nothing, "a"])) == 3

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

	e3 = ExtractCategorical(JsonGrinder.Entry(Dict(1=>1,2=>1), 2), true)
	@test e3.uniontypes == true
	@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 e3([]).data ≃ [missing missing missing]'

	e4 = ExtractCategorical(JsonGrinder.Entry(Dict(1.0=>1,2.0=>1), 2), true)
	@test e4.uniontypes == true
	@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 e4([]).data ≃ [missing missing missing]'

	e = ExtractCategorical(["a","b"], false)
	@test e.uniontypes == false
	@test e("a").data ≈ [1, 0, 0]
	@test e("b").data ≈ [0, 1, 0]
	@test e("z").data ≈ [0, 0, 1]
	@test_throws ErrorException e(nothing)
	@test_throws ErrorException e(missing)
	@test typeof(e("a").data) == OneHotMatrix{UInt32, UInt32(3), Vector{UInt32}}
	@test e(extractempty).data isa OneHotMatrix{UInt32, UInt32(3), Vector{UInt32}}
	@test nobs(e(extractempty)) == 0

	@test mapreduce(e, catobs, ["a", "b"]).data ≈ [1 0; 0 1; 0 0]
	@test_throws ErrorException e(["a", "b"]).data
	@test_throws ErrorException e(["a", missing])
	@test_throws ErrorException e(["a", missing, "x"])
	@test_throws ErrorException typeof(e(["a", "b"]).data)
	@test typeof(mapreduce(e, catobs, ["a", "b"]).data) == OneHotMatrix{UInt32, UInt32(3), Vector{UInt32}}
	@test_throws ErrorException e(["a", "b", nothing])

	@test isnothing(ExtractCategorical([], false))
	e2 = ExtractCategorical(JsonGrinder.Entry(Dict("a"=>1,"c"=>1), 2), false)
	@test e2.uniontypes == false
	@test e2("a").data ≈ [1, 0, 0]
	@test e2("c").data ≈ [0, 1, 0]
	@test e2("b").data ≈ [0, 0, 1]
	@test_throws ErrorException e2(nothing)
	@test_throws ErrorException e2(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]
	@test hcat(e("a").data, e("b").data) ≈ [1 0; 0 1; 0 0]
	@test_throws ErrorException e(Dict(1=>2))

	@test nobs(e("a")) == 1
	@test nobs(e("b")) == 1
	@test nobs(e("z")) == 1
	@test_throws ErrorException nobs(e(nothing))
	@test_throws ErrorException nobs(e(missing))
	@test_throws ErrorException nobs(e([missing, nothing]))
	@test_throws ErrorException nobs(e([missing, nothing, "a"]))

	e3 = ExtractCategorical(JsonGrinder.Entry(Dict(1=>1,2=>1), 2), false)
	@test e3.uniontypes == false
	@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_throws ErrorException isequal(e3([]).data, [missing missing missing]')

	e4 = ExtractCategorical(JsonGrinder.Entry(Dict(1.0=>1,2.0=>1), 2), false)
	@test e4.uniontypes == false
	@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_throws ErrorException isequal(e4([]).data, [missing missing missing]')

	@testset "type conversions" begin
		j1 = JSON.parse("""{"a": "4"}""")
		j2 = JSON.parse("""{"a": "11.5"}""")
		j3 = JSON.parse("""{"a": 7}""")
		j4 = JSON.parse("""{"a": 4.5}""")

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

		ext = suggestextractor(sch)

		@test ext(Dict("a"=>4)) == ext(Dict("a"=>4.0))
		@test ext(Dict("a"=>4)) == ext(Dict("a"=>4f0))
		@test ext(Dict("a"=>4)) == ext(Dict("a"=>"4"))
		@test ext(Dict("a"=>4)) == ext(Dict("a"=>"4.0"))
		# todo: add here metadata test
	end
end

@testset "ExtractString" begin
	e = ExtractString(true)
	ehello = e("Hello", store_input=false)
	ehellos = e("Hello", store_input=true)
	@test ehello.data.S == ["Hello"]
	@test ehello.data == ehellos.data
	@test e(Symbol("Hello")).data.S == ["Hello"]
	@test e(["Hello", "world"]).data.S ≃ [missing]
	@test mapreduce(e, catobs, ["Hello", "world"]).data.S == ["Hello", "world"]

	@test e(missing).data.S ≃ [missing]
	@test e(nothing).data.S ≃ [missing]
	@test isequal(e(Dict(1=>2)), e(missing))
	@test ehello isa ArrayNode{NGramMatrix{Union{Missing, String},Vector{Union{Missing, String}},Union{Missing, Int64}},Nothing}
	@test ehello.data isa NGramMatrix{Union{Missing, String},Vector{Union{Missing, String}},Union{Missing, Int64}}
	@test e(missing).data isa NGramMatrix{Union{Missing, String},Vector{Union{Missing, String}},Union{Missing, Int64}}
	@test e(nothing).data isa NGramMatrix{Union{Missing, String},Vector{Union{Missing, String}},Union{Missing, Int64}}
	@test ehellos.metadata == ["Hello"]
	@test nobs(e(extractempty)) == 0

	e = ExtractString(false)
	@test e("Hello").data.S == ["Hello"]
	@test e(Symbol("Hello")).data.S == ["Hello"]
	@test_throws ErrorException e(["Hello", "world"]).data.S
	@test mapreduce(e, catobs, ["Hello", "world"]).data.S == ["Hello", "world"]
	@test_throws ErrorException e(missing)
	@test_throws ErrorException e(nothing)
	@test_throws ErrorException e(Dict(1=>2))
	@test e("Hello") isa ArrayNode{NGramMatrix{String,Vector{String},Int64},Nothing}
	@test e("Hello").data isa NGramMatrix{String,Vector{String},Int64}
end

@testset "ExtractKeyAsField" 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))
	@test ext isa JsonGrinder.ExtractKeyAsField

	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

	with_emptyismissing(true) do
		b = ext(nothing)
		@test nobs(b) == 1
		@test ismissing(b.data)
		b = ext(Dict())
		@test nobs(b) == 1
		@test ismissing(b.data)
	end
	with_emptyismissing(false) do
		b = ext(nothing)
		@test nobs(b) == 1
		@test nobs(b.data) == 0
		@test b.data[:key].data isa NGramMatrix{Union{Missing, String},Vector{Union{Missing, String}},Union{Missing, Int64}}
		b = ext(Dict())
		@test nobs(b) == 1
		@test nobs(b.data) == 0
		@test b.data[:key].data isa NGramMatrix{Union{Missing, String},Vector{Union{Missing, String}},Union{Missing, Int64}}
	end

	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 NGramMatrix{Union{Missing, String},Vector{Union{Missing, String}},Union{Missing, 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
	@test isnothing(b.data[:key].metadata)
	@test isnothing(b.data[:item][:a].metadata)
	@test isnothing(b.data[:item][:b].metadata)

	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(js[1], store_input=true)
	k = only(keys(js[1]))
	@test b.data[:key].metadata == [k]
	@test b.data[:item][:a].metadata == fill(js[1][k]["a"],1,1)
	@test b.data[:item][:b].metadata == [js[1][k]["b"]]

	b = ext(Dict(), store_input=true)
	@test b.metadata == [Dict()]
	@test isnothing(b.data[:key].metadata)
	@test isnothing(b.data[:item].metadata)

	b = ext(nothing, store_input=true)
	@test b.metadata == [nothing]
	@test isnothing(b.data.data.key.metadata)
	@test isnothing(b.data.data.item.metadata)

	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 "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 ExtractCategorical
	@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, "g": "1"}""")
	j2 = JSON.parse("""{"a": "2", "b": "b", "c": "2",   "d": 2, "e": "1.3", "f": 2, "g": "2"}""")
	j3 = JSON.parse("""{"a": "3", "b": "c", "c": "2.3", "d": 2.3, "e": "1.4", "f": 3, "g": "3"}""")
	j4 = JSON.parse("""{"a": "4", "b": "c", "c": "5",   "d": 5, "e": "1.4", "f": 3, "g": "4"}""")
	j5_200 = [JSON.parse("""{"a": "5", "b": "$("c"^i)", "c": "5", "d": $(5+i/10), "e": "1.4", "f": 3, "g": "$i"}""") for i in 5:200]
	sch = JsonGrinder.schema([j1, j2, j3, j4, j5_200...])
	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 ExtractCategorical
	@test ext[:g] 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) <: Bool
	@test eltype(ext_j1[:g].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) <: Bool
	@test eltype(ext_j2[:g].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) <: Bool
	@test eltype(ext_j3[:g].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) <: Bool
	@test eltype(ext_j4[:g].data) <: Float32

	@test ext_j1["U"].data ≈ [0]
	@test ext_j2["U"].data ≈ [(2-1.1)/23.9]
	@test ext_j3["U"].data ≈ [(2.3-1.1)/23.9]
	@test ext_j4["U"].data ≈ [(5-1.1)/23.9]

	ext_j1s = ext(j1, store_input=true)
	ext_j2s = ext(j2, store_input=true)
	ext_j3s = ext(j3, store_input=true)
	ext_j4s = ext(j4, store_input=true)

	@test ext_j1["U"].data == ext_j1s["U"].data
	@test ext_j2["U"].data == ext_j2s["U"].data
	@test ext_j3["U"].data == ext_j3s["U"].data
	@test ext_j4["U"].data == ext_j4s["U"].data

	@test Mill.metadata.(values(ext_j1s.data)) == ([1], ["a"], ["1"], fill(1.1,1,1), ["1.2"], ["1.1"], fill("1",1,1))
	@test Mill.metadata.(values(ext_j2s.data)) == ([2], ["b"], ["2"], fill(2,1,1), ["1.3"], ["2"], fill("2",1,1))
	@test Mill.metadata.(values(ext_j3s.data)) == ([3], ["c"], ["3"], fill(2.3,1,1), ["1.4"], ["2.3"], fill("3",1,1))
	@test Mill.metadata.(values(ext_j4s.data)) == ([3], ["c"], ["4"], fill(5,1,1), ["1.4"], ["5"], fill("4",1,1))

	m = reflectinmodel(sch, ext)
	@test buf_printtree(m) == """
	ProductModel ↦ ArrayModel(Dense(52, 10)) \t# 2 arrays, 530 params, 2.148 KiB
	  ├── f: ArrayModel(Dense(4, 10)) \t# 2 arrays, 50 params, 280 bytes
	  ├── b: ArrayModel(Dense(2053, 10)) \t# 2 arrays, 20_540 params, 80.312 KiB
	  ├── a: ArrayModel(Dense(6, 10)) \t# 2 arrays, 70 params, 360 bytes
	  ├── d: ArrayModel(identity)
	  ├── e: ArrayModel(Dense(4, 10)) \t# 2 arrays, 50 params, 280 bytes
	  ├── c: ArrayModel(Dense(5, 10)) \t# 2 arrays, 60 params, 320 bytes
	  └── g: 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, testing_settings)

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

	@test buf_printtree(ext, trav=true) ==
	"""
	Dict [""]
	  ├── a: Float32 ["E"]
	  ├── b: FeatureVector with 3 items ["U"]
	  └── c: Array of ["k"]
	           └── Categorical d = 6 ["s"]
	"""

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

	ext_j1s = ext(j1, store_input=true)
	ext_j2s = ext(j2, store_input=true)
	ext_j3s = ext(j3, store_input=true)
	ext_j4s = ext(j4, store_input=true)

	@test ext_j1["E"].data ≈ [0]
	@test ext_j2["E"].data ≈ [1/3]
	@test ext_j3["E"].data ≈ [2/3]
	@test ext_j4["E"].data ≈ [1]

	@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 ≈ [
	1  0  0
 	0  1  0
 	0  0  1
 	0  0  0
 	0  0  0
 	0  0  0]
	@test ext_j2["s"].data ≈ [
	1  0  0  0
 	0  1  0  0
 	0  0  1  0
 	0  0  0  1
 	0  0  0  0
 	0  0  0  0
	]
	@test ext_j3["s"].data ≈ [
	1  0  0  0  0
	0  1  0  0  0
	0  0  1  0  0
	0  0  0  1  0
	0  0  0  0  1
	0  0  0  0  0
	]
	@test ext_j4["s"].data ≈ [
	1  0  0
 	0  1  0
 	0  0  1
 	0  0  0
 	0  0  0
 	0  0  0
	]

	for (a,b) in zip([ext_j1, ext_j2, ext_j3, ext_j4], [ext_j1s, ext_j2s, ext_j3s, ext_j4s])
		@test a["E"].data == b["E"].data
		@test a["U"].data == b["U"].data
		@test a["s"].data == b["s"].data
	end

	@test ext_j1s["U"].metadata == [[1, 2, 3]]
	@test ext_j2s["U"].metadata == [[2, 2, 3]]
	@test ext_j3s["U"].metadata == [[3, 2, 3]]
	@test ext_j4s["U"].metadata == [[2, 3, 4]]

	@test ext_j1s["s"].metadata == [1, 2, 3]
	@test ext_j2s["s"].metadata == [1, 2, 3, 4]
	@test ext_j3s["s"].metadata == [1, 2, 3, 4, 5]
	@test ext_j4s["s"].metadata == [1, 2, 3]
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, 4, 5]}""")

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

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

	@test ext_j1[:a].data.data isa Matrix{Float32}
	@test ext_j2[:a].data.data isa Matrix{Float32}
	@test ext_j3[:a].data.data isa Matrix{Float32}
	@test ext_j4[:a].data.data isa Matrix{Float32}
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, testing_settings)
	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
    # todo: add test for extraction of numeric scalars from string and also missings, if it's all matrix and catobsing properly
	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, testing_settings)
	@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 "merging numeric scalars" begin
	j1 = JSON.parse("""{"a": "1", "b":1}""")
	j2 = JSON.parse("""{"a": 4, "b":2}""")
	j3 = JSON.parse("""{"a": "3.1", "b":3}""")
	j4 = JSON.parse("""{"b":4}""")

	sch = JsonGrinder.schema([j1, j2, j3, j4])
	ext = suggestextractor(sch, testing_settings)
	@test nchildren(ext[:a]) == 1

	e1 = ext(j1)
	e2 = ext(j2)
	e3 = ext(j3)
	e4 = ext(j4)
	@test e1["M"].data ≈ [0]
	@test e2["M"].data ≈ [1]
	@test e3["M"].data ≈ [0.7]
	@test e4["M"].data ≃ fill(missing,1,1)
	@test e1["U"].data ≈ [0]
	@test e2["U"].data ≈ [1/3]
	@test e3["U"].data ≈ [2/3]
	@test e4["U"].data ≈ [1]
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, testing_settings)

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

	@test buf_printtree(ext, trav=true) ==
    """
	Dict [""]
	  └── a: MultiRepresentation ["U"]
	           ├── e1: FeatureVector with 5 items ["c"]
	           ├── e2: Dict ["k"]
	           │         └── Sylvanas is the worst warchief ever: String ["o"]
	           └── e3: Float32 ["s"]
	"""

	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 \t# 1 obs, 48 bytes
	  └── a: ProductNode \t# 1 obs, 48 bytes
	           ├── e1: ArrayNode(5×1 Array with Union{Missing, Float32} elements) \t# 1 obs, 73 bytes
	           ├── e2: ProductNode \t# 1 obs, 32 bytes
	           │         └── Sylvanas is the worst warchief ever: ArrayNode(2053×1 NGramMatrix with Union{Missing, Int64} elements) \t# 1 obs, 112 bytes
	           └── e3: ArrayNode(1×1 Array with Union{Missing, Float32} elements) \t# 1 obs, 53 bytes
	"""
	@test ext[:a][1] == ext["c"]
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, testing_settings)

	@test buf_printtree(sch) ==
	"""
	[Dict] \t# updated = 5
	  └── a: [MultiEntry] \t# updated = 5
	           ├── 1: [Scalar - String], 2 unique values \t# updated = 2
	           ├── 2: [Scalar - Float64,Int64], 2 unique values \t# updated = 2
	           └── 3: [List] \t# updated = 1
	                    └── [Scalar - Int64], 5 unique values \t# 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)

	@test ext(f, store_input=true) != ext(nothing, store_input=true)
	@test ext(f, store_input=true).data ≃ ext(nothing, store_input=true).data
	@test ext(a, store_input=true).metadata == ["a"]
	@test ext(b, store_input=true).metadata == ["b"]
	@test ext(c, store_input=true).metadata == [""]
	@test ext(d, store_input=true).metadata == ["d"]
	@test ext(e, store_input=true).metadata == [""]
	@test ext(f, store_input=true).metadata == [5.2]
	@test ext(nothing, store_input=true).metadata == [nothing]
	#todo: add tests for all extractors for store_input=true
end

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

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

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

@testset "Not 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, 4, 5]}""")
	j5 = JSON.parse("""{"a": [1, 2, 3, 6, 7]}""")

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

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

	ext_j2 = ext(j2)
	@test buf_printtree(ext_j2) ==
    """
	ProductNode \t# 1 obs, 16 bytes
	  └── a: BagNode \t# 1 obs, 80 bytes
	           └── ArrayNode(1×2 Array with Float32 elements) \t# 2 obs, 56 bytes
	"""

	m = reflectinmodel(sch, ext)
	@test buf_printtree(m) ==
    """
	ProductModel ↦ ArrayModel(identity)
	  └── a: BagModel ↦ BagCount([SegmentedMean(1); SegmentedMax(1)]) ↦ ArrayModel(Dense(3, 10)) \t# 4 arrays, 42 params, 328 bytes
	           └── ArrayModel(identity)
	"""
end

@testset "default_scalar_extractor" begin

end

@testset "type stability of extraction missings" begin
	# todo: make tests for categorical, scalar, string, vector
	j1 = JSON.parse("""{"a": 1}""")
	j2 = JSON.parse("""{"b": "a"}""")
	j3 = JSON.parse("""{"a": 3, "b": "b"}""")
	j4 = JSON.parse("""{"a": 4, "b": "c"}""")
	j5 = JSON.parse("""{"a": 5, "b": "d"}""")

	sch = schema([j1,j2,j3,j4,j5])
	ext = suggestextractor(sch, testing_settings)
	m = reflectinmodel(sch, ext)

	e1 = ext(j1)
	e2 = ext(j2)
	e3 = ext(j3)
	e12 = catobs(e1, e2)
	e23 = catobs(e2, e3)
	e13 = catobs(e1, e3)
	typeof(e1)
	typeof(e2)
	typeof(e3)
	typeof(e12)
	typeof(e23)
	typeof(e13)
	@test typeof(e1) == typeof(e2)
	@test typeof(e1) == typeof(e3)
	@test typeof(e1) == typeof(e12)
	@test typeof(e1) == typeof(e23)
	@test typeof(e1) == typeof(e13)
end


export AbstractBag
"""
    AbstractBag{BT, N}<:TagBlock{BT, N}

Abstract `Bag` is a wrapper of a block that conserves all properties.
Including `mat`, `apply!`, `ishermitian`, `isreflexive`, `isunitary`,
`occupied_locs`, `apply_back!` and `mat_back!`.
"""
abstract type AbstractBag{BT, N}<:TagBlock{BT, N} end

Yao.mat(::Type{T}, bag::AbstractBag{N}) where {T,N} = mat(T, content(bag))
Yao.apply!(reg::AbstractRegister, bag::AbstractBag) = apply!(reg, content(bag))
Yao.ishermitian(bag::AbstractBag) = ishermitian(content(bag))
Yao.isreflexive(bag::AbstractBag) = isreflexive(content(bag))
Yao.isunitary(bag::AbstractBag) = isunitary(content(bag))
Yao.occupied_locs(bag::AbstractBag) = occupied_locs(content(bag))

function Yao.AD.apply_back!(state, b::AbstractBag, collector)
    Yao.AD.apply_back!(state, content(b), collector)
end
function Yao.AD.mat_back!(::Type{T}, b::AbstractBag, adjy, collector) where T
    Yao.AD.mat_back!(T, content(b), adjy, collector)
end

export Bag, enable_block!, disable_block!, setcontent!, isenabled
"""
    Bag{N}<:TagBlock{AbstractBlock, N}

A bag is a trivil container, but can
    * `setcontent!(bag, content)`
    * `disable_block!(bag)`
    * `enable_block!(bag)`
"""
mutable struct Bag{N}<:AbstractBag{AbstractBlock, N}
    content::AbstractBlock{N}
    mask::Bool
end
Bag(b::AbstractBlock) = Bag(b,true)

Yao.content(bag::Bag{N}) where N = bag.mask ? bag.content : put(N, 1=>I2)
Yao.chcontent(bag::Bag, content) = Bag(content)
setcontent!(bag::Bag, content) = (bag.content = content; bag)
disable_block!(b::Bag) = (b.mask = false; b)
enable_block!(b::Bag) = (b.mask = true; b)
isenabled(b::Bag) = b.mask

function Yao.print_annotation(io::IO, bag::Bag)
    printstyled(io, isenabled(bag) ? "[⊙] " : "[⊗] "; bold=true, color=isenabled(bag) ? :green : :red)
end

function Base.show(io::IO, ::MIME"plain/text", blk::Bag)
    return print_tree(io, blk; title=false, compact=false)
end


function layoutable(::Type{Slider}, fig_or_scene; bbox = nothing, kwargs...)

    topscene = get_topscene(fig_or_scene)

    default_attrs = default_attributes(Slider, topscene).attributes
    theme_attrs = subtheme(topscene, :Slider)
    attrs = merge!(merge!(Attributes(kwargs), theme_attrs), default_attrs)

    decorations = Dict{Symbol, Any}()

    @extract attrs (
        halign, valign, horizontal, linewidth, snap,
        startvalue, value, color_active, color_active_dimmed, color_inactive
    )

    sliderrange = attrs.range

    protrusions = Node(GridLayoutBase.RectSides{Float32}(0, 0, 0, 0))
    layoutobservables = LayoutObservables{Slider}(attrs.width, attrs.height, attrs.tellwidth, attrs.tellheight,
        halign, valign, attrs.alignmode; suggestedbbox = bbox, protrusions = protrusions)

    onany(linewidth, horizontal) do lw, horizontal
        if horizontal
            layoutobservables.autosize[] = (nothing, Float32(lw))
        else
            layoutobservables.autosize[] = (Float32(lw), nothing)
        end
    end

    sliderbox = lift(identity, layoutobservables.computedbbox)

    endpoints = lift(sliderbox, horizontal) do bb, horizontal

        h = height(bb)
        w = width(bb)

        if horizontal
            y = bottom(bb) + h / 2
            [Point2f0(left(bb) + h/2, y),
             Point2f0(right(bb) - h/2, y)]
        else
            x = left(bb) + w / 2
            [Point2f0(x, bottom(bb) + w/2),
             Point2f0(x, top(bb) - w/2)]
        end
    end

    # this is the index of the selected value in the slider's range
    # selected_index = Node(1)
    # add the selected index to the attributes so it can be manipulated later
    attrs.selected_index = 1
    selected_index = attrs.selected_index

    # the fraction on the slider corresponding to the selected_index
    # this is only used after dragging
    sliderfraction = lift(selected_index, sliderrange) do i, r
        (i - 1) / (length(r) - 1)
    end

    dragging = Node(false)

    # what the slider actually displays currently (also during dragging when
    # the slider position is in an "invalid" position given the slider's range)
    displayed_sliderfraction = Node(0.0)

    on(sliderfraction) do frac
        # only update displayed fraction through sliderfraction if not dragging
        # dragging overrides the value so there is clear mouse interaction
        if !dragging[]
            displayed_sliderfraction[] = frac
        end
    end

    # when the range is changed, switch to closest value
    on(sliderrange) do rng
        selected_index[] = closest_index(rng, value[])
    end

    on(selected_index) do i
        value[] = sliderrange[][i]
    end

    # initialize slider value with closest from range
    selected_index[] = closest_index(sliderrange[], startvalue[])

    middlepoint = lift(endpoints, displayed_sliderfraction) do ep, sf
        Point2f0(ep[1] .+ sf .* (ep[2] .- ep[1]))
    end

    linepoints = lift(endpoints, middlepoint) do eps, middle
        [eps[1], middle, middle, eps[2]]
    end

    linecolors = lift(color_active_dimmed, color_inactive) do ca, ci
        [ca, ci]
    end

    endbuttons = scatter!(topscene, endpoints, color = linecolors,
        markersize = linewidth, strokewidth = 0, raw = true, inspectable = false)
    decorations[:endbuttons] = endbuttons

    linesegs = linesegments!(topscene, linepoints, color = linecolors,
        linewidth = linewidth, raw = true, inspectable = false)
    decorations[:linesegments] = linesegs

    button_magnification = Node(1.0)
    buttonsize = @lift($linewidth * $button_magnification)
    button = scatter!(topscene, middlepoint, color = color_active, strokewidth = 0,
        markersize = buttonsize, raw = true, inspectable = false)
    decorations[:button] = button

    mouseevents = addmouseevents!(topscene, layoutobservables.computedbbox)

    onmouseleftdrag(mouseevents) do event
        dragging[] = true
        dif = event.px - event.prev_px
        fraction = clamp(if horizontal[]
            (event.px[1] - endpoints[][1][1]) / (endpoints[][2][1] - endpoints[][1][1])
        else
            (event.px[2] - endpoints[][1][2]) / (endpoints[][2][2] - endpoints[][1][2])
        end, 0, 1)

        newindex = closest_fractionindex(sliderrange[], fraction)
        if snap[]
            fraction = (newindex - 1) / (length(sliderrange[]) - 1)
        end
        displayed_sliderfraction[] = fraction

        if selected_index[] != newindex
            selected_index[] = newindex
        end

        return true
    end

    onmouseleftdragstop(mouseevents) do event
        dragging[] = false
        # adjust slider to closest legal value
        sliderfraction[] = sliderfraction[]
        linecolors[] = [color_active_dimmed[], color_inactive[]]
        return true
    end

    onmouseleftdown(mouseevents) do event
        pos = event.px
        dim = horizontal[] ? 1 : 2
        frac = (pos[dim] - endpoints[][1][dim]) / (endpoints[][2][dim] - endpoints[][1][dim])
        selected_index[] = closest_fractionindex(sliderrange[], frac)
        # linecolors[] = [color_active[], color_inactive[]]
        return true
    end

    onmouseleftdoubleclick(mouseevents) do event
        selected_index[] = closest_index(sliderrange[], startvalue[])
        return true
    end

    onmouseenter(mouseevents) do event
        button_magnification[] = 1.25
        return false
    end

    onmouseout(mouseevents) do event
        button_magnification[] = 1.0
        linecolors[] = [color_active_dimmed[], color_inactive[]]
        return false
    end

    # trigger autosize through linewidth for first layout
    linewidth[] = linewidth[]

    Slider(fig_or_scene, layoutobservables, attrs, decorations)
end

function valueindex(sliderrange, value)
    for (i, val) in enumerate(sliderrange)
        if val == value
            return i
        end
    end
    nothing
end

function closest_fractionindex(sliderrange, fraction)
    n = length(sliderrange)
    onestepfrac = 1 / (n - 1)
    i = round(Int, fraction / onestepfrac) + 1
    min(max(i, 1), n)
end

function closest_index(sliderrange, value)
    for (i, val) in enumerate(sliderrange)
        if val == value
            return i
        end
    end
    # if the value wasn't found this way try inexact
    closest_index_inexact(sliderrange, value)
end

function closest_index_inexact(sliderrange, value)
    distance = Inf
    selected_i = 0
    for (i, val) in enumerate(sliderrange)
        newdist = abs(val - value)
        if newdist < distance
            distance = newdist
            selected_i = i
        end
    end
    selected_i
end

"""
Set the `slider` to the value in the slider's range that is closest to `value` and return this value.
"""
function set_close_to!(slider::Slider, value)
    closest = closest_index(slider.range[], value)
    slider.selected_index = closest
    slider.range[][closest]
end


@enum BracketState begin
    ISOTOPE = Int8(0)
    SYMBOL = Int8(1)
    MULTIPLICITY = Int8(2)
    CHARGE = Int8(3)
    FINISH = Int8(4)
end

struct BracketParseException <: Exception
    msg::String
    bracketcontents::String
    charloc::Int
end
Base.showerror(io::IO, e::BracketParseException) = print(io, "Failed to parse bracket(\"[$(e.bracketcontents)]\") on character ", e.charloc, "!\n", e.msg)

function ParseBracket(s::AbstractString)
    len = length(s)

    state = ISOTOPE
    symbol, isotope, aromatic, ringID, hydrogens, charge = nothing, nothing, false, [], 0, 0

    idx = 1
    while idx <= len
        cursor = s[idx]
        #Handle isotope, explicit hydrogen multiplicity
        if isdigit(cursor)
            if state == ISOTOPE
                isotope, idx = tryparseint16(s, idx)
                state = SYMBOL
            elseif state == MULTIPLICITY
                hydrogens, idx = tryparseint16(s, idx)
                state = CHARGE
            else
               throw(BracketParseException("Invalid SMILES. Charge must be denoted with +/-.", s, idx))
            end
        #Handle Element symbol parsing
        elseif isletter(cursor)
            if state == SYMBOL || state == ISOTOPE
                if isa(symbol, Nothing)
                    symbollist = isuppercase(cursor) ? bracket : bracket_aromatic
                    aromatic = islowercase(cursor)
                    symbol, idx = tryparseelement(s, idx, symbollist)
                    if isa(symbol, Nothing)
                        throw(BracketParseException("Invalid SMILES. Unrecognized chemical symbol.", s, idx))
                    end
                elseif cursor == 'H'
                    state = MULTIPLICITY
                    hydrogens = 1
                    idx = nextind(s, idx)
                else
                    throw(BracketParseException("Invalid SMILES. Bracket should contain one element (with or without Hydrogen).", s, idx))
                end
            else
                throw(BracketParseException("Invalid SMILES. Bracket should contain one element.", s, idx))
            end
        #Handle charge
        elseif ispm(cursor)
            if state == CHARGE || state == MULTIPLICITY || state == SYMBOL || state == ISOTOPE
                charge, idx = tryparsecharge(s, idx)
                state = SYMBOL
            else
                throw(BracketParseException("Invalid SMILES. Charge should follow element symbol.", s, idx))
            end
        #Handle operators
        elseif isspecialoperator(cursor)
            throw(BracketParseException("Invalid SMILES. Special operators (" * join(specialoperators, ",") * ") should not be contained in a bracket.", s, idx))
        elseif isbondoperator(cursor)
            if cursor == '@'
                @warn("Chirality is not yet supported, ignoring @ operator", maxlog=1)
                _, idx = tryparsechirality(s, idx)
            else
                throw(BracketParseException("Invalid SMILES. Bond operators(" * join(setdiff(bondoperators, '@'), ",") * ") should not be contained in a bracket.", s, idx))
            end
        end
    end
    return Element( symbol, isotope, aromatic, ringID, hydrogens, 0, charge )
end


# function to overwrite * and mul!

## with Base.threads
"""
    ThreadedSparseCSR.multithread_matmul()
    ThreadedSparseCSR.multithread_matmul(PolyesterThreads())

Overwrites `*` and `mul!` by their multithreaded versions, using Polyester.jl threading.

    ThreadedSparseCSR.multithread_matmul(BaseThreads())

Overwrites `*` and `mul!` by their multithreaded versions, using Threads.@spawn threading.
"""
function multithread_matmul(T::BaseThreads)

    @eval function  mul!(y::AbstractVector, A::SparseMatrixCSR, x::AbstractVector, alpha::Number, beta::Number)
        return tmul!(y, A, x, alpha, beta)
    end

    @eval function  mul!(y::AbstractVector, A::SparseMatrixCSR, x::AbstractVector)
        return tmul!(y, A, x)
    end

    @eval function  *(A::SparseMatrixCSR, x::AbstractVector)
        return tmul(A, x)
    end

end

## with Polyester
function multithread_matmul(T::PolyesterThreads)

    @eval function  mul!(y::AbstractVector, A::SparseMatrixCSR, x::AbstractVector, alpha::Number, beta::Number)
        return bmul!(y, A, x, alpha, beta)
    end

    @eval function  mul!(y::AbstractVector, A::SparseMatrixCSR, x::AbstractVector)
        return bmul!(y, A, x)
    end

    @eval function  *(A::SparseMatrixCSR, x::AbstractVector)
        return bmul(A, x)
    end

end

## with default threading backend
function  multithread_matmul()

    multithread_matmul(DefaultThreadingBackend())

end


"""
    module Vision

Data blocks, encodings and more for computer vision.

The most important [`Block`] is [`Image`](#)`{N}`.

Blocks:

- [`Image`](#)`{N}`: an `N`-dimensional color image
- [`Mask`](#)`{N}`: an `N`-dimensional categorical mask
- [`Keypoints`](#)`{N}`: a fixed number of `N`-dimensional keypoints

Encodings:

- [`OneHot`](#) is implemented for `Mask`s
- [`ImagePreprocessing`](#) prepares an `Image` for training by
    normalizing and expanding the color channels
- [`KeypointPreprocessing`](#) prepares `Keypoints` for training by
    normalizing them.

"""
module Vision

using ..FastAI
using ..FastAI:
    # blocks
    Block, WrapperBlock, AbstractBlock, OneHotTensor, OneHotTensorMulti, Label,
    LabelMulti, wrapped,
    # encodings
    Encoding, StatefulEncoding, OneHot,
    # visualization
    ShowText,
    # other
    FASTAI_METHOD_REGISTRY, registerlearningmethod!, Datasets
import FastAI.Datasets

# for tests
using ..FastAI: testencoding

# extending
import ..FastAI:
    blockmodel, blockbackbone, blocklossfn, encode, decode, checkblock,
    encodedblock, decodedblock, showblock!, mockblock, setup, encodestate,
    decodestate

import Colors: colormaps_sequential, Colorant, Color, Gray, Normed, RGB,
    alphacolor, deuteranopic, distinguishable_colors
import FixedPointNumbers: N0f8
import DataAugmentation
import DataAugmentation: apply, Identity, ToEltype, ImageToTensor, Normalize,
    BufferedThreadsafe, ScaleKeepAspect, PinOrigin, RandomCrop, CenterResizeCrop,
    AdjustBrightness, AdjustContrast, Maybe,
    ResizePadDivisible, itemdata
import ImageInTerminal
import IndirectArrays: IndirectArray
import Requires: @require
import StaticArrays: SVector
import Statistics: mean, std
import UnicodePlots

using InlineTest
using ShowCases

# Blocks
include("blocks/bounded.jl")
include("blocks/image.jl")
include("blocks/mask.jl")
include("blocks/keypoints.jl")

include("encodings/onehot.jl")
include("encodings/imagepreprocessing.jl")
include("encodings/keypointpreprocessing.jl")
include("encodings/projective.jl")

include("models/Models.jl")
include("models.jl")
include("learningmethods/utils.jl")
include("learningmethods/classification.jl")
include("learningmethods/segmentation.jl")
include("learningmethods/keypointregression.jl")
include("recipes.jl")

include("tests.jl")

function __init__()
    _registerrecipes()
    @require Makie="ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a" begin
        import .Makie
        import .Makie: @recipe, @lift
        import .FastAI: ShowMakie
        include("makie.jl")
    end
end

export Image, Mask, Keypoints, Bounded,
    # encodings
    ImagePreprocessing, KeypointPreprocessing, ProjectiveTransforms,
    # learning methods
    ImageClassificationSingle, ImageClassificationMulti,
    ImageKeypointRegression, ImageSegmentation

end


# *********************************************************************************
# REopt, Copyright (c) 2019-2020, Alliance for Sustainable Energy, LLC.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without modification,
# are permitted provided that the following conditions are met:
#
# Redistributions of source code must retain the above copyright notice, this list
# of conditions and the following disclaimer.
#
# Redistributions in binary form must reproduce the above copyright notice, this
# list of conditions and the following disclaimer in the documentation and/or other
# materials provided with the distribution.
#
# Neither the name of the copyright holder nor the names of its contributors may be
# used to endorse or promote products derived from this software without specific
# prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
# IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
# INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
# OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
# OF THE POSSIBILITY OF SUCH DAMAGE.
# *********************************************************************************
using Xpress

@testset "Thermal loads" begin
    m = Model(optimizer_with_attributes(Xpress.Optimizer, "OUTPUTLOG" => 0))
    results = run_reopt(m, "./scenarios/thermal_load.json")

    @test round(results["ExistingBoiler"]["year_one_boiler_fuel_consumption_mmbtu"], digits=0) ≈ 2905

    data = JSON.parsefile("./scenarios/thermal_load.json")
    data["DomesticHotWaterLoad"]["fuel_loads_mmbtu_per_hour"] = repeat([0.5], 8760)
    data["SpaceHeatingLoad"]["fuel_loads_mmbtu_per_hour"] = repeat([0.5], 8760)
    s = Scenario(data)
    inputs = REoptInputs(s)
    m = Model(optimizer_with_attributes(Xpress.Optimizer, "OUTPUTLOG" => 0))
    results = run_reopt(m, inputs)

    @test round(results["ExistingBoiler"]["year_one_boiler_fuel_consumption_mmbtu"], digits=0) ≈ 8760
end

@testset "CHP Sizing" begin
    # Sizing CHP with non-constant efficiency, no cost curve, no unavailability_periods
    data_sizing = JSON.parsefile("./scenarios/chp_sizing.json")
    s = Scenario(data_sizing)
    inputs = REoptInputs(s)
    m = Model(optimizer_with_attributes(Xpress.Optimizer, "MIPRELSTOP" => 0.01, "OUTPUTLOG" => 0))
    results = run_reopt(m, inputs)

    @test round(results["CHP"]["size_kw"], digits=0) ≈ 468.7 atol=1.0
    @test round(results["Financial"]["lcc"], digits=0) ≈ 1.3476e7 atol=1.0e7
end

@testset "CHP Cost Curve and Min Allowable Size" begin
    # Fixed size CHP with cost curve, no unavailability_periods
    data_cost_curve = JSON.parsefile("./scenarios/chp_sizing.json")
    data_cost_curve["CHP"] = Dict()
    data_cost_curve["CHP"]["prime_mover"] = "recip_engine"
    data_cost_curve["CHP"]["size_class"] = 2
    data_cost_curve["CHP"]["fuel_cost_per_mmbtu"] = 8.0
    data_cost_curve["CHP"]["min_kw"] = 0
    data_cost_curve["CHP"]["min_allowable_kw"] = 555.5
    data_cost_curve["CHP"]["max_kw"] = 1000
    data_cost_curve["CHP"]["installed_cost_per_kw"] = 1800.0
    data_cost_curve["CHP"]["installed_cost_per_kw"] = [2300.0, 1800.0, 1500.0]
    data_cost_curve["CHP"]["tech_sizes_for_cost_curve"] = [100.0, 300.0, 1140.0]

    data_cost_curve["CHP"]["federal_itc_pct"] = 0.1
    data_cost_curve["CHP"]["macrs_option_years"] = 0
    data_cost_curve["CHP"]["macrs_bonus_pct"] = 0.0
    data_cost_curve["CHP"]["macrs_itc_reduction"] = 0.0

    expected_x = data_cost_curve["CHP"]["min_allowable_kw"]
    cap_cost_y = data_cost_curve["CHP"]["installed_cost_per_kw"]
    cap_cost_x = data_cost_curve["CHP"]["tech_sizes_for_cost_curve"]
    slope = (cap_cost_x[3] * cap_cost_y[3] - cap_cost_x[2] * cap_cost_y[2]) / (cap_cost_x[3] - cap_cost_x[2])
    init_capex_chp_expected = cap_cost_x[2] * cap_cost_y[2] + (expected_x - cap_cost_x[2]) * slope
    lifecycle_capex_chp_expected = init_capex_chp_expected -
        REopt.npv(data_cost_curve["Financial"]["offtaker_discount_pct"],
        [0, init_capex_chp_expected * data_cost_curve["CHP"]["federal_itc_pct"]])

    #PV
    data_cost_curve["PV"]["min_kw"] = 1500
    data_cost_curve["PV"]["max_kw"] = 1500
    data_cost_curve["PV"]["installed_cost_per_kw"] = 1600
    data_cost_curve["PV"]["federal_itc_pct"] = 0.26
    data_cost_curve["PV"]["macrs_option_years"] = 0
    data_cost_curve["PV"]["macrs_bonus_pct"] = 0.0
    data_cost_curve["PV"]["macrs_itc_reduction"] = 0.0

    init_capex_pv_expected = data_cost_curve["PV"]["max_kw"] * data_cost_curve["PV"]["installed_cost_per_kw"]
    lifecycle_capex_pv_expected = init_capex_pv_expected -
        REopt.npv(data_cost_curve["Financial"]["offtaker_discount_pct"],
        [0, init_capex_pv_expected * data_cost_curve["PV"]["federal_itc_pct"]])

    s = Scenario(data_cost_curve)
    inputs = REoptInputs(s)
    m = Model(optimizer_with_attributes(Xpress.Optimizer, "MIPRELSTOP" => 0.01, "OUTPUTLOG" => 0))
    results = run_reopt(m, inputs)

    init_capex_total_expected = init_capex_chp_expected + init_capex_pv_expected
    lifecycle_capex_total_expected = lifecycle_capex_chp_expected + lifecycle_capex_pv_expected

    init_capex_total = results["Financial"]["initial_capital_costs"]
    lifecycle_capex_total = results["Financial"]["initial_capital_costs_after_incentives"]

    # Check initial CapEx (pre-incentive/tax) and life cycle CapEx (post-incentive/tax) cost with expect
    @test init_capex_total_expected ≈ init_capex_total atol=0.0001*init_capex_total_expected
    @test lifecycle_capex_total_expected ≈ lifecycle_capex_total atol=0.0001*lifecycle_capex_total_expected

    # Test CHP.min_allowable_kw - the size would otherwise be ~100 kW less by setting min_allowable_kw to zero
    @test results["CHP"]["size_kw"] ≈ data_cost_curve["CHP"]["min_allowable_kw"] atol=0.1
end

@testset "CHP Unavailability and Outage" begin
    """
    Validation to ensure that:
        1) CHP meets load during outage without exporting
        2) CHP never exports if chp.can_wholesale and chp.can_net_meter inputs are False (default)
        3) CHP does not "curtail", i.e. send power to a load bank when chp.can_curtail is False (default)
        4) CHP min_turn_down_pct is ignored during an outage
        5) **Not until cooling is added:** Cooling load gets zeroed out during the outage period
        6) Unavailability intervals that intersect with grid-outages get ignored
        7) Unavailability intervals that do not intersect with grid-outages result in no CHP production
    """
    # Sizing CHP with non-constant efficiency, no cost curve, no unavailability_periods
    data = JSON.parsefile("./scenarios/chp_unavailability_outage.json")

    # Add unavailability periods that 1) intersect (ignored) and 2) don't intersect with outage period
    data["CHP"]["unavailability_periods"] = [Dict([("month", 1), ("start_week_of_month", 2),
            ("start_day_of_week", 1), ("start_hour", 1), ("duration_hours", 8)]),
            Dict([("month", 1), ("start_week_of_month", 2),
            ("start_day_of_week", 3), ("start_hour", 9), ("duration_hours", 8)])]

    # Manually doing the math from the unavailability defined above
    unavail_1_start = 24 + 1
    unavail_1_end = unavail_1_start + 8 - 1
    unavail_2_start = 24*3 + 9
    unavail_2_end = unavail_2_start + 8 - 1

    # Specify the CHP.min_turn_down_pct which is NOT used during an outage
    data["CHP"]["min_turn_down_pct"] = 0.5
    # Specify outage period; outage timesteps are 1-indexed
    outage_start = unavail_1_start
    data["ElectricUtility"]["outage_start_time_step"] = outage_start
    outage_end = unavail_1_end
    data["ElectricUtility"]["outage_end_time_step"] = outage_end
    data["ElectricLoad"]["critical_load_pct"] = 0.25

    s = Scenario(data)
    inputs = REoptInputs(s)
    m = Model(optimizer_with_attributes(Xpress.Optimizer, "MIPRELSTOP" => 0.01, "OUTPUTLOG" => 0))
    results = run_reopt(m, inputs)

    tot_elec_load = results["ElectricLoad"]["load_series_kw"]
    chp_total_elec_prod = results["CHP"]["year_one_electric_production_series_kw"]
    chp_to_load = results["CHP"]["year_one_to_load_series_kw"]
    chp_export = results["CHP"]["year_one_to_grid_series_kw"]
    #cooling_elec_load = results["LoadProfileChillerThermal"]["year_one_chiller_electric_load_kw"]

    # The values compared to the expected values
    #@test sum([(chp_to_load[i] - tot_elec_load[i]) for i in outage_start:outage_end])) == 0.0
    critical_load = tot_elec_load[outage_start:outage_end] * data["ElectricLoad"]["critical_load_pct"]
    @test sum(chp_to_load[outage_start:outage_end]) ≈ sum(critical_load) atol=0.1
    @test sum(chp_export) == 0.0
    @test sum(chp_total_elec_prod) ≈ sum(chp_to_load) atol=1.0e-5*sum(chp_total_elec_prod)
    #@test sum(cooling_elec_load[outage_start:outage_end]) == 0.0
    @test sum(chp_total_elec_prod[unavail_2_start:unavail_2_end]) == 0.0
end

@testset "CHP Supplementary firing and standby" begin
    """
    Test to ensure that supplementary firing and standby charges work as intended.  The thermal and
    electrical loads are constant, and the CHP system size is fixed; the supplementary firing has a
    similar cost to the boiler and is purcahsed and used when the boiler efficiency is set to a lower
    value than that of the supplementary firing. The test also ensures that demand charges are
    correctly calculated when CHP is and is not allowed to reduce demand charges.
    """
    data = JSON.parsefile("./scenarios/chp_supplementary_firing.json")
    data["CHP"]["supplementary_firing_capital_cost_per_kw"] = 10000
    data["ElectricLoad"]["loads_kw"] = repeat([800.0], 8760)
    data["DomesticHotWaterLoad"]["fuel_loads_mmbtu_per_hour"] = repeat([6.0], 8760)
    data["SpaceHeatingLoad"]["fuel_loads_mmbtu_per_hour"] = repeat([6.0], 8760)
    #part 1: supplementary firing not used when less efficient than the boiler and expensive
    m1 = Model(optimizer_with_attributes(Xpress.Optimizer, "OUTPUTLOG" => 0))
    s = Scenario(data)
    inputs = REoptInputs(s)
    results = run_reopt(m1, inputs)
    @test results["CHP"]["size_kw"] == 800
    @test results["CHP"]["size_supplemental_firing_kw"] == 0
    @test results["CHP"]["year_one_electric_energy_produced_kwh"] ≈ 800*8760 rtol=1e-5
    @test results["CHP"]["year_one_thermal_energy_produced_mmbtu"] ≈ 800*(0.4418/0.3573)*8760/293.07107 rtol=1e-5
    @test results["ElectricTariff"]["lifecycle_demand_cost"] == 0

    #part 2: supplementary firing used when more efficient than the boiler and low-cost; demand charges not reduced by CHP
    data["CHP"]["supplementary_firing_capital_cost_per_kw"] = 10
    data["CHP"]["reduces_demand_charges"] = false
    data["ExistingBoiler"]["efficiency"] = 0.85
    m2 = Model(optimizer_with_attributes(Xpress.Optimizer, "OUTPUTLOG" => 0))
    s = Scenario(data)
    inputs = REoptInputs(s)
    results = run_reopt(m2, inputs)
    @test results["CHP"]["size_supplemental_firing_kw"] ≈ 278.73 atol=0.1
    @test results["CHP"]["year_one_thermal_energy_produced_mmbtu"] ≈ 138624 rtol=1e-5
    @test results["ElectricTariff"]["lifecycle_demand_cost"] ≈ 5212.7 rtol=1e-5
end

#=
add a time-of-export rate that is greater than retail rate for the month of January,
check to make sure that PV does NOT export unless the site load is met first for the month of January.
=#
@testset "Do not allow_simultaneous_export_import" begin
    model = Model(optimizer_with_attributes(Xpress.Optimizer, "OUTPUTLOG" => 0))
    data = JSON.parsefile("./scenarios/monthly_rate.json")

    # create wholesale_rate with compensation in January > retail rate
    jan_rate = data["ElectricTariff"]["monthly_energy_rates"][1]
    data["ElectricTariff"]["wholesale_rate"] =
        append!(repeat([jan_rate + 0.1], 31 * 24), repeat([0.0], 8760 - 31*24))
    data["ElectricTariff"]["monthly_demand_rates"] = repeat([0], 12)
    data["ElectricUtility"] = Dict("allow_simultaneous_export_import" => false)

    s = Scenario(data)
    inputs = REoptInputs(s)
    results = run_reopt(model, inputs)

    @test all(x == 0.0 for (i,x) in enumerate(results["ElectricUtility"]["year_one_to_load_series_kw"][1:744])
              if results["PV"]["year_one_to_grid_series_kw"][i] > 0)
end

@testset "Solar and Storage w/BAU" begin
    m1 = Model(optimizer_with_attributes(Xpress.Optimizer, "OUTPUTLOG" => 0))
    m2 = Model(optimizer_with_attributes(Xpress.Optimizer, "OUTPUTLOG" => 0))
    results = run_reopt([m1,m2], "./scenarios/pv_storage.json")

    @test results["PV"]["size_kw"] ≈ 216.6667 atol=0.01
    @test results["PV"]["lcoe_per_kwh"] ≈ 0.0483 atol = 0.001
    @test results["Financial"]["lcc"] ≈ 1.240037e7 rtol=1e-5
    @test results["Financial"]["lcc_bau"] ≈ 12766397 rtol=1e-5
    @test results["Storage"]["size_kw"] ≈ 55.9 atol=0.1
    @test results["Storage"]["size_kwh"] ≈ 78.9 atol=0.1
    proforma_npv = REopt.npv(results["Financial"]["offtaker_annual_free_cashflows"] -
        results["Financial"]["offtaker_annual_free_cashflows_bau"], 0.081)
    @test results["Financial"]["npv"] ≈ proforma_npv rtol=0.0001
end

@testset "Outage with Generator, outate simulator, BAU critical load outputs" begin
    m1 = Model(optimizer_with_attributes(Xpress.Optimizer, "OUTPUTLOG" => 0))
    m2 = Model(optimizer_with_attributes(Xpress.Optimizer, "OUTPUTLOG" => 0))
    p = REoptInputs("./scenarios/generator.json")
    results = run_reopt([m1,m2], p)
    @test results["Generator"]["size_kw"] ≈ 8.13 atol=0.01
    @test (sum(results["Generator"]["year_one_to_load_series_kw"][i] for i in 1:9) +
           sum(results["Generator"]["year_one_to_load_series_kw"][i] for i in 13:8760)) == 0
    @test results["ElectricLoad"]["bau_critical_load_met"] == false
    @test results["ElectricLoad"]["bau_critical_load_met_time_steps"] == 0

    simresults = simulate_outages(results, p)
    @test simresults["resilience_hours_max"] == 11
end

@testset "Minimize Unserved Load" begin
    m = Model(optimizer_with_attributes(Xpress.Optimizer, "OUTPUTLOG" => 0))
    results = run_reopt(m, "./scenarios/outage.json")

    @test results["Outages"]["expected_outage_cost"] ≈ 0
    @test sum(results["Outages"]["unserved_load_per_outage"]) ≈ 0
    @test value(m[:binMGTechUsed]["Generator"]) == 1
    @test value(m[:binMGTechUsed]["PV"]) == 0
    @test value(m[:binMGStorageUsed]) == 1
    @test results["Financial"]["lcc"] ≈ 7.3879557e7 atol=5e4

    #=
    Scenario with $0/kWh value_of_lost_load_per_kwh, 12x169 hour outages, 1kW load/hour, and min_resil_timesteps = 168
    - should meet 168 kWh in each outage such that the total unserved load is 12 kWh
    =#
    m = Model(optimizer_with_attributes(Xpress.Optimizer, "OUTPUTLOG" => 0))
    results = run_reopt(m, "./scenarios/nogridcost_minresilhours.json")
    @test sum(results["Outages"]["unserved_load_per_outage"]) ≈ 12

    # testing dvUnserved load, which would output 100 kWh for this scenario before output fix
    m = Model(optimizer_with_attributes(Xpress.Optimizer, "OUTPUTLOG" => 0))
    results = run_reopt(m, "./scenarios/nogridcost_multiscenario.json")
    @test sum(results["Outages"]["unserved_load_per_outage"]) ≈ 60

end

@testset "Multiple Sites" begin
    m = Model(optimizer_with_attributes(Xpress.Optimizer, "OUTPUTLOG" => 0))
    ps = [
        REoptInputs("./scenarios/pv_storage.json"),
        REoptInputs("./scenarios/monthly_rate.json"),
    ];
    results = run_reopt(m, ps)
    @test results[3]["Financial"]["lcc"] + results[10]["Financial"]["lcc"] ≈ 1.240037e7 + 437169.0 rtol=1e-5
end

@testset "MPC" begin
    model = Model(optimizer_with_attributes(Xpress.Optimizer, "OUTPUTLOG" => 0))
    r = run_mpc(model, "./scenarios/mpc.json")
    @test maximum(r["ElectricUtility"]["to_load_series_kw"][1:15]) <= 98.0
    @test maximum(r["ElectricUtility"]["to_load_series_kw"][16:24]) <= 97.0
    @test sum(r["PV"]["to_grid_series_kw"]) ≈ 0
end

@testset "Complex Incentives" begin
    """
    This test was compared against the API test:
        reo.tests.test_reopt_url.EntryResourceTest.test_complex_incentives
    when using the hardcoded levelization_factor in this package's REoptInputs function.
    The two LCC's matched within 0.00005%. (The Julia pkg LCC is  1.0971991e7)
    """
    m = Model(optimizer_with_attributes(Xpress.Optimizer, "OUTPUTLOG" => 0))
    results = run_reopt(m, "./scenarios/incentives.json")
    @test results["Financial"]["lcc"] ≈ 1.094596365e7 atol=5e4
end

@testset verbose = true "Rate Structures" begin

    @testset "Tiered Energy" begin
        m = Model(optimizer_with_attributes(Xpress.Optimizer, "OUTPUTLOG" => 0))
        results = run_reopt(m, "./scenarios/tiered_rate.json")
        @test results["ElectricTariff"]["year_one_energy_cost"] ≈ 2342.88
    end

    @testset "Lookback Demand Charges" begin
        m = Model(optimizer_with_attributes(Xpress.Optimizer, "OUTPUTLOG" => 0))
        results = run_reopt(m, "./scenarios/lookback_rate.json")
        @test results["ElectricTariff"]["year_one_demand_cost"] ≈ 721.99
    end

    @testset "Blended tariff" begin
        model = Model(optimizer_with_attributes(Xpress.Optimizer, "OUTPUTLOG" => 0))
        results = run_reopt(model, "./scenarios/no_techs.json")
        @test results["ElectricTariff"]["year_one_energy_cost"] ≈ 1000.0
        @test results["ElectricTariff"]["year_one_demand_cost"] ≈ 136.99
    end

    @testset "Coincident Peak Charges" begin
        model = Model(optimizer_with_attributes(Xpress.Optimizer, "OUTPUTLOG" => 0))
        results = run_reopt(model, "./scenarios/coincident_peak.json")
        @test results["ElectricTariff"]["year_one_coincident_peak_cost"] ≈ 15.0
        @test results["ElectricTariff"]["lifecycle_coincident_peak_cost"] ≈ 15.0 * 12.94887 atol=0.1
    end

    @testset "URDB sell rate" begin
        #= The URDB contains at least one "Customer generation" tariff that only has a "sell" key in the energyratestructure (the tariff tested here)
        =#
        model = Model(optimizer_with_attributes(Xpress.Optimizer, "OUTPUTLOG" => 0))
        p = REoptInputs("./scenarios/URDB_customer_generation.json")
        results = run_reopt(model, p)
        @test results["PV"]["size_kw"] ≈ p.max_sizes["PV"]
    end

    # # tiered monthly demand rate  TODO: expected results?
    # m = Model(optimizer_with_attributes(Xpress.Optimizer, "OUTPUTLOG" => 0))
    # data = JSON.parsefile("./scenarios/tiered_rate.json")
    # data["ElectricTariff"]["urdb_label"] = "59bc22705457a3372642da67"
    # s = Scenario(data)
    # inputs = REoptInputs(s)
    # results = run_reopt(m, inputs)

    # TODO test for tiered TOU demand rates
end

@testset "Wind" begin
    m = Model(optimizer_with_attributes(Xpress.Optimizer, "OUTPUTLOG" => 0))
    results = run_reopt(m, "./scenarios/wind.json")
    @test results["Wind"]["size_kw"] ≈ 3752 atol=0.1
    @test results["Financial"]["lcc"] ≈ 8.591017e6 rtol=1e-5
    #=
    0.5% higher LCC in this package as compared to API ? 8,591,017 vs 8,551,172
    - both have zero curtailment
    - same energy to grid: 5,839,317 vs 5,839,322
    - same energy to load: 4,160,683 vs 4,160,677
    - same city: Boulder
    - same total wind prod factor

    REopt.jl has:
    - bigger turbine: 3752 vs 3735
    - net_capital_costs_plus_om: 8,576,590 vs. 8,537,480

    TODO: will these discrepancies be addressed once NMIL binaries are added?
    =#
end

@testset "Multiple PVs" begin
    m1 = Model(optimizer_with_attributes(Xpress.Optimizer, "OUTPUTLOG" => 0))
    m2 = Model(optimizer_with_attributes(Xpress.Optimizer, "OUTPUTLOG" => 0))
    results = run_reopt([m1,m2], "./scenarios/multiple_pvs.json")

    ground_pv = results["PV"][findfirst(pv -> pv["name"] == "ground", results["PV"])]
    roof_west = results["PV"][findfirst(pv -> pv["name"] == "roof_west", results["PV"])]
    roof_east = results["PV"][findfirst(pv -> pv["name"] == "roof_east", results["PV"])]

    @test ground_pv["size_kw"] ≈ 15 atol=0.1
    @test roof_west["size_kw"] ≈ 7 atol=0.1
    @test roof_east["size_kw"] ≈ 4 atol=0.1
    @test ground_pv["lifecycle_om_cost_bau"] ≈ 782.0 atol=0.1
    @test roof_west["lifecycle_om_cost_bau"] ≈ 782.0 atol=0.1
    @test ground_pv["average_annual_energy_produced_kwh_bau"] ≈ 8844.19 atol=0.1
    @test roof_west["average_annual_energy_produced_kwh_bau"] ≈ 7440.1 atol=0.1
    @test ground_pv["average_annual_energy_produced_kwh"] ≈ 26533.54 atol=0.1
    @test roof_west["average_annual_energy_produced_kwh"] ≈ 10416.52 atol=0.1
    @test roof_east["average_annual_energy_produced_kwh"] ≈ 6482.37 atol=0.1
end

## equivalent REopt API Post for test 2:
#   NOTE have to hack in API levelization_factor to get LCC within 5e-5 (Mosel tol)
# {"Scenario": {
#     "Site": {
#         "longitude": -118.1164613,
#         "latitude": 34.5794343,
#         "roof_squarefeet": 5000.0,
#         "land_acres": 1.0,
#     "PV": {
#         "macrs_bonus_pct": 0.4,
#         "installed_cost_per_kw": 2000.0,
#         "tilt": 34.579,
#         "degradation_pct": 0.005,
#         "macrs_option_years": 5,
#         "federal_itc_pct": 0.3,
#         "module_type": 0,
#         "array_type": 1,
#         "om_cost_per_kw": 16.0,
#         "macrs_itc_reduction": 0.5,
#         "azimuth": 180.0,
#         "federal_rebate_per_kw": 350.0,
#         "dc_ac_ratio": 1.1
#     },
#     "LoadProfile": {
#         "doe_reference_name": "RetailStore",
#         "annual_kwh": 10000000.0,
#         "city": "LosAngeles"
#     },
#     "Storage": {
#         "total_rebate_per_kw": 100.0,
#         "macrs_option_years": 5,
#         "can_grid_charge": true,
#         "macrs_bonus_pct": 0.4,
#         "macrs_itc_reduction": 0.5,
#         "total_itc_pct": 0,
#         "installed_cost_per_kw": 1000.0,
#         "installed_cost_per_kwh": 500.0,
#         "replace_cost_per_kw": 460.0,
#         "replace_cost_per_kwh": 230.0
#     },
#     "ElectricTariff": {
#         "urdb_label": "5ed6c1a15457a3367add15ae"
#     },
#     "Financial": {
#         "escalation_pct": 0.026,
#         "offtaker_discount_pct": 0.081,
#         "owner_discount_pct": 0.081,
#         "analysis_years": 20,
#         "offtaker_tax_pct": 0.4,
#         "owner_tax_pct": 0.4,
#         "om_cost_escalation_pct": 0.025
#     }
# }}}


module LinearEnrollment
    using CurricularAnalyticsSimulationLibrary: Student

    function enroll!(currentTerm, simulation, max_credits)
        studentProgress = simulation.studentProgress
        terms = simulation.curriculum.terms

        for (termnum, term) in enumerate(terms)
            # Itterate through courses
            for course in term.courses
                # Clear the array of enrolled students for the course
                course.students = Student[]

                for student in simulation.enrolledStudents
                    # Enroll in coreqs first
                    for coreq in course.coreqs
                        if canEnroll(student, coreq, studentProgress, max_credits, currentTerm)
                            # Enroll the student in the course
                            push!(course.students, student)

                            # Increment the course's enrollment counters
                            course.enrolled += 1
                            course.termenrollment[currentTerm] += 1

                            # Increse the student's term credits
                            student.termcredits += course.credits
                        end
                    end

                    # Determine wheter the student can be enrolled in the current course.
                    if canEnroll(student, course, studentProgress, max_credits, currentTerm)

                        # Enroll the student in the course
                        push!(course.students, student)

                        # Increment the course's enrollment counters
                        course.enrolled += 1
                        course.termenrollment[currentTerm] += 1

                        # Increse the student's term credits
                        student.termcredits += course.credits
                    end
                end
            end
        end
    end

    # Function that determines wheter a student can enroll in a course
    function canEnroll(student, course, studentProgress, max_credits, term)
        # Find the prereq ids of the current course
        prereqIds = map(x -> x.id, course.prereqs)

        !in(student, course.students) &&
        (length(course.prereqs) == 0 || sum(studentProgress[student.id, prereqIds]) == length(course.prereqs)) &&       # No Prereqs or the student has completed them
        studentProgress[student.id, course.id] == 0.0 &&                                                                # The student has not already completed the course
        student.termcredits + course.credits <= max_credits &&                                                          # The student will not exceed the maximum number of credit hours
        course.termReq <= term &&                                                                                # The student must wait until the term req has been met
        enrolledInCoreqs(student, course, studentProgress)                                                              # The student is enrolled in or has completed coreqs
    end

    # Determines whether a student is enrolled in or has completed coreqs for a given course
    function enrolledInCoreqs(student, course, studentProgress)
        enrolled = true

        for coreq in course.coreqs
            enrolled =  enrolled && (in(student, coreq.students) || studentProgress[student.id, coreq.id] == 1.0)
        end

        return enrolled
    end
end


# VTK cell definitions

# Definitions copied from the vtkCellType.h file of the VTK source code:
# https://raw.githubusercontent.com/Kitware/VTK/master/Common/DataModel/vtkCellType.h

__precompile__()

module VTKCellTypes

export VTKCellType

struct VTKCellType
    vtk_name::String
    vtk_id::UInt8
    nodes::Int
end

const VTK_CELL_TYPES = Dict{UInt8, VTKCellType}()

# Construction of VTKCellType objects from the integer vtk id
function VTKCellType(vtkid::Integer)
    id = UInt8(vtkid)
    if !haskey(VTK_CELL_TYPES, id)
        throw(ArgumentError("unknown vtkid = $vtkid"))
    end
    VTK_CELL_TYPES[id]
end

# Define a VTK cell type
macro add(ctype, id, nodes)
    name = string(ctype)
    quote
        const $(esc(ctype)) = VTKCellType($name, $id, $nodes)
        VTK_CELL_TYPES[$id] = $(esc(ctype))
    end
end

# Linear cells
@add VTK_EMPTY_CELL 0 0
@add VTK_VERTEX 1  1
@add VTK_POLY_VERTEX 2 -1
@add VTK_LINE 3  2
@add VTK_POLY_LINE 4 -1
@add VTK_TRIANGLE 5  3
@add VTK_TRIANGLE_STRIP 6 -1
@add VTK_POLYGON 7 -1
@add VTK_PIXEL 8  4
@add VTK_QUAD 9  4
@add VTK_TETRA 10 4
@add VTK_VOXEL 11 8
@add VTK_HEXAHEDRON 12 8
@add VTK_WEDGE 13 6
@add VTK_PYRAMID 14 5
@add VTK_PENTAGONAL_PRISM 15 10
@add VTK_HEXAGONAL_PRISM 16 12

# Quadratic, isoparametric cells
@add VTK_QUADRATIC_EDGE 21 3
@add VTK_QUADRATIC_TRIANGLE 22 6
@add VTK_QUADRATIC_QUAD 23 8
@add VTK_QUADRATIC_POLYGON 36 -1
@add VTK_QUADRATIC_TETRA 24 10
@add VTK_QUADRATIC_HEXAHEDRON 25 20
@add VTK_QUADRATIC_WEDGE 26 15
@add VTK_QUADRATIC_PYRAMID 27 13
@add VTK_BIQUADRATIC_QUAD 28 9
@add VTK_TRIQUADRATIC_HEXAHEDRON 29 27
@add VTK_QUADRATIC_LINEAR_QUAD 30 6
@add VTK_QUADRATIC_LINEAR_WEDGE 31 12
@add VTK_BIQUADRATIC_QUADRATIC_WEDGE 32 18
@add VTK_BIQUADRATIC_QUADRATIC_HEXAHEDRON 33 24
@add VTK_BIQUADRATIC_TRIANGLE 34 7

# Cubic, isoparametric cell
@add VTK_CUBIC_LINE 35 4

# Special class of cells formed by convex group of points
@add VTK_CONVEX_POINT_SET 41 -1

# Polyhedron cell (consisting of polygonal faces)
@add VTK_POLYHEDRON 42 -1

# Higher order cells in parametric form
@add VTK_PARAMETRIC_CURVE 51 -1
@add VTK_PARAMETRIC_SURFACE 52 -1
@add VTK_PARAMETRIC_TRI_SURFACE 53 -1
@add VTK_PARAMETRIC_QUAD_SURFACE 54 -1
@add VTK_PARAMETRIC_TETRA_REGION 55 -1
@add VTK_PARAMETRIC_HEX_REGION 56 -1

# Higher order cells
@add VTK_HIGHER_ORDER_EDGE 60 -1
@add VTK_HIGHER_ORDER_TRIANGLE 61 -1
@add VTK_HIGHER_ORDER_QUAD 62 -1
@add VTK_HIGHER_ORDER_POLYGON 63 -1
@add VTK_HIGHER_ORDER_TETRAHEDRON 64 -1
@add VTK_HIGHER_ORDER_WEDGE 65 -1
@add VTK_HIGHER_ORDER_PYRAMID 66 -1
@add VTK_HIGHER_ORDER_HEXAHEDRON 67 -1

end # module