text:
module TestNullableArray

using Base.Test
using CategoricalArrays
using NullableArrays
using CategoricalArrays: DefaultRefType
using Compat

# == currently throws an error for Nullables
(==) = isequal

for ordered in (false, true)
    for R in (CategoricalArrays.DefaultRefType, UInt8, UInt, Int8, Int)
        # Vector with no null values
        for a in (["b", "a", "b"],
                  Nullable{String}["b", "a", "b"],
                  NullableArray(["b", "a", "b"]))
            x = NullableCategoricalVector{String, R}(a, ordered=ordered)
            na = eltype(a) <: Nullable ? a : convert(Array{Nullable{String}}, a)

            @test x == na
            @test isordered(x) === ordered
            @test levels(x) == sort(map(get, unique(na)))
            @test size(x) === (3,)
            @test length(x) === 3

            @test convert(NullableCategoricalArray, x) === x
            @test convert(NullableCategoricalArray{String}, x) === x
            @test convert(NullableCategoricalArray{String, 1}, x) === x
            @test convert(NullableCategoricalArray{String, 1, R}, x) === x
            @test convert(NullableCategoricalArray{String, 1, DefaultRefType}, x) == x
            @test convert(NullableCategoricalArray{String, 1, UInt8}, x) == x

            @test convert(NullableCategoricalVector, x) === x
            @test convert(NullableCategoricalVector{String}, x) === x
            @test convert(NullableCategoricalVector{String, R}, x) === x
            @test convert(NullableCategoricalVector{String, DefaultRefType}, x) == x
            @test convert(NullableCategoricalVector{String, UInt8}, x) == x

            @test convert(NullableCategoricalArray, a) == x
            @test convert(NullableCategoricalArray{String}, a) == x
            @test convert(NullableCategoricalArray{String, 1}, a) == x
            @test convert(NullableCategoricalArray{String, 1, R}, a) == x
            @test convert(NullableCategoricalArray{String, 1, DefaultRefType}, a) == x
            @test convert(NullableCategoricalArray{String, 1, UInt8}, a) == x

            @test convert(NullableCategoricalVector, a) == x
            @test convert(NullableCategoricalVector{String}, a) == x
            @test convert(NullableCategoricalVector{String, R}, a) == x
            @test convert(NullableCategoricalVector{String, DefaultRefType}, a) == x
            @test convert(NullableCategoricalVector{String, UInt8}, a) == x

            @test NullableCategoricalArray{String}(a, ordered=ordered) == x
            @test NullableCategoricalArray{String, 1}(a, ordered=ordered) == x
            @test NullableCategoricalArray{String, 1, R}(a, ordered=ordered) == x
            @test NullableCategoricalArray{String, 1, DefaultRefType}(a, ordered=ordered) == x
            @test NullableCategoricalArray{String, 1, UInt8}(a, ordered=ordered) == x

            @test NullableCategoricalVector(a, ordered=ordered) == x
            @test NullableCategoricalVector{String}(a, ordered=ordered) == x
            @test NullableCategoricalVector{String, R}(a, ordered=ordered) == x
            @test NullableCategoricalVector{String, DefaultRefType}(a, ordered=ordered) == x
            @test NullableCategoricalVector{String, UInt8}(a, ordered=ordered) == x

            for (y, R1, R2, comp) in ((a, DefaultRefType, UInt8, true),
                                      (a, DefaultRefType, DefaultRefType, false),
                                      (x, R, UInt8, true),
                                      (x, R, R, false))
                x2 = categorical(y, ordered=ordered)
                @test x2 == y
                if eltype(y) <: Nullable
                    @test isa(x2, NullableCategoricalVector{String, R1})
                else
                    @test isa(x2, CategoricalVector{String, R1})
                end
                @test isordered(x2) === ordered

                x2 = categorical(y, comp, ordered=ordered)
                @test x2 == y
                if eltype(y) <: Nullable
                    @test isa(x2, NullableCategoricalVector{String, R2})
                else
                    @test isa(x2, CategoricalVector{String, R2})
                end
                @test isordered(x2) === ordered
            end

            x2 = compress(x)
            @test x2 == x
            @test isa(x2, NullableCategoricalVector{String, UInt8})
            @test isordered(x2) === isordered(x)
            @test levels(x2) == levels(x)

            x2 = copy(x)
            @test x2 == x
            @test isordered(x2) === isordered(x)
            @test typeof(x2) === typeof(x)
            @test levels(x2) == levels(x)

            if !isordered(x)
                @test ordered!(x, true) === x
                @test isordered(x) === true
            end
            @test get(x[1] > x[2])
            @test get(x[3] > x[2])

            @test ordered!(x, false) === x
            @test isordered(x) === false
            @test_throws Exception x[1] > x[2]
            @test_throws Exception x[3] > x[2]

            @test x[1] === Nullable(x.pool.valindex[1])
            @test x[2] === Nullable(x.pool.valindex[2])
            @test x[3] === Nullable(x.pool.valindex[1])
            @test_throws BoundsError x[4]

            x2 = x[:]
            @test typeof(x2) === typeof(x)
            @test x2 == x
            @test x2 !== x
            @test levels(x2) == levels(x)
            @test levels(x2) !== levels(x)
            @test isordered(x2) == isordered(x)

            x2 = x[2:3]
            @test typeof(x2) === typeof(x)
            @test x2 == Nullable{String}["a", "b"]
            @test levels(x2) == levels(x)
            @test levels(x2) !== levels(x)
            @test isordered(x2) == isordered(x)

            x2 = x[1:1]
            @test typeof(x2) === typeof(x)
            @test x2 == Nullable{String}["b"]
            @test levels(x2) == levels(x)
            @test levels(x2) !== levels(x)
            @test isordered(x2) == isordered(x)

            x2 = x[2:1]
            @test typeof(x2) === typeof(x)
            @test isempty(x2)
            @test levels(x2) == levels(x)
            @test levels(x2) !== levels(x)
            @test isordered(x2) == isordered(x)

            x[1] = x[2]
            @test x[1] === Nullable(x.pool.valindex[2])
            @test x[2] === Nullable(x.pool.valindex[2])
            @test x[3] === Nullable(x.pool.valindex[1])

            x[3] = "c"
            @test x[1] === Nullable(x.pool.valindex[2])
            @test x[2] === Nullable(x.pool.valindex[2])
            @test x[3] === Nullable(x.pool.valindex[3])
            @test levels(x) == ["a", "b", "c"]

            x[2:3] = "b"
            @test x[1] === Nullable(x.pool.valindex[2])
            @test x[2] === Nullable(x.pool.valindex[1])
            @test x[3] === Nullable(x.pool.valindex[1])
            @test levels(x) == ["a", "b", "c"]

            @test droplevels!(x) === x
            @test levels(x) == ["a", "b"]
            @test x[1] === Nullable(x.pool.valindex[1])
            @test x[2] === Nullable(x.pool.valindex[2])
            @test x[3] === Nullable(x.pool.valindex[2])
            @test levels(x) == ["a", "b"]

            @test levels!(x, ["b", "a"]) === x
            @test levels(x) == ["b", "a"]
            @test x[1] === Nullable(x.pool.valindex[1])
            @test x[2] === Nullable(x.pool.valindex[2])
            @test x[3] === Nullable(x.pool.valindex[2])
            @test levels(x) == ["b", "a"]

            @test_throws ArgumentError levels!(x, ["a"])
            @test_throws ArgumentError levels!(x, ["e", "b"])
            @test_throws ArgumentError levels!(x, ["e", "a", "b", "a"])

            @test levels!(x, ["e", "a", "b"]) === x
            @test levels(x) == ["e", "a", "b"]
            @test x[1] === Nullable(x.pool.valindex[1])
            @test x[2] === Nullable(x.pool.valindex[2])
            @test x[3] === Nullable(x.pool.valindex[2])
            @test levels(x) == ["e", "a", "b"]

            x[1] = "c"
            @test x[1] === Nullable(x.pool.valindex[4])
            @test x[2] === Nullable(x.pool.valindex[2])
            @test x[3] === Nullable(x.pool.valindex[2])
            @test levels(x) == ["e", "a", "b", "c"]

            @test_throws ArgumentError levels!(x, ["e", "c"])
            @test levels!(x, ["e", "c"], nullok=true) === x
            @test levels(x) == ["e", "c"]
            @test x[1] === Nullable(x.pool.valindex[2])
            @test x[2] === eltype(x)()
            @test x[3] === eltype(x)()
            @test levels(x) == ["e", "c"]

            push!(x, "e")
            @test length(x) == 4
            @test isequal(x, NullableArray(["c", "", "", "e"], [false, true, true, false]))
            @test levels(x) == ["e", "c"]

            push!(x, "zz")
            @test length(x) == 5
            @test isequal(x, NullableArray(["c", "", "", "e", "zz"], [false, true, true, false, false]))
            @test levels(x) == ["e", "c", "zz"]

            push!(x, x[1])
            @test length(x) == 6
            @test isequal(x, NullableArray(["c", "", "", "e", "zz", "c"], [false, true, true, false, false, false]))
            @test levels(x) == ["e", "c", "zz"]

            push!(x, eltype(x)())
            @test length(x) == 7
            @test isequal(x, NullableArray(["c", "", "", "e", "zz", "c", ""], [false, true, true, false, false, false, true]))
            @test isnull(x[end])
            @test levels(x) == ["e", "c", "zz"]

            append!(x, x)
            @test isequal(x, NullableArray(["c", "", "", "e", "zz", "c", "", "c", "", "", "e", "zz", "c", ""], [false, true, true, false, false, false, true, false, true, true, false, false, false, true]))
            @test levels(x) == ["e", "c", "zz"]
            @test isordered(x) === false
            @test length(x) == 14

            b = ["z","y","x"]
            y = NullableCategoricalVector{String, R}(b)
            append!(x, y)
            @test length(x) == 17
            @test isordered(x) === false
            @test levels(x) == ["e", "c", "zz", "x", "y", "z"]
            @test isequal(x, NullableArray(["c", "", "", "e", "zz", "c", "", "c", "", "", "e", "zz", "c", "", "z", "y", "x"], [false, true, true, false, false, false, true, false, true, true, false, false, false, true, false, false, false]))

            empty!(x)
            @test isordered(x) === false
            @test length(x) == 0
            @test levels(x) == ["e", "c", "zz", "x", "y", "z"]
        end

        # Vector with null values
        for a in (Nullable{String}["a", "b", Nullable()],
                  NullableArray(Nullable{String}["a", "b", Nullable()]))
            x = NullableCategoricalVector{String, R}(a, ordered=ordered)

            @test x == a
            @test levels(x) == map(get, filter(x->!isnull(x), unique(a)))
            @test size(x) === (3,)
            @test length(x) === 3

            @test convert(NullableCategoricalArray, x) === x
            @test convert(NullableCategoricalArray{String}, x) === x
            @test convert(NullableCategoricalArray{String, 1}, x) === x
            @test convert(NullableCategoricalArray{String, 1, R}, x) === x
            @test convert(NullableCategoricalArray{String, 1, DefaultRefType}, x) == x
            @test convert(NullableCategoricalArray{String, 1, UInt8}, x) == x

            @test convert(NullableCategoricalVector, x) === x
            @test convert(NullableCategoricalVector{String}, x) === x
            @test convert(NullableCategoricalVector{String, R}, x) === x
            @test convert(NullableCategoricalVector{String, DefaultRefType}, x) == x
            @test convert(NullableCategoricalVector{String, UInt8}, x) == x

            @test convert(NullableCategoricalArray, a) == x
            @test convert(NullableCategoricalArray{String}, a) == x
            @test convert(NullableCategoricalArray{String, 1}, a) == x
            @test convert(NullableCategoricalArray{String, 1, R}, a) == x
            @test convert(NullableCategoricalArray{String, 1, DefaultRefType}, a) == x
            @test convert(NullableCategoricalArray{String, 1, UInt8}, a) == x

            @test convert(NullableCategoricalVector, a) == x
            @test convert(NullableCategoricalVector{String}, a) == x
            @test convert(NullableCategoricalVector{String, R}, a) == x
            @test convert(NullableCategoricalVector{String, DefaultRefType}, a) == x
            @test convert(NullableCategoricalVector{String, UInt8}, a) == x

            @test NullableCategoricalArray{String}(a, ordered=ordered) == x
            @test NullableCategoricalArray{String, 1}(a, ordered=ordered) == x
            @test NullableCategoricalArray{String, 1, R}(a, ordered=ordered) == x
            @test NullableCategoricalArray{String, 1, DefaultRefType}(a, ordered=ordered) == x
            @test NullableCategoricalArray{String, 1, UInt8}(a, ordered=ordered) == x

            @test NullableCategoricalVector(a, ordered=ordered) == x
            @test NullableCategoricalVector{String}(a, ordered=ordered) == x
            @test NullableCategoricalVector{String, R}(a, ordered=ordered) == x
            @test NullableCategoricalVector{String, DefaultRefType}(a, ordered=ordered) == x
            @test NullableCategoricalVector{String, UInt8}(a, ordered=ordered) == x

            for (y, R1, R2, comp) in ((a, DefaultRefType, UInt8, true),
                                      (a, DefaultRefType, DefaultRefType, false),
                                      (x, R, UInt8, true),
                                      (x, R, R, false))
                x2 = categorical(y, ordered=ordered)
                @test x2 == y
                @test isa(x2, NullableCategoricalVector{String, R1})
                @test isordered(x2) === ordered

                x2 = categorical(y, comp, ordered=ordered)
                @test x2 == y
                @test isa(x2, NullableCategoricalVector{String, R2})
                @test isordered(x2) === ordered
            end

            x2 = compress(x)
            @test x2 == x
            @test isa(x2, NullableCategoricalVector{String, UInt8})
            @test levels(x2) == levels(x)

            x2 = copy(x)
            @test x2 == x
            @test typeof(x2) === typeof(x)
            @test levels(x2) == levels(x)

            @test x[1] === Nullable(x.pool.valindex[1])
            @test x[2] === Nullable(x.pool.valindex[2])
            @test x[3] === eltype(x)()
            @test_throws BoundsError x[4]

            x2 = x[:]
            @test typeof(x2) === typeof(x)
            @test x2 == x
            @test x2 !== x
            @test levels(x2) == levels(x)
            @test levels(x2) !== levels(x)
            @test isordered(x2) == isordered(x)

            x2 = x[2:3]
            @test typeof(x2) === typeof(x)
            @test x2 == Nullable{String}["b", Nullable()]
            @test levels(x2) == levels(x)
            @test levels(x2) !== levels(x)
            @test isordered(x2) == isordered(x)

            x2 = x[1:1]
            @test typeof(x2) === typeof(x)
            @test x2 == Nullable{String}["a"]
            @test levels(x2) == levels(x)
            @test levels(x2) !== levels(x)
            @test isordered(x2) == isordered(x)

            x2 = x[2:1]
            @test typeof(x2) === typeof(x)
            @test isempty(x2)
            @test levels(x2) == levels(x)
            @test levels(x2) !== levels(x)
            @test isordered(x2) == isordered(x)

            x[1] = "b"
            @test x[1] === Nullable(x.pool.valindex[2])
            @test x[2] === Nullable(x.pool.valindex[2])
            @test x[3] === eltype(x)()

            x[3] = "c"
            @test x[1] === Nullable(x.pool.valindex[2])
            @test x[2] === Nullable(x.pool.valindex[2])
            @test x[3] === Nullable(x.pool.valindex[3])
            @test levels(x) == ["a", "b", "c"]

            x[1] = Nullable()
            @test x[1] === eltype(x)()
            @test x[2] === Nullable(x.pool.valindex[2])
            @test x[3] === Nullable(x.pool.valindex[3])
            @test levels(x) == ["a", "b", "c"]

            x[2:3] = Nullable()
            @test x[1] === eltype(x)()
            @test x[2] === eltype(x)()
            @test x[3] === eltype(x)()
            @test levels(x) == ["a", "b", "c"]
        end

        # Vector created from range (i.e. non-Array AbstractArray),
        # direct conversion to a vector with different eltype
        a = 0.0:0.5:1.5
        x = NullableCategoricalVector{Float64, R}(a, ordered=ordered)

        @test x == map(Nullable, a)
        @test isordered(x) === ordered
        @test levels(x) == unique(a)
        @test size(x) === (4,)
        @test length(x) === 4

        @test convert(NullableCategoricalArray, x) === x
        @test convert(NullableCategoricalArray{Float64}, x) === x
        @test convert(NullableCategoricalArray{Float64, 1}, x) === x
        @test convert(NullableCategoricalArray{Float64, 1, R}, x) === x
        @test convert(NullableCategoricalArray{Float64, 1, DefaultRefType}, x) == x
        @test convert(NullableCategoricalArray{Float64, 1, UInt8}, x) == x

        @test convert(NullableCategoricalVector, x) === x
        @test convert(NullableCategoricalVector{Float64}, x) === x
        @test convert(NullableCategoricalVector{Float64, R}, x) === x
        @test convert(NullableCategoricalVector{Float64, DefaultRefType}, x) == x
        @test convert(NullableCategoricalVector{Float64, UInt8}, x) == x

        @test convert(NullableCategoricalArray, a) == x
        @test convert(NullableCategoricalArray{Float64}, a) == x
        @test convert(NullableCategoricalArray{Float32}, a) == x
        @test convert(NullableCategoricalArray{Float64, 1}, a) == x
        @test convert(NullableCategoricalArray{Float32, 1}, a) == x
        @test convert(NullableCategoricalArray{Float64, 1, R}, a) == x
        @test convert(NullableCategoricalArray{Float32, 1, R}, a) == x
        @test convert(NullableCategoricalArray{Float64, 1, DefaultRefType}, a) == x
        @test convert(NullableCategoricalArray{Float32, 1, DefaultRefType}, a) == x
        @test convert(NullableCategoricalArray{Float64, 1, UInt8}, a) == x
        @test convert(NullableCategoricalArray{Float32, 1, UInt8}, a) == x

        @test convert(NullableCategoricalVector, a) == x
        @test convert(NullableCategoricalVector{Float64}, a) == x
        @test convert(NullableCategoricalVector{Float32}, a) == x
        @test convert(NullableCategoricalVector{Float64, R}, a) == x
        @test convert(NullableCategoricalVector{Float32, R}, a) == x
        @test convert(NullableCategoricalVector{Float64, DefaultRefType}, a) == x
        @test convert(NullableCategoricalVector{Float32, DefaultRefType}, a) == x
        @test convert(NullableCategoricalVector{Float64, UInt8}, a) == x
        @test convert(NullableCategoricalVector{Float32, UInt8}, a) == x

        @test NullableCategoricalArray{Float64}(a, ordered=ordered) == x
        @test NullableCategoricalArray{Float32}(a, ordered=ordered) == x
        @test NullableCategoricalArray{Float64, 1}(a, ordered=ordered) == x
        @test NullableCategoricalArray{Float32, 1}(a, ordered=ordered) == x
        @test NullableCategoricalArray{Float64, 1, R}(a, ordered=ordered) == x
        @test NullableCategoricalArray{Float32, 1, R}(a, ordered=ordered) == x
        @test NullableCategoricalArray{Float64, 1, DefaultRefType}(a, ordered=ordered) == x
        @test NullableCategoricalArray{Float32, 1, DefaultRefType}(a, ordered=ordered) == x

        @test NullableCategoricalVector(a, ordered=ordered) == x
        @test NullableCategoricalVector{Float64}(a, ordered=ordered) == x
        @test NullableCategoricalVector{Float32}(a, ordered=ordered) == x
        @test NullableCategoricalVector{Float64, R}(a, ordered=ordered) == x
        @test NullableCategoricalVector{Float32, R}(a, ordered=ordered) == x
        @test NullableCategoricalVector{Float64, DefaultRefType}(a, ordered=ordered) == x
        @test NullableCategoricalVector{Float32, DefaultRefType}(a, ordered=ordered) == x

        for (y, R1, R2, comp) in ((a, DefaultRefType, UInt8, true),
                                  (a, DefaultRefType, DefaultRefType, false),
                                  (x, R, UInt8, true),
                                  (x, R, R, false))
            x2 = categorical(y, ordered=ordered)
            @test x2 == collect(y)
            if eltype(y) <: Nullable
                @test isa(x2, NullableCategoricalVector{Float64, R1})
            else
                @test isa(x2, CategoricalVector{Float64, R1})
            end
            @test isordered(x2) === ordered

            x2 = categorical(y, comp, ordered=ordered)
            @test x2 == collect(y)
            if eltype(y) <: Nullable
                @test isa(x2, NullableCategoricalVector{Float64, R2})
            else
                @test isa(x2, CategoricalVector{Float64, R2})
            end
            @test isordered(x2) === ordered
        end

        x2 = compress(x)
        @test x2 == x
        @test isordered(x2) === isordered(x)
        @test isa(x2, NullableCategoricalVector{Float64, UInt8})
        @test levels(x2) == levels(x)

        x2 = copy(x)
        @test x2 == x
        @test isordered(x2) === isordered(x)
        @test typeof(x2) === typeof(x)
        @test levels(x2) == levels(x)

        @test x[1] === Nullable(x.pool.valindex[1])
        @test x[2] === Nullable(x.pool.valindex[2])
        @test x[3] === Nullable(x.pool.valindex[3])
        @test x[4] === Nullable(x.pool.valindex[4])
        @test_throws BoundsError x[5]

        x2 = x[:]
        @test typeof(x2) === typeof(x)
        @test x2 == x
        @test x2 !== x
        @test levels(x2) == levels(x)
        @test levels(x2) !== levels(x)
        @test isordered(x2) == isordered(x)

        x2 = x[1:2]
        @test typeof(x2) === typeof(x)
        @test x2 == Nullable{Float64}[0.0, 0.5]
        @test levels(x2) == levels(x)
        @test levels(x2) !== levels(x)
        @test isordered(x2) == isordered(x)

        x2 = x[1:1]
        @test typeof(x2) === typeof(x)
        @test x2 == Nullable{Float64}[0.0]
        @test levels(x2) == levels(x)
        @test levels(x2) !== levels(x)
        @test isordered(x2) == isordered(x)

        x2 = x[2:1]
        @test typeof(x2) === typeof(x)
        @test isempty(x2)
        @test levels(x2) == levels(x)
        @test levels(x2) !== levels(x)
        @test isordered(x2) == isordered(x)

        x[2] = 1
        @test x[1] === Nullable(x.pool.valindex[1])
        @test x[2] === Nullable(x.pool.valindex[3])
        @test x[3] === Nullable(x.pool.valindex[3])
        @test x[4] === Nullable(x.pool.valindex[4])
        @test levels(x) == unique(a)

        x[1:2] = -1
        @test x[1] === Nullable(x.pool.valindex[5])
        @test x[2] === Nullable(x.pool.valindex[5])
        @test x[3] === Nullable(x.pool.valindex[3])
        @test x[4] === Nullable(x.pool.valindex[4])
        @test levels(x) == vcat(unique(a), -1)

        push!(x, 2.0)
        @test length(x) == 5
        @test isequal(x, NullableArray([-1.0, -1.0, 1.0, 1.5, 2.0]))
        @test isordered(x) === ordered
        @test levels(x) == [0.0,  0.5,  1.0,  1.5, -1.0,  2.0]

        push!(x, x[1])
        @test length(x) == 6
        @test isequal(x, NullableArray([-1.0, -1.0, 1.0, 1.5, 2.0, -1.0]))
        @test isordered(x) === ordered
        @test levels(x) == [0.0,  0.5,  1.0,  1.5, -1.0,  2.0]

        append!(x, x)
        @test length(x) == 12
        @test isequal(x, NullableArray([-1.0, -1.0, 1.0, 1.5, 2.0, -1.0, -1.0, -1.0, 1.0, 1.5, 2.0, -1.0]))
        @test isordered(x) === ordered
        @test levels(x) == [0.0,  0.5,  1.0,  1.5, -1.0,  2.0]

        b = [2.5, 3.0, -3.5]
        y = NullableCategoricalVector{Float64, R}(b)
        append!(x, y)
        @test length(x) == 15
        @test isequal(x, NullableArray([-1.0, -1.0, 1.0, 1.5, 2.0, -1.0, -1.0, -1.0, 1.0, 1.5, 2.0, -1.0, 2.5, 3.0, -3.5]))
        @test isordered(x) === ordered
        @test levels(x) == [0.0, 0.5, 1.0, 1.5, -1.0, 2.0, -3.5, 2.5, 3.0]

        empty!(x)
        @test length(x) == 0
        @test isordered(x) === ordered
        @test levels(x) == [0.0, 0.5, 1.0, 1.5, -1.0, 2.0, -3.5, 2.5, 3.0]

        # Matrix with no null values
        for a in (["a" "b" "c"; "b" "a" "c"],
                  Nullable{String}["a" "b" "c"; "b" "a" "c"],
                  NullableArray(["a" "b" "c"; "b" "a" "c"]))
            na = eltype(a) <: Nullable ? a : convert(Array{Nullable{String}}, a)
            x = NullableCategoricalMatrix{String, R}(a, ordered=ordered)

            @test x == na
            @test isordered(x) === ordered
            @test levels(x) == map(get, unique(na))
            @test size(x) === (2, 3)
            @test length(x) === 6

            @test convert(NullableCategoricalArray, x) === x
            @test convert(NullableCategoricalArray{String}, x) === x
            @test convert(NullableCategoricalArray{String, 2}, x) === x
            @test convert(NullableCategoricalArray{String, 2, R}, x) === x
            @test convert(NullableCategoricalArray{String, 2, DefaultRefType}, x) == x
            @test convert(NullableCategoricalArray{String, 2, UInt8}, x) == x

            @test convert(NullableCategoricalMatrix, x) === x
            @test convert(NullableCategoricalMatrix{String}, x) === x
            @test convert(NullableCategoricalMatrix{String, R}, x) === x
            @test convert(NullableCategoricalMatrix{String, DefaultRefType}, x) == x
            @test convert(NullableCategoricalMatrix{String, UInt8}, x) == x

            @test convert(NullableCategoricalArray, a) == x
            @test convert(NullableCategoricalArray{String}, a) == x
            @test convert(NullableCategoricalArray{String, 2, R}, a) == x
            @test convert(NullableCategoricalArray{String, 2, DefaultRefType}, a) == x
            @test convert(NullableCategoricalArray{String, 2, UInt8}, a) == x

            @test convert(NullableCategoricalMatrix, a) == x
            @test convert(NullableCategoricalMatrix{String}, a) == x
            @test convert(NullableCategoricalMatrix{String, R}, a) == x
            @test convert(NullableCategoricalMatrix{String, DefaultRefType}, a) == x
            @test convert(NullableCategoricalMatrix{String, UInt8}, a) == x

            @test NullableCategoricalArray{String}(a, ordered=ordered) == x
            @test NullableCategoricalArray{String, 2}(a, ordered=ordered) == x
            @test NullableCategoricalArray{String, 2}(a, ordered=ordered) == x
            @test NullableCategoricalArray{String, 2, R}(a, ordered=ordered) == x
            @test NullableCategoricalArray{String, 2, DefaultRefType}(a, ordered=ordered) == x
            @test NullableCategoricalArray{String, 2, UInt8}(a, ordered=ordered) == x

            @test NullableCategoricalMatrix(a, ordered=ordered) == x
            @test NullableCategoricalMatrix{String}(a, ordered=ordered) == x
            @test NullableCategoricalMatrix{String, R}(a, ordered=ordered) == x
            @test NullableCategoricalMatrix{String, DefaultRefType}(a, ordered=ordered) == x
            @test NullableCategoricalMatrix{String, UInt8}(a, ordered=ordered) == x

            for (y, R1, R2, comp) in ((a, DefaultRefType, UInt8, true),
                                      (a, DefaultRefType, DefaultRefType, false),
                                      (x, R, UInt8, true),
                                      (x, R, R, false))
            x2 = categorical(y, ordered=ordered)
            @test x2 == y
            if eltype(y) <: Nullable
                @test isa(x2, NullableCategoricalMatrix{String, R1})
            else
                @test isa(x2, CategoricalMatrix{String, R1})
            end
            @test isordered(x2) === ordered

            x2 = categorical(y, comp, ordered=ordered)
            @test x2 == y
            if eltype(y) <: Nullable
                @test isa(x2, NullableCategoricalMatrix{String, R2})
            else
                @test isa(x2, CategoricalMatrix{String, R2})
            end
            @test isordered(x2) === ordered
            end

            x2 = compress(x)
            @test x2 == x
            @test isa(x2, NullableCategoricalMatrix{String, UInt8})
            @test isordered(x2) === isordered(x)
            @test levels(x2) == levels(x)

            x2 = copy(x)
            @test x2 == x
            @test typeof(x2) === typeof(x)
            @test isordered(x2) === isordered(x)
            @test levels(x2) == levels(x)

            @test x[1] === Nullable(x.pool.valindex[1])
            @test x[2] === Nullable(x.pool.valindex[2])
            @test x[3] === Nullable(x.pool.valindex[2])
            @test x[4] === Nullable(x.pool.valindex[1])
            @test x[5] === Nullable(x.pool.valindex[3])
            @test x[6] === Nullable(x.pool.valindex[3])
            @test_throws BoundsError x[7]

            @test x[1,1] === Nullable(x.pool.valindex[1])
            @test x[2,1] === Nullable(x.pool.valindex[2])
            @test x[1,2] === Nullable(x.pool.valindex[2])
            @test x[2,2] === Nullable(x.pool.valindex[1])
            @test x[1,3] === Nullable(x.pool.valindex[3])
            @test x[2,3] === Nullable(x.pool.valindex[3])
            @test_throws BoundsError x[1,4]
            @test_throws BoundsError x[4,1]
            @test_throws BoundsError x[4,4]

            @test x[1:2,:] == x
            @test typeof(x[1:2,:]) === typeof(x)
            @test x[1:2,1] == Nullable{String}["a", "b"]
            @test typeof(x[1:2,1]) === NullableCategoricalVector{String, R}

            x[1] = "z"
            @test x[1] === Nullable(x.pool.valindex[4])
            @test x[2] === Nullable(x.pool.valindex[2])
            @test x[3] === Nullable(x.pool.valindex[2])
            @test x[4] === Nullable(x.pool.valindex[1])
            @test x[5] === Nullable(x.pool.valindex[3])
            @test x[6] === Nullable(x.pool.valindex[3])
            @test levels(x) == ["a", "b", "c", "z"]

            x[1,:] = "a"
            @test x[1] === Nullable(x.pool.valindex[1])
            @test x[2] === Nullable(x.pool.valindex[2])
            @test x[3] === Nullable(x.pool.valindex[1])
            @test x[4] === Nullable(x.pool.valindex[1])
            @test x[5] === Nullable(x.pool.valindex[1])
            @test x[6] === Nullable(x.pool.valindex[3])
            @test levels(x) == ["a", "b", "c", "z"]

            x[1,1:2] = "z"
            @test x[1] === Nullable(x.pool.valindex[4])
            @test x[2] === Nullable(x.pool.valindex[2])
            @test x[3] === Nullable(x.pool.valindex[4])
            @test x[4] === Nullable(x.pool.valindex[1])
            @test x[5] === Nullable(x.pool.valindex[1])
            @test x[6] === Nullable(x.pool.valindex[3])
            @test levels(x) == ["a", "b", "c", "z"]
        end

        # Matrix with null values
        for a in (Nullable{String}["a" Nullable() "c"; "b" "a" Nullable()],
                  NullableArray(Nullable{String}["a" Nullable() "c"; "b" "a" Nullable()]))
            x = NullableCategoricalMatrix{String, R}(a, ordered=ordered)

            @test x == a
            @test isordered(x) === ordered
            @test levels(x) == map(get, filter(x->!isnull(x), unique(a)))
            @test size(x) === (2, 3)
            @test length(x) === 6

            @test convert(NullableCategoricalArray, x) === x
            @test convert(NullableCategoricalArray{String}, x) === x
            @test convert(NullableCategoricalArray{String, 2}, x) === x
            @test convert(NullableCategoricalArray{String, 2, R}, x) === x
            @test convert(NullableCategoricalArray{String, 2, DefaultRefType}, x) == x
            @test convert(NullableCategoricalArray{String, 2, UInt8}, x) == x

            @test convert(NullableCategoricalMatrix, x) === x
            @test convert(NullableCategoricalMatrix{String}, x) === x
            @test convert(NullableCategoricalMatrix{String, R}, x) === x
            @test convert(NullableCategoricalMatrix{String, DefaultRefType}, x) == x
            @test convert(NullableCategoricalMatrix{String, UInt8}, x) == x

            @test convert(NullableCategoricalArray, a) == x
            @test convert(NullableCategoricalArray{String}, a) == x
            @test convert(NullableCategoricalArray{String, 2, R}, a) == x
            @test convert(NullableCategoricalArray{String, 2, DefaultRefType}, a) == x
            @test convert(NullableCategoricalArray{String, 2, UInt8}, a) == x

            @test convert(NullableCategoricalMatrix, a) == x
            @test convert(NullableCategoricalMatrix{String}, a) == x
            @test convert(NullableCategoricalMatrix{String, R}, a) == x
            @test convert(NullableCategoricalMatrix{String, DefaultRefType}, a) == x
            @test convert(NullableCategoricalMatrix{String, UInt8}, a) == x

            @test NullableCategoricalArray{String}(a, ordered=ordered) == x
            @test NullableCategoricalArray{String, 2}(a, ordered=ordered) == x
            @test NullableCategoricalArray{String, 2}(a, ordered=ordered) == x
            @test NullableCategoricalArray{String, 2, R}(a, ordered=ordered) == x
            @test NullableCategoricalArray{String, 2, DefaultRefType}(a, ordered=ordered) == x
            @test NullableCategoricalArray{String, 2, UInt8}(a, ordered=ordered) == x

            @test NullableCategoricalMatrix(a, ordered=ordered) == x
            @test NullableCategoricalMatrix{String}(a, ordered=ordered) == x
            @test NullableCategoricalMatrix{String, R}(a, ordered=ordered) == x
            @test NullableCategoricalMatrix{String, DefaultRefType}(a, ordered=ordered) == x
            @test NullableCategoricalMatrix{String, UInt8}(a, ordered=ordered) == x

            for (y, R1, R2, comp) in ((a, DefaultRefType, UInt8, true),
                                      (a, DefaultRefType, DefaultRefType, false),
                                      (x, R, UInt8, true),
                                      (x, R, R, false))
                x2 = categorical(y, ordered=ordered)
                @test x2 == y
                @test isa(x2, NullableCategoricalMatrix{String, R1})
                @test isordered(x2) === ordered

                x2 = categorical(y, comp, ordered=ordered)
                @test x2 == y
                @test isa(x2, NullableCategoricalMatrix{String, R2})
                @test isordered(x2) === ordered
            end

            x2 = compress(x)
            @test x2 == x
            @test isa(x2, NullableCategoricalMatrix{String, UInt8})
            @test isordered(x2) === isordered(x)
            @test levels(x2) == levels(x)

            x2 = copy(x)
            @test x2 == x
            @test typeof(x2) === typeof(x)
            @test isordered(x2) === isordered(x)
            @test levels(x2) == levels(x)

            @test x[1] === Nullable(x.pool.valindex[1])
            @test x[2] === Nullable(x.pool.valindex[2])
            @test x[3] === eltype(x)()
            @test x[4] === Nullable(x.pool.valindex[1])
            @test x[5] === Nullable(x.pool.valindex[3])
            @test x[6] === eltype(x)()
            @test_throws BoundsError x[7]

            @test x[1,1] === Nullable(x.pool.valindex[1])
            @test x[2,1] === Nullable(x.pool.valindex[2])
            @test x[1,2] === eltype(x)()
            @test x[2,2] === Nullable(x.pool.valindex[1])
            @test x[1,3] === Nullable(x.pool.valindex[3])
            @test x[2,3] === eltype(x)()
            @test_throws BoundsError x[1,4]
            @test_throws BoundsError x[4,1]
            @test_throws BoundsError x[4,4]

            x2 = x[1:2,:]
            @test typeof(x2) === typeof(x)
            @test x2 == x
            @test levels(x2) == levels(x)
            @test levels(x2) !== levels(x)
            @test isordered(x2) == isordered(x)

            x2 = x[:,[1, 3]]
            @test typeof(x2) === typeof(x)
            @test x2 == Nullable{String}["a" "c"; "b" Nullable()]
            @test levels(x2) == levels(x)
            @test levels(x2) !== levels(x)
            @test isordered(x2) == isordered(x)

            x2 = x[1:1,2]
            @test isa(x2, NullableCategoricalVector{String, R})
            @test x2 == [Nullable()]
            @test levels(x2) == levels(x)
            @test levels(x2) !== levels(x)
            @test isordered(x2) == isordered(x)

            x2 = x[1:0,:]
            @test typeof(x2) === typeof(x)
            @test size(x2) == (0,3)
            @test levels(x2) == levels(x)
            @test levels(x2) !== levels(x)
            @test isordered(x2) == isordered(x)

            @test_throws BoundsError x[1:4, :]
            @test_throws BoundsError x[1:1, -1:1]
            @test_throws BoundsError x[4, :]

            x[1] = "z"
            @test x[1] === Nullable(x.pool.valindex[4])
            @test x[2] === Nullable(x.pool.valindex[2])
            @test x[3] === eltype(x)()
            @test x[4] === Nullable(x.pool.valindex[1])
            @test x[5] === Nullable(x.pool.valindex[3])
            @test x[6] === eltype(x)()
            @test levels(x) == ["a", "b", "c", "z"]

            x[1,:] = "a"
            @test x[1] === Nullable(x.pool.valindex[1])
            @test x[2] === Nullable(x.pool.valindex[2])
            @test x[3] === Nullable(x.pool.valindex[1])
            @test x[4] === Nullable(x.pool.valindex[1])
            @test x[5] === Nullable(x.pool.valindex[1])
            @test x[6] === eltype(x)()
            @test levels(x) == ["a", "b", "c", "z"]

            x[1,1:2] = Nullable("z")
            @test x[1] === Nullable(x.pool.valindex[4])
            @test x[2] === Nullable(x.pool.valindex[2])
            @test x[3] === Nullable(x.pool.valindex[4])
            @test x[4] === Nullable(x.pool.valindex[1])
            @test x[5] === Nullable(x.pool.valindex[1])
            @test x[6] === eltype(x)()
            @test levels(x) == ["a", "b", "c", "z"]

            x[1] = Nullable()
            @test x[1] === eltype(x)()
            @test x[2] === Nullable(x.pool.valindex[2])
            @test x[3] === Nullable(x.pool.valindex[4])
            @test x[4] === Nullable(x.pool.valindex[1])
            @test x[5] === Nullable(x.pool.valindex[1])
            @test x[6] === eltype(x)()
            @test levels(x) == ["a", "b", "c", "z"]

            x[1,1:2] = Nullable()
            @test x[1] === eltype(x)()
            @test x[2] === Nullable(x.pool.valindex[2])
            @test x[3] === eltype(x)()
            @test x[4] === Nullable(x.pool.valindex[1])
            @test x[5] === Nullable(x.pool.valindex[1])
            @test x[6] === eltype(x)()
            @test levels(x) == ["a", "b", "c", "z"]

            x[:,2] = Nullable()
            @test x[1] === eltype(x)()
            @test x[2] === Nullable(x.pool.valindex[2])
            @test x[3] === eltype(x)()
            @test x[4] === eltype(x)()
            @test x[5] === Nullable(x.pool.valindex[1])
            @test x[6] === eltype(x)()
            @test levels(x) == ["a", "b", "c", "z"]

            x[1,2] = Nullable("a")
            @test x[1] === eltype(x)()
            @test x[2] === Nullable(x.pool.valindex[2])
            @test x[3] === Nullable(x.pool.valindex[1])
            @test x[4] === eltype(x)()
            @test x[5] === Nullable(x.pool.valindex[1])
            @test x[6] === eltype(x)()
            @test levels(x) == ["a", "b", "c", "z"]

            x[2,1] = Nullable()
            @test x[1] === eltype(x)()
            @test x[2] === eltype(x)()
            @test x[3] === Nullable(x.pool.valindex[1])
            @test x[4] === eltype(x)()
            @test x[5] === Nullable(x.pool.valindex[1])
            @test x[6] === eltype(x)()
            @test levels(x) == ["a", "b", "c", "z"]

            # Constructor with values plus missingness array
            x = NullableCategoricalArray(1:3, [true, false, true], ordered=ordered)
            @test x == Nullable{Int}[Nullable(), 2, Nullable()]
            @test isordered(x) == ordered
            @test levels(x) == [2]

            x = NullableCategoricalVector(1:3, [true, false, true], ordered=ordered)
            @test x == Nullable{Int}[Nullable(), 2, Nullable()]
            @test isordered(x) === ordered
            @test levels(x) == [2]

            x = NullableCategoricalMatrix([1 2; 3 4], [true false; false true],
                                          ordered=ordered)
            @test x == Nullable{Int}[Nullable() 2; 3 Nullable()]
            @test isordered(x) === ordered
            @test levels(x) == [2, 3]
        end

        # Uninitialized array
        v = Any[NullableCategoricalArray(2, ordered=ordered),
                NullableCategoricalArray{String}(2, ordered=ordered),
                NullableCategoricalArray{String, 1}(2, ordered=ordered),
                NullableCategoricalArray{String, 1, R}(2, ordered=ordered),
                NullableCategoricalVector(2, ordered=ordered),
                NullableCategoricalVector{String}(2, ordered=ordered),
                NullableCategoricalVector{String, R}(2, ordered=ordered),
                NullableCategoricalArray(2, 3, ordered=ordered),
                NullableCategoricalArray{String}(2, 3, ordered=ordered),
                NullableCategoricalArray{String, 2}(2, 3, ordered=ordered),
                NullableCategoricalArray{String, 2, R}(2, 3, ordered=ordered),
                NullableCategoricalMatrix(2, 3, ordered=ordered),
                NullableCategoricalMatrix{String}(2, 3, ordered=ordered),
                NullableCategoricalMatrix{String, R}(2, 3, ordered=ordered)]

        for x in v
            @test isordered(x) === ordered
            @test isnull(x[1])
            @test isnull(x[2])
            @test levels(x) == []

            x2 = compress(x)
            @test x2 == x
            @test isa(x2, NullableCategoricalArray{String, ndims(x), UInt8})
            @test isordered(x2) === isordered(x)
            @test levels(x2) == []

            x2 = copy(x)
            @test x2 == x
            @test typeof(x2) === typeof(x)
            @test isordered(x2) === isordered(x)
            @test levels(x2) == []

            x[1] = "c"
            @test x[1] === Nullable(x.pool.valindex[1])
            @test isnull(x[2])
            @test levels(x) == ["c"]

            x[1] = "a"
            @test x[1] === Nullable(x.pool.valindex[2])
            @test isnull(x[2])
            @test levels(x) == ["c", "a"]

            x[2] = Nullable()
            @test x[1] === Nullable(x.pool.valindex[2])
            @test x[2] === eltype(x)()
            @test levels(x) == ["c", "a"]

            x[1] = Nullable("b")
            @test x[1] === Nullable(x.pool.valindex[3])
            @test x[2] === eltype(x)()
            @test levels(x) == ["c", "a", "b"]
        end
    end
