#!/usr/bin/env julia
#############################################################################
## validate_external.jl
## Out-of-sample validation via held-out expert observations
##
## Design:
##   - Text data stays 100% intact (same parties, segments, indices)
##   - 20% of expert/LR observations held out (stratified by source)
##   - Expert-only parties (no text data) are never held out
##   - Model trains on 80% expert + 100% text
##   - Held-out expert ratings compared to model predictions
##
## Usage:
##   julia scripts/validate_external.jl prepare
##   julia 01_run_model.jl --data-dir validation/external_split/
##   julia 02_post_estimation.jl  # on training run
##   julia scripts/validate_external.jl compute <model_positions.csv>
#############################################################################

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

const HOLDOUT_FRAC = 0.20
const SEED = 42

# =========================================================================
# STEP 1: Prepare train/test split
# =========================================================================

function prepare_holdout_data(base_dir::String=".")
    println("="^70)
    println("PREPARING OUT-OF-SAMPLE VALIDATION SPLIT")
    println("="^70)
    println()

    # Load data
    text_data = CSV.read(joinpath(base_dir, "text_data.csv"), DataFrame)
    expert = CSV.read(joinpath(base_dir, "expert.csv"), DataFrame)
    lr_data = CSV.read(joinpath(base_dir, "lr_data.csv"), DataFrame)

    println("Full dataset:")
    println("  text_data: $(nrow(text_data)) rows, $(length(unique(text_data.party))) parties")
    println("  expert:    $(nrow(expert)) rows, $(length(unique(expert.party))) parties")
    println("  lr_data:   $(nrow(lr_data)) rows, $(length(unique(lr_data.party))) parties")

    # Identify expert-only parties (no text data) — these are NEVER held out
    text_parties = Set(unique(text_data.party))
    expert_only_parties = Set(p for p in unique(vcat(expert.party, lr_data.party))
                              if !(p in text_parties))

    println()
    println("Expert-only parties (protected from holdout): $(length(expert_only_parties))")

    # Split expert data: stratified by source variable
    # Safeguard: ensure each party keeps at least one observation in training
    Random.seed!(SEED)
    expert.row_id = 1:nrow(expert)
    expert.is_holdout = falses(nrow(expert))

    for var_group in groupby(expert, :var)
        var_name = first(var_group.var)
        eligible = findall(row -> !(row.party in expert_only_parties), eachrow(var_group))

        n_holdout = round(Int, length(eligible) * HOLDOUT_FRAC)
        holdout_candidates = shuffle(eligible)

        # Track per-party counts to ensure at least 1 stays in training
        party_train_count = Dict{Int, Int}()
        for idx in eligible
            p = var_group.party[idx]
            party_train_count[p] = get(party_train_count, p, 0) + 1
        end

        n_held = 0
        for idx in holdout_candidates
            n_held >= n_holdout && break
            p = var_group.party[idx]
            if party_train_count[p] > 1  # keep at least 1 in training
                row_id = var_group.row_id[idx]
                expert.is_holdout[row_id] = true
                party_train_count[p] -= 1
                n_held += 1
            end
        end
    end

    # Split LR data: stratified by source variable (same safeguard)
    lr_data.row_id = 1:nrow(lr_data)
    lr_data.is_holdout = falses(nrow(lr_data))

    for var_group in groupby(lr_data, :var)
        var_name = first(var_group.var)
        eligible = findall(row -> !(row.party in expert_only_parties), eachrow(var_group))

        n_holdout = round(Int, length(eligible) * HOLDOUT_FRAC)
        holdout_candidates = shuffle(eligible)

        party_train_count = Dict{Int, Int}()
        for idx in eligible
            p = var_group.party[idx]
            party_train_count[p] = get(party_train_count, p, 0) + 1
        end

        n_held = 0
        for idx in holdout_candidates
            n_held >= n_holdout && break
            p = var_group.party[idx]
            if party_train_count[p] > 1
                row_id = var_group.row_id[idx]
                lr_data.is_holdout[row_id] = true
                party_train_count[p] -= 1
                n_held += 1
            end
        end
    end

    # Create train/test splits
    expert_train = expert[.!expert.is_holdout, Not([:row_id, :is_holdout])]
    expert_test = expert[expert.is_holdout, Not([:row_id, :is_holdout])]
    lr_train = lr_data[.!lr_data.is_holdout, Not([:row_id, :is_holdout])]
    lr_test = lr_data[lr_data.is_holdout, Not([:row_id, :is_holdout])]

    # Report split
    println()
    println("Split summary ($(round(100*HOLDOUT_FRAC))% holdout):")
    println("  expert train: $(nrow(expert_train)) rows ($(round(100*nrow(expert_train)/nrow(expert), digits=1))%)")
    println("  expert test:  $(nrow(expert_test)) rows ($(round(100*nrow(expert_test)/nrow(expert), digits=1))%)")
    println("  lr train:     $(nrow(lr_train)) rows ($(round(100*nrow(lr_train)/nrow(lr_data), digits=1))%)")
    println("  lr test:      $(nrow(lr_test)) rows ($(round(100*nrow(lr_test)/nrow(lr_data), digits=1))%)")

    # Report per-source breakdown
    println()
    println("Per-source breakdown (expert):")
    for var_name in sort(unique(expert.var))
        n_full = count(expert.var .== var_name)
        n_test = count(expert_test.var .== var_name)
        println("  $var_name: $(n_full - n_test) train / $n_test test")
    end
    println()
    println("Per-source breakdown (LR):")
    for var_name in sort(unique(lr_data.var))
        n_full = count(lr_data.var .== var_name)
        n_test = count(lr_test.var .== var_name)
        println("  $var_name: $(n_full - n_test) train / $n_test test")
    end

    # Verify: training set has same parties as full set
    train_parties_expert = Set(unique(expert_train.party))
    train_parties_lr = Set(unique(lr_train.party))
    full_parties_expert = Set(unique(expert.party))
    full_parties_lr = Set(unique(lr_data.party))

    lost_expert = setdiff(full_parties_expert, train_parties_expert)
    lost_lr = setdiff(full_parties_lr, train_parties_lr)

    println()
    if isempty(lost_expert) && isempty(lost_lr)
        println("✓ No parties lost from training set")
    else
        println("⚠ Parties lost from expert training: $(length(lost_expert))")
        println("⚠ Parties lost from LR training: $(length(lost_lr))")
    end

    # Save to output directory
    # Files use standard names so 01_run_model.jl can load with --data-dir
    output_dir = joinpath(base_dir, "validation", "external_split")
    mkpath(output_dir)

    # Training files (standard names for model loading)
    CSV.write(joinpath(output_dir, "text_data.csv"), text_data)  # UNCHANGED
    CSV.write(joinpath(output_dir, "expert.csv"), expert_train)
    CSV.write(joinpath(output_dir, "lr_data.csv"), lr_train)

    # Test files (for compute step)
    CSV.write(joinpath(output_dir, "expert_test.csv"), expert_test)
    CSV.write(joinpath(output_dir, "lr_data_test.csv"), lr_test)

    # Copy union mapping (needed by model)
    if isdir(joinpath(base_dir, "data"))
        mkpath(joinpath(output_dir, "data"))
        cp(joinpath(base_dir, "data", "union_mapping.csv"),
           joinpath(output_dir, "data", "union_mapping.csv"), force=true)
    end

    println()
    println("Files saved to: $output_dir")
    println("  text_data.csv  — IDENTICAL to original ($(nrow(text_data)) rows)")
    println("  expert.csv     — training only ($(nrow(expert_train)) rows)")
    println("  lr_data.csv    — training only ($(nrow(lr_train)) rows)")
    println("  expert_test.csv — held-out ($(nrow(expert_test)) rows)")
    println("  lr_data_test.csv — held-out ($(nrow(lr_test)) rows)")

    return output_dir
