text:
using Test
using Distributions
using StableRNGs
using StatsBase
using LinearAlgebra
using Random

using EnsembleKalmanProcesses.ParameterDistributions

@testset "ParameterDistributions" begin
    @testset "ParameterDistributionType" begin
        # Tests for the ParameterDistributionType
        d = Parameterized(Gamma(2.0, 0.8))
        @test d.distribution == Gamma(2.0, 0.8)

        d = Samples([1 2 3; 4 5 6])
        @test d.distribution_samples == [1 2 3; 4 5 6]
        d = Samples([1 4; 2 5; 3 6]; params_are_columns = false)
        @test d.distribution_samples == [1 2 3; 4 5 6]
        d = Samples([1, 2, 3, 4, 5, 6])
        @test d.distribution_samples == [1 2 3 4 5 6]
        d = Samples([1, 2, 3, 4, 5, 6]; params_are_columns = false)
        @test d.distribution_samples == reshape([1, 2, 3, 4, 5, 6], :, 1)

        # Vector container -
        d = VectorOfParameterized(repeat([Normal(0, 1)], 5))
        @test d.distribution == repeat([Normal(0, 1)], 5)
    end
    @testset "ConstraintType" begin
        # Tests for the ConstraintType
        # The predefined transforms
        c1 = bounded_below(0.2)
        @test isapprox(c1.constrained_to_unconstrained(1.0) - (log(1.0 - 0.2)), 0)
        @test isapprox(c1.unconstrained_to_constrained(0.0) - (exp(0.0) + 0.2), 0)

        c2 = bounded_above(0.2)
        @test isapprox(c2.constrained_to_unconstrained(-1.0) - (log(0.2 - -1.0)), 0)
        @test isapprox(c2.unconstrained_to_constrained(10.0) - (0.2 - exp(10.0)), 0)

        c3 = bounded(-0.1, 0.2)
        @test isapprox(c3.constrained_to_unconstrained(0.0) - (log((0.0 - -0.1) / (0.2 - 0.0))), 0)
        @test isapprox(c3.unconstrained_to_constrained(1.0) - ((0.2 * exp(1.0) + -0.1) / (exp(1.0) + 1)), 0)
        @test_throws DomainError bounded(0.2, -0.1)

        #an example with user defined invertible transforms
        c_to_u = (x -> 3 * x + 14)
        u_to_c = (x -> (x - 14) / 3)

        c4 = Constraint(c_to_u, u_to_c)
        @test isapprox(c4.constrained_to_unconstrained(5.0) - c_to_u(5.0), 0)
        @test isapprox(c4.unconstrained_to_constrained(5.0) - u_to_c(5.0), 0)

        #length, size
        @test length(c1) == 1
        @test size(c1) == (1,)

    end

    @testset "ParameterDistribution: Build and combine" begin

        # create single ParameterDistribution
        d = VectorOfParameterized([Normal(0, 1), Gamma(2.0, 0.8)])
        c_mismatch = no_constraint()
        c_wrongtype = [no_constraint(), 3.0]
        c = [no_constraint(), no_constraint()]
        name = "normal_and_gamma"

        @test_throws ArgumentError ParameterDistribution(d, c_wrongtype, name) #wrong type of constraint
        @test_throws DimensionMismatch ParameterDistribution(d, c_mismatch, name) #wrong number of constraints

        # test checks on stored information
        u = ParameterDistribution(d, c, name)
        @test u.distribution == [d]
        @test u.constraint == c #as c is already a vector
        @test u.name == [name]

        d = Parameterized(Normal(0, 1))
        c = [no_constraint()]
        u = ParameterDistribution(d, c, name)
        @test u.constraint == c #as c is already a vector

        # test concatenation
        d1 = Parameterized(Normal(0, 1))
        c1 = no_constraint()
        name1 = "unconstrained_normal"
        u1 = ParameterDistribution(d1, c1, name1)
        @test u1.constraint == [c1]

        d2 = VectorOfParameterized(repeat([MvNormal(3, 0.2)], 4))
        c2 = repeat([no_constraint()], 12) #3D distribution repeated 4 times has 12 constraints
        name2 = "three_unconstrained_MvNs"
        u2 = ParameterDistribution(d2, c2, name2)

        d3 = Samples([1.0 3.0 5.0 7.0; 9.0 11.0 13.0 15.0])
        c3 = [bounded(10, 15), no_constraint()]
        name3 = "constrained_sampled"
        u3 = ParameterDistribution(d3, c3, name3)

        @test_throws ArgumentError ParameterDistribution([u1, u2])

        u = combine_distributions([u1, u2, u3])
        @test u.distribution == [d1, d2, d3]
        @test u.constraint == cat([[c1], c2, c3]..., dims = 1)
        @test u.name == [name1, name2, name3]

    end

    @testset "ParameterDistribution: Dictionary interface" begin

        # dictionary - testing test checks on the distribution elements
        d = Parameterized(Normal(0, 1))
        c = no_constraint()
        name = "unconstrained_normal"

        param_dict_fail = [Dict("a" => 1), 1]
        @test_throws ArgumentError ParameterDistribution(param_dict_fail) # not an array of Dicts

        @test_throws ArgumentError ParameterDistribution(param_dict_fail) # not a Dict

        param_dict_fail = Dict("distribution" => d)
        @test_throws ArgumentError ParameterDistribution(param_dict_fail) # not all the keys
        param_dict_fail = Dict("distribution" => 1, "constraint" => c, "name" => name)
        @test_throws ArgumentError ParameterDistribution(param_dict_fail) # not right distribution type
        param_dict_fail = Dict("distribution" => d, "constraint" => 1, "name" => name)
        @test_throws ArgumentError ParameterDistribution(param_dict_fail) # not right constraint type
        param_dict_fail = Dict("distribution" => d, "constraint" => [1, no_constraint(), 3], "name" => name)
        @test_throws ArgumentError ParameterDistribution(param_dict_fail) # not right constraint type
        param_dict_fail = Dict("distribution" => d, "constraint" => c, "name" => 1)
        @test_throws ArgumentError ParameterDistribution(param_dict_fail) # not right name type
        param_dict_fail = Dict("distribution" => d, "constraint" => [no_constraint(), no_constraint()], "name" => name)
        @test_throws DimensionMismatch ParameterDistribution(param_dict_fail) # wrong number of constraints

        param_dict = Dict("distribution" => d, "constraint" => c, "name" => name)
        u = ParameterDistribution(param_dict)
        @test u.distribution == [d]
        @test u.constraint == [c]
        @test u.name == [name]

        c_arr = [no_constraint()]
        param_dict = Dict("distribution" => d, "constraint" => c_arr, "name" => name)
        u = ParameterDistribution(param_dict)
        @test u.constraint == [c]

        # Tests for the ParameterDistribution
        d = Parameterized(MvNormal(4, 0.1))
        c = [no_constraint(), bounded_below(-1.0), bounded_above(0.4), bounded(-0.1, 0.2)]
        name = "constrained_mvnormal"
        param_dict = Dict("distribution" => d, "constraint" => c, "name" => name)

        u = ParameterDistribution(param_dict)
        @test u.distribution == [d]
        @test u.constraint == c
        @test u.name == [name]
        param_dict_fail = param_dict
        param_dict_fail["constraint"] = c[1]
        @test_throws DimensionMismatch ParameterDistribution(param_dict_fail) # wrong number of constraints

        # Tests for the ParameterDistribution
        d1 = Parameterized(MvNormal(4, 0.1))
        c1 = [no_constraint(), bounded_below(-1.0), bounded_above(0.4), bounded(-0.1, 0.2)]
        name1 = "constrained_mvnormal"
        param_dict1 = Dict("distribution" => d1, "constraint" => c1, "name" => name1)

        d2 = Samples([1.0 3.0 5.0 7.0; 9.0 11.0 13.0 15.0])
        c2 = [bounded(10, 15), no_constraint()]
        name2 = "constrained_sampled"
        param_dict2 = Dict("distribution" => d2, "constraint" => c2, "name" => name2)

        param_dict_array = [param_dict1, param_dict2]
        u = ParameterDistribution(param_dict_array)
        @test u.distribution == [d1, d2]
        @test u.constraint == cat([c1, c2]..., dims = 1)
        @test u.name == [name1, name2]

    end

    @testset "getter functions" begin
        # setup for the tests:
        d1 = Parameterized(MvNormal(4, 0.1))
        c1 = [no_constraint(), bounded_below(-1.0), bounded_above(0.4), bounded(-0.1, 0.2)]
        name1 = "constrained_mvnormal"
        u1 = ParameterDistribution(d1, c1, name1)

        d2 = Samples([1 2 3 4])
        c2 = [bounded(10, 15)]
        name2 = "constrained_sampled"
        u2 = ParameterDistribution(d2, c2, name2)

        d3 = VectorOfParameterized(repeat([Beta(2, 2)], 3))
        c3 = repeat([no_constraint()], 3)
        name3 = "vector_beta"
        u3 = ParameterDistribution(d3, c3, name3)

        u = combine_distributions([u1, u2, u3])

        # Test for get_dimension(s)
        @test ndims(u1) == 4
        @test ndims(u2) == 1
        @test ndims(u) == 8

        @test get_dimensions(u1) == [4]
        @test get_dimensions(u2) == [1]
        @test get_dimensions(u) == [4, 1, 3]

        # Tests for get_name
        @test get_name(u1) == [name1]
        @test get_name(u) == [name1, name2, name3]

        # Tests for get_n_samples
        @test typeof(get_n_samples(u)[name1]) <: String
        @test typeof(get_n_samples(u)[name3]) <: String
        @test get_n_samples(u)[name2] == 4

        # Tests for get_distribution
        @test get_distribution(d1) == MvNormal(4, 0.1)
        @test get_distribution(u1)[name1] == MvNormal(4, 0.1)
        @test typeof(get_distribution(d2)) == Array{Int64, 2}
        @test typeof(get_distribution(u2)[name2]) == Array{Int64, 2}

        @test get_distribution(d3) == repeat([Beta(2, 2)], 3)
        @test get_distribution(u3)[name3] == repeat([Beta(2, 2)], 3)

        d = get_distribution(u)
        @test d[name1] == MvNormal(4, 0.1)
        @test typeof(d[name2]) == Array{Int64, 2}

        # Test for get_all_constraints
        @test get_all_constraints(u) == cat([c1, c2, c3]..., dims = 1)
    end

    @testset "statistics functions" begin

        # setup for the tests:
        d1 = Parameterized(MvNormal(4, 0.1))
        c1 = [no_constraint(), bounded_below(-1.0), bounded_above(0.4), bounded(-0.1, 0.2)]
        name1 = "constrained_mvnormal"

        u1 = ParameterDistribution(d1, c1, name1)

        d2 = Samples([1 2 3 4])
        c2 = [bounded(10, 15)]
        name2 = "constrained_sampled"
        u2 = ParameterDistribution(d2, c2, name2)

        d3 = VectorOfParameterized([Beta(2, 2), MvNormal(2, 0.1)])
        c3 = [no_constraint(), no_constraint(), no_constraint()] # d3 has 3 dimensions
        name3 = "unconstrained_beta_and_MvN"
        u3 = ParameterDistribution(d3, c3, name3)

        u = combine_distributions([u1, u2])

        d4 = Samples([1 2 3 4 5 6 7 8; 8 7 6 5 4 3 2 1])
        c4 = [no_constraint(), no_constraint()]
        name4 = "constrained_MVsampled"
        u4 = ParameterDistribution(d4, c4, name4)

        v = combine_distributions([u1, u2, u3, u4])

        # Tests for sample distribution
        seed = 2020
        Random.seed!(seed)
        s1 = rand(MvNormal(4, 0.1), 1)
        Random.seed!(seed)
        @test sample(u1) == s1

        Random.seed!(seed)
        s1 = rand(MvNormal(4, 0.1), 3)
        Random.seed!(seed)
        @test sample(u1, 3) == s1

        Random.seed!(seed)
        idx = StatsBase.sample(collect(1:size(d2.distribution_samples)[2]), 1)
        s2 = d2.distribution_samples[:, idx]
        Random.seed!(seed)
        @test sample(u2) == s2

        Random.seed!(seed)
        idx = StatsBase.sample(collect(1:size(d2.distribution_samples)[2]), 3)
        s2 = d2.distribution_samples[:, idx]
        Random.seed!(seed)
        @test sample(u2, 3) == s2

        Random.seed!(seed)
        s3 = zeros(3, 1)
        s3[1] = rand(Beta(2, 2))
        s3[2:3] = rand(MvNormal(2, 0.1))
        Random.seed!(seed)
        @test sample(u3) == s3

        Random.seed!(seed)
        s1 = sample(u1, 3)
        s2 = sample(u2, 3)
        s3 = sample(u3, 3)
        s4 = sample(u4, 3)
        Random.seed!(seed)
        s = sample(v, 3)
        @test s == cat([s1, s2, s3, s4]..., dims = 1)

        #Test for get_logpdf
        @test_throws ErrorException get_logpdf(u, zeros(ndims(u)))
        x_in_bd = [0.5, 0.5, 0.5]
        Random.seed!(seed)
        lpdf3 = sum([logpdf(Beta(2, 2), x_in_bd[1])[1], logpdf(MvNormal(2, 0.1), x_in_bd[2:3])[1]]) #throws deprecated warning without "."

        Random.seed!(seed)
        @test isapprox(get_logpdf(u3, x_in_bd) - lpdf3, 0.0; atol = 1e-6)
        @test_throws DimensionMismatch get_logpdf(u3, [0.5, 0.5])

        #Test for cov, var
        block_cov = cat([cov(d1), var(d2), cov(d3), cov(d4)]..., dims = (1, 2))
        @test isapprox(cov(v) - block_cov, zeros(ndims(v), ndims(v)); atol = 1e-6)
        block_var = [block_cov[i, i] for i in 1:size(block_cov)[1]]
        @test isapprox(var(v) - block_var, zeros(ndims(v)); atol = 1e-6)

        #Test for mean
        means = cat([mean(d1), mean(d2), mean(d3), mean(d4)]..., dims = 1)
        @test isapprox(mean(v) - means, zeros(ndims(v)); atol = 1e-6)

    end

    @testset "statistics functions: explict RNG" begin

        # setup for the tests:
        rng_seed = 1234
        test_d = MvNormal(4, 0.1)
        test_d3a = Beta(2, 2)
        test_d3b = MvNormal(2, 0.1)
        d0 = Parameterized(test_d)

        d1 = Parameterized(test_d)
        c1 = [no_constraint(), bounded_below(-1.0), bounded_above(0.4), bounded(-0.1, 0.2)]
        name1 = "constrained_mvnormal"
        u1 = ParameterDistribution(d1, c1, name1)

        d2 = Samples([1 2 3 4])
        c2 = [bounded(10, 15)]
        name2 = "constrained_sampled"
        u2 = ParameterDistribution(d2, c2, name2)

        d3 = VectorOfParameterized([Beta(2, 2), MvNormal(2, 0.1)])
        c3 = repeat([no_constraint()], 3)
        name3 = "beta_and_normal"
        u3 = ParameterDistribution(d3, c3, name3)

        # Tests for sample distribution
        rng1 = Random.MersenneTwister(rng_seed)
        @test sample(copy(rng1), d0) == rand(copy(rng1), test_d, 1)
        @test sample(copy(rng1), d0, 3) == rand(copy(rng1), test_d, 3)

        @test sample(copy(rng1), u1) == rand(copy(rng1), test_d, 1)
        @test sample(copy(rng1), u1, 3) == rand(copy(rng1), test_d, 3)

        idx = StatsBase.sample(copy(rng1), collect(1:size(d2.distribution_samples)[2]), 1)
        s2 = d2.distribution_samples[:, idx]
        @test sample(copy(rng1), u2) == s2
        @test sample(copy(rng1), u2, 1) == s2
        @test sample(copy(rng1), d2) == s2
        @test sample(copy(rng1), d2, 1) == s2

        rng1_new = copy(rng1) # (copy here, as internally u3 samples rng1 sequentially)
        @test sample(copy(rng1), u3) == cat([rand(rng1_new, test_d3a, 1), rand(rng1_new, test_d3b, 1)]..., dims = 1)
        rng1_new = copy(rng1)
        @test sample(copy(rng1), u3, 3) ==
              cat([reshape(rand(rng1_new, test_d3a, 3), :, 3), rand(rng1_new, test_d3b, 3)]..., dims = 1)
        rng1_new = copy(rng1)
        @test sample(copy(rng1), d3) == cat([rand(rng1_new, test_d3a, 1), rand(rng1_new, test_d3b, 1)]..., dims = 1)
        rng1_new = copy(rng1)
        @test sample(copy(rng1), d3, 1) == cat([rand(rng1_new, test_d3a, 1), rand(rng1_new, test_d3b, 1)]..., dims = 1)

        # try it again with different RNG; use StableRNG since Random doesn't provide a
        # second seedable algorithm on julia <=1.7
        rng2 = StableRNG(rng_seed)
        @test sample(copy(rng2), d0) == rand(copy(rng2), test_d, 1)
        @test sample(copy(rng2), d0, 3) == rand(copy(rng2), test_d, 3)

        @test sample(copy(rng2), u1) == rand(copy(rng2), test_d, 1)
        @test sample(copy(rng2), u1, 3) == rand(copy(rng2), test_d, 3)

        idx = StatsBase.sample(copy(rng2), collect(1:size(d2.distribution_samples)[2]), 1)
        s2 = d2.distribution_samples[:, idx]
        @test sample(copy(rng2), u2) == s2

        rng2_new = copy(rng2)
        @test sample(copy(rng2), u3) == cat([rand(rng2_new, test_d3a, 1), rand(rng2_new, test_d3b, 1)]..., dims = 1)
        rng2_new = copy(rng2)
        @test sample(copy(rng2), u3, 3) ==
              cat([reshape(rand(rng2_new, test_d3a, 3), :, 3), rand(rng2_new, test_d3b, 3)]..., dims = 1)

        # test that optional parameter defaults to Random.GLOBAL_RNG, for all methods.
        # reset the global seed instead of copying the rng object's state
        rng_seed = 2468
        Random.seed!(rng_seed)
        test_lhs = sample(d0)
        Random.seed!(rng_seed)
        @test test_lhs == rand(test_d, 1)

        Random.seed!(rng_seed)
        test_lhs = sample(d0, 3)
        Random.seed!(rng_seed)
        @test test_lhs == rand(test_d, 3)

        Random.seed!(rng_seed)
        test_lhs = sample(u1)
        Random.seed!(rng_seed)
        @test test_lhs == rand(test_d, 1)

        Random.seed!(rng_seed)
        test_lhs = sample(u1, 3)
        Random.seed!(rng_seed)
        @test test_lhs == rand(test_d, 3)

        Random.seed!(rng_seed)
        test_lhs = sample(d0)
        Random.seed!(rng_seed)
        @test test_lhs == rand(test_d, 1)

        Random.seed!(rng_seed)
        idx = StatsBase.sample(collect(1:size(d2.distribution_samples)[2]), 1)
        test_lhs = d2.distribution_samples[:, idx]
        Random.seed!(rng_seed)
        @test test_lhs == sample(u2)

        Random.seed!(rng_seed)
        test_lhs = sample(d3)
        Random.seed!(rng_seed)
        @test test_lhs == cat([rand(test_d3a, 1), rand(test_d3b, 1)]..., dims = 1)

    end

    @testset "transform functions" begin
        #setup for the tests
        tol = 1e-8
        d1 = Parameterized(MvNormal(4, 0.1))
        c1 = [no_constraint(), bounded_below(-1.0), bounded_above(0.4), bounded(-0.1, 0.2)]
        name1 = "constrained_mvnormal"
        param_dict1 = Dict("distribution" => d1, "constraint" => c1, "name" => name1)
        u1 = ParameterDistribution(param_dict1)

        d2 = Samples([1.0 3.0 5.0 7.0; 9.0 11.0 13.0 15.0])
        c2 = [bounded(10, 15), no_constraint()]
        name2 = "constrained_sampled"
        param_dict2 = Dict("distribution" => d2, "constraint" => c2, "name" => name2)
        u2 = ParameterDistribution(param_dict2)

        param_dict = [param_dict1, param_dict2]
        u = ParameterDistribution(param_dict)

        x_unbd = rand(MvNormal(6, 3), 1000)  #6 x 1000
        # Tests for transforms
        x_real_constrained1 = mapslices(x -> transform_unconstrained_to_constrained(u1, x), x_unbd[1:4, :]; dims = 1)
        @test isapprox(
            x_unbd[1:4, :] -
            mapslices(x -> transform_constrained_to_unconstrained(u1, x), x_real_constrained1; dims = 1),
            zeros(size(x_unbd[1:4, :]));
            atol = tol,
        )

        x_real_constrained2 = mapslices(x -> transform_unconstrained_to_constrained(u2, x), x_unbd[5:6, :]; dims = 1)
        @test isapprox(
            x_unbd[5:6, :] -
            mapslices(x -> transform_constrained_to_unconstrained(u2, x), x_real_constrained2; dims = 1),
            zeros(size(x_unbd[5:6, :]));
            atol = tol,
        )

        x_real = mapslices(x -> transform_unconstrained_to_constrained(u, x), x_unbd; dims = 1)
        x_unbd_tmp = mapslices(x -> transform_constrained_to_unconstrained(u, x), x_real; dims = 1)
        @test isapprox(x_unbd - x_unbd_tmp, zeros(size(x_unbd)); atol = tol)

        # Tests transforms for other input structures
        @test isapprox(transform_unconstrained_to_constrained(u1, x_unbd[1:4, :]), x_real_constrained1; atol = tol)
        @test isapprox(
            x_unbd[1:4, :] - transform_constrained_to_unconstrained(u1, x_real_constrained1),
            zeros(size(x_unbd[1:4, :]));
            atol = tol,
        )
        @test isapprox(transform_unconstrained_to_constrained(u2, x_unbd[5:6, :]), x_real_constrained2; atol = tol)
        @test isapprox(
            x_unbd[5:6, :] - transform_constrained_to_unconstrained(u2, x_real_constrained2),
            zeros(size(x_unbd[5:6, :]));
            atol = tol,
        )
        @test isapprox(
            transform_unconstrained_to_constrained(u2, [x_unbd[5:6, :], x_unbd[5:6, :]]),
            [x_real_constrained2, x_real_constrained2];
            atol = tol,
        )
    end
