# This file is a part of Julia. License is MIT: https://julialang.org/license

module RangesTest

using Test
using Dates

using InteractiveUtils: subtypes

let
    for T in (Dates.Date, Dates.DateTime)
        f1 = T(2014); l1 = T(2013, 12, 31)
        f2 = T(2014); l2 = T(2014)
        f3 = T(-2000); l3 = T(2000)
        f4 = typemin(T); l4 = typemax(T)

        for P in subtypes(Dates.DatePeriod)
            for pos_step in (P(1), P(2), P(50), P(2048), P(10000))
                # empty range
                dr = f1:pos_step:l1
                len = length(dr)
                @test len == 0
                @test isa(len, Int64)
                @test isempty(dr)
                @test first(dr) == f1
                @test last(dr) < f1
                @test length([i for i in dr]) == 0
                @test_throws ArgumentError minimum(dr)
                @test_throws ArgumentError maximum(dr)
                @test_throws BoundsError dr[1]
                @test findall(in(dr), dr) == Int64[]
                @test [dr;] == T[]
                @test isempty(reverse(dr))
                @test length(reverse(dr)) == 0
                @test first(reverse(dr)) < f1
                @test last(reverse(dr)) >= f1
                @test issorted(dr)
                @test sortperm(dr) == 1:1:0
                @test !(f1 in dr)
                @test !(l1 in dr)
                @test !(f1 - pos_step in dr)
                @test !(l1 + pos_step in dr)
                @test dr == []
                @test hash(dr) == hash([])

                for (f, l) in ((f2, l2), (f3, l3), (f4, l4))
                    dr = f:pos_step:l
                    len = length(dr)
                    @test len > 0
                    @test isa(len, Int64)
                    @test !isempty(dr)
                    @test first(dr) == f
                    @test last(dr) <= l
                    @test minimum(dr) == first(dr)
                    @test maximum(dr) == last(dr)
                    @test dr[1] == f
                    @test dr[end] <= l
                    @test iterate(dr) == (first(dr), (length(dr), 1))

                    if len < 10000
                        dr1 = [i for i in dr]
                        @test length(dr1) == len
                        @test findall(in(dr), dr) == [1:len;]
                        @test length([dr;]) == len
                        @test dr == dr1
                        @test hash(dr) == hash(dr1)
                    end
                    @test !isempty(reverse(dr))
                    @test length(reverse(dr)) == len
                    @test last(reverse(dr)) == f
                    @test issorted(dr)
                    @test f in dr

                end
            end
            for neg_step in (P(-1), P(-2), P(-50), P(-2048), P(-10000))
                # empty range
                dr = l1:neg_step:f1
                len = length(dr)
                @test len == 0
                @test isa(len, Int64)
                @test isempty(dr)
                @test first(dr) == l1
                @test last(dr) > l1
                @test length([i for i in dr]) == 0
                @test_throws ArgumentError minimum(dr)
                @test_throws ArgumentError maximum(dr)
                @test_throws BoundsError dr[1]
                @test findall(in(dr), dr) == Int64[]
                @test [dr;] == T[]
                @test isempty(reverse(dr))
                @test length(reverse(dr)) == 0
                @test first(reverse(dr)) > l1
                @test last(reverse(dr)) <= l1
                @test issorted(dr)
                @test sortperm(dr) == 1:1:0
                @test !(l1 in dr)
                @test !(l1 in dr)
                @test !(l1 - neg_step in dr)
                @test !(l1 + neg_step in dr)
                @test dr == []
                @test hash(dr) == hash([])

                for (f, l) in ((f2, l2), (f3, l3), (f4, l4))
                    dr = l:neg_step:f
                    len = length(dr)
                    @test len > 0
                    @test isa(len, Int64)
                    @test !isempty(dr)
                    @test first(dr) == l
                    @test last(dr) >= f
                    @test minimum(dr) == last(dr)
                    @test maximum(dr) == first(dr)
                    @test dr[1] == l
                    @test dr[end] >= f
                    @test iterate(dr) == (first(dr), (length(dr), 1))

                    if len < 10000
                        dr1 = [i for i in dr]
                        @test length(dr1) == len
                        @test findall(in(dr), dr) == [1:len;]
                        @test length([dr;]) == len
                        @test dr == dr1
                        @test hash(dr) == hash(dr1)
                    end
                    @test !isempty(reverse(dr))
                    @test length(reverse(dr)) == len
                    @test issorted(dr) == (len <= 1)
                    @test l in dr
                end
            end
        end
        if T == Dates.DateTime
            for P in subtypes(Dates.TimePeriod)
                P in (Dates.Microsecond, Dates.Nanosecond) && continue
                for pos_step in (P(1), P(2), P(50), P(2048), P(10000))
                    # empty range
                    dr = f1:pos_step:l1
                    len = length(dr)
                    @test len == 0
                    @test isa(len, Int64)
                    @test isempty(dr)
                    @test first(dr) == f1
                    @test last(dr) < f1
                    @test length([i for i in dr]) == 0
                    @test_throws ArgumentError minimum(dr)
                    @test_throws ArgumentError maximum(dr)
                    @test_throws BoundsError dr[1]
                    @test findall(in(dr), dr) == Int64[]
                    @test [dr;] == T[]
                    @test isempty(reverse(dr))
                    @test length(reverse(dr)) == 0
                    @test first(reverse(dr)) < f1
                    @test last(reverse(dr)) >= f1
                    @test issorted(dr)
                    @test sortperm(dr) == 1:1:0
                    @test !(f1 in dr)
                    @test !(l1 in dr)
                    @test !(f1 - pos_step in dr)
                    @test !(l1 + pos_step in dr)
                    @test dr == []
                    @test hash(dr) == hash([])

                    for (f, l) in ((f2, l2), (f3, l3), (f4, l4))
                        dr = f:pos_step:l
                        len = length(dr)
                        @test len > 0
                        @test isa(len, Int64)
                        @test !isempty(dr)
                        @test first(dr) == f
                        @test last(dr) <= l
                        @test minimum(dr) == first(dr)
                        @test maximum(dr) == last(dr)
                        @test dr[1] == f
                        @test dr[end] <= l
                        @test iterate(dr) == (first(dr), (length(dr), 1))

                        if len < 10000
                            dr1 = [i for i in dr]
                            @test length(dr1) == len
                            @test findall(in(dr), dr) == [1:len;]
                            @test length([dr;]) == len
                            @test dr == dr1
                            @test hash(dr) == hash(dr1)
                        end
                        @test !isempty(reverse(dr))
                        @test length(reverse(dr)) == len
                        @test last(reverse(dr)) == f
                        @test issorted(dr)
                        @test f in dr

                    end
                end
                for neg_step in (P(-1), P(-2), P(-50), P(-2048), P(-10000))
                    # empty range
                    dr = l1:neg_step:f1
                    len = length(dr)
                    @test len == 0
                    @test isa(len, Int64)
                    @test isempty(dr)
                    @test first(dr) == l1
                    @test last(dr) > l1
                    @test length([i for i in dr]) == 0
                    @test_throws ArgumentError minimum(dr)
                    @test_throws ArgumentError maximum(dr)
                    @test_throws BoundsError dr[1]
                    @test findall(in(dr), dr) == Int64[]
                    @test [dr;] == T[]
                    @test isempty(reverse(dr))
                    @test length(reverse(dr)) == 0
                    @test first(reverse(dr)) > l1
                    @test last(reverse(dr)) <= l1
                    @test issorted(dr)
                    @test sortperm(dr) == 1:1:0
                    @test !(l1 in dr)
                    @test !(l1 in dr)
                    @test !(l1 - neg_step in dr)
                    @test !(l1 + neg_step in dr)
                    @test dr == []
                    @test hash(dr) == hash([])

                    for (f, l) in ((f2, l2), (f3, l3), (f4, l4))
                        dr = l:neg_step:f
                        len = length(dr)
                        @test len > 0
                        @test isa(len, Int64)
                        @test !isempty(dr)
                        @test first(dr) == l
                        @test last(dr) >= f
                        @test minimum(dr) == last(dr)
                        @test maximum(dr) == first(dr)
                        @test dr[1] == l
                        @test dr[end] >= f
                        @test iterate(dr) == (first(dr), (length(dr), 1))

                        if len < 10000
                            dr1 = [i for i in dr]
                            @test length(dr1) == len
                            @test findall(in(dr), dr) == [1:len;]
                            @test length([dr;]) == len
                            @test dr == dr1
                            @test hash(dr) == hash(dr1)
                        end
                        @test !isempty(reverse(dr))
                        @test length(reverse(dr)) == len
                        @test last(reverse(dr)) >= l
                        @test issorted(dr) == (len <= 1)
                        @test l in dr

                    end
                end
            end
        end
    end
end

# Dates are physical units, and ranges should require an explicit step.
# See #19896 and https://discourse.julialang.org/t/type-restriction-on-unitrange/6557/12
if VERSION >= v"1.0.0"
    @test_throws MethodError Dates.DateTime(2013, 1, 1):Dates.DateTime(2013, 2, 1)
end

# All the range representations we want to test
# Date ranges
dr  = Dates.DateTime(2013, 1, 1):Dates.Day(1):Dates.DateTime(2013, 2, 1)
dr1 = Dates.DateTime(2013, 1, 1):Dates.Day(1):Dates.DateTime(2013, 1, 1)
dr2 = Dates.DateTime(2013, 1, 1):Dates.Day(1):Dates.DateTime(2012, 2, 1) # empty range
dr3 = Dates.DateTime(2013, 1, 1):Dates.Day(-1):Dates.DateTime(2012) # negative step
# Big ranges
dr4 = Dates.DateTime(0):Dates.Day(1):Dates.DateTime(20000, 1, 1)
dr5 = Dates.DateTime(0):Dates.Day(1):Dates.DateTime(200000, 1, 1)
dr6 = Dates.DateTime(0):Dates.Day(1):Dates.DateTime(2000000, 1, 1)
dr7 = Dates.DateTime(0):Dates.Day(1):Dates.DateTime(20000000, 1, 1)
dr8 = Dates.DateTime(0):Dates.Day(1):Dates.DateTime(200000000, 1, 1)
dr9 = typemin(Dates.DateTime):Dates.Day(1):typemax(Dates.DateTime)
# Other steps
dr10 = typemax(Dates.DateTime):Dates.Day(-1):typemin(Dates.DateTime)
dr11 = typemin(Dates.DateTime):Dates.Week(1):typemax(Dates.DateTime)

dr12 = typemin(Dates.DateTime):Dates.Month(1):typemax(Dates.DateTime)
dr13 = typemin(Dates.DateTime):Dates.Year(1):typemax(Dates.DateTime)

dr14 = typemin(Dates.DateTime):Dates.Week(10):typemax(Dates.DateTime)
dr15 = typemin(Dates.DateTime):Dates.Month(100):typemax(Dates.DateTime)
dr16 = typemin(Dates.DateTime):Dates.Year(1000):typemax(Dates.DateTime)
dr17 = typemax(Dates.DateTime):Dates.Week(-10000):typemin(Dates.DateTime)
dr18 = typemax(Dates.DateTime):Dates.Month(-100000):typemin(Dates.DateTime)
dr19 = typemax(Dates.DateTime):Dates.Year(-1000000):typemin(Dates.DateTime)
dr20 = typemin(Dates.DateTime):Dates.Day(2):typemax(Dates.DateTime)

drs = Any[dr, dr1, dr2, dr3, dr4, dr5, dr6, dr7, dr8, dr9, dr10,
          dr11, dr12, dr13, dr14, dr15, dr16, dr17, dr18, dr19, dr20]
drs2 = map(x->Dates.Date(first(x)):step(x):Dates.Date(last(x)), drs)

@test map(length, drs) == map(x->size(x)[1], drs)
@test map(length, drs) == map(x->length(Dates.Date(first(x)):step(x):Dates.Date(last(x))), drs)
@test map(length, drs) == map(x->length(reverse(x)), drs)
@test all(x->findall(in(x), x)==[1:length(x);], drs[1:4])
@test isempty(dr2)
@test all(x->reverse(x) == range(last(x), step=-step(x), length=length(x)), drs)
@test all(x->minimum(x) == (step(x) < zero(step(x)) ? last(x) : first(x)), drs[4:end])
@test all(x->maximum(x) == (step(x) < zero(step(x)) ? first(x) : last(x)), drs[4:end])
@test all(drs[1:3]) do dd
    for (i, d) in enumerate(dd)
        @test d == (first(dd) + Dates.Day(i - 1))
    end
    true
end
@test_throws MethodError dr .+ 1
a = Dates.DateTime(2013, 1, 1)
b = Dates.DateTime(2013, 2, 1)
@test map!(x->x + Dates.Day(1), Vector{Dates.DateTime}(undef, 32), dr) == [(a + Dates.Day(1)):Dates.Day(1):(b + Dates.Day(1));]
@test map(x->x + Dates.Day(1), dr) == [(a + Dates.Day(1)):Dates.Day(1):(b + Dates.Day(1));]

