#!/usr/bin/env julia
#############################################################################
## validate_uncertainty.jl
## Uncertainty validation: Posterior Predictive Coverage (PPC)
##
## Following Claassen (2019), this script computes:
## - PPC: What % of expert values fall within 95% posterior predictive interval?
## - Also computes 80% PPC for direct Claassen comparison (he reports 60.3%)
## - Wilson score CI for coverage proportion
## - Breakdown by survey source and decade
##
## Unlike credible interval coverage (which checks θ-CIs), posterior predictive
## coverage simulates what a new expert observation would look like given the
## model's beta likelihood, accounting for both position uncertainty AND
## measurement noise. Well-calibrated models should yield ~95% PPC at 95%.
##
## No model re-run needed: reads θ, γ, and φ from existing chain CSV files.
#############################################################################

using CSV, DataFrames, Statistics, Dates, Printf, Random, JSON

# =============================================================================
# Utility: Wilson score CI for a proportion
# =============================================================================

function wilson_ci(p, n; alpha=0.05)
    if n == 0
        return (lower=NaN, upper=NaN, se=NaN)
    end
    z = 1.96  # For 95% CI
    denominator = 1 + z^2/n
    center = (p + z^2/(2n)) / denominator
    margin = z * sqrt((p*(1-p) + z^2/(4n))/n) / denominator
    se = sqrt(p * (1-p) / n)
    return (lower=center - margin, upper=center + margin, se=se)
end

# =============================================================================
# STEP 0: Find latest model run
# =============================================================================

function find_latest_run(base_dir::String="model_outputs")
    if !isdir(base_dir)
        error("Model outputs directory not found: $base_dir")
    end
    runs = filter(d -> startswith(d, "run_") && isdir(joinpath(base_dir, d)), readdir(base_dir))
    if isempty(runs)
        error("No runs found in $base_dir")
    end
    sort!(runs, rev=true)
    latest = joinpath(base_dir, runs[1])
    println("Using latest run: $latest")
    return latest
end

# =============================================================================
# STEP 1: Load expert_dim.csv from model run data
# =============================================================================

function load_expert_dim(run_dir::String)
    expert_dim_file = joinpath(run_dir, "data", "expert_dim.csv")
    if !isfile(expert_dim_file)
        error("expert_dim.csv not found in $run_dir/data/")
    end
    expert_dim = CSV.read(expert_dim_file, DataFrame)
    println("Loaded expert_dim.csv: $(nrow(expert_dim)) observations")
    println("  Unique rr values: $(length(unique(expert_dim.rr_exp_dim)))")
    println("  Item indices (var_exp_dim): $(sort(unique(expert_dim.var_exp_dim)))")
    println("  Dimensions (dim_idx_exp): $(sort(unique(expert_dim.dim_idx_exp)))")
    return expert_dim
end

# =============================================================================
# STEP 2: Selectively load chain columns
# =============================================================================

function load_chains_selective(run_dir::String, needed_rr::Set{Int}, K::Int)
    """Load only the chain columns we need for posterior predictive checks."""
    chains_dir = joinpath(run_dir, "chains")
    chain_files = sort(filter(f -> endswith(f, ".csv") && startswith(f, "chain_"), readdir(chains_dir)))

    if isempty(chain_files)
        error("No chain files found in $chains_dir")
    end
    println("\nLoading $(length(chain_files)) chain files (selective columns)...")

    # Build the set of column names we need
    needed_cols = Set{String}()

    # theta columns: economic_lr.{rr} and galtan.{rr} (cultural cosmopolitan--traditionalist) for each unique rr
    for rr in needed_rr
        push!(needed_cols, "economic_lr.$rr")
        push!(needed_cols, "galtan.$rr")
    end

    # Item parameters: gamma_exp_intercept.1-K, gamma_exp_slope.1-K
    for k in 1:K
        push!(needed_cols, "gamma_exp_intercept.$k")
        push!(needed_cols, "gamma_exp_slope.$k")
    end

    # Precision parameter
    push!(needed_cols, "phi_exp_dim")

    println("  Need $(length(needed_cols)) columns ($(length(needed_rr)) rr × 2 dims + $(2*K) item params + 1 phi)")

    # Read the header from first chain to identify column indices
    first_chain_path = joinpath(chains_dir, chain_files[1])
    header_line = ""
    open(first_chain_path) do f
        for line in eachline(f)
            if !startswith(line, "#")
                header_line = line
                break
            end
        end
    end

    all_cols = split(header_line, ",")
    col_indices = Int[]
    col_names = String[]

    for (i, col) in enumerate(all_cols)
        if col in needed_cols
            push!(col_indices, i)
            push!(col_names, col)
        end
    end

    println("  Found $(length(col_indices))/$(length(needed_cols)) columns in chains")

    if length(col_indices) < length(needed_cols)
        missing_cols = setdiff(needed_cols, Set(col_names))
        n_missing = length(missing_cols)
        sample = collect(missing_cols)[1:min(5, n_missing)]
        println("  WARNING: Missing columns (showing $( min(5, n_missing))/$n_missing): $sample")
    end

    # Build a type specification for selective reading
    # We'll use CSV.read with select parameter
    select_symbols = Symbol.(col_names)

    all_chains = DataFrame[]
    for (i, cf) in enumerate(chain_files)
        path = joinpath(chains_dir, cf)
        print("  Loading chain $i: $(cf)... ")
        t = @elapsed begin
            chain = CSV.read(path, DataFrame; comment="#", select=select_symbols)
        end
        println("$(nrow(chain)) samples, $(round(t, digits=1))s")
        push!(all_chains, chain)
    end

    combined = vcat(all_chains...)
    println("Combined: $(nrow(combined)) total posterior draws")
    return combined