end

# Test vcat with nulls
ca1 = NullableCategoricalArray(["a", "b"], [false, true])
ca2 = NullableCategoricalArray(["b", "a"], [true, false])
r = vcat(ca1, ca2)
@test r == NullableCategoricalArray(["a", "", "", "a"], [false, true, true, false])
@test levels(r) == ["a"]
@test !isordered(r)
ordered!(ca1,true)
@test !isordered(vcat(ca1, ca2))
ordered!(ca2,true)
@test isordered(vcat(ca1, ca2))
ordered!(ca1,false)
@test !isordered(vcat(ca1, ca2))

# vcat with all nulls
ca1 = NullableCategoricalArray(["a", "b"], [false, true])
ca2 = NullableCategoricalArray(["a", "b"], [true, true])
r = vcat(ca1, ca2)
@test isequal(r, NullableCategoricalArray(["a", "", "", ""], [false, true, true, true]))
@test levels(r) == ["a"]
@test !isordered(r)

# vcat with all nulls preserves isordered
# needed for instance when expanding an array with nulls
# such as vcat of DataFrame with missing columns
ordered!(ca1, true)
@test isempty(levels(ca2))
r = vcat(ca1, ca2)
@test isordered(r)

# vcat with all empty array
ca1 = NullableCategoricalArray(0)
ca2 = NullableCategoricalArray(["a", "b"], [true, false])
r = vcat(ca1, ca2)
@test isequal(r, NullableCategoricalArray(["", "b"], [true, false]))
@test levels(r) == ["b"]
@test !isordered(r)

