#!/usr/bin/env julia
#
# bench_model.jl — microbenchmark the classifier in isolation, with no server,
# no queue, and no disk in the way.
#
# bin/bench.jl measures the *pipeline*: it reports stage 1 as one number, the
# wall time of `handle_classify_job`, which is feature reads + inference + a
# rename + a log line, under whatever thread contention the other three pools are
# creating. That number is the right one for capacity planning and the wrong one
# for answering "is the model slow?". This script answers that question by taking
# the model apart:
#
#   read_features   open, read 16 bytes, seek, read 16 bytes, scale
#   Lux.apply       the network itself, on a feature vector already in memory
#   classify        both together — what stage 1 actually calls per file
#
# Three properties are worth checking beyond the raw per-file cost:
#
#   * Feature reads should be flat in file size. read_features seeks to the tail
#     rather than slurping, so a 1 GiB file should cost the same as a 1 KiB one.
#     (This is the same claim bin/bench.jl makes about memory, on the CPU axis.)
#   * Batching should be much cheaper per file. A 32x1 matmul wastes most of a
#     BLAS call; if batch-64 inference is many times cheaper per file, that is
#     the headroom a batching stage-1 would buy — worth knowing before building
#     one, since today the pipeline classifies strictly one file at a time.
#   * Inference should scale across threads. `Classifier` is shared read-only by
#     the whole stage-1 pool on the claim that Lux inference is pure. If per-file
#     cost degrades as tasks are added, that claim holds but BLAS threading is
#     fighting the worker pool, and stage-1 workers are contending, not scaling.
#
# Usage:
#   julia --project=. -t auto bin/bench_model.jl [options]
#
#   --model PATH     classifier artifact (default: $FS_MODEL_PATH or model/classifier.jld2)
#   --reps N         inference calls per timed trial (default: 20000)
#   --trials N       timed trials; the minimum is reported (default: 5)
#   --batches LIST   batch sizes to sweep, comma-separated (default: 1,8,64,512)
#   --sizes LIST     file sizes for the read_features sweep (default: 1k,64k,4m,256m)
#   --no-threads     skip the thread-scaling sweep
#   --json PATH      also write the results as JSON
#
# Reported times are the *minimum* over trials: for a microbenchmark the floor is
# the signal and everything above it is scheduler and GC noise.

using JSON3
using Random
using Statistics
using Printf

# The script is run directly, not as part of the package, so pull in exactly the
# pieces the classifier needs. `Lux`/`JLD2` first — model.jl and classify.jl both
# assume the including scope already has them (see the note at the top of model.jl).
using Lux
using JLD2
using LinearAlgebra

const SRC = joinpath(dirname(@__DIR__), "src")
include(joinpath(SRC, "model.jl"))
include(joinpath(SRC, "classify.jl"))

# ---------------------------------------------------------------- option parsing

const DEFAULTS = Dict{String,Any}(
    "model"       => get(ENV, "FS_MODEL_PATH", "model/classifier.jld2"),
    "reps"        => 20_000,
    "trials"      => 5,
    "batches"     => "1,8,64,512",
    "sizes"       => "1k,64k,4m,256m",
    "no-threads"  => false,
    "json"        => nothing,
)

const FLAGS = ("no-threads",)

function parse_size(s::AbstractString)::Int
    m = match(r"^(\d+(?:\.\d+)?)\s*([kKmMgG]?)[bB]?$", strip(s))
    m === nothing && error("bad size: $s (expected e.g. 512, 64k, 8m, 1g)")
    mult = Dict('k' => 1024, 'm' => 1024^2, 'g' => 1024^3)
    scale = isempty(m[2]) ? 1 : mult[lowercase(m[2])[1]]
    return round(Int, parse(Float64, m[1]) * scale)
end

function parse_args(argv)
    opts = copy(DEFAULTS)
    i = 1
    while i <= length(argv)
        a = argv[i]
        startswith(a, "--") || error("unexpected argument: $a")
        key = a[3:end]
        haskey(opts, key) || error("unknown option: $a")
        if key in FLAGS
            opts[key] = true; i += 1; continue
        end
        i + 1 <= length(argv) || error("option --$key needs a value")
        opts[key] = key in ("reps", "trials") ? parse(Int, argv[i+1]) : argv[i+1]
        i += 2
    end
    return opts
end

# ------------------------------------------------------------------- measurement

# Every timed loop stores its result here. Without a visible side effect the
# compiler is free to hoist a pure call out of the loop and we would be timing an
# empty `for`.
const SINK = Ref{Any}(nothing)

