text:
export HLAData
export load_data

function load_data(save_file::String)
    isfile(save_file) || error("File $save_file does not exist.")

    data = load(save_file)["data"]

    return data
end

mutable struct HLAData <: AbstractHLAData
    name::String
    records::Vector{FASTX.FASTA.Record}
    hla_types::Vector{HLAType}
    tree::Union{Missing, PhylogeneticTree}
    stan_input::Union{Missing, Dict{String, Any}}
    allele_depth::Int
    replacement_mincount::Int

    function HLAData(
        name::String,
        records::Vector{FASTX.FASTA.Record},
        hla_types::Vector{HLAType},
        tree::Union{Missing, PhylogeneticTree},
        stan_input::Union{Missing, Dict{String, Any}},
        allele_depth::Int,
        replacement_mincount::Int
    )

        if length(hla_types) != length(records)
            msg = "Vector of HLATypes has size $(length(hla_types)), " *
                "expected $(length(records))"
            error(msg)
        end

        if !ismissing(tree)
            m = Escape.matching(tree, records)
            m isa Exception && throw(m)
        end

        if !ismissing(stan_input)
            m = ErrorException("stan_input is not in a valid format.")
            is_valid(stan_input) || throw(m)
        end

        if !(allele_depth == 1 || allele_depth == 2)
            throw(ErrorException("allele_depth must be 1 (for 2-digit accuracy " *
                "or 2 (for 4-digit accuravy)."))
        end

        if !(replacement_mincount > 0)
            throw(ErrorException("replacement_mincount must be at least 1."))
        end

        if replacement_mincount > length(records)
            throw(ErrorException("replacement_mincount must not exceed the number of" *
                "available sequences."))
        end

        new(name, records, hla_types, tree, stan_input, allele_depth, replacement_mincount)
    end
end

function HLAData(name::String, fasta_file::String, hla_types::Vector{HLAType};
        allele_depth::Int = 1, replacement_mincount::Int = 1)

    records = Escape.records(fasta_file)

    return HLAData(name, records, hla_types, missing, missing, allele_depth,
        replacement_mincount)
end

function HLAData(
    name::String,
    fasta_file::String,
    hla_types::Vector{HLAType},
    tree::Union{Missing, PhylogeneticTree},
    stan_input::Union{Missing, Dict{String, Any}};
    allele_depth::Int = 1,
    replacement_mincount::Int = 1
)
    records = Escape.records(fasta_file)

    return HLAData(name, records, hla_types, tree, stan_input, allele_depth,
        replacement_mincount)
end

function HLAData(;
    alignment_file::String = "",
    tree_file::String = "",
    hla_annotation_file::String = "",
    name = splitext(basename(alignment_file))[1],
    allele_depth = 1,
    replacement_mincount = 1,
    save_file = ""
)

    if save_file != ""
        save(save_file, Dict("data" => "test"))
        rm(save_file)
    end

    tree = phylogenetic_tree(readline(tree_file))

    if hla_annotation_file == ""
        check_inputs(alignment_file, tree_file)
        hla_types = let
            r = records(alignment_file)
            hla_types = HLAType[]

            for record in r
                annotation = FASTA.identifier(record)
                push!(hla_types, HLAType(Tuple(parse_allele.(split(annotation, '_')))))
            end

            hla_types
        end
    else
        check_inputs(alignment_file, tree_file, hla_annotation_file)
        df = hla_annotation_df(hla_annotation_file)
        hla_types = let
            hla_types = HLAType[]

            for row in eachrow(df)
                push!(hla_types, HLAType(tuple(
                    row[:A1], row[:A2], row[:B1], row[:B2], row[:C1], row[:C2]
                )))
            end

            hla_types
        end
    end

    data = Escape.HLAData(name, alignment_file, hla_types, tree, missing,
        allele_depth = allele_depth, replacement_mincount = replacement_mincount)
    stan_input = Escape.compute_stan_input(
        data, allele_depth = allele_depth, replacement_mincount = replacement_mincount
    )

    data = HLAData(name, alignment_file, hla_types, tree, stan_input,
        allele_depth = allele_depth, replacement_mincount = replacement_mincount)

    if save_file != ""
        save(save_file, Dict("data" => data))
    end

    return data
end

function Base.show(io::IO, x::HLAData)
    print(io, """HLAData("$(name(x))", $(length(records(x))) records)""")
end

function hla_annotation_df(hla_annotation_file::String)
    df = CSV.read(hla_annotation_file, DataFrame)

    if names(df) != ["identifier", "A1", "A2", "B1", "B2", "C1", "C2"]
        error("Expected column names are [:identifier, :A1, :A2, :B1, :B2, :C1, :C2], " *
            "got $(names(df)).")
    end

    N = nrow(df)
    parsed_df = DataFrame(
        :identifier => Vector{String}(undef, N),
        :A1 => Vector{HLAAllele}(undef, N), :A2 => Vector{HLAAllele}(undef, N),
        :B1 => Vector{HLAAllele}(undef, N), :B2 => Vector{HLAAllele}(undef, N),
        :C1 => Vector{HLAAllele}(undef, N), :C2 => Vector{HLAAllele}(undef, N),
    )

    parsed_df[!, :identifier] = df[!, :identifier]

    for col in (:A1, :A2, :B1, :B2, :C1, :C2)
        parsed_df[!, col] = parse_allele.(df[!, col])

        if mean(ismissing.(df[!, col])) > 0.3
            @warn "HLA annotation column $col contains more than 30% missing values. " *
                "Please check if HLA alleles are parsed correctly using the " *
                """parse_allele(::String) function, e.g. parse_allele("A03")"""
        end
    end

    return parsed_df
end

function check_inputs(alignment_file, tree_file, hla_annotation_file)
    alignment_error_string = "Alignment file $alignment_file does not exist."
    tree_error_string  = "Phylogenetic tree file $tree_file does not exist."
    hla_annotation_string = "HLA annotation file $hla_annotation_file does not exist."

    isfile(alignment_file) || throw(ArgumentError(alignment_error_string))
    isfile(tree_file) || throw(ArgumentError(tree_error_string))
    isfile(hla_annotation_file) || throw(ArgumentError(hla_annotation_string))

    df = hla_annotation_df(hla_annotation_file)
    r = records(alignment_file)
    tree = phylogenetic_tree(readline(tree_file))
    leaf_names = Set(map(x -> get_property(tree, x, :name), leaves(tree)))

    if length(leaf_names) != length(r)
        throw(ArgumentError("Number of leaves of the phylogenetic tree " *
            "($(length(leaf_names))) does not match number of records " *
            "($(length(r)))"))
    end

    idx_labeled_leaves = Set(leaf_names) == Set(string.(1:length(r)))

    for record in r
        annotation = FASTA.identifier(record)

        if !(annotation ∈ df[!, :identifier])
            throw(ArgumentError("Sequence with ID $annotation does not " *
                "exist in the HLA annotation file."))
        end

        if !(idx_labeled_leaves) & !(annotation ∈ leaf_names)
            throw(ArgumentError("Sequence with ID $annotation does not " *
                "exist in phylogenetic tree."))
        end
    end

    return nothing
end

