#!/usr/bin/env julia
#############################################################################
## 04_model_execution_4dim.jl
## Stan model compilation and execution for 4D latent trait model
## Based on v9 execution but adapted for four dimensions
#############################################################################

using StanSample
using ProgressMeter
using Dates
using Statistics: mean

const RUN_METRICS_FILENAME = "run_metrics.json"

#############################################################################
## INTERNAL HELPERS
#############################################################################

json_escape_char(c::Char) = c == '"' ? "\\\"" : c == '\\' ? "\\\\" : c == '\n' ? "\\n" :
    c == '\r' ? "\\r" : c == '\t' ? "\\t" : string(c)

function json_escape(str::AbstractString)
    sprint() do io
        for c in str
            print(io, json_escape_char(c))
        end
    end
end

function json_serialize(value)
    if value === nothing
        return "null"
    elseif value isa Bool
        return value ? "true" : "false"
    elseif value isa Integer || value isa AbstractFloat
        return string(value)
    elseif value isa AbstractString
        return string("\"", json_escape(value), "\"")
    elseif value isa AbstractVector
        return "[" * join(json_serialize.(value), ",") * "]"
    elseif value isa Dict
        pairs_serialized = (string("\"", json_escape(String(k)), "\":", json_serialize(v)) for (k, v) in value)
        return string("{", join(pairs_serialized, ","), "}")
    elseif value isa NamedTuple
        return json_serialize(Dict{String,Any}(String(k) => getfield(value, k) for k in keys(value)))
    else
        return json_serialize(string(value))
    end
end

function write_json_file(path::AbstractString, data)
    open(path, "w") do io
        write(io, json_serialize(data))
        write(io, '\n')
    end
end

function safe_write_json(path::AbstractString, data)
    try
        mkpath(dirname(path))
        write_json_file(path, data)
        return true
    catch err
        println("  Warning: Failed to write JSON metrics to $(path): $(err)")
        return false
    end
end

function parse_elapsed_seconds(text::AbstractString)
    if (m = match(r"([0-9eE\.+-]+)\s*seconds", text)) !== nothing
        try
            return parse(Float64, m.captures[1])
        catch
            return nothing
        end
    end
    return nothing
end

function discover_stan_csvs(paths::Vector{String})
    found = String[]
    for root in unique(paths)
        if isdir(root)
            for (dirpath, _, files) in walkdir(root)
                for file in files
                    if endswith(file, ".csv")
                        push!(found, joinpath(dirpath, file))
                    end
                end
            end
        end
    end
    return unique(found)
end

function parse_chain_id_from_path(path::AbstractString, default::Int)
    name = basename(path)
    if (m = match(r"[-_]([0-9]+)\.csv$", name)) !== nothing
        try
            return parse(Int, m.captures[1])
        catch
            return default
        end
    end
    return default
end

function parse_chain_csv_metrics(path::AbstractString; max_depth::Union{Nothing,Int}=nothing)
    step_size = nothing
    warmup_seconds = nothing
    sampling_seconds = nothing
    total_seconds = nothing
    divergences = 0
    treedepth_hits = 0
    iterations_recorded = 0
    header = nothing
    divergent_idx = nothing
    treedepth_idx = nothing
    leapfrog_idx = nothing
    leapfrog_total = 0.0

    try
        open(path, "r") do io
            for raw_line in eachline(io)
                line = strip(raw_line)
                isempty(line) && continue

                if startswith(line, "#")
                    comment = strip(line[2:end])
                    if occursin("Step size", comment) || occursin("step_size__", comment)
                        parts = split(comment, "=")
                        if !isempty(parts)
                            candidate = strip(parts[end])
                            step_size = try
                                parse(Float64, candidate)
                            catch
                                step_size
                            end
                        end
                    elseif occursin("Warm-up", comment)
                        warmup_seconds = parse_elapsed_seconds(comment)
                    elseif occursin("(Sampling)", comment)
                        sampling_seconds = parse_elapsed_seconds(comment)
                    elseif occursin("(Total)", comment)
                        total_seconds = parse_elapsed_seconds(comment)
                    end
                    continue
                end

                if header === nothing
                    header = split(line, ",")
                    header = [strip(col) for col in header]
                    for (idx, col) in enumerate(header)
                        col == "divergent__" && (divergent_idx = idx)
                        col == "treedepth__" && (treedepth_idx = idx)
                        col == "n_leapfrog__" && (leapfrog_idx = idx)
                    end
                else
                    values = split(line, ",")
                    divergent_idx !== nothing && divergent_idx <= length(values) && begin
                        val = try
                            parse(Float64, strip(values[divergent_idx]))
                        catch
                            0.0
                        end
                        divergences += Int(round(val))
                    end

                    treedepth_idx !== nothing && treedepth_idx <= length(values) && begin
                        depth = try
                            parse(Float64, strip(values[treedepth_idx]))
                        catch
                            -Inf
                        end
                        if max_depth !== nothing && depth >= max_depth
                            treedepth_hits += 1
                        end
                    end

                    leapfrog_idx !== nothing && leapfrog_idx <= length(values) && begin
                        leapfrog_total += try
                            parse(Float64, strip(values[leapfrog_idx]))
                        catch
                            0.0
                        end
                    end

                    iterations_recorded += 1
                end
            end
        end
    catch
        return nothing
    end

    return Dict{String,Any}(
        "step_size" => step_size,
        "warmup_seconds" => warmup_seconds,
        "sampling_seconds" => sampling_seconds,
        "total_seconds" => total_seconds,
        "divergences" => divergences,
        "treedepth_hits" => treedepth_hits,
        "iterations_recorded" => iterations_recorded,
        "mean_leapfrog" => iterations_recorded > 0 ? leapfrog_total / iterations_recorded : nothing
    )