@test map(x->a in x, drs[1:4]) == [true, true, false, true]
@test a in dr
@test b in dr
@test Dates.DateTime(2013, 1, 3) in dr
@test Dates.DateTime(2013, 1, 15) in dr
@test Dates.DateTime(2013, 1, 26) in dr
@test !(Dates.DateTime(2012, 1, 1) in dr)

@test all(x->sort(x) == (step(x) < zero(step(x)) ? reverse(x) : x), drs)
@test all(x->step(x) < zero(step(x)) ? issorted(reverse(x)) : issorted(x), drs)

@test length(b:Dates.Day(-1):a) == 32
@test length(b:Dates.Day(1):a) == 0
@test length(b:Dates.Day(1):a) == 0
@test length(a:Dates.Day(2):b) == 16
@test last(a:Dates.Day(2):b) == Dates.DateTime(2013, 1, 31)
@test length(a:Dates.Day(7):b) == 5
@test last(a:Dates.Day(7):b) == Dates.DateTime(2013, 1, 29)
@test length(a:Dates.Day(32):b) == 1
@test last(a:Dates.Day(32):b) == Dates.DateTime(2013, 1, 1)
@test (a:Dates.Day(1):b)[1] == Dates.DateTime(2013, 1, 1)
@test (a:Dates.Day(1):b)[2] == Dates.DateTime(2013, 1, 2)
@test (a:Dates.Day(1):b)[7] == Dates.DateTime(2013, 1, 7)
@test (a:Dates.Day(1):b)[end] == b
@test first(a:Dates.Day(1):Dates.DateTime(20000, 1, 1)) == a
@test first(a:Dates.Day(1):Dates.DateTime(200000, 1, 1)) == a
@test first(a:Dates.Day(1):Dates.DateTime(2000000, 1, 1)) == a
@test first(a:Dates.Day(1):Dates.DateTime(20000000, 1, 1)) == a
@test first(a:Dates.Day(1):Dates.DateTime(200000000, 1, 1)) == a
@test first(a:Dates.Day(1):typemax(Dates.DateTime)) == a
@test first(typemin(Dates.DateTime):Dates.Day(1):typemax(Dates.DateTime)) == typemin(Dates.DateTime)

# Date ranges
dr  = Dates.Date(2013, 1, 1):Dates.Day(1):Dates.Date(2013, 2, 1)
dr1 = Dates.Date(2013, 1, 1):Dates.Day(1):Dates.Date(2013, 1, 1)
dr2 = Dates.Date(2013, 1, 1):Dates.Day(1):Dates.Date(2012, 2, 1) # empty range
dr3 = Dates.Date(2013, 1, 1):Dates.Day(-1):Dates.Date(2012, 1, 1) # negative step
# Big ranges
dr4 = Dates.Date(0):Dates.Day(1):Dates.Date(20000, 1, 1)
dr5 = Dates.Date(0):Dates.Day(1):Dates.Date(200000, 1, 1)
dr6 = Dates.Date(0):Dates.Day(1):Dates.Date(2000000, 1, 1)
dr7 = Dates.Date(0):Dates.Day(1):Dates.Date(20000000, 1, 1)
dr8 = Dates.Date(0):Dates.Day(1):Dates.Date(200000000, 1, 1)
dr9 = typemin(Dates.Date):Dates.Day(1):typemax(Dates.Date)
# Other steps
dr10 = typemax(Dates.Date):Dates.Day(-1):typemin(Dates.Date)
dr11 = typemin(Dates.Date):Dates.Week(1):typemax(Dates.Date)
dr12 = typemin(Dates.Date):Dates.Month(1):typemax(Dates.Date)
dr13 = typemin(Dates.Date):Dates.Year(1):typemax(Dates.Date)
dr14 = typemin(Dates.Date):Dates.Week(10):typemax(Dates.Date)
dr15 = typemin(Dates.Date):Dates.Month(100):typemax(Dates.Date)
dr16 = typemin(Dates.Date):Dates.Year(1000):typemax(Dates.Date)
dr17 = typemax(Dates.Date):Dates.Week(-10000):typemin(Dates.Date)
dr18 = typemax(Dates.Date):Dates.Month(-100000):typemin(Dates.Date)
dr19 = typemax(Dates.Date):Dates.Year(-1000000):typemin(Dates.Date)
dr20 = typemin(Dates.Date):Dates.Day(2):typemax(Dates.Date)

drs = Any[dr, dr1, dr2, dr3, dr4, dr5, dr6, dr7, dr8, dr9, dr10,
          dr11, dr12, dr13, dr14, dr15, dr16, dr17, dr18, dr19, dr20]

@test map(length, drs) == map(x->size(x)[1], drs)
@test all(x->findall(in(x), x) == [1:length(x);], drs[1:4])
@test isempty(dr2)
@test all(x->reverse(x) == last(x): - step(x):first(x), drs)
@test all(x->minimum(x) == (step(x) < zero(step(x)) ? last(x) : first(x)), drs[4:end])
@test all(x->maximum(x) == (step(x) < zero(step(x)) ? first(x) : last(x)), drs[4:end])
@test all(drs[1:3]) do dd
    for (i, d) in enumerate(dd)
        @test d == (first(dd) + Dates.Day(i - 1))
    end
    true
end
@test_throws MethodError dr .+ 1
a = Dates.Date(2013, 1, 1)
b = Dates.Date(2013, 2, 1)
@test map!(x->x + Dates.Day(1), Vector{Dates.Date}(undef, 32), dr) == [(a + Dates.Day(1)):Dates.Day(1):(b + Dates.Day(1));]
@test map(x->x + Dates.Day(1), dr) == [(a + Dates.Day(1)):Dates.Day(1):(b + Dates.Day(1));]

@test map(x->a in x, drs[1:4]) == [true, true, false, true]
@test a in dr
@test b in dr
@test Dates.Date(2013, 1, 3) in dr
@test Dates.Date(2013, 1, 15) in dr
@test Dates.Date(2013, 1, 26) in dr
@test !(Dates.Date(2012, 1, 1) in dr)

@test all(x->sort(x) == (step(x) < zero(step(x)) ? reverse(x) : x), drs)
@test all(x->step(x) < zero(step(x)) ? issorted(reverse(x)) : issorted(x), drs)

@test length(b:Dates.Day(-1):a) == 32
@test length(b:Dates.Day(1):a) == 0
@test length(a:Dates.Day(2):b) == 16
@test last(a:Dates.Day(2):b) == Dates.Date(2013, 1, 31)
@test length(a:Dates.Day(7):b) == 5
@test last(a:Dates.Day(7):b) == Dates.Date(2013, 1, 29)
@test length(a:Dates.Day(32):b) == 1
@test last(a:Dates.Day(32):b) == Dates.Date(2013, 1, 1)
@test (a:Dates.Day(1):b)[1] == Dates.Date(2013, 1, 1)
@test (a:Dates.Day(1):b)[2] == Dates.Date(2013, 1, 2)
@test (a:Dates.Day(1):b)[7] == Dates.Date(2013, 1, 7)
@test (a:Dates.Day(1):b)[end] == b
@test first(a:Dates.Day(1):Dates.Date(20000, 1, 1)) == a
@test first(a:Dates.Day(1):Dates.Date(200000, 1, 1)) == a
@test first(a:Dates.Day(1):Dates.Date(2000000, 1, 1)) == a
@test first(a:Dates.Day(1):Dates.Date(20000000, 1, 1)) == a
@test first(a:Dates.Day(1):Dates.Date(200000000, 1, 1)) == a
@test first(a:Dates.Day(1):typemax(Dates.Date)) == a
@test first(typemin(Dates.Date):Dates.Day(1):typemax(Dates.Date)) == typemin(Dates.Date)

@test length(typemin(Dates.Date):Dates.Week(1):typemax(Dates.Date)) == 26351950414948059
# Big Month/Year ranges
@test length(typemin(Dates.Date):Dates.Month(1):typemax(Dates.Date)) == 6060531933867600
@test length(typemin(Dates.Date):Dates.Year(1):typemax(Dates.Date)) == 505044327822300
@test length(typemin(Dates.DateTime):Dates.Month(1):typemax(Dates.DateTime)) == 3507324288
@test length(typemin(Dates.DateTime):Dates.Year(1):typemax(Dates.DateTime)) == 292277024

@test length(typemin(Dates.DateTime):Dates.Week(1):typemax(Dates.DateTime)) == 15250284420
@test length(typemin(Dates.DateTime):Dates.Day(1):typemax(Dates.DateTime)) == 106751990938
@test length(typemin(Dates.DateTime):Dates.Hour(1):typemax(Dates.DateTime)) == 2562047782512
@test length(typemin(Dates.DateTime):Dates.Minute(1):typemax(Dates.DateTime)) == 153722866950720
@test length(typemin(Dates.DateTime):Dates.Second(1):typemax(Dates.DateTime)) == 9223372017043200
@test length(typemin(DateTime):Dates.Millisecond(1):typemax(DateTime)) == 9223372017043199001

c = Dates.Date(2013, 6, 1)
@test length(a:Dates.Month(1):c) == 6
@test [a:Dates.Month(1):c;] == [a + Dates.Month(1)*i for i in 0:5]
@test [a:Dates.Month(2):Dates.Date(2013, 1, 2);] == [a]
@test [c:Dates.Month(-1):a;] == reverse([a:Dates.Month(1):c;])

@test length(range(Date(2000), step=Dates.Day(1), length=366)) == 366
let n = 100000
    local a = Dates.Date(2000)
    for i = 1:n
        @test length(range(a, step=Dates.Day(1), length=i)) == i
    end
    return a + Dates.Day(n)
end

@test typeof(step(Dates.DateTime(2000):Dates.Day(1):Dates.DateTime(2001))) == Dates.Day

a = Dates.Date(2013, 1, 1)
b = Dates.Date(2013, 2, 1)
d = Dates.Date(2020, 1, 1)
@test length(a:Dates.Year(1):d) == 8
@test first(a:Dates.Year(1):d) == a
@test last(a:Dates.Year(1):d) == d
@test length(a:Dates.Month(12):d) == 8
@test first(a:Dates.Month(12):d) == a
@test last(a:Dates.Month(12):d) == d
@test length(a:Dates.Week(52):d) == 8
@test first(a:Dates.Week(52):d) == a
@test last(a:Dates.Week(52):d) == Dates.Date(2019, 12, 24)
@test length(a:Dates.Day(365):d) == 8
@test first(a:Dates.Day(365):d) == a
@test last(a:Dates.Day(365):d) == Dates.Date(2019, 12, 31)

a = Dates.Date(2013, 1, 1)
b = Dates.Date(2013, 2, 1)
@test length(a:Dates.Year(1):Dates.Date(2020, 2, 1)) == 8
@test length(a:Dates.Year(1):Dates.Date(2020, 6, 1)) == 8
@test length(a:Dates.Year(1):Dates.Date(2020, 11, 1)) == 8
@test length(a:Dates.Year(1):Dates.Date(2020, 12, 31)) == 8
@test length(a:Dates.Year(1):Dates.Date(2021, 1, 1)) == 9
@test length(Dates.Date(2000):Dates.Year(-10):Dates.Date(1900)) == 11
@test length(Dates.Date(2000, 6, 23):Dates.Year(-10):Dates.Date(1900, 2, 28)) == 11
@test length(Dates.Date(2000, 1, 1):Dates.Year(1):Dates.Date(2000, 2, 1)) == 1

let n = 100000
    local a, b
    a= b = Dates.Date(0)
    for i = 1:n
        @test length(a:Dates.Year(1):b) == i
        b += Dates.Year(1)
    end
end

let n = 10000,
    a = Dates.Date(1985, 12, 5),
    b = Dates.Date(1986, 12, 27),
    c = Dates.DateTime(1985, 12, 5),
    d = Dates.DateTime(1986, 12, 27)
    for i = 1:n
        @test length(a:Dates.Month(1):b) == 13
        @test length(a:Dates.Year(1):b) == 2
        @test length(c:Dates.Month(1):d) == 13
        @test length(c:Dates.Year(1):d) == 2
        a += Dates.Day(1)
        b += Dates.Day(1)
    end
end

let n = 100000
    local a, b
    a = b = Dates.Date(2000)
    for i = 1:n
        @test length(a:Dates.Month(1):b) == i
        b += Dates.Month(1)
    end
end

@test length(Dates.Year(1):Dates.Year(1):Dates.Year(10)) == 10
@test length(Dates.Year(10):Dates.Year(-1):Dates.Year(1)) == 10
@test length(Dates.Year(10):Dates.Year(-2):Dates.Year(1)) == 5
@test_throws OverflowError length(typemin(Dates.Year):Dates.Year(1):typemax(Dates.Year))
@test_throws MethodError Dates.Date(0):Dates.DateTime(2000)
@test_throws MethodError Dates.Date(0):Dates.Year(10)
@test length(range(Dates.Date(2000), step=Dates.Day(1), length=366)) == 366
@test last(range(Dates.Date(2000), step=Dates.Day(1), length=366)) == Dates.Date(2000, 12, 31)
@test last(range(Dates.Date(2001), step=Dates.Day(1), length=365)) == Dates.Date(2001, 12, 31)
@test last(range(Dates.Date(2000), step=Dates.Day(1), length=367)) == last(range(Dates.Date(2000), step=Dates.Month(12), length=2)) == last(range(Dates.Date(2000), step=Dates.Year(1), length=2))
@test last(range(Dates.DateTime(2000), step=Dates.Day(366), length=2)) == last(range(Dates.DateTime(2000), step=Dates.Hour(8784), length=2))

