Code Reference¶

def build_crosswalk(cycles: list, descriptions: dict, cutoff: float = 0.6) -> dict:
    """
    Build a crosswalk mapping reference-cycle variable names to each cycle's
    matching variable name, by fuzzy-matching variable descriptions.

    The last entry in `cycles` is treated as the reference cycle: every
    other cycle's variable is matched against each reference variable's
    description via difflib.get_close_matches (single best match). This is
    a heuristic textual match, not an authoritative concordance - unmatched
    or ambiguous descriptions are recorded as None and should be reviewed.

    Args:
        cycles: Cycle years in order, with the reference cycle last
        descriptions: Dict mapping cycle -> {variable_name: description}
        cutoff: difflib similarity cutoff (0-1) for a match to count

    Returns:
        Dict mapping reference_var -> {cycle: cycle_specific_var_or_None}
    """
    reference_cycle = cycles[-1]
    reference_vars = descriptions[reference_cycle]

    crosswalk = {}
    for ref_var, ref_desc in reference_vars.items():
        crosswalk[ref_var] = {reference_cycle: ref_var}
        for cycle in cycles[:-1]:
            candidates = descriptions[cycle]
            best_match = difflib.get_close_matches(ref_desc, candidates.values(), n=1, cutoff=cutoff)
            if best_match:
                for var, desc in candidates.items():
                    if desc == best_match[0]:
                        crosswalk[ref_var][cycle] = var
                        break
            else:
                crosswalk[ref_var][cycle] = None

    return crosswalk

def auto_harmonize(label: str) -> str:
    """
    Normalize a raw codebook category label into a shared harmonized category.

    Order matters: "not stated", "valid skip", "don't know", and "female"
    each contain "no" or "male" as a substring (e.g. "**no**t stated",
    "fe**male**"), so the more specific phrases must be checked before the
    shorter "no"/"male" rules or they get misclassified.
    """
    l = label.lower()
    if "not stated" in l:
        return "Not stated"
    if "valid skip" in l:
        return "Valid skip"
    if "don’t know" in l or "don't know" in l:
        return "Don't know"
    if "female" in l:
        return "Female"
    if "male" in l:
        return "Male"
    if "yes" in l:
        return "Yes"
    if "no" in l:
        return "No"
    return label.strip()

def get_common_harmonized_vars(cycles: list, crosswalk: dict, data_dict: dict) -> list:
    """
    Get harmonized variables that exist in all specified cycles.

    Args:
        cycles: List of cycle years
        crosswalk: Crosswalk dictionary
        data_dict: Dictionary mapping cycle -> DataFrame (with cycle-specific column names)

    Returns:
        List of harmonized variable names available in all cycles
    """
    common_vars = []

    for harmonized_var, cycle_mapping in crosswalk.items():
        available_in_all = True
        for cycle in cycles:
            cycle_specific_var = cycle_mapping.get(cycle)
            if not cycle_specific_var or cycle_specific_var not in data_dict[cycle].columns:
                available_in_all = False
                break

        if available_in_all:
            common_vars.append(harmonized_var)

    return common_vars

def parse_codebook_lines(lines: Iterable[str]) -> dict:
    """
    Parse codebook text lines into a variable -> {description, categories} dict.

    Expects the Statistics Canada CCHS Data Dictionary/Freqs layout: each
    variable starts with a "Variable Name:" line, has a "Concept:" line for
    its description, and an "Answer Categories" section listing "label code
    frequency frequency%" rows until a blank line or a Note:/Source:/
    Universe:/Total line ends the section.

    Args:
        lines: An iterable of text lines, in document order (may span
            multiple PDF pages - state carries across page boundaries the
            same way it does across lines within a page).

    Returns:
        Dict mapping variable_name -> {"description": str, "categories": {code: label}}
    """
    variables = {}
    current_variable = None
    current_description = None
    codes = {}
    in_answer_categories = False

    for line in lines:
        if "Variable Name:" in line:
            if current_variable:
                variables[current_variable] = {
                    "description": current_description,
                    "categories": codes.copy() if codes else {}
                }
            current_variable = line.split(":")[1].strip().replace(" Length", "")
            current_description = None
            codes = {}
            in_answer_categories = False
        elif "Concept:" in line:
            current_description = line.split(":")[1].strip()
        elif "Answer Categories" in line:
            in_answer_categories = True
            continue
        elif in_answer_categories:
            match = CATEGORY_ROW.match(line)
            if match:
                meaning = match.group(1).strip()
                code = match.group(2).strip()
                codes[code] = meaning
            elif line.strip() == "" or SECTION_END.match(line):
                in_answer_categories = False
                continue
        elif "Note:" in line or "Source:" in line or "Universe:" in line:
            in_answer_categories = False
            continue

    if current_variable:
        variables[current_variable] = {
            "description": current_description,
            "categories": codes.copy() if codes else {}
        }

    return variables