end


text:
__map(f, itr; kwargs...) =
    tcollect(Map(f), itr; basesize = default_basesize(itr), kwargs...)
__map(f, itrs...; kwargs...) =
    tcollect(MapSplat(f), zip(itrs...); basesize = default_basesize(itrs[1]), kwargs...)

reshape_as(ys, xs) = reshape_as(ys, xs, IteratorSize(xs))
reshape_as(ys, xs, ::IteratorSize) = ys
reshape_as(ys, xs, ::HasShape) = reshape(ys, size(xs))
reshape_as(::Empty{T}, xs, isize::HasShape) where {T<:AbstractVector} =
    reshape_as(T(undef, length(xs)), xs, isize)

function _map(f, itr, itrs...; kwargs...)
    ys = __map(f, itr, itrs...; kwargs...)
    isempty(ys) && return map(f, itr, itrs...)
    return reshape_as(ys, itr)
end

ThreadsX.map(f, itr, itrs...; kwargs...) = _map(f, itr, itrs...; kwargs...)

function ThreadsX.map(
    f,
    array::AbstractArray{<:Any,N},
    arrays::AbstractArray{<:Any,N}...;
    kwargs...,
) where {N}
    dims = size(array)
    if !all(a -> size(a) == dims, arrays)
        throw(ArgumentError("shape of arrays do not match"))
    end
    output = _map(f, array, arrays...; kwargs...)
    return reshape(output, dims)
end

function ThreadsX.map!(f, dest, array, arrays...; kw...)
    @inline map_body!(y, xs...) = y[] = f(xs...)
    ThreadsX.foreach(map_body!, referenceable(dest), array, arrays...; kw...)
    return dest
end

struct ConvertTo{T} end
(::ConvertTo{T})(x) where {T} = convert(T, x)

ThreadsX.collect(::Type{T}, itr; kwargs...) where {T} = reshape_as(
    tcopy(Map(ConvertTo{T}()), Vector{T}, itr; basesize = default_basesize(itr), kwargs...),
    itr,
)

ThreadsX.collect(itr; kwargs...) =
    reshape_as(tcollect(itr; basesize = default_basesize(itr), kwargs...), itr)

