#!/usr/bin/env julia
#
# Phase-A calibration for stage-5 header clustering (model/DESIGN_clustering.md
# §7). Runs labeled known files through the exact clustering pipeline, scores the
# recovered partition against magic-collapsed ground truth with ARI / V-measure,
# grid-tunes (α, β, bg_mass, n), and cross-checks the winning config against a
# model-free NCD (gzip) baseline (§8). The settings printed here are the ones the
# machine rediscovers known formats at — copy the winner into config.jl.
#
#   julia --project=. bin/cluster_calibrate.jl [training_set_dir]
#
# Defaults to ../training_set. Prints a report; writes nothing.

using Random
using Printf

include(joinpath(@__DIR__, "..", "src", "cluster.jl"))

# --- ground truth: magic-collapsed classes, NOT extensions (DESIGN §7.2) -----

"""
    truth_label(path) -> String

The magic-collapsed format class of a file, read from its actual bytes (so
docx≡zip and the whole ELF family merge, exactly the answer we want the
clustering to reproduce). `tar` is detected by the `ustar` magic at offset 257 —
outside the model's front window, so tars are the accepted blind spot that
scatters to background.
"""
function truth_label(path::AbstractString)
    b = zeros(UInt8, 262)
    open(path) do io
        chunk = read(io, 262)
        copyto!(b, 1, chunk, 1, length(chunk))
    end
    b[1] == 0x1f && b[2] == 0x8b && return "gzip"
    b[1] == 0x50 && b[2] == 0x4b && return "pkzip"
    b[1] == 0x25 && b[2] == 0x50 && b[3] == 0x44 && b[4] == 0x46 && return "pdf"
    b[1] == 0xff && b[2] == 0xd8 && b[3] == 0xff && return "jpeg"
    b[1] == 0x7f && b[2] == 0x45 && b[3] == 0x4c && b[4] == 0x46 && return "elf"
    (b[258] == 0x75 && b[259] == 0x73 && b[260] == 0x74 && b[261] == 0x61 && b[262] == 0x72) && return "tar"
    return "other"
end

# --- NCD (Normalized Compression Distance) baseline, model-free (DESIGN §8) ---

"gzip-compressed size of a byte buffer, via the gzip CLI (no CodecZlib dep)."
function gz_size(bytes::Vector{UInt8})
    out = IOBuffer()
    open(pipeline(`gzip -c`; stdout=out); write=true) do io
        write(io, bytes)
    end
    return length(take!(out))
end

"NCD(x,y) = (C(xy) - min(C(x),C(y))) / max(C(x),C(y)) — 0 = identical, ~1 = unrelated."
function ncd(xb, yb, cx, cy)
    cxy = gz_size(vcat(xb, yb))
    return (cxy - min(cx, cy)) / max(cx, cy)
end

"""
    ncd_1nn_purity(paths, truth; head_bytes) -> Float64

Fraction of files whose NCD-nearest neighbour shares its true label — a cheap,
O(N²) sanity read on how well raw gzip-similarity alone separates formats on the
same input. The Bayesian clusters should broadly agree; a big gap is a red flag
(DESIGN §10.3). Uses each file's first `head_bytes` so the giant files don't
dominate compression time.
"""
function ncd_1nn_purity(paths::Vector{String}, truth::Vector{String}; head_bytes::Int=4096)
    bufs = map(paths) do p
        open(io -> read(io, head_bytes), p)
    end
    csz = gz_size.(bufs)
    N = length(paths)
    correct = 0
    for i in 1:N
        best_j = 0; best_d = Inf
        for j in 1:N
            i == j && continue
            d = ncd(bufs[i], bufs[j], csz[i], csz[j])
            if d < best_d
                best_d = d; best_j = j
            end
        end
        best_j != 0 && truth[best_j] == truth[i] && (correct += 1)
    end
    return correct / N
end

"1-NN label purity of a *cluster* assignment vs truth (same yardstick as NCD's)."
function cluster_1nn_purity(pred::Vector{Int}, truth::Vector{String})
    # For each file, its 'nearest neighbour' is any other file in the same
    # cluster; purity = P(a random same-cluster neighbour shares the true label).
    groups = Dict{Int,Vector{Int}}()
    for (i, k) in enumerate(pred)
        push!(get!(groups, k, Int[]), i)
    end
    correct = 0; total = 0
    for (_, idxs) in groups
        length(idxs) < 2 && continue
        for i in idxs
            same = count(j -> j != i && truth[j] == truth[i], idxs)
            total += 1
            same > 0 && (correct += 1)
        end
    end
    return total == 0 ? 0.0 : correct / total
end

# --- data ---------------------------------------------------------------------

function load_corpus(dir::AbstractString)
    paths = String[]
    for name in readdir(dir; join=true)
        isfile(name) && push!(paths, name)
    end
    truth = truth_label.(paths)
    return paths, truth
end

# --- grid search --------------------------------------------------------------

function evaluate(X, truth; α, β, bg_mass, sweeps, restarts, seed)
    r = gibbs_cluster(X; α=α, β=β, bg_mass=bg_mass, sweeps=sweeps,
                      restarts=restarts, rng=MersenneTwister(seed))
    pred = r.assignments
    ari = adjusted_rand_index(truth, pred)
    keep = truth .!= "tar"
    ari_notar = adjusted_rand_index(truth[keep], pred[keep])
    v, h, comp = v_measure(truth, pred)
    return (; ari, ari_notar, v, h, comp, k=length(r.clusters),
            bg=count(==(0), pred), result=r)
