#!/usr/bin/env julia
#############################################################################
## run_model.jl
## Main runner for latent trait model
## Executes the complete pipeline: data loading → preparation → model fitting
##
## Supports two model versions:
##   - "2dim": 2D bipolar model (V1) - estimates economic_lr and galtan directly
##   - "4dim": 4D unipolar model (V10) - estimates 4 traits, derives 2 scales
##
## Default is 2D model (better identification, faster convergence)
#############################################################################

using Dates

#############################################################################
## EXECUTION CONFIGURATION (Change these values as needed)
#############################################################################
const MODEL_VERSION = "2dim"                                # "2dim" (recommended) or "4dim"
const STAN_MODEL_FILE = MODEL_VERSION == "2dim" ? "models/stan_model_2dim_v6.stan" : "models/stan_model_4dim_v10.stan"
const NUM_CHAINS = 4                                        # Number of MCMC chains (run in parallel)
const NUM_WARMUP = 1000                                     # Number of warmup iterations
const NUM_SAMPLES = 2000                                    # Number of sampling iterations
const ADAPT_DELTA = 0.95                                    # Target acceptance probability
const MAX_DEPTH = 15                                        # Maximum tree depth
const START_YEAR = 1944                                     # First year to include (avoid sparse early data)

println("=" ^ 70)
println("Latent Trait Model - Estimation Pipeline")
println("=" ^ 70)
println("Started at: ", Dates.now())
println("Configuration:")
println("  Model version: $(MODEL_VERSION)")
println("  Stan model: $(STAN_MODEL_FILE)")
println("  Chains: $(NUM_CHAINS)")
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 depth: $(MAX_DEPTH)")
println("  Start year: $(START_YEAR)")
if MODEL_VERSION == "2dim"
    println("\n  2D MODEL: Estimates economic_lr and galtan directly")
    println("           (Half the parameters, better convergence)")
else
    println("\n  4D MODEL: Estimates 4 traits, derives 2 scales")
    println("           (Known identification issues - see VERSION_HISTORY.md)")
end
println("=" ^ 70)

# Include pipeline modules
include("pipeline/00_validation.jl")        # Validation checks
include("pipeline/02_data_loading.jl")
include("pipeline/03_data_preparation.jl")
include("pipeline/04_model_execution.jl")
include("pipeline/05_results_processing.jl")

# Load robust save module
include("pipeline/06_save_model.jl")
import .RobustSave: robust_save_model