end

# =============================================================================
# STEP 3: Compute posterior predictive coverage
# =============================================================================

function compute_posterior_predictive_cic(chains::DataFrame, expert_dim::DataFrame;
                                          ci_level::Float64=0.95, seed::Int=42)
    """
    Compute posterior predictive coverage for expert dimension observations.

    For each expert observation n with observed value y_n:
    1. For each posterior draw s:
       - Get theta_s = theta[dim, rr] (on logit scale, but chains store inv_logit)
       - Compute mu_s = invlogit(gamma_intercept[k] + gamma_slope[k] * logit(theta_s))
       - V4 (Beta):          Draw y_pred_s ~ Beta(phi * mu_s, phi * (1 - mu_s))
       - V5 (Beta-Binomial): Draw y_pred_s ~ Beta(phi * K * mu_s, phi * K * (1 - mu_s))
         where K = n_experts for that observation
    2. Compute quantile interval of y_pred draws
    3. Check if y_n falls within interval

    Returns DataFrame with one row per observation plus coverage indicator.
    """
    rng = MersenneTwister(seed)

    alpha_lower = (1 - ci_level) / 2
    alpha_upper = 1 - alpha_lower

    N = nrow(expert_dim)
    S = nrow(chains)  # total posterior draws

    # Detect V5 (Beta-Binomial with K-scaling) by presence of n_experts column
    has_k_scaling = hasproperty(expert_dim, :n_experts)
    if has_k_scaling
        k_vec = expert_dim.n_experts
        println("\nV5 detected: using Beta(phi*K*mu, phi*K*(1-mu)) with per-observation K")
    else
        println("\nV4 detected: using Beta(phi*mu, phi*(1-mu))")
    end

    println("Computing posterior predictive coverage ($(round(Int, 100*ci_level))% level)")
    println("  Expert observations: $N")
    println("  Posterior draws: $S")

    # Pre-extract phi vector
    phi_vec = chains[!, :phi_exp_dim]

    # Pre-extract gamma vectors for each item k
    K = maximum(expert_dim.var_exp_dim)
    gamma_int = Dict{Int, Vector{Float64}}()
    gamma_slope = Dict{Int, Vector{Float64}}()
    for k in 1:K
        col_int = Symbol("gamma_exp_intercept.$k")
        col_slope = Symbol("gamma_exp_slope.$k")
        if hasproperty(chains, col_int) && hasproperty(chains, col_slope)
            gamma_int[k] = chains[!, col_int]
            gamma_slope[k] = chains[!, col_slope]
        end
    end

    # Allocate result columns
    covered = BitVector(undef, N)
    pred_lower = Vector{Float64}(undef, N)
    pred_upper = Vector{Float64}(undef, N)
    pred_median = Vector{Float64}(undef, N)

    # Pre-allocate per-observation draw buffer
    y_pred = Vector{Float64}(undef, S)

    prog_interval = max(1, N ÷ 20)

    for n in 1:N
        if n % prog_interval == 0 || n == N
            pct = round(100 * n / N, digits=1)
            print("\r  Progress: $pct% ($n / $N)")
        end

        rr = expert_dim.rr_exp_dim[n]
        dim = expert_dim.dim_idx_exp[n]
        k = expert_dim.var_exp_dim[n]
        y_obs = expert_dim.val[n]

        # Get theta column (chains store inv_logit(theta), i.e. on [0,1] scale)
        theta_col = dim == 1 ? Symbol("economic_lr.$rr") : Symbol("galtan.$rr")

        if !hasproperty(chains, theta_col) || !haskey(gamma_int, k)
            # Missing chain data — mark as not covered
            covered[n] = false
            pred_lower[n] = NaN
            pred_upper[n] = NaN
            pred_median[n] = NaN
            continue
        end

        theta_star_vec = chains[!, theta_col]  # inv_logit(theta), i.e. on [0,1]
        g_int = gamma_int[k]
        g_slope = gamma_slope[k]

        # Effective concentration: phi for V4, phi * n_experts for V5
        k_mult = has_k_scaling ? Float64(k_vec[n]) : 1.0

        # For each posterior draw, simulate a predictive observation
        for s in 1:S
            theta_star = theta_star_vec[s]

            # Convert back to latent scale for linear predictor
            # theta_star is inv_logit(theta), so theta = logit(theta_star)
            # Clamp to avoid Inf
            theta_star_clamped = clamp(theta_star, 1e-10, 1 - 1e-10)
            theta_latent = log(theta_star_clamped / (1 - theta_star_clamped))

            # Linear predictor
            lin = g_int[s] + g_slope[s] * theta_latent

            # Mean of beta
            mu = 1 / (1 + exp(-lin))
            mu = clamp(mu, 1e-6, 1 - 1e-6)

            # Beta parameters: phi * K * mu for V5, phi * mu for V4
            phi = phi_vec[s] * k_mult
            a = phi * mu
            b = phi * (1 - mu)

            # Draw from Beta(a, b) via gamma method (no Distributions.jl needed)
            y_pred[s] = _rand_beta(rng, a, b)
        end

        # Compute predictive interval
        sort!(y_pred)
        idx_lo = max(1, round(Int, alpha_lower * S))
        idx_hi = min(S, round(Int, alpha_upper * S))
        idx_med = round(Int, 0.5 * S)

        pred_lower[n] = y_pred[idx_lo]
        pred_upper[n] = y_pred[idx_hi]
        pred_median[n] = y_pred[idx_med]
        covered[n] = (y_obs >= pred_lower[n]) && (y_obs <= pred_upper[n])
    end
    println()  # newline after progress

    # Add results to a copy of expert_dim
    result = DataFrame(
        rr = expert_dim.rr_exp_dim,
        dim_idx = expert_dim.dim_idx_exp,
        var_idx = expert_dim.var_exp_dim,
        val = expert_dim.val,
        party = expert_dim.party,
        country = expert_dim.country,
        year = expert_dim.year,
        project = expert_dim.project,
        var = expert_dim.var,
        pred_lower = pred_lower,
        pred_upper = pred_upper,
        pred_median = pred_median,
        covered = covered
    )

    return result