ThreadsX.mapi(f, itr; kwargs...) = reshape_as(
    itr |>
    NondeterministicThreading(; kwargs...) |>
    Map(f) |>
    Map(SingletonVector) |>
    foldxl(append!!; init = EmptyVector()),
    itr,
)

ThreadsX.mapi(f, itr1, itrs...; kwargs...) =
    ThreadsX.mapi(Base.splat(f), zip(itr1, itrs...); kwargs...)

ext: jl


text:
# SAM Meta-Information
# ====================

type SAMMetaInfo
    # data and filled range
    data::Vector{UInt8}
    filled::UnitRange{Int}
    # indexes
    tag::UnitRange{Int}
    val::UnitRange{Int}
    dictkey::Vector{UnitRange{Int}}
    dictval::Vector{UnitRange{Int}}
end

function SAMMetaInfo(data::Vector{UInt8}=UInt8[])
    metainfo = SAMMetaInfo(data, 1:0, 1:0, 1:0, UnitRange{Int}[], UnitRange{Int}[])
    if !isempty(data)
        index!(metainfo)
    end
    return metainfo
end

function SAMMetaInfo(data::AbstractString)
    return SAMMetaInfo(convert(Vector{UInt8}, data))
end

function initialize!(metainfo::SAMMetaInfo)
    metainfo.filled = 1:0
    metainfo.tag = 1:0
    metainfo.val = 1:0
    empty!(metainfo.dictkey)
    empty!(metainfo.dictval)
    return metainfo
end

function isfilled(metainfo::SAMMetaInfo)
    return !isempty(metainfo.filled)
end

function datarange(metainfo::SAMMetaInfo)
    return metainfo.filled
end

function checkfilled(metainfo::SAMMetaInfo)
    if !isfilled(metainfo)
        throw(ArgumentError("unfilled SAM metainfo"))
    end
end

function SAMMetaInfo(tag::AbstractString, value)
    buf = IOBuffer()
    if tag == "CO"  # comment
        if !isa(value, AbstractString)
            throw(ArgumentError("value must be a string"))
        end
        write(buf, "@CO\t", value)
    elseif ismatch(r"[A-Z][A-Z]", tag)
        print(buf, '@', tag)
        for (key, val) in value
            print(buf, '\t', key, ':', val)
        end
    else
        throw(ArgumentError("tag must match r\"[A-Z][A-Z]\""))
    end
    return SAMMetaInfo(takebuf_array(buf))
end

function Base.:(==)(metainfo1::SAMMetaInfo, metainfo2::SAMMetaInfo)
    if isfilled(metainfo1) == isfilled(metainfo2) == true
        r1 = datarange(metainfo1)
        r2 = datarange(metainfo2)
        return length(r1) == length(r2) && memcmp(pointer(metainfo1.data, first(r1)), pointer(metainfo2.data, first(r2)), length(r1)) == 0
    else
        return isfilled(metainfo1) == isfilled(metainfo2) == false
    end
end

function isequalkey(metainfo::SAMMetaInfo, key::AbstractString)
    if !isfilled(metainfo) || sizeof(key) != 2
        return false
    end
    k1, k2 = UInt8(key[1]), UInt8(key[2])
    return metainfo.data[metainfo.tag[1]] == k1 && metainfo.data[metainfo.tag[2]] == k2
end

function iscomment(metainfo::SAMMetaInfo)
    return isequalkey(metainfo, "CO")
end

function Bio.metainfotag(metainfo::SAMMetaInfo)
    checkfilled(metainfo)
    return String(metainfo.data[metainfo.tag])
end

function Bio.metainfoval(metainfo::SAMMetaInfo)
    checkfilled(metainfo)
    return String(metainfo.data[metainfo.val])
end

function Base.keys(metainfo::SAMMetaInfo)
    checkfilled(metainfo)
    if iscomment(metainfo)
        throw(ArgumentError("not a dictionary"))
    end
    return [String(metainfo.data[r]) for r in metainfo.dictkey]
end

function Base.values(metainfo::SAMMetaInfo)
    checkfilled(metainfo)
    if iscomment(metainfo)
        throw(ArgumentError("not a dictionary"))
    end
    return [String(metainfo.data[r]) for r in metainfo.dictval]
end

function Base.haskey(metainfo::SAMMetaInfo, key::AbstractString)
    return findkey(metainfo, key) > 0
end

function Base.getindex(metainfo::SAMMetaInfo, key::AbstractString)
    i = findkey(metainfo, key)
    if i == 0
        throw(KeyError(key))
    end
    return String(metainfo.data[metainfo.dictval[i]])
end

function findkey(metainfo::SAMMetaInfo, key::AbstractString)
    checkfilled(metainfo)
    if sizeof(key) != 2
        return 0
    end
    t1, t2 = UInt8(key[1]), UInt8(key[2])
    for (i, k) in enumerate(metainfo.dictkey)
        if metainfo.data[first(k)] == t1 && metainfo.data[first(k)+1] == t2
            return i
        end
    end
    return 0
end

function Base.show(io::IO, metainfo::SAMMetaInfo)
    print(io, summary(metainfo), ':')
    if isfilled(metainfo)
        println(io)
        println(io, "    tag: ", metainfotag(metainfo))
          print(io, "  value:")
        if !iscomment(metainfo)
            for (key, val) in zip(keys(metainfo), values(metainfo))
                print(io, ' ', key, '=', val)
            end
        else
            print(io, ' ', metainfoval(metainfo))
        end
    else
        print(io, " <not filled>")
    end
end

function Base.print(io::IO, metainfo::SAMMetaInfo)
    write(io, metainfo)
    return nothing
end

function Base.write(io::IO, metainfo::SAMMetaInfo)
    checkfilled(metainfo)
    r = datarange(metainfo)
    return unsafe_write(io, pointer(metainfo.data, first(r)), length(r))
end


text:
module TransportUndirectedModule

import ..GridWorlds as GW
import Random
import ReinforcementLearningBase as RLBase

#####
##### game logic
#####

const NUM_OBJECTS = 4
const AGENT = 1
const WALL = 2
const GEM = 3
const TARGET = 4
const NUM_ACTIONS = 6

mutable struct TransportUndirected{R, RNG} <: GW.AbstractGridWorld
    tile_map::BitArray{3}
    agent_position::CartesianIndex{2}
    reward::R
    rng::RNG
    done::Bool
    terminal_reward::R
    gem_position::CartesianIndex{2}
    target_position::CartesianIndex{2}
    has_gem::Bool
end

function TransportUndirected(; R = Float32, height = 8, width = 8, rng = Random.GLOBAL_RNG)
    tile_map = falses(NUM_OBJECTS, height, width)

    tile_map[WALL, 1, :] .= true
    tile_map[WALL, height, :] .= true
    tile_map[WALL, :, 1] .= true
    tile_map[WALL, :, width] .= true

    agent_position = GW.sample_empty_position(rng, tile_map)
    tile_map[AGENT, agent_position] = true

    gem_position = GW.sample_empty_position(rng, tile_map)
    tile_map[GEM, gem_position] = true

    target_position = GW.sample_empty_position(rng, tile_map)
    tile_map[TARGET, target_position] = true

    reward = zero(R)
    done = false
    terminal_reward = one(R)
    has_gem = false

    env = TransportUndirected(tile_map, agent_position, reward, rng, done, terminal_reward, gem_position, target_position, has_gem)

    GW.reset!(env)

    return env
end

function GW.reset!(env::TransportUndirected)
    tile_map = env.tile_map
    rng = env.rng

    tile_map[AGENT, env.agent_position] = false
    tile_map[GEM, env.gem_position] = false
    tile_map[TARGET, env.target_position] = false

    new_agent_position = GW.sample_empty_position(rng, tile_map)
    env.agent_position = new_agent_position
    tile_map[AGENT, new_agent_position] = true

    new_gem_position = GW.sample_empty_position(rng, tile_map)
    env.gem_position = new_gem_position
    tile_map[GEM, new_gem_position] = true

    new_target_position = GW.sample_empty_position(rng, tile_map)
    env.target_position = new_target_position
    tile_map[TARGET, new_target_position] = true

    env.reward = zero(env.reward)
    env.done = false
    env.has_gem = false

    return nothing
end

function GW.act!(env::TransportUndirected, action)
    @assert action in Base.OneTo(NUM_ACTIONS) "Invalid action $(action). Action must be in Base.OneTo($(NUM_ACTIONS))"

    tile_map = env.tile_map
    agent_position = env.agent_position

    if action in Base.OneTo(4)
        if action == 1
            new_agent_position = GW.move_up(agent_position)
        elseif action == 2
            new_agent_position = GW.move_down(agent_position)
        elseif action == 3
            new_agent_position = GW.move_left(agent_position)
        else
            new_agent_position = GW.move_right(agent_position)
        end

        if !tile_map[WALL, new_agent_position]
            tile_map[AGENT, agent_position] = false
            env.agent_position = new_agent_position
            tile_map[AGENT, new_agent_position] = true
        end
    elseif action == 5
        if tile_map[GEM, agent_position]
            tile_map[GEM, agent_position] = false
            env.has_gem = true
        end
    else
        if env.has_gem
            env.has_gem = false
            env.gem_position = agent_position
            tile_map[GEM, agent_position] = true
        end
    end

    if tile_map[GEM, env.target_position]
        env.reward = env.terminal_reward
        env.done = true
    else
        env.reward = zero(env.reward)
        env.done = false
    end

    return nothing
end

#####
##### miscellaneous
#####

GW.get_height(env::TransportUndirected) = size(env.tile_map, 2)
GW.get_width(env::TransportUndirected) = size(env.tile_map, 3)

GW.get_action_names(env::TransportUndirected) = (:MOVE_UP, :MOVE_DOWN, :MOVE_LEFT, :MOVE_RIGHT, :PICK_UP, :DROP)
GW.get_object_names(env::TransportUndirected) = (:AGENT, :WALL, :GEM, :TARGET)

function GW.get_pretty_tile_map(env::TransportUndirected, position::CartesianIndex{2})
    characters = ('☻', '█', '♦', '✖', '⋅')

    object = findfirst(@view env.tile_map[:, position])
    if isnothing(object)
        return characters[end]
    else
        return characters[object]
    end
end

function GW.get_pretty_sub_tile_map(env::TransportUndirected, window_size, position::CartesianIndex{2})
    tile_map = env.tile_map
    agent_position = env.agent_position

    characters = ('☻', '█', '♦', '✖', '⋅')

    sub_tile_map = GW.get_sub_tile_map(tile_map, agent_position, window_size)

    object = findfirst(@view sub_tile_map[:, position])
    if isnothing(object)
        return characters[end]
    else
        return characters[object]
    end
end

function Base.show(io::IO, ::MIME"text/plain", env::TransportUndirected)
    str = "tile_map:\n"
    str = str * GW.get_pretty_tile_map(env)
    str = str * "\nsub_tile_map:\n"
    str = str * GW.get_pretty_sub_tile_map(env, GW.get_window_size(env))
    str = str * "\nreward: $(env.reward)"
    str = str * "\ndone: $(env.done)"
    str = str * "\nhas_gem: $(env.has_gem)"
    str = str * "\naction_names: $(GW.get_action_names(env))"
    str = str * "\nobject_names: $(GW.get_object_names(env))"
    print(io, str)
    return nothing
end

GW.get_action_keys(env::TransportUndirected) = ('w', 's', 'a', 'd', 'p', 'l')

#####
##### RLBase API
#####

RLBase.StateStyle(env::GW.RLBaseEnv{E}) where {E <: TransportUndirected} = RLBase.InternalState{Any}()
RLBase.state_space(env::GW.RLBaseEnv{E}, ::RLBase.InternalState) where {E <: TransportUndirected} = nothing
RLBase.state(env::GW.RLBaseEnv{E}, ::RLBase.InternalState) where {E <: TransportUndirected} = env.env.tile_map

RLBase.reset!(env::GW.RLBaseEnv{E}) where {E <: TransportUndirected} = GW.reset!(env.env)

RLBase.action_space(env::GW.RLBaseEnv{E}) where {E <: TransportUndirected} = Base.OneTo(NUM_ACTIONS)
(env::GW.RLBaseEnv{E})(action) where {E <: TransportUndirected} = GW.act!(env.env, action)

RLBase.reward(env::GW.RLBaseEnv{E}) where {E <: TransportUndirected} = env.env.reward
RLBase.is_terminated(env::GW.RLBaseEnv{E}) where {E <: TransportUndirected} = env.env.done

end # module