# Issue 5
lastdaysofmonth = [Dates.Date(2014, i, Dates.daysinmonth(2014, i)) for i=1:12]
@test [Date(2014, 1, 31):Dates.Month(1):Date(2015);] == lastdaysofmonth

# Range addition/subtraction:
let d = Dates.Day(1)
    @test (Dates.Date(2000):d:Dates.Date(2001)) + d == (Dates.Date(2000) + d:d:Dates.Date(2001) + d)
    @test (Dates.Date(2000):d:Dates.Date(2001)) - d == (Dates.Date(2000) - d:d:Dates.Date(2001) - d)
end

# Time ranges
dr  = Dates.Time(23, 1, 1):Dates.Second(1):Dates.Time(23, 2, 1)
dr1 = Dates.Time(23, 1, 1):Dates.Second(1):Dates.Time(23, 1, 1)
dr2 = Dates.Time(23, 1, 1):Dates.Second(1):Dates.Time(22, 2, 1) # empty range
dr3 = Dates.Time(23, 1, 1):Dates.Minute(-1):Dates.Time(22, 1, 1) # negative step
# Big ranges
dr8 = typemin(Dates.Time):Dates.Second(1):typemax(Dates.Time)
dr9 = typemin(Dates.Time):Dates.Nanosecond(1):typemax(Dates.Time)
# Other steps
dr10 = typemax(Dates.Time):Dates.Microsecond(-1):typemin(Dates.Time)
dr11 = typemin(Dates.Time):Dates.Millisecond(1):typemax(Dates.Time)
dr12 = typemin(Dates.Time):Dates.Minute(1):typemax(Dates.Time)
dr13 = typemin(Dates.Time):Dates.Hour(1):typemax(Dates.Time)
dr14 = typemin(Dates.Time):Dates.Millisecond(10):typemax(Dates.Time)
dr15 = typemin(Dates.Time):Dates.Minute(100):typemax(Dates.Time)
dr16 = typemin(Dates.Time):Dates.Hour(1000):typemax(Dates.Time)
dr17 = typemax(Dates.Time):Dates.Millisecond(-10000):typemin(Dates.Time)
dr18 = typemax(Dates.Time):Dates.Minute(-100):typemin(Dates.Time)
dr19 = typemax(Dates.Time):Dates.Hour(-10):typemin(Dates.Time)
dr20 = typemin(Dates.Time):Dates.Microsecond(2):typemax(Dates.Time)

drs = Any[dr, dr1, dr2, dr3, dr8, dr9, dr10,
          dr11, dr12, dr13, dr14, dr15, dr16, dr17, dr18, dr19, dr20]

@test map(length, drs) == map(x->size(x)[1], drs)
@test all(x->findall(in(x), x) == [1:length(x);], drs[1:4])
@test isempty(dr2)
@test all(x->reverse(x) == last(x): - step(x):first(x), drs)
@test all(x->minimum(x) == (step(x) < zero(step(x)) ? last(x) : first(x)), drs[4:end])
@test all(x->maximum(x) == (step(x) < zero(step(x)) ? first(x) : last(x)), drs[4:end])
@test_throws MethodError dr .+ 1

a = Dates.Time(23, 1, 1)
@test map(x->a in x, drs[1:4]) == [true, true, false, true]
@test a in dr

@test all(x->sort(x) == (step(x) < zero(step(x)) ? reverse(x) : x), drs)
@test all(x->step(x) < zero(step(x)) ? issorted(reverse(x)) : issorted(x), drs)

@test !(1 in Date(2017, 01, 01):Dates.Day(1):Date(2017, 01, 05))
@test !(Complex(1, 0) in Date(2017, 01, 01):Dates.Day(1):Date(2017, 01, 05))
@test !(π in Date(2017, 01, 01):Dates.Day(1):Date(2017, 01, 05))
@test !("a" in Date(2017, 01, 01):Dates.Day(1):Date(2017, 01, 05))

@test hash(Any[Date("2018-1-03"), Date("2018-1-04"), Date("2018-1-05")]) ==
      hash([Date("2018-1-03"), Date("2018-1-04"), Date("2018-1-05")]) ==
      hash(Date("2018-1-03"):Day(1):Date("2018-1-05"))

end


for TYP in (:ReverseOrientation,:Reverse)
    WRAP = Meta.parse(string(TYP)*"Wrapper")
    @eval begin
        abstract type $TYP{T} <: Operator{T} end

        struct $WRAP{OS,T} <: Operator{T}
            op::OS
        end

        $WRAP(op::Operator) = $WRAP{typeof(op),eltype(op)}(op)
        convert(::Type{Operator{T}},op::$TYP) where {T} = $TYP{T}()
        convert(::Type{Operator{T}},op::$WRAP) where {T} = $WRAP(Operator{T}(op.op))::Operator{T}

        @wrapper $WRAP
    end
end


mutable struct MaxNNode
    snakes::Vector{Snake}
    food::Food
    height::Int
    width::Int
    depth::Int
    action::Union{Direction, Nothing}
    winner::Union{Int, Nothing}
    dead::Set{Int}
end

function iterative_deepening_maxn(board::Boardstate, simulation_time::Float64, height::Int, width::Int, you_index::Int)
        simulation_time = Millisecond(floor(Int,simulation_time))
        snakes = board.snakes
        food = board.food
        root = MaxNNode(snakes, food, height, width, 0, nothing, nothing, Set{Int}())
        action_nodes = possible_successors(root, you_index)
        best_action = nothing
        depth = 1

        times = []
        start = now()
        last_time = Millisecond(0)
        values = Tuple{Vector{Float64}, MaxNNode}[]
        while now() - start + last_time < simulation_time
            start_run = now()
            empty!(values)

            for a_node in action_nodes
                if now() - start < simulation_time
                    vector = maxn(a_node, depth, next_player(you_index, length(a_node.snakes)) , length(a_node.snakes), you_index)
                    push!(values, (vector, a_node))
                end
            end
            depth += 1
            if length(values) == length(action_nodes)
                sort!(values, by= x -> x[1][you_index], rev=true)
                best_action = values[1][2].action
                action_nodes = [v[2] for v in values]
            end
            last_time = now() - start_run
            push!(times, last_time)
        end
        println("$depth")
        println("times $times")
        println("Time used" ,now() - start)
        # println([(v[1], v[2].action) for v in values])
        return best_action
end

function maxn(node::MaxNNode, depth::Int, player::Int, num_players::Int, you_index::Int)
    if depth == 0 || node.winner !== nothing || length(node.dead) == length(node.snakes)
        return score(node)
    end

    while player in node.dead
        player = next_player(player, num_players)
    end
    next_p = next_player(player, num_players)
    max_vector = fill(-Inf, 4)
    next_nodes = possible_successors(node, player)
    for successor in next_nodes
        vector = maxn(successor, depth - 1, next_p, length(node.snakes), you_index)
        if vector[player] > max_vector[player]
            max_vector = vector
        end
    end

    return max_vector
end

function possible_successors(node::MaxNNode, player::Int)
    if player in node.dead
        return []
    end
    successors = MaxNNode[]
    snake = node.snakes[player]
    for a in possible_actions(snake)
        new_snek, new_food = deepcopy(snake), deepcopy(node.food)
        move_snake!(new_snek, new_food, a)
        if !check_bounds(get_head(new_snek), node.height, node.width)
            continue
        end
        snakes = deepcopy(node.snakes)
        snakes[player] = new_snek
        dead_for_a = deepcopy(node.dead)
        winner = check_winner(node, snakes, dead_for_a)
        n = MaxNNode(snakes, new_food, node.height, node.width, node.depth + 1, a, winner, dead_for_a)
        push!(successors, n)
    end
    return successors
end

function check_winner(node::MaxNNode, snakes::Vector{Snake}, dead::Set{Int})

    for (p, snek) in enumerate(snakes)
        # starve
        if snek.health == 0
            push!(dead, p)
        end

        for (o, opponent) in enumerate(snakes)
            if o == p
                continue
            end

            if get_head(snek) in opponent.body
                # head2head
                if get_head(snek) == get_head(opponent)
                    # RAM
                    if length(snek.body) > length(opponent.body)
                        push!(dead, o)
                    # draw
                    elseif length(snek.body) == length(opponent.body)
                        push!(dead, o)
                        push!(dead, p)
                    else
                        # lost (head2head) collision with opponent
                        push!(dead, p)
                    end
                else
                    # ran into opponent
                    push!(dead, p)
                end
            end

            # collision with self
            if get_head(snek) in get_body_no_head(snek)
                push!(dead, p)
            end
            if !check_bounds(get_head(snek), node.height, node.width)
                push!(dead, p)
            end

            if length(dead) == length(snakes) - 1
                winset = setdiff(Set(collect(1:length(snakes))), dead)
                return minimum(winset)
            end
        end
    end
end

function score(node::MaxNNode)
    result = zeros(length(node.snakes))
    health = [snek.health for snek in node.snakes]
    lens = [length(snek.body) for snek in node.snakes]
    flood_values = maxN_floodfill(node.snakes, node.width, node.height)
    for i in 1:length(node.snakes)
        score = 1 * health[i] + 4 * lens[i] + 1 * flood_values[i]
        result[i] += score
    end

    if !isnothing(node.winner)
        result[node.winner] = 1000
    end

    for d in node.dead
        result[d] = -1000
    end

    return result

end


module GmshParser

export extract_mesh

function getNodes(meshString::Array{String})
    """
    takes meshString and returns all Nodes with their corresponding tags
    as 4th element in the Array element
    """

    nodeStart = findall(meshString .== "\$Nodes")[1] #\$ escape string key $
    nodeEnd = findall(meshString .== "\$EndNodes")[1]
    _, nnodes,min_node_tag,max_node_tag = parse.(Int, split(meshString[nodeStart+1]))
    @assert max_node_tag == nnodes "nodes tags are not contiguous"
    nodes = []
    node = []
    tag = []
    nodeBlock::Int = 0
    n_entity_nodes = 0
    entity_dim = 0
    entity_tag = 0
    tmp = 0
    first = true

    for nodeLine in meshString[nodeStart+2:nodeEnd-1]
        if nodeBlock == 0 && parse(Int, split(nodeLine)[4]) == 0
            continue
        elseif nodeBlock == 0 && parse(Int, split(nodeLine)[4]) != 0
            entity_dim, entity_tag, _, n_entity_nodes = parse.(Int, split(nodeLine))
            #nodeBlock = parse(Int, split(nodeLine)[4])
            nodeBlock = 2 * n_entity_nodes
            tmp = nodeBlock
            if first == false
                #block = push!.(node, tag)
                block = vcat.(tag, node)
                push!(nodes, block)
                tag = []
                node = []
            end
            first = false
        else
            if nodeBlock > tmp // 2
                push!(tag, [entity_dim, entity_tag, parse(Int, nodeLine)])
            else
                push!(node, parse.(Float64, split(nodeLine)))
            end
            nodeBlock -= 1.0
        end
    end
    # The arithmetic above does not account the last node Block
    # hence we need to explicitely append it
    #block = push!.(node, tag)
    block = vcat.(tag, node)
    push!(nodes, block)
    # flatten the array
    return collect(Iterators.flatten(nodes)) # [entDim entTag nodeTag X Y Z]
end

function map_entities_to_physical_names(meshstring)
    linestart = findall(meshstring .== "\$Entities")[1] #\$ escape string key $
    linestop = findall(meshstring .== "\$EndEntities")[1]
    phys_tag2name = map_physical_tags_to_names(meshstring)
    npoints, ncurves, nsurfs, nvols = parse.(Int, split(meshstring[linestart+1]))
    limits = cumsum((npoints, ncurves, nsurfs, nvols))
    entities_tag2name = (Dict{Int,String}(), Dict{Int,String}(), Dict{Int,String}(), Dict{Int,String}())
    dim_id = 1
    for (i,line) in enumerate(meshstring[linestart+2 : linestop-1])
        dim_id = get_dim_id_recursively(i,dim_id,limits)
        if dim_id == 1
            pointline = split(line)
            if length(pointline) > 5
                n_phys_tags = parse(Int,pointline[5])
                for itag in 1:n_phys_tags
                    entity_tag = parse(Int,pointline[1])
                    phys_tag = parse(Int,pointline[5+itag])
                    push!(entities_tag2name[dim_id], entity_tag => phys_tag2name[phys_tag])
                end
            end
        elseif dim_id >= 2
            line_vec = split(line)
            if length(line_vec) > 8
                n_phys_tags = parse(Int,line_vec[8])
                for itag in 1:n_phys_tags
                    entity_tag = parse(Int,line_vec[1])
                    phys_tag = parse(Int,line_vec[8+itag])
                    push!(entities_tag2name[dim_id], entity_tag => phys_tag2name[phys_tag])
                end
            end
        end
    end
    return entities_tag2name
end

