# Stage-5: unknown-format discovery by Bayesian header clustering.
#
# See model/DESIGN_clustering.md for the full rationale. In brief: files that
# stage-3 sorted into `binary/` are the `:unknown` sink — genuinely
# unrecognized bytes. This stage clusters them by *file format* (not producer)
# using the first `HEADER_N` header bytes, modeled as a Dirichlet-process
# mixture of per-position categoricals over a 257-symbol alphabet
# (byte 0–255, plus symbol 257 = "past EOF"). Each cluster's signature is a
# magic-number template that can be promoted into the classifier's fast path.
#
# This file is deliberately dependency-light: everything below is base Julia
# (only `log`, no `SpecialFunctions`), so it drops into the existing module and
# the offline calibration script alike without touching the Manifest. The model
# is categorical on purpose — do NOT reuse model.jl's [0,1] byte scaling here
# (that metric is meaningful for the Lux net and meaningless for header bytes,
# where 0x89 and 0x88 are not "close"; see DESIGN §4).

"Number of leading header bytes modeled per file (the feature window). DESIGN §4b."
const HEADER_N = 32

"Alphabet size: byte values 0–255 plus one extra symbol for 'past end of file'."
const ALPHABET = 257

"The 'past EOF' symbol (1-based index `ALPHABET`). A file shorter than a given
position emits this here — real, discriminative signal for fixed-length formats,
and it avoids colliding zero-padding with genuine 0x00 header bytes (DESIGN §4)."
const PAST_EOF = ALPHABET

# ---------------------------------------------------------------------------
# Feature extraction
# ---------------------------------------------------------------------------

"""
    header_symbols(path; n=HEADER_N) -> Vector{Int}

Read the first `n` bytes of the file at `path` and map them to a length-`n`
vector of 1-based categorical symbols: byte value `b` → `b + 1` (so `1..256`),
and every position at or beyond end-of-file → `PAST_EOF` (`257`). Never reads
more than `n` bytes, so memory stays flat regardless of file size.
"""
function header_symbols(path::AbstractString; n::Integer=HEADER_N)
    syms = fill(PAST_EOF, n)
    open(path, "r") do io
        bytes = read(io, n)
        @inbounds for i in eachindex(bytes)
            syms[i] = Int(bytes[i]) + 1
        end
    end
    return syms
end

"""
    header_matrix(paths; n=HEADER_N) -> Matrix{Int}

Stack `header_symbols` for every path into an `n × length(paths)` matrix (one
column per file), the input layout the Gibbs sampler and predictive scorer both
consume.
"""
function header_matrix(paths::AbstractVector{<:AbstractString}; n::Integer=HEADER_N)
    X = Matrix{Int}(undef, n, length(paths))
    for (j, p) in enumerate(paths)
        X[:, j] = header_symbols(p; n=n)
    end
    return X
end

# ---------------------------------------------------------------------------
# Model: DP mixture of per-position categoricals (Dirichlet-Categorical)
# ---------------------------------------------------------------------------

"""
    ClusterStats

Sufficient statistics for one cluster: a per-position count table `counts`
(`n × ALPHABET`; `counts[i, v]` = how many member files show symbol `v` at
position `i`) and the member count `members`. These are exactly what phase-B
persists per catalog entry, and everything the collapsed predictive needs.
A slot with `members == 0` is inactive (reusable) — the Gibbs sweep prunes
emptied clusters without renumbering, so surviving cluster ids stay stable.
"""
mutable struct ClusterStats
    counts::Matrix{Int}   # n × ALPHABET
    members::Int
end

ClusterStats(n::Integer) = ClusterStats(zeros(Int, n, ALPHABET), 0)

"Add file `x` (a length-n symbol vector) into cluster `c`'s sufficient stats."
function add!(c::ClusterStats, x::AbstractVector{<:Integer})
    @inbounds for i in eachindex(x)
        c.counts[i, x[i]] += 1
    end
    c.members += 1
    return c
end

"Remove file `x` from cluster `c`'s sufficient stats (inverse of `add!`)."
function remove!(c::ClusterStats, x::AbstractVector{<:Integer})
    @inbounds for i in eachindex(x)
        c.counts[i, x[i]] -= 1
    end
    c.members -= 1
    return c
end

