#!/usr/bin/env julia
#############################################################################
## performance_monitoring.jl
## Post-run performance diagnostics for Stan sampling jobs
#############################################################################

using StanSample
using Statistics: mean, median
using Dates

function _safe_column(df, candidates::Vector{String})
    for candidate in candidates
        sym = Symbol(candidate)
        if sym in names(df)
            return df[!, sym]
        elseif candidate in names(df)
            return df[!, candidate]
        end
    end
    return nothing
end

function _clean_values(vec)
    cleaned = Float64[]
    for v in vec
        if v isa Missing || v === nothing
            continue
        end
        try
            value = Float64(v)
            isfinite(value) && push!(cleaned, value)
        catch
            continue
        end
    end
    return cleaned
end

function _summarize_vector(vec)
    if vec === nothing
        return Dict{String,Any}("available" => false)
    end
    cleaned = _clean_values(vec)
    if isempty(cleaned)
        return Dict{String,Any}("available" => false)
    end
    return Dict{String,Any}(
        "available" => true,
        "count" => length(cleaned),
        "mean" => mean(cleaned),
        "median" => median(cleaned),
        "min" => minimum(cleaned),
        "max" => maximum(cleaned)
    )
end

function monitor_sampling_performance!(stanmodel;
    run_metrics::Union{Nothing,Dict{String,Any}}=nothing,
    metrics_path::Union{Nothing,String}=nothing,
    csv_paths::Union{Nothing,Vector{String}}=nothing,
    aggregate_metrics::Union{Nothing,Dict{String,Any}}=nothing,
    max_depth::Union{Nothing,Int}=nothing)

    csv_paths === nothing && (csv_paths = discover_stan_csvs([stanmodel.tmpdir]))
    aggregate_metrics === nothing && begin
        _, aggregate_metrics = collect_run_metrics(csv_paths; max_depth=max_depth)
    end

    summary_df = nothing
    try
        summary_df = read_summary(stanmodel)
    catch e
        println("Warning: could not read Stan summary: $e")
    end

    performance = Dict{String,Any}(
        "generated_at" => Dates.format(Dates.now(), "yyyy-mm-ddTHH:MM:SS"),
        "csv_paths" => csv_paths
    )

    if summary_df !== nothing
        performance["ess_bulk"] = _summarize_vector(_safe_column(summary_df, ["ess_bulk", "ess"]))
        performance["ess_tail"] = _summarize_vector(_safe_column(summary_df, ["ess_tail"]))
        performance["ess_per_sec"] = _summarize_vector(_safe_column(summary_df, ["ess_per_sec", "n_eff/s"]))
        performance["r_hat"] = _summarize_vector(_safe_column(summary_df, ["r_hat"]))

        if haskey(performance["r_hat"], "available") && performance["r_hat"]["available"]
            performance["r_hat"]["max"] = maximum(_clean_values(_safe_column(summary_df, ["r_hat"])))
        end

        performance["parameters_considered"] = size(summary_df, 1)
    else
        performance["ess_bulk"] = Dict{String,Any}("available" => false)
        performance["ess_tail"] = Dict{String,Any}("available" => false)
        performance["ess_per_sec"] = Dict{String,Any}("available" => false)
        performance["r_hat"] = Dict{String,Any}("available" => false)
        performance["parameters_considered"] = 0
    end

    divergences = get(aggregate_metrics, "divergences", 0)
    total_samples = stanmodel.num_samples * stanmodel.num_chains
    divergence_rate = total_samples > 0 ? divergences / total_samples : nothing

    performance["divergences"] = Dict{String,Any}(
        "count" => divergences,
        "rate" => divergence_rate,
        "total_draws" => total_samples
    )

    performance["leapfrog"] = Dict{String,Any}(
        "mean" => get(aggregate_metrics, "mean_leapfrog", nothing)
    )

    performance["step_size"] = Dict{String,Any}(
        "mean" => get(aggregate_metrics, "mean_step_size", nothing)
    )

    sampling_seconds = get(aggregate_metrics, "sampling_seconds", nothing)
    if sampling_seconds !== nothing && sampling_seconds > 0
        performance["throughput"] = Dict{String,Any}(
            "samples_per_second" => (total_samples / sampling_seconds),
            "seconds_sampling" => sampling_seconds
        )
    else
        performance["throughput"] = Dict{String,Any}(
            "samples_per_second" => nothing,
            "seconds_sampling" => sampling_seconds
        )
    end

    if run_metrics !== nothing
        run_metrics["performance"] = performance
        if metrics_path !== nothing
            safe_write_json(metrics_path, run_metrics)
        end
    elseif metrics_path !== nothing
        temp_metrics = Dict{String,Any}("performance" => performance)
        safe_write_json(metrics_path, temp_metrics)
    end

    println("\nPERFORMANCE SUMMARY")
    println("  ESS bulk (mean): $(get(performance["ess_bulk"], "mean", "n/a"))")
    println("  ESS/sec (mean): $(get(performance["ess_per_sec"], "mean", "n/a"))")
    println("  Divergences: $(divergences)")
    println("  Divergence rate: $(divergence_rate === nothing ? "n/a" : round(divergence_rate, digits=6))")
    println("  Mean leapfrog steps: $(get(performance["leapfrog"], "mean", "n/a"))")

    return performance
end