function map_physical_tags_to_names(meshstring)
    linestart = findall(meshstring .== "\$PhysicalNames")[1] #\$ escape string key $
    linestop = findall(meshstring .== "\$EndPhysicalNames")[1]
    phys_tag2name = Dict{Int,String}()
    for line in meshstring[linestart+2 : linestop-1]
        entity_line = split(line)
        phys_tag2name[parse(Int,entity_line[2])] = entity_line[3][2:end-1]
    end
    return phys_tag2name
end

function get_dim_id_recursively(i,dim_id,limits)
    if i > limits[dim_id]
        return get_dim_id_recursively(i,dim_id + 1,limits)
    else
        return dim_id
    end
end

function getElements(meshString::Array{String})
    dim = get_dim(meshString)
    elementStart = findall(meshString .== "\$Elements")[1] #\$ escape string key $
    elementStop = findall(meshString .== "\$EndElements")[1]
    elementTotal = parse(Int, split(meshString[elementStart+1])[2]) #no. of Nodes
    nodeElements = []
    boundaryElements = []
    domainElements = []
    ele = []
    eleBlock::Int = 0
    tmp = 0
    eleDim = 0
    first = true
    eletype = 0
    entity_tag = 1

    for eleLine in meshString[elementStart+2:elementStop-1]
        if eleBlock == 0 && parse(Int, split(eleLine)[4]) == 0 # Entity definition and containing no elements
            continue
        elseif eleBlock == 0 && parse(Int, split(eleLine)[4]) != 0 # entity def and contains elements
            eleBlock = parse(Int, split(eleLine)[4]) # number of element for this entity
            tmp = eleBlock
            if first == false
                if eleDim == dim-1
                    push!(boundaryElements, pushfirst!.(ele,entity_tag))
                elseif eleDim == dim
                    push!(domainElements, ele)
                end
            end
            eleDim, entity_tag, eletype, param = parse.(Int, split(eleLine)) # dimension of the entity
            ele = []
            first = false # TODO move that before loop ends
        else
            push!(ele, parse.(Int, split(eleLine)))  # elem id + nodes global indices
            eleBlock -= 1.0
        end
    end
    if eleDim == dim - 2
        push!(boundaryElements, pushfirst!.(ele,entity_tag))
    elseif eleDim == dim - 1
        push!(boundaryElements, pushfirst!.(ele,entity_tag)) # [ entity_tag  elem id  nodes_gidx... ]
    elseif eleDim == dim
        push!(domainElements, ele) # [ elem id  nodes_idx... ]
    end

    # flatten the array
    nodeElements = collect(Iterators.flatten(nodeElements))
    boundaryElements = collect(Iterators.flatten(boundaryElements))
    domainElements = collect(Iterators.flatten(domainElements))
    return nodeElements, boundaryElements, domainElements
end

# function nodeToDataFrame(nodes)
#     df = DataFrame(tag = [], x = [], y = [], z = [])
#     for node in nodes
#         push!(df, [node[1], node[2], node[3], node[4]])
#     end
#     return df
# end

function get_facesets(meshstring, elements, el_type)
    nodeElements, boundaryElements, domainElements = elements
    n_nodeElements = size(nodeElements,1)
    n_boundaryElements = size(boundaryElements,1)
    n_domainElements, nnodes_domainElements = size(domainElements)
    nnodes_domainElements -= 1
    boundary_faces = Tuple{String,Int,Int}[]
    n_boundaries_entities = length(unique(boundaryElements[:,1]))
    #boundaries_names = get_boundaries_entities_names(meshstring, dim)
    curves_tag2name = map_entities_to_physical_names(meshstring)[2] # get this for curves entities
    for i = 1:n_boundaryElements
        boundary_entity_tag = boundaryElements[i,1]
        boundary_name = curves_tag2name[boundary_entity_tag]
        boundary_nodes = boundaryElements[i,3:end]
        for j = 1:n_domainElements
            element_nodes = domainElements[j,2:end]
            #println("element type is : $(domainElements[j,1]) $el_type")
            if issubset(boundary_nodes, element_nodes)
                faces_nodes_indices = faces_indices(el_type)
                face_id = findall([issubset(element_nodes[face_idxs], boundary_nodes) for face_idxs in faces_nodes_indices])
                @assert length(face_id) == 1 "length(face_id) = $(length(face_id))"
                push!(boundary_faces, (boundary_name, j, face_id[1])) # pushes (physical_entity_tag, elem_id, face_number)
            end # domainElements[j,1] - (n_nodeElements + n_boundaryElements)
        end
    end

    facesets = Dict{String,Set{Tuple{Int,Int}}}()
    for name in unique([face[1] for face in boundary_faces])#unique(two_c[:,1])
        slice = filter(x->x[1]==name, boundary_faces)
        facesets[name] = Set{Tuple{Int,Int}}(getindex.(slice,[2:3]))
    end
    return facesets
end

function get_nodesets(meshstring, nodes) #[entDim entTag nodeTag X Y Z]
    nodes_in_nodalentity = [node for node in nodes if node[1] == 0]#nodes[nodes[:,1] .== 0 ,:]
    nodesets = Dict{String,Set{Int}}()
    points_tag2name = map_entities_to_physical_names(meshstring)[1]
    nnodessets = length(points_tag2name)
    for (entity_tag,entity_name) in pairs(points_tag2name)
        nodes_w_entity_tag = [node for node in nodes_in_nodalentity if node[2] == entity_tag]#[nodes_in_nodalentity[:,2] .== entity_tag,:]
        nodes_tags = Set([node[3] for node in nodes_w_entity_tag])
        push!(nodesets, entity_name => nodes_tags)
    end
    return nodesets
end

function faces_indices(c::String)
    if c in ("Line","QuadraticLine")
        return [1, 2]
    elseif c in ("Triangle","QuadraticTriangle")
        return ([1,2], [2,3], [3,1])
    elseif c in ("Quadrilateral","QuadraticQuadrilateral")
        return ([1,2], [2,3], [3,4], [4,1])
    elseif c in ("Tetrahedron","QuadraticTetrahedron")
        return ([1,2,3], [1,2,4], [2,3,4], [1,4,3])
    elseif c in ("Hexahedron","QuadraticHexahedron")
        return ([1,4,3,2], [1,2,6,5], [2,3,7,6], [3,4,8,7], [1,5,8,4], [5,6,7,8])
    else

    end
end

get_el_type_str(meshstring,higher_order_elem) = _get_el_type_str(get_dim(meshstring),length(higher_order_elem[end,2:end]))

function _get_el_type_str(dim, nnodes)
    if dim == 2
        if nnodes == 3
            el_type = "Triangle" #TRI3
        elseif nnodes == 6
            el_type = "QuadraticTriangle" #TRI3
        elseif nnodes == 4
            el_type = "Quadrilateral" #QUAD4
        elseif nnodes == 9
            el_type = "QuadraticQuadrilateral" #QUAD4
        else
            @error "element type identifier $el_type_number from .msh doesn't have any JuAFEM correspondance ."
        end
    else
        @error "dimension != 2 are not implemented yet"
    end
    return el_type
end

function get_dim(meshstring)
    linestart = findall(meshstring .== "\$Entities")[1] #\$ escape string key $
    ndim_entities = parse.(Int, split(meshstring[linestart+1]))
    domain_dim = findlast(ndim_entities .!= 0) - 1
    return domain_dim
end

function extract_mesh(meshfile)
    io = open(meshfile)
    lines = readlines(io)
    close(io)

    meshnodes = getNodes(lines)

    el_0D, el_1D, el_2D = getElements(lines)

    el_0D_mat = hcat(el_0D...)'
    el_1D_mat = hcat(el_1D...)'
    el_2D_mat = hcat(el_2D...)'
    elements_mat = (el_0D_mat, el_1D_mat, el_2D_mat)

    higher_order_el_mat = elements_mat[findlast(.!isempty.(elements_mat))]
    #println(higher_order_el_mat)
    el_type_str = get_el_type_str(lines,higher_order_el_mat)

    facesets = get_facesets(lines,elements_mat,el_type_str)
    nodesets = get_nodesets(lines, meshnodes)

    return el_type_str, higher_order_el_mat, meshnodes, facesets, nodesets

end

end # module


@testset "smarts.smarts" begin

@testset "chain" begin
    nullmol = parse(SMARTS, "")
    @test nodecount(nullmol) == 0
    @test edgecount(nullmol) == 0

    aliphatic = parse(SMARTS, "C")
    @test nodeattr(aliphatic, 1).query == QueryFormula(:and, Set([
        QueryFormula(:atomsymbol, :C),
        QueryFormula(:isaromatic, false)
    ]))
    @test edgecount(aliphatic) == 0

    carbonyl = parse(SMARTS, "[CX3]=[OX1]")
    @test nodeattr(carbonyl, 1).query == QueryFormula(:and, Set([
        QueryFormula(:atomsymbol, :C),
        QueryFormula(:isaromatic, false),
        QueryFormula(:connectivity, 3)
    ]))
    @test nodeattr(carbonyl, 2).query == QueryFormula(:and, Set([
        QueryFormula(:atomsymbol, :O),
        QueryFormula(:isaromatic, false),
        QueryFormula(:connectivity, 1)
    ]))

    ether = parse(SMARTS, "[#6][OD2][#6]")
    @test nodeattr(ether, 1).query == QueryFormula(:atomsymbol, :C)
    @test nodeattr(ether, 2).query == QueryFormula(:and, Set([
        QueryFormula(:atomsymbol, :O),
        QueryFormula(:isaromatic, false),
        QueryFormula(:nodedegree, 2)
    ]))

    notH = parse(SMARTS, "[!#1]")
    @test nodeattr(notH, 1).query == QueryFormula(:not, QueryFormula(:atomsymbol, :H))

    imine = parse(SMARTS, raw"[$([CX3]([#6])[#6]),$([CX3H][#6])]=[$([NX2][#6]),$([NX2H])]")
    @test nodeattr(imine, 1).query == QueryFormula(:or, Set([
        QueryFormula(:recursive, "[CX3]([#6])[#6]"),
        QueryFormula(:recursive, "[CX3H][#6]")
    ]))
    @test nodeattr(imine, 2).query == QueryFormula(:or, Set([
        QueryFormula(:recursive, "[NX2][#6]"),
        QueryFormula(:recursive, "[NX2H]")
    ]))
end

end # smarts.SMARTS


"""
Load metadata for selected stations.
"""
function load_metadata(stat_sel)

    file = joinpath(dirname(pathof(IcanProj)), "..", "data", "stations_metadata.xlsx")

    df_all = readxlsheet(DataFrame ,file, "Ark1")

    df_all[:regine_main] = string.(convert.(Int, df_all[:regine_area])) .* "." .* string.(convert.(Int, df_all[:main_no]))

    df_sel = @from i in df_all begin
        @where i.regine_main in stat_sel
        @select i
        @collect DataFrame
    end

    return df_sel

end

"""
Print nice table with metadata.
"""
function clean_metadata(df_sel)

    df_nice = @from i in df_sel begin
        @select {i.regine_main, i.station_name, i.area_total, i.utm_east_z33, i.utm_north_z33, i.perc_agricul,
                 i.perc_bog, i.perc_forest, i.perc_glacier, i.perc_lake, i.perc_mountain, i.perc_urban,
                 i.height_minimum, i.height_maximum}
        @collect DataFrame
    end

    return df_nice

end

"""
Compute indicies for different resolutions.
"""
function resolution_ind(res)
    @assert 50%res == 0
    ind = zeros(Int, 50, 50)
    counter = 1
    for irow = 1:res:50, icol = 1:res:50
        ind[irow:irow+res-1, icol:icol+res-1] = counter
        counter += 1
    end
    return ind
end

function resolution_ind(res, nrow, ncol)
    @assert 50%res==0
    ind = fill(-9999, nrow, ncol)
    nrow = nrow-nrow%res-res+1
    ncol = ncol-ncol%res-res+1
    counter = 1
    for irow = 1:res:nrow, icol = 1:res:ncol
        ind[irow:irow+res-1, icol:icol+res-1] = counter
        counter += 1
    end
    return ind
end