end

# =============================================================================
# Beta random variate without Distributions.jl
# =============================================================================

"""
    _rand_beta(rng, a, b)

Generate a Beta(a, b) random variate using the Gamma method:
Beta(a,b) = X/(X+Y) where X ~ Gamma(a), Y ~ Gamma(b).
Uses Marsaglia & Tsang (2000) for Gamma generation.
"""
function _rand_beta(rng::AbstractRNG, a::Float64, b::Float64)
    x = _rand_gamma(rng, a)
    y = _rand_gamma(rng, b)
    return x / (x + y)
end

"""
    _rand_gamma(rng, shape)

Generate Gamma(shape, 1) random variate using Marsaglia & Tsang (2000).
For shape < 1, uses the rejection method with shape+1 then scales.
"""
function _rand_gamma(rng::AbstractRNG, shape::Float64)
    if shape < 1.0
        # Gamma(a) = Gamma(a+1) * U^(1/a) where U ~ Uniform(0,1)
        return _rand_gamma(rng, shape + 1.0) * rand(rng)^(1.0 / shape)
    end

    # Marsaglia & Tsang (2000) for shape >= 1
    d = shape - 1.0/3.0
    c = 1.0 / sqrt(9.0 * d)

    while true
        local x::Float64
        local v::Float64

        while true
            x = randn(rng)
            v = 1.0 + c * x
            if v > 0.0
                break
            end
        end

        v = v * v * v
        u = rand(rng)

        if u < 1.0 - 0.0331 * x^2 * x^2
            return d * v
        end

        if log(u) < 0.5 * x^2 + d * (1.0 - v + log(v))
            return d * v
        end
    end