def extract_variables_with_categories(pdf_path) -> dict:
    """Extract variable/category metadata from a CCHS codebook PDF."""
    from pypdf import PdfReader

    lines = []
    with open(pdf_path, "rb") as file:
        reader = PdfReader(file)
        for page in reader.pages:
            lines.extend(page.extract_text().splitlines())

    return parse_codebook_lines(lines)

def build_harmonization_mapping(crosswalk: dict, cycle: str, available_columns) -> tuple[dict, list]:
    """
    Build a safe rename mapping for one cycle.

    If multiple harmonized variables point at the same source column, prefer the
    identity mapping (e.g. GEODVHR4 -> GEODVHR4) and otherwise keep the first
    mapping encountered. This prevents exact-match geography columns from being
    renamed away by alias entries later in the crosswalk.
    """
    rename_dict = {}
    available_vars = []
    available_column_set = set(available_columns)

    for harmonized_var, cycle_mapping in crosswalk.items():
        cycle_specific_var = cycle_mapping.get(cycle)
        if not cycle_specific_var or cycle_specific_var == "Not Available":
            continue
        if cycle_specific_var not in available_column_set:
            continue

        existing_target = rename_dict.get(cycle_specific_var)
        if existing_target is None:
            rename_dict[cycle_specific_var] = harmonized_var
            available_vars.append(harmonized_var)
            continue

        if existing_target == cycle_specific_var:
            continue

        if harmonized_var == cycle_specific_var:
            rename_dict[cycle_specific_var] = harmonized_var
            if existing_target in available_vars:
                available_vars.remove(existing_target)
            if harmonized_var not in available_vars:
                available_vars.append(harmonized_var)

    return rename_dict, available_vars

def restore_geography_aliases(df: pd.DataFrame) -> pd.DataFrame:
    """
    Restore expected geography columns when an older harmonization pass renamed
    them to alias fields.
    """
    restored = df.copy()

    if 'GEODVHR4' not in restored.columns and 'GEODVOHR' in restored.columns:
        restored['GEODVHR4'] = restored['GEODVOHR']

    return restored

@st.cache_data
def load_cycle_data(cycle: str) -> tuple[pd.DataFrame, pd.DataFrame]:
    """Load main data and bootstrap data for a specific cycle."""
    data_file = os.path.join(DATA_PATH, f"hs{cycle}_on_distr.parquet")
    bootstrap_file = os.path.join(DATA_PATH, f"hs{cycle}_on_bootwt.parquet")

    if os.path.exists(data_file) and os.path.exists(bootstrap_file):
        data = pd.read_parquet(data_file)
        bootstrap_data = pd.read_parquet(bootstrap_file)
        return data, bootstrap_data
    else:
        st.error(f"Data files for cycle {cycle} are missing. Please check the 'data/' directory.")
        return None, None

@st.cache_data
def load_variable_descriptions(cycle: str) -> tuple[pd.DataFrame, dict]:
    """Load variable descriptions from harmonized JSON file for a specific cycle."""
    try:
        # Load from harmonized JSON file first
        json_file = os.path.join("harmonization", f"CCHS_{cycle}.json")
        if os.path.exists(json_file):
            with open(json_file, "r", encoding="utf-8") as f:
                var_dict = json.load(f)

            # Extract descriptions from JSON structure
            descriptions_data = []
            desc_dict = {}

            for var_name, var_info in var_dict.items():
                description = var_info.get("description", "")
                if description:
                    descriptions_data.append({
                        'Variable': var_name,
                        'Description': description
                    })
                    desc_dict[var_name] = description

            # Ensure all variables are included, even if description is missing
            for var_name in var_dict.keys():
                if var_name not in desc_dict:
                    descriptions_data.append({
                        'Variable': var_name,
                        'Description': "No description available"
                    })
                    desc_dict[var_name] = "No description available"

            # Create DataFrame from extracted data
            if descriptions_data:
                desc_df = pd.DataFrame(descriptions_data)
                return desc_df, desc_dict

        # Fallback to CSV file if JSON doesn't exist or is empty
        desc_file = os.path.join(DATA_PATH, f"CCHS_{cycle}_Recoded_Variables.csv")
        if os.path.exists(desc_file):
            descriptions = pd.read_csv(desc_file)
            desc_dict = dict(zip(descriptions['Variable'], descriptions['Description']))
            return descriptions, desc_dict

        # If neither exists, return empty
        return None, {}

    except Exception as e:
        st.error(f"Error loading variable descriptions for cycle {cycle}: {e}")
        return None, {}