end

function collect_run_metrics(csv_paths::Vector{String}; max_depth::Union{Nothing,Int}=nothing)
    per_chain = Vector{Dict{String,Any}}()
    warmup_total = 0.0
    sampling_total = 0.0
    total_total = 0.0
    step_sizes = Float64[]
    leapfrog_means = Float64[]
    divergences_total = 0
    treedepth_total = 0

    warmup_count = 0
    sampling_count = 0
    total_count = 0

    for (idx, path) in enumerate(sort(csv_paths))
        metrics = parse_chain_csv_metrics(path; max_depth=max_depth)
        metrics === nothing && continue

        chain_id = parse_chain_id_from_path(path, idx)
        metrics["chain_id"] = chain_id
        metrics["csv_path"] = path

        if (ws = metrics["warmup_seconds"]) !== nothing
            warmup_total += ws
            warmup_count += 1
        end
        if (ss = metrics["sampling_seconds"]) !== nothing
            sampling_total += ss
            sampling_count += 1
        end
        if (ts = metrics["total_seconds"]) !== nothing
            total_total += ts
            total_count += 1
        end
        if (step = metrics["step_size"]) !== nothing
            push!(step_sizes, step)
        end

        divergences_total += get(metrics, "divergences", 0)
        treedepth_total += get(metrics, "treedepth_hits", 0)
        if (lf = metrics["mean_leapfrog"]) !== nothing
            push!(leapfrog_means, lf)
        end

        push!(per_chain, metrics)
    end

    aggregate = Dict{String,Any}(
        "warmup_seconds" => warmup_count > 0 ? warmup_total : nothing,
        "sampling_seconds" => sampling_count > 0 ? sampling_total : nothing,
        "total_seconds" => total_count > 0 ? total_total : nothing,
        "divergences" => divergences_total,
        "treedepth_hits" => treedepth_total,
        "mean_step_size" => !isempty(step_sizes) ? mean(step_sizes) : nothing,
        "mean_leapfrog" => !isempty(leapfrog_means) ? mean(leapfrog_means) : nothing
    )

    return per_chain, aggregate
end

function summarize_indices(idxs::Vector{Int}, example_fn; max_examples::Int=10)
    count = length(idxs)
    if count == 0
        return Dict("count" => 0, "examples" => Vector{Any}())
    end
    examples = Vector{Any}()
    limit = min(count, max_examples)
    for i in 1:limit
        idx = idxs[i]
        push!(examples, example_fn(idx))
    end
    return Dict("count" => count, "examples" => examples)
end

function collect_boundary_diagnostics(dat_4dim::Dict{String,Any}; epsilon::Float64=1e-6, max_examples::Int=10)
    diagnostics = Dict{String,Any}()

    if haskey(dat_4dim, "sample") && haskey(dat_4dim, "positive")
        sample = dat_4dim["sample"]
        positive = dat_4dim["positive"]
        zero_sample_idx = findall(==(0), sample)
        diagnostics["manifesto_zero_samples"] = summarize_indices(zero_sample_idx, i -> Dict(
            "index" => i,
            "positive" => positive[i],
            "sample" => sample[i]
        ); max_examples=max_examples)

        pos_zero_idx = [i for i in eachindex(sample) if sample[i] > 0 && positive[i] == 0]
        diagnostics["manifesto_all_negative"] = summarize_indices(pos_zero_idx, i -> Dict(
            "index" => i,
            "positive" => positive[i],
            "sample" => sample[i]
        ); max_examples=max_examples)

        pos_full_idx = [i for i in eachindex(sample) if sample[i] > 0 && positive[i] == sample[i]]
        diagnostics["manifesto_all_positive"] = summarize_indices(pos_full_idx, i -> Dict(
            "index" => i,
            "positive" => positive[i],
            "sample" => sample[i]
        ); max_examples=max_examples)
    end

    if haskey(dat_4dim, "val_dim")
        vals = dat_4dim["val_dim"]
        idxs = [i for i in eachindex(vals) if vals[i] <= epsilon || vals[i] >= 1 - epsilon]
        diagnostics["expert_dim_boundary"] = summarize_indices(idxs, i -> Dict(
            "index" => i,
            "value" => vals[i]
        ); max_examples=max_examples)
    end

    if haskey(dat_4dim, "val_lr")
        vals = dat_4dim["val_lr"]
        idxs = [i for i in eachindex(vals) if vals[i] <= epsilon || vals[i] >= 1 - epsilon]
        diagnostics["expert_lr_boundary"] = summarize_indices(idxs, i -> Dict(
            "index" => i,
            "value" => vals[i]
        ); max_examples=max_examples)
    end

    return diagnostics
end

function derive_error_hints(error_message::AbstractString, boundary_info::Dict{String,Any})
    hints = String[]
    if occursin("gamma_lpdf", error_message)
        push!(hints, "Gamma prior hit zero shape; check expert precision priors and boundary diagnostics.")
    end
    if occursin("beta_lpdf", error_message)
        push!(hints, "Beta likelihood saw 0/1; inspect expert boundary diagnostics.")
    end
    if occursin("lkj_corr_cholesky", error_message)
        push!(hints, "Correlation Cholesky invalid; investigate covariance priors and initial values.")
    end

    for (key, info) in boundary_info
        if info isa Dict
            count = get(info, "count", 0)
            if count isa Integer && count > 0
                push!(hints, "Boundary observations detected in $(key): $(count) matches.")
            end
        end
    end

    return unique(hints)
