# TODO: Remove when `MatrixReshaped` is removed
using Distributions, Test, Random, LinearAlgebra

rng = MersenneTwister(123456)

@testset "matrixreshaped.jl" begin

function test_matrixreshaped(rng, d1, sizes)
    x1 = rand(rng, d1)
    d1s = [@test_deprecated(MatrixReshaped(d1, s...)) for s in sizes]

@testset "MatrixReshaped $(nameof(typeof(d1))) tests" begin
    @testset "MatrixReshaped constructor" begin
        for d in d1s
            @test d isa MatrixReshaped
        end
    end
    @testset "MatrixReshaped constructor errors" begin
        @test_deprecated(@test_throws ArgumentError MatrixReshaped(d1, length(d1), 2))
        @test_deprecated(@test_throws ArgumentError MatrixReshaped(d1, length(d1)))
        @test_deprecated(@test_throws ArgumentError MatrixReshaped(d1, -length(d1), -1))
    end
    @testset "MatrixReshaped size" begin
        for (d, s) in zip(d1s[1:end-1], sizes[1:end-1])
            @test size(d) == s
        end
    end
    @testset "MatrixReshaped length" begin
        for d in d1s
            @test length(d) == length(d1)
        end
    end
    @testset "MatrixReshaped rank" begin
        for (d, s) in zip(d1s, sizes)
            @test rank(d) == minimum(s)
        end
    end
    @testset "MatrixReshaped insupport" begin
        for (i, d) in enumerate(d1s[1:end-1])
            for (j, s) in enumerate(sizes[1:end-1])
                @test (i == j) ⊻ !insupport(d, reshape(x1, s))
            end
        end
    end
    @testset "MatrixReshaped mean" begin
        for (d, s) in zip(d1s[1:end-1], sizes[1:end-1])
            @test mean(d) == reshape(mean(d1), s)
        end
    end
    @testset "MatrixReshaped mode" begin
        for (d, s) in zip(d1s[1:end-1], sizes[1:end-1])
            @test mode(d) == reshape(mode(d1), s)
        end
    end
    @testset "MatrixReshaped covariance" begin
        for (d, (n, p)) in zip(d1s[1:end-1], sizes[1:end-1])
            @test cov(d) == cov(d1)
            @test cov(d, Val(false)) == reshape(cov(d1), n, p, n, p)
        end
    end
    @testset "MatrixReshaped variance" begin
        for (d, s) in zip(d1s[1:end-1], sizes[1:end-1])
            @test var(d) == reshape(var(d1), s)
        end
    end
    @testset "MatrixReshaped params" begin
        for (d, s) in zip(d1s[1:end-1], sizes[1:end-1])
            @test params(d) == (d1, s)
        end
    end
    @testset "MatrixReshaped partype" begin
        for d in d1s
            @test partype(d) === partype(d1)
        end
    end
    @testset "MatrixReshaped eltype" begin
        for d in d1s
            @test eltype(d) === eltype(d1)
        end
    end
    @testset "MatrixReshaped logpdf" begin
        for (d, s) in zip(d1s[1:end-1], sizes[1:end-1])
            x = reshape(x1, s)
            @test logpdf(d, x) == logpdf(d1, x1)
        end
    end
    @testset "MatrixReshaped rand" begin
        for d in d1s
            x = rand(rng, d)
            @test insupport(d, x)
            @test insupport(d1, vec(x))
            @test logpdf(d, x) == logpdf(d1, vec(x))
        end
    end
    @testset "MatrixReshaped vec" begin
        for d in d1s
            @test vec(d) === d1
        end
    end
end
end

    # MvNormal
    σ = rand(rng, 16, 16)
    μ = rand(rng, 16)
    d1 = MvNormal(μ, σ * σ')
    sizes = [(4, 4), (8, 2), (2, 8), (1, 16), (16, 1), (4,)]
    test_matrixreshaped(rng, d1, sizes)

    # Dirichlet
    α = rand(rng, 36) .+ 1 # mode is only defined if all alpha > 1
    d1 = Dirichlet(α)
    sizes = [(6, 6), (4, 9), (9, 4), (3, 12), (12, 3), (1, 36), (36, 1), (6,)]
    test_matrixreshaped(rng, d1, sizes)
end
