#!/usr/bin/env julia
#############################################################################
## 05_results_processing.jl
## Extract and process 4D model results with diagnostics
## Adapted from old_project for latent traits only (no election effects)
#############################################################################

using StanSample, DataFrames, Statistics

function extract_model_results_4dim(stanmodel)
    """
    Extract model results for 4D latent trait model
    Simplified version - no election effects (pure latent traits)
    """
    println("Extracting 4D model results...")

    try
        println("Model completed successfully - extracting results")

        # For 4D latent trait model, we save the full stanmodel object
        # Post-estimation will extract specific parameters later

        return (
            samples = stanmodel,  # Save the full stanmodel with all MCMC samples
            extraction_status = "success"
        )

    catch e
        println("Error in result extraction: $e")
        return (
            samples = "error",
            extraction_status = "error",
            error_message = string(e)
        )
    end
end

function compute_model_diagnostics(stanmodel_result)
    """
    Compute convergence diagnostics from Stan model
    Returns R-hat, ESS statistics, and overall convergence assessment

    stanmodel_result can be either:
    - A SampleModel object directly
    - A named tuple from run_4dim_stan_model containing .stanmodel
    """
    println("Computing model diagnostics...")

    try
        # Handle both direct SampleModel and named tuple from run_4dim_stan_model
        stanmodel = if hasproperty(stanmodel_result, :stanmodel)
            stanmodel_result.stanmodel
        else
            stanmodel_result
        end

        # Get REAL diagnostics using StanSample.jl
        diagnostics_summary = read_summary(stanmodel)

        # Extract real Rhat and ESS values. Stan summary column names differ
        # across CmdStan/StanSample versions, so resolve aliases explicitly.
        summary_names = names(diagnostics_summary)
        rhat_col = if "r_hat" in summary_names
            "r_hat"
        elseif "R_hat" in summary_names
            "R_hat"
        elseif "RHat" in summary_names
            "RHat"
        else
            error("No R-hat column found in summary. Columns: $(join(summary_names, ", "))")
        end
        ess_col = if "ess_bulk" in summary_names
            "ess_bulk"
        elseif "ess" in summary_names
            "ess"
        elseif "ESS_bulk" in summary_names
            "ESS_bulk"
        elseif "n_eff" in summary_names
            "n_eff"
        else
            error("No ESS column found in summary. Columns: $(join(summary_names, ", "))")
        end

        rhat_vals = diagnostics_summary[!, rhat_col]
        ess_bulk_vals = diagnostics_summary[!, ess_col]

        # Compute real statistics (handle NaN values properly)
        # Use isfinite to exclude both missing and NaN values
        valid_rhat = filter(isfinite, rhat_vals)
        valid_ess = filter(isfinite, ess_bulk_vals)

        mean_rhat = length(valid_rhat) > 0 ? mean(valid_rhat) : NaN
        max_rhat = length(valid_rhat) > 0 ? maximum(valid_rhat) : NaN
        mean_ess = length(valid_ess) > 0 ? mean(valid_ess) : NaN
        min_ess = length(valid_ess) > 0 ? minimum(valid_ess) : NaN

        # Count problematic parameters (use isfinite for consistent counting)
        high_rhat_count = count(x -> isfinite(x) && x > 1.1, rhat_vals)
        moderate_rhat_count = count(x -> isfinite(x) && x > 1.05, rhat_vals)
        low_ess_count = count(x -> isfinite(x) && x < 400, ess_bulk_vals)
        very_low_ess_count = count(x -> isfinite(x) && x < 100, ess_bulk_vals)

        # Total parameter count
        total_params = length(valid_rhat)

        # Overall assessment (handle NaN values)
        if isnan(max_rhat) || total_params == 0
            convergence_status = "insufficient_data"
        else
            excellent_convergence = max_rhat < 1.05 && high_rhat_count == 0 && very_low_ess_count == 0
            good_convergence = max_rhat < 1.1 && high_rhat_count < 5 && very_low_ess_count < total_params * 0.1
            acceptable_convergence = max_rhat < 1.2 && high_rhat_count < total_params * 0.1

            if excellent_convergence
                convergence_status = "excellent"
            elseif good_convergence
                convergence_status = "good"
            elseif acceptable_convergence
                convergence_status = "acceptable"
            else
                convergence_status = "poor"
            end
        end

        println("\nDiagnostics computed:")
        println("  Total parameters: $total_params")
        println("  Mean R-hat: $(round(mean_rhat, digits=4))")
        println("  Max R-hat: $(round(max_rhat, digits=4))")
        println("  High R-hat count (>1.1): $high_rhat_count")
        println("  Mean ESS: $(round(mean_ess, digits=0))")
        println("  Min ESS: $(round(min_ess, digits=0))")
        println("  Very low ESS count (<100): $very_low_ess_count")
        println("  Convergence status: $convergence_status")

        return (
            diagnostics_summary = diagnostics_summary,
            convergence_status = convergence_status,
            mean_rhat = mean_rhat,
            max_rhat = max_rhat,
            mean_ess = mean_ess,
            min_ess = min_ess,
            high_rhat_count = high_rhat_count,
            moderate_rhat_count = moderate_rhat_count,
            low_ess_count = low_ess_count,
            very_low_ess_count = very_low_ess_count,
            total_params = total_params
        )

    catch e
        println("Error in diagnostics computation: $e")
        println("Stack trace:")
        showerror(stdout, e, catch_backtrace())

        return (
            diagnostics_summary = "error",
            convergence_status = "error",
            mean_rhat = 999.0,
            max_rhat = 999.0,
            mean_ess = 0.0,
            min_ess = 0.0,
            high_rhat_count = 999,
            moderate_rhat_count = 999,
            low_ess_count = 999,
            very_low_ess_count = 999,
            total_params = 0,
            error_message = string(e)
        )
    end
end

# Execute if run directly
if abspath(PROGRAM_FILE) == @__FILE__
    println("Run from main run_model.jl to execute the full pipeline")
end