end

function copy_log_files(log_files::Vector{String}, destination::String)
    stored = String[]
    for log_path in log_files
        if isfile(log_path)
            try
                dest = joinpath(destination, basename(log_path))
                cp(log_path, dest; force=true)
                push!(stored, dest)
            catch err
                println("  Warning: Failed to copy Stan log $(log_path): $(err)")
            end
        end
    end
    return stored
end

function archive_run_diagnostics(final_run_dir::AbstractString,
                                 metrics_path::AbstractString,
                                 stan_log_paths::Vector{String})
    diagnostics_dir = joinpath(final_run_dir, "diagnostics")
    mkpath(diagnostics_dir)

    archived_logs = String[]
    if isfile(metrics_path)
        cp(metrics_path, joinpath(diagnostics_dir, basename(metrics_path)); force=true)
    end

    for log_path in stan_log_paths
        if isfile(log_path)
            dest = joinpath(diagnostics_dir, basename(log_path))
            cp(log_path, dest; force=true)
            push!(archived_logs, dest)
        end
    end

    return diagnostics_dir, archived_logs
end

function read_log_tail(log_paths::Vector{String}; bytes::Int=4000)
    tails = Dict{String,String}()
    for path in log_paths
        if isfile(path)
            try
                open(path, "r") do io
                    seekend(io)
                    pos = max(position(io) - bytes, 0)
                    seek(io, pos)
                    content = read(io, String)
                    tails[path] = content
                end
            catch err
                println("  Warning: Failed to read Stan log $(path): $(err)")
            end
        end
    end
    return tails
end

function validate_cmdstan_log_config(log_paths::Vector{String};
                                     num_samples::Int,
                                     num_warmup::Int,
                                     adapt_delta::Real,
                                     max_depth::Int,
                                     refresh::Int)
    if isempty(log_paths)
        println("  Warning: No Stan log files available for command-configuration verification")
        return false
    end

    required = [
        "num_samples = $(num_samples)",
        "num_warmup = $(num_warmup)",
        "delta = $(adapt_delta)",
        "max_depth = $(max_depth)",
        "refresh = $(refresh)"
    ]

    combined = IOBuffer()
    ok = true
    for path in log_paths
        if !isfile(path)
            println("  Warning: Missing Stan log for config verification: $path")
            ok = false
            continue
        end
        text = read(path, String)
        write(combined, text)
        write(combined, '\n')
    end

    combined_text = String(take!(combined))
    missing = [entry for entry in required if !occursin(entry, combined_text)]
    if isempty(missing) && ok
        println("  ✓ CmdStan config verified in Stan logs")
    else
        !isempty(missing) && println("  ✗ CmdStan config mismatch in logs; missing: $(join(missing, ", "))")
        ok = false
    end

    return ok
end

# WINDOWS FIX: Set CMDSTAN_HOME using StanSample's function BEFORE any model creation
if haskey(ENV, "CMDSTAN_HOME") && !isempty(ENV["CMDSTAN_HOME"])
    set_cmdstan_home!(ENV["CMDSTAN_HOME"])
    println("CMDSTAN_HOME configured: $(ENV["CMDSTAN_HOME"])")
end

function compile_4dim_stan_model(model_file)
    """Compile 4D Stan model from file"""

    if !isfile(model_file)
        error("Stan model file not found: $model_file")
    end

    println("Compiling 4D Stan model: $model_file")

    # Read Stan model content
    stan_content = read(model_file, String)

    # Create a local temp directory for compilation (Windows-compatible)
    local_tmpdir = mktempdir()
    println("  Using local compilation directory: $local_tmpdir")

    # Create Stan model using StanSample.jl with local tmpdir
    stanmodel = SampleModel(
        splitext(basename(model_file))[1],  # Use filename without extension as model name
        stan_content,
        local_tmpdir  # Use local temp directory (Windows-compatible)
    )

    println("4D Stan model compiled successfully")

    return stanmodel
end

"""
Create initialization function for either 2D or 4D model.
Dispatches based on model_version parameter.
"""
function create_init_function(dat_4dim, J, P, R, T_year, N_ciy; model_version::String="v1_2dim")
    if startswith(model_version, "v1_2dim")
        return create_2dim_init_function(dat_4dim, J, P, R, T_year, N_ciy)
    else
        return create_4dim_init_function(dat_4dim, J, P, R, T_year, N_ciy; model_version=model_version)
    end
end