# vcat with all nulls and empty
ca1 = NullableCategoricalArray(0)
ca2 = NullableCategoricalArray(["a", "b"], [true, true])
r = vcat(ca1, ca2)
@test isequal(r, NullableCategoricalArray(["", ""], [true, true]))
@test levels(r) == String[]
@test !isordered(r)

ordered!(ca1, true)
@test isempty(levels(ca2))
r = vcat(ca1, ca2)
@test isordered(r)

ca1 = NullableCategoricalArray(["a", "b"], [false, true])
ca2 = NullableCategoricalArray{String}(2)
ordered!(ca1, true)
@test isempty(levels(ca2))
r = vcat(ca1, ca2)
@test isequal(r, NullableCategoricalArray(["a", "", "", ""], [false, true, true, true]))
@test isordered(r)

# Test unique() and levels()

x = NullableCategoricalArray(["Old", "Young", "Middle", Nullable(), "Young"])
@test levels(x) == ["Middle", "Old", "Young"]
@test unique(x) == NullableArray(["Middle", "Old", "Young", Nullable()])
@test levels!(x, ["Young", "Middle", "Old"]) === x
@test levels(x) == ["Young", "Middle", "Old"]
@test unique(x) == NullableArray(["Young", "Middle", "Old", Nullable()])
@test levels!(x, ["Young", "Middle", "Old", "Unused"]) === x
@test levels(x) == ["Young", "Middle", "Old", "Unused"]
@test unique(x) == NullableArray(["Young", "Middle", "Old", Nullable()])
@test levels!(x, ["Unused1", "Young", "Middle", "Old", "Unused2"]) === x
@test levels(x) == ["Unused1", "Young", "Middle", "Old", "Unused2"]
@test unique(x) == NullableArray(["Young", "Middle", "Old", Nullable()])