function check_inputs(alignment_file, tree_file)
    alignment_error_string = "Alignment file $alignment_file does not exist."
    tree_error_string  = "Phylogenetic tree file $tree_file does not exist."

    isfile(alignment_file) || throw(ArgumentError(alignment_error_string))
    isfile(tree_file) || throw(ArgumentError(tree_error_string))

    r = records(alignment_file)
    tree = phylogenetic_tree(readline(tree_file))
    leaf_names = Set(map(x -> get_property(tree, x, :name), leaves(tree)))

    if length(leaf_names) != length(r)
        throw(ArgumentError("Number of leaves of the phylogenetic tree " *
            "($(length(leaf_names))) does not match number of records " *
            "($(length(r)))"))
    end

    if Set(leaf_names) != Set(string.(1:length(r)))
        throw(ArgumentError("Leaves of the phylogenetic tree are not labeled with " *
            "1, 2, ..., $(length(r))."))
    end

    for record in r
        annotation = FASTA.identifier(record)
        if length(split(annotation, '_')) != 6
            throw(ArgumentError("Sequence with annotation $annotation does not provide
                6 HLA alleles (use 'missing' to indicate missing alleles)."))
        end
    end

    return nothing
end

name(data::AbstractHLAData) = getfield(data, :name)
records(data::AbstractHLAData) = getfield(data, :records)
hla_types(data::AbstractHLAData) = getfield(data, :hla_types)
stan_input(data::AbstractHLAData) = getfield(data, :stan_input)
allele_depth(data::AbstractHLAData) = getfield(data, :allele_depth)
replacement_mincount(data::AbstractHLAData) = getfield(data, :replacement_mincount)

function compute_stan_input(
    data::AbstractHLAData;
    allele_depth::Int = 1, replacement_mincount::Int = 1
)
    hla_types = Escape.hla_types(data)

    X = Float64.(hla_matrix(hla_types; allele_depth = allele_depth))
    check_hla_matrix(X)
    no_hlas = map(x -> all(x .== 0), eachrow(X))
    with_hla = .!no_hlas

    r = replacements(data)
    phy = phylogeny_information(data, r)
    Z = epitope_information(data, r, allele_depth)

    y = Int[]
    rs = Int[]
    idx = Int[]

    for i in 1:length(r)
        y_i = targets(r[i], data)
        c = .!ismissing.(y_i) .& with_hla
        append!(y, y_i[c])
        append!(rs, repeat([i], sum(c)))
        append!(idx, findall(c))
    end

    d = Dict(
        "N" => length(y),
        "S" => size(X, 1),
        "D" => size(X, 2),
        "R" => length(r),
        "X" => X,
        "y" => y,
        "rs" => rs,
        "phy" => phy,
        "Z" => Z,
        "idx" => idx
    )

    return d
end

function check_hla_matrix(X::AbstractMatrix)
    size(X, 2) != 0 || error("Could not parse any HLA alleles with " *
        "allele_depth = $allele_depth. " *
        "Try setting allele_depth = 1 for HLA alleles with 2-digit accuracy.")
    no_hlas = map(x -> all(x .== 0), eachrow(X))

    if any(no_hlas)
        @warn "No HLA alleles parsed for sequences with index: " *
            "$(findall(no_hlas))). " *
            "These sequences are not included in the model."
    end

    return nothing
end

function epitope_information(
    data::AbstractHLAData, r::Vector{Replacement},
    allele_depth::Int
)
    N = length(unique_alleles(data.hla_types, allele_depth = allele_depth))
    Z = Escape.epitope_feature_matrix(
        data, allele_depth = allele_depth
    )[map(x -> x.position, r), :]

    return Z
end

function phylogeny_information(
    data::AbstractHLAData, r::Vector{Replacement}
)

    if nworkers() == 1
        @warn "Calculating phylogeny information on a single worker only. It is strongly " *
            "advised to add more workers with: \n using Distributed \n " *
            "addprocs($(Sys.CPU_THREADS)) \n @everywhere using Escape \n"
    end

    R = length(r)
    N = length(data.hla_types)
    tree = phylogenetic_tree(data, verbose = true)

    phy = @showprogress "phylogeny " @distributed hcat for i in 1:R
        phylogeny_probabilities(r[i], data, tree)
    end
    phy = phy'

    return phy
end

function is_valid(input::Dict{String, Any})
    m = ErrorException("stan_input is not in a valid format.")

    required_keys = ("Z", "rs", "S","idx", "X", "y", "N", "D", "R", "phy")
    all(map(x -> x in keys(input), required_keys)) || return false
    size(input["Z"]) == (input["R"], input["D"]) || return false
    maximum(input["rs"]) <= input["R"] || return false
    maximum(input["idx"]) <= input["S"] || return false
    size(input["X"]) == (input["S"], input["D"]) || return false
    size(input["phy"]) == (input["R"], input["S"]) || return false

    return true
end

function sequence_length(data::AbstractHLAData)
    records = Escape.records(data)
    sequence_length = maximum([length(FASTA.sequence(r)) for r in records])

    return sequence_length
end

function records(fasta_file::String)
    records = FASTA.Record[]

    open(FASTA.Reader, fasta_file) do reader
        for record in reader
            push!(records, record)
        end
    end

    return records
end


text:
module FinEtoolsFlexBeamsTutorials

end # module

ext: jl


text:
### Fitting by KDE

function calcrange(::Type{KernelDensity.UnivariateKDE}, x::AbstractVector, binNumber)
    bw = KernelDensity.default_bandwidth(x)
    bb = KernelDensity.kde_boundary(x, bw)
    mps = Float64[]
    i = 1
    while length(mps) < binNumber # range correction
        stp = (bb[2]-bb[1])/(binNumber-i)
        mps = bb[1]:stp:bb[2]+eps()
        # println("$i: $binNumber => $(bb[1]):$stp:$(bb[2]) ", length(mps))
        i -=1
    end
    return mps
end

function fit(::Type{EmpiricalLinearManifold}, ::Type{KernelDensity.UnivariateKDE},
             μ::AbstractVector{T}, B::AbstractMatrix{T}, X::AbstractMatrix{T}; kvargs...) where {T <: Real}

    K = size(B,2)

    nquants = 1024 # default
    for (k,v) in kvargs
        if k == :nquants
            nquants = Int(v)
        end
    end

    estimates = Array{KernelDensity.UnivariateKDE}(undef, K+1)

    zk = B'*(X .- μ)
    for k in 1:K
        estimates[k] = KernelDensity.kde(zk[k,:], calcrange(KernelDensity.UnivariateKDE, zk[k,:], nquants))
    end

    d = LMCLUS.distance_to_manifold(X, μ, B)
    estimates[K+1] = KernelDensity.kde(d, calcrange(KernelDensity.UnivariateKDE, d, nquants))

	return EmpiricalLinearManifold(μ, B, estimates)
end

fit(::Type{EmpiricalLinearManifold}, ::Type{KernelDensity.UnivariateKDE},
    LMC::LMCLUS.Manifold, X::AbstractMatrix{T}; kvargs...) where {T <: Real} =
    fit(EmpiricalLinearManifold, KernelDensity.UnivariateKDE, LMCLUS.mean(LMC), LMCLUS.projection(LMC), X; kvargs...)

### Evaluation

rectint(y, r) = y*r
rectint(y, x1, x2) = rectint(y, x2 - x1)
trapint(y1, y2, r) = r.*(y1.+y2)/2
trapint(y1, y2, x1, x2) = trapint(y1, y2, x2 .- x1)
simpint(y1, y2, y3, x1, x2) = (x2 .- x1).*(y1 .+ 4*y2 .+ y3)/6

function _pdf(d::EmpiricalLinearManifold{T, KernelDensity.UnivariateKDE},
              x::AbstractVector{T}) where {T<:Real}
    K = size(d.B,2)
    prob = one(T)
    z = d.B'*(x - d.μ)
    dist = LMCLUS.distance_to_manifold(collect(x), d.μ, d.B) |> first
    for k in 1:K+1
        val = k > K ? dist : z[k] # Orthonormal subspace distances for K+1
        stp = step(d.estimate[k].x)

        # Search-based
        # binIndex = searchsortedlast(d.estimate[k].x, val)
        # dv, x1, x2 = if binIndex < 1
        #     d.estimate[k].density[1], d.estimate[k].x[1]-stp, d.estimate[k].x[1]
        # elseif binIndex > length(d.estimate[k].x)
        #     d.estimate[k].density[end], d.estimate[k].x[end], d.estimate[k].x[end]+stp
        # else
        #     d.estimate[k].density[binIndex], d.estimate[k].x[binIndex], d.estimate[k].x[binIndex]+stp
        # end
        # tmp = rectint(dv, x1, x2)

        # Approximation-based
        x1 = val-stp/2
        x2 = val
        x3 = val+stp/2
        # tmp = rectint(pdf(d.estimate[k], x2), x1, x3)
        tmp = trapint(pdf(d.estimate[k], x1), pdf(d.estimate[k], x3), x1, x3)
        # tmp = simpint(pdf(d.estimate[k], x1), pdf(d.estimate[k], x2), pdf(d.estimate[k], x3), x1, x3)

        # if probability is negative
        if tmp > 0.
            prob *= tmp
        else
            prob = 0.
        end
    end
    return prob
end

function _pdf!(r::AbstractArray{T}, d::EmpiricalLinearManifold{T,KernelDensity.UnivariateKDE},
               X::AbstractMatrix{T}) where {T<:Real}
    K = size(d.B,2)
    N, n = size(X)
    Z = (d.B'*(X .- d.μ))'
    dist = LMCLUS.distance_to_manifold(X, d.μ, d.B)
    for i in 1 : n
        @inbounds r[i] = one(T)
    end
    for k in 1 : K+1
        val = k > K ? dist : Z[:,k] # Orthonormal subspace distances for K+1
        stp = step(d.estimate[k].x)
        x1 = val .- stp/2
        x3 = val .+ stp/2
        tmp = trapint(pdf(d.estimate[k], x1), pdf(d.estimate[k], x3), x1, x3)
        # tmp = simpint(pdf(d.estimate[k], x1), pdf(d.estimate[k], val), pdf(d.estimate[k], x3), x1, x3)
        for i in 1 : n
            @inbounds r[i] *= tmp[i] < zero(T) ? zero(T) : tmp[i]
        end
    end
    return r
end

### Entropy

# function Distributions.entropy(d::EmpiricalLinearManifold{T, KernelDensity.UnivariateKDE}) where {T<:Real}
#     H = 0.0
#     for est in d.estimate
#         tot = length(est.density)
#         ds = filter(x->x != zero(T), est.density)
#         H += -sum(log(ds))/tot
#     end
#     return H
# end

function entropy(d::EmpiricalLinearManifold{T, KernelDensity.UnivariateKDE}) where {T<:Real}
    H = 0.0
    qb = first(d.estimate).x |> length
    dp = ones(qb-1)
    for est in d.estimate
        stp = step(est.x)
        tmp = [trapint(est.density[i], est.density[i+1], est.x[i], est.x[i+1]) for i in 1:length(est.density)-1]
        dp = dp .* tmp
    end
    H = entropy(dp)
    @debug "Empirical LM model entropy" EPDF=dp H=H
    return H
end

function entropy(M::MixtureModel{Q,R,EmpiricalLinearManifold{T, KernelDensity.UnivariateKDE}}) where {Q,R,T<:Real}
    qb = first(first(components(M)).estimate).x |> length
    c = length(components(M))
    H = zero(T)
    for i in 1:qb-1
        Pᵢⱼ = zero(T)
        for j in 1:c
            d = components(M)[j]
            pⱼ = probs(M)[j]
            prob = one(T)
            for est in d.estimate
                prob *= trapint(est.density[i], est.density[i+1], est.x[i], est.x[i+1])
            end
            Pᵢⱼ += pⱼ*prob
            @debug "Pᵢⱼ" i j Pᵢⱼ
        end
        H += Pᵢⱼ > zero(T) ? -log(Pᵢⱼ) : zero(T)
        # H += Pᵢⱼ > zero(T) ? -Pᵢⱼ*log(Pᵢⱼ) : zero(T)
        @debug "H" i H
    end
    @debug "Empirical LM mixture model entropy" H=H
    return H
end

### Sampling

function _rand!(s::EmpiricalLinearManifoldSampler{T, KernelDensity.UnivariateKDE}, x::AbstractVector{T}) where {T <: Real}
    return x
end


text:
#Todo:
#PkgBenchmark
#HAC s.e.s from CovariancesMatrices.jl?
#Float16/32 don't seem to work anymore. Problem in Optim?
#support missing data? timeseries?
#implement lrtest
#allow uninititalized constructors for UnivariateVolatilitySpec, MeanSpec and StandardizedDistribution? If so, then be consistent with how they are defined
#  (change for meanspec and dist ), document, and test. Also, NaN is prob. safer than undef.
#logconst needs to return the correct type
"""
The ARCHModels package for Julia. For documentation, see https://s-broda.github.io/ARCHModels.jl/dev.
"""
module ARCHModels
using Reexport
@reexport using StatsBase
using StatsFuns: normcdf, normccdf, normlogpdf, norminvcdf, log2π, logtwo, RFunctions.tdistrand, RFunctions.tdistinvcdf, RFunctions.gammarand, RFunctions.gammainvcdf
using GLM: modelmatrix, response, LinearModel
using SpecialFunctions: beta, gamma, digamma #, lgamma

# work around https://github.com/JuliaMath/SpecialFunctions.jl/issues/186
# until https://github.com/JuliaDiff/ForwardDiff.jl/pull/419/ is merged
# remove test in runtests.jl as well when this gets fixed
using Base.Math: libm
using ForwardDiff: Dual, value, partials
@inline lgamma(x::Float64) = ccall((:lgamma, libm), Float64, (Float64,), x)
@inline lgamma(x::Float32) = ccall((:lgammaf, libm), Float32, (Float32,), x)
@inline lgamma(d::Dual{T}) where T = Dual{T}(lgamma(value(d)), digamma(value(d)) * partials(d))

using Optim
using ForwardDiff
using FiniteDiff
using Distributions
using HypothesisTests
using Roots
using LinearAlgebra
using DataStructures: CircularBuffer
using DelimitedFiles
using Statistics: cov

import Distributions: quantile
import Base: show, showerror, eltype
import Statistics: mean
import Random: rand, AbstractRNG
import HypothesisTests: HypothesisTest, testname, population_param_of_interest, default_tail, show_params, pvalue
import StatsBase: StatisticalModel, stderror, loglikelihood, nobs, fit, fit!, confint, aic,
                  bic, aicc, dof, coef, coefnames, coeftable, CoefTable,
				  informationmatrix, islinear, score, vcov, residuals, predict
import StatsModels: TableRegressionModel
export ARCHModel, UnivariateARCHModel, UnivariateVolatilitySpec, StandardizedDistribution, Standardized, MeanSpec,
       simulate, simulate!, selectmodel, StdNormal, StdT, StdGED, StdSkewT, Intercept, Regression,
       NoIntercept, ARMA, AR, MA, BG96, volatilities, mean, quantile, VaRs, pvalue, means, VolatilitySpec,
	   MultivariateVolatilitySpec, MultivariateStandardizedDistribution, MultivariateARCHModel, MultivariateStdNormal,
	   EGARCH, ARCH, GARCH, TGARCH, ARCHLMTest, DQTest,
	   DOW29, DCC, CCC, covariances, correlations

include("utils.jl")
include("general.jl")
include("univariatearchmodel.jl")
include("meanspecs.jl")
include("univariatestandardizeddistributions.jl")
include("EGARCH.jl")
include("TGARCH.jl")
include("tests.jl")
include("multivariatearchmodel.jl")
include("multivariatestandardizeddistributions.jl")
include("DCC.jl")
end#module


text:
using Distributions, NLsolve, Statistics

function effectsize(xs::AbstractVector, ys::AbstractVector)
    if length(xs) != length(ys)
        throw(ArgumentError("samples must have the same number of observations"))
    end
    n, m = length(xs), length(ys)
    x̄, ȳ = mean(xs), mean(ys)
    σx², σy² = var(xs), var(ys)
    σ = sqrt(0.5*(σx² + σy²))
    abs(x̄ - ȳ)/σ
end

function power(xs::AbstractVector, ys::AbstractVector, n::Int=length(xs), p::Float64=0.05; unequalvar=true)
    dof = if unequalvar
        varx, vary = var(xs), var(ys)
        nx, ny = length(xs), length(ys)
        (varx / nx + vary / ny)^2 / ((varx / nx)^2 / (nx - 1) + (vary / ny)^2 / (ny - 1))
    else
        length(xs) + length(ys) - 2
    end
    power(effectsize(xs, ys), n, dof, p)
end

function power(d::Float64, n::Int=length(xs), dof=2n-2, α::Float64=0.05)
    δ = d*sqrt(n/2)
    tc = quantile(TDist(dof), 1 - α/2)
    ccdf(NoncentralT(dof, δ), tc) + cdf(NoncentralT(dof, δ), -tc)
end

function power(d::Float64, n::Float64=length(xs), dof=2n-2, α::Float64=0.05)
    δ = d*sqrt(n/2)
    tc = quantile(TDist(dof), 1 - α/2)
    ccdf(NoncentralT(dof, δ), tc) + cdf(NoncentralT(dof, δ), -tc)
end

function solvepower(d::Float64, dof::Float64, β::Float64=0.8, α::Float64=0.05)
    first(nlsolve((F,x) -> F[1] = power(d, x[1], dof, α) - β, [2.0]).zero)
end

function solvepower(xs::AbstractVector, ys::AbstractVector, β::Float64=0.8, α::Float64=0.05; unequalvar=true)
    dof = if unequalvar
        varx, vary = var(xs), var(ys)
        nx, ny = length(xs), length(ys)
        (varx / nx + vary / ny)^2 / ((varx / nx)^2 / (nx - 1) + (vary / ny)^2 / (ny - 1))
    else
        length(xs) + length(ys) - 2
    end
    solvepower(effectsize(xs, ys), dof, β, α)
end

max_stars_repo_name: ELIFE-ASU/dglmoore

max_issues_repo_name: ELIFE-ASU/dglmoore

max_forks_repo_name: ELIFE-ASU/dglmoore


text:
const ToT = UInt8
const ChannelID = UInt8
const DOMID = UInt32
const Floor = UInt8
const DU = UInt8
const HitTime = Float64

struct OMKey
    dom_id::UInt32
    channel_id::UInt8
end

struct Position{T} <: FieldVector{3, T}
    x::T
    y::T
    z::T

    Position(x::T, y::T, z::T) where {T} = new{T}(x, y, z)
end

Position(x, y, z) = Position(promote(x, y, z)...)

# Base.*(x::Vector3D, y::Vector3D ) = Vector3D(SVector(y)* SVector(x))

struct Direction{T} <: FieldVector{3, T}
    x::T
    y::T
    z::T
end

Direction(x, y, z) = Direction(promote(x, y, z)...)

function Direction(ϕ, θ)
    Direction(cos(ϕ)*cos(θ), sin(ϕ)*cos(θ), sin(θ))
end

Base.show(io::IO, p::Position) = begin
    s = Printf.@sprintf "%.2f %.2f %.2f" p.x p.y p.z
    print(io, s)
end

Base.show(io::IO, d::Direction) = begin
    s = Printf.@sprintf "%.2f %.2f %.2f" d.x d.y d.z
    print(io, s)
end

struct MCEventInfo
    event_id::Int64
    group_id::Int64
    mc_id::Int64
    mc_time::Float64
    nanoseconds::Int64
    run_id::Int64
    timestamp::Int64
    weight_w1::Float64
    weight_w2::Float64
    weight_w3::Float64
    weight_w4::Float64
end

Base.show(io::IO, e::MCEventInfo) = begin
    print(io, "MCEventInfo: id($(e.event_id)), mc_id($(e.mc_id)), " *
          "mc_time($(e.mc_time) ($(e.nanoseconds)), ts($(e.timestamp)), run_id($(e.run_id))")
end

struct DAQEventInfo
    det_id::Int64
    event_id::Int64
    frame_index::Int64
    group_id::Int64
    mc_run_id::Int64
    mc_time::Float64
    nanoseconds::Int64
    overlays::Int64
    run_id::Int64
    timestamp::Int64
    trigger_counter::Int64
    trigger_mask::Int64
    weight_w1::Float64
    weight_w2::Float64
    weight_w3::Float64
    weight_w4::Float64
end

Base.show(io::IO, e::DAQEventInfo) = begin
    print(io, "DAQEventInfo: id($(e.event_id)), det_id($(e.det_id)), " *
          "timestamp($(e.timestamp) ($(e.nanoseconds)), run_id($(e.run_id))")
end

struct JMuon
    JENERGY_CHI2::Float64
    JENERGY_ENERGY::Float64
    JENERGY_MUON_RANGE_METRES::Float64
    JENERGY_NDF::Float64
    JENERGY_NOISE_LIKELIHOOD::Float64
    JENERGY_NUMBER_OF_HITS::Float64
    JGANDALF_BETA0_RAD::Float64
    JGANDALF_BETA1_RAD::Float64
    JGANDALF_CHI2::Float64
    JGANDALF_LAMBDA::Float64
    JGANDALF_NUMBER_OF_HITS::Float64
    JGANDALF_NUMBER_OF_ITERATIONS::Float64
    JMUONENERGY::Bool
    JMUONGANDALF::Bool
    JMUONPREFIT::Bool
    JMUONSIMPLEX::Bool
    JMUONSTART::Bool
    JSTART_LENGTH_METRES::Float64
    JSTART_NPE_MIP::Float64
    JSTART_NPE_MIP_TOTAL::Float64
    JVETO_NPE::Float64
    JVETO_NUMBER_OF_HITS::Float64
    dir_x::Float64
    dir_y::Float64
    dir_z::Float64
    energy::Float64
    group_id::Int64
    id::Int64
    length::Float64
    likelihood::Float64
    pos_x::Float64
    pos_y::Float64
    pos_z::Float64
    rec_type::Int64
    time::Float64
end

struct TimesliceInfo
    frame_index::UInt32
    slice_id::UInt32
    timestamp::UInt32
    nanoseconds::UInt32
    n_frames::UInt32
    group_id::UInt32
end

struct Track
    dir::Direction
    pos::Position
    time
end

# Fit
abstract type AbstractRecoTrack end

struct RecoTrack<:AbstractRecoTrack
    dir::Direction
    pos::Position
    time::Float64
end

struct NoRecoTrack<:AbstractRecoTrack end

# MC
struct MCTrack
    bjorken_y::Float64
    dir_x::Float64
    dir_y::Float64
    dir_z::Float64
    E::Float64
    group_id::Int64
    id::Int64
    interaction_channel::Int64
    is_cc::Bool
    length::Float64
    pos_x::Float64
    pos_y::Float64
    pos_z::Float64
    t::Float64
    particle_type::Int64
end

Track(t::MCTrack) = Track([t.dir_x, t.dir_y, t.dir_z], [t.pos_x, t.pos_y, t.pos_z], t.t)

Base.show(io::IO, t::MCTrack) = begin
    E = Printf.@sprintf "%0.1f" t.E
    bjorken_y = Printf.@sprintf "%0.2f" t.bjorken_y
    print(io, "MCTrack: bjorken_y($(bjorken_y)), t($(t.t)), " *
              "pos($(t.pos_x), $(t.pos_y), $(t.pos_z)), " *
              "dir($(t.dir_x), $(t.dir_y), $(t.dir_z)), " *
              "E($(E)), type($(t.particle_type))")
end

# Hardware
struct PMT
    channel_id::ChannelID
    pos::Position
    dir::Direction
end

struct DOM
    id::UInt32
    floor::Floor
    du::DU
    pmts::Vector{PMT}
end

struct Calibration
    det_id::Int32
    pos::Dict{Int32,Vector{NeRCA.Position}}
    dir::Dict{Int32,Vector{NeRCA.Direction}}
    t0::Dict{Int32,Vector{Float64}}
    du::Dict{Int32,DU}
    floor::Dict{Int32,Floor}
    omkeys::Dict{Int32,OMKey}
    max_z
    n_dus
end

Base.show(io::IO, c::Calibration) = begin
    print(io, "Calibration data for detector '$(c.det_id)' " *
              "with $(length(c.pos)) modules.")
end

# Signal
abstract type AbstractHit end
abstract type DAQHit<:AbstractHit end

Base.isless(lhs::AbstractHit, rhs::AbstractHit) = lhs.t < rhs.t

struct Hit <: DAQHit
    channel_id::ChannelID
    dom_id::DOMID
    t::HitTime
    tot::ToT
    triggered::Bool
end

Base.show(io::IO, h::DAQHit) = begin
    print(io, "$(typeof(h)): DOM ID $(h.dom_id), channel ID, $(h.channel_id), t=$(h.t), tot=$(h.tot)")
end

struct McHit <: AbstractHit
    a::Float32
    origin::UInt32
    pmt_id::UInt32
    t::HitTime
end

mutable struct Multiplicity
    count::Int32
    id::Int64
end

struct CalibratedHit <: DAQHit
    channel_id::ChannelID
    dom_id::UInt32
    du::DU
    floor::Floor
    t::HitTime
    tot::ToT
    pos::Position
    dir::Direction
    t0::HitTime
    triggered::Bool
    multiplicity::Multiplicity
end

struct TimesliceHit <: DAQHit
    channel_id::Int8
    dom_id::UInt32
    t::Int32
    tot::Int16
end

Hit(hit::HDF5.HDF5Compound{5}) = begin
    Hit(hit.data...)
end

const TriggerMask = Int64

struct DAQTriggeredHit <: DAQHit
    dom_id::Int32
    channel_id::UInt8
    t::Int32
    tot::UInt8
    trigger_mask::TriggerMask
end

struct DAQEvent
    det_id::Int32
    run_id::Int32
    timeslice_id::Int32
    timestamp::Int32
    ticks::Int32
    trigger_counter::Int64
    trigger_mask::Int64
    overlays::Int32
    n_triggered_hits::Int32
    triggered_hits::Vector{DAQTriggeredHit}
    n_hits::Int32
    hits::Vector{Hit}
end

Base.show(io::IO, d::DAQEvent) = begin
    print(io, "DAQEvent: $(d.n_triggered_hits) triggered hits, " *
              "$(d.n_hits) snapshot hits")
end

struct DAQEventFile
    filename
    n_events
    _fobj::HDF5.HDF5File
    _hit_indices
    _event_infos::Vector{DAQEventInfo}

    function DAQEventFile(filename)
        fobj = HDF5.h5open(filename, "r")
        event_infos = read_compound(fobj, "/event_info", DAQEventInfo)
        hit_indices = read_indices(fobj, "/hits")

        new(filename, length(event_infos), fobj, hit_indices, event_infos)
    end
end

Base.show(io::IO, f::DAQEventFile) = begin
    print(io, "DAQEventFile(\"$(f.filename)\")")
end

function read_io(io::IOBuffer, t::T) where T
    length = read(io, Int32)
    type = read(io, Int32)
    det_id = read(io, Int32)
    run_id = read(io, Int32)
    timeslice_id = read(io, Int32)
    timestamp = read(io, Int32)
    ticks = read(io, Int32)
    trigger_counter = read(io, Int64)
    trigger_mask = read(io, Int64)
    overlays = read(io, Int32)

    n_triggered_hits = read(io, Int32)
    triggered_hits = Vector{DAQTriggeredHit}()
    sizehint!(triggered_hits, n_triggered_hits)
    triggered_map = Dict{Tuple{Int32, UInt8, Int32, UInt8}, Int64}()
    @inbounds for i ∈ 1:n_triggered_hits
        dom_id = read(io, Int32)
        channel_id = read(io, UInt8)
        time = bswap(read(io, Int32))
        tot = read(io, UInt8)
        trigger_mask = read(io, Int64)
        triggered_map[(dom_id, channel_id, time, tot)] = trigger_mask
        push!(triggered_hits, DAQTriggeredHit(dom_id, channel_id, time, tot, trigger_mask))
    end

    n_hits = read(io, Int32)
    hits = Vector{Hit}()
    sizehint!(hits, n_hits)
    @inbounds for i ∈ 1:n_hits
        dom_id = read(io, Int32)
        channel_id = read(io, UInt8)
        time = bswap(read(io, Int32))
        tot = read(io, UInt8)
        key = (dom_id, channel_id, time, tot)
        triggered = false
        if haskey(triggered_map, key)
            triggered = true
        end
        push!(hits, Hit(channel_id, dom_id, time, tot, triggered))
    end

    DAQEvent(det_id, run_id, timeslice_id, timestamp, ticks, trigger_counter, trigger_mask, overlays, n_triggered_hits, triggered_hits, n_hits, hits)
end

struct MCEventReader
    filename::AbstractString
    detx::AbstractString
    _fobj::HDF5.HDF5File
    _calib::Calibration
    _event_infos::Vector{MCEventInfo}
    _mc_tracks::Dict{Int64, Vector{MCTrack}}
    _length::UInt64

    function MCEventReader(filename, detx)
        fobj = h5open(filename, "r")
        calib = read_calibration(detx)
        event_infos = read_compound(fobj, "/event_info", MCEventInfo)

        mc_tracks = Dict{Int64}{Vector{MCTrack}}()
        for track in read_compound(fobj, "/mc_tracks", MCTrack)
            group_id = track.group_id + 1
            if !haskey(mc_tracks, group_id)
                mc_tracks[group_id] = Vector{MCTrack}()
            end
            push!(mc_tracks[group_id], track)
        end

        n_events = length(event_infos)
        new(filename, detx, fobj, calib, event_infos, mc_tracks, n_events)
    end
end

Base.show(io::IO, e::MCEventReader) = begin
    print(io, join(["MCEventReader: $(e.filename)",
          "       detx: $(e.detx)",
          "     events: $(length(e))"], "\n"))
end

Base.length(e::MCEventReader) = e._length
Base.firstindex(E::MCEventReader) = 0
Base.lastindex(E::MCEventReader) = length(E)

mutable struct MCEvent
    hits::Vector{Hit}
    mc_tracks::Vector{MCTrack}
    info::MCEventInfo
    calib::Calibration
end

function Base.iterate(iter::MCEventReader)
    group_id = 0
    (iter[group_id], group_id)
end

function Base.iterate(iter::MCEventReader, state)
    group_id = state + 1
    if group_id >= length(iter)
        return nothing
    end
    return (iter[group_id], group_id)
end

function Base.getindex(E::MCEventReader, i::Int)
    0 <= i < length(E) || throw(BoundsError(E, i))
    hits = read_hits(E._fobj, i)
    mc_tracks = E._mc_tracks[i+1]
    event_info = E._event_infos[i+1]
    MCEvent(hits, mc_tracks, event_info, E._calib)
end


text:
#
# BlockSparseTensor (Tensor using BlockSparse storage)
#

const BlockSparseTensor{ElT,N,StoreT,IndsT} =
  Tensor{ElT,N,StoreT,IndsT} where {StoreT<:BlockSparse}

nonzeros(T::Tensor) = data(T)

# Special version for BlockSparseTensor
# Generic version doesn't work since BlockSparse us parametrized by
# the Tensor order
function similartype(
  ::Type{<:Tensor{ElT,NT,<:BlockSparse{ElT,VecT},<:Any}}, ::Type{IndsR}
) where {NT,ElT,VecT,IndsR}
  NR = length(IndsR)
  return Tensor{ElT,NR,BlockSparse{ElT,VecT,NR},IndsR}
end

function similartype(
  ::Type{<:Tensor{ElT,NT,<:BlockSparse{ElT,VecT},<:Any}}, ::Type{IndsR}
) where {NT,ElT,VecT,IndsR<:NTuple{NR}} where {NR}
  return Tensor{ElT,NR,BlockSparse{ElT,VecT,NR},IndsR}
end

function BlockSparseTensor(
  ::Type{ElT}, ::UndefInitializer, boffs::BlockOffsets, inds
) where {ElT<:Number}
  nnz_tot = nnz(boffs, inds)
  storage = BlockSparse(ElT, undef, boffs, nnz_tot)
  return tensor(storage, inds)
end

function BlockSparseTensor(
  ::Type{ElT}, ::UndefInitializer, blocks::Vector{BlockT}, inds
) where {ElT<:Number,BlockT<:Union{Block,NTuple}}
  boffs, nnz = blockoffsets(blocks, inds)
  storage = BlockSparse(ElT, undef, boffs, nnz)
  return tensor(storage, inds)
end

"""
    BlockSparseTensor(::UndefInitializer, blocks, inds)

Construct a block sparse tensor with uninitialized memory
from indices and locations of non-zero blocks.
"""
function BlockSparseTensor(::UndefInitializer, blockoffsets, inds)
  return BlockSparseTensor(Float64, undef, blockoffsets, inds)
end

function BlockSparseTensor(
  ::Type{ElT}, blockoffsets::BlockOffsets{N}, inds
) where {ElT<:Number,N}
  nnz_tot = nnz(blockoffsets, inds)
  storage = BlockSparse(ElT, blockoffsets, nnz_tot)
  return tensor(storage, inds)
end

function BlockSparseTensor(blockoffsets::BlockOffsets, inds)
  return BlockSparseTensor(Float64, blockoffsets, inds)
end

"""
    BlockSparseTensor(inds)

Construct a block sparse tensor with no blocks.
"""
BlockSparseTensor(inds) = BlockSparseTensor(BlockOffsets{length(inds)}(), inds)

function BlockSparseTensor(::Type{ElT}, inds) where {ElT<:Number,N}
  return BlockSparseTensor(ElT, BlockOffsets{length(inds)}(), inds)
end

"""
    BlockSparseTensor(inds)

Construct a block sparse tensor with no blocks.
"""
function BlockSparseTensor(inds::Vararg{DimT,N}) where {DimT,N}
  return BlockSparseTensor(BlockOffsets{N}(), inds)
end

"""
    BlockSparseTensor(blocks::Vector{Block{N}}, inds)

Construct a block sparse tensor with the specified blocks.
Defaults to setting structurally non-zero blocks to zero.
"""
function BlockSparseTensor(blocks::Vector{BlockT}, inds) where {BlockT<:Union{Block,NTuple}}
  return BlockSparseTensor(Float64, blocks, inds)
end

function BlockSparseTensor(
  ::Type{ElT}, blocks::Vector{BlockT}, inds
) where {ElT<:Number,BlockT<:Union{Block,NTuple}}
  boffs, nnz = blockoffsets(blocks, inds)
  storage = BlockSparse(ElT, boffs, nnz)
  return tensor(storage, inds)
end

#complex(::Type{BlockSparseTensor{ElT,N,StoreT,IndsT}}) where {ElT<:Number,N,StoreT<:BlockSparse
#  = Tensor{ElT,N,StoreT,IndsT} where {StoreT<:BlockSparse}

function randn(
  ::Type{<:BlockSparseTensor{ElT,N}}, blocks::Vector{<:BlockT}, inds
) where {ElT,BlockT<:Union{Block{N},NTuple{N,<:Integer}}} where {N}
  boffs, nnz = blockoffsets(blocks, inds)
  storage = randn(BlockSparse{ElT}, boffs, nnz)
  return tensor(storage, inds)
end

# XXX: use the syntax:
# BlockSparseTensor(randn, ElT, blocks, inds)
function randomBlockSparseTensor(
  ::Type{ElT}, blocks::Vector{<:BlockT}, inds
) where {ElT,BlockT<:Union{Block{N},NTuple{N,<:Integer}}} where {N}
  return randn(BlockSparseTensor{ElT,N}, blocks, inds)
end

function randomBlockSparseTensor(blocks::Vector, inds)
  return randomBlockSparseTensor(Float64, blocks, inds)
end

"""
BlockSparseTensor(blocks::Vector{Block{N}},
                  inds::BlockDims...)

Construct a block sparse tensor with the specified blocks.
Defaults to setting structurally non-zero blocks to zero.
"""
function BlockSparseTensor(
  blocks::Vector{BlockT}, inds::Vararg{BlockDim,N}
) where {BlockT<:Union{Block{N},NTuple{N,<:Integer}}} where {N}
  return BlockSparseTensor(blocks, inds)
end

function similar(
  ::BlockSparseTensor{ElT,N}, blockoffsets::BlockOffsets{N}, inds
) where {ElT,N}
  return BlockSparseTensor(ElT, undef, blockoffsets, inds)
end

function similar(
  ::Type{<:BlockSparseTensor{ElT,N}}, blockoffsets::BlockOffsets{N}, inds
) where {ElT,N}
  return BlockSparseTensor(ElT, undef, blockoffsets, inds)
end

# This version of similar creates a tensor with no blocks
function similar(::Type{TensorT}, inds::Tuple) where {TensorT<:BlockSparseTensor}
  return similar(TensorT, BlockOffsets{ndims(TensorT)}(), inds)
end

function zeros(
  ::BlockSparseTensor{ElT,N}, blockoffsets::BlockOffsets{N}, inds
) where {ElT,N}
  return BlockSparseTensor(ElT, blockoffsets, inds)
end

function zeros(
  ::Type{<:BlockSparseTensor{ElT,N}}, blockoffsets::BlockOffsets{N}, inds
) where {ElT,N}
  return BlockSparseTensor(ElT, blockoffsets, inds)
end

function zeros(::BlockSparseTensor{ElT,N}, inds) where {ElT,N}
  return BlockSparseTensor(ElT, inds)
end

function zeros(::Type{<:BlockSparseTensor{ElT,N}}, inds) where {ElT,N}
  return BlockSparseTensor(ElT, inds)
end

# Basic functionality for AbstractArray interface
IndexStyle(::Type{<:BlockSparseTensor}) = IndexCartesian()

# Get the CartesianIndices for the range of indices
# of the specified
function blockindices(T::BlockSparseTensor{ElT,N}, block) where {ElT,N}
  return CartesianIndex(blockstart(T, block)):CartesianIndex(blockend(T, block))
end

"""
indexoffset(T::BlockSparseTensor,i::Int...) -> offset,block,blockoffset

Get the offset in the data of the specified
CartesianIndex. If it falls in a block that doesn't
exist, return nothing for the offset.
Also returns the block the index is found in and the offset
within the block.
"""
function indexoffset(T::BlockSparseTensor{ElT,N}, i::Vararg{Int,N}) where {ElT,N}
  index_within_block, block = blockindex(T, i...)
  block_dims = blockdims(T, block)
  offset_within_block = LinearIndices(block_dims)[CartesianIndex(index_within_block)]
  offset_of_block = offset(T, block)
  offset_of_i = isnothing(offset_of_block) ? nothing : offset_of_block + offset_within_block
  return offset_of_i, block, offset_within_block
end

# TODO: Add a checkbounds
# TODO: write this nicer in terms of blockview?
#       Could write:
#       block,index_within_block = blockindex(T,i...)
#       return blockview(T,block)[index_within_block]
@propagate_inbounds function getindex(
  T::BlockSparseTensor{ElT,N}, i::Vararg{Int,N}
) where {ElT,N}
  offset, _ = indexoffset(T, i...)
  isnothing(offset) && return zero(ElT)
  return storage(T)[offset]
end

@propagate_inbounds function getindex(T::BlockSparseTensor{ElT,0}) where {ElT}
  nnzblocks(T) == 0 && return zero(ElT)
  return storage(T)[]
end

# These may not be valid if the Tensor has no blocks
#@propagate_inbounds getindex(T::BlockSparseTensor{<:Number,1},ind::Int) = storage(T)[ind]

#@propagate_inbounds getindex(T::BlockSparseTensor{<:Number,0}) = storage(T)[1]

# Add the specified block to the BlockSparseTensor
# Insert it such that the blocks remain ordered.
# Defaults to adding zeros.
# Returns the offset of the new block added.
# XXX rename to insertblock!, no need to return offset
function insertblock_offset!(T::BlockSparseTensor{ElT,N}, newblock::Block{N}) where {ElT,N}
  newdim = blockdim(T, newblock)
  newoffset = nnz(T)
  insert!(blockoffsets(T), newblock, newoffset)
  # Insert new block into data
  splice!(data(storage(T)), (newoffset + 1):newoffset, zeros(ElT, newdim))
  return newoffset
end

function insertblock!(T::BlockSparseTensor{<:Number,N}, block::Block{N}) where {N}
  insertblock_offset!(T, block)
  return T
end

insertblock!(T::BlockSparseTensor, block) = insertblock!(T, Block(block))

# Insert missing diagonal blocks as zero blocks
function insert_diag_blocks!(T::AbstractArray{ElT}) where {ElT}
  for b in eachdiagblock(T)
    blockT = blockview(T, b)
    if isnothing(blockT)
      # Block was not found in the list, insert it
      insertblock!(T, b)
    end
  end
end

# TODO: Add a checkbounds
@propagate_inbounds function setindex!(
  T::BlockSparseTensor{ElT,N}, val, i::Vararg{Int,N}
) where {ElT,N}
  offset, block, offset_within_block = indexoffset(T, i...)
  if isnothing(offset)
    offset_of_block = insertblock_offset!(T, block)
    offset = offset_of_block + offset_within_block
  end
  storage(T)[offset] = val
  return T
end

hasblock(T::Tensor, block::Block) = isassigned(blockoffsets(T), block)

@propagate_inbounds function setindex!(
  T::BlockSparseTensor{ElT,N}, val, b::Block{N}
) where {ElT,N}
  if !hasblock(T, b)
    insertblock!(T, b)
  end
  Tb = T[b]
  Tb .= val
  return T
end

getindex(T::BlockSparseTensor, block::Block) = blockview(T, block)

to_indices(T::Tensor{<:Any,N}, b::Tuple{Block{N}}) where {N} = blockindices(T, b...)

function blockview(T::BlockSparseTensor, block::Block)
  return blockview(T, block, offset(T, block))
end

function blockview(T::BlockSparseTensor, block::Block, offset::Integer)
  return blockview(T, BlockOffset(block, offset))
end

# Case where the block isn't found, return nothing
function blockview(T::BlockSparseTensor, block::Block, ::Nothing)
  return nothing
end

blockview(T::BlockSparseTensor, block) = blockview(T, Block(block))

function blockview(T::BlockSparseTensor, bof::BlockOffset)
  blockT, offsetT = bof
  blockdimsT = blockdims(T, blockT)
  blockdimT = prod(blockdimsT)
  dataTslice = @view data(storage(T))[(offsetT + 1):(offsetT + blockdimT)]
  return tensor(Dense(dataTslice), blockdimsT)
end

view(T::BlockSparseTensor, b::Block) = blockview(T, b)

# convert to Dense
function dense(T::TensorT) where {TensorT<:BlockSparseTensor}
  R = zeros(dense(TensorT), inds(T))
  for block in keys(blockoffsets(T))
    # TODO: make sure this assignment is efficient
    R[blockindices(T, block)] = blockview(T, block)
  end
  return R
end

#
# Operations
#

# TODO: extend to case with different block structures
function +(T1::BlockSparseTensor{<:Number,N}, T2::BlockSparseTensor{<:Number,N}) where {N}
  inds(T1) ≠ inds(T2) &&
    error("Cannot add block sparse tensors with different block structure")
  R = copy(T1)
  return permutedims!!(R, T2, ntuple(identity, Val(N)), +)
end

function permutedims(T::BlockSparseTensor{<:Number,N}, perm::NTuple{N,Int}) where {N}
  blockoffsetsR, indsR = permutedims(blockoffsets(T), inds(T), perm)
  R = similar(T, blockoffsetsR, indsR)
  permutedims!(R, T, perm)
  return R
end

function _permute_combdims(combdims::NTuple{NC,Int}, perm::NTuple{NP,Int}) where {NC,NP}
  res = MVector{NC,Int}(undef)
  iperm = invperm(perm)
  for i in 1:NC
    res[i] = iperm[combdims[i]]
  end
  return Tuple(res)
end

#
# These are functions to help with combining and uncombining
#

# Note that combdims is expected to be contiguous and ordered
# smallest to largest
function combine_dims(blocks::Vector{Block{N}}, inds, combdims::NTuple{NC,Int}) where {N,NC}
  nblcks = nblocks(inds, combdims)
  blocks_comb = Vector{Block{N - NC + 1}}(undef, length(blocks))
  for (i, block) in enumerate(blocks)
    blocks_comb[i] = combine_dims(block, inds, combdims)
  end
  return blocks_comb
end

function combine_dims(block::Block, inds, combdims::NTuple{NC,Int}) where {NC}
  nblcks = nblocks(inds, combdims)
  slice = getindices(block, combdims)
  slice_comb = LinearIndices(nblcks)[slice...]
  block_comb = deleteat(block, combdims)
  block_comb = insertafter(block_comb, tuple(slice_comb), minimum(combdims) - 1)
  return block_comb
end

# In the dimension dim, permute the blocks
function perm_blocks(blocks::Blocks{N}, dim::Int, perm) where {N}
  blocks_perm = Blocks{N}(undef, nnzblocks(blocks))
  iperm = invperm(perm)
  for (i, block) in enumerate(blocks)
    blocks_perm[i] = setindex(block, iperm[block[dim]], dim)
  end
  return blocks_perm
end

# In the dimension dim, permute the block
function perm_block(block::Block, dim::Int, perm) where {N}
  iperm = invperm(perm)
  return setindex(block, iperm[block[dim]], dim)
end

# In the dimension dim, combine the specified blocks
function combine_blocks(blocks::Blocks, dim::Int, blockcomb::Vector{Int})
  blocks_comb = copy(blocks)
  nnz_comb = nnzblocks(blocks)
  for (i, block) in enumerate(blocks)
    dimval = block[dim]
    blocks_comb[i] = setindex(block, blockcomb[dimval], dim)
  end
  unique!(blocks_comb)
  return blocks_comb
end

function permutedims_combine_output(
  T::BlockSparseTensor{ElT,N},
  is,
  perm::NTuple{N,Int},
  combdims::NTuple{NC,Int},
  blockperm::Vector{Int},
  blockcomb::Vector{Int},
) where {ElT,N,NC}
  # Permute the indices
  indsT = inds(T)
  inds_perm = permute(indsT, perm)

  # Now that the indices are permuted, compute
  # which indices are now combined
  combdims_perm = sort(_permute_combdims(combdims, perm))

  # Permute the nonzero blocks (dimension-wise)
  blocks = nzblocks(T)
  blocks_perm = permutedims(blocks, perm)

  # Combine the nonzero blocks (dimension-wise)
  blocks_perm_comb = combine_dims(blocks_perm, inds_perm, combdims_perm)

  # Permute the blocks (within the newly combined dimension)
  comb_ind_loc = minimum(combdims_perm)
  blocks_perm_comb = perm_blocks(blocks_perm_comb, comb_ind_loc, blockperm)
  blocks_perm_comb = sort(blocks_perm_comb; lt=isblockless)

  # Combine the blocks (within the newly combined and permuted dimension)
  blocks_perm_comb = combine_blocks(blocks_perm_comb, comb_ind_loc, blockcomb)

  return BlockSparseTensor(ElT, blocks_perm_comb, is)
end

function permutedims_combine(
  T::BlockSparseTensor{ElT,N},
  is,
  perm::NTuple{N,Int},
  combdims::NTuple{NC,Int},
  blockperm::Vector{Int},
  blockcomb::Vector{Int},
) where {ElT,N,NC}
  R = permutedims_combine_output(T, is, perm, combdims, blockperm, blockcomb)

  # Permute the indices
  inds_perm = permute(inds(T), perm)

  # Now that the indices are permuted, compute
  # which indices are now combined
  combdims_perm = sort(_permute_combdims(combdims, perm))
  comb_ind_loc = minimum(combdims_perm)

  # Determine the new index before combining
  inds_to_combine = getindices(inds_perm, combdims_perm)
  ind_comb = ⊗(inds_to_combine...)
  ind_comb = permuteblocks(ind_comb, blockperm)

  for bof in pairs(blockoffsets(T))
    Tb = blockview(T, bof)
    b = nzblock(bof)
    b_perm = permute(b, perm)
    b_perm_comb = combine_dims(b_perm, inds_perm, combdims_perm)
    b_perm_comb = perm_block(b_perm_comb, comb_ind_loc, blockperm)
    b_in_combined_dim = b_perm_comb[comb_ind_loc]
    new_b_in_combined_dim = blockcomb[b_in_combined_dim]
    offset = 0
    pos_in_new_combined_block = 1
    while b_in_combined_dim - pos_in_new_combined_block > 0 &&
      blockcomb[b_in_combined_dim - pos_in_new_combined_block] == new_b_in_combined_dim
      offset += blockdim(ind_comb, b_in_combined_dim - pos_in_new_combined_block)
      pos_in_new_combined_block += 1
    end
    b_new = setindex(b_perm_comb, new_b_in_combined_dim, comb_ind_loc)

    Rb_total = blockview(R, b_new)
    dimsRb_tot = dims(Rb_total)
    subind = ntuple(
      i -> if i == comb_ind_loc
        range(1 + offset; stop=offset + blockdim(ind_comb, b_in_combined_dim))
      else
        range(1; stop=dimsRb_tot[i])
      end,
      N - NC + 1,
    )
    Rb = @view array(Rb_total)[subind...]

    # XXX Are these equivalent?
    #Tb_perm = permutedims(Tb,perm)
    #copyto!(Rb,Tb_perm)

    # XXX Not sure what this was for
    Rb = reshape(Rb, permute(dims(Tb), perm))
    Tbₐ = convert(Array, Tb)
    @strided Rb .= permutedims(Tbₐ, perm)
  end

  return R
end

# TODO: optimize by avoiding findfirst
function _number_uncombined(blockval::Integer, blockcomb::Vector)
  if blockval == blockcomb[end]
    return length(blockcomb) - findfirst(==(blockval), blockcomb) + 1
  end
  return findfirst(==(blockval + 1), blockcomb) - findfirst(==(blockval), blockcomb)
end

# TODO: optimize by avoiding findfirst
function _number_uncombined_shift(blockval::Integer, blockcomb::Vector)
  if blockval == 1
    return 0
  end
  ncomb_shift = 0
  for i in 1:(blockval - 1)
    ncomb_shift += findfirst(==(i + 1), blockcomb) - findfirst(==(i), blockcomb) - 1
  end
  return ncomb_shift
end

# Uncombine the blocks along the dimension dim
# according to the pattern in blockcomb (for example, blockcomb
# is [1,2,2,3] and dim = 2, so the blocks (1,2),(2,3) get
# split into (1,2),(1,3),(2,4))
function uncombine_blocks(blocks::Blocks{N}, dim::Int, blockcomb::Vector{Int}) where {N}
  blocks_uncomb = Blocks{N}()
  ncomb_tot = 0
  for i in 1:length(blocks)
    block = blocks[i]
    blockval = block[dim]
    ncomb = _number_uncombined(blockval, blockcomb)
    ncomb_shift = _number_uncombined_shift(blockval, blockcomb)
    push!(blocks_uncomb, setindex(block, blockval + ncomb_shift, dim))
    for j in 1:(ncomb - 1)
      push!(blocks_uncomb, setindex(block, blockval + ncomb_shift + j, dim))
    end
  end
  return blocks_uncomb
end

function uncombine_block(block::Block{N}, dim::Int, blockcomb::Vector{Int}) where {N}
  blocks_uncomb = Blocks{N}()
  ncomb_tot = 0
  blockval = block[dim]
  ncomb = _number_uncombined(blockval, blockcomb)
  ncomb_shift = _number_uncombined_shift(blockval, blockcomb)
  push!(blocks_uncomb, setindex(block, blockval + ncomb_shift, dim))
  for j in 1:(ncomb - 1)
    push!(blocks_uncomb, setindex(block, blockval + ncomb_shift + j, dim))
  end
  return blocks_uncomb
end

function uncombine_output(
  T::BlockSparseTensor{ElT,N},
  is,
  combdim::Int,
  blockperm::Vector{Int},
  blockcomb::Vector{Int},
) where {ElT<:Number,N}
  ind_uncomb_perm = ⊗(setdiff(is, inds(T))...)
  inds_uncomb_perm = insertat(inds(T), ind_uncomb_perm, combdim)
  # Uncombine the blocks of T
  blocks_uncomb = uncombine_blocks(nzblocks(T), combdim, blockcomb)
  blocks_uncomb_perm = perm_blocks(blocks_uncomb, combdim, invperm(blockperm))
  boffs_uncomb_perm, nnz_uncomb_perm = blockoffsets(blocks_uncomb_perm, inds_uncomb_perm)
  T_uncomb_perm = tensor(
    BlockSparse(ElT, boffs_uncomb_perm, nnz_uncomb_perm), inds_uncomb_perm
  )
  R = reshape(T_uncomb_perm, is)
  return R
end

function reshape(blockT::Block{NT}, indsT, indsR) where {NT}
  nblocksT = nblocks(indsT)
  nblocksR = nblocks(indsR)
  blockR = Tuple(
    CartesianIndices(nblocksR)[LinearIndices(nblocksT)[CartesianIndex(blockT)]]
  )
  return blockR
end

function uncombine(
  T::BlockSparseTensor{<:Number,NT},
  is,
  combdim::Int,
  blockperm::Vector{Int},
  blockcomb::Vector{Int},
) where {NT}
  NR = length(is)
  R = uncombine_output(T, is, combdim, blockperm, blockcomb)
  invblockperm = invperm(blockperm)

  # This is needed for reshaping the block
  # It is already calculated in uncombine_output, use it from there
  ind_uncomb_perm = ⊗(setdiff(is, inds(T))...)
  ind_uncomb = permuteblocks(ind_uncomb_perm, blockperm)
  # Same as inds(T) but with the blocks uncombined
  inds_uncomb = insertat(inds(T), ind_uncomb, combdim)
  inds_uncomb_perm = insertat(inds(T), ind_uncomb_perm, combdim)
  for bof in pairs(blockoffsets(T))
    b = nzblock(bof)
    Tb_tot = blockview(T, bof)
    dimsTb_tot = dims(Tb_tot)
    bs_uncomb = uncombine_block(b, combdim, blockcomb)
    offset = 0
    for i in 1:length(bs_uncomb)
      b_uncomb = bs_uncomb[i]
      b_uncomb_perm = perm_block(b_uncomb, combdim, invblockperm)
      b_uncomb_perm_reshape = reshape(b_uncomb_perm, inds_uncomb_perm, is)
      Rb = blockview(R, b_uncomb_perm_reshape)
      b_uncomb_in_combined_dim = b_uncomb_perm[combdim]
      start = offset + 1
      stop = offset + blockdim(ind_uncomb_perm, b_uncomb_in_combined_dim)
      subind = ntuple(
        i -> i == combdim ? range(start; stop=stop) : range(1; stop=dimsTb_tot[i]), NT
      )
      offset = stop
      Tb = @view array(Tb_tot)[subind...]

      # Alternative (but maybe slower):
      #copyto!(Rb,Tb)

      if length(Tb) == 1
        Rb[1] = Tb[1]
      else
        # XXX: this used to be:
        # Rbₐᵣ = ReshapedArray(parent(Rbₐ), size(Tb), ())
        # however that doesn't work with subarrays
        Rbₐ = convert(Array, Rb)
        Rbₐᵣ = ReshapedArray(Rbₐ, size(Tb), ())
        @strided Rbₐᵣ .= Tb
      end
    end
  end
  return R
end

function copyto!(R::BlockSparseTensor, T::BlockSparseTensor)
  for bof in pairs(blockoffsets(T))
    copyto!(blockview(R, nzblock(bof)), blockview(T, bof))
  end
  return R
end

# TODO: handle case where:
# f(zero(ElR),zero(ElT)) != promote_type(ElR,ElT)
function permutedims!!(
  R::BlockSparseTensor{ElR,N},
  T::BlockSparseTensor{ElT,N},
  perm::NTuple{N,Int},
  f::Function=(r, t) -> t,
) where {ElR,ElT,N}
  RR = convert(promote_type(typeof(R), typeof(T)), R)
  #@timeit_debug timer "block sparse permutedims!!" begin
  bofsRR = blockoffsets(RR)
  bofsT = blockoffsets(T)

  # Determine if bofsRR has been copied
  copy_bofsRR = false

  new_nnz = nnz(RR)
  for (blockT, offsetT) in pairs(bofsT)
    blockTperm = permute(blockT, perm)
    if !isassigned(bofsRR, blockTperm)
      if !copy_bofsRR
        bofsRR = deepcopy(bofsRR)
        copy_bofsRR = true
      end
      insert!(bofsRR, blockTperm, new_nnz)
      new_nnz += blockdim(T, blockT)
    end
  end

  ## RR = BlockSparseTensor(promote_type(ElR,ElT), undef,
  ##                        bofsRR, inds(R))
  ## # Directly copy the data since it is the same blocks
  ## # and offsets
  ## copyto!(data(RR), data(R))

  if new_nnz > nnz(RR)
    dataRR = append!(data(RR), zeros(new_nnz - nnz(RR)))
    RR = Tensor(BlockSparse(dataRR, bofsRR), inds(RR))
  end

  permutedims!(RR, T, perm, f)
  return RR
  #end
end

# <fermions>
scale_blocks!(T, compute_fac::Function=(b) -> 1) = T

# <fermions>
function scale_blocks!(
  T::BlockSparseTensor{<:Number,N}, compute_fac::Function=(b) -> 1
) where {N}
  for blockT in keys(blockoffsets(T))
    fac = compute_fac(blockT)
    if fac != 1
      Tblock = blockview(T, blockT)
      scale!(Tblock, fac)
    end
  end
  return T
end

# <fermions>
permfactor(perm, block, inds) = 1.0

# Version where it is known that R has the same blocks
# as T
function permutedims!(
  R::BlockSparseTensor{<:Number,N},
  T::BlockSparseTensor{<:Number,N},
  perm::NTuple{N,Int},
  f::Function=(r, t) -> t,
) where {N}
  for blockT in keys(blockoffsets(T))
    # Loop over non-zero blocks of T/R
    Tblock = blockview(T, blockT)
    Rblock = blockview(R, permute(blockT, perm))

    # <fermions>
    pfac = permfactor(perm, blockT, inds(T))
    fac_f = (r, t) -> f(r, pfac * t)

    permutedims!(Rblock, Tblock, perm, fac_f)
  end
  return R
end

#
# Contraction
#

# TODO: complete this function: determine the output blocks from the input blocks
# Also, save the contraction list (which block-offsets contract with which),
# may not be generic with other contraction functions!
function contraction_output(
  T1::TensorT1, T2::TensorT2, indsR::IndsR
) where {TensorT1<:BlockSparseTensor,TensorT2<:BlockSparseTensor,IndsR}
  TensorR = contraction_output_type(TensorT1, TensorT2, IndsR)
  return similar(TensorR, blockoffsetsR, indsR)
end

"""
find_matching_positions(t1,t2) -> t1_to_t2

In a tuple of length(t1), store the positions in t2
where the element of t1 is found. Otherwise, store 0
to indicate that the element of t1 is not in t2.

For example, for all t1[pos1] == t2[pos2], t1_to_t2[pos1] == pos2,
otherwise t1_to_t2[pos1] == 0.
"""
function find_matching_positions(t1, t2)
  t1_to_t2 = @MVector zeros(Int, length(t1))
  for pos1 in 1:length(t1)
    for pos2 in 1:length(t2)
      if t1[pos1] == t2[pos2]
        t1_to_t2[pos1] = pos2
      end
    end
  end
  return Tuple(t1_to_t2)
end

function contract_labels(labels1, labels2, labelsR)
  labels1_to_labels2 = find_matching_positions(labels1, labels2)
  labels1_to_labelsR = find_matching_positions(labels1, labelsR)
  labels2_to_labelsR = find_matching_positions(labels2, labelsR)
  return labels1_to_labels2, labels1_to_labelsR, labels2_to_labelsR
end

function are_blocks_contracted(
  block1::Block{N1}, block2::Block{N2}, labels1_to_labels2::NTuple{N1,Int}
) where {N1,N2}
  t1 = Tuple(block1)
  t2 = Tuple(block2)
  for i1 in 1:N1
    i2 = @inbounds labels1_to_labels2[i1]
    if i2 > 0
      # This dimension is contracted
      if @inbounds t1[i1] != @inbounds t2[i2]
        return false
      end
    end
  end
  return true
end

function contract_blocks(
  block1::Block{N1}, labels1_to_labelsR, block2::Block{N2}, labels2_to_labelsR, ::Val{NR}
) where {N1,N2,NR}
  blockR = ntuple(_ -> UInt(0), Val(NR))
  t1 = Tuple(block1)
  t2 = Tuple(block2)
  for i1 in 1:N1
    iR = @inbounds labels1_to_labelsR[i1]
    if iR > 0
      blockR = @inbounds setindex(blockR, t1[i1], iR)
    end
  end
  for i2 in 1:N2
    iR = @inbounds labels2_to_labelsR[i2]
    if iR > 0
      blockR = @inbounds setindex(blockR, t2[i2], iR)
    end
  end
  return Block{NR}(blockR)
end

function _contract_blockoffsets(
  boffs1::BlockOffsets{N1},
  inds1,
  labels1,
  boffs2::BlockOffsets{N2},
  inds2,
  labels2,
  indsR,
  labelsR,
) where {N1,N2}
  NR = length(labelsR)
  ValNR = ValLength(labelsR)
  labels1_to_labels2, labels1_to_labelsR, labels2_to_labelsR = contract_labels(
    labels1, labels2, labelsR
  )
  blockoffsetsR = BlockOffsets{NR}()
  nnzR = 0
  contraction_plan = Tuple{Block{N1},Block{N2},Block{NR}}[]
  # Reserve some capacity
  # In theory the maximum is length(boffs1) * length(boffs2)
  # but in practice that is too much
  sizehint!(contraction_plan, max(length(boffs1), length(boffs2)))
  for block1 in keys(boffs1)
    for block2 in keys(boffs2)
      if are_blocks_contracted(block1, block2, labels1_to_labels2)
        blockR = contract_blocks(
          block1, labels1_to_labelsR, block2, labels2_to_labelsR, ValNR
        )
        push!(contraction_plan, (block1, block2, blockR))
        if !isassigned(blockoffsetsR, blockR)
          insert!(blockoffsetsR, blockR, nnzR)
          nnzR += blockdim(indsR, blockR)
        end
      end
    end
  end
  return blockoffsetsR, contraction_plan
end

function _threaded_contract_blockoffsets(
  boffs1::BlockOffsets{N1},
  inds1,
  labels1,
  boffs2::BlockOffsets{N2},
  inds2,
  labels2,
  indsR,
  labelsR,
) where {N1,N2}
  NR = length(labelsR)
  ValNR = ValLength(labelsR)
  labels1_to_labels2, labels1_to_labelsR, labels2_to_labelsR = contract_labels(
    labels1, labels2, labelsR
  )
  contraction_plans = [Tuple{Block{N1},Block{N2},Block{NR}}[] for _ in 1:nthreads()]

  #
  # Reserve some capacity
  # In theory the maximum is length(boffs1) * length(boffs2)
  # but in practice that is too much
  #for contraction_plan in contraction_plans
  #  sizehint!(contraction_plan, max(length(boffs1), length(boffs2)))
  #end
  #

  blocks1 = keys(boffs1)
  blocks2 = keys(boffs2)
  if length(blocks1) > length(blocks2)
    @sync for blocks1_partition in
              Iterators.partition(blocks1, max(1, length(blocks1) ÷ nthreads()))
      @spawn for block1 in blocks1_partition
        for block2 in blocks2
          if are_blocks_contracted(block1, block2, labels1_to_labels2)
            blockR = contract_blocks(
              block1, labels1_to_labelsR, block2, labels2_to_labelsR, ValNR
            )
            push!(contraction_plans[threadid()], (block1, block2, blockR))
          end
        end
      end
    end
  else
    @sync for blocks2_partition in
              Iterators.partition(blocks2, max(1, length(blocks2) ÷ nthreads()))
      @spawn for block2 in blocks2_partition
        for block1 in blocks1
          if are_blocks_contracted(block1, block2, labels1_to_labels2)
            blockR = contract_blocks(
              block1, labels1_to_labelsR, block2, labels2_to_labelsR, ValNR
            )
            push!(contraction_plans[threadid()], (block1, block2, blockR))
          end
        end
      end
    end
  end

  contraction_plan = vcat(contraction_plans...)
  blockoffsetsR = BlockOffsets{NR}()
  nnzR = 0
  for (_, _, blockR) in contraction_plan
    if !isassigned(blockoffsetsR, blockR)
      insert!(blockoffsetsR, blockR, nnzR)
      nnzR += blockdim(indsR, blockR)
    end
  end

  return blockoffsetsR, contraction_plan
end

function contract_blockoffsets(args...)
  if using_threaded_blocksparse() && nthreads() > 1
    return _threaded_contract_blockoffsets(args...)
  end
  return _contract_blockoffsets(args...)
end

function contraction_output(
  T1::TensorT1, labelsT1, T2::TensorT2, labelsT2, labelsR
) where {TensorT1<:BlockSparseTensor,TensorT2<:BlockSparseTensor}
  indsR = contract_inds(inds(T1), labelsT1, inds(T2), labelsT2, labelsR)
  TensorR = contraction_output_type(TensorT1, TensorT2, typeof(indsR))
  blockoffsetsR, contraction_plan = contract_blockoffsets(
    blockoffsets(T1),
    inds(T1),
    labelsT1,
    blockoffsets(T2),
    inds(T2),
    labelsT2,
    indsR,
    labelsR,
  )
  R = similar(TensorR, blockoffsetsR, indsR)
  return R, contraction_plan
end

function contract(
  T1::BlockSparseTensor{<:Any,N1},
  labelsT1,
  T2::BlockSparseTensor{<:Any,N2},
  labelsT2,
  labelsR=contract_labels(labelsT1, labelsT2),
) where {N1,N2}
  #@timeit_debug timer "Block sparse contract" begin
  R, contraction_plan = contraction_output(T1, labelsT1, T2, labelsT2, labelsR)
  R = contract!(R, labelsR, T1, labelsT1, T2, labelsT2, contraction_plan)
  return R
  #end
end

# <fermions>
function compute_alpha(
  ElR, labelsR, blockR, indsR, labelsT1, blockT1, indsT1, labelsT2, blockT2, indsT2
)
  return one(ElR)
end

# XXX: this is not thread safe, divide into groups of
# contractions that contract into the same block
function _threaded_contract!(
  R::BlockSparseTensor{ElR,NR},
  labelsR,
  T1::BlockSparseTensor{ElT1,N1},
  labelsT1,
  T2::BlockSparseTensor{ElT2,N2},
  labelsT2,
  contraction_plan,
) where {ElR,ElT1,ElT2,N1,N2,NR}
  # Sort the contraction plan by the output blocks
  # This is to help determine which output blocks are the result
  # of multiple contractions
  sort!(contraction_plan; by=last)

  # Ranges of contractions to the same block
  repeats = Vector{UnitRange{Int}}(undef, nnzblocks(R))
  ncontracted = 1
  posR = last(contraction_plan[1])
  posR_unique = posR
  for n in 1:(nnzblocks(R) - 1)
    start = ncontracted
    while posR == posR_unique
      ncontracted += 1
      posR = last(contraction_plan[ncontracted])
    end
    posR_unique = posR
    repeats[n] = start:(ncontracted - 1)
  end
  repeats[end] = ncontracted:length(contraction_plan)

  contraction_plan_blocks = Vector{Tuple{Tensor,Tensor,Tensor}}(
    undef, length(contraction_plan)
  )
  for ncontracted in 1:length(contraction_plan)
    block1, block2, blockR = contraction_plan[ncontracted]
    T1block = T1[block1]
    T2block = T2[block2]
    Rblock = R[blockR]
    contraction_plan_blocks[ncontracted] = (T1block, T2block, Rblock)
  end

  α = one(ElR)
  @sync for repeats_partition in
            Iterators.partition(repeats, max(1, length(repeats) ÷ nthreads()))
    @spawn for ncontracted_range in repeats_partition
      # Overwrite the block since it hasn't been written to
      # R .= α .* (T1 * T2)
      β = zero(ElR)
      for ncontracted in ncontracted_range
        blockT1, blockT2, blockR = contraction_plan_blocks[ncontracted]
        # R .= α .* (T1 * T2) .+ β .* R

        # <fermions>:
        α = compute_alpha(
          ElR,
          labelsR,
          blockR,
          inds(R),
          labelsT1,
          block1,
          inds(T1),
          labelsT2,
          block2,
          inds(T2),
        )

        contract!(blockR, labelsR, blockT1, labelsT1, blockT2, labelsT2, α, β)
        # Now keep adding to the block, since it has
        # been written to
        # R .= α .* (T1 * T2) .+ R
        β = one(ElR)
      end
    end
  end
  return R
end

function contract!(
  R::BlockSparseTensor{ElR,NR},
  labelsR,
  T1::BlockSparseTensor{ElT1,N1},
  labelsT1,
  T2::BlockSparseTensor{ElT2,N2},
  labelsT2,
  contraction_plan,
) where {ElR,ElT1,ElT2,N1,N2,NR}
  if isempty(contraction_plan)
    return R
  end
  if using_threaded_blocksparse() && nthreads() > 1
    _threaded_contract!(R, labelsR, T1, labelsT1, T2, labelsT2, contraction_plan)
    return R
  end
  already_written_to = Dict{Block{NR},Bool}()
  # In R .= α .* (T1 * T2) .+ β .* R
  for (block1, block2, blockR) in contraction_plan

    #<fermions>
    α = compute_alpha(
      ElR, labelsR, blockR, inds(R), labelsT1, block1, inds(T1), labelsT2, block2, inds(T2)
    )

    T1block = T1[block1]
    T2block = T2[block2]
    Rblock = R[blockR]
    β = one(ElR)
    if !haskey(already_written_to, blockR)
      already_written_to[blockR] = true
      # Overwrite the block of R
      β = zero(ElR)
    end
    contract!(Rblock, labelsR, T1block, labelsT1, T2block, labelsT2, α, β)
  end
  return R
end

const IntTuple = NTuple{N,Int} where {N}
const IntOrIntTuple = Union{Int,IntTuple}

function permute_combine(inds::IndsT, pos::Vararg{IntOrIntTuple,N}) where {IndsT,N}
  IndT = eltype(IndsT)
  # Using SizedVector since setindex! doesn't
  # work for MVector when eltype not isbitstype
  newinds = SizedVector{N,IndT}(undef)
  for i in 1:N
    pos_i = pos[i]
    newind_i = inds[pos_i[1]]
    for p in 2:length(pos_i)
      newind_i = newind_i ⊗ inds[pos_i[p]]
    end
    newinds[i] = newind_i
  end
  IndsR = similartype(IndsT, Val{N})
  indsR = IndsR(Tuple(newinds))
  return indsR
end

"""
Indices are combined according to the grouping of the input,
for example (1,2),3 will combine the first two indices.
"""
function combine(inds::IndsT, com::Vararg{IntOrIntTuple,N}) where {IndsT,N}
  IndT = eltype(IndsT)
  # Using SizedVector since setindex! doesn't
  # work for MVector when eltype not isbitstype
  newinds = SizedVector{N,IndT}(undef)
  i_orig = 1
  for i in 1:N
    newind_i = inds[i_orig]
    i_orig += 1
    for p in 2:length(com[i])
      newind_i = newind_i ⊗ inds[i_orig]
      i_orig += 1
    end
    newinds[i] = newind_i
  end
  IndsR = similartype(IndsT, Val{N})
  indsR = IndsR(Tuple(newinds))
  return indsR
end

function permute_combine(
  boffs::BlockOffsets, inds::IndsT, pos::Vararg{IntOrIntTuple,N}
) where {IndsT,N}
  perm = flatten(pos...)
  boffsp, indsp = permutedims(boffs, inds, perm)
  indsR = combine(indsp, pos...)
  boffsR = reshape(boffsp, indsp, indsR)
  return boffsR, indsR
end

function reshape(boffsT::BlockOffsets{NT}, indsT, indsR) where {NT}
  NR = length(indsR)
  boffsR = BlockOffsets{NR}()
  nblocksT = nblocks(indsT)
  nblocksR = nblocks(indsR)
  for (blockT, offsetT) in pairs(boffsT)
    blockR = Block(
      CartesianIndices(nblocksR)[LinearIndices(nblocksT)[CartesianIndex(blockT)]]
    )
    insert!(boffsR, blockR, offsetT)
  end
  return boffsR
end

function reshape(boffsT::BlockOffsets{NT}, blocksR::Vector{Block{NR}}) where {NR,NT}
  boffsR = BlockOffsets{NR}()
  # TODO: check blocksR is ordered and are properly reshaped
  # versions of the blocks of boffsT
  for (i, (blockT, offsetT)) in enumerate(boffsT)
    blockR = blocksR[i]
    boffsR[blockR] = offsetT
  end
  return boffsR
end

reshape(T::BlockSparse, boffsR::BlockOffsets) = BlockSparse(data(T), boffsR)

function reshape(T::BlockSparseTensor, boffsR::BlockOffsets, indsR)
  storeR = reshape(storage(T), boffsR)
  return tensor(storeR, indsR)
end

function reshape(T::BlockSparseTensor, indsR)
  # TODO: add some checks that the block dimensions
  # are consistent (e.g. nnzblocks(T) == nnzblocks(R), etc.)
  boffsR = reshape(blockoffsets(T), inds(T), indsR)
  R = reshape(T, boffsR, indsR)
  return R
end

function permute_combine(
  T::BlockSparseTensor{ElT,NT,IndsT}, pos::Vararg{IntOrIntTuple,NR}
) where {ElT,NT,IndsT,NR}
  boffsR, indsR = permute_combine(blockoffsets(T), inds(T), pos...)

  perm = flatten(pos...)

  length(perm) ≠ NT && error("Index positions must add up to order of Tensor ($NT)")
  isperm(perm) || error("Index positions must be a permutation")

  if !is_trivial_permutation(perm)
    Tp = permutedims(T, perm)
  else
    Tp = copy(T)
  end
  NR == NT && return Tp
  R = reshape(Tp, boffsR, indsR)
  return R
end

#
# Print block sparse tensors
#

#function summary(io::IO,
#                 T::BlockSparseTensor{ElT,N}) where {ElT,N}
#  println(io,Base.dims2string(dims(T))," ",typeof(T))
#  for (dim,ind) in enumerate(inds(T))
#    println(io,"Dim $dim: ",ind)
#  end
#  println(io,"Number of nonzero blocks: ",nnzblocks(T))
#end

#function summary(io::IO,
#                 T::BlockSparseTensor{ElT,N}) where {ElT,N}
#  println(io,typeof(T))
#  println(io,Base.dims2string(dims(T))," ",typeof(T))
#  for (dim,ind) in enumerate(inds(T))
#    println(io,"Dim $dim: ",ind)
#  end
#  println("Number of nonzero blocks: ",nnzblocks(T))
#end

function _range2string(rangestart::NTuple{N,Int}, rangeend::NTuple{N,Int}) where {N}
  s = ""
  for n in 1:N
    s = string(s, rangestart[n], ":", rangeend[n])
    if n < N
      s = string(s, ", ")
    end
  end
  return s
end

function show(io::IO, mime::MIME"text/plain", T::BlockSparseTensor)
  summary(io, T)
  for (n, block) in enumerate(keys(blockoffsets(T)))
    blockdimsT = blockdims(T, block)
    println(io, block)
    println(io, " [", _range2string(blockstart(T, block), blockend(T, block)), "]")
    print_tensor(io, blockview(T, block))
    n < nnzblocks(T) && print(io, "\n\n")
  end
end

show(io::IO, T::BlockSparseTensor) = show(io, MIME("text/plain"), T)


text:
using BenchmarkTools

time = zeros(21)
memory = zeros(21)

include("atlas.jl")
t = @benchmarkable simulate!($mech)
r = BenchmarkTools.minimum(run(t, samples = 1))
time[1] = r.time
memory[1] = r.memory

include("chain_in_chain.jl")
t = @benchmarkable simulate!($mech)
r = BenchmarkTools.minimum(run(t, samples = 1))
time[2] = r.time
memory[2] = r.memory

include("dice_nofriction.jl")
t = @benchmarkable simulate!($mech)
r = BenchmarkTools.minimum(run(t, samples = 1))
time[3] = r.time
memory[3] = r.memory

include("dice_tiltedplane.jl")
t = @benchmarkable simulate!($mech)
r = BenchmarkTools.minimum(run(t, samples = 1))
time[4] = r.time
memory[4] = r.memory

include("dice.jl")
t = @benchmarkable simulate!($mech)
r = BenchmarkTools.minimum(run(t, samples = 1))
time[5] = r.time
memory[5] = r.memory

include("disconnected_bodies.jl")
t = @benchmarkable simulate!($mech)
r = BenchmarkTools.minimum(run(t, samples = 1))
time[6] = r.time
memory[6] = r.memory

include("doublependulum_3d.jl")
t = @benchmarkable simulate!($mech)
r = BenchmarkTools.minimum(run(t, samples = 1))
time[7] = r.time
memory[7] = r.memory

include("inverted_pyramid_plane.jl")
t = @benchmarkable simulate!($mech)
r = BenchmarkTools.minimum(run(t, samples = 1))
time[8] = r.time
memory[8] = r.memory

include("joint_force.jl")
t = @benchmarkable simulate!($mech)
r = BenchmarkTools.minimum(run(t, samples = 1))
time[9] = r.time
memory[9] = r.memory

include("joint_torque.jl")
t = @benchmarkable simulate!($mech)
r = BenchmarkTools.minimum(run(t, samples = 1))
time[10] = r.time
memory[10] = r.memory

include("n_fourbars.jl")
t = @benchmarkable simulate!($mech)
r = BenchmarkTools.minimum(run(t, samples = 1))
time[11] = r.time
memory[11] = r.memory

include("n_pendulum.jl")
t = @benchmarkable simulate!($mech)
r = BenchmarkTools.minimum(run(t, samples = 1))
time[12] = r.time
memory[12] = r.memory

include("pendulum_controlled.jl")
t = @benchmarkable simulate!($mech)
r = BenchmarkTools.minimum(run(t, samples = 1))
time[13] = r.time
memory[13] = r.memory

include("pendulum_forced.jl")
t = @benchmarkable simulate!($mech)
r = BenchmarkTools.minimum(run(t, samples = 1))
time[14] = r.time
memory[14] = r.memory

include("pendulum.jl")
t = @benchmarkable simulate!($mech)
r = BenchmarkTools.minimum(run(t, samples = 1))
time[15] = r.time
memory[15] = r.memory

include("planar_example.jl")
t = @benchmarkable simulate!($mech)
r = BenchmarkTools.minimum(run(t, samples = 1))
time[16] = r.time
memory[16] = r.memory

include("scissor_lift.jl")
t = @benchmarkable simulate!($mech)
r = BenchmarkTools.minimum(run(t, samples = 1))
time[17] = r.time
memory[17] = r.memory

include("slider_crank.jl")
t = @benchmarkable simulate!($mech)
r = BenchmarkTools.minimum(run(t, samples = 1))
time[18] = r.time
memory[18] = r.memory

include("slider_crank3d.jl")
t = @benchmarkable simulate!($mech)
r = BenchmarkTools.minimum(run(t, samples = 1))
time[19] = r.time
memory[19] = r.memory

include("urdf_doublependulum.jl")
t = @benchmarkable simulate!($mech)
r = BenchmarkTools.minimum(run(t, samples = 1))
time[20] = r.time
memory[20] = r.memory

# include("wheel.jl")
# t = @benchmarkable simulate!($mech)
# r = BenchmarkTools.minimum(run(t,samples=1))
# time[21] = r.time
# memory[21] = r.memory

time /= 1e6


text:
function run_pandamodels_vd(json_path)
    pm = load_pm_from_json(json_path)
    active_powermodels_silence!(pm)
    pm = check_powermodels_data!(pm)
    model = get_model(pm["pm_model"])
    solver = get_solver(pm)

    result = _run_vd(
        pm,
        model,
        solver,
        setting = Dict("output" => Dict("branch_flows" => true)),
        ext = extract_params!(pm),
    )
    return result
end


text:
using Documenter
using CVChannel

makedocs(;
    modules=[CVChannel],
    authors="Brian Doolittle <brian.d.doolittle@gmail.com> and Ian George <igeorge3@illinois.edu>",
    repo="https://github.com/ChitambarLab/CVChannel.jl/blob/{commit}{path}#L{line}",
    sitename="CVChannel.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://ChitambarLab.github.io/CVChannel.jl",
        assets=String["assets/custom.css"],
    ),
    pages=[
        "Home" => "index.md",
        "CV Background" => "background.md",
        "CV Optimizations" => "communication_value.md",
        "CV Multiplicativity" => "cv_multiplicativity.md",
        "Utilities" => [
            "Channels, States, and Operations" => "channel_states.md",
            "Optimization Backends" => "optimizer_interface.md",
        ],
    ],
)

deploydocs(
    repo="github.com/ChitambarLab/CVChannel.jl",
    devbranch = "main",
)

max_forks_repo_name: ChitambarLab/cv-channel


text:
struct Range
    state::Int
    stream::AbstractStream
end

function Base.iterate(range::Range, state = 1)
    if state > range.state
        return nothing
    end

    data = listen(range.stream)

    return isempty(data) ? nothing : (data, state+1)
end

function range(value::Int, stream::AbstractStream)
    return Range(value, stream)
end


text:
##  Conclusion: make my life easier and just define a GraphNode and a TreeNode

print("Hi!")
module tmp
    using GeometryTypes
    struct Node2
        id::Int
        iPrev::Int #parent
        state::Nullable{Point2f0}
    end

    #Node2(id::Int, iPrev::Int) = Node2(id, iPrev, Point2f0())
    Node2(id::Int, iPrev::Int) = Node2(id, iPrev, Point(999,999))
    # Node2(id::Int, iPrev::Int) = Node2(id, iPrev, Point())
#
	# Fails with;
	# WARNING: replacing module tmp
	# ERROR: LoadError: MethodError: Cannot `convert` an object of type Tuple{} to an object of type Float32
	# This may have arisen from a call to the constructor Float32(...),
	# since type constructors fall back to convert methods.

end

    x = tmp.Node2(1,2,Point(2,3))
    y = tmp.Node2(1,2)


text:
using Dates
using Distributed
using InteractiveUtils
using JSON
using Test

using Discord
using Discord:
    EVENTS_FIRED,
    Conn,
    Empty,
    Handler,
    JobQueue,
    Limiter,
    Response,
    Snowflake,
    allhandlers,
    alwaystrue,
    datetime,
    dec!,
    enqueue!,
    get_channel_message,
    handler,
    handlers,
    hasdefault,
    increment,
    insert_or_update!,
    iscollecting,
    isexpired,
    logkws,
    mock,
    parse_endpoint,
    predicate,
    process_id,
    readjson,
    results,
    should_put,
    snowflake,
    snowflake2datetime,
    split_message,
    trywritejson,
    validate_fetch,
    worker_id,
    wrapfn!,
    writejson,
    @boilerplate,
    @constructors,
    @lower,
    @merge

c = Client("token")

@testset "Discord.jl" begin
    include("client.jl")
    include("parsing.jl")
    include("json.jl")
    include("helpers.jl")
    include("rest.jl")
    include("ratelimits.jl")
    include("boilerplate.jl")
    include("handlers.jl")
    include("commands.jl")
    include("state.jl")
end


text:
"""
    adaptiveHRVFilter(signal::Vector{<:Real}; removeOutliers::Bool = true, replaceNonNormal::Bool = true, adaptiveControllingCoef::Real = 0.05, proportionalLimit::Real = 10/100, controllingFilterCoef::Real = 3, controllingBasicVariability::Real = 20, physiological_values::Tuple{Real, Real} = (200, 2000))

An adaptive filter for HRV data. Filters out nonphysiological beats in the RR series with the ability to replace them.

Based on: Wessel, N., Voss, A., Malberg, H., Ziehmann, Ch., Voss, H. U., Schirdewan, A., Meyerfeldt, U.,Kurths, J.: Nonlinear analysis of complex phenomena in cardiological data, Herzschr. Elektrophys., 11(3), 2000, 159-173, doi:10.1007/s003990070035.

# Input

* signal: HRV in ms

# Keyword Arguments

* removeOutliers: option if nonphysiological outliers shall be removed, defaults to true
* replaceNonNormal: if true, non normal HRV values are replaced with a random value sampled from an adaptive interval to not falsely decrease variability.
* physiological_values: Definition of the physiological value range, scheme: (min, max), defaults to (200, 2000)
* adaptiveControllingCoef: Controlling coefficient for the adaptive mean, ∈ [0,1] and defaults to 0.05
* proportionalLimit: Proportional limit for the exclusion rule, defaults to 10%
* controllingFilterCoef: Filter coefficient in the final controlling procedure , defaults to 3
* controllingBasicVariability: Basic variability in the controlling procedure to reduce filtering errors for time series with low variability, defaults to 20 ms

# Return
Return value depends on keyword args

```julia
julia> adaptiveHRVFilter(signal)
```
"""
function adaptiveHRVFilter(signal::Vector{<:Real}; removeOutliers::Bool = true, replaceNonNormal::Bool = true, adaptiveControllingCoef::Real = 0.05, proportionalLimit::Real = 10/100, controllingFilterCoef::Real = 3, controllingBasicVariability::Real = 20, physiological_values::Tuple{Real, Real} = (200, 2000))
    # some input checking
    0 < adaptiveControllingCoef < 1 || throw(DomainError("adaptiveControllingCoef has to be in range 0:1."))
    physiological_values[1] < physiological_values[2] || throw(DomainError("physiological_values is a tuple with minumum first and maximum second, is: $physiological_values."))
    # TODO

    # Remove values that are out of range of Physiological Heart beats
    if removeOutliers
        ok = (physiological_values[1] .<= signal .<= physiological_values[2])
        removed_nonphysiological_outliers = .!ok
        signal = signal[ok]
        @assert length(signal) >= 1 "Signal contains only outliers, recheck parameters."
    end
    # Calculate helper matrixes for the Binominal Filter
    bin_coeff = [1; 6; 15; 20; 15; 6; 1]
    coeff_sum = sum(bin_coeff)
    # normalize filter coefs
    bin_coeff =  bin_coeff .* (1/coeff_sum)
    coeff_count = length(bin_coeff)
    @assert (mod(coeff_count, 2) == 1) "Nono"

    # signalLen = length(signal)
    # filter_values = repeat(bin_coeff', signalLen)'
    # value_index_column = (1:coeff_count) .- Int(ceil(coeff_count/2))
    # hrv_ind_offset_val = repeat(value_index_column', signalLen)
    # hrv_ind_base_val = repeat(collect(1:signalLen)', coeff_count)'

    # #The first and last elements will be out of range work around by just (re-)using first and last elements...
    # hrv_idx_mat = hrv_ind_offset_val .+ hrv_ind_base_val
    # replace!(x -> x <=0 ? 1 : x, hrv_idx_mat)
    # replace!(x -> x > signalLen ? signalLen : x, hrv_idx_mat)

    # Calculate first filtered Signal (Through Binominal Filter)
    # filtered_signal = vec(sum(signal[hrv_idx_mat] .* filter_values', dims = 2))
    filtered_signal = filtfilt(bin_coeff, signal)

    # Apply "Adaptive Percent Filter" => fixed_hrv_timeseries
    adaptive_mean, adaptive_sigma = adaptive_moments(filtered_signal, coeff_count, adaptiveControllingCoef)
    adaptive_sigma_mean = mean(adaptive_sigma)

    last_good_value = signal[1] #Maybe filtered_signal(1)
    last_good_range = proportionalLimit * last_good_value + controllingFilterCoef * adaptive_sigma_mean

    hrv_diff = diff(signal)
    normal_values = repeat([true], length(filtered_signal))
    for i in 2: length(signal)
        current_diff = abs(hrv_diff[i - 1])
        current_max_range = proportionalLimit * signal[i-1] + controllingFilterCoef * adaptive_sigma_mean
        current_value_is_normal= current_diff <= current_max_range || current_diff <= last_good_range
        if current_value_is_normal
            last_good_value = signal[i]
            last_good_range = proportionalLimit * last_good_value + controllingFilterCoef * adaptive_sigma_mean
        end
        normal_values[i] = current_value_is_normal
    end
    # Just change reference, no copy
    fixed_signal = signal
    nonnormal_values = broadcast(!, normal_values)
    if eltype(signal) <: Int
        fixed_signal[nonnormal_values] = round.(Int, adaptive_mean[nonnormal_values] + (randn(sum(nonnormal_values)) .- 0.5) .* adaptive_sigma[nonnormal_values])
    else
        fixed_signal[nonnormal_values] = adaptive_mean[nonnormal_values] + (randn(sum(nonnormal_values)) .- 0.5) .* adaptive_sigma[nonnormal_values]
    end
    # Calculate second filtered Signal (Through Binominal Filter)
    #filtered_fixed_signal =  vec(sum(fixed_signal[hrv_idx_mat] .* filter_values', dims = 2))
    filtered_fixed_signal = filtfilt(bin_coeff, fixed_signal)

    # Apply "Adaptive Controlling Procedure" => fixed_fixed_signal
    adaptive_fixed_mean, adaptive_fixed_sigma = adaptive_moments(filtered_fixed_signal, coeff_count, adaptiveControllingCoef)

    normal_fixed_hrv_values = abs.(fixed_signal - adaptive_fixed_mean) .<= (controllingFilterCoef .* adaptive_fixed_sigma .+ controllingBasicVariability)

    fixed_fixed_signal = fixed_signal
    nonnormal_fixed_values = .!normal_fixed_hrv_values
    if eltype(signal) <: Int
        fixed_fixed_signal[nonnormal_fixed_values] = round.(Int, adaptive_fixed_mean[nonnormal_fixed_values] + (randn(sum(nonnormal_fixed_values)) .- 0.5) .* adaptive_fixed_sigma[nonnormal_fixed_values])
    else
        fixed_fixed_signal[nonnormal_fixed_values] = adaptive_fixed_mean[nonnormal_fixed_values] + (randn(sum(nonnormal_fixed_values)) .- 0.5) .* adaptive_fixed_sigma[nonnormal_fixed_values]
    end

    # Returns
    if replaceNonNormal
        if removeOutliers
            return fixed_fixed_signal, removed_nonphysiological_outliers, nonnormal_values
        else
            return fixed_fixed_signal, nonnormal_values
        end
    else
        if removeOutliers
            return removed_nonphysiological_outliers, nonnormal_values
        else
            return nonnormal_values
        end
    end
end

# Helper to calculate adaptive moments (mean, std)
function adaptive_moments(signal::Vector{<:Real}, init_length::Int, adaptiveControllingCoef::Real)
    adaptive_mean = Float64[]
    adaptive_variance = Float64[]
    push!(adaptive_mean, mean(signal[1:init_length]))
    push!(adaptive_variance, var(signal[1:init_length]))
    for i in 2:length(signal)
        push!(adaptive_mean, (adaptive_mean[end] - adaptiveControllingCoef * (adaptive_mean[end] - signal[i-1])))
        last_variance_item = (adaptive_mean[end-1] - signal[i-1]) ^ 2
        push!(adaptive_variance, adaptive_variance[end] - adaptiveControllingCoef * (adaptive_variance[end] - last_variance_item))
    end
   return adaptive_mean, sqrt.(adaptive_variance)
end


text:
## Mass segregation computation in stellar clusters.
## References : Cartwright et al. 2004, 2009 ; Allison et al. 2009 ; Olczak et al. 2011 ; Yu et al. 2011
##

## Create a graph with the selected stars
## ind: subset to select the x,y
## dtyp= "2d" or "3d"
##

##!! WARNING !! 2d is for the projection YZ since X is the distance with GAIA
##
function mst_graph(ind,x,y,z=[] , dtyp="2d")
    nxy= length(ind)
    if dtyp == "2d"
        A= Array{Float64}(undef,2,nxy)
    elseif dtyp == "3d"
       A= Array{Float64}(undef,3,nxy)
    end

    for i in 1:nxy
        if dtyp == "2d"
            A[:,i]= [y[ind[i]] z[ind[i]]]
        elseif dtyp == "3d"
            A[:,i]= [x[ind[i]] y[ind[i]] z[ind[i]]]
        end
    end

    d= Euclidean()
    p= pairwise(d, A, A, dims=2)

    src= Array{Int}(undef,0)
    dst= Array{Int}(undef,0)
    wgt= Array{Float64}(undef,0)

    for i in 1:nxy
        for j in 1:nxy
            if i != j
                push!(src,i)
                push!(dst,j)
                push!(wgt,p[i,j])
            end
        end
    end

    G= SimpleWeightedGraph(src, dst, wgt)
    kr= kruskal_mst(G, G.weights)

    return(kr)
end

## Computation of the λMST using either arithmetic or geometric average (Olczak et al. 2011)
##
## dtype= "ari" or "geo"
function lambda_mst(edges, dtyp="ari")
    nedg= length(edges)
    ltot= 0

    for w in edges
        if dtyp== "ari"
            ltot+= w.weight
        elseif dtyp== "geo"
            ltot+= log(w.weight)
        end
    end

    if dtyp== "ari"
        λmst= ltot/nedg
    elseif dtyp== "geo"
        λmst= exp(ltot/nedg)
    end
    return(λmst)
end

## estimate the sample size for the stats given n and N for the massive and total sample
## p= 99%
## Olczak et al. 2011)
function sample_size(n, N, p=0.99)
    k= log(1-p)/log(1-n/N)
    return(convert(Int,ceil(k)))
end

## Compute the mass segregation metric comparing two population set (2nd is random)
##
## nrandom: number of random set (default=1)
##
function kappa_ms(indMass, x, y , z, nrandom=1, dtyp="2d", daver= "geo")
    npop= length(indMass)
    ntot= length(x)

    ## Massive stars
    krMass= mst_graph(indMass, x,y,z , dtyp)
    λMass= lambda_mst(krMass,daver)
    ## Random set
    λArr= []
    κArr= []
    for i in 1:nrandom
        perm= randperm(ntot)
        indRef= perm[1:npop]
        krRef= mst_graph(indRef, x,y,z, dtyp)
        λRef= lambda_mst(krRef,daver)
        push!(λArr,λRef)
        push!(κArr, log(λRef/λMass))
    end
    λMean= mean(λArr)
    σκ= std(κArr)

    ## ratio (log) of the reference over massive ones
    κms= log(λMean/λMass)

    return(κms , σκ)
end

## Select the N percentile most massive stars (brighter)
##
## reverse:true -> light stars
function select_massivestars(mag, percentile=15 , reverse=false)
    nstar= length(mag)
    isort= sortperm(mag, rev=reverse)
    nselect= convert(Int, floor(nstar*percentile/100))

    return(isort[1:nselect])
end

## return the MS kappa and kappa_err
## IN
## mag: magnitude array
## X,Y,Z: positions array
## perc: percentile of massive stars
## dtyp: "3d" or "2d"
## daver: "geo" or "ari"
## reverse: false:massive, true:light - stars
##
function get_kappaMS(mag, X, Y, Z, perc, dtyp="3d", daver="geo", reverse=false)
    imass= select_massivestars(mag, perc, reverse)
    ns= sample_size(length(imass), length(X))
    κms, κms_err= kappa_ms(imass, X, Y, Z, ns, dtyp, daver)
    return(κms, κms_err)
end

#####################
## Compute the Q ratio structure to measure the clustering of the OC.
## See Cartwright et al. (2004)
## directly adapted to 3D
##

### !!!!!!! Normalization to be checked on λ , Cartwright & 2004
function get_Q(X, Y , Z)
    nxyz= length(X)
    A= Array{Float64}(undef,3,nxyz)
    for i in 1:nxyz
            A[:,i]= [X[i] Y[i] Z[i]]
    end

    d= Euclidean()
    p= pairwise(d, A, A, dims=2)

    src= Array{Int}(undef,0)
    dst= Array{Int}(undef,0)
    wgt= Array{Float64}(undef,0)

    for i in 1:nxyz
        for j in 1:nxyz
            if i != j
                push!(src,i)
                push!(dst,j)
                push!(wgt,p[i,j])
            end
        end
    end

    G= SimpleWeightedGraph(src, dst, wgt)
    kr= kruskal_mst(G, G.weights)

    ## compute the mean (arithmtic) edge of MST and of star separation. We do not apply other normalization
    λmst= 0
    for w in kr
        λmst+= w.weight
    end
    ## not the original normalization of Cartwright!
    k= (4π / 3 / nxyz )^(1/3)
    λmst /= k*length(kr)

    npair= nxyz*(nxyz-1)/2
    sbar= sum(p) / npair

    Q= λmst / sbar
    return(Q)
end

max_stars_repo_name: bosscha/gaia-shock

max_issues_repo_name: bosscha/gaia-shock

max_forks_repo_name: bosscha/gaia-shock