#!/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, because 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, and 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