end

function main()
    dir = length(ARGS) >= 1 ? ARGS[1] : joinpath(@__DIR__, "..", "..", "training_set")
    isdir(dir) || error("training set dir not found: $dir")
    paths, truth = load_corpus(dir)
    classes = sort(unique(truth))
    counts = [(c, count(==(c), truth)) for c in classes]
    @printf("corpus: %d files from %s\n", length(paths), dir)
    println("magic-collapsed truth classes: ", join(["$c=$n" for (c, n) in counts], "  "))
    println()

    sweeps = 150
    restarts = 6
    seed = 20260703

    # Grid. n is expensive to re-featurize, so loop it outermost. Ranges are
    # centred where the coarse sweep found the optimum: small β (peaked
    # per-position priors) is what separates formats whose headers differ in only
    # a few magic bytes; large β over-merges. bg_mass barely moves the result
    # here (almost nothing lands in background on this corpus), so it is fixed.
    αs   = [1.0, 2.0]
    βs   = [0.05, 0.08, 0.1, 0.15, 0.2]
    bgs  = [5.0]
    ns   = [32, 64]

    println("grid search (sweeps=$sweeps, restarts=$restarts):")
    @printf("  %-4s %-5s %-5s %-6s | %-6s %-8s %-6s %-6s %-6s %-4s %-4s\n",
            "n", "alpha", "beta", "bgmss", "ARI", "ARI-tar", "V", "homog", "compl", "k", "bg")
    results = Vector{Any}()
    for n in ns
        X = header_matrix(paths; n=n)
        for α in αs, β in βs, bg in bgs
            e = evaluate(X, truth; α=α, β=β, bg_mass=bg, sweeps=sweeps, restarts=restarts, seed=seed)
            push!(results, (; n, α, β, bg, e))
            @printf("  %-4d %-5.1f %-5.2f %-6.1f | %-6.3f %-8.3f %-6.3f %-6.3f %-6.3f %-4d %-4d\n",
                    n, α, β, bg, e.ari, e.ari_notar, e.v, e.h, e.comp, e.k, e.bg)
        end
    end

    # Rank by ARI-excluding-tar (tar is the accepted blind spot; scoring it would
    # penalise the correct answer of scattering tars to background — DESIGN §7.2).
    sort!(results; by=r -> r.e.ari_notar, rev=true)
    best = results[1]
    println()
    @printf("BEST (by ARI excl. tar): n=%d α=%.1f β=%.2f bg_mass=%.1f\n",
            best.n, best.α, best.β, best.bg)
    @printf("  ARI=%.3f  ARI(excl tar)=%.3f  V=%.3f  homogeneity=%.3f  completeness=%.3f  clusters=%d  background=%d\n",
            best.e.ari, best.e.ari_notar, best.e.v, best.e.h, best.e.comp, best.e.k, best.e.bg)

    # Per-cluster composition of the winning partition, and promotion nominations.
    pred = best.e.result.assignments
    println("\nwinning partition — cluster composition (truth breakdown):")
    for (id, c) in sort(collect(best.e.result.clusters); by=x -> -x[2].members)
        members = [truth[i] for i in eachindex(pred) if pred[i] == id]
        comp = sort([(l, count(==(l), members)) for l in unique(members)]; by=x -> -x[2])
        sig = signature(c)
        promo = is_promotable(c, sig; min_members=20, min_magic=3) ? "  ✓NOMINATE" : ""
        @printf("  cluster %-4d  n=%-3d  magic=%-2d  %s%s\n",
                id, c.members, magic_positions(sig),
                join(["$l:$k" for (l, k) in comp], " "), promo)
    end
    nbg = count(==(0), pred)
    bg_truth = [truth[i] for i in eachindex(pred) if pred[i] == 0]
    bgc = sort([(l, count(==(l), bg_truth)) for l in unique(bg_truth)]; by=x -> -x[2])
    @printf("  background   n=%-3d          %s\n", nbg, join(["$l:$k" for (l, k) in bgc], " "))

    # NCD baseline cross-check on a subsample (O(N²), so keep it small).
    println("\nNCD (gzip) baseline cross-check:")
    subn = min(150, length(paths))
    sub = shuffle(MersenneTwister(seed), collect(1:length(paths)))[1:subn]
    subpaths = paths[sub]; subtruth = truth[sub]
    ncd_pur = ncd_1nn_purity(subpaths, subtruth)
    Xsub = header_matrix(subpaths; n=best.n)
    rsub = gibbs_cluster(Xsub; α=best.α, β=best.β, bg_mass=best.bg,
                         sweeps=sweeps, restarts=restarts, rng=MersenneTwister(seed))
    bay_pur = cluster_1nn_purity(rsub.assignments, subtruth)
    @printf("  subsample=%d  NCD 1-NN label purity=%.3f   Bayesian same-cluster purity=%.3f\n",
            subn, ncd_pur, bay_pur)
    println("  (both high ⇒ header-byte signal agrees with model-free gzip similarity — DESIGN §10.3)")
end

main()