x = NullableCategoricalArray([Nullable{String}()])
@test isa(levels(x), Vector{String}) && isempty(levels(x))
@test unique(x) == NullableArray{String}(1)
@test levels!(x, ["Young", "Middle", "Old"]) === x
@test levels(x) == ["Young", "Middle", "Old"]
@test unique(x) == NullableArray{String}(1)

# To test short-circuit after 1000 elements
x = NullableCategoricalArray(repeat(1:1500, inner=10))
@test levels(x) == collect(1:1500)
@test unique(x) == NullableArray(1:1500)
@test levels!(x, [1600:-1:1; 2000]) === x
x[3] = Nullable()
@test levels(x) == [1600:-1:1; 2000]
@test unique(x) == NullableArray([1500:-1:3; 2; 1; Nullable()])

end


text:
#### Domain Decomposition Functions

# Provides a set of helper functions that operate on StateVec.

export grid_mean!, distribute!, diagnose_environment!, total_covariance!

"""
    grid_mean!(sv::StateVec, sv_a::StateVec, ϕ, a, grid::Grid)

Compute domain average of field(s) `ϕ`, using weights `a`,
in state vector `sv`, the grid `grid` respectively.

A domain-averaged variable ``⟨ϕ⟩`` is computed from

``⟨ϕ⟩ = Σ_i a_i \\overline{ϕ}_i``

Where variable ``\\overline{ϕ}_i`` represents ``ϕ`` decomposed across multiple
sub-domains, which are weighted by area fractions ``a_i``.

Note that `grid_mean!` is the inverse function of `distribute!`.
"""
function grid_mean!(sv::StateVec, sv_a::StateVec, a, var_names, grid::Grid)
    gm, en, ud, sd, al = allcombinations(sv)
    !(var_names isa Tuple) && (var_names = (var_names,))
    @inbounds for k in over_elems_real(grid)
        for ϕ in var_names
            sv[ϕ, k, gm] = 0
            @inbounds for i in sd
                sv[ϕ, k, gm] += sv[ϕ, k, i] * sv_a[a, k, i]
            end
        end
    end
    extrap_0th_order!(sv, var_names, grid, gm)
