using StatsBase: geomean
using ThreadsX

# available Î± estimation methods
const Î±_METHODS = [
    "constant",
    "jacobian_upper_bound",
    "batch_maximum",
    "product_upper_bound",
    "explog_approx",
    "arithmetic_mean",
]

"""
    compute_Î±(tfrec, ZÌ‚, est_method)

Compute `Î±` given recursion wrapper `tfrec`, forcing signals `ZÌ‚` and
an `Î±` estimation method `est_method`. `ZÌ‚` is an `M Ã— S Ã— T` array where `M`
denotes the latent dimension, `S` is the batch dimension and `T` is the time dimension.
Throws an error if `est_method` is not one of the methods in `Î±_METHODS`.

Returns `Î±` and spectral norm `ğ�’©` used to compute it.
"""
function compute_Î±(
    tfrec,
    ZÌ‚::AbstractArray{T, 3},
    est_method::String;
    return_norm::Bool = false,
) where {T}
    # model
    â„³ = tfrec.model

    # select method
    if est_method == "jacobian_upper_bound"
        Î±, ğ�’© = compute_Î±_jacobian_upper_bound(â„³)
    elseif est_method == "batch_maximum"
        Î±, ğ�’© = compute_Î±_batch_maximum(â„³, ZÌ‚)
    elseif est_method == "product_upper_bound"
        Î±, ğ�’© = compute_Î±_product_upper_bound(â„³, ZÌ‚)
    elseif est_method == "explog_approx"
        Î±, ğ�’© = compute_Î±_explog_approx(â„³, ZÌ‚)
    elseif est_method == "arithmetic_mean"
        Î±, ğ�’© = compute_Î±_arithmetic_mean(â„³, ZÌ‚)
    elseif est_method == "constant"
        Î±, ğ�’© = tfrec.Î±, zero(tfrec.Î±)
    else
        error(
            "Unknown Î± estimation method: $(est_method). Please choose one of $(Î±_METHODS).",
        )
    end
    return return_norm ? (Î±, ğ�’©) : Î±
end

"""
    compute_Î±_jacobian_upper_bound(â„³)
	
Compute upper bound on `Î±` via parameters of model `â„³`.
"""
function compute_Î±_jacobian_upper_bound(â„³)
    ğ�’© = norm_upper_bound(â„³)
    return opt_Î±(ğ�’©), ğ�’©
end

"""
    compute_Î±_batch_maximum(â„³, ZÌ‚)

Compute `Î±` via the maximum Jacobian spectral norm of a given batch of sequences.
"""
function compute_Î±_batch_maximum(â„³, ZÌ‚::AbstractArray{T, 3}) where {T}
    # flatten batch and time dimensions
    ZÌ‚_flat = reshape(ZÌ‚, size(ZÌ‚, 1), :)

    # compute jacobian spectral norms
    ğ�’©s = ThreadsX.map(z -> opnorm(jacobian(â„³, z)), eachcol(ZÌ‚_flat))

    # max spectral norm across entire batch
    ğ�’© = maximum(ğ�’©s)
    return opt_Î±(ğ�’©), ğ�’©
end

"""
    compute_Î±_explog_approx(â„³, ZÌ‚)

Compute `Î±` using and the "exp-log" approximation.
"""
function compute_Î±_explog_approx(â„³, ZÌ‚::AbstractArray{T, 3}) where {T}
    # geomean of spectral norms along sequences
    ğ�’©s = ThreadsX.map(seq -> opnorm(geomean(jacobian(â„³, seq))), eachslice(ZÌ‚, dims = 2))
    ğ�’© = maximum(ğ�’©s)
    return opt_Î±(ğ�’©), ğ�’©
end

"""
    compute_Î±_product_upper_bound(â„³, ZÌ‚)

Compute `Î±` using the upper bound on the Jacobian product norm along
a training sequence.
"""
function compute_Î±_product_upper_bound(â„³, ZÌ‚::AbstractArray{T, 3}) where {T}
    # geomean of spectral norms along sequences
    ğ�’©s = ThreadsX.map(seq -> geomean(opnorm.(jacobian(â„³, seq))), eachslice(ZÌ‚, dims = 2))
    ğ�’© = maximum(ğ�’©s)
    return opt_Î±(ğ�’©), ğ�’©
end

"""
    compute_Î±_arithmetic_mean(â„³, ZÌ‚)

Compute `Î±` using the arithmetic mean of Jacobians along a training sequence.
"""
function compute_Î±_arithmetic_mean(â„³, ZÌ‚::AbstractArray{T, 3}) where {T}
    # geomean of spectral norms along sequences
    ğ�’©s = ThreadsX.map(seq -> opnorm(mean(jacobian(â„³, seq))), eachslice(ZÌ‚, dims = 2))
    ğ�’© = maximum(ğ�’©s)
    return opt_Î±(ğ�’©), ğ�’©
end

"""
    opt_Î±(ğ�’©)

Compute optimal GTF `Î±` for given spectral norm `ğ�’©`.
"""
opt_Î±(ğ�’©) = max(0, 1 - 1 / ğ�’©)