@st.cache_data
def load_json_variable_descriptions(cycle: str) -> dict:
    """Load JSON variable descriptions for a specific cycle (legacy function for compatibility)."""
    json_file = os.path.join("harmonization", f"CCHS_{cycle}.json")
    if os.path.exists(json_file):
        with open(json_file, "r", encoding="utf-8") as f:
            var_dict = json.load(f)
        return {k: v.get("description", "") for k, v in var_dict.items()}
    return {}

@st.cache_data
def load_cycle_variable_info(cycle: str) -> dict:
    """Load complete variable information (descriptions AND categories) for a specific cycle."""
    json_file = os.path.join("harmonization", f"CCHS_{cycle}.json")
    if os.path.exists(json_file):
        with open(json_file, "r", encoding="utf-8") as f:
            return json.load(f)
    return {}

@st.cache_data
def load_crosswalk() -> dict:
    """Load harmonization crosswalk."""
    crosswalk_file = os.path.join("harmonization", "crosswalk.json")
    if os.path.exists(crosswalk_file):
        with open(crosswalk_file, "r", encoding="utf-8") as f:
            return json.load(f)
    return {}

@st.cache_data
def load_categories() -> dict:
    """Load harmonization categories."""
    categories_file = os.path.join("harmonization", "categories.json")
    if os.path.exists(categories_file):
        with open(categories_file, "r", encoding="utf-8") as f:
            return json.load(f)
    return {}

@st.cache_data
def load_ontario_csd_lookup() -> dict:
    """Load the Ontario census subdivision code-to-name lookup."""
    lookup_file = os.path.join("harmonization", "ontario_csd_lookup.json")
    if os.path.exists(lookup_file):
        with open(lookup_file, "r", encoding="utf-8") as f:
            return json.load(f)
    return {}

@st.cache_data
def load_ontario_official_municipalities() -> dict:
    """Load the official Ontario municipalities lookup keyed by CSD code."""
    lookup_file = os.path.join("harmonization", "ontario_official_municipalities.json")
    if os.path.exists(lookup_file):
        with open(lookup_file, "r", encoding="utf-8") as f:
            return json.load(f)
    return {}

@st.cache_data
def merge_data(filtered_data, bootstrap_data):
    """Merge survey records with bootstrap weights without pooling cycles.

    Multi-cycle trend analysis keeps each survey cycle independent. When a
    ``CYCLE`` column is present, it must be present on both frames and becomes
    part of the join key. The one-to-one validation prevents repeated IDs from
    silently multiplying records and biasing estimates.
    """
    survey_has_cycle = 'CYCLE' in filtered_data.columns
    bootstrap_has_cycle = 'CYCLE' in bootstrap_data.columns

    if survey_has_cycle != bootstrap_has_cycle:
        raise ValueError(
            "Cycle-aware analysis requires a CYCLE column on both survey and "
            "bootstrap data. Regenerate precomputed files if necessary."
        )

    join_columns = ['ONT_ID']
    if survey_has_cycle:
        join_columns.insert(0, 'CYCLE')

    return pd.merge(
        filtered_data,
        bootstrap_data,
        on=join_columns,
        how='left',
        validate='one_to_one',
    )