end

"""
    distribute!(sv::StateVec, grid::Grid, ϕ)

Distributes values in the state vector `sv`, from the grid-mean
domain to the sub-domains.

A domain-decomposed variable ``\\overline{ϕ}_i`` is computed from

``\\overline{ϕ}_i = ⟨ϕ⟩``

Where variable ``⟨ϕ⟩`` is the grid-mean variable, computed
across multiple sub-domains.

Note that `distribute!` is the inverse function of `grid_mean!`.
"""
function distribute!(sv::StateVec, grid::Grid, var_names)
    !(var_names isa Tuple) && (var_names = (var_names,))
    gm, en, ud, sd, al = allcombinations(sv)
    @inbounds for k in over_elems_real(grid)
        @inbounds for i in sd
            @inbounds for ϕ in var_names
                sv[ϕ, k, i] = sv[ϕ, k, gm]
            end
        end
    end
    @inbounds for i in sd
        extrap_0th_order!(sv, var_names, grid, i)
    end
end

"""
    diagnose_environment!(sv::StateVec, grid::Grid, names)

Diagnose environment values in the state vector `sv`, from the grid-mean
domain and updraft variables from the decomposition constraint.

Formulaically:

``ϕ_en = (⟨ϕ⟩-Σ_{i=ud} a_i ϕ_i)/a_en``
"""
function diagnose_environment!(sv::StateVec, grid::Grid, a, var_names)
    !(var_names isa Tuple) && (var_names = (var_names,))
    gm, en, ud, sd, al = allcombinations(sv)
    @inbounds for k in over_elems(grid)
        sv[a, k, en] = sv[a, k, gm] - sum([sv[a, k, i] for i in ud]...)
        a_en = sv[a, k, en]
        @inbounds for ϕ in var_names
            sv[ϕ, k, en] =
                (
                    sv[ϕ, k, gm] -
                    sum([sv[a, k, i] * sv[ϕ, k, i] for i in ud]...)
                ) / a_en
        end
    end
    extrap_0th_order!(sv, var_names, grid, en)
end

"""
    total_covariance!(dst::StateVec, src::StateVec, cv::StateVec, weights::StateVec,
                      dst_idxs, src_idxs, cv_idxs, a,
                      grid::Grid, decompose_ϕ_ψ::Function)

Computes the total covariance in state vector `dst`, given
 - `src` source state vector
 - `cv` state vector containing co-variances
 - `weights` state vector containing weights
 - `dst_idxs` indexes for destination state vector
 - `cv_idxs` indexes for state vector containing co-variances
 - `a` indexes for state vector containing weights
 - `grid` the grid
 - `decompose_ϕ_ψ` a function that receives the covariance index and
                   returns the indexes for each variable. For example:
                   `:ϕ_idx, :ψ_idx = decompose_ϕ_ψ(:cv_ϕ_ψ)`

A total covariance between variables ``ϕ`` and ``ψ`` is computed from

``⟨ϕ^*ψ^*⟩ = Σ_i a_i \\overline{ϕ_i'ψ_i'} + Σ_i Σ_j a_i a_j \\overline{ϕ}_i (\\overline{ψ}_i - \\overline{ψ}_j)``

Where variable ``\\overline{ϕ}_i`` represents ``ϕ`` decomposed across multiple
sub-domains, which are weighted by area fractions ``a_i``.
"""
function total_covariance!(
    aux::StateVec,
    sv::StateVec,
    cv::StateVec,
    tcv_ϕψ::Symbol,
    cv_ϕψ::Symbol,
    a::Symbol,
    grid::Grid,
    decompose_ϕ_ψ::Function,
)
    gm, en, ud, sd, al = allcombinations(sv)
    @inbounds for k in over_elems(grid)
        _ϕ, _ψ = decompose_ϕ_ψ(tcv_ϕψ)
        aux[tcv_ϕψ, k, gm] = 0
        @inbounds for i in sd
            aux[tcv_ϕψ, k, gm] += sv[a, k, i] * cv[cv_ϕψ, k, i]
            @inbounds for j in sd
                aux[tcv_ϕψ, k, gm] +=
                    sv[a, k, i] *
                    sv[a, k, j] *
                    sv[_ϕ, k, i] *
                    (sv[_ψ, k, i] - sv[_ψ, k, j])
            end
        end
    end
    extrap_0th_order!(aux, tcv_ϕψ, grid, gm)
end

ext: jl


text:
using Elly
using Test
using Random

function test_ugi()
    @info("test UserGroupInformation...")
    ugi = UserGroupInformation()
    @test !isempty(ugi.userinfo.realUser)
    iob = IOBuffer()
    show(iob, ugi)
    @test !isempty(take!(iob))

    ugi = UserGroupInformation(; proxy=true)
    @test !isempty(ugi.userinfo.realUser)
    @test !isempty(ugi.userinfo.effectiveUser)
    iob = IOBuffer()
    show(iob, ugi)
    @test !isempty(take!(iob))

    proxyuser = ugi.userinfo.realUser * "proxy"
    ugi = UserGroupInformation(; proxy=true, proxyuser=proxyuser)
    @test !isempty(ugi.userinfo.realUser)
    @test !isempty(ugi.userinfo.effectiveUser)
    @test ugi.userinfo.effectiveUser == proxyuser
    iob = IOBuffer()
    show(iob, ugi)
    @test !isempty(take!(iob))

    user = ENV["USER"]
    delete!(ENV, "USER")
    @test_throws Exception Elly.default_username()
    ENV["USER"] = user

    nothing
end