end

# =============================================================================
# STEP 4: Summarize and save results
# =============================================================================

function summarize_coverage(result::DataFrame, ci_level::Float64)
    level_pct = round(Int, 100 * ci_level)
    println("\n" * "="^60)
    println("POSTERIOR PREDICTIVE COVERAGE ($level_pct%)")
    println("="^60)

    # Overall by dimension
    dim_names = Dict(1 => "economic_lr", 2 => "galtan")
    display_dim_names = Dict(1 => "economic left-right", 2 => "cultural cosmopolitan--traditionalist")

    summary_rows = []
    for dim in sort(unique(result.dim_idx))
        subset = filter(r -> r.dim_idx == dim, result)
        n = nrow(subset)
        n_covered = sum(subset.covered)
        ppc = n_covered / n
        ci = wilson_ci(ppc, n)
        dim_name = dim_names[dim]

        println(@sprintf("\n  %-40s: %.1f%% [%.1f%%, %.1f%%]  (%d/%d)",
            display_dim_names[dim], 100*ppc, 100*ci.lower, 100*ci.upper, n_covered, n))

        push!(summary_rows, (
            dimension = dim_name,
            cic = ppc,
            cic_pct = round(100 * ppc, digits=1),
            ci_lower = ci.lower,
            ci_upper = ci.upper,
            n = n,
            covered = n_covered
        ))

        # By project
        println("\n  By survey source:")
        by_project = combine(groupby(subset, :project)) do df
            nc = sum(df.covered)
            DataFrame(n = nrow(df), covered = nc, cic = nc / nrow(df))
        end
        sort!(by_project, :n, rev=true)

        @printf("    %-12s %6s %8s\n", "Project", "N", "PPC")
        for row in eachrow(by_project)
            @printf("    %-12s %6d %7.1f%%\n", row.project, row.n, 100*row.cic)
        end

        # By decade
        subset_with_decade = copy(subset)
        subset_with_decade.decade = div.(subset_with_decade.year, 10) .* 10
        println("\n  By decade:")
        by_decade = combine(groupby(subset_with_decade, :decade)) do df
            nc = sum(df.covered)
            DataFrame(n = nrow(df), covered = nc, cic = nc / nrow(df))
        end
        sort!(by_decade, :decade)

        @printf("    %-8s %6s %8s\n", "Decade", "N", "PPC")
        for row in eachrow(by_decade)
            @printf("    %-8d %6d %7.1f%%\n", row.decade, row.n, 100*row.cic)
        end
    end

    return summary_rows
end

function save_results(result_95::DataFrame, summary_95, summary_80,
                      by_project_95::Dict, output_dir::String="validation")
    if !isdir(output_dir)
        mkpath(output_dir)
    end

    timestamp = Dates.format(now(), "yyyy-mm-dd_HH-MM-SS")

    # Summary table (95%)
    if !isempty(summary_95)
        summary_df = DataFrame(summary_95)
        summary_file = joinpath(output_dir, "uncertainty_cic_summary_$timestamp.csv")
        CSV.write(summary_file, summary_df)
        println("\nSaved: $summary_file")
    end

    # Also save 80% summary
    if !isempty(summary_80)
        summary80_df = DataFrame(summary_80)
        summary80_file = joinpath(output_dir, "uncertainty_cic_80pct_summary_$timestamp.csv")
        CSV.write(summary80_file, summary80_df)
        println("Saved: $summary80_file")
    end

    # By-project tables (95%)
    dim_names = Dict(1 => "economic_lr", 2 => "galtan")
    for (dim, bp) in by_project_95
        project_file = joinpath(output_dir, "uncertainty_$(dim_names[dim])_by_project_$timestamp.csv")
        CSV.write(project_file, bp)
        println("Saved: $project_file")
    end

    return summary_95
end