function run_model(;
    num_chains=NUM_CHAINS,
    num_warmup=NUM_WARMUP,
    num_samples=NUM_SAMPLES,
    adapt_delta=ADAPT_DELTA,
    max_depth=MAX_DEPTH,
    model_file=STAN_MODEL_FILE,
    start_year=START_YEAR,
    data_dir="data"
)
    """Run the complete latent trait model pipeline (2D or 4D based on MODEL_VERSION)"""

    try
        # Step 1: Load and preprocess data
        println("\n" * "="^50)
        println("STEP 1: DATA LOADING")
        println("="^50)

        manifesto, expert_dim, expert_lr, year0, union_to_constituents, constituent_to_union = load_and_preprocess_4dim_data(start_year; data_dir=data_dir)

        # Step 2: Prepare Stan data structure
        println("\n" * "="^50)
        println("STEP 2: DATA PREPARATION")
        println("="^50)

        # Prepare indices and mappings (V4: union-aware)
        data_prep = prepare_4dim_stan_data(manifesto, expert_dim, expert_lr, year0;
                                           union_to_constituents=union_to_constituents,
                                           constituent_to_union=constituent_to_union)

        # Finalize Stan data dictionary (V4: includes constituent arrays)
        final_data = finalize_4dim_stan_data(
            data_prep.manifesto, data_prep.expert_dim, data_prep.expert_lr,
            data_prep.segment_year, data_prep.segment_info,
            data_prep.all_parties, data_prep.all_groups,
            data_prep.group_to_index, year0, data_prep.S, data_prep.J, data_prep.P, data_prep.R,
            data_prep.N_ciy, data_prep.len_theta_ts, data_prep.segment_country_idx,
            data_prep.F, data_prep.segment_family_idx, data_prep.anchor_segment_idx;
            N_const_man_total=data_prep.N_const_man_total,
            n_const_man=data_prep.n_const_man,
            const_offset_man=data_prep.const_offset_man,
            const_rr_man=data_prep.const_rr_man,
            N_const_exp_dim_total=data_prep.N_const_exp_dim_total,
            n_const_exp_dim=data_prep.n_const_exp_dim,
            const_offset_exp_dim=data_prep.const_offset_exp_dim,
            const_rr_exp_dim=data_prep.const_rr_exp_dim,
            N_const_exp_lr_total=data_prep.N_const_exp_lr_total,
            n_const_exp_lr=data_prep.n_const_exp_lr,
            const_offset_exp_lr=data_prep.const_offset_exp_lr,
            const_rr_exp_lr=data_prep.const_rr_exp_lr
        )

        dat_4dim = final_data.dat_4dim

        # Step 3: Validate data BEFORE running Stan
        println("\n" * "="^50)
        println("STEP 3: DATA VALIDATION")
        println("="^50)

        if !validate_stan_data(dat_4dim; verbose=true)
            error("Data validation failed - see errors above")
        end

        estimate_memory_requirements(dat_4dim; verbose=true)

        # Step 4: Create initialization function
        println("\n" * "="^50)
        println("STEP 4: MODEL INITIALIZATION")
        println("="^50)

        # Determine model version for initialization
        model_init_version = MODEL_VERSION == "2dim" ? "v1_2dim" : "v10"

        # Use S (segments) for initialization, not J (parties)
        init_fn = create_init_function(dat_4dim, data_prep.S, data_prep.P,
                                       data_prep.R, final_data.T_year, data_prep.N_ciy;
                                       model_version=model_init_version)

        # Validate initialization values
        println("\nValidating initialization for chain 1...")
        test_init = init_fn()
        if !validate_init_values(test_init; verbose=true)
            error("Initialization validation failed - see errors above")
        end

        # Step 5: Run Stan model
        println("\n" * "="^50)
        println("STEP 5: MODEL EXECUTION")
        println("="^50)

        # Create temp folder for output
        temp_folder = mktempdir()
        println("Temporary folder for Stan output: $temp_folder")

        # Run Stan model
        stanmodel = run_4dim_stan_model(
            dat_4dim, init_fn, temp_folder;
            num_chains=num_chains,
            num_warmup=num_warmup,
            num_samples=num_samples,
            adapt_delta=adapt_delta,
            max_depth=max_depth,
            model_file=model_file
        )

        # Step 6: Results Processing & Diagnostics
        println("\n" * "="^50)
        println("STEP 6: RESULTS PROCESSING & DIAGNOSTICS")
        println("="^50)

        results = extract_model_results_4dim(stanmodel)
        diagnostics = compute_model_diagnostics(stanmodel)

        println("\nCONVERGENCE DIAGNOSTICS:")
        println("  Max R-hat: $(round(diagnostics.max_rhat, digits=4))")
        println("  Mean R-hat: $(round(diagnostics.mean_rhat, digits=4))")
        println("  High R-hat count: $(diagnostics.high_rhat_count)")
        println("  Min ESS: $(round(diagnostics.min_ess, digits=0))")
        println("  Mean ESS: $(round(diagnostics.mean_ess, digits=0))")
        println("  Convergence Status: $(diagnostics.convergence_status)")

        # Step 7: Save results
        println("\n" * "="^50)
        println("STEP 7: SAVING RESULTS")
        println("="^50)
        println("  Using save-local-then-move strategy...")

        # Prepare data for saving (SINGLE copy of stanmodel, not multiple!)
        model_data_to_save = Dict{String, Any}(
            # StanModel object (contains all MCMC samples)
            "stanmodel_object" => stanmodel,

            # Processed diagnostics
            "diagnostics_summary" => diagnostics.diagnostics_summary,
            "data_dict" => dat_4dim,

            # Original data for reference
            "manifesto" => final_data.manifesto,
            "expert_dim" => final_data.expert_dim,
            "expert_lr" => final_data.expert_lr,
            "segment_year" => final_data.segment_year,  # V10: segment-year mapping
            "segment_info" => final_data.segment_info,  # V10: segment metadata (party_id, segment_num, year range)

            # Metadata with convergence info
            "model_info" => Dict(
                "timestamp" => Dates.format(Dates.now(), "yyyy-mm-dd_HH-MM-SS"),
                "max_rhat" => diagnostics.max_rhat,
                "mean_rhat" => diagnostics.mean_rhat,
                "min_ess" => diagnostics.min_ess,
                "mean_ess" => diagnostics.mean_ess,
                "convergence_status" => diagnostics.convergence_status,
                "model_file" => model_file,
                "model_version" => MODEL_VERSION,
                "num_chains" => num_chains,
                "num_warmup" => num_warmup,
                "num_samples" => num_samples,
                "adapt_delta" => adapt_delta,
                "max_depth" => max_depth,
                "year0" => year0,
                "dimensions" => MODEL_VERSION == "2dim" ?
                    ["economic_lr", "galtan"] :
                    ["pro_market", "pro_welfare", "cosmopolitan", "traditional"]
            )
        )

        # Save using CSV-first robust system. Chains have already been secured
        # by model execution; robust_save_model adds metadata/data and verifies.
        save_dir = data_dir != "data" ? joinpath(data_dir, "model_run") : "outputs/model_outputs"
        output_file = robust_save_model(
            stanmodel,
            model_data_to_save,
            save_dir;
            compress=true,       # Ignored by CSV-first save implementation
            keep_local_backups=2  # Ignored; chains are already saved before this step
        )

        # Robust save module already verified everything!

        println("\n" * "="^70)
        println("MODEL EXECUTION COMPLETED SUCCESSFULLY!")
        println("="^70)
        println("   Max R-hat: $(round(diagnostics.max_rhat, digits=4))")
        println("   Mean R-hat: $(round(diagnostics.mean_rhat, digits=4))")
        println("   Convergence: $(diagnostics.convergence_status)")
        println("   Output file: $output_file")

        # Print summary statistics
        println("\nModel Summary ($(MODEL_VERSION == "2dim" ? "2D Direct Bipolar" : "4D Unipolar")):")
        println("  Segments: $(data_prep.S)")
        println("  Parties with valid segments: $(data_prep.J)")
        println("  Countries: $(data_prep.P)")
        println("  Segment-year combinations: $(data_prep.R)")
        println("  Years: $(final_data.T_year)")
        println("  Manifesto observations: $(dat_4dim["N_man"])")
        println("  Expert dimension-specific observations: $(dat_4dim["N_exp_dim"])")
        println("  Expert general L-R observations: $(dat_4dim["N_exp_lr"])")
        if MODEL_VERSION == "2dim"
            println("  Dimensions estimated: 2 (economic_lr, galtan)")
            println("  Theta parameters: $(2 * data_prep.R) (2 × R)")
        else
            println("  Dimensions estimated: 4 (pro_market, pro_welfare, cosmopolitan, traditional)")
            println("  Theta parameters: $(4 * data_prep.R) (4 × R)")
        end
        println("=" ^ 70)

        # Cleanup temp folder after successful save
        println("\nCLEANING UP TEMPORARY FILES...")
        try
            if isdir(temp_folder)
                rm(temp_folder, recursive=true, force=true)
                println("  Removed temporary folder: $temp_folder")
            end
        catch cleanup_error
            println("  Warning: Could not remove temp folder: $cleanup_error")
            println("  (This won't affect your saved results)")
        end

        return true

    catch e
        println("\nERROR in model pipeline: $e")
        println("Stack trace:")
        showerror(stdout, e, catch_backtrace())
        rethrow(e)
    end