"""
Collect metadata for selected watershed.
"""
function get_watershed_data(df_sel)

    # Get digital elevation data

    file = joinpath(dirname(pathof(IcanProj)), "..", "raw/elevation.asc")

    dem = read_esri_raster(file)

    elev = dem["data"]

    # Get information about senorge extent

    ind_senorge, xcoord, ycoord = senorge_info()

    # Get information about gridcells with valid input data

    joinpath(dirname(pathof(IcanProj)), "..", "data", "InnenforNorge_20170516.txt")

    tmp = readdlm(Pkg.dir("IcanProj", "data", "InnenforNorge_20170516.txt"), ';'; header=true)

    has_metdata = convert.(Int64, tmp[1][:,2]) + 1

    valid_cells = fill(false, size(elev))
    valid_cells = similar(elev, Bool)
    valid_cells[findin(ind_senorge, has_metdata)] = true

    # Catchment information

    dbk_ind = read_dbk_ind()

    # Indices for different resolutions

    ind_1km  = resolution_ind(1)
    ind_5km  = resolution_ind(5)
    ind_10km = resolution_ind(10)
    ind_25km = resolution_ind(25)
    ind_50km = resolution_ind(50)

    # Loop over stations

    wsh_info = WatershedData[]

    for row in eachrow(df_sel)

        # Basic metadata

        name = row[:station_name]
        regine_main = row[:regine_main]
        dbk = row[:drainage_basin_key][1]

        # Find a grid box around the watershed with valid cells

        wsh_ind = dbk_ind[dbk]

        xcoord_center = mean(xcoord[findin(ind_senorge, wsh_ind)])
        ycoord_center = mean(ycoord[findin(ind_senorge, wsh_ind)])

        dist = sqrt.((xcoord - xcoord_center).^2 + (ycoord - ycoord_center).^2)

        xmin_ind, ymin_ind = findn(dist .== minimum(dist))

        xmin_ind, ymin_ind = xmin_ind[1], ymin_ind[1]

        shift_vec = [0 -1 1 -2 2 -3 3 -4 4 -5 5 -6 6 -7 7 -8 8 -9 9 -10 10 -11 11 -12 12 -13 13 -14 14]

        xrange, yrange = [], []

        for xshift = shift_vec, yshift = shift_vec

            xrange = (xmin_ind-25+xshift):(xmin_ind+24+xshift)
            yrange = (ymin_ind-25+yshift):(ymin_ind+24+yshift)

            if sum(valid_cells[xrange, yrange]) == 2500
                break
            end

        end

        println("Number valid cells for $(name)/$(regine_main): $(sum(valid_cells[xrange, yrange]))")
        println("Percentage glacier for $(name)/$(regine_main): $(round(row[:perc_glacier], 2))")

        # Save data if enough input data is available

        if sum(valid_cells[xrange, yrange]) == 2500

            println("Saved $(name)/$(regine_main)")

            wsh_tmp = WatershedData(
                name,
                regine_main,
                dbk,
                xrange,
                yrange,
                ind_senorge[xrange, yrange] - 1,
                ind_1km,
                ind_5km,
                ind_10km,
                ind_25km,
                ind_50km,
                elev[xrange, yrange]
            )

            push!(wsh_info, wsh_tmp)

        else

            println("Skipped $(name)/$(regine_main): not enough input data")

        end

    end

    return wsh_info

end

"""
Link drainage basin key to watershed name.
"""
function get_wsh_name(wsh_info)

    wsh_name = Dict()

    for wsh_one in wsh_info

        name = wsh_one.name
        regine_main = wsh_one.regine_main

        wsh_name[regine_main] = name

    end

    return wsh_name

end


const Maybe{T} = Union{T, Nothing}

abstract type Interpretation end
struct Nonstandard <: Interpretation end
struct Standard <: Interpretation end
struct Replayed <: Interpretation end
struct Conditioned <: Interpretation end
struct Blocked <: Interpretation end
struct Deterministic <: Interpretation end
struct Proposed <: Interpretation end
struct Empirical <: Interpretation end
struct Input <: Interpretation end

const NONSTANDARD = Nonstandard()
const STANDARD = Standard()
const REPLAYED = Replayed()
const BLOCKED = Blocked()
const CONDITIONED = Conditioned()
const DETERMINISTIC = Deterministic()
const PROPOSED = Proposed()
const EMPIRICAL = Empirical()
const INPUT = Input()

abstract type Node end

@doc raw"""
    mutable struct ParametricNode{A, D, T} <: Node
        address :: A
        dist :: D
        value :: Maybe{T}
        logprob :: Float64
        logprob_sum :: Float64
        observed :: Bool
        pa :: Array{Node, 1}
        ch :: Array{Node, 1}
        interpretation :: Union{Interpretation, Vector{Interpretation}}
        last_interpretation :: Union{Interpretation, Vector{Interpretation}}
    end

A `Node` that can be used with arbitrary code for which `rand` and `Distributionss.logpdf` are defined.
"""
mutable struct ParametricNode{A, D, T} <: Node
    address :: A
    dist :: D
    value :: Maybe{T}
    logprob :: Float64
    logprob_sum :: Float64
    observed :: Bool
    pa :: Array{Node, 1}
    ch :: Array{Node, 1}
    interpretation :: Union{Interpretation, Vector{Interpretation}}
    last_interpretation :: Union{Interpretation, Vector{Interpretation}}
end

@doc raw"""
    mutable struct SampleableNode{A, T} <: Node
        address :: A
        dist :: Sampleable
        value :: Maybe{T}
        logprob :: Float64
        logprob_sum :: Float64
        observed :: Bool
        pa :: Array{Node, 1}
        ch :: Array{Node, 1}
        interpretation :: Union{Interpretation, Vector{Interpretation}}
        last_interpretation :: Union{Interpretation, Vector{Interpretation}}
    end

A `Node` that is restricted to be used with any `Sampleable` from `Distributions.jl`.
"""
mutable struct SampleableNode{A, T} <: Node
    address :: A
    dist :: Sampleable
    value :: Maybe{T}
    logprob :: Float64
    logprob_sum :: Float64
    observed :: Bool
    pa :: Array{Node, 1}
    ch :: Array{Node, 1}
    interpretation :: Union{Interpretation, Vector{Interpretation}}
    last_interpretation :: Union{Interpretation, Vector{Interpretation}}
end

@doc raw"""
    function node(T :: DataType, address :: A, dist :: D, is_obs :: Bool, i :: Interpretation) where {A, D}

Outer constructor for `Node` where no data is passed during construction.
"""
function node(T :: DataType, address :: A, dist :: D, is_obs :: Bool, i :: Interpretation) where {A, D}
    ParametricNode{A, D, T}(address, dist, nothing, 0.0, 0.0, is_obs, Array{Node, 1}(), Array{Node, 1}(), i, i)
end

@doc raw"""
    function node(T :: DataType, address :: A, dist :: D, is_obs :: Bool, i :: Interpretation) where {A, D}

Outer constructor for `Node` where no data is passed during construction.
"""
function node(T :: DataType, address :: A, dist :: D, is_obs :: Bool, i :: Interpretation) where {A, D <: Sampleable}
    SampleableNode{A, T}(address, dist, nothing, 0.0, 0.0, is_obs, Array{Node, 1}(), Array{Node, 1}(), i, i)
end

function setlp!(n::Node, lp)
    n.logprob = lp
    n.logprob_sum = lp
end

# default fallback to distributions method
logprob(d, v) = Distributions.logpdf(d, v)

@doc raw"""
    function node(value, address :: A, dist :: D, is_obs :: Bool, i :: Interpretation) where {A, D}

Outer constructor for `Node` where data is passed during construction. Data type is inferred from the passed data.
"""
function node(value, address :: A, dist :: D, is_obs :: Bool, i :: Interpretation) where {A, D}
    T = typeof(value)
    lp = logprob(dist, value)
    ParametricNode{A, D, T}(address, dist, value, lp, lp, is_obs, Array{Node, 1}(), Array{Node, 1}(), i, i)
end

@doc raw"""
    function node(value, address :: A, dist :: D, is_obs :: Bool, i :: Interpretation) where {A, D}

Outer constructor for `Node` where data is passed during construction. Data type is inferred from the passed data.
"""
function node(value, address :: A, dist :: D, is_obs :: Bool, i :: Interpretation) where {A, D <: Sampleable}
    T = typeof(value)
    lp = logprob(dist, value)
    SampleableNode{A, T}(address, dist, value, lp, lp, is_obs, Array{Node, 1}(), Array{Node, 1}(), i, i)
end

# pass through input into program without scoring it
logprob(::Input, value) = 0.0

@doc raw"""
    abstract type Trace end

Base type for all traces. `Trace`s support the following `Base` methods: `setindex!`, `getindex`,
`keys`, `values`, and `length`.
"""
abstract type Trace end

@doc raw"""
    mutable struct UntypedTrace
        trace :: OrderedDict{Any, Node}
        logprob_sum :: Float64
    end

Trace that can hold nodes with all address and value types.
"""
mutable struct UntypedTrace <: Trace
    trace :: OrderedDict{Any, Node}
    logprob_sum :: Float64
end

@doc raw"""
    mutable struct TypedTrace{A, T} <: Trace
        trace :: OrderedDict{A, SampleableNode{A, T}}
        logprob_sum :: Float64
    end

Trace that can hold nodes of the specific address (`A`) and value (`T`) types.
"""
mutable struct TypedTrace{A, T} <: Trace
    trace :: OrderedDict{A, SampleableNode{A, T}}
    logprob_sum :: Float64
end

@doc raw"""
    trace()

This is the recommended way to construct a new trace.
"""
trace() = UntypedTrace(OrderedDict{Any, Node}(), 0.0)

@doc raw"""
    trace(A, T)

This is the recommended way to construct a new typed trace.
`A` is the address type, `T` is the value type.
"""
trace(A, T) = TypedTrace(OrderedDict{A, SampleableNode{A, T}}(), 0.0)

Base.setindex!(t :: Trace, k, v) = setindex!(t.trace, k, v)
Base.getindex(t :: Trace, k) = t.trace[k]
Base.values(t :: Trace) = values(t.trace)
Base.keys(t :: Trace) = keys(t.trace)
Base.length(t :: Trace) = length(t.trace)

@doc raw"""
    function interpret_latent!(t :: Trace, i :: Interpretation)

Changes the interpretation of all latent nodes in `t` to have `interpretation == i`
"""
function interpret_latent!(t :: Trace, i :: Interpretation)
    for a in keys(t)
        if !t[a].observed && !(t[a].interpretation == CONDITIONED)
            t[a].interpretation = i
        end
    end
end

@doc raw"""
    function logprob(t :: Trace)

Computes and returns the joint log probability of the trace:

```math
\log p(t) = \sum_{a \in \text{keys}(t)}\log p(t[a])
```
"""
function logprob(t :: Trace)
    l = 0.0
    for v in values(t)
        l += v.logprob_sum
    end
    l
end

@doc raw"""
    function logprob!(t :: Trace)

Computes the joint log probability to the trace and assigns it to `t.logprob_sum`.
"""
function logprob!(t :: Trace)
    l = logprob(t)
    t.logprob_sum = l
end

@doc raw"""
    function loglikelihood(t :: Trace)

Computes and returns the log likelihood of the observed data under the model:

```math
\ell(t) =
\sum_{a:\ [a \in \text{keys}(t)] \wedge [\text{interpretation}(a) = \text{Standard}]}
\log p(t[a])
```
"""
function loglikelihood(t :: Trace)
    l = 0.0
    for v in values(t)
        if v.observed
            l += v.logprob_sum
        end
    end
    l
end

@doc raw"""
    function pa_from_trace(t :: Trace, pa)

Collects nodes in trace corresponding to an iterable of parent addresses `pa`.
"""
function pa_from_trace(t :: Trace, pa)
    this_pa = Array{Node, 1}()
    for that_address in pa
        append!(this_pa, (t[that_address],))
    end
    this_pa
end

@doc raw"""
    function connect_pa_ch!(t :: Trace, pa, a)

Connects parent and child nodes. Adds child nodes to parent's `ch` and
parent nodes to child's `pa`.
"""
function connect_pa_ch!(t :: Trace, pa, a)
    append!(t[a].pa, pa_from_trace(t, pa))
    for that_n in t[a].pa
        append!(that_n.ch, (t[a],))
    end
end

@doc raw"""
    function sample(t :: Trace, a, d, ii :: Array{Interpretation, 1}; pa = ())

Sequentially apply sample statements with interpretations as given in `ii`. This is
used to depth-first traverse the interpretation graph.
"""
function sample(t :: Trace, a, d, ii :: Vector{Interpretation}; pa = ())
    local s
    for i in ii
        s = sample(t, a, d, i; pa = pa)
    end
    s
end

@doc raw"""
    function sample(t :: Trace, a, d, i :: Nonstandard; pa = ())

Samples from distribution `d` into trace `t` at address `a`.

1. Samples a value from `d`
2. Creates a sample node
3. Adds the sample node to trace `t` at value `a`
4. Optionally adds nodes corresponding to the addresses in `pa` as parent nodes
5. Returns the sampled value
"""
function sample(t :: Trace, a, d, i :: Nonstandard; pa = ())
    s = rand(d)
    n = node(s, a, d, false, i)
    t[a] = n
    connect_pa_ch!(t, pa, a)
    s
end

@doc raw"""
    function plate(t::Trace, op::F, a, d, s::Int64, i::Nonstandard; pa = ()) where F<:Function

Plate over latent variables.
"""
function plate(t::Trace, op::F, a, d, s::Int64, i::Nonstandard; pa = ()) where F<:Function
    n = node(Vector{eltype(d)}, a, d, false, i)
    rvals = [rand(d) for _ in 1:s]
    lp = 0.0
    for r in rvals
        lp += logpdf(d, r)
    end
    setlp!(n, lp)
    n.value = rvals
    t[a] = n
    connect_pa_ch!(t, pa, a)
    rvals
end

