#!/usr/bin/env julia
#############################################################################
## 03_data_preparation_4dim.jl
## V10: Segment-based indexing to fix long gap interpolation issues
##
## Key change: Split parties at gaps > MAX_GAP years into independent segments.
## Each segment has its own random walk (restarts at segment start).
## Segments with < MIN_OBS observations are dropped.
#############################################################################

using DataFrames, CSV, CategoricalArrays, Statistics, StatsFuns

# Import configuration from data loading module
include("02_data_loading.jl")

"""
    split_party_years_into_segments(years::Vector{Int}, max_gap::Int)

Split a party's observation years into segments based on gaps.
Returns a vector of vectors, where each inner vector contains consecutive years
with max `max_gap` years between observations.
"""
function split_party_years_into_segments(years::Vector{Int}, max_gap::Int)
    if isempty(years)
        return Vector{Vector{Int}}()
    end

    sorted_years = sort(unique(years))
    segments = [Int[sorted_years[1]]]

    for y in sorted_years[2:end]
        if y - segments[end][end] <= max_gap
            push!(segments[end], y)
        else
            # Gap too large - start new segment
            push!(segments, [y])
        end
    end

    return segments
end

function prepare_4dim_stan_data(manifesto, expert_dim, expert_lr, year0;
                                union_to_constituents=Dict{Int,Vector{Int}}(),
                                constituent_to_union=Dict{Int,Int}())
    println("Preparing data for Stan model (Segment-based indexing)...")
    println("  MAX_GAP = $MAX_GAP years, MIN_OBS = $MIN_OBS observations")

    has_unions = !isempty(union_to_constituents)
    if has_unions
        println("  Union mapping: $(length(union_to_constituents)) unions, $(length(constituent_to_union)) constituents")
    else
        println("  No union mapping - standard party indexing")
    end

    # =========================================================================
    # STEP 1: Collect observation years per party (union-aware)
    # For union parties: create segments for each CONSTITUENT, not the union.
    # Union manifesto years are assigned to all constituents.
    # =========================================================================

    # Identify which party IDs in data are unions vs standalone
    # NOTE: levels() returns raw types (Int64 for integer party IDs).
    # We consistently use String keys for all party lookups to avoid type mismatches.
    manifesto_party_strs = Set(string.(levels(manifesto.party)))
    union_ids_in_data = Set{String}()
    if has_unions
        for uid in keys(union_to_constituents)
            uid_str = string(uid)
            if uid_str in manifesto_party_strs
                push!(union_ids_in_data, uid_str)
            end
        end
        println("  Union party IDs found in manifesto data: $(length(union_ids_in_data))")
    end

    # Collect all party IDs that need segments
    # For unions: constituents get segments; union itself does NOT
    # For standalone: party gets segment as before
    party_obs_years = Dict{String, Set{Int}}()

    # First pass: collect non-union parties from all data sources (as strings)
    all_data_parties = Set{String}()
    for p in levels(manifesto.party)
        push!(all_data_parties, string(p))
    end
    for p in levels(expert_dim.party)
        push!(all_data_parties, string(p))
    end
    for p in levels(expert_lr.party)
        push!(all_data_parties, string(p))
    end

    # Initialize observation years for standalone parties and constituents
    for p in all_data_parties
        p_int = tryparse(Int, string(p))
        if p_int !== nothing && has_unions && haskey(union_to_constituents, p_int) && string(p) in union_ids_in_data
            # This is a union ID in manifesto - skip it, create entries for constituents instead
            continue
        end
        party_obs_years[p] = Set{Int}()
    end

    # For unions: ensure all constituents have entries
    if has_unions
        for (uid, constituents) in union_to_constituents
            uid_str = string(uid)
            if uid_str in union_ids_in_data
                for cid in constituents
                    cid_str = string(cid)
                    if !haskey(party_obs_years, cid_str)
                        party_obs_years[cid_str] = Set{Int}()
                    end
                end
            end
        end
    end

    # Add years from manifesto
    for row in eachrow(manifesto)
        p_str = string(row.party)
        p_int = tryparse(Int, p_str)
        if p_int !== nothing && has_unions && haskey(union_to_constituents, p_int) && p_str in union_ids_in_data
            # Union manifesto: add year to ALL constituents
            for cid in union_to_constituents[p_int]
                cid_str = string(cid)
                if haskey(party_obs_years, cid_str)
                    push!(party_obs_years[cid_str], row.Year)
                end
            end
        else
            if haskey(party_obs_years, p_str)
                push!(party_obs_years[p_str], row.Year)
            end
        end
    end

    # Add years from expert_dim (individual party data - direct)
    for row in eachrow(expert_dim)
        p_str = string(row.party)
        if haskey(party_obs_years, p_str)
            push!(party_obs_years[p_str], row.Year)
        end
    end

    # Add years from expert_lr (individual party data - direct)
    for row in eachrow(expert_lr)
        p_str = string(row.party)
        if haskey(party_obs_years, p_str)
            push!(party_obs_years[p_str], row.Year)
        end
    end

    J_original = length(party_obs_years)
    println("Total number of parties to create segments for: $J_original")

    # =========================================================================
    # STEP 2: Split parties into segments and filter by MIN_OBS
    # =========================================================================
    println("\nCreating segments (splitting at gaps > $MAX_GAP years)...")

    segment_data = []
    segment_id = 0
    parties_split = 0
    segments_dropped = 0
    observations_dropped = 0

    for (p, years_set) in party_obs_years
        years = collect(years_set)
        if isempty(years)
            continue
        end

        segments = split_party_years_into_segments(years, MAX_GAP)

        if length(segments) > 1
            parties_split += 1
        end

        for (seg_num, seg_years) in enumerate(segments)
            n_obs = length(seg_years)
            if n_obs >= MIN_OBS
                segment_id += 1
                push!(segment_data, (
                    segment_id = segment_id,
                    party_id = p,
                    segment_num = seg_num,
                    year_start = minimum(seg_years),
                    year_end = maximum(seg_years),
                    n_obs = n_obs
                ))
            else
                segments_dropped += 1
                observations_dropped += n_obs
            end
        end
    end

    segment_info = DataFrame(segment_data)
    S = nrow(segment_info)  # Number of valid segments

    println("  Segments created: $S (from $J_original parties)")
    println("  Parties split into multiple segments: $parties_split")
    println("  Segments dropped (< $MIN_OBS obs): $segments_dropped")
    println("  Observations dropped: $observations_dropped")

    # Get unique parties that have at least one valid segment
    all_parties = unique(segment_info.party_id)
    J = length(all_parties)
    println("  Parties with valid segments: $J")

    # Create party-to-index mapping for the valid parties
    party_to_index = Dict(all_parties .=> 1:J)

    # =========================================================================
    # STEP 3: Create segment-year index space (R) - consecutive within segment
    # =========================================================================
    println("\nCreating segment-year index space...")

    segment_year_data = []
    for row in eachrow(segment_info)
        for y in row.year_start:row.year_end
            push!(segment_year_data, (
                segment_id = row.segment_id,
                party_id = row.party_id,
                Year = y
            ))
        end
    end

    segment_year = DataFrame(segment_year_data)
    segment_year.rr = 1:nrow(segment_year)
    R = nrow(segment_year)
    println("Total segment-year positions (R): $R")

    # Compute len_theta_ts for segments (years per segment)
    len_theta_ts = [row.year_end - row.year_start + 1 for row in eachrow(segment_info)]
    @assert sum(len_theta_ts) == R "sum(len_theta_ts)=$(sum(len_theta_ts)) must equal R=$R"

    # =========================================================================
    # STEP 4: Map observations to segment-year indices (union-aware)
    # =========================================================================
    println("\nMapping observations to segment-year indices...")

    # Create lookup: (party_str, year) -> segment_id (for valid segments only)
    party_year_to_segment = Dict{Tuple{String, Int}, Int}()
    for row in eachrow(segment_info)
        for y in row.year_start:row.year_end
            party_year_to_segment[(string(row.party_id), y)] = row.segment_id
        end
    end

    # --- MANIFESTO: union-aware mapping ---
    # For union manifesto obs: map to first constituent's segment (for ss_man).
    # The actual theta averaging is handled via constituent arrays.
    manifesto_segment_ids = Union{Int, Missing}[]
    for row in eachrow(manifesto)
        p_str = string(row.party)
        p_int = tryparse(Int, p_str)
        key = (p_str, row.Year)

        if haskey(party_year_to_segment, key)
            # Direct mapping (non-union or constituent with own segment)
            push!(manifesto_segment_ids, party_year_to_segment[key])
        elseif p_int !== nothing && has_unions && haskey(union_to_constituents, p_int)
            # Union party: use first constituent's segment
            found = false
            for cid in union_to_constituents[p_int]
                ckey = (string(cid), row.Year)
                if haskey(party_year_to_segment, ckey)
                    push!(manifesto_segment_ids, party_year_to_segment[ckey])
                    found = true
                    break
                end
            end
            if !found
                push!(manifesto_segment_ids, missing)
            end
        else
            push!(manifesto_segment_ids, missing)
        end
    end
    manifesto.segment_id = manifesto_segment_ids
    n_manifesto_before = nrow(manifesto)
    manifesto = manifesto[.!ismissing.(manifesto.segment_id), :]
    manifesto.segment_id = Int.(manifesto.segment_id)
    println("  Manifesto: $(nrow(manifesto))/$n_manifesto_before observations (dropped $(n_manifesto_before - nrow(manifesto)) in invalid segments)")

    # --- EXPERT DIM: direct mapping (individual party data) ---
    expert_dim_segment_ids = Union{Int, Missing}[]
    for row in eachrow(expert_dim)
        key = (string(row.party), row.Year)
        if haskey(party_year_to_segment, key)
            push!(expert_dim_segment_ids, party_year_to_segment[key])
        else
            push!(expert_dim_segment_ids, missing)
        end
    end
    expert_dim.segment_id = expert_dim_segment_ids
    n_expert_dim_before = nrow(expert_dim)
    expert_dim = expert_dim[.!ismissing.(expert_dim.segment_id), :]
    expert_dim.segment_id = Int.(expert_dim.segment_id)
    println("  Expert dim: $(nrow(expert_dim))/$n_expert_dim_before observations (dropped $(n_expert_dim_before - nrow(expert_dim)) in invalid segments)")

    # --- EXPERT LR: direct mapping (individual party data) ---
    expert_lr_segment_ids = Union{Int, Missing}[]
    for row in eachrow(expert_lr)
        key = (string(row.party), row.Year)
        if haskey(party_year_to_segment, key)
            push!(expert_lr_segment_ids, party_year_to_segment[key])
        else
            push!(expert_lr_segment_ids, missing)
        end
    end
    expert_lr.segment_id = expert_lr_segment_ids
    n_expert_lr_before = nrow(expert_lr)
    expert_lr = expert_lr[.!ismissing.(expert_lr.segment_id), :]
    expert_lr.segment_id = Int.(expert_lr.segment_id)
    println("  Expert LR: $(nrow(expert_lr))/$n_expert_lr_before observations (dropped $(n_expert_lr_before - nrow(expert_lr)) in invalid segments)")

    # =========================================================================
    # STEP 5: Create segment indices (ss) for each observation
    # ss indexes into 1:S (segment space), used for party-level parameters
    # =========================================================================

    # Create segment_id to ss mapping
    segment_to_ss = Dict(row.segment_id => i for (i, row) in enumerate(eachrow(segment_info)))

    manifesto.ss_man = [segment_to_ss[sid] for sid in manifesto.segment_id]
    expert_dim.ss_exp_dim = [segment_to_ss[sid] for sid in expert_dim.segment_id]
    expert_lr.ss_exp_lr = [segment_to_ss[sid] for sid in expert_lr.segment_id]

    # Validate segment indices
    @assert all(1 .<= manifesto.ss_man .<= S)
    @assert all(1 .<= expert_dim.ss_exp_dim .<= S)
    @assert all(1 .<= expert_lr.ss_exp_lr .<= S)

    # =========================================================================
    # STEP 6: Map rr indices (segment-year) to datasets via leftjoin
    # =========================================================================

    # Create (segment_id, Year) -> rr lookup
    seg_year_to_rr = Dict{Tuple{Int, Int}, Int}()
    for row in eachrow(segment_year)
        seg_year_to_rr[(row.segment_id, row.Year)] = row.rr
    end

    # Join manifesto with segment_year to get rr indices
    manifesto = leftjoin(manifesto, segment_year, on=[:segment_id, :Year])
    rename!(manifesto, :rr => :rr_man)

    expert_dim = leftjoin(expert_dim, segment_year, on=[:segment_id, :Year])
    rename!(expert_dim, :rr => :rr_exp_dim)

    expert_lr = leftjoin(expert_lr, segment_year, on=[:segment_id, :Year])
    rename!(expert_lr, :rr => :rr_exp_lr)

    # Validate rr mappings
    @assert all(!ismissing, manifesto.rr_man) "Some manifesto observations have no rr_man mapping"
    @assert all(!ismissing, expert_dim.rr_exp_dim) "Some expert_dim observations have no rr_exp_dim mapping"
    @assert all(!ismissing, expert_lr.rr_exp_lr) "Some expert_lr observations have no rr_exp_lr mapping"

    # Convert to Int
    manifesto.rr_man = Int.(manifesto.rr_man)
    expert_dim.rr_exp_dim = Int.(expert_dim.rr_exp_dim)
    expert_lr.rr_exp_lr = Int.(expert_lr.rr_exp_lr)

    # Validate bounds
    @assert all(1 .<= manifesto.rr_man .<= R) "rr_man out of bounds"
    @assert all(1 .<= expert_dim.rr_exp_dim .<= R) "rr_exp_dim out of bounds"
    @assert all(1 .<= expert_lr.rr_exp_lr .<= R) "rr_exp_lr out of bounds"

    # Print diagnostics
    n_observed = length(unique(vcat(manifesto.rr_man, expert_dim.rr_exp_dim, expert_lr.rr_exp_lr)))
    println("\nSegment-years with data: $n_observed / $R ($(round(100*n_observed/R, digits=1))%)")
    println("Segment-years for interpolation: $(R - n_observed)")

    # =========================================================================
    # STEP 6b: Build constituent arrays for union manifesto/expert observations
    # For each manifesto obs: store list of constituent rr indices for averaging
    # Non-union obs: single rr (n_const=1)
    # Union obs: multiple rr values (n_const=len(constituents))
    # =========================================================================
    println("\nBuilding constituent arrays for mean-constituent model...")

    # --- Manifesto constituent arrays ---
    n_const_man_vec = Int[]       # n_const for each manifesto obs
    const_rr_man_vec = Int[]      # flat array of constituent rr values
    const_offset_man_vec = Int[]  # offset into const_rr for each obs

    offset = 1
    for row in eachrow(manifesto)
        p_str = string(row.party)
        p_int = tryparse(Int, p_str)

        if p_int !== nothing && has_unions && haskey(union_to_constituents, p_int) && p_str in union_ids_in_data
            # Union manifesto obs: find rr for each constituent in this year
            constituent_rrs = Int[]
            for cid in union_to_constituents[p_int]
                ckey = (string(cid), row.Year)
                if haskey(party_year_to_segment, ckey)
                    sid = party_year_to_segment[ckey]
                    rr_key = (sid, row.Year)
                    if haskey(seg_year_to_rr, rr_key)
                        push!(constituent_rrs, seg_year_to_rr[rr_key])
                    end
                end
            end
            if isempty(constituent_rrs)
                # Fallback: use the rr_man already assigned
                push!(constituent_rrs, row.rr_man)
            end
            push!(n_const_man_vec, length(constituent_rrs))
            push!(const_offset_man_vec, offset)
            append!(const_rr_man_vec, constituent_rrs)
            offset += length(constituent_rrs)
        else
            # Non-union: single constituent (itself)
            push!(n_const_man_vec, 1)
            push!(const_offset_man_vec, offset)
            push!(const_rr_man_vec, row.rr_man)
            offset += 1
        end
    end

    N_const_man_total = length(const_rr_man_vec)
    n_union_man = count(x -> x > 1, n_const_man_vec)
    println("  Manifesto: $(nrow(manifesto)) obs, $n_union_man union obs, $N_const_man_total total constituent entries")

    # --- Expert dim constituent arrays ---
    # Individual expert obs always have n_const=1 (direct mapping)
    # Union-level expert obs (if any) would average — but typically expert data is at individual party level
    n_const_exp_dim_vec = Int[]
    const_rr_exp_dim_vec = Int[]
    const_offset_exp_dim_vec = Int[]

    offset = 1
    for row in eachrow(expert_dim)
        p_str = string(row.party)
        p_int = tryparse(Int, p_str)

        if p_int !== nothing && has_unions && haskey(union_to_constituents, p_int) && p_str in union_ids_in_data
            # Union-level expert obs: average over constituents
            constituent_rrs = Int[]
            for cid in union_to_constituents[p_int]
                ckey = (string(cid), row.Year)
                if haskey(party_year_to_segment, ckey)
                    sid = party_year_to_segment[ckey]
                    rr_key = (sid, row.Year)
                    if haskey(seg_year_to_rr, rr_key)
                        push!(constituent_rrs, seg_year_to_rr[rr_key])
                    end
                end
            end
            if isempty(constituent_rrs)
                push!(constituent_rrs, row.rr_exp_dim)
            end
            push!(n_const_exp_dim_vec, length(constituent_rrs))
            push!(const_offset_exp_dim_vec, offset)
            append!(const_rr_exp_dim_vec, constituent_rrs)
            offset += length(constituent_rrs)
        else
            # Individual party obs
            push!(n_const_exp_dim_vec, 1)
            push!(const_offset_exp_dim_vec, offset)
            push!(const_rr_exp_dim_vec, row.rr_exp_dim)
            offset += 1
        end
    end

    N_const_exp_dim_total = length(const_rr_exp_dim_vec)
    n_union_exp_dim = count(x -> x > 1, n_const_exp_dim_vec)
    println("  Expert dim: $(nrow(expert_dim)) obs, $n_union_exp_dim union obs, $N_const_exp_dim_total total constituent entries")

    # --- Expert LR constituent arrays ---
    n_const_exp_lr_vec = Int[]
    const_rr_exp_lr_vec = Int[]
    const_offset_exp_lr_vec = Int[]

    offset = 1
    for row in eachrow(expert_lr)
        p_str = string(row.party)
        p_int = tryparse(Int, p_str)

        if p_int !== nothing && has_unions && haskey(union_to_constituents, p_int) && p_str in union_ids_in_data
            constituent_rrs = Int[]
            for cid in union_to_constituents[p_int]
                ckey = (string(cid), row.Year)
                if haskey(party_year_to_segment, ckey)
                    sid = party_year_to_segment[ckey]
                    rr_key = (sid, row.Year)
                    if haskey(seg_year_to_rr, rr_key)
                        push!(constituent_rrs, seg_year_to_rr[rr_key])
                    end
                end
            end
            if isempty(constituent_rrs)
                push!(constituent_rrs, row.rr_exp_lr)
            end
            push!(n_const_exp_lr_vec, length(constituent_rrs))
            push!(const_offset_exp_lr_vec, offset)
            append!(const_rr_exp_lr_vec, constituent_rrs)
            offset += length(constituent_rrs)
        else
            push!(n_const_exp_lr_vec, 1)
            push!(const_offset_exp_lr_vec, offset)
            push!(const_rr_exp_lr_vec, row.rr_exp_lr)
            offset += 1
        end
    end

    N_const_exp_lr_total = length(const_rr_exp_lr_vec)
    n_union_exp_lr = count(x -> x > 1, n_const_exp_lr_vec)
    println("  Expert LR: $(nrow(expert_lr)) obs, $n_union_exp_lr union obs, $N_const_exp_lr_total total constituent entries")

    # =========================================================================
    # STEP 7: Filter manifesto items with sufficient observations
    # =========================================================================
    var_man_counts_df = combine(groupby(manifesto, :var), :var => length => :n_obs)
    var_man_counts_df = var_man_counts_df[var_man_counts_df.n_obs .>= 2, :]
    var_man_counts = var_man_counts_df.var

    var_exp_dim_counts_df = combine(groupby(expert_dim, :var), :var => length => :n_obs)
    var_exp_dim_counts_df = var_exp_dim_counts_df[var_exp_dim_counts_df.n_obs .>= 2, :]
    var_exp_dim_counts = var_exp_dim_counts_df.var

    var_exp_lr_counts_df = combine(groupby(expert_lr, :var), :var => length => :n_obs)
    var_exp_lr_counts_df = var_exp_lr_counts_df[var_exp_lr_counts_df.n_obs .>= 2, :]
    var_exp_lr_counts = var_exp_lr_counts_df.var

    manifesto = manifesto[in.(manifesto.var, Ref(var_man_counts)), :]
    manifesto.var_man = levelcode.(categorical(manifesto.var, levels=unique(var_man_counts)))

    expert_dim = expert_dim[in.(expert_dim.var, Ref(var_exp_dim_counts)), :]
    expert_dim.var_exp_dim = levelcode.(categorical(expert_dim.var, levels=unique(var_exp_dim_counts)))

    expert_lr = expert_lr[in.(expert_lr.var, Ref(var_exp_lr_counts)), :]
    expert_lr.var_exp_lr = levelcode.(categorical(expert_lr.var, levels=unique(var_exp_lr_counts)))

    # =========================================================================
    # STEP 8: Country (group) indexing
    # =========================================================================
    all_groups = unique(vcat(levels(manifesto.country), levels(expert_dim.country), levels(expert_lr.country)))
    P = length(all_groups)
    println("Total number of unique countries (P): $P")

    group_to_index = Dict(all_groups .=> 1:P)
    manifesto.pp_man = [group_to_index[c] for c in manifesto.country]
    expert_dim.pp_exp_dim = [group_to_index[c] for c in expert_dim.country]
    expert_lr.pp_exp_lr = [group_to_index[c] for c in expert_lr.country]

    @assert all(1 .<= manifesto.pp_man .<= P)
    @assert all(1 .<= expert_dim.pp_exp_dim .<= P)
    @assert all(1 .<= expert_lr.pp_exp_lr .<= P)

    # =========================================================================
    # STEP 9: Country-item-year combinations for zero-inflation model
    # =========================================================================
    manifesto.ciy_key = string.(manifesto.country, "_", manifesto.var_man, "_", manifesto.Year)
    ciy_keys = unique(manifesto.ciy_key)
    N_ciy = length(ciy_keys)
    println("Total unique country-item-year combinations (N_ciy): $N_ciy")

    ciy_key_to_index = Dict(ciy_keys .=> 1:N_ciy)
    manifesto.ciy_idx = [ciy_key_to_index[k] for k in manifesto.ciy_key]
    @assert all(1 .<= manifesto.ciy_idx .<= N_ciy)

    # =========================================================================
    # STEP 10: Segment-country mapping (each segment inherits from party)
    # For constituents: inherit country from their union's manifesto data
    # =========================================================================
    party_country_dict = Dict{String, Int}()

    # From data directly
    for row in eachrow(manifesto)
        p_str = string(row.party)
        p_int = tryparse(Int, p_str)
        c_idx = group_to_index[string(row.country)]
        party_country_dict[p_str] = c_idx

        # If union, also assign country to all constituents
        if p_int !== nothing && has_unions && haskey(union_to_constituents, p_int)
            for cid in union_to_constituents[p_int]
                party_country_dict[string(cid)] = c_idx
            end
        end
    end
    for row in eachrow(expert_dim)
        party_country_dict[string(row.party)] = group_to_index[string(row.country)]
    end
    for row in eachrow(expert_lr)
        party_country_dict[string(row.party)] = group_to_index[string(row.country)]
    end

    # Map each segment to its party's country
    segment_country_idx = Int[]
    for row in eachrow(segment_info)
        pid = string(row.party_id)
        if haskey(party_country_dict, pid)
            push!(segment_country_idx, party_country_dict[pid])
        else
            # Fallback: try to find via union mapping
            pf_int = tryparse(Int, pid)
            if pf_int !== nothing && haskey(constituent_to_union, pf_int)
                uid = constituent_to_union[pf_int]
                if haskey(party_country_dict, string(uid))
                    push!(segment_country_idx, party_country_dict[string(uid)])
                else
                    error("Cannot find country for constituent $pid (union $uid)")
                end
            else
                error("Cannot find country for party $pid")
            end
        end
    end
    @assert all(1 .<= segment_country_idx .<= P)
    @assert length(segment_country_idx) == S

    # Persist canonical country on segment mapping tables
    segment_country = [all_groups[idx] for idx in segment_country_idx]
    segment_info.country = segment_country

    segment_country_by_id = Dict(row.segment_id => segment_country[i] for (i, row) in enumerate(eachrow(segment_info)))
    segment_year.country = [segment_country_by_id[sid] for sid in segment_year.segment_id]

    # =========================================================================
    # STEP 11: Load party family data and map to segments
    # =========================================================================
    println("\nLoading party family data...")
    party_families_file = joinpath("data", "party_families.csv")
    if !isfile(party_families_file)
        error("party_families.csv not found at $party_families_file")
    end
    party_families_df = CSV.read(party_families_file, DataFrame)

    pf_to_family = Dict(row.partyfacts_id => row.family for row in eachrow(party_families_df))

    all_family_names = unique(party_families_df.family)
    family_to_idx = Dict(f => i for (i, f) in enumerate(all_family_names))
    F = length(all_family_names)
    println("Total number of party families (F): $F")
    println("Family categories: ", join(all_family_names, ", "))

    # Map each segment to its party's family
    # For constituents: try own family first, fall back to union's family
    segment_family_idx = Int[]
    unmatched_segments = String[]
    default_family_idx = haskey(family_to_idx, "other") ? family_to_idx["other"] : 1

    for row in eachrow(segment_info)
        pf_id = tryparse(Int, string(row.party_id))
        if !isnothing(pf_id) && haskey(pf_to_family, pf_id)
            family_name = pf_to_family[pf_id]
            push!(segment_family_idx, family_to_idx[family_name])
        elseif !isnothing(pf_id) && haskey(constituent_to_union, pf_id) && haskey(pf_to_family, constituent_to_union[pf_id])
            # Fall back to union's family
            family_name = pf_to_family[constituent_to_union[pf_id]]
            push!(segment_family_idx, family_to_idx[family_name])
        else
            push!(segment_family_idx, default_family_idx)
            push!(unmatched_segments, "$(row.party_id)_seg$(row.segment_num)")
        end
    end

    if !isempty(unmatched_segments)
        println("  Warning: $(length(unmatched_segments)) segments not matched to families (assigned to 'other')")
        if length(unmatched_segments) <= 10
            println("    Unmatched: ", join(unmatched_segments, ", "))
        end
    end

    @assert all(1 .<= segment_family_idx .<= F)
    @assert length(segment_family_idx) == S

    # =========================================================================
    # STEP 12: Find anchor segment
    # With unions: use CDU (1375) as anchor (individual constituent)
    # Without unions: use CDU/CSU (211) as anchor (union ID)
    # =========================================================================
    anchor_party_id = has_unions ? 1375 : 211
    anchor_label = has_unions ? "CDU" : "CDU/CSU"
    anchor_segments = filter(row -> tryparse(Int, string(row.party_id)) == anchor_party_id, segment_info)

    if nrow(anchor_segments) > 0
        # Pick segment with most observations
        anchor_segment_idx = anchor_segments[argmax(anchor_segments.n_obs), :segment_id]
        # Convert to ss index (1:S)
        anchor_segment_ss = segment_to_ss[anchor_segment_idx]
        println("  Anchor segment ($anchor_label, ID $anchor_party_id): segment $anchor_segment_ss ($(anchor_segments[argmax(anchor_segments.n_obs), :n_obs]) obs)")
    else
        println("  Warning: Anchor party $anchor_label (ID $anchor_party_id) not found, using segment 1")
        anchor_segment_ss = 1
    end

    # =========================================================================
    # STEP 13: Validation - check no long gaps remain within segments
    # =========================================================================
    println("\nValidating segment structure...")
    rr_to_segment_year = Dict(row.rr => (segment_id=row.segment_id, year=row.Year) for row in eachrow(segment_year))
    seg_obs_years = Dict(row.segment_id => Set{Int}() for row in eachrow(segment_info))

    for row in eachrow(manifesto)
        push!(seg_obs_years[row.segment_id], row.Year)
    end
    for row in eachrow(expert_dim)
        push!(seg_obs_years[row.segment_id], row.Year)
    end
    for row in eachrow(expert_lr)
        push!(seg_obs_years[row.segment_id], row.Year)
    end

    # Union manifesto rows contribute to every constituent through const_rr_man_vec,
    # even though row.segment_id stores only a representative segment for indexing.
    # Include these constituent rr values so validation matches the actual Stan data.
    for rr in const_rr_man_vec
        if haskey(rr_to_segment_year, rr)
            sy = rr_to_segment_year[rr]
            push!(seg_obs_years[sy.segment_id], sy.year)
        end
    end

    max_internal_gap = 0
    for (i, row) in enumerate(eachrow(segment_info))
        seg_obs = collect(seg_obs_years[row.segment_id])
        if length(seg_obs) > 1
            gaps = diff(sort(unique(seg_obs)))
            if !isempty(gaps)
                max_gap_in_seg = maximum(gaps)
                max_internal_gap = max(max_internal_gap, max_gap_in_seg)
                if max_gap_in_seg > MAX_GAP
                    @warn "Segment $i (party $(row.party_id)) has internal gap of $max_gap_in_seg years!"
                end
            end
        end
    end
    println("  Maximum internal gap within segments: $max_internal_gap years (limit: $MAX_GAP)")

    # Validate all segments have >= MIN_OBS
    for row in eachrow(segment_info)
        @assert row.n_obs >= MIN_OBS "Segment $(row.segment_id) has only $(row.n_obs) obs (min: $MIN_OBS)"
    end
    println("  All segments have >= $MIN_OBS observations: PASS")

    println("\nSegment-based indexing created successfully")
    println("  S (segments): $S")
    println("  R (segment-years): $R")
    println("  J (parties with valid segments): $J")

    return (manifesto=manifesto, expert_dim=expert_dim, expert_lr=expert_lr,
            segment_year=segment_year, segment_info=segment_info,
            all_parties=all_parties, all_groups=all_groups,
            S=S, J=J, P=P, R=R, N_ciy=N_ciy, len_theta_ts=len_theta_ts,
            segment_country_idx=segment_country_idx, group_to_index=group_to_index,
            F=F, segment_family_idx=segment_family_idx, anchor_segment_idx=anchor_segment_ss,
            # Constituent arrays for mean-constituent model
            N_const_man_total=N_const_man_total,
            n_const_man=n_const_man_vec,
            const_offset_man=const_offset_man_vec,
            const_rr_man=const_rr_man_vec,
            N_const_exp_dim_total=N_const_exp_dim_total,
            n_const_exp_dim=n_const_exp_dim_vec,
            const_offset_exp_dim=const_offset_exp_dim_vec,
            const_rr_exp_dim=const_rr_exp_dim_vec,
            N_const_exp_lr_total=N_const_exp_lr_total,
            n_const_exp_lr=n_const_exp_lr_vec,
            const_offset_exp_lr=const_offset_exp_lr_vec,
            const_rr_exp_lr=const_rr_exp_lr_vec,
            union_to_constituents=union_to_constituents,
            constituent_to_union=constituent_to_union)