@st.cache_data
def load_multi_cycle_data(cycles: list, crosswalk: dict, categories: dict):
    """
    Load and harmonize data from multiple cycles using pre-computed crosswalk.
    This uses simple column renaming (no value transformation) for performance.

    Args:
        cycles: List of cycle years to load (e.g., ["2021", "2022", "2023"])
        crosswalk: Crosswalk dictionary mapping harmonized_var -> {cycle: cycle_specific_var}
        categories: Categories dictionary (not used, kept for compatibility)

    Returns:
        Tuple of (combined_data, combined_bootstrap_data) or (None, None) if error
    """
    if not cycles:
        st.error("No cycles specified for multi-cycle loading.")
        return None, None

    if not crosswalk:
        st.warning("No crosswalk provided. Loading cycles without harmonization.")

    combined_data_list = []
    combined_bootstrap_list = []

    for cycle in cycles:
        # Load raw data for this cycle
        data, bootstrap_data = load_cycle_data(cycle)
        if data is None or bootstrap_data is None:
            st.warning(f"Skipping cycle {cycle} due to missing data files.")
            continue

        rename_dict = {}
        if crosswalk:
            rename_dict, _ = build_harmonization_mapping(crosswalk, cycle, data.columns)

        # Apply harmonization - just column rename, no value transformation (fast!)
        harmonized_data = data.rename(columns=rename_dict) if rename_dict else data.copy()
        harmonized_data = restore_geography_aliases(harmonized_data)
        harmonized_data['CYCLE'] = cycle

        cycle_bootstrap = bootstrap_data.copy()
        cycle_bootstrap['CYCLE'] = cycle

        combined_data_list.append(harmonized_data)
        combined_bootstrap_list.append(cycle_bootstrap)

    if not combined_data_list:
        st.error("No valid cycles could be loaded.")
        return None, None

    # Before the concat line, check for and remove duplicate columns
    for i, df in enumerate(combined_data_list):
        duplicates = df.columns[df.columns.duplicated()].unique()
        if len(duplicates) > 0:
            import streamlit as st
            st.warning(f"Cycle {cycles[i]} has duplicate columns: {duplicates.tolist()}")
            print(f"Duplicate columns in cycle {cycles[i]}: {duplicates.tolist()}")
            # Remove duplicate columns, keeping first occurrence
            combined_data_list[i] = df.loc[:, ~df.columns.duplicated()]

    # Same check for bootstrap data if it exists
    if combined_bootstrap_list:
        for i, df in enumerate(combined_bootstrap_list):
            duplicates = df.columns[df.columns.duplicated()].unique()
            if len(duplicates) > 0:
                import streamlit as st
                st.warning(f"Cycle {cycles[i]} bootstrap data has duplicate columns: {duplicates.tolist()}")
                combined_bootstrap_list[i] = df.loc[:, ~df.columns.duplicated()]

    # Combine all cycles (fast concat operation)
    combined_data = pd.concat(combined_data_list, ignore_index=True)
    combined_bootstrap = pd.concat(combined_bootstrap_list, ignore_index=True)

    return combined_data, combined_bootstrap

@st.cache_data
def load_data():
    """Legacy function - loads 2021 data by default for backward compatibility."""
    return load_cycle_data("2021")

def check_precompute_status(cycles: List[str], save_dir: Path = PRECOMPUTE_DIR) -> Dict[str, bool]:
    """
    Check which cycles have been precomputed.

    Args:
        cycles: List of cycle years (e.g., ["2021", "2022", "2023"])
        save_dir: Directory where precomputed data is stored

    Returns:
        Dictionary mapping cycle -> True/False (precomputed status)
    """
    status = {}

    for cycle in cycles:
        data_path = save_dir / f"harmonized_data_{cycle}.parquet"
        bootstrap_path = save_dir / f"harmonized_bootstrap_{cycle}.parquet"
        meta_path = save_dir / f"metadata_{cycle}.json"
        status[cycle] = (
            data_path.exists()
            and bootstrap_path.exists()
            and meta_path.exists()
        )

    return status

def load_precomputed_data(cycle: str, save_dir: Path = PRECOMPUTE_DIR) -> pd.DataFrame:
    """
    Load precomputed harmonized data for a cycle.

    Args:
        cycle: Cycle year (e.g., "2021")
        save_dir: Directory where precomputed data is stored

    Returns:
        DataFrame with harmonized variables

    Raises:
        FileNotFoundError: If precomputed data doesn't exist
    """
    path = save_dir / f"harmonized_data_{cycle}.parquet"

    if not path.exists():
        raise FileNotFoundError(f"Precomputed data not found: {path}")

    return pd.read_parquet(path)

def load_precomputed_bootstrap(cycle: str, save_dir: Path = PRECOMPUTE_DIR) -> pd.DataFrame:
    """
    Load precomputed bootstrap weights for a cycle.

    Args:
        cycle: Cycle year (e.g., "2021")
        save_dir: Directory where precomputed data is stored

    Returns:
        DataFrame with bootstrap weights (ONT_ID and BSW columns)

    Raises:
        FileNotFoundError: If precomputed data doesn't exist
    """
    path = save_dir / f"harmonized_bootstrap_{cycle}.parquet"

    if not path.exists():
        raise FileNotFoundError(f"Precomputed bootstrap data not found: {path}")

    return pd.read_parquet(path)