@doc raw"""
    function sample(t :: Trace, a, d, i :: Replayed; pa = ())

Replays the sampled node through the trace.

1. If `a` is not in `t`'s address set, calls `sample(t, a, d, NONSTANDARD; pa = pa)`.
2. Creates a sample node that copies the value from the last node stored in the trace at address `a`.
3. Adds the sample node to trace `t` at value `a`
4. Optionally adds nodes corresponding to the addresses in `pa` as parent nodes
5. Resets the node's interpretation to the original interpretation
"""
function sample(t :: Trace, a, d, i :: Replayed; pa = ())
    if a in keys(t)
        last_i = t[a].last_interpretation
        n = node(t[a].value, a, d, false, i)
        t[a] = n
        connect_pa_ch!(t, pa, a)
        r = t[a].value
        t[a].interpretation = last_i
    else
        r = sample(t, a, d, NONSTANDARD; pa = pa)
    end
    r
end

@doc raw"""
    function plate(t::Trace, op::F, a, d, s::Int64, i::Replayed; pa = ()) where F<:Function

Plate over replayed variables. Note that this method assumes *and does not check* that the value
to be replayed `v` satisfies `length(v) == s`.
"""
function plate(t::Trace, op::F, a, d, s::Int64, i::Replayed; pa = ()) where F<:Function
    n = node(Vector{eltype(d)}, a, d, false, i)
    last_i = t[a].last_interpretation
    lp = 0.0
    for r in t[a].value
        lp += logpdf(d, r)
    end
    setlp!(n, lp)
    n.value = t[a].value
    t[a] = n
    connect_pa_ch!(t, pa, a)
    t[a].interpretation = last_i
    t[a].value
end

@doc raw"""
    function sample(t :: Trace, a, d, i :: Blocked; pa = ())

Samples from `d`, deletes the node stored at address `a` from trace `t`, and returns the
sampled value.
"""
function sample(t :: Trace, a, d, i :: Blocked; pa = ())
    s = rand(d)
    delete!(t.trace, a)
    s
end

@doc raw"""
    function plate(t::Trace, op::F, a, d, s::Int64, i::Blocked; pa = ()) where F<:Function

Plate over blocked variables.
"""
function plate(t::Trace, op::F, a, d, s::Int64, i::Blocked; pa = ()) where F<:Function
    n = node(Vector{eltype(d)}, a, d, false, i)
    rvals = [rand(d) for _ in 1:s]
    delete!(t.trace, a)
    rvals
end

@doc raw"""
    function sample(t :: Trace, a, d, params, i :: Blocked; pa = ())

Samples from `d` passing the optional arguments `params`, deletes the node stored at address `a`
from trace `t`, and returns the sampled value.
"""
function sample(t :: Trace, a, d, params, i :: Blocked; pa = ())
    s = rand(d, params...)
    delete!(t.trace, a)
    s
end

@doc raw"""
    function sample(t :: Trace, a, d, params, ii :: Array{Interpretation, 1}; pa = ())

Sequentially apply sample statements with interpretations as given in `ii`. This is
used to depth-first traverse the interpretation graph.
"""
function sample(t :: Trace, a, d, params, ii :: Array{Interpretation, 1}; pa = ())
    local s
    for i in ii
        s = sample(t, a, d, params, i; pa = pa)
    end
    s
end

@doc raw"""
    function sample(t :: Trace, a, d, params, i :: Replayed; pa = ())

Replays the sampled node through the trace.

1. If `a` is not in `t`'s address set, calls `sample(t, a, d, NONSTANDARD; pa = pa)`.
2. Creates a sample node that copies the value from the last node stored in the trace at address `a`.
3. Adds the sample node to trace `t` at value `a`
4. Optionally adds nodes corresponding to the addresses in `pa` as parent nodes
5. Resets the node's interpretation to the original interpretation
"""
function sample(t :: Trace, a, d, params, i :: Replayed; pa = ())
    sample(t, a, d, i; pa = pa)
end

@doc raw"""
    function sample(t :: Trace, a, d, params, i :: Nonstandard; pa = ())

Samples from distribution `d` into trace `t` at address `a`.

1. Samples a value from `d` passing the optional arguments `params`
2. Creates a sample node
3. Adds the sample node to trace `t` at value `a`
4. Optionally adds nodes corresponding to the addresses in `pa` as parent nodes
5. Returns the sampled value
"""
function sample(t :: Trace, a, d, params, i :: Nonstandard; pa = ())
    s = rand(d, params...)
    n = node(s, a, d, false, i)
    t[a] = n
    connect_pa_ch!(t, pa, a)
    s
end

@doc raw"""
    function sample(t :: Trace, a, d, s, i :: Standard; pa = ())

Scores an observed value `s` against the distribution `d`, storing the value in trace `t` at
address `a` and optionally adds nodes corresponding to the addresses in `pa` as parent nodes.
"""
function sample(t :: Trace, a, d, s, i :: Standard; pa = ())
    n = node(s, a, d, true, i)
    t[a] = n
    connect_pa_ch!(t, pa, a)
    s
end

@doc raw"""
    function plate(t::Trace, op::F, a, d, v::Vector{T}; pa = ()) where {T, F<:Function}

Plate over observed variables, i.e., a plated component of model likelihood. `v` is the
vector of observations, while `op` is likely `observe`.

Example usage: instead of

```
for (i, d) in enumerate(data)
    observe(t, "data $i", Normal(loc, scale), d)
end
```

we can write

```
plate(t, observe, "data", Normal(loc, scale), data)
```
"""
function plate(t::Trace, op::F, a, d, v::Vector{T}; pa = ()) where {T, F<:Function}
    n = node(typeof(v), a, d, true, STANDARD)
    lp = 0.0
    for r in v
        lp += logpdf(d, r)
    end
    setlp!(n, lp)
    n.value = v
    t[a] = n
    connect_pa_ch!(t, pa, a)
    v
end

@doc raw"""
    function plate(t::Trace, op::F, a, d, v::Vector{T}, params; pa = ()) where {T, F<:Function}

Plate over observed variables with different values but identical distribution, i.e.,
``p(x|z) = \prod_n p(x_n | z_n)``.
This is as opposed to `plate(t::Trace, op::F, a, d, v::Vector{T}; pa = ())`, which is equivalent to
``p(x|z) = \prod_n p(x_n | z)``.

`params` must have the same length as `v`. Each element of `params` corresponds to a vector of that
particular component of the `params`, i.e., ``z = (z_1, ..., z_D)`` where each ``z_d`` has length ``N``,
the number of observed datapoints, and ``D`` is the cardinality of the parameterization of the
distribution.

E.g., replace
```
locs = sample(t, "locs", MvNormal(D, 1.0))
for (i, (loc, d)) in enumerate(zip(locs, data))
    observe(t, "data $i", Normal(loc, 1.0), d)
end
```
with
```
locs = sample(t, "locs", MvNormal(D, 1.0))
plate(t, observe, "data", Normal, data, (locs, ones(D)))
```
"""
function plate(t::Trace, op::F, a, d, v::Vector{T}, params; pa = ()) where {T, F<:Function}
    # accumulate logprob from values with same distribution
    lp = 0.0
    for (datapoint, param_group) in zip(v, zip(params...))
        lp += logpdf(d(param_group...), datapoint)
    end
    n = node(typeof(v), a, d, true, STANDARD)
    setlp!(n, lp)
    n.value = v
    t[a] = n
    connect_pa_ch!(t, pa, a)
    v
end

function plate(t::Trace, op::F, a, d, v::Matrix{T}, params; pa = ()) where {T, F<:Function}
    # accumulate logprob from values with same distribution
    lp = 0.0
    for (i, param_group) in enumerate(zip(params...))
        lp += logpdf(d(param_group...), vec(v[:, i]))
    end
    n = node(typeof(v), a, d, true, STANDARD)
    setlp!(n, lp)
    n.value = v
    t[a] = n
    connect_pa_ch!(t, pa, a)
    v
end

@doc raw"""
    function sample(t :: Trace, a, d, i :: Union{Standard,Conditioned}; pa = ())

Scores an observed value against the distribution `d`, storing the value in trace `t` at
address `a` and optionally adds nodes corresponding to the addresses in `pa` as parent nodes.

This method is used by the `condition` effect. It will probably not be used by most
users.
"""
function sample(t :: Trace, a, d, i :: Union{Standard,Conditioned}; pa = ())
    n = node(t[a].value, a, d, true, i)
    t[a] = n
    connect_pa_ch!(t, pa, a)
    t[a].value
end

@doc raw"""
    function plate(t::Trace, op::F, a, d, s::Int64, i::Conditioned; pa = ()) where F<:Function

Plate over conditioned variables.
"""
function plate(t::Trace, op::F, a, d, s::Int64, i::Conditioned; pa = ()) where F<:Function
    n = node(Vector{eltype(d)}, a, d, false, i)
    lp = 0.0
    for r in t[a].value
        lp += logpdf(d, r)
    end
    setlp!(n, lp)
    n.value = t[a].value
    t[a] = n
    connect_pa_ch!(t, pa, a)
    t[a].value
end

@doc raw"""
    function sample(t :: Trace, a, f, v, i :: Deterministic; pa = ())

Creates a deterministic node mapping the tuple of data `v` through function `f`,
storing the value in trace `t` at address `a`.

1. Maps tuple of data `v` through function `f`, yielding `r = f(v...)`
2. Creates a deterministic node and stores it in `t` at address `a`
3. Optionally adds nodes corresponding to the addresses in `pa` as parent nodes
4. Returns `r`
"""
function sample(t :: Trace, a, f, v, i :: Deterministic; pa = ())
    # TODO: add better testing of this, just pushed two hotfixes from applied work
    r = f(v...)
    T = typeof(r)
    n = node(T, a, f, false, i)
    n.value = r
    t[a] = n
    connect_pa_ch!(t, pa, a)
    r
end

function sample(t :: Trace, a, f, v::String, i :: Deterministic; pa = ())
    r = f(v)
    T = typeof(r)
    n = node(T, a, f, false, i)
    n.value = r
    t[a] = n
    connect_pa_ch!(t, pa, a)
    r
end

function sample(t :: Trace, a, d, i :: Input)
    n = node(d, a, INPUT, true, i)
    t[a] = n
    connect_pa_ch!(t, (), a)
    d
end

@doc raw"""
    input(t :: Trace, a, d)

Track a model input.
"""
input(t :: Trace, a, d) = sample(t, a, d, INPUT)

@doc raw"""
    function transform(t :: Trace, a, f :: F, v; pa = ()) where F <: Function

Alias for `sample(t, a, f, v, DETERMINISTIC; pa = pa)`.
"""
function transform(t :: Trace, a, f :: F, v; pa = ()) where F <: Function
    sample(t, a, f, v, DETERMINISTIC; pa = pa)
end

@doc raw"""
    function sample(t :: Trace, a, d; pa = ())

If `a` is in the set of trace addresses, calls `sample` using `t[a]`'s interpretation.
Otherwise, calls `sample` using nonstandard interpretation.
"""
function sample(t :: Trace, a, d; pa = ())
    if a in keys(t)
        sample(t, a, d, t[a].interpretation; pa = pa)
    else
        sample(t, a, d, NONSTANDARD; pa = pa)
    end
end

@doc raw"""
    function sample(t :: Trace, a, d, params; pa = ())

If `a` is in the set of trace addresses, calls `sample` using `t[a]`'s interpretation.
Otherwise, calls `sample` using nonstandard interpretation.
"""
function sample(t :: Trace, a, d, params; pa = ())
    if a in keys(t)
        sample(t, a, d, params, t[a].interpretation; pa = pa)
    else
        sample(t, a, d, params, NONSTANDARD; pa = pa)
    end
end

@doc raw"""
    function observe(t :: Trace, a, d, s; pa = ())

If `d` is not `nothing` an alias for calling `sample` with standard interpretation.
Otherwise, an alias for calling `sample` with nonstandard interpretation.
"""
function observe(t :: Trace, a, d, s; pa = ())
    if s === nothing
        sample(t, a, d, NONSTANDARD; pa = pa)
    else
        sample(t, a, d, s, STANDARD; pa = pa)
    end
end

# maybe Proposed will have meaning later, stub out just in case
# e.g., right now new trace is made in mh, but is this always true?

sample(t :: Trace, a, d, i :: Proposed) = sample(t, a, d, NONSTANDARD)

@doc raw"""
    propose(t :: Trace, a, d)

Propose a value for the address `a` in trace `t` from the distribution `d`.
"""
propose(t :: Trace, a, d) = sample(t, a, d, PROPOSED)

@doc raw"""
    function plate(t::Trace, op::F, a, d, s::Int64; pa = ()) where F<:Function

Sample or observe a vector of random variables at a single site instead of multiple.
This can speed up inference since the number of sites in the model will no longer scale
with dataset size (though numerical value computation is still linear in dataset size).

Example usage: instead of

```
vals = [sample(t, "val $i", Geometric()) for i in 1:N]
```

we can write

```
vals = plate(t, sample, "val", Geometric(), N)
```

Mathematically, this is equivalent to the product ``p(z) = \prod_n p(z_n)``
and treating it as the single object ``p(z)`` instead of the ``N`` objects ``p(z_n)``.
"""
function plate(t::Trace, op::F, a, d, s::Int64; pa = ()) where F<:Function
    if a in keys(t)
        plate(t, op, a, d, s, t[a].interpretation; pa = pa)
    else
        plate(t, op, a, d, s, NONSTANDARD; pa = pa)
    end