end

function main(args=ARGS)
    # Parse --data-dir argument
    data_dir = "data"
    for (i, arg) in enumerate(args)
        if arg == "--data-dir" && i < length(args)
            data_dir = args[i + 1]
        elseif startswith(arg, "--data-dir=")
            data_dir = split(arg, "=", limit=2)[2]
        end
    end

    if data_dir != "."
        println("Using data directory: $data_dir")
    end
    println("Executing $(MODEL_VERSION) latent trait model pipeline...")

    # Check that required data files exist (in data_dir)
    required_data = [joinpath(data_dir, f) for f in ["text_data.csv", "expert.csv", "lr_data.csv"]]
    required_files = vcat(required_data, [STAN_MODEL_FILE])
    missing_files = []

    for file in required_files
        if !isfile(file)
            push!(missing_files, file)
        end
    end

    if !isempty(missing_files)
        println("ERROR: Missing required files:")
        for file in missing_files
            println("  - $file")
        end

        if any(f -> endswith(f, "text_data.csv") || endswith(f, "expert.csv") || endswith(f, "lr_data.csv"), missing_files)
            println("\nTo generate data files, run:")
            println("  bash scripts/01_prepare_data.sh")
        end

        error("Cannot proceed without required files")
    end

    # Run the complete pipeline
    results = run_model(data_dir=data_dir)

    println("\n$(MODEL_VERSION) latent trait model pipeline completed successfully!")
    println("Check outputs/model_outputs/latest/ for chain CSVs and metadata.")
end

# Main execution
if abspath(PROGRAM_FILE) == @__FILE__
    main()
end