end

# =========================================================================
# STEP 2: Compute held-out validation metrics
# =========================================================================

function compute_holdout_metrics(model_file::String, test_dir::String)
    println()
    println("="^70)
    println("COMPUTING HELD-OUT VALIDATION METRICS")
    println("="^70)

    # Load model output
    model = CSV.read(model_file, DataFrame)
    party_col = hasproperty(model, :party_id) ? :party_id : :party
    println("Model output: $(nrow(model)) party-years")

    # Load test data
    expert_test = CSV.read(joinpath(test_dir, "expert_test.csv"), DataFrame)
    lr_test = CSV.read(joinpath(test_dir, "lr_data_test.csv"), DataFrame)
    println("Held-out expert: $(nrow(expert_test)) observations")
    println("Held-out LR: $(nrow(lr_test)) observations")

    # Build lookup: (party, year) → model estimates
    model_lookup = Dict{Tuple{Int,Int}, NamedTuple}()
    for row in eachrow(model)
        key = (row[party_col], row.year)
        model_lookup[key] = (
            economic_lr = row.economic_lr,
            galtan = row.galtan,
            economic_lr_se = hasproperty(row, :economic_lr_se) ? row.economic_lr_se : missing,
            galtan_se = hasproperty(row, :galtan_se) ? row.galtan_se : missing,
            economic_lr_q025 = hasproperty(row, :economic_lr_q025) ? row.economic_lr_q025 : missing,
            economic_lr_q975 = hasproperty(row, :economic_lr_q975) ? row.economic_lr_q975 : missing,
            galtan_q025 = hasproperty(row, :galtan_q025) ? row.galtan_q025 : missing,
            galtan_q975 = hasproperty(row, :galtan_q975) ? row.galtan_q975 : missing,
        )
    end

    # Map expert variables to dimensions
    econ_vars = Set(["lrecon_ches", "lrecon_poppa", "lrecon_gps", "lrecon_vparty", "welf_vparty"])
    galtan_vars = Set(["galtan_ches", "libcon_gps", "immig_vparty", "lgbt_vparty",
                       "culsup_vparty", "relig_vparty", "gender_vparty"])

    # Process expert test observations
    results = NamedTuple[]

    for row in eachrow(expert_test)
        key = (row.party, row.year)
        haskey(model_lookup, key) || continue
        m = model_lookup[key]

        if row.var in econ_vars
            dim = "economic_lr"
            model_val = m.economic_lr
            model_q025 = m.economic_lr_q025
            model_q975 = m.economic_lr_q975
        elseif row.var in galtan_vars
            dim = "galtan"
            model_val = m.galtan
            model_q025 = m.galtan_q025
            model_q975 = m.galtan_q975
        else
            continue
        end

        covered = !ismissing(model_q025) && !ismissing(model_q975) &&
                  row.val >= model_q025 && row.val <= model_q975

        push!(results, (
            party = row.party,
            country = row.country,
            year = row.year,
            var = row.var,
            dimension = dim,
            expert_val = row.val,
            model_val = model_val,
            error = row.val - model_val,
            abs_error = abs(row.val - model_val),
            covered_95 = ismissing(model_q025) ? missing : covered,
        ))
    end

    if isempty(results)
        println("ERROR: No matching test observations found")
        return nothing
    end

    results_df = DataFrame(results)

    # Compute and report metrics
    println()
    println("-"^70)
    println("HELD-OUT VALIDATION RESULTS")
    println("-"^70)

    # Overall
    overall_r = cor(results_df.expert_val, results_df.model_val)
    overall_mae = mean(results_df.abs_error)
    overall_rmse = sqrt(mean(results_df.error .^ 2))
    println()
    println(@sprintf("Overall: r=%.4f, MAE=%.4f, RMSE=%.4f, n=%d",
                     overall_r, overall_mae, overall_rmse, nrow(results_df)))

    # By dimension
    println()
    println("By dimension:")
    for dim in sort(unique(results_df.dimension))
        d = filter(r -> r.dimension == dim, results_df)
        r_val = cor(d.expert_val, d.model_val)
        mae = mean(d.abs_error)
        rmse = sqrt(mean(d.error .^ 2))
        cov = count(skipmissing(d.covered_95)) / count(!ismissing, d.covered_95)
        println(@sprintf("  %-12s: r=%.4f, MAE=%.4f, RMSE=%.4f, CIC95=%.1f%%, n=%d",
                         dim, r_val, mae, rmse, 100*cov, nrow(d)))
    end

    # By source
    println()
    println("By source:")
    for var in sort(unique(results_df.var))
        v = filter(r -> r.var == var, results_df)
        nrow(v) < 5 && continue
        r_val = cor(v.expert_val, v.model_val)
        mae = mean(v.abs_error)
        println(@sprintf("  %-20s: r=%.4f, MAE=%.4f, n=%d", var, r_val, mae, nrow(v)))
    end

    return results_df