end
plate(t::Trace, op::F, a, d, v::Vector{Nothing}; pa = ()) where F <: Function = plate(t, op, a, d, length(v); pa = pa)
export plate

function Base.show(io::IO, t::Trace)
    s = "\nTrace with $(length(t)) nodes:"
    for node in values(t)
        s *= "\n$(node.address) = ($(node.value), $(node.dist), $(node.logprob_sum), $(node.interpretation))"
    end
    print(io, s)
end

export Node, ParametricNode, SampleableNode, node
export Trace, UntypedTrace, TypedTrace, trace, logprob, logprob!, sample, observe, input, loglikelihood
export node_info, graph, transform
export Interpretation
export Nonstandard, Standard, Replayed, Conditioned, Deterministic, Input, Empirical
export NONSTANDARD, STANDARD, REPLAYED, CONDITIONED, DETERMINISTIC, INPUT, EMPIRICAL


@testset "Conversions Utilities" begin
    _voxels = zeros(Float32, 32, 32, 32, 2)
    _voxels[1:15, 2:10, 18:32, :] .= 1

    _v = VoxelGrid(_voxels)
    _m = load_trimesh(
        joinpath(@__DIR__, "./assets/teapot.obj"),
        joinpath(@__DIR__, "./assets/sphere.obj"),
    )
    _p = PointCloud(sample_points(_m, 1024))

    res = 28
    points = 512
    thresh = 0.9

    for algo in [:Exact, :MarchingCubes, :MarchingTetrahedra]
        m1 = TriMesh(_p, res; algo = algo)
        @test m1 isa TriMesh{Float32,UInt32,Array}

        m2 = TriMesh(_v, thresh = thresh, algo = algo)
        @test m2 isa TriMesh{Float32,UInt32,Array}

        p1 = PointCloud(_v, points, thresh = thresh, algo = algo)
        @test p1 isa PointCloud{Float32}
        @test size(p1.points) == (3, points, 2)
    end

    p2 = PointCloud(_m, points)
    @test p2 isa PointCloud{Float32}
    @test size(p2.points) == (3, points, 2)

    v1 = VoxelGrid(_m, res)
    @test v1 isa VoxelGrid{Float32}
    @test size(v1.voxels) == (res, res, res, 2)

    v2 = VoxelGrid(_p, res)
    @test v2 isa VoxelGrid{Float32}
    @test size(v2.voxels) == (res, res, res, 2)

end

# Validating the geometry of each figures is not possible, so use below code to
# to cross check the dimension and figure of different conversion.

# using Makie
# Makie.inline!(true)
# Makie.set_theme!(show_axis = true)
#
# _voxels = zeros(Float32,32,32,32,2)
# _voxels[1:15,2:10,18:32,:] .= 1
#
# _v = VoxelGrid(_voxels)
# _m = load_trimesh(["assets/teapot.obj","assets/sphere.obj"])
# _p = PointCloud(sample_points(_m,1024))
#
# visualize(_v,1,algo=:Exact)
# visualize(_v,2,algo=:MarchingCubes)
#
# m1 = TriMesh(_v, algo=:MarchingCubes)
# visualize(m1,1)
# visualize(m1,2)
#
# p1 = PointCloud(_v,1000,algo=:MarchingCubes)
# visualize(p1,1,markersize=0.1)
# visualize(p1,2, markersize=0.06)
#
# v2 = VoxelGrid(_m)
# visualize(v2,1,algo=:MarchingCubes)
# visualize(v2,2,algo=:MarchingCubes)
#
# p2 = PointCloud(_m)
# visualize(p2,1,markersize=0.3)
# visualize(p2,2)
#
# m3 = TriMesh(_p)
# visualize(m3,1,color=:green)
# visualize(m3,2)
#
# v3 = VoxelGrid(_p)
# visualize(v3,1)
# visualize(v3,2)


module NeuralProcesses

using CUDA
using Distributions
using Flux
using LinearAlgebra
using NNlib
using Printf
using Random
using Statistics
using StatsBase
using Tracker

if CUDA.functional()
    include("gpu.jl")

    randn_gpu = CUDA.randn
    zeros_gpu = CUDA.zeros
    ones_gpu = CUDA.ones
else
    const CuOrVector = Vector
    const CuOrMatrix = Matrix
    const CuOrArray = Array

    randn_gpu = randn
    zeros_gpu = zeros
    ones_gpu = ones
end

const AV = AbstractVector
const AM = AbstractMatrix
const AA = AbstractArray

const MaybeAV = Union{Nothing, AbstractVector}
const MaybeAM = Union{Nothing, AbstractMatrix}
const MaybeAA = Union{Nothing, AbstractArray}

export Chain

include("util.jl")
include("data.jl")
include("conv.jl")
include("discretisation.jl")
include("parallel.jl")
include("nn.jl")
include("distribution.jl")

include("model/setconv.jl")
include("model/attention.jl")
include("model/coding.jl")
include("model/coder.jl")
include("model/noise.jl")
include("model/model.jl")

include("model/architectures/cnp.jl")
include("model/architectures/np.jl")
include("model/architectures/acnp.jl")
include("model/architectures/anp.jl")
include("model/architectures/convcnp.jl")
include("model/architectures/corconvcnp.jl")
include("model/architectures/convnp.jl")

include("experiment/experiment.jl")

end


# This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT).

"""
    AbstractArrayOfSimilarArrays{T,M,N} <: AbstractArray{<:AbstractArray{T,M},N}

An array that contains arrays that have the same size/axes. The array is
internally stored in flattened form as some kind of array of dimension
`M + N`. The flattened form can be accessed via `flatview(A)`.

Subtypes must implement (in addition to typical array operations):

    flatview(A::SomeArrayOfSimilarArrays)::AbstractArray{T,M+N}

The following type aliases are defined:

* `AbstractVectorOfSimilarArrays{T,M} = AbstractArrayOfSimilarArrays{T,M,1}`
* `AbstractArrayOfSimilarVectors{T,N} = AbstractArrayOfSimilarArrays{T,1,N}`
* `AbstractVectorOfSimilarVectors{T} = AbstractArrayOfSimilarArrays{T,1,1}`
"""
abstract type AbstractArrayOfSimilarArrays{T,M,N} <: AbstractArray{Array{T,M},N} end
export AbstractArrayOfSimilarArrays

const AbstractVectorOfSimilarArrays{T,M} = AbstractArrayOfSimilarArrays{T,M,1}
export AbstractVectorOfSimilarArrays

const AbstractArrayOfSimilarVectors{T,N} = AbstractArrayOfSimilarArrays{T,1,N}
export AbstractArrayOfSimilarVectors

const AbstractVectorOfSimilarVectors{T} = AbstractArrayOfSimilarArrays{T,1,1}
export AbstractVectorOfSimilarVectors

"""
    ArrayOfSimilarArrays{T,M,N,L,P} <: AbstractArrayOfSimilarArrays{T,M,N}

Represents a view of an array of dimension `L = M + N` as an array of
dimension M with elements that are arrays with dimension N. All element arrays
implicitly have equal size/axes.

Constructors:

    ArrayOfSimilarArrays{T,M,N}(flat_data::AbstractArray)
    ArrayOfSimilarArrays{T,M}(flat_data::AbstractArray)

The following type aliases are defined:

* `VectorOfSimilarArrays{T,M} = AbstractArrayOfSimilarArrays{T,M,1}`
* `ArrayOfSimilarVectors{T,N} = AbstractArrayOfSimilarArrays{T,1,N}`
* `VectorOfSimilarVectors{T} = AbstractArrayOfSimilarArrays{T,1,1}`

`VectorOfSimilarArrays` supports `push!()`, etc., provided the underlying
array supports resizing of it's last dimension (e.g. an `ElasticArray`).

The nested array can also be created using the function [`nestedview`](@ref)
and the wrapped flat array can be accessed using [`flatview`](@ref)
afterwards:

```julia
A_flat = rand(2,3,4,5,6)
A_nested = nestedview(A_flat, 2)
A_nested isa AbstractArray{<:AbstractArray{T,2},3} where T
flatview(A_nested) === A_flat
```
"""
struct ArrayOfSimilarArrays{
    T, M, N, L,
    P<:AbstractArray{T,L}
} <: AbstractArrayOfSimilarArrays{T,M,N}
    data::P

    function ArrayOfSimilarArrays{T,M,N}(flat_data::AbstractArray{U,L}) where {T,M,N,L,U}
        size_inner, size_outer = split_tuple(size(flat_data), Val{M}())
        require_ndims(flat_data, _add_vals(Val{M}(), Val{N}()))
        conv_parent = _convert_elype(T, flat_data)
        P = typeof(conv_parent)
        new{T,M,N,L,P}(conv_parent)
    end

    function ArrayOfSimilarArrays{T,M}(flat_data::AbstractArray{U,L}) where {T,M,L,U}
        size_inner, size_outer = split_tuple(size(flat_data), Val{M}())
        N = length(size_outer)
        conv_parent = _convert_elype(T, flat_data)
        P = typeof(conv_parent)
        new{T,M,N,L,P}(conv_parent)
    end
end

export ArrayOfSimilarArrays

function ArrayOfSimilarArrays{T,M,N}(A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U}
    B = ArrayOfSimilarArrays{T,M,N}(Array{T}(undef, innersize(A)..., size(A)...))
    copyto!(B, A)
end

ArrayOfSimilarArrays{T}(A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} =
    ArrayOfSimilarArrays{T,M,N}(A)

ArrayOfSimilarArrays(A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N} =
    ArrayOfSimilarArrays{T,M,N}(A)

Base.convert(R::Type{ArrayOfSimilarArrays{T,M,N}}, flat_data::AbstractArray{U,L}) where {T,M,N,L,U} = R(flat_data)
Base.convert(R::Type{ArrayOfSimilarArrays{T,M}}, flat_data::AbstractArray{U,L}) where {T,M,L,U} = R(flat_data)

Base.convert(R::Type{ArrayOfSimilarArrays{T,M,N}}, A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} = R(A)
Base.convert(R::Type{ArrayOfSimilarArrays{T}}, A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} = R(A)
Base.convert(R::Type{ArrayOfSimilarArrays}, A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N} = R(A)

@inline function innersize(A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N}
    front_tuple(size(A.data), Val{M}())
end

@inline function _innerlength(A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N}
    prod(innersize(A))
end

import Base.==
(==)(A::ArrayOfSimilarArrays{T,M,N}, B::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} =
    (A.data == B.data)

"""
    flatview(A::ArrayOfSimilarArrays{T,M,N,L,P})::P

Returns the array of dimensionality `L = M + N` wrapped by `A`. The shape of
the result may be freely changed without breaking the inner consistency of
`A`.
"""
flatview(A::ArrayOfSimilarArrays) = A.data

Base.size(A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} = split_tuple(size(A.data), Val{M}())[2]

Base.@propagate_inbounds Base.getindex(A::ArrayOfSimilarArrays{T,M,N}, idxs::Vararg{Integer,N}) where {T,M,N} =
    view(A.data, _ncolons(Val{M}())..., idxs...)

Base.@propagate_inbounds Base.setindex!(A::ArrayOfSimilarArrays{T,M,N}, x::AbstractArray{U,M}, idxs::Vararg{Integer,N}) where {T,M,N,U} =
    setindex!(A.data, x, _ncolons(Val{M}())..., idxs...)

Base.@propagate_inbounds function Base.unsafe_view(A::ArrayOfSimilarArrays{T,M,N}, idxs::Vararg{Union{Real, AbstractArray},N}) where {T,M,N}
    dataview = view(A.data, _ncolons(Val{M}())..., idxs...)
    L = length(size(dataview))
    N_view = L - M
    ArrayOfSimilarArrays{T,M,N_view}(dataview)
end

@inline function Base.resize!(A::ArrayOfSimilarArrays{T,M,N}, dims::Vararg{Integer,N}) where {T,M,N}
    resize!(A.data, innersize(A)..., dims...)
    A
end

function Base.similar(A::ArrayOfSimilarArrays{T,M,N}, ::Type{<:AbstractArray{U}}, dims::Dims) where {T,M,N,U}
    data = A.data
    size_inner, size_outer = split_tuple(size(data), Val{M}())
    # ToDo: Don't use similar if data is an ElasticArray?
    ArrayOfSimilarArrays{T,M,N}(similar(data, U, size_inner..., dims...))
end

function Base.deepcopy(A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N}
    ArrayOfSimilarArrays{T,M,N}(deepcopy(A.data))
end

function Base.copyto!(dest::ArrayOfSimilarArrays{T,M,N}, src::ArrayOfSimilarArrays{U,M,N}) where {T,M,N,U}
    copyto!(dest.data, src.data)
    dest
end

function Base.append!(dest::ArrayOfSimilarArrays{T,M,N}, src::ArrayOfSimilarArrays{U,M,N}) where {T,M,N,U}
    innersize(dest) != innersize(src) && throw(DimensionMismatch("Can't append, shape of element arrays of source and dest are not equal"))
    append!(dest.data, src.data)
    dest
end