function test_hdfs(host="localhost", port=9000)
    limitedtestenv = (get(ENV, "CI", "false") == "true")

    ugi = UserGroupInformation()
    hdfsclnt = HDFSClient(host, port, ugi)

    exists(hdfsclnt, "/tmp") || mkdir(hdfsclnt, "/tmp")

    result = readdir(hdfsclnt, "/")
    @info("listing files in root folder", result)
    @test "tmp" in result

    result = hdfs_server_defaults(hdfsclnt)
    @info("server defaults", result)
    @test hdfs_default_block_size(hdfsclnt) > 0
    @test hdfs_default_replication(hdfsclnt) > 0

    result = hdfs_status(hdfsclnt)
    @info("file system status", result)

    result = hdfs_capacity(hdfsclnt)
    @info("hdfs capacity", result)
    @test result > 0

    result = isdir(hdfsclnt, "/tmp")
    @info("/tmp should be a directory", result)
    @test result

    result = du(hdfsclnt, "/")
    @info("du should be >= 0", result)
    @test result >= 0

    st = stat(hdfsclnt, "/tmp")
    @info("stat /tmp", st)
    @test st.name == "/tmp"

    @info("create a temporary dir...")
    t1 = time()
    cd(hdfsclnt, "/tmp")
    foo_dir = HDFSFile(hdfsclnt, "foo")
    mkdir(foo_dir, Elly.DEFAULT_FOLDER_MODE)
    @info("...done in $(time() - t1) secs")

    @info("create a temporary file...")
    cd(hdfsclnt, "foo")
    bar_file = HDFSFile(hdfsclnt, "bar")
    teststr = "hello world\n"
    nloops = 1000
    t1 = time()
    open(bar_file, "w") do f
        for idx in 1:nloops
            write(f, teststr)
        end
    end
    @info("...done in $(time() - t1) secs")
    expected_len = length(teststr)*nloops
    actual_len = filesize(bar_file)
    @info("verify file size", expected_len, actual_len=Int(actual_len))
    @test expected_len == actual_len

    @info("touch, move and delete file...")
    touch(bar_file)
    @test isfile(bar_file)
    st = stat(bar_file)
    @info("stat file", st)
    @test st.name == "bar"
    mv(bar_file, "/tmp/foo/bar2")
    bar2_file = HDFSFile("hdfs://$(host):$(port)/tmp/foo/bar2"; ugi=ugi)
    @test isfile(bar2_file)
    rm(bar2_file)

    @info("touch and delete a new file...")
    touch(bar_file)
    @test exists(bar_file)
    @test isfile(bar_file)
    rm(bar_file)

    @info("create a large file...")
    size_bytes = limitedtestenv ? (128 * 10 * 1000) : (128 * 1000 * 1000)
    nloops = limitedtestenv ? 2 : 5
    A = rand(UInt8, size_bytes)
    open(bar_file, "w"; blocksz=UInt64(size_bytes)) do f
        for idx in 1:nloops
            t1 = time()
            write(f, A)
            @info("...block written in $(time() - t1) secs")
        end
    end
    expected_size = sizeof(A) * nloops
    actual_size = filesize(bar_file)
    @info("verify file size", expected_size, actual_size)
    @test expected_size == actual_size

    @info("read and verify...")
    B = Array{UInt8}(undef, size_bytes)
    open(bar_file, "r") do f
        for idx in 1:nloops
            t1 = time()
            read!(f, B)
            @info("...block read in $(time() - t1) secs")
            @test A == B
        end
    end

    @info("read and verify with crc...")
    B = Array{UInt8}(undef, size_bytes)
    open(bar_file, "r"; crc=true) do f
        for idx in 1:nloops
            t1 = time()
            read!(f, B)
            @info("...block read in $(time() - t1) secs")
            @test A == B
        end
    end

    blocks = hdfs_blocks(bar_file)
    @info("file blocks", file=bar_file, blocks)

    @info("setting replication factor for $bar_file")
    @test hdfs_set_replication(bar_file, 2)

    cd(hdfsclnt, "/tmp/foo")
    NFILES = limitedtestenv ? 10 : 1000
    @info("create many ($NFILES) files...")
    for fidx in 1:NFILES
        bar_file = HDFSFile(hdfsclnt, "bar$fidx")
        open(bar_file, "w") do f
            for idx in 1:nloops
                write(f, teststr)
            end
        end
        ((fidx % 10) == 0) && @info("...created file #$fidx")
    end
    @info("reading directory...")
    allfiles = readdir(hdfsclnt, "/tmp/foo")
    @info("delete many ($NFILES) files...")
    for idx in 1:NFILES
        bar_file = HDFSFile(hdfsclnt, "bar$idx")
        rm(bar_file)
        ((idx % 10) == 0) && @info("...deleted file #$idx")
    end
    @test length(allfiles) >= NFILES

    @info("test renew lease")
    hdfs_renewlease(hdfsclnt)
    iob = IOBuffer()
    show(iob, hdfsclnt)
    @test !isempty(take!(iob))

    nothing
end


text:
#############################################################
function mylogpdf_Cov(x, μ, Σ)
#############################################################

    @assert(length(x) == length(μ) == size(Σ,1) == size(Σ,2))

    d = length(μ)

    L = chol(Hermitian(Σ))'

    TERM1 = -0.5 * 2.0*sum(log.(diag(L))) # consult eq. A.18 in RW
    TERM2 = -0.5 * d * log(2π)
    TERM3 = -0.5 * sum(( (L\(x-μ)).^2 ))

    return TERM1 + TERM2 + TERM3

end

#############################################################
function mylogpdf_scalCov(x, μ, σ2)
#############################################################

    @assert(length(x) == length(μ))

    d = length(μ)

    TERM1 = -0.5 * d * log(σ2)
    TERM2 = -0.5 * d * log(2π)
    TERM3 = -0.5 * sum((x-μ).^2)/σ2

    return TERM1 + TERM2 + TERM3

end


text:
@testset "Local quantities" begin

    # Test Dynamics
    T = Float64
    dt = 0.01
    model = DoubleIntegratorGame(p=3, d=2)
    x = rand(SVector{model.n,T})
    u = rand(SVector{model.m,T})
    x_dot = RobotDynamics.dynamics(model, x, u)
    x1 = x + dt*x_dot

    z = KnotPoint(x,u,dt)
    @test typeof(z) <: AbstractKnotPoint
    @test norm(dynamics_residual(model, z, x1), 1) < 1e-3
    @test (@ballocated dynamics_residual($model, $z, $x1)) == 0

    N = 10
    k = 5
    probsize = ProblemSize(N,model)
    pdtraj = PrimalDualTraj(probsize, dt)
    @test (@ballocated dynamics_residual($model, $pdtraj, $k)) == 0

    # Test dynamics Jacobian
    function ∇dynamics_easy(∇f, model, z)
        dt = z.dt
        x = state(z)
        u = control(z)
        function local_dyn_x(x)
            z = KnotPoint(x, u, dt)
            return discrete_dynamics(RK2, model, z)
        end
        function local_dyn_u(u)
            z = KnotPoint(x, u, dt)
            return discrete_dynamics(RK2, model, z)
        end
        ∇f[:,1:n] = ForwardDiff.jacobian(local_dyn_x, x)
        ∇f[:,n .+ (1:m)] = ForwardDiff.jacobian(local_dyn_u, u)
        return ∇f
    end

    T = Float64
    dt = 0.2
    model = UnicycleGame(p=3)
    x = rand(SVector{model.n,T})
    u = rand(SVector{model.m,T})
    z = KnotPoint(x,u,dt)

    n = model.n
    m = model.m
    ∇f1 = zeros(MMatrix{n,n+m,T,n*(n+m)})
    ∇f2 = zeros(MMatrix{n,n+m,T,n*(n+m)})
    ∇f3 = zeros(MMatrix{n,n+m,T,n*(n+m)})
    RobotDynamics.discrete_jacobian!(RK2,∇f1,model,z)
    ∇dynamics_easy(∇f2, model, z)
    ∇dynamics!(∇f3, model, z)
    @test norm(∇f1 - ∇f2, 1) < 1e-10
    @test norm(∇f1 - ∇f3, 1) < 1e-10
    @test (@ballocated ∇dynamics!($∇f3, $model, $z)) == 0
    @test (@ballocated ∇dynamics!($∇f3, $model, $pdtraj, $k)) == 0

end


text:
# Power cone constrained problems
using Test, LinearAlgebra, SparseArrays
using COSMO

# This test is precision agnostic and will be run with T precision
if !@isdefined(UnitTestFloats)
    UnitTestFloats = [Float64] #if not run in full test setup, just do it for one float type
end

@testset "Power cone problems" begin

for T in UnitTestFloats

    @testset "Power cone problems (T = $(T))" begin
        if precision(T) >= precision(Float64)
            @testset "Feasible" begin

                # solve the following power cone problem
                # max  x1^0.6 y^0.4 + x2^0.1
                # s.t. x1, y, x2 >= 0
                #      x1 + 2y  + 3x2 == 3
                # which is equivalent to
                # max z1 + z2
                # s.t. (x1, y, z1) in K_pow(0.6)
                #      (x2, 1, z2) in K_pow(0.1)
                #      x1 + 2y + 3x2 == 3

                # x = (x1, y, z1, x2, y2, z2)
                n = 6
                P = spzeros(T, n, n)
                q = zeros(T, 6)
                q[3] = q[6] = -one(T)

                # (x1, y, z1) in K_pow(0.6)
                A1 = spdiagm(0 => ones(T, 3))
                b1 = zeros(T, 3)
                cs1 = COSMO.Constraint(A1, b1, COSMO.PowerCone(T(0.6)), 6, 1:3)

                # (x2, 2, z2) in K_pow(0.1)
                A2 = spdiagm(0 => ones(T, 3))
                b2 = zeros(T, 3)
                cs2 = COSMO.Constraint(A2, b2, COSMO.PowerCone(T(0.1)),6, 4:6)

                # x1 + 2y + 3x2 == 3
                cs3 = COSMO.Constraint(T[1.0 2.0 0 3.0 0 0], T[-3.], COSMO.ZeroSet)

                # y2 == 1
                cs4 = COSMO.Constraint(T[0 0 0 0 1.0 0], T[-1.], COSMO.ZeroSet)

                model = COSMO.Model{T}()
                assemble!(model, P, q, [cs1; cs2; cs3; cs4], settings = COSMO.Settings{T}(max_iter = 5000))

                res = COSMO.optimize!(model)
                @test res.status == :Solved
                @test isapprox(res.obj_val, -1.8458, atol = 1e-3)

                # using JuMP
                # m = Model(with_optimizer(COSMO.Optimizer))
                # @variable(m, x1)
                # @variable(m, y)
                # @variable(m, z1)
                # @variable(m, x2)
                # @variable(m, y2 == 1.)
                # @variable(m, z2)

                # @objective(m, Max, z1 + z2)
                # @constraint(m, x1 + 2 * y + 3 * x2 == 3)
                # @constraint(m, [x1; y; z1] in MOI.PowerCone(0.6))
                # @constraint(m, [x2; y2; z2] in MOI.PowerCone(0.1))
                # status = JuMP.optimize!(m)

            end
    end
        # Run infeasibility tests only if precision(T) is high enough
        if precision(T) >= precision(Float64)
            @testset "Primal infeasible" begin
                # max  z
                # s.t. x^0.8 * y^0.2 >= z
                #       x = y == 1
                #       z = 2

                P = spzeros(T, 3, 3)
                q = T[0; 0; -1.0]

                cs1 = COSMO.Constraint(spdiagm(0 => ones(T, 3)), zeros(T, 3), COSMO.PowerCone{T}(0.8))
                cs2 = COSMO.Constraint(spdiagm(0 => ones(T, 3)), T[-1.; -1.; -2.], COSMO.ZeroSet)

                model = COSMO.Model{T}()
                assemble!(model, P, q, [cs1; cs2])

                res = COSMO.optimize!(model)

                @test res.status == :Primal_infeasible
            end

            @testset "Dual infeasible" begin
                # min z
                # s.t. x^0.8 * y^0.2 >= z

                P = spzeros(T, 3, 3)
                q = T[0; 0; 1.0]

                cs1 = COSMO.Constraint(spdiagm(0 => ones(T, 3)), zeros(T, 3), COSMO.PowerCone(T(0.8)))

                model = COSMO.Model{T}()
                assemble!(model, P, q, cs1)

                res = COSMO.optimize!(model)
                @test res.status == :Dual_infeasible
            end
        end

    end

    @testset "Dual power cone problems" begin
        @testset "Feasible" begin
            # max  z
            # s.t. (x/0.8)^0.8 * (y/0.2)^0.2 >= z
            #       x == 0.8,  y == 0.2
            P = spzeros(T, 3, 3)
            q = T[0.; 0.; -1]

            # (x, y, z) in K^*_pow(0.8)
            A1 = spdiagm( 0 => ones(T, 3))
            b1 = zeros(T, 3)
            cs1 = COSMO.Constraint(A1, b1, COSMO.DualPowerCone(T(0.8)))

            # x == 0.8, y == 0.2
            cs2 = COSMO.Constraint(T[1. 0 0; 0 1 0], T[-0.8; -0.2], COSMO.ZeroSet)

            model = COSMO.Model{T}()
            assemble!(model, P, q, [cs1; cs2])

            res = COSMO.optimize!(model)
            @test res.status == :Solved
            @test isapprox(res.obj_val, -one(T), atol=1e-3)

        end
    end
end
end
nothing


text:
module Rendering

export render, RenderSettings

using CuArrays, CUDAnative