"""
Create initialization function for 2D bipolar model (V1).
Returns a function that generates initial values for Stan.
"""
function create_2dim_init_function(dat_4dim, J, P, R, T_year, N_ciy)
    println("Creating 2D initialization function...")

    F = get(dat_4dim, "F", 1)
    K_man = dat_4dim["K_man"]
    K_exp_dim = dat_4dim["K_exp_dim"]
    K_exp_lr = dat_4dim["K_exp_lr"]
    S = get(dat_4dim, "S", J)
    D_lr = get(dat_4dim, "D_lr", 1)

    function init_fn()
        Dict{String, Any}(
            # 2D latent trait parameters - Vector of Vectors for correct JSON
            "theta_ncp" => [zeros(R) for _ in 1:2],              # 2 rows of R elements
            "theta_init_raw" => [zeros(S) for _ in 1:2],         # 2 rows of S elements
            "sigma_theta_init" => ones(2),

            # Three-level hierarchical variance (2D)
            "mu_sigma_global_raw" => fill(-0.8, 2),
            "tau_sigma_country" => fill(0.1, 2),
            "sigma_country_raw" => [zeros(P) for _ in 1:2],
            "tau_sigma_family" => fill(0.05, 2),
            "sigma_family_raw" => [zeros(F) for _ in 1:2],

            # Country-item intercepts (coding convention differences)
            "country_item_raw" => [zeros(K_man) for _ in 1:P],
            "sigma_country_item" => 0.2,

            # Zero-sample parameters
            "alpha_zs" => -1.5,
            "year_effect_raw" => zeros(T_year),
            "sigma_year_effect" => 0.5,
            "segment_zs_raw" => zeros(S),
            "sigma_segment_zs" => 0.3,
            "ciy_zs_raw" => zeros(N_ciy),
            "sigma_ciy_zs" => 0.3,

            # 2D item parameters (simplified from 4D)
            "gamma_man_intercept_raw" => zeros(K_man),
            "gamma_man_loading" => ones(K_man),

            "gamma_exp_intercept_raw" => zeros(K_exp_dim),
            "gamma_exp_slope" => ones(K_exp_dim),

            "gamma_lr_intercept_raw" => zeros(K_exp_lr),
            "gamma_lr_slope" => ones(K_exp_lr),

            # V6: Hierarchical L-R weights (replace simplex)
            "lr_weight_global" => 0.0,
            "lr_country_offset_raw" => zeros(P),
            "lr_source_offset_raw" => zeros(K_exp_lr),
            "lr_decade_offset_raw" => zeros(D_lr),
            "sigma_lr_country" => 0.1,
            "sigma_lr_source" => 0.1,
            "sigma_lr_decade" => 0.1,

            # Precision parameters
            "phi_exp_dim" => 100.0,
            "phi_exp_lr" => 20.0,

            # Scale parameters
            "sigma_intercept_man" => 1.0,
            "sigma_loading_man" => 0.3,
            "sigma_intercept_exp" => 1.0,
            "sigma_slope_exp" => 0.3,
            "sigma_intercept_lr" => 1.0,
            "sigma_slope_lr" => 0.3,

            # Mean parameters
            "mu_lambda_man" => dat_4dim["mn_resp_log_man"],
            "mu_lambda_exp_dim" => dat_4dim["mn_resp_log_exp_dim"],
            "mu_lambda_exp_lr" => dat_4dim["mn_resp_log_exp_lr"]
        )
    end

    println("2D initialization function created")
    return init_fn
end