end

function finalize_4dim_stan_data(manifesto, expert_dim, expert_lr, segment_year, segment_info,
                                all_parties, all_groups, group_to_index, year0, S, J, P, R, N_ciy,
                                len_theta_ts, segment_country_idx, F, segment_family_idx, anchor_segment_idx;
                                N_const_man_total=0, n_const_man=Int[], const_offset_man=Int[], const_rr_man=Int[],
                                N_const_exp_dim_total=0, n_const_exp_dim=Int[], const_offset_exp_dim=Int[], const_rr_exp_dim=Int[],
                                N_const_exp_lr_total=0, n_const_exp_lr=Int[], const_offset_exp_lr=Int[], const_rr_exp_lr=Int[])
    println("Finalizing 4D Stan data structure (V10: segment-based)...")

    # Map years for temporal indexing
    years_all = sort(unique(vcat(manifesto.Year, expert_dim.Year, expert_lr.Year)))
    T_year = length(years_all)
    year_map = DataFrame(year_rel=years_all, year_ix=1:T_year)

    # Apply year mapping to all datasets
    manifesto = leftjoin(manifesto, year_map, on=[:Year => :year_rel])
    rename!(manifesto, :year_ix => :year_for_man)

    expert_dim = leftjoin(expert_dim, year_map, on=[:Year => :year_rel])
    rename!(expert_dim, :year_ix => :year_for_exp_dim)

    expert_lr = leftjoin(expert_lr, year_map, on=[:Year => :year_rel])
    rename!(expert_lr, :year_ix => :year_for_exp_lr)

    # Validate year assignments
    @assert all(.!ismissing.(manifesto.year_for_man))
    @assert all(.!ismissing.(expert_dim.year_for_exp_dim))
    @assert all(.!ismissing.(expert_lr.year_for_exp_lr))
    @assert all(1 .<= manifesto.year_for_man .<= T_year)
    @assert all(1 .<= expert_dim.year_for_exp_dim .<= T_year)
    @assert all(1 .<= expert_lr.year_for_exp_lr .<= T_year)

    # Add small epsilon to prevent exact zeros and ones
    epsilon = 1e-6
    expert_dim.val = clamp.(expert_dim.val, epsilon, 1.0 - epsilon)
    expert_lr.val = clamp.(expert_lr.val, epsilon, 1.0 - epsilon)

    # Calculate prior means
    man_positive_sample = manifesto.positive[manifesto.sample .> 0] ./ manifesto.sample[manifesto.sample .> 0]
    mn_resp_log_man = StatsFuns.logit(mean(man_positive_sample))
    mn_resp_log_exp_dim = StatsFuns.logit(mean(expert_dim.val))
    mn_resp_log_exp_lr = StatsFuns.logit(mean(expert_lr.val))

    println("Prior means calculated:")
    println("  Manifesto: $(round(mn_resp_log_man, digits=3))")
    println("  Expert dimension-specific: $(round(mn_resp_log_exp_dim, digits=3))")
    println("  Expert general L-R: $(round(mn_resp_log_exp_lr, digits=3))")

    # V6: Decade indexing for hierarchical L-R weights
    expert_lr_actual_years = expert_lr.Year .+ year0
    expert_lr_decade_raw = div.(expert_lr_actual_years, 10)
    all_lr_decades = sort(unique(expert_lr_decade_raw))
    lr_decade_to_index = Dict(all_lr_decades .=> 1:length(all_lr_decades))
    dd_exp_lr = [lr_decade_to_index[d] for d in expert_lr_decade_raw]
    D_lr = length(all_lr_decades)
    println("  Decade indexing (V6): $D_lr decades, range $(minimum(all_lr_decades)*10)s-$(maximum(all_lr_decades)*10)s")

    # Create Stan data dictionary - V10 uses S (segments) instead of J (parties)
    dat_4dim = Dict(
        # Common data - V10: S = number of segments
        "S" => S,                          # NEW: Number of segments (was J)
        "J" => J,                          # Keep J for reference (parties with valid segments)
        "P" => P,
        "R" => R,
        "T_year" => T_year,
        "len_theta_ts" => Int.(len_theta_ts),

        # Segment-country mapping (V10: segments inherit country from party)
        "segment_country" => segment_country_idx,

        # Segment family data (V10: segments inherit family from party)
        "F" => F,
        "segment_family" => segment_family_idx,

        # Anchor segment for identification (CDU/CSU segment)
        "anchor_segment" => anchor_segment_idx,

        # Manifesto data - use ss_man (segment index) instead of jj_man (party index)
        "N_man" => nrow(manifesto),
        "K_man" => length(unique(manifesto.var_man)),
        "kk_man" => manifesto.var_man,
        "ss_man" => manifesto.ss_man,      # V10: segment index (was jj_man)
        "rr_man" => manifesto.rr_man,
        "pp_man" => manifesto.pp_man,
        "positive" => manifesto.positive,
        "sample" => manifesto.sample,
        "year_for_man" => manifesto.year_for_man,
        "type_high_idx_man" => manifesto.type_high_idx,
        "type_low_idx_man" => manifesto.type_low_idx,

        # V1 (2D model): dimension index and direction for text data
        "dim_idx_man" => manifesto.dim_idx_man,
        "direction_man" => manifesto.direction_man,

        # Country-item-year data for zero-inflation
        "N_ciy" => N_ciy,
        "ciy_idx" => manifesto.ciy_idx,

        # Expert dimension-specific data
        "N_exp_dim" => nrow(expert_dim),
        "K_exp_dim" => length(unique(expert_dim.var_exp_dim)),
        "kk_exp_dim" => expert_dim.var_exp_dim,
        "ss_exp_dim" => expert_dim.ss_exp_dim,  # V10: segment index
        "rr_exp_dim" => expert_dim.rr_exp_dim,
        "pp_exp_dim" => expert_dim.pp_exp_dim,
        # V5 K-scaling: use rounded sum (mean × K × n_scale) and total trials (K × n_scale)
        "val_dim_int" => Int.(clamp.(round.(expert_dim.val .* expert_dim.n_scale .* expert_dim.n_experts), 0, expert_dim.n_scale .* expert_dim.n_experts)),
        "n_total_exp_dim" => expert_dim.n_scale .* expert_dim.n_experts,
        "n_experts_exp_dim" => expert_dim.n_experts,
        "type_high_idx" => expert_dim.type_high_idx,
        "type_low_idx" => expert_dim.type_low_idx,

        # V1 (2D model): dimension index for expert dimension data
        "dim_idx_exp" => expert_dim.dim_idx_exp,

        # Expert general L-R data
        "N_exp_lr" => nrow(expert_lr),
        "K_exp_lr" => length(unique(expert_lr.var_exp_lr)),
        "kk_exp_lr" => expert_lr.var_exp_lr,
        "ss_exp_lr" => expert_lr.ss_exp_lr,  # V10: segment index
        "rr_exp_lr" => expert_lr.rr_exp_lr,
        "pp_exp_lr" => expert_lr.pp_exp_lr,
        # V5 K-scaling: use rounded sum (mean × K × n_scale) and total trials (K × n_scale)
        "val_lr_int" => Int.(clamp.(round.(expert_lr.val .* expert_lr.n_scale .* expert_lr.n_experts), 0, expert_lr.n_scale .* expert_lr.n_experts)),
        "n_total_exp_lr" => expert_lr.n_scale .* expert_lr.n_experts,
        "n_experts_exp_lr" => expert_lr.n_experts,

        # V6: Decade indexing for hierarchical L-R weights
        "D_lr" => D_lr,
        "dd_exp_lr" => dd_exp_lr,

        # Prior information
        "mn_resp_log_man" => mn_resp_log_man,
        "mn_resp_log_exp_dim" => mn_resp_log_exp_dim,
        "mn_resp_log_exp_lr" => mn_resp_log_exp_lr,

        # Constituent arrays for mean-constituent model (V4)
        "N_const_man_total" => max(1, N_const_man_total),
        "n_const_man" => isempty(n_const_man) ? ones(Int, nrow(manifesto)) : n_const_man,
        "const_offset_man" => isempty(const_offset_man) ? collect(1:nrow(manifesto)) : const_offset_man,
        "const_rr_man" => isempty(const_rr_man) ? manifesto.rr_man : const_rr_man,

        "N_const_exp_dim_total" => max(1, N_const_exp_dim_total),
        "n_const_exp_dim" => isempty(n_const_exp_dim) ? ones(Int, nrow(expert_dim)) : n_const_exp_dim,
        "const_offset_exp_dim" => isempty(const_offset_exp_dim) ? collect(1:nrow(expert_dim)) : const_offset_exp_dim,
        "const_rr_exp_dim" => isempty(const_rr_exp_dim) ? expert_dim.rr_exp_dim : const_rr_exp_dim,

        "N_const_exp_lr_total" => max(1, N_const_exp_lr_total),
        "n_const_exp_lr" => isempty(n_const_exp_lr) ? ones(Int, nrow(expert_lr)) : n_const_exp_lr,
        "const_offset_exp_lr" => isempty(const_offset_exp_lr) ? collect(1:nrow(expert_lr)) : const_offset_exp_lr,
        "const_rr_exp_lr" => isempty(const_rr_exp_lr) ? expert_lr.rr_exp_lr : const_rr_exp_lr
    )

    println("4D Stan data dictionary created with $(length(dat_4dim)) elements")

    # Print summary statistics
    println("\nData summary (V10: Segment-based):")
    println("  Segments: $(dat_4dim["S"])")
    println("  Parties with valid segments: $(dat_4dim["J"])")
    println("  Segment-year combinations: $(dat_4dim["R"])")
    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"])")
    println("  Unique manifesto items: $(dat_4dim["K_man"])")
    println("  Unique expert dimension-specific items: $(dat_4dim["K_exp_dim"])")
    println("  Unique expert general L-R items: $(dat_4dim["K_exp_lr"])")
    println("  Years: $(dat_4dim["T_year"])")

    return (dat_4dim=dat_4dim, manifesto=manifesto, expert_dim=expert_dim,
            expert_lr=expert_lr, T_year=T_year, segment_year=segment_year,
            segment_info=segment_info)
end

# Execute if run directly
if abspath(PROGRAM_FILE) == @__FILE__
    println("Run from main script to execute the full 4D pipeline")
end