Base.append!(dest::ArrayOfSimilarArrays{T,M,N}, src::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} =
    append!(dest, ArrayOfSimilarArrays(src))

function Base.prepend!(dest::ArrayOfSimilarArrays{T,M,N}, src::ArrayOfSimilarArrays{U,M,N}) where {T,M,N,U}
    innersize(dest) != innersize(src) && throw(DimensionMismatch("Can't prepend, shape of element arrays of source and dest are not equal"))
    prepend!(dest.data, src.data)
    dest
end

Base.prepend!(dest::ArrayOfSimilarArrays{T,M,N}, src::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} =
    prepend!(dest, ArrayOfSimilarArrays(src))

UnsafeArrays.unsafe_uview(A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} =
    ArrayOfSimilarArrays{T,M,N}(uview(A.data))

function innermap(f::Base.Callable, A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N}
    new_data = map(f, A.data)
    U = eltype(new_data)
    ArrayOfSimilarArrays{U,M,N}(new_data)
end

function deepmap(f::Base.Callable, A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N}
    new_data = deepmap(f, A.data)
    U = eltype(new_data)
    ArrayOfSimilarArrays{U,M,N}(new_data)
end

Base.@pure _result_is_nested(idxs_outer::Tuple, idxs_inner::Tuple) =
    Val{!(Base.index_dimsum(idxs_outer...) isa Tuple{}) && !(Base.index_dimsum(idxs_inner...) isa Tuple{})}()

Base.@pure ndims_after_getindex(idxs::Tuple) = Val{length(Base.index_dimsum(idxs...))}()

Base.@propagate_inbounds function deepgetindex(A::ArrayOfSimilarArrays{T,M,N,L}, idxs::Vararg{Any,L}) where {T,M,N,L}
    idxs_outer, idxs_inner = split_tuple(idxs, Val{N}())
    nested = _result_is_nested(idxs_outer, idxs_inner)
    _deepgetindex_impl_aosa(A, idxs_outer, idxs_inner, nested)
end

Base.@propagate_inbounds _deepgetindex_impl_aosa(A::ArrayOfSimilarArrays, idxs_outer::Tuple, idxs_inner::Tuple, nested::Val{false}) =
    getindex(A.data, idxs_inner..., idxs_outer...)

Base.@propagate_inbounds function _deepgetindex_impl_aosa(A::ArrayOfSimilarArrays, idxs_outer::Tuple, idxs_inner::Tuple, nested::Val{true})
    new_data = getindex(A.data, idxs_inner..., idxs_outer...)
    nestedview(new_data, ndims_after_getindex(idxs_inner))
end

Base.@propagate_inbounds function deepsetindex!(A::ArrayOfSimilarArrays{T,M,N,L}, x, idxs::Vararg{Any,L}) where {T,M,N,L}
    idxs_outer, idxs_inner = split_tuple(idxs, Val{N}())
    _deepsetindex_impl_aosa!(A, x, idxs_outer, idxs_inner, _result_is_nested(idxs_outer, idxs_inner))
    A
end

Base.@propagate_inbounds _deepsetindex_impl_aosa!(A::ArrayOfSimilarArrays, x, idxs_outer::Tuple, idxs_inner::Tuple, nested::Val{false}) =
    setindex!(A.data, x, idxs_inner..., idxs_outer...)

Base.@propagate_inbounds _deepsetindex_impl_aosa!(A::ArrayOfSimilarArrays, x, idxs_outer::Tuple, idxs_inner::Tuple, nested::Val{true}) =
    _deepsetindex_impl!(A, x, idxs_outer, idxs_inner)

# TODO: Specialized method _deepsetindex_impl_aosa!(A::ArrayOfSimilarArrays, x::ArrayOfSimilarArrays, idxs_outer::Tuple, idxs_inner::Tuple, nested::Val{true}) =

Base.@propagate_inbounds function deepview(A::ArrayOfSimilarArrays{T,M,N,L}, idxs::Vararg{Any,L}) where {T,M,N,L}
    idxs_outer, idxs_inner = split_tuple(idxs, Val{N}())
    nested = _result_is_nested(idxs_outer, idxs_inner)
    _deepview_impl_aosa(A, idxs_outer, idxs_inner, nested)
end

Base.@propagate_inbounds _deepview_impl_aosa(A::ArrayOfSimilarArrays, idxs_outer::Tuple, idxs_inner::Tuple, nested::Val{false}) =
    view(A.data, idxs_inner..., idxs_outer...)

Base.@propagate_inbounds function _deepview_impl_aosa(A::ArrayOfSimilarArrays, idxs_outer::Tuple, idxs_inner::Tuple, nested::Val{true})
    new_data = view(A.data, idxs_inner..., idxs_outer...)
    nestedview(new_data, ndims_after_getindex(idxs_inner))
end

const VectorOfSimilarArrays{
    T, M, L,
    P<:AbstractArray{T,L}
} = ArrayOfSimilarArrays{T,M,1,L,P}

export VectorOfSimilarArrays

VectorOfSimilarArrays{T}(flat_data::AbstractArray{U,L}) where {T,U,L} =
    ArrayOfSimilarArrays{T,length(Base.front(size(flat_data))),1}(flat_data)

VectorOfSimilarArrays(flat_data::AbstractArray{T,L}) where {T,L} =
    ArrayOfSimilarArrays{T,length(Base.front(size(flat_data))),1}(flat_data)

VectorOfSimilarArrays{T}(A::AbstractVector{<:AbstractArray{U,M}}) where {T,M,U} =
    VectorOfSimilarArrays{T,M}(A)

VectorOfSimilarArrays(A::AbstractVector{<:AbstractArray{T,M}}) where {T,M} =
    VectorOfSimilarArrays{T,M}(A)

Base.convert(R::Type{VectorOfSimilarArrays{T}}, flat_data::AbstractArray{U,L}) where {T,U,L} = R(flat_data)
Base.convert(R::Type{VectorOfSimilarArrays}, flat_data::AbstractArray{T,L}) where {T,L} = R(flat_data)
Base.convert(R::Type{VectorOfSimilarArrays{T}}, A::AbstractVector{<:AbstractArray{U,M}}) where {T,M,U} = R(A)
Base.convert(R::Type{VectorOfSimilarArrays}, A::AbstractVector{<:AbstractArray{T,M}}) where {T,M} = R(A)

@inline Base.IndexStyle(A::VectorOfSimilarArrays) = IndexLinear()

function Base.push!(V::VectorOfSimilarArrays{T,M}, x::AbstractArray{U,M}) where {T,M,U}
    size(x) != Base.front(size(V.data)) && throw(DimensionMismatch("Can't push, shape of source and elements of target is incompatible"))
    append!(V.data, x)
    V
end

function Base.pop!(V::VectorOfSimilarArrays)
    isempty(V) && throw(ArgumentError("array must be non-empty"))
    x = V[end]
    resize!(V, size(V, 1) - 1)
    x
end

function Base.pushfirst!(V::VectorOfSimilarArrays{T,M}, x::AbstractArray{U,M}) where {T,M,U}
    size(x) != Base.front(size(V.data)) && throw(DimensionMismatch("Can't push, shape of source and elements of target is incompatible"))
    prepend!(V.data, x)
    V
end

# Will need equivalent of resize! that resizes in front of data instead of in back:
# popfirst!(V::ArrayOfSimilarArrays) = ...

function Base.empty(A::VectorOfSimilarArrays{T,M}, ::Type{<:AbstractArray{U}}) where {T,M,U}
    data = A.data
    size_inner, size_outer = split_tuple(size(data), Val{1}())
    empty_size_outer = map(x -> zero(x), size_outer)

    # ToDo: Don't use similar if data is an ElasticArray?
    VectorOfSimilarArrays{T,M}(similar(data, U, size_inner..., empty_size_outer...))
end

Base.empty(A::VectorOfSimilarArrays{T,M}) where {T,M} = empty(A, Array{T})

const ArrayOfSimilarVectors{
    T, N, L,
    P<:AbstractArray{T,L}
} = ArrayOfSimilarArrays{T,1,N,L,P}

export ArrayOfSimilarVectors

ArrayOfSimilarVectors{T}(flat_data::AbstractArray{U,L}) where {T,U,L} =
    ArrayOfSimilarArrays{T,1,length(Base.front(size(flat_data)))}(flat_data)

ArrayOfSimilarVectors(flat_data::AbstractArray{T,L}) where {T,L} =
    ArrayOfSimilarArrays{T,1,length(Base.front(size(flat_data)))}(flat_data)

ArrayOfSimilarVectors{T}(A::AbstractArray{<:AbstractVector{U},N}) where {T,N,U} =
    ArrayOfSimilarVectors{T,N}(A)

ArrayOfSimilarVectors(A::AbstractArray{<:AbstractVector{T},N}) where {T,N} =
    ArrayOfSimilarVectors{T,N}(A)

Base.convert(R::Type{ArrayOfSimilarVectors{T}}, flat_data::AbstractArray{U,L}) where {T,U,L} = R(flat_data)
Base.convert(R::Type{ArrayOfSimilarVectors}, flat_data::AbstractArray{T,L}) where {T,L} = R(flat_data)
Base.convert(R::Type{ArrayOfSimilarVectors{T}}, A::AbstractArray{<:AbstractVector{U},N}) where {T,N,U} = R(A)
Base.convert(R::Type{ArrayOfSimilarVectors}, A::AbstractArray{<:AbstractVector{T},N}) where {T,N} = R(A)

const VectorOfSimilarVectors{
    T,
    P<:AbstractArray{T,2}
} = ArrayOfSimilarArrays{T,1,1,2,P}

export VectorOfSimilarVectors

VectorOfSimilarVectors{T}(flat_data::AbstractArray{U,2}) where {T,U} =
    ArrayOfSimilarArrays{T,1,1}(flat_data)

VectorOfSimilarVectors(flat_data::AbstractArray{T,2}) where {T} =
    VectorOfSimilarVectors{T}(flat_data)

VectorOfSimilarVectors{T}(A::AbstractVector{<:AbstractVector{U}}) where {T,U} =
    ArrayOfSimilarArrays{T,1}(A)

VectorOfSimilarVectors(A::AbstractVector{<:AbstractVector{T}}) where {T} =
    VectorOfSimilarVectors{T}(A)

Base.convert(R::Type{VectorOfSimilarVectors{T}}, flat_data::AbstractArray{U,2}) where {T,U} = R(flat_data)
Base.convert(R::Type{VectorOfSimilarVectors}, flat_data::AbstractArray{T,2}) where {T} = R(flat_data)
Base.convert(R::Type{VectorOfSimilarVectors{T}}, A::AbstractVector{<:AbstractVector{U}}) where {T,U} = R(A)
Base.convert(R::Type{VectorOfSimilarVectors}, A::AbstractVector{<:AbstractVector{T}}) where {T} = R(A)

@inline Base.IndexStyle(A::VectorOfSimilarVectors) = IndexLinear()

"""
    sum(X::AbstractVectorOfSimilarArrays)
    sum(X::AbstractVectorOfSimilarArrays, w::StatsBase.AbstractWeights)

Compute the sum of the elements vectors of `X`. Equivalent to `sum` of
`flatview(X)` along the last dimension.
"""
Base.sum(X::AbstractVectorOfSimilarArrays{T,M}) where {T,M} = sum(flatview(X); dims = M + 1)

"""
    mean(X::AbstractVectorOfSimilarArrays)
    mean(X::AbstractVectorOfSimilarArrays, w::StatsBase.AbstractWeights)

Compute the mean of the elements vectors of `X`. Equivalent to `mean` of
`flatview(X)` along the last dimension.
"""
Statistics.mean(X::AbstractVectorOfSimilarArrays{T,M}) where {T,M} =
    mean(flatview(X); dims = M + 1)

"""
    var(X::AbstractVectorOfSimilarArrays; corrected::Bool = true)
    var(X::AbstractVectorOfSimilarArrays, w::StatsBase.AbstractWeights; corrected::Bool = true)

Compute the sample variance of the elements vectors of `X`. Equivalent to
`var` of `flatview(X)` along the last dimension.
"""
Statistics.var(X::AbstractVectorOfSimilarArrays{T,M}; corrected::Bool = true) where {T,M} =
    var(flatview(X); dims = M + 1, corrected = corrected)

"""
    cov(X::AbstractVectorOfSimilarVectors; corrected::Bool = true)
    cov(X::AbstractVectorOfSimilarVectors, w::StatsBase.AbstractWeights; corrected::Bool = true)

Compute the covariance matrix between the elements of the elements of `X`
along `X`. Equivalent to `cov` of `flatview(X)` along dimension 2.
"""
Statistics.cov(X::AbstractVectorOfSimilarVectors; corrected::Bool = true) =
    cov(flatview(X); dims = 2, corrected = corrected)

"""
    cor(X::AbstractVectorOfSimilarVectors)
    cor(X::AbstractVectorOfSimilarVectors, w::StatsBase.AbstractWeights)

Compute the Pearson correlation matrix between the elements of the elements of
 `X` along `X`. Equivalent to `cor` of `flatview(X)` along dimension 2.
"""
Statistics.cor(X::AbstractVectorOfSimilarVectors) =
    cor(flatview(X); dims = 2)