"""
    log_predictive(c, x, β) -> Float64

Log probability that file `x` was produced by cluster `c` under the collapsed
Dirichlet-Categorical predictive, given `c`'s current counts: at each position
`i`, `p(x_i | c) = (counts[i, x_i] + β) / (members + ALPHABET·β)`, summed in log
space over positions. Call with `c` NOT containing `x` (Gibbs excludes the point
being resampled), so an emptied cluster reduces to the uniform prior `1/ALPHABET`
per position — identical to a brand-new cluster, as it should be.
"""
function log_predictive(c::ClusterStats, x::AbstractVector{<:Integer}, β::Float64)
    denom = log(c.members + ALPHABET * β)
    s = 0.0
    @inbounds for i in eachindex(x)
        s += log(c.counts[i, x[i]] + β) - denom
    end
    return s
end

"Log likelihood of `x` under the fixed uniform component (each position uniform
over the 257 symbols): `n · log(1/ALPHABET)`. Used for both the never-adaptive
background 'junk drawer' and the prior predictive of a fresh cluster (DESIGN §4a)."
log_uniform(n::Integer) = -n * log(ALPHABET)

# Lanczos approximation to log Γ(x) for x > 0, so partition scoring (below) needs
# the Dirichlet-multinomial marginal's gamma terms without pulling in
# SpecialFunctions — keeping this stage dependency-flat (no Manifest churn).
# g = 7, standard coefficients; accurate to ~1e-14 over the range we use.
const _LANCZOS_G = 7
const _LANCZOS_C = (0.99999999999980993, 676.5203681218851, -1259.1392167224028,
                    771.32342877765313, -176.61502916214059, 12.507343278686905,
                    -0.13857109526572012, 9.9843695780195716e-6, 1.5056327351493116e-7)

function loggamma(x::Float64)
    x < 0.5 && return log(π / sin(π * x)) - loggamma(1.0 - x)  # reflection
    x -= 1.0
    a = _LANCZOS_C[1]
    t = x + _LANCZOS_G + 0.5
    @inbounds for i in 1:_LANCZOS_G + 1
        a += _LANCZOS_C[i + 1] / (x + i)
    end
    return 0.5 * log(2π) + (x + 0.5) * log(t) - t + log(a)
end

# ---------------------------------------------------------------------------
# Phase A: collapsed Gibbs sampler (offline — the science)
# ---------------------------------------------------------------------------

"""
    GibbsResult

Output of `gibbs_cluster`: `assignments` (one per input file; `0` = absorbed by
the background junk drawer, positive ints = cluster id), the surviving
`clusters` keyed by id, and `score` (the partition's collapsed pseudo-likelihood,
used to rank restarts).
"""
struct GibbsResult
    assignments::Vector{Int}
    clusters::Dict{Int,ClusterStats}
    score::Float64
end

"""
    gibbs_cluster(X; α, β, bg_mass, sweeps, restarts, rng) -> GibbsResult

Cluster the columns of `X` (an `n × N` header-symbol matrix) with a collapsed
Gibbs sampler over a CRP/Dirichlet-Categorical mixture plus a fixed uniform
background (DESIGN §5). Unknown *k* falls out of the CRP natively.

Per point, per sweep, the point is removed from its cluster and reassigned by
sampling from the CRP-predictive weights:

  * existing cluster `k`: `members_k · exp(log_predictive)`
  * background:           `bg_mass · (1/ALPHABET)^n`   (never adapts)
  * a fresh cluster:      `α · (1/ALPHABET)^n`

A uniform/high-entropy blob matches no structured cluster, and background vs.
fresh is then decided by `bg_mass` vs. `α`; with `bg_mass ≥ α` such blobs are
absorbed rather than minting singletons. `restarts` independent runs are made
from different seeds and the highest-scoring partition is returned (a cheap,
base-only stand-in for the offline Binder/VI point-summary the design defers).
"""
function gibbs_cluster(X::AbstractMatrix{<:Integer};
                       α::Float64=1.0, β::Float64=0.5, bg_mass::Float64=5.0,
                       sweeps::Integer=80, restarts::Integer=4,
                       rng::AbstractRNG=Random.default_rng())
    best = nothing
    for _ in 1:restarts
        r = _gibbs_once(X; α=α, β=β, bg_mass=bg_mass, sweeps=sweeps, rng=rng)
        if best === nothing || r.score > best.score
            best = r
        end
    end
    return best
end

