#!/usr/bin/env julia
#############################################################################
## 02_data_loading.jl
## Load and preprocess data for latent trait model
## Loads three datasets: text_data (manifesto + PolDem), expert dimension-specific, expert general L-R
##
## Supports both:
##   - 4D model (V10): type_high/type_low columns for bipolar bridges
##   - 2D model (V1): dim_idx + direction columns for direct estimation
#############################################################################

using DataFrames, CSV, CategoricalArrays, Statistics

#############################################################################
## UNION MAPPING: Individual party estimates via mean-constituent model
## Loads data/union_mapping.csv and builds lookup structures
#############################################################################

"""
    load_union_mapping(project_root::String)

Load union_mapping.csv and build lookup dictionaries.
Returns (union_to_constituents, constituent_to_union) dicts.
If file is missing or empty, returns empty dicts (backwards compatible).
"""
function load_union_mapping(project_root::String=".")
    mapping_file = joinpath(project_root, "data", "union_mapping.csv")

    union_to_constituents = Dict{Int, Vector{Int}}()
    constituent_to_union = Dict{Int, Int}()

    if !isfile(mapping_file)
        println("  No union_mapping.csv found - running without union decomposition")
        return union_to_constituents, constituent_to_union
    end

    df = CSV.read(mapping_file, DataFrame)
    if nrow(df) == 0
        println("  union_mapping.csv is empty - running without union decomposition")
        return union_to_constituents, constituent_to_union
    end

    for row in eachrow(df)
        union_id = row.manifesto_pf_id
        expert_id = row.expert_pf_id

        if !haskey(union_to_constituents, union_id)
            union_to_constituents[union_id] = Int[]
        end
        if !(expert_id in union_to_constituents[union_id])
            push!(union_to_constituents[union_id], expert_id)
        end
        constituent_to_union[expert_id] = union_id
    end

    println("  Union mapping loaded: $(length(union_to_constituents)) unions, $(length(constituent_to_union)) constituents")
    return union_to_constituents, constituent_to_union
end

#############################################################################
## SEGMENT-BASED INDEXING CONFIGURATION
## Split parties at gaps > MAX_GAP years to avoid flat posteriors
#############################################################################
const MAX_GAP = 7      # Maximum years between observations within a segment
const MIN_OBS = 2      # Minimum observations per segment (drop segments with fewer)

#############################################################################
## 2D MODEL MAPPING CONFIGURATION
## Maps type_high/type_low pairs to dim_idx + direction
#############################################################################
const TYPE_TO_DIM_DIRECTION = Dict(
    # Economic dimension: pro_market = right (+1), pro_welfare = left (-1)
    ("pro_market", "pro_welfare") => (dim_idx=1, direction=1),   # Right
    ("pro_welfare", "pro_market") => (dim_idx=1, direction=-1),  # Left

    # Cultural dimension: traditional = TAN (+1), cosmopolitan = GAL (-1)
    ("traditional", "cosmopolitan") => (dim_idx=2, direction=1), # TAN
    ("cosmopolitan", "traditional") => (dim_idx=2, direction=-1) # GAL
)

# Expert dimension mapping (lrecon -> economic, galtan/cultural -> galtan)
const EXPERT_VAR_TO_DIM = Dict(
    "lrecon_ches" => 1,
    "lrecon_vparty" => 1,
    "welf_vparty" => 1,
    "lrecon_gps" => 1,
    "lrecon_poppa" => 1,
    "galtan_ches" => 2,
    "libcon_gps" => 2,
    "immig_vparty" => 2,
    "lgbt_vparty" => 2,
    "culsup_vparty" => 2,
    "relig_vparty" => 2,
    "gender_vparty" => 2
)

