# Tests for Equality of Variances

export OneWayANOVATest, LeveneTest, BrownForsytheTest, FlignerKilleenTest

struct VarianceEqualityTest{TD <: ContinuousUnivariateDistribution} <: HypothesisTest
    Nᵢ::Vector{Int}
    SStᵢ::Vector{Float64}
    SSeᵢ::Vector{Float64}
    DFt::Int
    DFe::Int
    # test name, parameter of interest, test statistic name
    description::Tuple{String, String, String}
end

population_param_of_interest(t::VarianceEqualityTest) =
    (t.description[2], "all equal", NaN)
testname(t::VarianceEqualityTest) = t.description[1]
teststatisticname(t::VarianceEqualityTest{TD}) where {TD <: ContinuousDistribution} =
    length(t.description[3]) != 0 ? t.description[3] : (TD <: FDist ? "F" : "χ²")

StatsAPI.nobs(t::VarianceEqualityTest) = t.Nᵢ
StatsAPI.dof(t::VarianceEqualityTest{Chisq}) = t.DFt
StatsAPI.dof(t::VarianceEqualityTest{FDist}) = (t.DFt, t.DFe)

teststatistic(t::VarianceEqualityTest{FDist}) = (t.DFe/t.DFt)*sum(t.SStᵢ)/sum(t.SSeᵢ)
function teststatistic(t::VarianceEqualityTest{Chisq})
    y = sum(t.SStᵢ)/sum(t.SSeᵢ)
    y*(t.DFe+t.DFt)/(1 + y) # sum(t.SStᵢ)/t.s²
end
StatsAPI.pvalue(t::VarianceEqualityTest{TD}; tail=:right) where {TD <: ContinuousDistribution} =
    pvalue(TD(dof(t)...), teststatistic(t), tail=tail)

function show_params(io::IO, t::VarianceEqualityTest{TD},
                     indent="") where {TD <: ContinuousDistribution}
    println(io, indent, "number of observations: ", nobs(t))
    println(io, indent, rpad("$(teststatisticname(t)) statistic:", 24), teststatistic(t))
    println(io, indent, "degrees of freedom:     ", dof(t))
end

function Base.show(io::IOContext, t::VarianceEqualityTest)
    if !get(io, :table, false) # No table
        show(io.io, t)
    else
        println(io, testname(t))
        println(io, repeat("-", 55))
        SSt = sum(t.SStᵢ)
        SSe = sum(t.SSeᵢ)
        MSt = SSt/t.DFt
        MSe = SSe/t.DFe
        println(io, "Source            SS    DF        MS         F  P-value")
        println(io, repeat("-", 55))
        @printf(io, "Treatments  %8.3f  %4d  %8.3f  %8.3f  %7.5f\n",
                SSt, t.DFt, MSt, MSt/MSe, pvalue(t))
        @printf(io, "Error       %8.3f  %4d  %8.3f\n", SSe, t.DFe, MSe)
        println(io, repeat("-", 55))
        @printf(io, "Total       %8.3f  %4d\n", SSt+SSe, t.DFt+t.DFe)
    end
end

function anova(scores)
    if isempty(scores) || !all(x -> x isa AbstractVector{<:Real}, scores)
        throw(ArgumentError("`anova` requires a non-empty collection of vectors of `Real` numbers"))
    end

    Nᵢ = [length(s) for s in scores]
    Z̄ᵢ = map(mean, scores)
    Z̄ = dot(Z̄ᵢ, Nᵢ) / sum(Nᵢ)
    SStᵢ = Nᵢ .* abs2.(Z̄ᵢ .- Z̄)
    SSeᵢ = [sum(z -> abs2(z - z̄ᵢ), scoresᵢ) for (scoresᵢ, z̄ᵢ) in zip(scores, Z̄ᵢ)]
    (Nᵢ, SStᵢ, SSeᵢ)
end

anova(scores::AbstractVector{<:Real}...) = anova(scores)

"""
    OneWayANOVATest(groups)
    OneWayANOVATest(groups::AbstractVector{<:Real}...)

Perform one-way analysis of variance test of the hypothesis that that the `groups`
means are equal.

The one-way analysis of variance (one-way ANOVA) is a technique that can be used
to compare means of two or more samples. The ANOVA tests the null hypothesis, which states
that samples in all groups are drawn from populations with the same mean values.
To do this, two estimates are made of the population variance.
The ANOVA produces an F-statistic, the ratio of the variance calculated among
the means to the variance within the samples.

Implements: [`pvalue`](@ref)

# External links

  * [One-way analysis of variance on Wikipedia
    ](https://en.wikipedia.org/wiki/One-way_analysis_of_variance)
"""
function OneWayANOVATest(groups)
    Nᵢ, SStᵢ, SSeᵢ = anova(groups)
    k = length(Nᵢ)
    VarianceEqualityTest{FDist}(Nᵢ, SStᵢ, SSeᵢ, k-1, sum(Nᵢ)-k,
        ("One-way analysis of variance (ANOVA) test", "Means", "F"))
end

OneWayANOVATest(groups::AbstractVector{<:Real}...) = OneWayANOVATest(groups)