function create_4dim_init_function(dat_4dim, J, P, R, T_year, N_ciy; model_version::String="v7")
    println("Creating 4D initialization function (model: $model_version)...")

    # Get F from data if available (V8/V9/V10 model)
    F = get(dat_4dim, "F", 1)
    K_man = dat_4dim["K_man"]
    K_exp_dim = dat_4dim["K_exp_dim"]
    K_exp_lr = dat_4dim["K_exp_lr"]

    # V10: Use S (segments) instead of J (parties) for theta_init
    S = get(dat_4dim, "S", J)  # Fall back to J for older models

    function init_fn()
        # Stan matrix[4, R] expects JSON as array of 4 arrays with R elements each
        # Julia Matrix serialization iterates columns, causing dimension mismatch
        # SOLUTION: Use explicit Vector{Vector} to guarantee correct JSON structure
        # V10: theta_init_raw has S rows (segments), not J rows (parties)
        base_init = Dict{String, Any}(
            # 4D latent trait parameters - Vector of Vectors for correct JSON
            "theta_ncp" => [zeros(R) for _ in 1:4],              # 4 rows of R elements
            "theta_init_raw" => [zeros(S) for _ in 1:4],         # 4 rows of S elements (V10: segments)
            "sigma_theta_init" => ones(4),           # SD per dimension
        )

        # Model-specific variance parameters
        if model_version in ["v8", "v9", "v10"]
            # V8/V9/V10: Three-level hierarchical variance (global -> country -> family)
            merge!(base_init, Dict(
                "mu_sigma_global_raw" => fill(-0.8, 4),    # Unconstrained, softplus(-0.8) = 0.37
                "tau_sigma_country" => fill(0.1, 4),       # Country deviation scale
                "sigma_country_raw" => [zeros(P) for _ in 1:4],  # 4 rows of P elements
                "tau_sigma_family" => fill(0.05, 4),       # Family deviation scale
                "sigma_family_raw" => [zeros(F) for _ in 1:4],   # 4 rows of F elements
            ))
        else
            # V7: Simple hierarchical variance
            merge!(base_init, Dict(
                "mu_sigma_theta" => 0.3,                 # Mean variance across dimensions
                "tau_sigma_theta" => 0.05,               # SD of variance (modest)
                "sigma_theta_raw" => zeros(4),           # Non-centered parameters (can be any real)
            ))
        end

        # Common parameters for all model versions
        merge!(base_init, Dict(
            # Strategic deviation parameters
            "ciy_strategic_raw" => zeros(N_ciy),
            "party_strategic_raw" => zeros(dat_4dim["N_man"]),
            "sigma_ciy_strategic" => 0.2,
            "sigma_party_strategic" => 0.1,

            # Zero-sample parameters
            # V10: segment_zs_raw replaces party_zs_raw (S segments instead of J parties)
            "alpha_zs" => -1.5,
            "year_effect_raw" => zeros(T_year),
            "sigma_year_effect" => 0.5,
            # V10 uses segment_zs_raw (S elements), older versions use party_zs_raw (J elements)
            (model_version == "v10" ? "segment_zs_raw" : "party_zs_raw") => zeros(S),
            (model_version == "v10" ? "sigma_segment_zs" : "sigma_party_zs") => 0.3,
            "ciy_zs_raw" => zeros(N_ciy),
            "sigma_ciy_zs" => 0.3,

            # L-R weights
            "lr_weights_raw" => [0.5, -0.5, 0.3, -0.3],  # Informed initialization
            "phi_exp_dim" => 100.0,  # High precision
            "phi_exp_lr" => 20.0,    # Low precision

            # Mean parameters
            "mu_lambda_man" => dat_4dim["mn_resp_log_man"],
            "mu_lambda_exp_dim" => dat_4dim["mn_resp_log_exp_dim"],
            "mu_lambda_exp_lr" => dat_4dim["mn_resp_log_exp_lr"]
        ))

        # V9/V10: Sign-constrained item parameters (separate from NCP structure)
        if model_version in ["v9", "v10"]
            merge!(base_init, Dict(
                # Text data item parameters - V9 sign-constrained
                "gamma_man_intercept_raw" => zeros(K_man),
                "gamma_man_loading_high" => ones(K_man),      # Positive, initialized at 1.0
                "gamma_man_loading_low_abs" => ones(K_man),   # Positive, will be negated in model

                # Expert dimension item parameters - V9 sign-constrained
                "gamma_exp_intercept_raw" => zeros(K_exp_dim),
                "gamma_exp_loading_high" => ones(K_exp_dim),
                "gamma_exp_loading_low_abs" => ones(K_exp_dim),

                # Expert L-R item parameters - V9 sign-constrained
                "gamma_lr_intercept_raw" => zeros(K_exp_lr),
                "gamma_lr_slope" => ones(K_exp_lr),           # Positive slope

                # V9: Scale parameters (simpler structure than V8)
                "sigma_intercept_man" => 1.0,
                "sigma_loading_man" => 0.3,
                "sigma_intercept_exp" => 1.0,
                "sigma_loading_exp" => 0.3,
                "sigma_intercept_lr" => 1.0,
                "sigma_slope_lr" => 0.3
            ))
        else
            # V4-V8: NCP-based item parameters
            merge!(base_init, Dict(
                # Manifesto item parameters - bipolar bridge structure (K_man x 3)
                # Stan matrix[K_man, 3] -> K_man rows of 3 elements
                "Gamma_man_ncp" => [zeros(3) for _ in 1:K_man],

                # Expert dimension-specific item parameters
                "Gamma_exp_dim_ncp" => [zeros(3) for _ in 1:K_exp_dim],

                # Expert general L-R item parameters
                "Gamma_exp_lr_ncp" => [zeros(2) for _ in 1:K_exp_lr],

                # Cholesky factors (for NCP models) - 3x3 identity as 3 rows
                "L_Omega_man" => [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]],
                "L_Omega_dim" => [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]],
                "L_Omega_lr" => [[1.0, 0.0], [0.0, 1.0]],

                # Scale parameters (for NCP models)
                "tau_man" => [1.0, 1.0, 1.0],
                "tau_dim" => [1.0, 1.0, 1.0],
                "tau_lr" => [1.0, 1.0]
            ))
        end

        return base_init
    end

    println("4D initialization function created (model: $model_version)")
    return init_fn
end