"""
    measure(f, reps; trials) -> (ns_per_op, bytes_per_op)

Time `f` over `reps` calls, `trials` times, and report the fastest trial.

The first call is thrown away: it pays Julia's JIT compilation, which on a
function this small is orders of magnitude more than the thing being measured.
"""
function measure(f, reps::Int; trials::Int = 5)
    SINK[] = f()                      # warm up (compile), and keep the result
    best = Inf
    for _ in 1:trials
        GC.gc()
        t0 = time_ns()
        for _ in 1:reps
            SINK[] = f()
        end
        best = min(best, (time_ns() - t0) / reps)
    end
    bytes = @allocated(SINK[] = f())  # one call, after warmup
    return (Float64(best), Float64(bytes))
end

# ------------------------------------------------------------------- formatting

function human_time(ns::Real)
    ns < 1_000      && return @sprintf("%.0f ns", ns)
    ns < 1_000_000  && return @sprintf("%.2f µs", ns / 1e3)
    ns < 1e9        && return @sprintf("%.2f ms", ns / 1e6)
    return @sprintf("%.2f s", ns / 1e9)
end

human_bytes(b::Real) = b < 1024 ? @sprintf("%.0f B", b) :
                       b < 1024^2 ? @sprintf("%.1f KiB", b / 1024) :
                       @sprintf("%.1f MiB", b / 1024^2)

rate(ns::Real) = 1e9 / max(ns, 1e-9)     # calls per second

function human_rate(r::Real)
    r >= 1e6 && return @sprintf("%.2fM/s", r / 1e6)
    r >= 1e3 && return @sprintf("%.1fk/s", r / 1e3)
    return @sprintf("%.0f/s", r)
end

fmt2(x::Real) = @sprintf("%.2f", x)

# ------------------------------------------------------------- scaling control

"""
    control_kernel(x) -> Float64

Pure arithmetic, no allocation, no library call — deliberately dependent
(each step needs the last) so the compiler can't vectorize it away, and sized to
land in the same microsecond neighbourhood as one `Lux.apply`.
"""
function control_kernel(x::Float64)
    a = x
    @inbounds for i in 1:600
        a = sqrt(a + i)
    end
    return a
end

"""
    control_scaling(trials) -> Vector

Run `control_kernel` over the same task counts as the model sweep. This is the
machine's own ceiling for perfectly parallel work: if the control scales and the
model doesn't, the shortfall is the model's, and no amount of `FS_WORKERS` will
recover it.
"""
function control_scaling(trials::Int)
    rows = []
    base = 0.0
    per_task = 200_000
    for k in unique([1; 2; 4; 8; Threads.nthreads()])
        k > Threads.nthreads() && continue
        best = Inf
        for _ in 1:trials
            GC.gc()
            t0 = time_ns()
            @sync for _ in 1:k
                Threads.@spawn begin
                    local acc = 0.0
                    for i in 1:per_task
                        acc += control_kernel(i % 97 + 1.0)
                    end
                    SINK[] = acc
                end
            end
            best = min(best, Float64(time_ns() - t0))
        end
        r = k * per_task / (best / 1e9)
        k == 1 && (base = r)
        push!(rows, (; tasks = k, ops_per_sec = r, speedup = r / base))
    end
    return rows
end

# ------------------------------------------------------------------------- corpus

"""
Write a file of exactly `size` random bytes, in bounded chunks.

Content is random rather than zeros so the classifier sees a realistic input —
and so the filesystem can't cheat with a sparse file, which would make the tail
`seek` unrepresentatively fast.
"""
function write_file(path::AbstractString, size::Int, rng)
    chunk = 1024 * 1024
    open(path, "w") do io
        remaining = size
        while remaining > 0
            n = min(chunk, remaining)
            write(io, rand(rng, UInt8, n))
            remaining -= n
        end
    end
    return path
end

# ------------------------------------------------------------------------- main