end

# =========================================================================
# Main
# =========================================================================

function main()
    args = ARGS

    if isempty(args) || args[1] == "prepare"
        output_dir = prepare_holdout_data()

        println()
        println("="^70)
        println("NEXT STEPS")
        println("="^70)
        println("""
1. Run model on training data:
   julia 01_run_model.jl --data-dir $output_dir

2. Run post-estimation on training run output:
   julia 02_post_estimation.jl

3. Compute held-out metrics:
   julia scripts/validate_external.jl compute <party_positions_file.csv>
""")

    elseif args[1] == "compute" && length(args) >= 2
        model_file = args[2]
        test_dir = joinpath("validation", "external_split")

        isfile(model_file) || error("Model file not found: $model_file")
        isdir(test_dir) || error("Test data not found. Run 'prepare' first.")

        results_df = compute_holdout_metrics(model_file, test_dir)

        if results_df !== nothing
            timestamp = Dates.format(now(), "yyyy-mm-dd_HH-MM-SS")
            output_file = joinpath("validation", "external_validation_$(timestamp).csv")
            CSV.write(output_file, results_df)
            println()
            println("Detailed results saved: $output_file")
        end

    else
        println("Usage:")
        println("  julia scripts/validate_external.jl prepare")
        println("  julia scripts/validate_external.jl compute <party_positions.csv>")
    end
end

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