text:
function(a::Vector{Float64}, b::Vector{Float64}, c::Vector{Float64},
         d::Vector{Float64}, soln::Vector{Float64})
    # Setting initial elements
    c[0] = c[0] / b[0]
    d[0] = d[0] / b[0]

    for i = 1:n
        # Scale factor is for c and d
        scale = 1 / (b[i] - c[i-1]*a[i])
        c[i] = c[i] * scale
        d[i] = (d[i] - a[i] * d[i-1]) * scale
    end

    # Set the last solution for back-substitution
    soln[n-1] = d[n-1]

    # Back-substitution
    for i = n-2:0
        soln[i] = d[i] - c[i] * soln[i+1]
    end

end


text:
function drawstairs(n)
  buffer = IOBuffer()
  write(buffer,'I')
  for i in 2:n
    write(buffer,'\n',' '^(i-1),'I')
  end
  return String(take!(buffer))
end

max_stars_repo_name: UlrichBerntien/Codewars-Katas

max_issues_repo_name: UlrichBerntien/Codewars-Katas

max_forks_repo_name: UlrichBerntien/Codewars-Katas


text:
"""
    NaiveMultivariatePoissonProcess{R}

Homogeneous temporal multivariate Poisson process.

This naive implementation demonstrates the use of our general `TemporalPointProcess` interface.

# Fields
- `λ::Vector{R}`: event rates.
"""
struct NaiveMultivariatePoissonProcess{R} <: TemporalPointProcess{Int}
    λ::Vector{R}
end

intensity(pp::NaiveMultivariatePoissonProcess) = pp.λ

## Conversion utilities

NaiveMultivariatePoissonProcess{R}(nt::NamedTuple) where {R} =
    NaiveMultivariatePoissonProcess(nt.λ)

function build_transform(pp::NaiveMultivariatePoissonProcess)
    M = length(pp.λ)
    return as((λ = as(Vector, asℝ₊, M),))
end

## Intensity functions

function intensity(pp::NaiveMultivariatePoissonProcess, h::TemporalHistory, t, m)
    return pp.λ[m]
end

function mark_distribution(pp::NaiveMultivariatePoissonProcess, h::TemporalHistory, t)
    return Dists.Categorical(pp.λ / sum(pp.λ))
end

function ground_intensity(pp::NaiveMultivariatePoissonProcess, h::TemporalHistory, t)
    return sum(pp.λ)
end

function ground_intensity_bound(pp::NaiveMultivariatePoissonProcess, h::TemporalHistory, t)::Tuple{Float64, Float64}
    λg = ground_intensity(pp, h, t)
    return λg, Inf
end


text:
using TrackRunning
using Base.Test

# write your own tests here, run by: Pkg.test("TrackRunning")

try
    global hike = readrun(joinpath(Pkg.dir("TrackRunning"),
        "example", "data", "hiking.csv"))
  catch
    error("Failed to read 'hiking.csv'")
end

@testset "Run Tests" begin
    @test length(hike) == 24673
    @test length(hike.lap) == 24673
    @test length(hike.date) == 24673
    @test length(hike.lon) == 24673
    @test length(hike.lat) == 24673
    @test length(hike.ele) == 0
    @test length(hike.hr) == 0

    @test hike[end].date == [DateTime("2016-10-16T17:55:05")]
    @test hike[(end-1):end].date == [DateTime("2016-10-16T17:55:04");
                                     DateTime("2016-10-16T17:55:05")]

    @test start(hike) == 1
    @test next(hike, 1)[2] == 2
    @test done(hike, length(hike)) == false
    @test done(hike, length(hike) + 1) == true
end

sec = time_sec(hike)
@testset "Time Tests" begin
    @test length(sec) == length(hike.lon) - 1
    @test sec[1] ≈ 1.0
    @test sec[end] ≈ 24678.0
end

m = dist_m(hike)
@testset "Distance Tests" begin
    @test length(m) == length(hike.lon) - 1
    @test m[1] ≈ 1.2038457030465493
    @test m[end] ≈ 21865.947388700704
end

kmh = speed_kmh(hike)
@testset "Speed Tests" begin
    @test length(kmh) == length(hike.lon) - 1
    @test kmh[1] ≈ 4.3338445309675775
    @test kmh[end] ≈ 0.0
end

minkm = pace_minkm(hike)
@testset "Pace Tests" begin
    @test length(minkm) == length(hike.lon) - 1
    @test minkm[1] ≈ 13.844520626263526
    @test minkm[end] ≈ Inf
end

@test mapzoom(hike) == 13

mapcoor = mapprojection(hike)
@testset "MapCoordinates Tests" begin
    @test mapcoor.lon_center ≈ -0.01981505
    @test mapcoor.lat_center ≈ 42.6412988
    @test mapcoor.lon_left ≈ -0.07474669062500539
    @test mapcoor.lat_top ≈ 42.681693887264494
    @test mapcoor.lon_right ≈ 0.035116590624994615
    @test mapcoor.lat_bottom ≈ 42.600877469417284
    @test mapcoor.x_left ≈ 4094.299097528889
    @test mapcoor.y_top ≈ 3020.022953178298
    @test mapcoor.x_right ≈ 4096.7990975288885
    @test mapcoor.y_bottom ≈ 3022.522953178298
    @test mapcoor.zoom ≈ 13
end

mapfile = downloadmap(mapcoor)
@test isa(mapfile, AbstractString)


text:
include("core.jl")

function main()
    mat = CCVDenseMatrix(load("img_209.jpg"))
    cascade = scd_classifier_cascade_read(FACE_CASCADE_FILE)
    faces = scd_detect_objects(mat, [cascade])
                               # SCDParamType(5, 1, 4, CCVSize(250, 250)))
    length(faces)

end


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