function main(argv)
    opts = parse_args(argv)
    modelpath = String(opts["model"])
    isfile(modelpath) || (println(stderr, "model artifact not found: $modelpath"); return 1)

    reps, trials = opts["reps"], opts["trials"]
    batches = [parse(Int, s) for s in split(String(opts["batches"]), ",")]
    sizes   = [parse_size(s) for s in split(String(opts["sizes"]), ",")]

    clf = load_classifier(modelpath)
    println("model             $modelpath")
    println("architecture      $(FEATURE_DIM) → 64 → 16 → 2  (Dense/relu, raw logits)")
    println("julia threads     $(Threads.nthreads())   BLAS threads $(BLAS.get_num_threads())")
    println("timing            min of $trials trials × $reps reps")
    println("=" ^ 72)

    results = Dict{String,Any}()

    # --- 1. inference alone, one file at a time: the number the pipeline pays.
    x1 = rand(Float32, FEATURE_DIM, 1)
    infer_ns, infer_bytes = measure(reps; trials) do
        Lux.apply(clf.model, x1, clf.ps, clf.st)
    end
    println()
    println("INFERENCE (Lux.apply, batch 1 — features already in memory)")
    println("  per call        $(human_time(infer_ns))   $(human_rate(rate(infer_ns)))")
    println("  allocations     $(human_bytes(infer_bytes)) per call")
    results["inference_batch1"] = (; ns = infer_ns, bytes = infer_bytes, per_sec = rate(infer_ns))

    # --- 2. batching: how much of that is per-call overhead rather than math?
    println()
    println("INFERENCE BATCHED (same net, N files per apply)")
    println("     batch      per batch      per file        files/s   speedup")
    batch_rows = []
    for b in batches
        xb = rand(Float32, FEATURE_DIM, b)
        ns, _ = measure(max(1, reps ÷ b); trials) do
            Lux.apply(clf.model, xb, clf.ps, clf.st)
        end
        per_file = ns / b
        @printf("  %8d  %13s  %12s  %13s  %7.1fx\n",
                b, human_time(ns), human_time(per_file),
                human_rate(rate(per_file)), infer_ns / per_file)
        push!(batch_rows, (; batch = b, ns_per_batch = ns, ns_per_file = per_file,
                             files_per_sec = rate(per_file), speedup = infer_ns / per_file))
    end
    results["batched"] = batch_rows
    println("  (a large speedup is headroom a batching stage 1 could claim; the pipeline")
    println("   classifies one file per job today, so it pays the batch-1 row above)")

    # --- 3. feature reads: should be flat in file size (seek, not slurp).
    println()
    println("FEATURE READS (read_features: 16 head + 16 tail bytes, scaled)")
    println("      file size        per call         calls/s     allocations")
    read_rows = []
    dir = mktempdir(; prefix = "fsmodel-")
    try
        rng = MersenneTwister(1234)
        for sz in sizes
            path = write_file(joinpath(dir, "f-$sz.bin"), sz, rng)
            # Fewer reps for the big files: this touches the page cache, and the
            # point is the shape of the curve, not another digit of precision.
            r = max(200, reps ÷ 20)
            ns, bytes = measure(() -> read_features(path), r; trials)
            @printf("  %13s  %14s  %14s  %14s\n",
                    human_bytes(sz), human_time(ns), human_rate(rate(ns)), human_bytes(bytes))
            push!(read_rows, (; size_bytes = sz, ns, bytes, per_sec = rate(ns)))
        end
    finally
        rm(dir; recursive = true, force = true)
    end
    results["read_features"] = read_rows
    flat = length(read_rows) > 1 ?
        maximum(r.ns for r in read_rows) / minimum(r.ns for r in read_rows) : 1.0
    @printf("  spread across a %.0fx size range: %.1fx — %s\n",
            maximum(sizes) / minimum(sizes), flat,
            flat < 3 ? "flat, as designed (it seeks to the tail)" :
                       "NOT flat: something is reading more than 32 bytes")

    # --- 4. classify(): what stage 1 calls, I/O and inference together.
    println()
    println("CLASSIFY (read_features + Lux.apply — one whole stage-1 file)")
    dir2 = mktempdir(; prefix = "fsmodel-")
    classify_ns = 0.0
    try
        path = write_file(joinpath(dir2, "sample.bin"), 64 * 1024, MersenneTwister(7))
        classify_ns, classify_bytes = measure(() -> classify(clf, path), max(200, reps ÷ 20); trials)
        println("  per file        $(human_time(classify_ns))   $(human_rate(rate(classify_ns)))")
        println("  allocations     $(human_bytes(classify_bytes)) per file")
        @printf("  split           %.0f%% feature read, %.0f%% inference\n",
                100 * (classify_ns - infer_ns) / classify_ns, 100 * infer_ns / classify_ns)
        results["classify"] = (; ns = classify_ns, bytes = classify_bytes, per_sec = rate(classify_ns))
    finally
        rm(dir2; recursive = true, force = true)
    end

    # --- 5. thread scaling: does the shared read-only Classifier actually scale?
    if !opts["no-threads"] && Threads.nthreads() > 1
        println()
        println("THREAD SCALING (concurrent Lux.apply on the one shared Classifier)")
        println("     tasks        files/s     per file   speedup   efficiency      GC")
        thread_rows = []
        base = 0.0
        for k in unique([1; 2; 4; 8; Threads.nthreads()])
            k > Threads.nthreads() && continue
            # Work per task is held *constant* as tasks are added, so total work
            # scales with `k`. Splitting a fixed total instead would shrink each
            # task as the pool grows until `@spawn`/`@sync` overhead dominated,
            # and the resulting curve would show a collapse that is the
            # measurement's fault rather than the model's.
            per_task = max(reps, 20_000)
            # Each task gets its own input so we measure the model, not cache
            # line ping-pong on a shared buffer.
            xs = [rand(Float32, FEATURE_DIM, 1) for _ in 1:k]
            best, best_gc = Inf, 0.0
            for _ in 1:trials
                GC.gc()
                # Julia's GC stops the world, so it is the one cost that cannot
                # be parallelized away: measuring its share here is what turns a
                # bad efficiency number into a diagnosis (see the note below).
                gc0 = Base.gc_num().total_time
                t0 = time_ns()
                @sync for t in 1:k
                    Threads.@spawn begin
                        local acc = 0.0f0
                        for _ in 1:per_task
                            y, _ = Lux.apply(clf.model, xs[t], clf.ps, clf.st)
                            acc += y[1]        # consume the result
                        end
                        SINK[] = acc
                    end
                end
                elapsed = Float64(time_ns() - t0)
                if elapsed < best
                    best = elapsed
                    best_gc = Float64(Base.gc_num().total_time - gc0)
                end
            end
            files = k * per_task
            fps = files / (best / 1e9)
            k == 1 && (base = fps)
            @printf("  %8d  %13s  %11s  %7.2fx  %10.0f%%  %5.0f%%\n",
                    k, human_rate(fps), human_time(best / files), fps / base,
                    100 * fps / base / k, 100 * best_gc / best)
            push!(thread_rows, (; tasks = k, files_per_sec = fps,
                                  ns_per_file = best / files, speedup = fps / base,
                                  gc_fraction = best_gc / best))
        end
        results["thread_scaling"] = thread_rows

        # A poor scaling curve has two possible authors — the model or the box —
        # and the table alone can't tell them apart. So run the same sweep on a
        # kernel that is pure arithmetic with no allocation and no library
        # underneath: whatever *it* achieves is this machine's ceiling for
        # embarrassingly parallel work, and the gap between the two curves is
        # the part that belongs to Lux.apply.
        ctrl = control_scaling(trials)
        results["control_scaling"] = ctrl
        top = last(ctrl)
        println("  control         pure-compute kernel, same sweep: " *
                "$(fmt2(top.speedup))x at $(top.tasks) tasks " *
                "($(round(Int, 100 * top.speedup / top.tasks))% efficiency)")
        model_top = last(thread_rows)
        if model_top.speedup < 0.6 * top.speedup
            println("                  → the machine parallelizes; Lux.apply does not. Stage-1")
            println("                    workers past ~4 buy little, whatever FS_WORKERS says.")
        else
            println("                  → inference tracks the machine's own scaling ceiling.")
        end
        gc_top = maximum(r.gc_fraction for r in thread_rows)
        gc_top > 0.15 && println("                  ! GC is $(round(Int, 100 * gc_top))% of the " *
                                 "worst case: apply allocates per call, and\n" *
                                 "                    collection stops every thread.")
    end

    println()
    println("=" ^ 72)
    println("Stage 1's cost per file in bin/bench.jl is this classify() figure plus a")
    println("rename, a log line, and whatever contention the other three pools create.")
    println("A large gap between the two is pipeline overhead, not the model.")

    if opts["json"] !== nothing
        results["meta"] = (; model = modelpath, feature_dim = FEATURE_DIM,
                             julia_threads = Threads.nthreads(),
                             blas_threads = BLAS.get_num_threads(),
                             reps, trials)
        open(String(opts["json"]), "w") do io
            JSON3.write(io, results)
        end
        println("\nwrote $(opts["json"])")
    end
    return 0
end

if abspath(PROGRAM_FILE) == @__FILE__
    exit(main(ARGS))
end