function _gibbs_once(X::AbstractMatrix{<:Integer};
                     α::Float64, β::Float64, bg_mass::Float64,
                     sweeps::Integer, rng::AbstractRNG)
    n, N = size(X)
    # Seed every file in its own singleton (NOT the background). Cold-starting
    # from the background deadlocks: at a format's first file, a fresh cluster
    # and the background are equally uniform, so with bg_mass ≥ α the background
    # always wins and no real cluster is ever seeded. Singleton init sidesteps
    # this — same-format singletons merge and snowball, while a lone
    # random-blob singleton dissolves on resample and is reclaimed by the
    # (stickier) background. See DESIGN §4a.
    z = collect(1:N)
    clusters = Dict{Int,ClusterStats}()
    for j in 1:N
        c = ClusterStats(n)
        add!(c, view(X, :, j))
        clusters[j] = c
    end
    next_id = N + 1
    log_u = log_uniform(n)
    log_bg = log(bg_mass) + log_u
    log_new = log(α) + log_u

    # Reused across every point-visit so the sampler's hot loop allocates nothing
    # per step (2M+ visits per run): `idbuf[t]` is the cluster id whose weight is
    # `logw[t+1]` (logw[1] = background, logw[end] = fresh). Rebuilding these with
    # fresh `Vector`/`collect(keys(...))` each step was both slow and enough GC
    # churn to trip a Julia GC segfault on long grid runs.
    idbuf = Int[]
    logw = Float64[]
    for _ in 1:sweeps
        for j in 1:N
            x = view(X, :, j)

            # Remove point j from its current component.
            zj = z[j]
            if zj > 0
                c = clusters[zj]
                remove!(c, x)
                if c.members == 0
                    delete!(clusters, zj)    # prune emptied cluster; id retired
                end
            end

            # Candidate log-weights: background, each live cluster, fresh.
            empty!(idbuf); empty!(logw)
            push!(logw, log_bg)
            for (k, c) in clusters
                push!(idbuf, k)
                push!(logw, log(c.members) + log_predictive(c, x, β))
            end
            push!(logw, log_new)

            # Gumbel-max sample from the categorical over components.
            pick = _gumbel_argmax(logw, rng)

            if pick == 1
                z[j] = 0                     # background
            elseif pick == length(logw)
                id = next_id; next_id += 1   # fresh cluster
                c = ClusterStats(n)
                add!(c, x)
                clusters[id] = c
                z[j] = id
            else
                id = idbuf[pick - 1]
                add!(clusters[id], x)
                z[j] = id
            end
        end
    end

    return GibbsResult(z, clusters, partition_logmarginal(X, z, clusters, α, β))
end

"Argmax of `logw .+ Gumbel noise` — an exact draw from softmax(logw) without
normalizing (numerically safe for the tiny header-likelihood magnitudes)."
function _gumbel_argmax(logw::AbstractVector{Float64}, rng::AbstractRNG)
    best_i = 1
    best_v = -Inf
    @inbounds for i in eachindex(logw)
        g = logw[i] - log(-log(rand(rng)))
        if g > best_v
            best_v = g
            best_i = i
        end
    end
    return best_i
end

"""
    partition_logmarginal(X, z, clusters, α, β) -> Float64

The joint log-evidence `log p(z, X)` of a partition under the CRP prior and the
collapsed Dirichlet-Categorical likelihood — the principled score for ranking
Gibbs restarts (higher = better). It is the sum of:

  * the Dirichlet-multinomial **marginal** of each cluster's per-position counts,
    `lΓ(Aβ) − lΓ(mₖ+Aβ) + Σ_v [lΓ(c_v+β) − lΓ(β)]`, whose normalizer supplies the
    Occam penalty that a plain product-of-predictives lacks — it is what makes a
    *merged, heterogeneous* cluster score **worse** than two clean ones (an
    earlier pseudo-likelihood scorer omitted this and wrongly rewarded merging);
  * the CRP prior over the clustered points, `K·log α + Σₖ lΓ(mₖ) + lΓ(α) −
    lΓ(α+N_clustered)`, penalizing gratuitous extra clusters; and
  * the fixed uniform term for background-assigned files.
"""
function partition_logmarginal(X::AbstractMatrix{<:Integer}, z::AbstractVector{<:Integer},
                               clusters::Dict{Int,ClusterStats}, α::Float64, β::Float64)
    n, N = size(X)
    Aβ = ALPHABET * β
    lg_Aβ = loggamma(Aβ)
    lg_β = loggamma(β)
    s = 0.0
    # Dirichlet-Categorical marginal likelihood, per cluster × position.
    for (_, c) in clusters
        lg_denom = loggamma(c.members + Aβ)
        @inbounds for i in 1:n, v in 1:ALPHABET
            cv = c.counts[i, v]
            cv > 0 && (s += loggamma(cv + β) - lg_β)
        end
        s += n * (lg_Aβ - lg_denom)
    end
    # CRP prior over the partition of the clustered points.
    n_bg = count(==(0), z)
    n_clustered = N - n_bg
    K = length(clusters)
    s += K * log(α) + loggamma(α) - loggamma(α + n_clustered)
    for (_, c) in clusters
        s += loggamma(float(c.members))
    end
    # Background files: independent, uniform.
    s += n_bg * log_uniform(n)
    return s