function print_claassen_comparison(summary_95, summary_80)
    println("\n" * "="^60)
    println("COMPARISON WITH CLAASSEN (2019) BENCHMARKS")
    println("="^60)

    println("\nClaassen's result:")
    println("  CIC (80% CI): 60.3%")
    println("  (Using credible intervals for θ, not posterior predictive)")
    println()
    println("Our results (posterior predictive):")
    println()

    println("-"^60)
    @printf("%-15s %10s %10s %8s\n", "Dimension", "PPC 95%", "PPC 80%", "Status")
    println("-"^60)

    dim_map_80 = Dict(r.dimension => r for r in summary_80)

    for r in summary_95
        ppc80 = haskey(dim_map_80, r.dimension) ? dim_map_80[r.dimension].cic : NaN
        # Well-calibrated: 95% PPC should be near 95%
        status = r.cic >= 0.90 ? "GOOD" : (r.cic >= 0.80 ? "OK" : "LOW")
        @printf("%-15s %9.1f%% %9.1f%% %8s\n",
            r.dimension, 100*r.cic, 100*ppc80, status)
    end
    println("-"^60)
    println()
    println("Interpretation:")
    println("  95% PPC ~95% = well-calibrated uncertainty")
    println("  80% PPC > 60% = exceeds Claassen (2019) benchmark")
end

# =============================================================================
# MAIN
# =============================================================================

function main()
    println("="^60)
    println("UNCERTAINTY VALIDATION: Posterior Predictive Coverage")
    println("="^60)
    println("Following Claassen (2019) validation framework")
    println("Posterior predictive intervals account for both position")
    println("uncertainty AND observation-level measurement noise.")
    println()

    # Step 0: Parse options and find run directory
    run_dir = nothing
    quick_mode = get(ENV, "QUICK_VALIDATION", "0") == "1"
    for (i, arg) in enumerate(ARGS)
        if arg == "--run-dir" && i < length(ARGS)
            run_dir = ARGS[i + 1]
        elseif startswith(arg, "--run-dir=")
            run_dir = split(arg, "=", limit=2)[2]
        elseif arg == "--quick"
            quick_mode = true
        end
    end
    if run_dir === nothing
        run_dir = find_latest_run()
    else
        println("Using specified run directory: $run_dir")
    end

    # Step 1: Load expert_dim.csv
    expert_dim = load_expert_dim(run_dir)

    # Step 2: Selectively load chains
    needed_rr = Set(expert_dim.rr_exp_dim)
    K = maximum(expert_dim.var_exp_dim)
    chains = load_chains_selective(run_dir, needed_rr, K)

    # Step 3a: Compute 95% posterior predictive coverage
    result_95 = compute_posterior_predictive_cic(chains, expert_dim; ci_level=0.95)
    summary_95 = summarize_coverage(result_95, 0.95)

    # Step 3b: Compute 80% posterior predictive coverage (Claassen benchmark)
    # Recompute coverage from the same predictive draws but with 80% quantiles
    if quick_mode
        println("\nQUICK MODE: skipping 80% PPC recomputation")
        summary_80 = [(dimension=r.dimension, n=r.n, covered=r.covered, cic=NaN, ci_level=0.80) for r in summary_95]
    else
        println("\n" * "="^60)
        println("RECOMPUTING WITH 80% LEVEL (Claassen comparison)")
        println("="^60)
        result_80 = compute_posterior_predictive_cic(chains, expert_dim; ci_level=0.80, seed=42)
        summary_80 = summarize_coverage(result_80, 0.80)
    end

    # Build by-project tables for 95%
    dim_names = Dict(1 => "economic_lr", 2 => "galtan")
    by_project_95 = Dict{Int, DataFrame}()
    for dim in sort(unique(result_95.dim_idx))
        subset = filter(r -> r.dim_idx == dim, result_95)
        bp = combine(groupby(subset, :project)) do df
            nc = sum(df.covered)
            DataFrame(n = nrow(df), covered = nc, cic = nc / nrow(df))
        end
        sort!(bp, :n, rev=true)
        by_project_95[dim] = bp
    end

    # Step 4: Save results
    save_results(result_95, summary_95, summary_80, by_project_95)

    # Step 5: Print Claassen comparison
    print_claassen_comparison(summary_95, summary_80)

    println("\n" * "="^60)
    println("VALIDATION COMPLETE")
    println("="^60)

    return (summary_95=summary_95, summary_80=summary_80)
end

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