function load_and_preprocess_4dim_data(start_year=1950; data_dir::String=".")
    println("Loading party-position data files...")
    println("Start year filter: $start_year")
    data_dir != "." && println("Data directory: $data_dir")

    # Load union mapping (check data_dir first, fall back to project root)
    println("\nLoading union mapping...")
    union_mapping_dir = isfile(joinpath(data_dir, "data", "union_mapping.csv")) ? data_dir : "."
    union_to_constituents, constituent_to_union = load_union_mapping(union_mapping_dir)

    # Load the three datasets
    text_data_raw = CSV.read(joinpath(data_dir, "text_data.csv"), DataFrame)
    expert_raw = CSV.read(joinpath(data_dir, "expert.csv"), DataFrame)
    lr_data_raw = CSV.read(joinpath(data_dir, "lr_data.csv"), DataFrame)

    # Filter to start year BEFORE calculating year0
    text_data_raw = text_data_raw[text_data_raw.year .>= start_year, :]
    expert_raw = expert_raw[expert_raw.year .>= start_year, :]
    lr_data_raw = lr_data_raw[lr_data_raw.year .>= start_year, :]

    println("Data filtered to $start_year onwards:")
    println("  Text data (manifesto + PolDem): $(nrow(text_data_raw)) observations")
    println("  Expert: $(nrow(expert_raw)) observations")
    println("  L-R data: $(nrow(lr_data_raw)) observations")

    # Define base year for relative time indexing
    year0 = Int(minimum([minimum(text_data_raw.year), minimum(expert_raw.year), minimum(lr_data_raw.year)])) - 1
    println("Base year set to: $year0")

    # Create type mapping for the four dimensions
    type_map = Dict(
        "pro_market" => 1,
        "pro_welfare" => 2,
        "cosmopolitan" => 3,
        "traditional" => 4
    )

    println("Type mapping: pro_market=1, pro_welfare=2, cosmopolitan=3, traditional=4")

    # Process text data (manifesto + PolDem media)
    text_data = copy(text_data_raw)
    text_data = text_data[text_data.year .> year0, :]

    # Add type indices for text items (V4/V10: bipolar bridge structure)
    if !("type_high" in names(text_data)) || !("type_low" in names(text_data))
        error("Text data must contain 'type_high' and 'type_low' columns with values: pro_market, pro_welfare, cosmopolitan, traditional")
    end

    text_data.type_high_idx = [type_map[t] for t in text_data.type_high]
    text_data.type_low_idx = [type_map[t] for t in text_data.type_low]

    # V1 (2D model): Add dim_idx and direction columns
    # Maps type_high/type_low to single dimension + direction
    dim_idx_man = Int[]
    direction_man = Int[]
    for row in eachrow(text_data)
        key = (row.type_high, row.type_low)
        if haskey(TYPE_TO_DIM_DIRECTION, key)
            mapping = TYPE_TO_DIM_DIRECTION[key]
            push!(dim_idx_man, mapping.dim_idx)
            push!(direction_man, mapping.direction)
        else
            # Unknown mapping - this should not happen with valid data
            error("Unknown type_high/type_low pair: $(row.type_high) / $(row.type_low)")
        end
    end
    text_data.dim_idx_man = dim_idx_man
    text_data.direction_man = direction_man

    # Standard processing
    text_data.country = categorical(text_data.country)
    text_data.party = categorical(text_data.party)
    text_data.var = categorical(text_data.var)
    text_data.Year = Int.(text_data.year) .- year0
    sort!(text_data, [:country, :party, :year, :var])
    println("Text data processed: $(nrow(text_data)) observations with bipolar bridge structure")

    # Process expert dimension-specific data (bipolar items like lrecon_ches, galtan_ches)
    expert_dim_vars = ["lrecon_ches", "galtan_ches", "lrecon_vparty", "welf_vparty",
                       "lrecon_gps", "libcon_gps", "lrecon_poppa",
                       "immig_vparty", "lgbt_vparty", "culsup_vparty", "relig_vparty", "gender_vparty"]
    expert_dim = expert_raw[in.(expert_raw.var, Ref(expert_dim_vars)), :]
    expert_dim = expert_dim[(expert_dim.year .> year0) .& (expert_dim.val .>= 0) .& (expert_dim.val .<= 1), :]

    # V5: Load integer observations, scale sizes, and expert counts for beta-binomial likelihood
    expert_dim.val_int = Int.(expert_dim.val_int)
    expert_dim.n_scale = Int.(expert_dim.n_scale)
    expert_dim.n_experts = Int.(expert_dim.n_experts)

    # Add type mappings for dimension-specific expert data
    if !("type_low" in names(expert_dim)) || !("type_high" in names(expert_dim))
        error("Expert data must contain 'type_low' and 'type_high' columns")
    end

    expert_dim.type_high_idx = [type_map[t] for t in expert_dim.type_high]
    expert_dim.type_low_idx = [type_map[t] for t in expert_dim.type_low]

    # V1 (2D model): Add dim_idx for expert dimension data
    dim_idx_exp = Int[]
    for row in eachrow(expert_dim)
        var_name = string(row.var)
        if haskey(EXPERT_VAR_TO_DIM, var_name)
            push!(dim_idx_exp, EXPERT_VAR_TO_DIM[var_name])
        else
            # Fallback: infer from type_high/type_low
            key = (row.type_high, row.type_low)
            if haskey(TYPE_TO_DIM_DIRECTION, key)
                push!(dim_idx_exp, TYPE_TO_DIM_DIRECTION[key].dim_idx)
            else
                error("Unknown expert variable: $var_name with type pair $(row.type_high) / $(row.type_low)")
            end
        end
    end
    expert_dim.dim_idx_exp = dim_idx_exp

    expert_dim.country = categorical(expert_dim.country)
    expert_dim.party = categorical(expert_dim.party)
    expert_dim.var = categorical(expert_dim.var)
    expert_dim.Year = Int.(expert_dim.year) .- year0
    sort!(expert_dim, [:country, :party, :year, :var])
    println("Expert dimension-specific data processed: $(nrow(expert_dim)) observations")

    # Process expert general L-R data (cross-dimensional anchoring)
    lr_vars = ["lr_ches", "lr_poppa", "lr_morgan"]  # General left-right items
    expert_lr = lr_data_raw[in.(lr_data_raw.var, Ref(lr_vars)), :]
    expert_lr = expert_lr[(expert_lr.year .> year0) .& (expert_lr.val .>= 0) .& (expert_lr.val .<= 1), :]

    # V5: Load integer observations, scale sizes, and expert counts for beta-binomial likelihood
    expert_lr.val_int = Int.(expert_lr.val_int)
    expert_lr.n_scale = Int.(expert_lr.n_scale)
    expert_lr.n_experts = Int.(expert_lr.n_experts)

    expert_lr.country = categorical(expert_lr.country)
    expert_lr.party = categorical(expert_lr.party)
    expert_lr.var = categorical(expert_lr.var)
    expert_lr.Year = Int.(expert_lr.year) .- year0
    sort!(expert_lr, [:country, :party, :year, :var])
    println("Expert general L-R data processed: $(nrow(expert_lr)) observations")

    # Validate data integrity
    println("\nData validation:")

    # Check text data dimension pair distribution (V4: bipolar bridges)
    type_pair_counts = combine(groupby(text_data, [:type_high, :type_low]), nrow => :count)
    for row in eachrow(type_pair_counts)
        println("  $(row.type_high) ↔ $(row.type_low): $(row.count) text data observations")
    end

    # Check expert dimension-specific type pairs
    type_pair_counts = combine(groupby(expert_dim, [:type_high, :type_low]), nrow => :count)
    for row in eachrow(type_pair_counts)
        println("  $(row.type_high) - $(row.type_low): $(row.count) expert dimension-specific observations")
    end

    # Check general L-R items
    lr_var_counts = combine(groupby(expert_lr, :var), nrow => :count)
    for row in eachrow(lr_var_counts)
        println("  $(row.var): $(row.count) general L-R observations")
    end

    # Check overlapping parties across datasets
    text_data_parties = Set(text_data.party)
    expert_dim_parties = Set(expert_dim.party)
    expert_lr_parties = Set(expert_lr.party)

    all_parties = union(text_data_parties, expert_dim_parties, expert_lr_parties)
    println("\nParty coverage:")
    println("  Total unique parties: $(length(all_parties))")
    println("  In text data: $(length(text_data_parties))")
    println("  In expert dimension-specific: $(length(expert_dim_parties))")
    println("  In expert general L-R: $(length(expert_lr_parties))")
    println("  In all three datasets: $(length(intersect(text_data_parties, expert_dim_parties, expert_lr_parties)))")

    return text_data, expert_dim, expert_lr, year0, union_to_constituents, constituent_to_union
end

# Execute if run directly
if abspath(PROGRAM_FILE) == @__FILE__
    text_data, expert_dim, expert_lr, year0, u2c, c2u = load_and_preprocess_4dim_data()
    println("4D data loading test completed successfully")
end