using CthulhuVision.Random
using CthulhuVision.Math
using CthulhuVision.Light
using CthulhuVision.Image
using CthulhuVision.Camera
using CthulhuVision.Materials
using CthulhuVision.Triangles
using CthulhuVision.Scenes
using CthulhuVision.BVH
using CthulhuVision.Worlds

@inline function color(r::Ray, acceleration::BVHAcceleration, world::World, settings::SceneSettings, rng::UniformRNG) :: RGB
    maxbounces = 50

    attenuation = RGB(1.0f0, 1.0f0, 1.0f0)
    result = RGB(0.0f0, 0.0f0, 0.0f0)
    ray = r

    for i = 1:maxbounces
        rec = hitacceleration(acceleration, 0.001f0, typemax(Float32), ray, world)

        if rec.ishit
            scattered = scatter(ray, rec, rng)
            if !scattered.isreflected
                break
            end

            attenuation = attenuation * scattered.attenuation
            result += rec.material.emission * attenuation
            ray = scattered.ray
        else
            result += attenuation * settings.ambientemission
            break
        end
    end

    result
end

@inline function makeprng() :: UniformRNG
    index = threadIdx().x + (blockIdx().x - 1) * blockDim().x
    uniformfromindex(index)
end

struct RenderSettings
    nsamples::UInt32
end

function gpurender(a, camera, width, height, scenesettings, rendersettings::RenderSettings, bvhnodes, vertexes, meshtriangles, instances)
    y = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    x = (blockIdx().y - 1) * blockDim().y + threadIdx().y

    rng = makeprng()

    acceleration = BVHAcceleration(bvhnodes)
    world = World(vertexes, meshtriangles, instances)

    if x <= width && y <= height
        col = RGB(0.0f0, 0.0f0, 0.0f0)

        for s = 1:rendersettings.nsamples
            dx = next(rng)
            dy = next(rng)
            u = Float32((x + dx) / width)
            v = Float32((y + dy) / height)
            ray = getray(camera, u, v, rng)
            col = col + color(ray, acceleration, world, scenesettings, rng)
        end

        col /= Float32(rendersettings.nsamples)

        @inbounds a[y, x] = col
    end

    return nothing
end

function render(image::PPM, camera, scene::Scene, rendersettings::RenderSettings)
    CuArrays.@allowscalar false

    rng = uniformfromindex(0)
    bvhnodes = buildacceleration(scene.hitables, rng)

    pixels = CuArray{RGB}(undef, image.dimension.height, image.dimension.width)
    bvhnodes_d = CuArray{BVHNode}(bvhnodes)
    vertexes_d = CuArray{Vector3}(scene.vertexes)
    meshtriangles_d = CuArray{MeshTriangle}(scene.meshtriangles)
    instances_d = CuArray{MeshInstance}(scene.instances)

    blocks = ceil(Int, image.dimension.height / 16), ceil(Int, image.dimension.width / 16)
    threads = (16, 16)

    CuArrays.@sync begin
        @cuda threads=threads blocks=blocks gpurender(pixels, camera, image.dimension.width, image.dimension.height, scene.settings, rendersettings, bvhnodes_d, vertexes_d, meshtriangles_d, instances_d)
    end

    cpuarray = Array{RGB, 2}(pixels)
    for y = 0:image.dimension.height-1, x = 0:image.dimension.width-1
        px = Pixel(x, y)
        c = cpuarray[y + 1, x + 1]

        pixel(image, px, c)
    end
end

end


text:
# bradi-tests.jl -
#
# Tests for BRADI algorithm.
#

module BraDiTests

using Test
import OptimPackNextGen: BraDi

# Simple parabola.  To be minimized over [-1,2].
parabola(x) = x*x

# Brent's 5th function.  To be minimized over [-10,10].
brent5(x) = (x - sin(x))*exp(-x*x)

# Michalewicz's 1st function.  To be minimized over [-1,2].
michalewicz1(x) = x*sin(10.0*x)

# Michalewicz's 2nd function.  To be maximized over [0,pi].
function michalewicz2(x)
    s = 0.0
    a = sin(x)
    b = x*x/pi
    for i in 1:10
        s += a*(sin(b*i)^20)
    end
    return s
end

function runtests(; quiet::Bool=false)
    tol = sqrt(eps(Float64))
    @testset "BraDi algorithm" begin
        quiet || println("\n# Simple parabola:")
        (xbest, fbest) = BraDi.minimize(parabola, range(-1, stop=2, length=2))
        println("x = $xbest, f(x) = $fbest")
        @test xbest ≈ 0.0 atol=tol

        quiet || println("\n# Brent's 5th function:")
        (xbest, fbest) = BraDi.minimize(brent5, range(-10, stop=10, length=5))
        quiet || println("x = $xbest, f(x) = $fbest")
        @test xbest ≈ -1.1951366418407416 rtol=tol

        quiet || println("\n# Michalewicz's 1st function:")
        (xbest, fbest) = BraDi.minimize(michalewicz1, range(-1, stop=2, length=21))
        quiet || println("x = $xbest, f(x) = $fbest")
        @test xbest ≈ 1.7336377815999973 rtol=tol

        quiet || println("\n# Michalewicz's 2nd function:")
        (xbest, fbest) = BraDi.maximize(michalewicz2, range(0, stop=pi, length=60))
        quiet || println("x = $xbest, f(x) = $fbest")
        @test xbest ≈ 2.2208651539586493 rtol=tol
    end
end
runtests()

end # module


text:
module V
const NAMES = [
    "M"
    "E"
    "CD"
    "CE"
    "R"
    "RP"
    "RE"
]

for (idx,name) in enumerate(NAMES)
    eval(Meta.parse("const $name = $idx"))
end

const NUM = length(NAMES)
end  # module

max_stars_repo_name: himoto/Bifurcation

max_issues_repo_name: himoto/Bifurcation

max_forks_repo_name: himoto/Bifurcation


text:
# ps3000_stop.jl

export ps3000_stop

function ps3000_stop(handle:: Int16)
  ps_status = ccall((:ps3000_stop, ps3000driver), Int16,(Int16,),handle)
  if ps_status == 0
    error("ps3000 Error: Handle Not Valid")
  end
  return nothing
end


text:
# # Scattering: Mie Phase Function Tutorial
#
# ### Introduction
# In the following tutorial, we will walk through how to compute phase function matrices from
# Mie theory.

# Along the way, we will also calculate greek coefficients (which are essentially the
# legendre decomposition of the phase matrix components), which will be required for radiative
# transfer calculations.

# ---

# ### Performing the phase function calculations

# First let's use the required packages

using vSmartMOM.Scattering
using Distributions
using Plots
#----------------------------------------------------------------------------

# Now, we define the aerosol size distribution and properties. We only support univariate
# aerosols for now, but will add multivariate aerosols soon.

μ     = 0.3;             ## Log-normal median radius [μm]
σ     = 2.0;             ## Log-normal stddev of radius
nᵣ    = 1.3;             ## Real part of refractive index
nᵢ    = 0.0;             ## Imag part of refractive index (sign changed, use only + here)
r_max = 30.0;            ## Maximum radius [μm]
nquad_radius  = 2500;    ## Number of quadrature points for integrating of size dist.

## Create a Size Distribution (using Julia's Distributions package)
size_distribution = LogNormal(log(μ), log(σ));

#----------------------------------------------------------------------------

# Create the aerosol
aero = Aerosol(size_distribution, nᵣ, nᵢ)
#----------------------------------------------------------------------------

# Now, we define some scattering and truncation properties:

λ = 0.55;                             ## Incident wavelength [μm]
polarization_type = Stokes_IQUV();    ## Polarization type

## Trunction length for legendre terms
l_max = 20;

## Exclusion angle for forward peak (in fitting procedure) `[degrees]`
Δ_angle = 2;

## Truncation type
truncation_type   = δBGE(l_max, Δ_angle);
#----------------------------------------------------------------------------

## Create a Mie model, using the Siewert method NAI2
model_NAI2 = make_mie_model(NAI2(), aero, λ, polarization_type, truncation_type, r_max, nquad_radius);
#----------------------------------------------------------------------------

## Compute aerosol optical properties:
aerosol_optics_NAI2 = compute_aerosol_optical_properties(model_NAI2);
#----------------------------------------------------------------------------

# ### Plotting the greek coefficients
# (These are basically giving us the legendre decomposition of the phase matrix components)

using Parameters
@unpack α,β,γ,δ,ϵ,ζ = aerosol_optics_NAI2.greek_coefs;
p1 = plot(α,  title="α");
p2 = plot(β,  title="β");
p3 = plot(γ,  title="γ");
p4 = plot(δ,  title="δ");
p5 = plot(ϵ,  title="ϵ");
p6 = plot(ζ,  title="ζ");
plot(p1, p2, p3, p4, p5, p6, layout=(3, 2), legend=false)
xlims!(0,100)
#----------------------------------------------------------------------------

# ### Reconstructing the Phase Functions from the greek coefficients

using FastGaussQuadrature
μ_quad, w_μ = gausslegendre(1000)
scattering_matrix = Scattering.reconstruct_phase(aerosol_optics_NAI2.greek_coefs, μ_quad);
@unpack f₁₁, f₁₂, f₂₂, f₃₃, f₃₄, f₄₄ = scattering_matrix
#----------------------------------------------------------------------------

# Plot only phase function for I (f₁₁) and the I -> Q transition in the phase matrix (f₁₂) for the Stokes Vector [I,Q,U,V]

p1 = plot(μ_quad, f₁₁, yscale=:log10, title="f₁₁")
p2 = plot(μ_quad, f₁₂ ./ f₁₁,  title="f₁₂/f₁₁")

plot(p1, p2, layout=(2, 1), legend=false)
xlabel!("cos(Θ)")
#----------------------------------------------------------------------------

p1 = plot([acos.(μ_quad); -reverse(acos.(μ_quad))], log10.([f₁₁ ; reverse(f₁₁)]), proj=:polar, yscale=:log10, title="f₁₁", lims=(-3,4.2), yaxis=false)
p2 = plot([acos.(μ_quad); -reverse(acos.(μ_quad))], [f₁₂ ./ f₁₁ ; reverse(f₁₂ ./ f₁₁)], proj=:polar, title="f₁₂/f₁₁", lims=(0,0.6), yaxis=false)

plot(p1, p2, layout=(1, 2), legend=false)
xlabel!("Θ")

#----------------------------------------------------------------------------

anim = @animate for r = 0.03:0.05:4.3
    @show r
    local size_distribution = LogNormal(log(r), log(σ))
    ## Create the aerosol
    local aero       = Aerosol(size_distribution, nᵣ, nᵢ)
    local model_NAI2 = make_mie_model(NAI2(), aero, λ, polarization_type, truncation_type, r_max, nquad_radius)
    local aerosol_optics_NAI2 = compute_aerosol_optical_properties(model_NAI2);
    local scattering_matrix = Scattering.reconstruct_phase(aerosol_optics_NAI2.greek_coefs, μ_quad);
    @unpack f₁₁, f₁₂, f₂₂, f₃₃, f₃₄, f₄₄ = scattering_matrix
    ## @show f₁₁[1]

    p1 = plot(μ_quad, f₁₁, yscale=:log10, title="f₁₁", label="r(μm)=$r")
    ylims!(1e-3, 1e3)
    p2 = plot(μ_quad, f₁₂ ./ f₁₁,  title="f₁₂/f₁₁", label="Q/I")
    ylims!(-1.1, 1.1)
    plot(p1, p2, layout=(2, 1))

    p3 = plot([acos.(μ_quad); -reverse(acos.(μ_quad))], log10.([f₁₁ ; reverse(f₁₁)]), proj=:polar, yscale=:log10, title="f₁₁, r = $(r)μm", lims=(-3,4.2), yaxis=false)
    p4 = plot([acos.(μ_quad); -reverse(acos.(μ_quad))], [f₁₂ ./ f₁₁ ; reverse(f₁₂ ./ f₁₁)], proj=:polar, title="f₁₂/f₁₁, r = $(r)μm", lims=(0,1.0), yaxis=false)

    plot(plot(p1, p2, layout=(2,1), legend=false), plot(p3, p4, layout=(1,2), legend=false), layout=(2, 1), legend=false)

    plot!(size=(600,800))

