#!/usr/bin/env julia
#=
02_enrich_output.jl - Enrich party positions CSV with model-input-derived metadata

Fast post-processing script that adds union membership status from the model-ready
inputs. Operates purely on CSV files (no chain loading). Takes seconds, not minutes.

Usage:
    julia 02_enrich_output.jl                # enriches latest party_positions_*.csv
    julia 02_enrich_output.jl somefile.csv    # enriches a specific file

Adds columns:
    in_union            - 1 if party's union had a joint manifesto that year, 0 otherwise
=#

using CSV
using DataFrames
using Dates

function find_latest_output()
    outdir = "outputs/estimations/latest"
    files = filter(f -> startswith(f, "party_positions_") && endswith(f, ".csv") &&
                        !contains(f, "metadata") && !contains(f, "tables"),
                   readdir(outdir))
    isempty(files) && error("No party_positions_*.csv found in $outdir/")
    sort!(files, rev=true)
    return joinpath(outdir, files[1])
end

function enrich(input_file::String)
    println("="^60)
    println("ENRICH OUTPUT")
    println("="^60)
    println("Input: $input_file")

    output = CSV.read(input_file, DataFrame)
    println("  Rows: $(nrow(output)), Columns: $(ncol(output))")

    # --- Union mapping (for in_union) ---
    union_mapping_file = joinpath("data", "union_mapping.csv")
    constituent_to_union = Dict{Int, Int}()
    if isfile(union_mapping_file)
        union_df = CSV.read(union_mapping_file, DataFrame)
        for row in eachrow(union_df)
            constituent_to_union[row.expert_pf_id] = row.manifesto_pf_id
        end
    end

    # --- in_union dummy (year-varying) ---
    text_data_file = "data/text_data.csv"
    output.in_union = zeros(Int, nrow(output))

    if isfile(text_data_file) && !isempty(constituent_to_union)
        text_df = CSV.read(text_data_file, DataFrame)

        # Build set of (party_pf_id, year) pairs for manifesto data
        manifesto_text = filter(r -> r.project == "Manifesto Project", text_df)
        manifesto_party_years = Set{Tuple{Int, Int}}()
        for row in eachrow(manifesto_text)
            push!(manifesto_party_years, (row.party, row.year))
        end

        # Process each (party, segment) group
        gdf = groupby(output, [:party_id, :segment_num])
        for subdf in gdf
            pid = subdf.party_id[1]
            !haskey(constituent_to_union, pid) && continue
            union_id = constituent_to_union[pid]

            # Get row indices in the full output for this group
            idxs = parentindices(subdf)[1]

            # Determine in_union at election years
            election_year_vals = Dict{Int, Int}()
            for (j, row) in enumerate(eachrow(subdf))
                if (union_id, row.year) in manifesto_party_years
                    election_year_vals[row.year] = 1
                elseif (pid, row.year) in manifesto_party_years
                    election_year_vals[row.year] = 0
                end
            end

            # Forward-fill within segment
            sorted_pairs = sort(collect(zip(subdf.year, idxs)))
            last_val = 0
            for (yr, idx) in sorted_pairs
                if haskey(election_year_vals, yr)
                    last_val = election_year_vals[yr]
                end
                output.in_union[idx] = last_val
            end
        end

        n_in_union = count(x -> x == 1, output.in_union)
        println("  in_union: $n_in_union rows flagged as union members")
    else
        println("  WARNING: Could not compute in_union (missing files)")
    end

    # --- Reorder columns ---
    estimate_cols = Symbol[]
    for base in ["economic_lr", "galtan", "pro_market", "pro_welfare", "cosmopolitan", "traditional"]
        sym = Symbol(base)
        if hasproperty(output, sym)
            push!(estimate_cols, sym)
            push!(estimate_cols, Symbol("$(base)_se"))
            push!(estimate_cols, Symbol("$(base)_q025"))
            push!(estimate_cols, Symbol("$(base)_q975"))
        end
    end

    # Fix country code: MO (Macau) → MK (North Macedonia) — GPS uses wrong ISO2
    output.country = replace(output.country, "MO" => "MK")

    col_order = vcat(
        [:party_id, :country, :year, :segment_num, :union_party_id, :in_union],
        estimate_cols
    )
    col_order = filter(c -> hasproperty(output, c), col_order)
    select!(output, col_order)

    # --- Write back ---
    CSV.write(input_file, output)
    println("\n  Wrote: $input_file")
    println("  Columns ($(ncol(output))): $(join(string.(names(output)), ", "))")

    return output
end

function main(args=ARGS)
    input = length(args) >= 1 ? args[1] : find_latest_output()
    enrich(input)
    println("\nDone.")
end

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