_dot(v::AbstractVector{<:Real}, x::AbstractVecOrMat{<:Real}) = (v' * x)

# x and y may alias:
function householder_trafo!(y::AbstractVecOrMat{<:Real}, v::AbstractVector{<:Real}, x::AbstractVecOrMat{<:Real})
    k = _dot(v, x) / _dot(v, v)
    y .= muladd.(-2 .* k, v, x)
end

function householder_trafo(v::AbstractVector{T}, x::AbstractVecOrMat{U}) where {T<:Real,U<:Real}
    R = promote_type(T, U)
    y = similar(x, R)
    householder_trafo!(y, v, x)
    return y
end

function householder_trafo_pullback_v(v, x, ΔΩ)
    inrm = inv(norm(v))

    # Readable version:
    #=
    # w == normalize(v):
    w = inrm .* v
    # pullback for I - 2 * (w*w')) * x:
    ∂w = -2 .* (ΔΩ .* dot(w, x) .+ x .* dot(w, ΔΩ))
    # pullback for normalize(v):
    ∂v = inrm .* (∂w .- w .* dot(∂w, w))
    =#

    inrm_2 = inrm * inrm
    w_x = inrm * _dot(v, x)
    w_ΔΩ = inrm * _dot(v, ΔΩ)
    ∂w_v = sum(-2 .* v .* (w_ΔΩ .* x .+ w_x .* ΔΩ), dims = 1)
    sum(inrm .* (-2 .* (w_x .* ΔΩ  .+ w_ΔΩ .* x) .- inrm_2 .* ∂w_v .* v), dims = 2)
end

#householder_trafo_pullback_x(v, x, ΔΩ) = householder_trafo(v, ΔΩ)

function householder_trafo_pbacc_x!(∂x, v, x, ΔΩ)
    k = -2 .* _dot(v, ΔΩ) ./ _dot(v, v)
    ∂x += muladd.(k, v, ΔΩ)
end

function householder_trafo_pullback_x(v, x, ΔΩ)
    R = promote_type(eltype(v), eltype(ΔΩ))
    ∂x = fill!(similar(x, R), 0)
    householder_trafo_pbacc_x!(∂x, v, x, ΔΩ)
end

function rrule(::typeof(householder_trafo), v::AbstractVector{T}, x::AbstractVecOrMat{U}) where {T<:Real,U<:Real}
    y = householder_trafo(v, x)
    function householder_trafo_pullback(ΔΩ)
        ∂v = @thunk(householder_trafo_pullback_v(v, x, unthunk(ΔΩ)))
        ∂x = InplaceableThunk(
            ∂x -> householder_trafo_pbacc_x!(∂x, v, x, unthunk(ΔΩ)),
            @thunk(householder_trafo_pullback_x(v, x, unthunk(ΔΩ)))
        )
        (NoTangent(), ∂v, ∂x)
    end
    return y, householder_trafo_pullback
end

function chained_householder_trafo!(y::AbstractVecOrMat{<:Real}, V::AbstractMatrix{<:Real}, x::AbstractVecOrMat{<:Real})
    y .= x
    for i in axes(V, 2)
        v = view(V, :, i)
        householder_trafo!(y, v, y)
    end
    return y
end

function chained_householder_trafo(V::AbstractMatrix{T}, x::AbstractVecOrMat{U}) where {T<:Real,U<:Real}
    R = promote_type(T, U)
    y = similar(x, R)
    chained_householder_trafo!(y, V, x)
    return y
end

function chained_householder_trafo_pullback_V(V, x, y, ΔΩ)
    # @assert y == chained_householder_trafo(V, x)
    ∂V = similar(V)
    z = deepcopy(y)
    Δ = similar(x)
    Δ .= ΔΩ
    for i in reverse(axes(V, 2))
        v = view(V, :, i)
        ∂v = view(∂V, :, i)
        householder_trafo!(z, v, z)
        ∂v .= householder_trafo_pullback_v(v, z, Δ)
        householder_trafo!(Δ, v, Δ)
    end
    @assert z ≈ x
    return ∂V
end

function chained_householder_trafo_pullback_x(V, x, y, ΔΩ)
    # @assert y == chained_householder_trafo(V, x)
    ∂x = similar(x)
    ∂x .= ΔΩ
    for i in reverse(axes(V, 2))
        v = view(V, :, i)
        householder_trafo!(∂x, v, ∂x)
    end
    return ∂x
end

function rrule(::typeof(chained_householder_trafo), V::AbstractMatrix{T}, x::AbstractVecOrMat{U}) where {T<:Real,U<:Real}
    y = chained_householder_trafo(V, x)
    function householder_trafo_pullback(ΔΩ)
        ∂V = @thunk(chained_householder_trafo_pullback_V(V, x, y, unthunk(ΔΩ)))
        ∂x = @thunk(chained_householder_trafo_pullback_x(V, x, y, unthunk(ΔΩ)))
        (NoTangent(), ∂V, ∂x)
    end
    return y, householder_trafo_pullback
end

@with_kw struct HouseholderTrafo{T<:AbstractVecOrMat{<:Real}} <:Function
    V::T
end

# ToDo: Normalize in HouseholderTrafo as well?

_ht_normalize(V::AbstractVector{<:Real}) = normalize(V)
function _ht_normalize(V::AbstractMatrix{<:Real})
    V_norm = deepcopy(V)
    normalize!.(eachcol(V_norm))
    V_norm
end

function Base.convert(::Type{HouseholderTrafo}, nt::NamedTuple{(:V,),<:Tuple{AbstractVecOrMat{<:Real}}})
    HouseholderTrafo(_ht_normalize(nt.V))
end
Base.convert(::Type{NamedTuple}, x::HouseholderTrafo) = (V = x.V,)

Functors.functor(::Type{<:HouseholderTrafo}, x) = convert(NamedTuple, x), y -> convert(HouseholderTrafo, y)

Base.:(==)(a::HouseholderTrafo, b::HouseholderTrafo) = a.V == b.V
Base.isequal(a::HouseholderTrafo, b::HouseholderTrafo) = isequal(a.V, b.V)
Base.hash(x::HouseholderTrafo, h::UInt) = hash(x.V, hash(:HouseholderTrafo, hash(:EuclidianNormalizingFlows, h)))

inverse(f::HouseholderTrafo{<:AbstractVector}) = f
inverse(f::HouseholderTrafo{<:AbstractMatrix}) = HouseholderTrafo(reverse(f.V, dims = 2))

(f::HouseholderTrafo{<:AbstractVector})(x::AbstractVecOrMat{<:Real}) = householder_trafo(f.V, x)
(f::HouseholderTrafo{<:AbstractMatrix})(x::AbstractVecOrMat{<:Real}) = chained_householder_trafo(f.V, x)

with_logabsdet_jacobian(f::HouseholderTrafo, x::AbstractVector{T}) where {T<:Real} = (f(x), zero(T))
with_logabsdet_jacobian(f::HouseholderTrafo, x::AbstractMatrix{T}) where {T<:Real} = (f(x), similar_zeros(x, (size(x, 2),))')


text:
using Test
using PolynomialBasis
using ImplicitDomainQuadrature
# using Revise
using CutCell
include("useful_routines.jl")

polyorder = 1
numqp = 2

basis = TensorProductBasis(2, polyorder)
levelset = InterpolatingPolynomial(1, basis)
nf = CutCell.number_of_basis_functions(basis)
mesh = CutCell.Mesh([0.0, 0.0], [3.0, 1.0], [3, 1], nf)
nodalcoordinates = CutCell.nodal_coordinates(mesh)
nodalconnectivity = CutCell.nodal_connectivity(mesh)
cellmaps = CutCell.cell_maps(mesh)

normal = [1.0, 0.0]
x0 = [0.5, 0.0]
tol = 1e-3
perturbation = 0.0
levelsetcoeffs = plane_distance_function(nodalcoordinates, normal, x0)

cellsign = CutCell.cell_sign!(levelset, levelsetcoeffs, nodalconnectivity, tol, perturbation)
@test allequal(cellsign, [0, 1, 1])

normal = [1.0, 0.0]
x0 = [1.0, 0.0]
tol = 1e-3
perturbation = 1e-3
levelsetcoeffs = plane_distance_function(nodalcoordinates, normal, x0)
cellsign = CutCell.cell_sign!(levelset,levelsetcoeffs,nodalconnectivity,tol,perturbation)

@test allapprox(levelsetcoeffs[1:4],[-1.0+perturbation,-1.0+perturbation,perturbation, perturbation])
@test allapprox(levelsetcoeffs[[3,4,5,6]],[perturbation,perturbation,1.0,1.0])
@test allapprox(levelsetcoeffs[[5,6,7,8]],[1.,1.,2.,2.])

posactivenodeids = CutCell.active_node_ids(+1, cellsign, nodalconnectivity)
@test allequal(posactivenodeids, 1:8)
negactivenodeids = CutCell.active_node_ids(-1, cellsign, nodalconnectivity)
@test allequal(negactivenodeids, 1:4)

totalnumnodes = CutCell.number_of_nodes(mesh)
cutmeshnodeids =
    CutCell.cut_mesh_nodeids(posactivenodeids, negactivenodeids, totalnumnodes)
testcutmeshnodeids = [
    1 2 3 4 5 6 7 8
    9 10 11 12 0 0 0 0
]
@test allequal(cutmeshnodeids, testcutmeshnodeids)

cutmesh = CutCell.CutMesh(levelset, levelsetcoeffs, mesh)
@test allequal(CutCell.nodal_connectivity(cutmesh, +1, 1), [1, 2, 3, 4])
@test allequal(CutCell.nodal_connectivity(cutmesh, -1, 1), [9, 10, 11, 12])


text:
# This file contains the recipes for plotting collections of simplices filtered by birth
# time. This also includes filtrations.

"""
    FilteredFaces

Abstract type for [`FilteredVertices`](@ref), [`FilteredEdges`](@ref), and
[`FilteredTriangles`](@ref).

The structures contain plotting data and a vector of times. Indexing into them with a time
returns an object that can be plotted without conversion.
"""
abstract type FilteredFaces end

Base.firstindex(fs::FilteredFaces) = first(fs.times)
Base.lastindex(fs::FilteredFaces) = last(fs.times)
Base.length(fs::FilteredFaces) = length(fs.times)

function Base.show(io::IO, fs::F) where {F<:FilteredFaces}
    print(io, nameof(F), " t=[", firstindex(fs), ", ", lastindex(fs), "], n=", length(fs))
    return nothing
end

"""
    FilteredVertices{V, T}

See: [`FilteredFaces`](@ref).
"""
struct FilteredVertices{V,T} <: FilteredFaces
    vertices::Vector{V}
    times::Vector{T}
end
function Base.getindex(v::FilteredVertices, t)
    i = searchsortedlast(v.times, t)
    return v.vertices[1:i]
end

"""
    FilteredEdges{V, T}

See: [`FilteredFaces`](@ref).
"""
struct FilteredEdges{V,T} <: FilteredFaces
    edges::Vector{V}
    times::Vector{T}
end
function Base.getindex(v::FilteredEdges, t)
    i = searchsortedlast(v.times, t)
    return v.edges[1:(2i)]
end

"""
    FilteredTriangles{F, D, T}

See: [`FilteredFaces`](@ref).
"""
struct FilteredTriangles{F,D,T} <: FilteredFaces
    faces::Vector{F}
    data::D
    times::Vector{T}
end
function Base.getindex(m::FilteredTriangles, t)
    i = searchsortedlast(m.times, t)
    return GeometryBasics.Mesh(m.data, m.faces[1:i])
end

"""
    FilteredChain

Contains plotting data to plot a collection of simplices. To construct it, use
`FilteredChain(simplices_or_filtration, data)`. The `data` should be something Makie can
convert to points.

# Example

julia> using Ripserer

julia> simplices = [Simplex{2}(1, 1), Simplex{2}(2, 2)];

julia> data = [(0, 0), (0, 1), (1, 0), (1, 1)];

julia> fc = FilteredChain(simplices, data)
FilteredChain(nv=4, ne=5, nt=2)

"""
struct FilteredChain{V<:FilteredVertices,E<:FilteredEdges,T<:FilteredTriangles}
    vertices::V
    edges::E
    triangles::T
end
function FilteredChain(flt::AbstractFiltration, data)
    points = convert_arguments(Scatter, data)[1]

    vertex_sxs = sort!([simplex(flt, Val(0), (v,)) for v in Ripserer.vertices(flt)])
    vertex_times = birth.(vertex_sxs)
    vertices = points[only.(Ripserer.vertices.(vertex_sxs))]

    edge_sxs = sort!(Ripserer.edges(flt))
    edge_times = birth.(edge_sxs)
    edges = points[collect(Iterators.flatten(Ripserer.vertices.(edge_sxs)))]

    triangle_sxs = sort!(collect(Ripserer.columns_to_reduce(flt, edge_sxs)))
    triangle_times = birth.(triangle_sxs)
    triangle_faces = typeof(GLTriangleFace(1, 2, 3))[]
    for sx in triangle_sxs
        vs = Ripserer.vertices(sx)
        push!(triangle_faces, GLTriangleFace(vs[1], vs[2], vs[3]))
    end

    return FilteredChain(
        FilteredVertices(vertices, vertex_times),
        FilteredEdges(edges, edge_times),
        FilteredTriangles(triangle_faces, points, triangle_times),
    )
end

function FilteredChain(sxs, data)
    points = convert_arguments(Scatter, data)[1]
    sxs = sort(sxs)

    # *_idx is used to prevent plotting faces multiple times
    vertex_times = Float64[]
    vertices = eltype(points)[]
    vertex_idx = Dict{Int,Int}()
    edge_times = Float64[]
    edges = eltype(points)[]
    edge_idx = Dict{Tuple{Int,Int},Int}()
    triangle_times = Float64[]
    triangle_faces = typeof(GLTriangleFace(1, 2, 3))[]
    triangle_idx = Dict{Tuple{Int,Int,Int},Int}()

    for sx in sxs
        sx_vertices = Ripserer.vertices(sx)
        t = birth(sx)
        for v in sx_vertices
            if !haskey(vertex_idx, v)
                push!(vertices, points[v])
                push!(vertex_times, t)
                vertex_idx[v] = lastindex(vertex_times)
            else
                i = vertex_idx[v]
                vertex_times[i] = min(vertex_times[i], t)
            end
        end
        for (u, v) in IterTools.subsets(sx_vertices, Val(2))
            if !haskey(edge_idx, (u, v))
                append!(edges, (points[u], points[v]))
                push!(edge_times, t)
                edge_idx[(u, v)] = lastindex(edge_times)
            else
                i = edge_idx[(u, v)]
                edge_times[i] = min(edge_times[i], t)
            end
        end
        for (u, v, w) in IterTools.subsets(sx_vertices, Val(3))
            if !haskey(triangle_idx, (u, v, w))
                push!(triangle_faces, GLTriangleFace(u, v, w))
                push!(triangle_times, t)
                triangle_idx[(u, v, w)] = lastindex(triangle_times)
            else
                i = triangle_idx[(u, v, w)]
                triangle_times[i] = min(triangle_times[i], t)
            end
        end
    end

    return FilteredChain(
        FilteredVertices(vertices, vertex_times),
        FilteredEdges(edges, edge_times),
        FilteredTriangles(triangle_faces, points, triangle_times),
    )
end

function FilteredChain(sx::AbstractSimplex, data)
    return FilteredChain([sx], data)
end
function FilteredChain(sx::AbstractChainElement, data)
    return FilteredChain([simplex(sx)], data)
end
function FilteredChain(sx::AbstractVector{<:AbstractChainElement}, data)
    return FilteredChain(simplex.(sx), data)
end

function Base.show(io::IO, chain::FilteredChain)
    nv = length(chain.vertices)
    ne = length(chain.edges)
    nt = length(chain.triangles)
    print(io, "FilteredChain(nv=$nv, ne=$ne, nt=$nt)")
    return nothing
end

@recipe(ChainPlot, chain) do scene
    Theme(;
        vertexcolor=1,
        edgecolor=:black,
        trianglecolor=2,
        shading=false,
        transparency=false,
        palette=DEFAULT_PALETTE,
        markersize=10,
        linewidth=1,
        triangles=true,
        edges=true,
        time=Inf,
    )
end

function AbstractPlotting.plot!(p::ChainPlot)
    chain = to_value(p[:chain])
    time = p[:time]
    draw_triangles = @lift length(chain.triangles) > 0 ? $(p[:triangles]) : false
    draw_edges = @lift length(chain.edges) > 0 ? $(p[:edges]) : false

    drawn_vertices = @lift chain.vertices[$time]
    drawn_edges = @lift chain.edges[$draw_edges ? $time : -Inf]
    drawn_triangles = @lift chain.triangles[$draw_triangles ? $time : -Inf]

    mesh!(
        p,
        drawn_triangles;
        color=get_color(p, p[:trianglecolor]),
        shading=p[:shading],
        transparency=p[:transparency],
    )
    linesegments!(
        p,
        drawn_edges;
        color=get_color(p, p[:edgecolor]),
        shading=p[:shading],
        linewidth=p[:linewidth],
    )
    return scatter!(
        p, drawn_vertices; color=get_color(p, p[:vertexcolor]), markersize=p[:markersize]
    )
end

function _no_data_error(arg::T) where {T}
    return throw(ArgumentError("No data provided. To plot $T, use `plot(::$T, data)"))
end

for T in (
    AbstractFiltration,
    AbstractSimplex,
    AbstractChainElement,
    AbstractVector{<:AbstractSimplex},
    AbstractVector{<:AbstractChainElement},
)
    @eval begin
        function AbstractPlotting.plottype(::$T, ::AbstractVector)
            return ChainPlot
        end
        function AbstractPlotting.convert_arguments(::Type{<:ChainPlot}, x::$T, data)
            return (FilteredChain(x, data),)
        end
        function AbstractPlotting.plottype(x::$T)
            return _no_data_error(x)
        end
    end
end

"""
    chainplot(x, data)

Plot a simplex, collection of simplices, or filtration. Data is used to determine vertex
locations. Setting the `time` filters plotted simplices by birth.

Names or `Int`s can be used to set colors.

# Attributes

    vertexcolor   1
    edgecolor     :black
    trianglecolor 2
    shading       false
    transparency  false
    palette       :default
    markersize    10
    linewidth     1
    triangles     true
    edges         true
    time          Inf

"""
chainplot

"""
    chainplot!(x, data)

Plot a simplex, collection of simplices, or filtration. Data is used to determine vertex
locations. Setting the `time` filters plotted simplices by birth.

Names or `Int`s can be used to set colors.

# Attributes

    vertexcolor   1
    edgecolor     :black
    trianglecolor 2
    shading       false
    transparency  false
    palette       :default
    markersize    10
    linewidth     1
    triangles     true
    edges         true
    time          Inf

"""
chainplot!


text:
using Test
using CGcoefficient
using CGcoefficient:check_couple,check_jm,HalfInt

include("test_special.jl")
include("test_with_gsl.jl")

@testset "special condition: CG" begin test_special_CG(1 // 2:1 // 2:10) end
@testset "special condition: 3j" begin test_special_3j(1 // 2:1 // 2:10) end
@testset "special condition: 6j" begin test_special_6j(1 // 2:1 // 2:5) end
@testset "special condition: 9j" begin test_special_9j(1 // 2:1 // 2:5) end

try
    gsl3j(1, 1, 1, 0, 0, 0)
catch err
    if isa(err, ErrorException)
        println(err.msg)
        println("If you want to run test with libgsl, please download gsl library and rerun the test.")
        exit()
    end
    error(err)
end

@testset "test with gsl: 3j" begin test_3j_with_gsl(1:5) end
@testset "test_with gsl: 6j" begin test_6j_with_gsl(1:5) end
@testset "test_with gsl: 9j" begin test_9j_with_gsl(1:5) end


text:
##############################################################################
##
## Graveyard of Alternative Functions
##
## The proper setting of epsilon depends on whether we're doing central
##  differencing or forward differencing. See, for example, section 5.7 of
##  Numerical Recipes.
##
##############################################################################

##############################################################################
##
## Derivative of f: C -> R
##
##############################################################################

macro forwardrule(x, e)
    x, e = esc(x), esc(e)
    quote
        $e = sqrt(eps(eltype($x))) * max(one(eltype($x)), abs($x))
    end
end

macro centralrule(x, e)
    x, e = esc(x), esc(e)
    quote
        $e = cbrt(eps(eltype($x))) * max(one(eltype($x)), abs($x))
    end
end

macro hessianrule(x, e)
    x, e = esc(x), esc(e)
    quote
        $e = eps(eltype($x))^(1/4) * max(one(eltype($x)), abs($x))
    end
end

macro complexrule(x, e)
    x, e = esc(x), esc(e)
    quote
        $e = eps($x)
    end
end

function finite_difference{T <: Number}(f::Function,
                                        x::T,
                                        dtype::Symbol = :central)
    if dtype == :forward
        @forwardrule x epsilon
        xplusdx = x + epsilon
        return (f(xplusdx) - f(x)) / epsilon
    elseif dtype == :central
        @centralrule x epsilon
        xplusdx, xminusdx = x + epsilon, x - epsilon
        return (f(xplusdx) - f(xminusdx)) / (epsilon + epsilon)
    elseif dtype == :complex
        @complexrule x epsilon
        xplusdx = x + epsilon * im
        return imag(f(xplusdx)) / epsilon
    else
        error("dtype must be :forward, :central or :complex")
    end
end

##############################################################################
##
## Complex Step Finite Differentiation Tools
##
## Martins, Sturdza, and Alonso (2003) suggest the only non-analytic
##  fuction of which complex step finite difference approximation
##  will fail and finite difference will not is abs().
##  They suggest redefining as follows for z = x + im*y
##
##  if x < 0
##      -x - im * y
##  else
##      x + im * y
##
## This is provided below as complex_differentiable_abs (renaming encouraged!)
##
## Also, if your fuctions has control flow using < or >, you must compare
## real(z) for your control flow.
##
##############################################################################

function complex_differentiable_abs{T <: Complex}(z::T)
    if real(z) < 0
        return -real(z) - im * imag(z)
    else
        return real(z) + im * imag(z)
    end
end

##############################################################################
##
## Gradient of f: R^n -> R
##
##############################################################################

function finite_difference!{S <: Number, T <: Number}(f::Function,
                                                      x::Vector{S},
                                                      g::Vector{T},
                                                      dtype::Symbol)
    # What is the dimension of x?
    n = length(x)

    # Iterate over each dimension of the gradient separately.
    # Use xplusdx to store x + dx instead of creating a new vector on each pass.
    # Use xminusdx to store x - dx instead of creating a new vector on each pass.
    if dtype == :forward
        # Establish a baseline value of f(x).
        f_x = f(x)
        for i = 1:n
            @forwardrule x[i] epsilon
            oldx = x[i]
            x[i] = oldx + epsilon
            f_xplusdx = f(x)
            x[i] = oldx
            g[i] = (f_xplusdx - f_x) / epsilon
        end
    elseif dtype == :central
        for i = 1:n
            @centralrule x[i] epsilon
            oldx = x[i]
            x[i] = oldx + epsilon
            f_xplusdx = f(x)
            x[i] = oldx - epsilon
            f_xminusdx = f(x)
            x[i] = oldx
            g[i] = (f_xplusdx - f_xminusdx) / (epsilon + epsilon)
        end
    else
        error("dtype must be :forward or :central")
    end

    return
end
function finite_difference{T <: Number}(f::Function,
                                        x::Vector{T},
                                        dtype::Symbol = :central)
    # Allocate memory for gradient
    g = Array(Float64, length(x))

    # Mutate allocated gradient
    finite_difference!(f, float(x), g, dtype)

    # Return mutated gradient
    return g
end

##############################################################################
##
## Jacobian derivative of f: R^n -> R^m
##
##############################################################################

function finite_difference_jacobian!{R <: Number,
                                     S <: Number,
                                     T <: Number}(f::Function,
                                                  x::Vector{R},
                                                  f_x::Vector{S},
                                                  J::Array{T},
                                                  dtype::Symbol = :central)
    # What is the dimension of x?
    m, n = size(J)

    # Iterate over each dimension of the gradient separately.
    if dtype == :forward
        for i = 1:n
            @forwardrule x[i] epsilon
            oldx = x[i]
            x[i] = oldx + epsilon
            f_xplusdx = f(x)
            x[i] = oldx
            J[:, i] = (f_xplusdx - f_x) / epsilon
        end
    elseif dtype == :central
        for i = 1:n
            @centralrule x[i] epsilon
            oldx = x[i]
            x[i] = oldx + epsilon
            f_xplusdx = f(x)
            x[i] = oldx - epsilon
            f_xminusdx = f(x)
            x[i] = oldx
            J[:, i] = (f_xplusdx - f_xminusdx) / (epsilon + epsilon)
        end
    else
        error("dtype must :forward or :central")
    end

    return
end
function finite_difference_jacobian{T <: Number}(f::Function,
                                                 x::Vector{T},
                                                 dtype::Symbol = :central)
    # Establish a baseline for f_x
    f_x = f(x)

    # Allocate space for the Jacobian matrix
    J = zeros(length(f_x), length(x))

    # Compute Jacobian inside allocated matrix
    finite_difference_jacobian!(f, x, f_x, J, dtype)

    # Return Jacobian
    return J
end

##############################################################################
##
## Second derivative of f: R -> R
##
##############################################################################

function finite_difference_hessian{T <: Number}(f::Function,
                                                x::T)
    @hessianrule x epsilon
    (f(x + epsilon) - 2*f(x) + f(x - epsilon))/epsilon^2
end
function finite_difference_hessian(f::Function,
                                   g::Function,
                                   x::Number,
                                   dtype::Symbol = :central)
    finite_difference(g, x, dtype)
end

##############################################################################
##
## Hessian of f: R^n -> R
##
##############################################################################

function finite_difference_hessian!{S <: Number,
                                    T <: Number}(f::Function,
                                                 x::Vector{S},
                                                 H::Array{T})
    # What is the dimension of x?
    n = length(x)

    epsilon = NaN
    # TODO: Remove all these copies
    xpp, xpm, xmp, xmm = copy(x), copy(x), copy(x), copy(x)
    fx = f(x)
    for i = 1:n
        xi = x[i]
        @hessianrule x[i] epsilon
        xpp[i], xmm[i] = xi + epsilon, xi - epsilon
        H[i, i] = (f(xpp) - 2*fx + f(xmm)) / epsilon^2
        @centralrule x[i] epsiloni
        xp = xi + epsiloni
        xm = xi - epsiloni
        xpp[i], xpm[i], xmp[i], xmm[i] = xp, xp, xm, xm
        for j = i+1:n
            xj = x[j]
            @centralrule x[j] epsilonj
            xp = xj + epsilonj
            xm = xj - epsilonj
            xpp[j], xpm[j], xmp[j], xmm[j] = xp, xm, xp, xm
            H[i, j] = (f(xpp) - f(xpm) - f(xmp) + f(xmm))/(4*epsiloni*epsilonj)
            xpp[j], xpm[j], xmp[j], xmm[j] = xj, xj, xj, xj
        end
        xpp[i], xpm[i], xmp[i], xmm[i] = xi, xi, xi, xi
    end
    Base.LinAlg.copytri!(H,'U')
end
function finite_difference_hessian{T <: Number}(f::Function,
                                                x::Vector{T})
    # What is the dimension of x?
    n = length(x)

    # Allocate an empty Hessian
    H = Array(Float64, n, n)

    # Mutate the allocated Hessian
    finite_difference_hessian!(f, x, H)

    # Return the Hessian
    return H
end
function finite_difference_hessian{T <: Number}(f::Function,
                                                g::Function,
                                                x::Vector{T},
                                                dtype::Symbol = :central)
    finite_difference_jacobian(g, x, dtype)
end

##############################################################################
##
## Taylor Series based estimates of first and second derivatives
##
## TODO: Fill in multivariate equivalents
##
##############################################################################

# Higher precise finite difference method based on Taylor series approximation.
# h is the stepsize
function taylor_finite_difference(f::Function,
                                  x::Real,
                                  dtype::Symbol = :central,
                                  h::Real = 10e-4)
    if dtype == :forward
        f_x = f(x)
        d = 2^3 * (2^2 * (f(x + h) - f_x) - (f(x + 2 * h) - f_x))
        d += - (2^2 * (f(x + 2 * h) - f_x) - (f(x + 4 * h) - f_x))
        d /= 3 * 2^2 * h
    elseif dtype == :central
        d = 4^5 * (2^3 * (f(x + h) - f(x - h)) - (f(x + 2 * h) - f(x - 2 * h)))
        d -= 2^3 * (f(x + 4 * h) - f(x - 4 * h)) - (f(x + 8 * h) - f(x - 8 * h))
        d /= (4^5 * (2^4 - 2^2) - (2^6 - 2^4)) * h
    else
        error("dtype must be :forward or :central")
    end
    return d
end

function taylor_finite_difference_hessian(f::Function,
                                          x::Real,
                                          h::Real)
    f_x = f(x)
    d = 4^6 * (2^4 * (f(x + h) + f(x - h) - 2 * f_x) - (f(x + 2 * h) + f(x - 2 * h) - 2 * f_x))
    d += - (2^4 * (f(x + 4 * h) + f(x - 4 * h) - 2 * f_x) - (f(x + 8 * h) + f(x - 8 * h) - 2 * f_x))
    return d / (3 * 2^6 * (2^8 - 1) * h^2)
end

##############################################################################
##
## TODO: Implement directional_derivative()
##
##############################################################################

# The function "dirderivative" calculates directional derivatives in the direction v.
# The function supplied must have the form Vector{Float64} -> Float64
# function dirderivative(f::Function, v::Vector{Float64}, x0::Vector{Float64}, h::Float64, twoside::Bool)
#     derivative(t::Float64 -> f(x0 + v*t) / norm(v), 0.0, h, twoside)
# end
# function dirderivative(f::Function, v::Vector{Float64}, x0::Vector{Float64}, h::Float64)
#     dirderivative(f, v, x0, h, true)
# end
# function dirderivative(f::Function, v::Vector{Float64}, x0::Vector{Float64}, )
#     derivative(f, v, x0, 0.0001)
# end
# function dirderivative(f::Function, v::Vector{Float64})
#     x -> dirderivative(f, v, x)
# end


text:
# k-space trajectory utilities

export KSpaceCartesianSampling
export kspace_sampling, kspace_Cartesian_sampling, downscale, downscale_phase_encode_index

## k-space trajectory

coord(K::AbstractKSpaceSampling) = K.K

struct KSpaceFixedSizeSampling{T}<:AbstractKSpaceFixedSizeSampling{T}
    K::AbstractArray{T,3} # Array (nt,nk,3)
end

Base.getindex(K::AbstractKSpaceFixedSizeSampling, t::Integer) = K.K[t,:,:]
Base.size(K::AbstractKSpaceFixedSizeSampling) = size(K.K)[1:2]

kspace_sampling(K::AbstractArray{T,3}) where {T<:Real} = KSpaceFixedSizeSampling{T}(K)

struct KSpaceCartesianSampling{T}<:AbstractKSpaceCartesianSampling{T}
    X::RegularCartesianSpatialSampling{T}
    K::AbstractArray{T,3} # Array (nt,nk,3)
    phase_encoding::NTuple{2,Integer}
    readout::Integer
    subsampling::AbstractVector{<:Integer}
end

function kspace_Cartesian_sampling(X::RegularCartesianSpatialSampling{T}; phase_encoding::NTuple{2,Integer}=(1,2), subsampling::Union{Nothing,AbstractVector{<:Integer}}=nothing) where {T<:Real}

    # Set phase-encode/readout dimensions
    nx, ny, nz = X.n
    hx, hy, hz = X.h
    ((phase_encoding == (1,2)) || (phase_encoding == (2,1))) && (readout = 3)
    ((phase_encoding == (1,3)) || (phase_encoding == (3,1))) && (readout = 2)
    ((phase_encoding == (2,3)) || (phase_encoding == (3,2))) && (readout = 1)
    nt = (X.n[phase_encoding[1]], X.n[phase_encoding[2]]); nk = X.n[readout]
    perm = (phase_encoding..., readout)

    # Full k-space grid
    kx = T(pi)/hx*collect(coord_norm(nx))
    ky = T(pi)/hy*collect(coord_norm(ny))
    kz = T(pi)/hz*collect(coord_norm(nz))
    Kx = reshape(permutedims(repeat(reshape(kx,:,1,1); outer=(1,ny,nz)), perm), prod(nt), nk)
    Ky = reshape(permutedims(repeat(reshape(ky,1,:,1); outer=(nx,1,nz)), perm), prod(nt), nk)
    Kz = reshape(permutedims(repeat(reshape(kz,1,1,:); outer=(nx,ny,1)), perm), prod(nt), nk)
    K = cat(Kx, Ky, Kz; dims=3)

    # Subsampling and reordering
    isnothing(subsampling) && (subsampling = 1:prod(nt))
    K = K[subsampling, :, :]

    return KSpaceCartesianSampling{T}(X, K, phase_encoding, readout, subsampling)

end

function coord_norm(n::Integer)
    (mod(n,2) == 0) ? (c = -div(n, 2):div(n, 2)-1) : (c = -div(n-1, 2):div(n-1, 2))
    return c/norm(c, Inf)
end

function downscale(K::KSpaceCartesianSampling{T}; fact::Integer=1) where {T<:Real}
    (fact == 0) && (return K)
    X_h = downscale(K.X; fact=fact)
    phase_encode_idx_h, readout_idx_h = downscale_phase_encode_index(K; fact=fact, readout=true)
    K_h = K.K[phase_encode_idx_h, readout_idx_h, :]
    return KSpaceCartesianSampling{T}(X_h, K_h, K.phase_encoding, K.readout, K.subsampling)
end

function downscale_phase_encode_index(K::KSpaceCartesianSampling{T}; fact::Integer=1, readout::Bool=false) where {T<:Real}
    k_max = T(pi)./K.X.h
    pe = K.phase_encoding
    r = K.readout
    phase_encode_idx_h = findall((K.K[:,1,pe[1]] .< k_max[pe[1]]/2^fact) .&& (K.K[:,1,pe[1]] .>= -k_max[pe[1]]/2^fact) .&& (K.K[:,1,pe[2]] .< k_max[pe[2]]/2^fact) .&& (K.K[:,1,pe[2]] .>= -k_max[pe[2]]/2^fact))
    readout_idx_h = findall((K.K[1,:,r] .< k_max[r]/2^fact) .&& (K.K[1,:,r] .>= -k_max[r]/2^fact))
    readout ? (return (phase_encode_idx_h, readout_idx_h)) : (return phase_encode_idx_h)
end

function downscale(d::AbstractArray{CT,2}, K::KSpaceCartesianSampling{T}; fact::Integer=1, flat::Bool=false, coeff::Real=1) where {T<:Real,CT<:RealOrComplex{T}}
    (fact == 0) && (return d)
    i_pe, i_r = downscale_phase_encode_index(K; fact=fact, readout=true)
    return anti_aliasing_filter(K; fact=fact, flat=flat, coeff=coeff).*d[i_pe, i_r]/T(sqrt(2.0^(3*fact)))
end

max_stars_repo_name: grizzuti/UtilitiesForMRI

max_issues_repo_name: grizzuti/UtilitiesForMRI

max_forks_repo_name: grizzuti/UtilitiesForMRI


text:
Statistics.var(rr::RegressionModel) = rss(rr) / dof_residual(rr)
Statistics.std(rr::RegressionModel) = sqrt(var(rr))

function bh_return(vals::AbstractVector{Float64})
    out = 1.0
    @simd for x in vals
        out *= (1 + x)
    end
    out - 1
end
function bh_return(vals::AbstractVector{Union{Missing, Float64}})
    out = 1.0
    @simd for x in vals
        if ismissing(x)
            out *= 1
        else
            out *= (1+x)
        end
    end
    out - 1
end

function bh_return(pred::AbstractMatrix, coef)
    #@assert size(pred, 2) == length(coef) "Got Matrix of size $(size(pred)) and coefficients of $coef $pred"
    out = 1.0
    @simd for i in 1:size(pred, 1)
        out *= (fast_pred(pred, coef, i) + 1)
    end
    out - 1
end

bhar(resp, pred, coef) = bh_return(resp) - bh_return(pred, coef)

bhar(x::AbstractVector, y::AbstractVector) = bh_return(x) - bh_return(y)

function cumulative_return(pred::AbstractMatrix, coef)
    #@assert size(pred, 2) == length(coef) "Got Matrix of size $(size(pred)) and coefficients of $coef $pred"
    out = 0.0
    @simd for i in 1:size(pred, 1)
        out += fast_pred(pred, coef, i)
    end
    out
end
car(resp, pred, coef) = sum(resp) - cumulative_return(pred, coef)
car(x::AbstractVector, y::AbstractVector) = sum(skipmissing(x)) - sum(skipmissing(y))

var_diff(x, y) = var(x) + var(y) - 2 * cov(x, y)
# for firm data
"""

Calculates the buy and hold returns (also called geometric return) for TimelineData. If an Integer
is passed, then it is calculated based on the FIRM_DATA_CACHE (for the integer provided), otherwise
is calculated for the MARKET_DATA_CACHE.

These functions treat missing returns in the period implicitly as a zero return.
"""
function bh_return(data::TimelineTable, col; minobs=0.8)
    col = TimelineColumn(col)
    select!(data, [col])
    dates = dates_min_max(norm_dates(data), data_dates(data))
    if nnz(data_missing_bdays(data)) == 0
        x = view(data[col], dates)
    else
        new_missings = get_missing_bdays(data, dates)
        x = view(data[col], dates, new_missings)
    end
    if length(x) < adjust_minobs(minobs, calendar(data), norm_dates(data))
        missing
    else
        bh_return(x)
    end
end

function fill_vector_bh_return!(
    out_vector::Vector{Union{Missing, Float64}},
    iter_data::IterateTimelineTable,
    col::TimelineColumn;
    minobs=0.8
)
    @assert validate_iterator(iter_data, out_vector) "Length of out_vector does not match the number of indexes in the iterator"

    Threads.@threads for (u_id, iter_indexes) in iter_data
        data = parent(iter_data)[u_id, [col], AllowMissing{true}]
        resp = data[col]
        for iter in iter_indexes
            cur_dates = dates_min_max(data_dates(data), iter_dates(iter))
            if nnz(data_missing_bdays(data)) == 0
                x = view(resp, cur_dates)
            else
                new_mssngs = get_missing_bdays(data, cur_dates)
                x = view(resp, cur_dates, new_mssngs)
            end

            length(x) < adjust_minobs(minobs, calendar(parent(iter_data)), iter_dates(iter)) && continue

            @inbounds out_vector[iter_index(iter)] = bh_return(x)

        end
    end
    out_vector
end

function bh_return(
    data::IterateTimelineTable,
    col;
    minobs=0.8
)
    col = TimelineColumn(col)
    out = Vector{Union{Missing, Float64}}(missing, total_length(data))
    fill_vector_bh_return!(
        out,
        data,
        col;
        minobs
    )
    if any(ismissing.(out))
        out
    else
        disallowmissing(out)
    end
end

function sum_return(data::TimelineTable, col; minobs=0.8)
    col = TimelineColumn(col)
    select!(data, [col])
    dates = dates_min_max(norm_dates(data), data_dates(data))
    if nnz(data_missing_bdays(data)) == 0
        x = view(data[col], dates)
    else
        new_missings = get_missing_bdays(data, dates)
        x = view(data[col], dates, new_missings)
    end
    if length(x) < adjust_minobs(minobs, calendar(data), norm_dates(data))
        missing
    else
        sum(x)
    end
end

function simple_diff(
    data::TimelineTable,
    firm_col,
    mkt_col,
    fun;
    minobs=0.8
)

    select!(data, [firm_col, mkt_col])
    dates = dates_min_max(norm_dates(data), data_dates(data))
    if nnz(data_missing_bdays(data)) == 0
        x = view(data[firm_col], dates)
        y = view(data[mkt_col], dates)
    else
        new_missings = get_missing_bdays(data, dates)
        x = view(data[firm_col], dates, new_missings)
        y = view(data[mkt_col], dates, new_missings)
    end
    if length(x) < adjust_minobs(minobs, calendar(data), norm_dates(data))
        missing
    else
        fun(x, y)
    end
end

function pred_diff(
    data::TimelineTable,
    rr::RegressionModel,
    fun;
    minobs::Float64=0.8
)
    if !isdefined(rr, :coef)
        return missing
    end
    f = rr.formula
    sch = apply_schema(f, schema(f, data))
    select!(data, internal_termvars(sch))
    if length(data) < adjust_minobs(minobs, calendar(data), norm_dates(data))
        return missing
    end
    # since the values in a continuous term of mean/var/min/max are not used here,
    # this just creates a schema from the available values without

    resp, pred = modelcols(sch, data)
    fun(resp, pred, coef(rr))
end

function fill_vector_simple!(
    out_vector::Vector{Union{Missing, Float64}},
    iter_data::IterateTimelineTable,
    cache::DataVector,
    col::TimelineColumn,
    fun;
    minobs=0.8
)
    @assert validate_iterator(iter_data, out_vector) "Length of out_vector does not match the number of indexes in the iterator"

    Threads.@threads for (u_id, iter_indexes) in iter_data
        data = parent(iter_data)[u_id, [col], AllowMissing{true}]
        resp = data[col]
        for iter in iter_indexes
            cur_dates = dates_min_max(data_dates(data), iter_dates(iter))
            if nnz(data_missing_bdays(data)) == 0
                x = view(resp, cur_dates)
                y = view(cache, cur_dates)
            else
                new_mssngs = get_missing_bdays(data, cur_dates)
                x = view(resp, cur_dates, new_mssngs)
                y = view(cache, cur_dates, new_mssngs)
            end

            length(x) < adjust_minobs(minobs, calendar(parent(iter_data)), iter_dates(iter)) && continue

            @inbounds out_vector[iter_index(iter)] = fun(
                x,
                y,
            )

        end
    end
    out_vector
end

function vector_simple_diff(
    data::IterateTimelineTable,
    col1::TimelineColumn,
    col2::TimelineColumn,
    fun;
    minobs=0.8
)
    cache = parent(data)[first(data)[1], [col2], AllowMissing{true}][col2]

    out = Vector{Union{Missing, Float64}}(missing, total_length(data))
    fill_vector_simple!(
        out,
        data,
        cache,
        col1,
        fun;
        minobs
    )
    if any(ismissing.(out))
        out
    else
        disallowmissing(out)
    end
end

function fill_vector_pred!(
    out_vector::Vector{Union{Missing, Float64}},
    iter_data::IterateTimelineTable{T},
    cache::DataMatrix,
    sch,
    rrs::Vector{BasicReg{L,R}},
    fun;
    minobs=0.8
) where {T, L, R}
    @assert validate_iterator(iter_data, out_vector) "Length of out_vector does not match the number of indexes in the iterator"
    @assert length(out_vector) == length(rrs) "Length of out_vector is not the same as number of regressions"
    cols = internal_termvars(sch)
    Threads.@threads for (u_id, iter_indexes) in iter_data
        data = parent(iter_data)[u_id, cols, AllowMissing{true}]
        resp = datavector_modelcols(int_lhs(sch), data)
        for iter in iter_indexes
            isdefined(rrs[iter_index(iter)], :coef) || continue
            cur_dates = dates_min_max(data_dates(data), iter_dates(iter))
            if nnz(data_missing_bdays(data)) == 0
                x = view(resp, cur_dates)
                y = view(cache, cur_dates)
            else
                new_mssngs = get_missing_bdays(data, cur_dates)
                x = view(resp, cur_dates, new_mssngs)
                y = view(cache, cur_dates, new_mssngs)
            end

            length(x) < adjust_minobs(minobs, calendar(parent(iter_data)), iter_dates(iter)) && continue

            @inbounds out_vector[iter_index(iter)] = fun(
                x,
                y,
                coef(rrs[iter_index(iter)])
            )

        end
    end
    out_vector
end

function vector_pred_diff(
    data::IterateTimelineTable,
    rrs::AbstractVector{BasicReg{L, R}},
    fun;
    minobs=0.8
) where {L, R}
    f = rrs[1].formula
    sch = apply_schema(f, schema(f, parent(data)))
    cache = create_pred_matrix(data, sch)

    out = Vector{Union{Missing, Float64}}(missing, total_length(data))
    fill_vector_pred!(
        out,
        data,
        cache,
        sch,
        rrs,
        fun;
        minobs
    )
    if any(ismissing.(out))
        out
    else
        disallowmissing(out)
    end
end

"""
    bhar(
        data::TimelineTable{Mssng, T, MNames, FNames},
        firm_col=FNames[1],
        mkt_col=MNames[1];
        minobs=0.8
    ) where {Mssng, T, MNames, FNames}

    bhar(
        data::TimelineTable,
        rr::RegressionModel;
        minobs=0.8
    )

    bhar(
        data::IterateTimelineTable{T, MNames, FNames},
        firm_col=FNames[1],
        mkt_col=MNames[1];
        minobs=0.8
    ) where {T, MNames, FNames}

    bhar(
        data::IterateTimelineTable,
        rrs::AbstractVector{<:BasicReg};
        minobs=0.8
    )

Calculates the difference between buy and hold returns (also referred to as geometric returns) for a firm and a benchmark.
If a regression is passed, then the benchmark is based on the coefficients from that regression and the performance of the benchmarks
in the regression. These are sometimes called Fama-French abnormal returns. Simple abnormal returns use a market index as the benchmark
(such as the S&P 500 or a value weighted return of all firms).

Similar to constructing the regression, passing an `IterateTimelineTable` will return a Vector and uses a more optimized method.
"""
function bhar(
    data::TimelineTable{Mssng, T, MNames, FNames},
    firm_col=FNames[1],
    mkt_col=MNames[1];
    minobs=0.8
) where {Mssng, T, MNames, FNames}
    simple_diff(data, TimelineColumn(firm_col), TimelineColumn(mkt_col), bhar; minobs)
end

function bhar(
    data::TimelineTable,
    rr::RegressionModel;
    minobs=0.8
)
    pred_diff(data, rr, bhar; minobs)
end

function bhar(
    data::IterateTimelineTable{T, MNames, FNames},
    firm_col=FNames[1],
    mkt_col=MNames[1];
    minobs=0.8
) where {T, MNames, FNames}
    vector_simple_diff(data, TimelineColumn(firm_col), TimelineColumn(mkt_col), bhar; minobs)
end

function bhar(
    data::IterateTimelineTable,
    rrs::AbstractVector{<:BasicReg};
    minobs=0.8
)
    vector_pred_diff(data, rrs, bhar; minobs)
end

"""
    car(
        data::TimelineTable{Mssng, T, MNames, FNames},
        firm_col=FNames[1],
        mkt_col=MNames[1];
        minobs=0.8
    ) where {Mssng, T, MNames, FNames}

    car(
        data::TimelineTable,
        rr::RegressionModel;
        minobs=0.8
    )

    car(
        data::IterateTimelineTable{T, MNames, FNames},
        firm_col=FNames[1],
        mkt_col=MNames[1];
        minobs=0.8
    ) where {T, MNames, FNames}

    car(
        data::IterateTimelineTable,
        rrs::AbstractVector{<:BasicReg};
        minobs=0.8
    )

Calculates the cumulative returns of a firm over a benchmark (through addition of each return).
If a regression is passed, then the benchmark is based on the coefficients from that regression and the performance of the benchmarks
in the regression. These are sometimes called Fama-French abnormal returns. Simple abnormal returns use a market index as the benchmark
(such as the S&P 500 or a value weighted return of all firms).

Similar to constructing the regression, passing an `IterateTimelineTable` will return a Vector and uses a more optimized method.
"""
function car(
    data::TimelineTable{Mssng, T, MNames, FNames},
    firm_col=FNames[1],
    mkt_col=MNames[1];
    minobs=0.8
) where {Mssng, T, MNames, FNames}
    simple_diff(data, TimelineColumn(firm_col), TimelineColumn(mkt_col), car; minobs)
end

function car(
    data::TimelineTable,
    rr::RegressionModel;
    minobs=0.8
)
    pred_diff(data, rr, car; minobs)
end

function car(
    data::IterateTimelineTable{T, MNames, FNames},
    firm_col=FNames[1],
    mkt_col=MNames[1];
    minobs=0.8
) where {T, MNames, FNames}
    vector_simple_diff(data, TimelineColumn(firm_col), TimelineColumn(mkt_col), car; minobs)
end

function car(
    data::IterateTimelineTable{T},
    rrs::AbstractVector{<:BasicReg};
    minobs=0.8
) where {T}
    vector_pred_diff(data, rrs, car; minobs)
end

function Statistics.var(
    data::TimelineTable{Mssng, T, MNames, FNames},
    firm_col=FNames[1],
    mkt_col=MNames[1];
    minobs=0.8
) where {Mssng, T, MNames, FNames}
    simple_diff(data, TimelineColumn(firm_col), TimelineColumn(mkt_col), var_diff; minobs)
end

function Statistics.std(
    data::TimelineTable{Mssng, T, MNames, FNames},
    firm_col=FNames[1],
    mkt_col=MNames[1];
    minobs=0.8
) where {Mssng, T, MNames, FNames}
    sqrt(var(data, col_firm, col_market; minobs))
end

function get_coefficient_val(rr::RegressionModel, coefname::String...)
    ismissing(coefnames(rr)) && return missing
    for x in coefname
        if x ∈ coefnames(rr)
            return coef(rr)[findfirst(x .== coefnames(rr))]
        end
    end
    @error("None of $(coefname) is in the RegressionModel model.")
end
"""
    alpha(rr::RegressionModel, coefname::String...="intercept")

"alpha" in respect to the the CAPM model, i.e., the intercept in the model.
This is the alpha from the estimation period.

This function finds the position of the coefficient name provided, defaults to "intercept".
If the coefname is not in the regression, then this function returns an error.
"""
alpha(rr::RegressionModel, coefname::String...="(Intercept)") = get_coefficient_val(rr, coefname...)

"""
    beta(rr::RegressionModel, coefname::String...=["mkt", "mktrf", "vwretd", "ewretd"])

"beta" in respect to the CAPM model, i.e., the coefficient on the market return minus the risk free rate.
This is the beta from the estimation period.

This function finds the position of the coefficient name provided, defaults to several common market returns.
If the coefname is not in the regression, then this function returns an error.
"""
beta(rr::RegressionModel, coefname::String...=["mkt", "mktrf", "vwretd", "ewretd"]...) = get_coefficient_val(rr, coefname...)


text:
mutable struct Point
    x::Float64
    y::Float64
end

mutable struct Vector2D
    x::Float64
    y::Float64
end

p = Point(2, 5) #> Point(2.0,5.0)
v = Vector2D(3, 2) #> Vector2D(3.0,2.0)

# Example with sum:
# +(p, v) #> MethodError: ERROR: `+` has no method matching +(::Point, ::Vector2D)

import Base.+
+(p::Point,    q::Point) = Point(p.x + q.x, p.y + q.y)
+(u::Vector2D, v::Vector2D) = Point(u.x + v.x, u.y + v.y)
+(u::Vector2D, p::Point) = Point(u.x + p.x, u.y + p.y)
# +(p, v) #> ERROR: MethodError: `+` has no method matching +(::Point, ::Vector2D)
+(p::Point,    v::Vector2D) = Point(p.x + v.x, p.y + v.y)
+(p, v) #> Point(5.0,7.0)

# Example with dot product:
# *(p, v) #> ERROR: MethodError: `*` has no method matching *(::Point, ::Vector2D)

import Base.*
*(p::Point,    q::Point) = p.x * q.x + p.y * q.y
*(u::Vector2D, v::Vector2D) = u.x * v.x + u.y * v.y
*(u::Vector2D, p::Point) = u.x * p.x + u.y * p.y
# *(p, v) #> ERROR: MethodError: `*` has no method matching *(::Point, ::Vector2D)
*(p::Point,    v::Vector2D) = p.x * v.x + p.y * v.y
*(p, v) #> 16.0

VecOrPoint = Union{Vector2D, Point}

isa(p, VecOrPoint) #> true
isa(v, VecOrPoint) #> true

+(u::VecOrPoint, v:: VecOrPoint) = VecOrPoint(u.x + v.x, u.y + v.y)
+(p, v) #> Point(5.0,7.0)

*(u::VecOrPoint, v:: VecOrPoint) = u.x * v.x + u.y * v.y
*(p, v) #> 16.0


text:
using EBayes
using InteractiveUtils

ash_prs = [T() for  T in subtypes(EBayes.AshPrior)]

iw_prs = [T() for  T in subtypes(EBayes.IWPrior)]


text:
export zoom_Nwt
function zoom_Nwt(h :: LineModel,
                  t₀ :: Float64,
                  t₁ :: Float64;
                  c₁ :: Float64=0.01,
                  tol :: Float64=1e-7,
                  ϵ :: Float64=1e-10,
                  maxiter :: Int=100,
                  verbose :: Bool=false)

        if obj(h,t₀)<obj(h,t₁)
          tl=t₀
          th=t₁
        else
          tl=t₁
          th=t₀
        end

        hl=obj(h,tl)
        dhl=grad(h,tl)
        hh=obj(h,th)
        dhh=grad(h,th)

        h₀=obj(h,t₀)
        dh₀=grad(h,t₀)

        γ=0.8
        t=t₁
        tp=t₀
        tqnp=t₀
        hk=0
        hkm1=0
        gkm1=0
        hp=0
        i=0

        hk=obj(h,t)
        gk=grad(h,t)
        hkm1=obj(h,tqnp)
        gkm1=grad(h,tqnp)

        verbose && @printf(" iter        tₗ        tₕ         t        hₗ        hₕ         hk         gk\n")
        verbose && @printf(" %7.2e %7.2e  %7.2e  %7.2e  %7.2e %7.2e %7.2e %7.2e\n", i,tl,th,t,hl,hh,hk,gk)

        while ((i<maxiter) & (abs(gk)>tol)) || i==0

          if (hk>h₀+c₁*(t-t₀)*dh₀) || (hl<=hk)
            th=t
            tlast=th
            hlast=hh
            dhlast=dhh
            hh=hk
            dhh=gk
          else
            #gk=grad(h,t)
            if (abs(gk)<ϵ)
              topt=t
              iter=i
              return (topt,iter)
            elseif gk*(th-tl)>=0
              th=tl
              hh=hl
              dhh=dhl
            end
            tl=t
            tlast=tl
            hlast=hl
            dhlast=dhl
            hl=hk
            dhl=gk
          end

          kₖ=hess(h,t)
          dN=-gk/kₖ #direction de Newton

          if ((tp-t)*dN>0) & (dN/(tp-t)<γ)
            tplus = t + dN
            hplus = obj(h, tplus)
            gplus = grad(h,tplus)
            verbose && println("N")
          else
            tplus = (t+tp)/2
            hplus = obj(h, tplus)
            gplus = grad(h,tplus)
            verbose && println("B")
          end

          if t>tp
            if gplus<0
              tp=t
              tqnp=t
              t=tplus
            else
              tqnp=t
              t=tplus
            end
          else
            if gplus>0
              tp=t
              tqnp=t
              t=tplus
            else
              tqnp=t
              t=tplus
            end
          end

          #mise à jour des valeurs
          hkm1=hk
          gkm1=gk
          hk=hplus
          gk=gplus
          i=i+1
          #println("on continue d'itérer")
          verbose && @printf(" %7.2e %7.2e  %7.2e  %7.2e %7.2e %7.2e %7.2e %7.2e\n", i,tl,th,t,hl,hh,hk,gk)

        end
        topt=t
        iter=i

        return (topt,iter)
end