end

# ---------------------------------------------------------------------------
# Phase B: sequential CRP-predictive assignment (online — the catalog)
# ---------------------------------------------------------------------------

"""
    assign_file(x, clusters, ids; α, β, bg_mass) -> Int

Deterministically assign a single file `x` against an existing catalog: the
same CRP-predictive rule as Gibbs but at the **argmax** (no sampling) with the
current assignments held fixed (DESIGN §5B). Returns the id of the chosen
cluster, `0` for the background, or `-1` to signal "mint a new cluster". `ids`
is the caller's stable ordering of `keys(clusters)`.

A new cluster is minted (`-1`) only when the fresh-cluster weight strictly wins.
Fresh and background share the same `(1/ALPHABET)^n` likelihood (one file, however
structured, is indistinguishable from a uniform blob until a *second* like file
appears), so this reduces to `α > bg_mass`. Under the calibrated `bg_mass > α`,
minting is therefore effectively off on the live path **by design**: a novel file
that matches nothing parks in the background, and genuinely new formats are
discovered by the periodic **offline Gibbs compaction** re-clustering that
residue (DESIGN §5), not by single-file minting. Because ids are frozen at birth
by the caller, there is no label switching.
"""
function assign_file(x::AbstractVector{<:Integer}, clusters::Dict{Int,ClusterStats},
                     ids::AbstractVector{<:Integer};
                     α::Float64=1.0, β::Float64=0.5, bg_mass::Float64=5.0)
    n = length(x)
    log_u = log_uniform(n)
    best_kind = :bg          # :bg, :existing, :new
    best_id = 0
    best = log(bg_mass) + log_u
    new_w = log(α) + log_u
    if new_w > best
        best = new_w; best_kind = :new
    end
    for k in ids
        c = clusters[k]
        w = log(c.members) + log_predictive(c, x, β)
        if w > best
            best = w; best_kind = :existing; best_id = k
        end
    end
    return best_kind === :bg ? 0 : best_kind === :new ? -1 : best_id
end

# ---------------------------------------------------------------------------
# Signatures and promotion (closing the loop to the classifier — DESIGN §6)
# ---------------------------------------------------------------------------

"""
    signature(c; peak_threshold=0.9, β=0.5) -> Vector{Union{Int,Nothing}}

Turn a cluster's counts into a libmagic-style template: at each position, if the
modal symbol's posterior probability exceeds `peak_threshold`, that byte is
*required* (returned as the 0–255 byte value, or `PAST_EOF`); otherwise the
position is a wildcard (`nothing`). The vector of required bytes IS the
magic-number template — the whole point of the categorical model (DESIGN §4).

Peakedness is a **Bernoulli** question ("is this position fixed to byte `v`, or
not?"), so it uses a 2-way posterior mean `(count + β)/(members + 2β)` — NOT the
257-way mixture predictive. The alphabet-wide version would smear the estimate
across 257 symbols (`members + 257β` in the denominator), pulling even a
unanimous position below any sane threshold once β is small — which would make
promotion impossible. This decouples signature detection from the clustering
pseudocount and the alphabet size.
"""
function signature(c::ClusterStats; peak_threshold::Float64=0.9, β::Float64=0.5)
    n = size(c.counts, 1)
    sig = Vector{Union{Int,Nothing}}(nothing, n)
    c.members == 0 && return sig
    denom = c.members + 2β
    @inbounds for i in 1:n
        v = argmax(view(c.counts, i, :))
        p = (c.counts[i, v] + β) / denom
        if p > peak_threshold
            sig[i] = v == PAST_EOF ? PAST_EOF : v - 1   # back to raw byte value
        end
    end
    return sig