def load_precomputed_metadata(cycle: str, save_dir: Path = PRECOMPUTE_DIR) -> Dict:
    """
    Load metadata for precomputed cycle.

    Args:
        cycle: Cycle year (e.g., "2021")
        save_dir: Directory where precomputed data is stored

    Returns:
        Dictionary with metadata (available_vars, record_count, etc.)

    Raises:
        FileNotFoundError: If metadata doesn't exist
    """
    path = save_dir / f"metadata_{cycle}.json"

    if not path.exists():
        raise FileNotFoundError(f"Metadata not found: {path}")

    with open(path, 'r', encoding='utf-8') as f:
        return json.load(f)

def get_common_variables(cycles: List[str], save_dir: Path = PRECOMPUTE_DIR) -> List[str]:
    """
    Get variables available across all selected cycles using precomputed metadata.

    Args:
        cycles: List of cycle years
        save_dir: Directory where precomputed data is stored

    Returns:
        Sorted list of common harmonized variable names
    """
    if not cycles:
        return []

    # Load metadata for each cycle
    all_vars = []
    for cycle in cycles:
        try:
            meta = load_precomputed_metadata(cycle, save_dir)
            all_vars.append(set(meta['available_vars']))
        except FileNotFoundError:
            return []  # Missing precomputed data

    # Find intersection
    common = set.intersection(*all_vars) if all_vars else set()
    return sorted(list(common))

@st.cache_data
def precompute_cycle_data(
    cycle: str,
    crosswalk: Dict,
    data_path: str = "data",
    save_dir: Path = PRECOMPUTE_DIR
) -> Tuple[pd.DataFrame, pd.DataFrame, Dict]:
    """
    Precompute harmonized data for a single cycle.

    Creates:
    - harmonized_data_{cycle}.parquet: Harmonized survey data
    - harmonized_bootstrap_{cycle}.parquet: Bootstrap weights with ONT_ID
    - metadata_{cycle}.json: Variable metadata and availability

    Category value labels are not baked in here - they're resolved at
    display time from categories.json via get_value_label().

    Args:
        cycle: Cycle year (e.g., "2021")
        crosswalk: Crosswalk dictionary for variable harmonization
        data_path: Path to raw data files
        save_dir: Directory to save precomputed files

    Returns:
        Tuple of (harmonized_data, harmonized_bootstrap, metadata)
    """
    import os
    print(f"Precomputing {cycle}...")

    # Create output directory if needed
    save_dir.mkdir(parents=True, exist_ok=True)

    # Load raw data
    data_file = os.path.join(data_path, f"hs{cycle}_on_distr.parquet")
    bootstrap_file = os.path.join(data_path, f"hs{cycle}_on_bootwt.parquet")

    if not os.path.exists(data_file) or not os.path.exists(bootstrap_file):
        raise FileNotFoundError(f"Raw data files not found for cycle {cycle}")

    data = pd.read_parquet(data_file)
    bootstrap_data = pd.read_parquet(bootstrap_file)

    # Create harmonization mapping for this cycle
    rename_dict, available_vars = build_harmonization_mapping(
        crosswalk, cycle, data.columns
    )

    # Apply harmonization (rename columns)
    harmonized_data = data.rename(columns=rename_dict)
    harmonized_data = restore_geography_aliases(harmonized_data)

    # Handle duplicate columns (keep first occurrence only)
    # This can happen when multiple cycle-specific variables map to same harmonized name
    if harmonized_data.columns.duplicated().any():
        duplicate_cols = harmonized_data.columns[harmonized_data.columns.duplicated()].unique().tolist()
        print(f"   ⚠️  Warning: Removing {len(duplicate_cols)} duplicate columns: {duplicate_cols[:10]}{'...' if len(duplicate_cols) > 10 else ''}")
        harmonized_data = harmonized_data.loc[:, ~harmonized_data.columns.duplicated(keep='first')]

    # Add CYCLE column for multi-cycle identification
    harmonized_data['CYCLE'] = cycle

    # Keep only harmonized columns plus core variables
    core_vars = ['ONT_ID', 'WTS_S', 'GEODVHR4', 'GEODVCSD']
    # Add age column (varies by cycle)
    if 'DHH_AGE' in harmonized_data.columns:
        core_vars.append('DHH_AGE')
    if 'AWCAGE' in harmonized_data.columns:
        core_vars.append('AWCAGE')

    # Keep harmonized vars + core vars that exist
    cols_to_keep = list(set(available_vars + core_vars + ['CYCLE']) & set(harmonized_data.columns))
    harmonized_data = harmonized_data[cols_to_keep]

    # Cycle-qualified bootstrap data keeps trend comparisons independent and
    # prevents repeated ONT_ID values in different years from cross-joining.
    harmonized_bootstrap = bootstrap_data.copy()
    harmonized_bootstrap['CYCLE'] = cycle

    # Save harmonized data
    output_path = save_dir / f"harmonized_data_{cycle}.parquet"
    harmonized_data.to_parquet(output_path, index=False)

    # Save harmonized bootstrap
    bootstrap_output_path = save_dir / f"harmonized_bootstrap_{cycle}.parquet"
    harmonized_bootstrap.to_parquet(bootstrap_output_path, index=False)

    # Save metadata
    metadata = {
        'cycle': cycle,
        'available_vars': available_vars,
        'record_count': len(harmonized_data),
        'core_vars': core_vars,
        'precompute_date': pd.Timestamp.now().isoformat(),
        'harmonization_mapping': rename_dict
    }

    metadata_path = save_dir / f"metadata_{cycle}.json"
    with open(metadata_path, 'w', encoding='utf-8') as f:
        json.dump(metadata, f, indent=2)

    print(f"✅ {cycle}: {len(available_vars)} variables, {len(harmonized_data):,} records")
    print(f"   Saved to: {output_path}")

    return harmonized_data, harmonized_bootstrap, metadata

