#!/usr/bin/env julia

using CSV
using DataFrames
using Statistics
using Distributions

# --------------------------------------------------------------------------
# CONFIGURATION
#
# Map from CSV file name (as it appears inside each parameter directory) to
# the list of data columns in that file that should be summarized. Files may
# have different columns from one another; the `iter` column is assumed to
# be present in all of them and is NOT listed here.
# --------------------------------------------------------------------------
const LOG_COLS = ["value", "relative_value", "discrepancy", "aux_discrepancy", "joint_discrepancy", "n_spikes"];
const AUX_COLS = ["k", "c1", "c2", "k_dist", "c_dist"]

const FILES_COLUMNS = Dict(
    "aux_radon_fb_log.txt" => LOG_COLS,
    "aux_radon_sliding_fb_log.txt" => LOG_COLS,
    "aux2_radon_fb_log.txt" => LOG_COLS,
    "aux2_radon_sliding_fb_log.txt" => LOG_COLS,
    "aux_radon_fb_aux.csv" => AUX_COLS,
    "aux_radon_sliding_fb_aux.csv" => AUX_COLS,
    "aux2_radon_fb_aux.csv" => AUX_COLS,
    "aux2_radon_sliding_fb_aux.csv" => AUX_COLS,
)

const ITER_COL = :iter
const CONF_LEVEL = 0.95

# --------------------------------------------------------------------------
# Compute the half-width of a 95% (or CONF_LEVEL) confidence interval for
# the mean of `values`, using the t-distribution. Falls back gracefully for
# degenerate sample sizes.
# --------------------------------------------------------------------------
function conf_margin(values::AbstractVector{<:Real})
    n = length(values)
    if n < 2
        return 0.0
    end
    s = std(values)
    if s == 0
        return 0.0
    end
    tdist = TDist(n - 1)
    tcrit = quantile(tdist, 1 - (1 - CONF_LEVEL) / 2)
    return tcrit * s / sqrt(n)
end

# --------------------------------------------------------------------------
# Read one configured file from one parameter directory, if it exists.
# Returns a DataFrame with at least `iter` and the configured columns, or
# `nothing` if the file is missing or unreadable.
# --------------------------------------------------------------------------
function read_dir_file(dir::AbstractString, filename::AbstractString, cols::Vector{String})
    path = joinpath(dir, filename)
    if !isfile(path)
        @warn "Missing file, skipping" path
        return nothing
    end

    local df
    try
        df = CSV.read(path, DataFrame)
    catch e
        @warn "Could not read CSV, skipping" path exception = e
        return nothing
    end

    if !hasproperty(df, ITER_COL)
        @warn "File has no '$(ITER_COL)' column, skipping" path
        return nothing
    end

    present_cols = String[]
    for c in cols
        if hasproperty(df, Symbol(c))
            push!(present_cols, c)
        else
            @warn "Column '$c' not found in file, skipping that column for this directory" path
        end
    end
    if isempty(present_cols)
        return nothing
    end

    keep = vcat([ITER_COL], Symbol.(present_cols))
    return df[:, keep]
end

# --------------------------------------------------------------------------
# Build the summary DataFrame for one configured file, across all `dirs`.
# --------------------------------------------------------------------------
function summarize_file(dirs::Vector{String}, filename::AbstractString, cols::Vector{String})
    parts = DataFrame[]
    for dir in dirs
        d = read_dir_file(dir, filename, cols)
        d === nothing || push!(parts, d)
    end

    if isempty(parts)
        @warn "No data found for file across any directory, skipping summary" filename
        return nothing
    end

    combined = vcat(parts...; cols=:union)
    grouped = groupby(combined, ITER_COL)

    iters = sort(collect(keys(grouped)))
    summary = DataFrame(iter=[k[ITER_COL] for k in iters])

    for col in cols
        colsym = Symbol(col)
        mins = Vector{Union{Missing,Float64}}(undef, length(iters))
        maxs = Vector{Union{Missing,Float64}}(undef, length(iters))
        means = Vector{Union{Missing,Float64}}(undef, length(iters))
        confmin = Vector{Union{Missing,Float64}}(undef, length(iters))
        confmax = Vector{Union{Missing,Float64}}(undef, length(iters))

        for (i, key) in enumerate(iters)
            sub = grouped[key]
            if !hasproperty(sub, colsym)
                mins[i] = maxs[i] = means[i] = confmin[i] = confmax[i] = missing
                continue
            end
            vals = collect(skipmissing(sub[!, colsym]))
            if isempty(vals)
                mins[i] = maxs[i] = means[i] = confmin[i] = confmax[i] = missing
                continue
            end
            m = mean(vals)
            margin = conf_margin(vals)
            mins[i] = minimum(vals)
            maxs[i] = maximum(vals)
            means[i] = m
            confmin[i] = m - margin
            confmax[i] = m + margin
        end

        summary[!, Symbol(col, "_min")] = mins
        summary[!, Symbol(col, "_max")] = maxs
        summary[!, Symbol(col, "_confmax")] = confmax
        summary[!, Symbol(col, "_confmin")] = confmin
        summary[!, Symbol(col, "_mean")] = means
    end

    return summary
end

# --------------------------------------------------------------------------
# Main
# --------------------------------------------------------------------------
function main()
    if isempty(ARGS)
        println(stderr, "Usage: julia summarize.jl DIR1 [DIR2 DIR3 ...]")
        exit(1)
    end

    dirs = String[]
    for d in ARGS
        if isdir(d)
            push!(dirs, d)
        else
            @warn "Not a directory, skipping" d
        end
    end

    if isempty(dirs)
        println(stderr, "No valid directories given.")
        exit(1)
    end

    for (filename, cols) in FILES_COLUMNS
        summary = summarize_file(dirs, filename, cols)
        summary === nothing && continue

        stem = first(splitext(filename))
        outpath = joinpath(pwd(), "$(stem)_summary.csv")
        CSV.write(outpath, summary)
        println("Wrote $outpath ($(nrow(summary)) rows)")
    end
end

main()