end

"Number of fixed (non-wildcard) positions in a signature — its 'magic length'."
magic_positions(sig::AbstractVector) = count(!isnothing, sig)

"""
    is_promotable(c, sig; min_members=20, min_magic=3) -> Bool

A cluster qualifies for *nomination* (still human-gated, DESIGN §6) when it has
at least `min_members` files AND at least `min_magic` fixed signature positions.
The background (id 0) is never passed here — it is never promotable by design.
"""
function is_promotable(c::ClusterStats, sig::AbstractVector;
                       min_members::Integer=20, min_magic::Integer=3)
    return c.members >= min_members && magic_positions(sig) >= min_magic
end

# ---------------------------------------------------------------------------
# Calibration metrics (DESIGN §7): agreement of recovered clusters vs. truth
# ---------------------------------------------------------------------------

"Map a label vector to consecutive integer ids and a group→indices table."
function _groups(labels::AbstractVector)
    g = Dict{Any,Vector{Int}}()
    for (i, l) in enumerate(labels)
        push!(get!(g, l, Int[]), i)
    end
    return g
end

"""
    adjusted_rand_index(a, b) -> Float64

Adjusted Rand Index between two labelings of the same items: 1.0 = identical
partitions (up to relabeling), ~0.0 = chance agreement, can go negative. This is
the §7 calibration objective — grid-tuning maximizes ARI of recovered-vs-truth
(magic-collapsed) labels. Hand-rolled to keep the dependency footprint flat;
matches `Clustering.randindex`.
"""
function adjusted_rand_index(a::AbstractVector, b::AbstractVector)
    length(a) == length(b) || throw(DimensionMismatch("label vectors differ in length"))
    n = length(a)
    n < 2 && return 1.0
    ga = collect(values(_groups(a)))
    gb = collect(values(_groups(b)))
    # Contingency-table sum of C(n_ij, 2).
    comb2(x) = x * (x - 1) / 2
    sa = Set.(ga)
    index = 0.0
    for A in sa, B in gb
        nij = count(in(A), B)
        index += comb2(nij)
    end
    sum_a = sum(comb2(length(g)) for g in ga)
    sum_b = sum(comb2(length(g)) for g in gb)
    total = comb2(n)
    expected = sum_a * sum_b / total
    maxi = (sum_a + sum_b) / 2
    denom = maxi - expected
    return denom == 0 ? 1.0 : (index - expected) / denom
end

"""
    v_measure(truth, pred; β=1.0) -> (v, homogeneity, completeness)

Entropy-based cluster agreement (Rosenberg & Hirschberg): homogeneity (each
predicted cluster holds one true class), completeness (each true class stays in
one predicted cluster), and their weighted harmonic mean `v`. Reported alongside
ARI in §7 calibration as a second, differently-biased view.
"""
function v_measure(truth::AbstractVector, pred::AbstractVector; β::Float64=1.0)
    n = length(truth)
    n == 0 && return (1.0, 1.0, 1.0)
    gt = _groups(truth)
    gp = _groups(pred)
    entropy(g) = -sum((length(v) / n) * log(length(v) / n) for v in values(g))
    H_C = entropy(gt)
    H_K = entropy(gp)
    # Conditional entropies via the contingency table.
    H_CK = 0.0  # H(truth | pred)
    H_KC = 0.0  # H(pred | truth)
    for (_, P) in gp
        Ps = Set(P)
        for (_, C) in gt
            nij = count(in(Ps), C)
            nij == 0 && continue
            H_CK -= (nij / n) * log(nij / length(P))
        end
    end
    for (_, C) in gt
        Cs = Set(C)
        for (_, P) in gp
            nij = count(in(Cs), P)
            nij == 0 && continue
            H_KC -= (nij / n) * log(nij / length(C))
        end
    end
    homogeneity = H_C == 0 ? 1.0 : 1 - H_CK / H_C
    completeness = H_K == 0 ? 1.0 : 1 - H_KC / H_K
    v = (homogeneity + completeness == 0) ? 0.0 :
        (1 + β) * homogeneity * completeness / (β * homogeneity + completeness)
    return (v, homogeneity, completeness)
end