def precompute_all_cycles(
    cycles: List[str],
    crosswalk: Dict,
    data_path: str = "data",
    save_dir: Path = PRECOMPUTE_DIR
) -> Dict[str, Dict]:
    """
    Precompute data for all specified cycles.

    Args:
        cycles: List of cycle years to precompute
        crosswalk: Crosswalk dictionary
        data_path: Path to raw data files
        save_dir: Directory to save precomputed files

    Returns:
        Dictionary mapping cycle -> {'data': DataFrame, 'bootstrap': DataFrame, 'metadata': Dict}
    """
    results = {}

    for cycle in cycles:
        try:
            data, bootstrap, meta = precompute_cycle_data(
                cycle, crosswalk, data_path, save_dir
            )
            results[cycle] = {
                'data': data,
                'bootstrap': bootstrap,
                'metadata': meta
            }
        except Exception as e:
            print(f"❌ Failed to precompute {cycle}: {e}")
            import traceback
            traceback.print_exc()

    return results

def run_precompute_workflow(
    cycles: List[str],
    crosswalk: Dict,
    data_path: str = "data",
    save_dir: Path = PRECOMPUTE_DIR
) -> Dict:
    """
    Run the full precompute workflow and return a compact status summary.

    Args:
        cycles: List of cycle years to precompute
        crosswalk: Crosswalk dictionary
        data_path: Path to raw data files
        save_dir: Directory to save precomputed files

    Returns:
        Dictionary with success flag, validation results, and per-cycle summary.
    """
    results = precompute_all_cycles(
        cycles=cycles,
        crosswalk=crosswalk,
        data_path=data_path,
        save_dir=save_dir,
    )

    if not results:
        return {
            "success": False,
            "results": {},
            "validation": {},
            "message": "No cycles were successfully precomputed.",
        }

    validation = validate_precomputed_data(list(results.keys()), save_dir)

    try:
        create_variable_availability_index(
            cycles=list(results.keys()),
            save_dir=save_dir,
        )
    except Exception as exc:
        print(f"⚠️ Failed to create availability index: {exc}")

    all_valid = all(validation.values()) if validation else False
    return {
        "success": all_valid,
        "results": {
            cycle: {
                "record_count": result["metadata"]["record_count"],
                "available_vars": len(result["metadata"]["available_vars"]),
            }
            for cycle, result in results.items()
        },
        "validation": validation,
        "message": "Precompute completed successfully." if all_valid else "Precompute completed, but validation failed for one or more cycles.",
    }

def load_variable_availability_index(save_dir: Path = PRECOMPUTE_DIR) -> Dict[str, List[str]]:
    """
    Load the variable availability index (which variables exist in which cycles).

    Args:
        save_dir: Directory where precomputed data is stored

    Returns:
        Dictionary mapping variable_name -> [list of cycles where it exists]
    """
    index_path = save_dir / "variable_availability.json"

    if not index_path.exists():
        raise FileNotFoundError(f"Variable availability index not found: {index_path}")

    with open(index_path, 'r') as f:
        return json.load(f)