function run_4dim_stan_model(dat_4dim, init_fn, temp_folder;
                              num_chains=4, num_warmup=1000, num_samples=1000,
                              adapt_delta=0.95, max_depth=15,
                              num_threads=1,
                              refresh=parse(Int, get(ENV, "STAN_REFRESH", "100")),
                              model_file="stan_model_4dim_traits_only.stan",
                              diagnostic_mode::Bool=false)

    if diagnostic_mode
        if num_chains != 1
            println("Diagnostic mode overriding num_chains from $num_chains to 1")
            num_chains = 1
        end
        if num_threads != 1
            println("Diagnostic mode overriding num_threads from $num_threads to 1")
            num_threads = 1
        end
    end

    available_threads = try
        Sys.CPU_THREADS
    catch
        1
    end

    total_requested = num_chains * num_threads
    if total_requested > available_threads
        proposed = max(1, fld(available_threads, num_chains))
        if proposed < 1
            proposed = 1
        end
        if proposed < num_threads
            println("Requested $(total_requested) workers exceeds available $(available_threads); reducing threads per chain to $(proposed)")
            num_threads = proposed
        end
    end

    println("Setting up 4D Stan model with StanSample.jl ($(model_file))...")
    println("Configuration:")
    println("  Model file: $model_file")
    println("  Chains: $num_chains (parallel execution)")
    println("  Warmup iterations: $num_warmup")
    println("  Sampling iterations: $num_samples")
    println("  Total iterations per chain: $(num_warmup + num_samples)")
    println("  Adapt delta: $adapt_delta")
    println("  Max tree depth: $max_depth")
    println("  Threads per chain: $num_threads")
    println("  Total requested workers: $(num_chains * num_threads) (available: $available_threads)")
    println("  Stan refresh: $refresh")

    # Compile Stan model
    stanmodel = compile_4dim_stan_model(model_file)

    # Create init values for all chains
    init_values = [init_fn() for _ in 1:num_chains]

    # Run the model with built-in progress monitoring
    # PRESERVE STAN TMPDIR before it gets deleted
    stan_tmpdir = stanmodel.tmpdir
    println("Starting 4D Stan sampling...")
    println("Progress shown every $refresh iteration(s)")
    println("  Stan tmpdir: $(stan_tmpdir)")

    # Copy Stan tmpdir path to a file so we can find it later
    backup_dir = joinpath(temp_folder, "stan_backup")
    mkpath(backup_dir)
    open(joinpath(backup_dir, "stan_tmpdir.txt"), "w") do f
        write(f, stan_tmpdir)
    end

    metrics_path = joinpath(temp_folder, RUN_METRICS_FILENAME)
    boundary_info = collect_boundary_diagnostics(dat_4dim)

    metrics_payload = Dict{String,Any}(
        "status" => "initializing",
        "timestamp" => Dates.format(Dates.now(), "yyyy-mm-ddTHH:MM:SS"),
        "config" => Dict(
            "model_file" => model_file,
            "num_chains" => num_chains,
            "num_warmup" => num_warmup,
            "num_samples" => num_samples,
            "adapt_delta" => adapt_delta,
            "max_depth" => max_depth,
            "num_threads" => num_threads,
            "refresh" => refresh,
            "total_workers" => num_chains * num_threads,
            "available_threads" => available_threads
        ),
        "cmdstan_tmpdir" => stan_tmpdir,
        "temp_folder" => temp_folder,
        "csv_paths" => String[],
        "chains" => Any[],
        "aggregate" => Dict{String,Any}(),
        "timing" => Dict{String,Any}(),
        "notes" => Dict{String,Any}(
            "stan_tmpdir_recorded" => isdir(stan_tmpdir)
        ),
        "diagnostic_mode" => diagnostic_mode,
        "data_checks" => boundary_info
    )

    start_time = time()
    run_artifacts = nothing
    stan_log_copies = String[]
    stan_log_tail = Dict{String,String}()

    previous_env = Dict(
        "STAN_NUM_THREADS" => get(ENV, "STAN_NUM_THREADS", nothing),
        "TBB_NUM_THREADS" => get(ENV, "TBB_NUM_THREADS", nothing),
        "OMP_NUM_THREADS" => get(ENV, "OMP_NUM_THREADS", nothing)
    )

    ENV["STAN_NUM_THREADS"] = string(max(1, num_threads))
    ENV["TBB_NUM_THREADS"] = string(max(1, num_threads))
    ENV["OMP_NUM_THREADS"] = string(max(1, num_threads))

    try
        # Run Stan with built-in progress monitoring
        # V5: use_cpp_chains=false (default) spawns 4 separate processes, each reads STAN_NUM_THREADS env var
        stanmodel.num_threads = num_threads

        stan_sample(
            stanmodel,
            data=dat_4dim,
            init=init_values,
            num_chains=num_chains,
            num_warmups=num_warmup,
            num_samples=num_samples,
            delta=adapt_delta,
            engine=:nuts,
            max_depth=max_depth,
            num_threads=num_threads,
            show_logging=true,
            refresh=refresh
        )

        # Calculate total time
        elapsed_seconds = time() - start_time
        total_minutes = elapsed_seconds / 60

        csv_paths = discover_stan_csvs([stan_tmpdir, temp_folder])

        # =================================================================
        # BULLETPROOF: Save chain CSVs to FINAL location IMMEDIATELY
        # This happens BEFORE diagnostics - even if everything else fails,
        # the chain CSVs will be safe in outputs/model_outputs/latest/
        # =================================================================
        timestamp = Dates.format(Dates.now(), "yyyy-mm-dd_HH-MM-SS")
        run_id = "run_$(timestamp)"
        final_run_dir = joinpath("outputs", "model_outputs", "latest", run_id)
        final_chains_dir = joinpath(final_run_dir, "chains")
        mkpath(final_chains_dir)

        println("\n" * "=" ^ 70)
        println("SECURING CHAIN CSV FILES TO FINAL LOCATION")
        println("=" ^ 70)
        println("Destination: $final_chains_dir")

        chain_csv_files = filter(p -> contains(basename(p), "chain"), csv_paths)

        if isempty(chain_csv_files)
            error("CRITICAL: No chain CSV files found in Stan output! Check Stan naming conventions.")
        end

        copied_csv_paths = String[]
        total_size_gb = 0.0

        for (i, csv_path) in enumerate(sort(chain_csv_files))
            src_size = filesize(csv_path)
            dest_filename = "chain_$i.csv"
            dest_path = joinpath(final_chains_dir, dest_filename)

            # Copy with error handling
            try
                cp(csv_path, dest_path, force=true)
            catch copy_err
                error("CRITICAL: Failed to copy $(basename(csv_path)): $copy_err\nSource: $csv_path\nDest: $dest_path")
            end

            # Verify size matches (detect truncation from disk full)
            dest_size = filesize(dest_path)
            if dest_size != src_size
                error("CRITICAL: Size mismatch for $(dest_filename)!\nSource: $src_size bytes\nDest: $dest_size bytes\nPossible disk full or copy error.")
            end

            size_gb = dest_size / (1024^3)
            total_size_gb += size_gb
            push!(copied_csv_paths, dest_path)
            println("  ✓ chain_$i.csv ($(round(size_gb, digits=2)) GB) - SIZE VERIFIED")
        end

        println("=" ^ 70)
        println("✓ $(length(copied_csv_paths)) CHAIN FILES SECURED ($(round(total_size_gb, digits=2)) GB total)")
        println("  Location: $final_chains_dir")
        println("  Status: SAFE - Data will NOT be lost even if subsequent steps fail")
        println("=" ^ 70)

        # Update csv_paths to point to final location
        csv_paths = copied_csv_paths

        chain_metrics, aggregate_metrics = collect_run_metrics(csv_paths; max_depth=max_depth)

        metrics_payload["status"] = "success"
        metrics_payload["completed_at"] = Dates.format(Dates.now(), "yyyy-mm-ddTHH:MM:SS")
        metrics_payload["csv_paths"] = csv_paths
        metrics_payload["chains"] = chain_metrics
        metrics_payload["aggregate"] = aggregate_metrics
        metrics_payload["timing"] = Dict(
            "walltime_seconds" => elapsed_seconds,
            "walltime_minutes" => total_minutes,
            "reported_warmup_seconds" => get(aggregate_metrics, "warmup_seconds", nothing),
            "reported_sampling_seconds" => get(aggregate_metrics, "sampling_seconds", nothing),
            "reported_total_seconds" => get(aggregate_metrics, "total_seconds", nothing)
        )

        if !isempty(stanmodel.log_file)
            stan_log_copies = copy_log_files(stanmodel.log_file, temp_folder)
            metrics_payload["stan_log_files"] = stan_log_copies
            stan_log_tail = read_log_tail(stan_log_copies)
            if !isempty(stan_log_tail)
                metrics_payload["stan_log_tail"] = stan_log_tail
            end
        end

        config_verified = validate_cmdstan_log_config(
            !isempty(stan_log_copies) ? stan_log_copies : stanmodel.log_file;
            num_samples=num_samples,
            num_warmup=num_warmup,
            adapt_delta=adapt_delta,
            max_depth=max_depth,
            refresh=refresh
        )
        metrics_payload["cmdstan_config_verified"] = config_verified
        if !config_verified
            error("CmdStan log configuration verification failed; refusing to treat run as valid")
        end

        safe_write_json(metrics_path, metrics_payload)
        diagnostics_dir, archived_logs = archive_run_diagnostics(
            final_run_dir,
            metrics_path,
            !isempty(stan_log_copies) ? stan_log_copies : stanmodel.log_file
        )
        metrics_payload["diagnostics_dir"] = diagnostics_dir
        metrics_payload["archived_stan_logs"] = archived_logs
        safe_write_json(joinpath(diagnostics_dir, RUN_METRICS_FILENAME), metrics_payload)

        println("\n4D Stan model completed successfully!")
        println("  Total time: $(round(total_minutes, digits=1)) minutes")
        println("  Chains: $num_chains")
        println("  Chain CSV files SECURED in: $final_chains_dir")
        if (warmup_seconds = get(aggregate_metrics, "warmup_seconds", nothing)) !== nothing
            println("  Reported warmup time: $(round(warmup_seconds / 60, digits=2)) minutes")
        end
        if (sampling_seconds = get(aggregate_metrics, "sampling_seconds", nothing)) !== nothing
            println("  Reported sampling time: $(round(sampling_seconds / 60, digits=2)) minutes")
        end
        println("  Divergences (post-warmup): $(get(aggregate_metrics, "divergences", 0))")
        println("  Tree depth hits (== $max_depth): $(get(aggregate_metrics, "treedepth_hits", 0))")
        println("  Metrics artifact: $(metrics_path)")
        println("  Archived diagnostics: $(diagnostics_dir)")

        run_artifacts = (
            stanmodel=stanmodel,
            metrics=metrics_payload,
            metrics_path=metrics_path,
            csv_paths=csv_paths,
            aggregate=aggregate_metrics,
            stan_logs=stan_log_copies,
            stan_log_tail=stan_log_tail,
            final_run_dir=final_run_dir,  # Directory where chains are ALREADY saved
            num_chains_saved=length(copied_csv_paths)
        )

        # Print parameter count estimate
        K_man = dat_4dim["K_man"]
        K_exp_dim = dat_4dim["K_exp_dim"]
        K_exp_lr = dat_4dim["K_exp_lr"]
        R = dat_4dim["R"]
        J = dat_4dim["J"]
        P = dat_4dim["P"]
        T_year = dat_4dim["T_year"]
        N_ciy = dat_4dim["N_ciy"]
        N_man = dat_4dim["N_man"]

        # Estimate parameter count for 4D model
        theta_params = 4 * R  # 4 dimensions ? R party-years
        theta_init_params = 4 * J  # 4 dimensions ? J parties
        item_params = 4 * K_man * 2 + K_exp_dim * 3 + K_exp_lr * 2  # Item parameters
        strategic_params = P * K_man  # Country-item intercepts
        zero_sample_params = T_year + J + N_ciy  # Zero-sample components
        other_params = 50  # Approximate count of other parameters

        total_estimated_params = (theta_params + theta_init_params + item_params +
                                strategic_params + zero_sample_params + other_params)

        println("  Estimated parameters: ~$(total_estimated_params)")

    catch e
        println("\n" * "="^70)
        println("? STAN SAMPLING FAILED")
        println("="^70)
        println("Error type: $(typeof(e))")
        println("Error message: $e")

        println("\n?? COPYING Stan files before cleanup...")
        println("  Stan tmpdir: $stan_tmpdir")

        backup_dir = joinpath(temp_folder, "stan_backup")
        error_log_path = joinpath(backup_dir, "stan_error_capture.txt")
        stan_errors = String[]

        if isdir(stan_tmpdir)
            println("  Copying files from Stan tmpdir...")
            for file in readdir(stan_tmpdir, join=true)
                try
                    dest = joinpath(backup_dir, basename(file))
                    cp(file, dest, force=true)
                    println("    Copied: $(basename(file))")
                catch copy_err
                    println("    Failed to copy $(basename(file)): $copy_err")
                end
            end
            println("  Files preserved in: $backup_dir")
        else
            println("  ??  Stan tmpdir already deleted: $stan_tmpdir")
        end

        csv_paths = discover_stan_csvs([stan_tmpdir, temp_folder, backup_dir])
        chain_metrics, aggregate_metrics = collect_run_metrics(csv_paths; max_depth=max_depth)

        failure_elapsed = time() - start_time

        metrics_payload["status"] = "error"
        metrics_payload["completed_at"] = Dates.format(Dates.now(), "yyyy-mm-ddTHH:MM:SS")
        metrics_payload["error"] = Dict(
            "type" => string(typeof(e)),
            "message" => sprint(showerror, e)
        )
        metrics_payload["error_hints"] = derive_error_hints(string(e), boundary_info)
        metrics_payload["csv_paths"] = csv_paths
        metrics_payload["chains"] = chain_metrics
        metrics_payload["aggregate"] = aggregate_metrics
        metrics_payload["timing"] = Dict(
            "walltime_seconds" => failure_elapsed,
            "walltime_minutes" => failure_elapsed / 60,
            "reported_warmup_seconds" => get(aggregate_metrics, "warmup_seconds", nothing),
            "reported_sampling_seconds" => get(aggregate_metrics, "sampling_seconds", nothing),
            "reported_total_seconds" => get(aggregate_metrics, "total_seconds", nothing)
        )

        if !isempty(stanmodel.log_file)
            stan_log_copies = copy_log_files(stanmodel.log_file, backup_dir)
            metrics_payload["stan_log_files"] = stan_log_copies
            stan_log_tail = read_log_tail(stan_log_copies)
            if !isempty(stan_log_tail)
                metrics_payload["stan_log_tail"] = stan_log_tail
            end
        end

        safe_write_json(metrics_path, metrics_payload)

        try
            if isfile(metrics_path)
                cp(metrics_path, joinpath(backup_dir, RUN_METRICS_FILENAME); force=true)
            end
        catch copy_metrics_err
            println("  Warning: Failed to copy run metrics artifact: $(copy_metrics_err)")
        end

        println("\n?? Analyzing copied Stan files...")

        if isdir(backup_dir)
            all_files = readdir(backup_dir, join=true)
            csv_files = filter(f -> endswith(f, ".csv"), all_files)
            log_files = filter(f -> endswith(f, ".log"), all_files)
            txt_files = filter(f -> endswith(f, ".txt") || endswith(f, ".out") || endswith(f, ".err"), all_files)

            println("  Files in temp folder:")
            println("    CSV files: $(length(csv_files))")
            println("    Log files: $(length(log_files))")
            println("    Text files: $(length(txt_files))")
            println("    Total files: $(length(all_files))")

            for file in vcat(csv_files, log_files, txt_files)
                try
                    content = read(file, String)
                    if occursin(r"Exception|Error|Syntax error|Failed|Segmentation fault|basic_ios"i, content)
                        println("\n  ??  STAN ERROR FOUND in $(basename(file)):")
                        println("  " * "="^68)
                        lines = split(content, "\n")
                        for (i, line) in enumerate(lines)
                            if occursin(r"Exception|Error|Syntax|Failed|Segmentation|basic_ios"i, line)
                                start_idx = max(1, i-2)
                                end_idx = min(length(lines), i+5)
                                error_context = join(lines[start_idx:end_idx], "\n  ")
                                println("  $error_context")
                                println("  " * "-"^68)
                                push!(stan_errors, error_context)
                            end
                        end
                    end
                catch
                end
            end

            if isempty(stan_errors)
                println("\n  No Stan error messages found in output files")
                println("  This suggests the chains crashed before writing any output")
            end

            if !isempty(stan_errors)
                try
                    open(error_log_path, "w") do f
                        write(f, "Stan Error Capture - $(Dates.now())\n")
                        write(f, "="^70 * "\n\n")
                        for (i, err) in enumerate(stan_errors)
                            write(f, "Error $i:\n")
                            write(f, err * "\n\n")
                        end
                    end
                    println("\n?? Error details saved to: $error_log_path")
                catch
                    println("\n??  Could not save error log")
                end
            end
        else
            println("  ??  Stan tmpdir not found or already cleaned up")
            println("  Checking user temp folder: $temp_folder")
            if isdir(temp_folder)
                all_files = readdir(temp_folder, join=true)
                println("  Files in user temp folder: $(length(all_files))")
                for f in all_files
                    println("    - $(basename(f))")
                end
            end
        end

        println("\n?? DIAGNOSTIC SUGGESTIONS:")
        println("  1. Check data validation output above for NaN/Inf/boundary issues")
        println("  2. Check initialization validation output for constraint violations")
        println("  3. Review boundary diagnostics embedded in run_metrics.json")
        println("  4. Stan tmpdir: $stan_tmpdir (may be cleaned up)")
        println("\n?? Temp folder: $temp_folder")
        if !isempty(stan_errors)
            println("?? Stan errors captured in: $error_log_path")
        end

        run_artifacts = (
            stanmodel=stanmodel,
            metrics=metrics_payload,
            metrics_path=metrics_path,
            csv_paths=csv_paths,
            aggregate=aggregate_metrics,
            stan_logs=stan_log_copies,
            stan_log_tail=stan_log_tail
        )

        error("4D Stan sampling failed: $e")
    finally
        for (key, val) in previous_env
            if val === nothing
                if haskey(ENV, key)
                    delete!(ENV, key)
                end
            else
                ENV[key] = val
            end
        end
    end

    return run_artifacts
end

# Execute if run directly
if abspath(PROGRAM_FILE) == @__FILE__
    println("Run from main 4D script to execute the full pipeline")
end