end
gif(anim, fps = 5)


text:
using Unitful, Statistics, StaticArrays, Rotations, CoordinateTransformations

function getdata(x)
    track = Track(x[:track])
    feeder = point(get(x, :feeder, x[:track].data[1,1:2]))
    pellet = pointcollection(get(x, :pellet, missing), x[:track].data[1,3])
    feeder, track, pellet
end
abstract type DungMethod end
struct ClosedNest <: DungMethod
    feeder::Point
    nest::Point
    track::Track
    pellet::PointCollection
end
function ClosedNest(x)
    feeder, track, pellet = getdata(x.data)
    nest = point(x.data[:nest])
    ClosedNest(feeder, nest, track, pellet)
end
struct Transfer <: DungMethod
    feeder::Point
    track::Track
    pellet::PointCollection
    south::Point
    north::Point
    nest2feeder::Float64
    azimuth::Float64
end
function _getvalueunit(txt, default)
    m = match(r"^(\d+)\s*(\D*)", txt)
    d, txtu = m.captures
    u = try
        _u = getfield(Unitful, Symbol(txtu))
        @assert _u isa Unitful.FreeUnits
        _u
    catch
        default
    end
    Float64(ustrip(uconvert(default, parse(Int, d)*u)))
end
function Transfer(x)
    feeder, track, pellet = getdata(x.data)
    south = point(x.data[:south])
    north = point(x.data[:north])
    nest2feeder = _getvalueunit(x.metadata.setup[:nest2feeder], u"cm")
    azimuth = _getvalueunit(x.metadata.setup[:azimuth], u"°")
    Transfer(feeder, track, pellet, south, north, nest2feeder, azimuth)
end
struct TransferNest <: DungMethod
    feeder::Point
    track::Track
    pellet::PointCollection
    south::Point
    north::Point
    nest2feeder::Float64
    azimuth::Float64
    originalnest::Point
end
function TransferNest(x)
    y = Transfer(x)
    originalnest = point(x.data[:originalnest])
    TransferNest(y.feeder, y.track, y.pellet, y.south, y.north, y.nest2feeder, y.azimuth, originalnest)
end
Transfer(x::TransferNest) = Transfer(x.feeder, x.track, x.pellet, x.south, x.north, x.nest2feeder, x.azimuth)
struct TransferNestBelen <: DungMethod
    feeder::Point
    track::Track
    pellet::PointCollection
    southbefore::Point
    northbefore::Point
    feederbefore::Point
    nestbefore::Point
    south::Point
    north::Point
end
function TransferNestBelen(x)
    feeder, track, pellet = getdata(x.data)
    southbefore = point(x.data[:southbefore])
    northbefore = point(x.data[:northbefore])
    feederbefore = point(x.data[:feederbefore])
    nestbefore = point(x.data[:nestbefore])
    south = point(x.data[:south])
    north = point(x.data[:north])
    TransferNestBelen(feeder, track, pellet, rightdowninitial, leftdowninitial, rightupinitial, leftupinitial, rightdownfinal, leftdownfinal, rightupfinal, leftupfinal)
end
struct DawaySandpaper <: DungMethod
    feeder::Point
    nest::Point
    track::Track
    pellet::PointCollection
    rightdowninitial::Point
    leftdowninitial::Point
    rightupinitial::Point
    leftupinitial::Point
    rightdownfinal::Point
    leftdownfinal::Point
    rightupfinal::Point
    leftupfinal::Point
end
function DawaySandpaper(x)
    feeder, track, pellet = getdata(x.data)
    nest = point(x.data[:nest])
    rightdowninitial = point(x.data[:rightdowninitial])
    leftdowninitial = point(x.data[:leftdowninitial])
    rightupinitial = point(x.data[:rightupinitial])
    leftupinitial = point(x.data[:leftupinitial])
    rightdownfinal = point(x.data[:rightdownfinal])
    leftdownfinal = point(x.data[:leftdownfinal])
    rightupfinal = point(x.data[:rightupfinal])
    leftupfinal = point(x.data[:leftupfinal])
    DawaySandpaper(feeder, nest, track, pellet, rightdowninitial, leftdowninitial, rightupinitial, leftupinitial, rightdownfinal, leftdownfinal, rightupfinal, leftupfinal)
end
struct DawayNest <: DungMethod
    feeder::Point
    nest::Point
    track::Track
    pellet::PointCollection
    pickup::Point
end
function DawayNest(x)
    feeder, track, pellet = getdata(x.data)
    nest = point(x.data[:nest])
    pickup = point(x.data[:pickup])
    DawayNest(feeder, nest, track, pellet, pickup)
end
struct Daway <: DungMethod
    feeder::Point
    nest::Point
    track::Track
    pellet::PointCollection
    initialfeeder::Point
end
function Daway(x)
    feeder, track, pellet = getdata(x.data)
    nest = point(x.data[:nest])
    initialfeeder = point(x.data[:initialfeeder])
    Daway(feeder, nest, track, pellet, initialfeeder)
end
struct Displacement <: DungMethod
    feeder::Point
    nest::Point
    track::Track
    pellet::PointCollection
    pickup::Point
    dropoff::Point
end
function Displacement(x)
    feeder, track, pellet = getdata(x.data)
    nest = point(x.data[:nest])
    pickup = point(x.data[:pickup])
    dropoff = point(x.data[:dropoff])
    Displacement(feeder, nest, track, pellet, pickup, dropoff)
end

######################### DungMethod methods ###########

DungMethod(x, displace_location, displace_direction::Missing, transfer::Missing, nest_coverage::Missing) = error("unidentified experimental setup")
DungMethod(x, displace_location::Missing, displace_direction::Missing, transfer::Missing, nest_coverage) = ClosedNest(x)
function DungMethod(x, displace_location::Missing, displace_direction::Missing, transfer, nest_coverage)
    if x.metadata.setup[:person] == "belen"
        TransferNestBelen(x)
    else
        if transfer == "back"
            TransferNest(x)
        elseif transfer ∈ ("far", "50m")
            Transfer(x)
        else
            error("unidentified experimental setup")
        end
    end
end
function DungMethod(x, displace_location, displace_direction, transfer::Missing, nest_coverage)
    if displace_location == "feeder"
        if x.metadata.setup[:person] == "belen"
            DawaySandpaper(x)
        else
            Daway(x)
        end
    elseif displace_location == "nest"
        DawayNest(x)
    elseif displace_location == "halfway"
        Displacement(x)
    else
        error("unidentified experimental setup")
    end
end

######################### Common methods ###########

#=mutable struct Common{N}
    feeder::Point
    nest::Point
    track::Track
    pellet::PointCollection
    originalnest::N
end=#
common(x::ClosedNest) = Common(x.feeder, x.nest, x.track, x.pellet, x.nest)
# nest(x::Common) = x.nest
# turning(x::Common) = x.track.homing[end]
# originalnest(x::Common) = x.originalnest

function common(x::TransferNestBelen)
    v = x.northbefore - x.southbefore
    u = x.nestbefore - x.feederbefore
    azimuth = atan(v[2], v[1]) - atan(u[2], u[1])
    nest2feeder = norm(x.nestbefore - x.feederbefore)
    v = x.north - x.south
    α = atan(v[2], v[1]) + azimuth
    u = Point(cos(α), sin(α))
    nest = x.feeder + u*nest2feeder
    Common(x.feeder, nest, x.track, x.pellet, missing)
end

function common(x::Transfer)
    v = x.north - x.south
    α = atan(v[2], v[1]) + x.azimuth - π
    u = Point(cos(α), sin(α))
    nest = x.feeder + u*x.nest2feeder
    Common(x.feeder, nest, x.track, x.pellet, missing)
end

function common(x::TransferNest)
    y = common(Transfer(x))
    Common(y.feeder, y.nest, y.track, y.pellet, x.originalnest)
end

function common(x::DawaySandpaper)
    originalnest = x.nest
    initial = mean(getproperty(x, k) for k in [:rightdowninitial, :leftdowninitial, :rightupinitial, :leftupinitial])
    final = mean(getproperty(x, k) for k in [:rightdownfinal, :leftdownfinal, :rightupfinal, :leftupfinal])
    v = final - initial
    nest = originalnest + v
    _feeder = x.feeder
    feeder = _feeder + v
    Common(feeder, nest, x.track, x.pellet, originalnest)
end

function common(x::DawayNest)
    originalnest = x.nest
    feeder = x.feeder
    v = originalnest - x.pickup
    nest = feeder + v
    Common(feeder, nest, x.track, x.pellet, originalnest)
end

function common(x::Daway)
    originalnest = x.nest
    v = x.feeder - x.initialfeeder
    nest = originalnest + v
    Common(x.feeder, nest, x.track, x.pellet, originalnest)
end

function common(x::Displacement)
    v = x.dropoff - x.pickup
    originalnest = x.nest
    nest = originalnest + v
    Common(x.feeder, nest, x.track, x.pellet, originalnest)
end
######################### END ######################

common(x) = common(DungMethod(x, get(x.metadata.setup, :displace_location, missing), get(x.metadata.setup, :displace_direction, missing), get(x.metadata.setup, :transfer, missing), get(x.metadata.setup, :nest_coverage, missing)))

function get_rotation(nest, feeder)
    v = feeder - nest
    α = -atan(v[2], v[1])# - π/2
    rot = LinearMap(Angle2d(α))
    rot, α
end

(p::LinearMap{Angle2d{Float64}})(x::Missing) = missing

# LinearMap{Angle2d{Float64}}(x::Missing) = missing
function rotate!(x)
    rot, α = get_rotation(x.nest, x.feeder)
    x.nest = rot(x.nest)
    x.feeder = rot(x.feeder)
    x.track.coords .= rot.(x.track.coords)
    # x.track.direction .+= α
    x.pellet.xy .= rot.(x.pellet.xy)
    x.originalnest = rot(x.originalnest)
end

function rotate2!(x)
    rot, α = get_rotation(c, x.feeder)
    x.nest = rot(x.nest)
    x.feeder = rot(x.feeder)
    x.track.coords .= rot.(x.track.coords)
    # x.track.direction .+= α
    x.pellet.xy .= rot.(x.pellet.xy)
    x.originalnest = rot(x.originalnest)
end

function center2!(x)
    x.nest -= c
    x.feeder -= c
    for i in eachindex(x.track.coords)
        x.track.coords[i] -= c
    end
    for i in eachindex(x.pellet)
        x.pellet[i].xy -= c
    end
    if !ismissing(x.originalnest)
        x.originalnest -= c
    end
end