"""
    LeveneTest(groups; scorediff=abs, statistic=mean)
    LeveneTest(groups::AbstractVector{<:Real}...; scorediff=abs, statistic=mean)

Perform Levene's test of the hypothesis that that the `groups` variances are equal.
By default the mean `statistic` is used for centering in each of the `groups`, but
other statistics are accepted: median or truncated mean, see [`BrownForsytheTest`](@ref).
By default the absolute value of the score difference, `scorediff`, is used, but
other functions are accepted: x² or √|x|.

The test statistic, ``W``, is equivalent to the ``F`` statistic, and is defined as follows:

```math
W = \\frac{(N-k)}{(k-1)} \\cdot \\frac{\\sum_{i=1}^k N_i (Z_{i\\cdot}-Z_{\\cdot\\cdot})^2}
    {\\sum_{i=1}^k \\sum_{j=1}^{N_i} (Z_{ij}-Z_{i\\cdot})^2},
```
where

* ``k`` is the number of different groups to which the sampled cases belong,
* ``N_i`` is the number of cases in the ``i``th group,
* ``N`` is the total number of cases in all groups,
* ``Y_{ij}`` is the value of the measured variable for the ``j``th case from the ``i``th group,
* ``Z_{ij} = |Y_{ij} - \\bar{Y}_{i\\cdot}|``, ``\\bar{Y}_{i\\cdot}`` is a mean of the  ``i``th group,
* ``Z_{i\\cdot} = \\frac{1}{N_i} \\sum_{j=1}^{N_i} Z_{ij}`` is the mean of the ``Z_{ij}`` for group ``i``,
* ``Z_{\\cdot\\cdot} = \\frac{1}{N} \\sum_{i=1}^k \\sum_{j=1}^{N_i} Z_{ij}`` is the mean of all ``Z_{ij}``.

The test statistic ``W`` is approximately ``F``-distributed with ``k-1`` and ``N-k`` degrees of freedom.

Implements: [`pvalue`](@ref)

# References

  * [Brown and Forsythe (1974). Robust Tests for the Equality of Variances](@cite brown1974)

# External links

  * [Levene's test on Wikipedia
    ](https://en.wikipedia.org/wiki/Levene%27s_test)
"""
function LeveneTest(groups; scorediff=abs, statistic=mean)
    # calculate scores
    Zᵢⱼ = map(g -> scorediff.(g .- statistic(g)), groups)
    # anova
    Nᵢ, SStᵢ, SSeᵢ = anova(Zᵢⱼ)
    k = length(Nᵢ)
    VarianceEqualityTest{FDist}(Nᵢ, SStᵢ, SSeᵢ, k-1, sum(Nᵢ)-k,
        ("Levene's test", "Variances", "W"))
end

LeveneTest(groups::AbstractVector{<:Real}...; kwargs...) =
    LeveneTest(groups; kwargs...)

"""
    BrownForsytheTest(groups)
    BrownForsytheTest(groups::AbstractVector{<:Real}...)

The Brown–Forsythe test is a statistical test for the equality of `groups` variances.

The Brown–Forsythe test is a modification of the Levene's test with the median instead
of the mean statistic for computing the spread within each group.

Implements: [`pvalue`](@ref)

# References

  * [Brown and Forsythe (1974). Robust Tests for the Equality of Variances](@cite brown1974)

# External links

  * [Brown–Forsythe test on Wikipedia
    ](https://en.wikipedia.org/wiki/Brown%E2%80%93Forsythe_test)
"""
BrownForsytheTest(groups) = LeveneTest(groups; statistic=median)
BrownForsytheTest(groups::AbstractVector{<:Real}...) = BrownForsytheTest(groups)

"""
    FlignerKilleenTest(groups)
    FlignerKilleenTest(groups::AbstractVector{<:Real}...)

Perform Fligner-Killeen median test of the null hypothesis that the `groups`
have equal variances, a test for homogeneity of variances.

This test is most robust against departures from normality, see references.
It is a ``k``-sample simple linear rank method that uses the ranks of the absolute values
of the centered samples and weights
```math
a_{N,i} = \\Phi^{-1}(1/2 + (i/2(N+1)))
```
The version implemented here uses median centering in each of the samples.

Implements: [`pvalue`](@ref)

# References

  * [Conover, Johnson, and Johnson (1981). A Comparative Study of Tests for Homogeneity of
     Variances, with Applications to the Outer Continental Shelf Bidding Data](@cite conover1981)

# External links

  * [Fligner-Killeen test on Statistical Analysis Handbook
    ](https://www.statsref.com/HTML/index.html?fligner-killeen_test.html)
"""
function FlignerKilleenTest(groups)
    # calculate scores
    Zᵢⱼ = map(g -> abs.(g .- median(g)), groups)
    # rank scores
    (ranks, tieadj) = tiedrank_adj(reduce(vcat, Zᵢⱼ))
    qᵢⱼ = quantile.(Normal(), 0.5 .+ ranks ./ 2(length(ranks) + 1))
    lastᵢ = 0
    Qᵢⱼ = map(groups) do gᵢ
        n = length(gᵢ)
        qᵢ = view(qᵢⱼ, (lastᵢ + 1):(lastᵢ + n))
        lastᵢ += n
        return qᵢ
    end
    # anova
    Nᵢ, SStᵢ, SSeᵢ = anova(Qᵢⱼ)
    k = length(Nᵢ)
    VarianceEqualityTest{Chisq}(Nᵢ, SStᵢ, SSeᵢ, k-1, sum(Nᵢ)-k,
        ("Fligner-Killeen test", "Variances", "FK"))
end

FlignerKilleenTest(groups::AbstractVector{<:Real}...) =
    FlignerKilleenTest(groups)
