API

def CharlsonComorbidityIndex(
    df: pl.DataFrame,
    id_col: str = "id",
    code_col: str = "code",
    age_col: str = "age",
    icd_version: str = "icd10",
    icd_version_col: str | None = None,
    year_col: str | None = None,
    implementation: str = "quan",
    weights: str = "charlson",
    return_categories: bool = False,
) -> pl.DataFrame:
    """Calculate the Charlson Comorbidity Index (CCI) using ICD codes and age.

    Args:
        df (pl.DataFrame): Input data containing at least ``id_col``, ``code_col``, and ``age_col``.
        id_col (str, optional): Column name containing unique identifiers. Defaults to ``"id"``.
        code_col (str, optional): Column name containing ICD codes. Defaults to ``"code"``.
        age_col (str, optional): Column name containing patient ages (age points are only added for ``weights="charlson"``). Defaults to ``"age"``.
        icd_version (str, optional): ICD version; one of ``"icd9"``, ``"icd10"``, or ``"icd9_10"``. Defaults to ``"icd10"``.
        icd_version_col (str, optional): Column name with ICD version labels when ``icd_version`` is ``"icd9_10"``. Defaults to ``None``.
        year_col (str, optional): Column name with each record's ICD-10-GM catalogue year. Required when ``implementation`` is ``"sokolowski"`` (ignored otherwise), since that mapping is year-specific. Defaults to ``None``.
        implementation (str, optional): Definition set to use; ``"quan"``, ``"deyo"``, ``"romano"``, ``"dhoore"``, ``"australia"``, ``"sweden"``, ``"rcs"``, ``"uk_shmi"``, ``"sokolowski"``, or ``"thygesen"``. Defaults to ``"quan"``.
        weights (str, optional): Weighting scheme; one of ``"charlson"``, ``"schneeweiss"``, ``"rcs"``, ``"quan"``, or ``"uk_shmi"``. The ``"rcs"`` and ``"uk_shmi"`` implementations always use their own weights. Defaults to ``"charlson"``.
        return_categories (bool, optional): If ``True``, includes indicator columns for each CCI category. Defaults to ``False``.

    Returns:
        pl.DataFrame: DataFrame containing ``id_col``, the ``"Charlson Age-Comorbidity Score"`` column, and, when ``return_categories`` is ``True``, category indicators and the ``"Age Score"`` column.

    Raises:
        AssertionError: If ``implementation`` or ``weights`` is unsupported, or if required columns are missing.
    """

    icd_version = coerce_icd_version(icd_version, implementation, ICD9_ONLY, ICD10_ONLY)

    # Input validation specific to Charlson
    assert implementation in DEFINITION_FILES, f"implementation must be one of: {list(DEFINITION_FILES)}."
    assert weights in [
        "charlson",
        "schneeweiss",
        "rcs",
        "quan",
        "uk_shmi",
    ], "weights must be one of: 'charlson', 'schneeweiss', 'rcs', 'quan', or 'uk_shmi'."
    assert age_col in df.columns, f"Column '{age_col}' (age) must be present in input DataFrame." # fmt: skip
    if implementation == "sokolowski":
        assert year_col is not None and year_col in df.columns, "Implementation 'sokolowski' requires a 'year_col' column (ICD-10-GM catalogue year) in the input DataFrame." # fmt: skip

    # STEP 0: diagnoses are handled by CustomComorbidityIndex

    # STEP 1: Calculate Age Score separately (on one age per patient, not per diagnosis row)
    #  < 50: 0
    #  < 60: 1
    #  < 70: 2
    #  < 80: 3
    # >= 80: 4
    age_scores = (
        df.group_by(id_col)
        # an unknown age (null or NaN) gives no age points
        .agg(pl.col(age_col).cast(pl.Float64).fill_nan(None).max().fill_null(0))
        .with_columns(
            pl.when(pl.col(age_col) < 50)
            .then(pl.lit(0))
            .when(pl.col(age_col) < 60)
            .then(pl.lit(1))
            .when(pl.col(age_col) < 70)
            .then(pl.lit(2))
            .when(pl.col(age_col) < 80)
            .then(pl.lit(3))
            .otherwise(pl.lit(4))
            .cast(int)
            .alias("Age Score")
        )
        .select(id_col, "Age Score")
    )

    # STEP 2: Calculate Comorbidity Score using generalized function
    # Determine weight column and score column names based on weights argument
    # rcs/uk_shmi implementations always use their own matching weights
    if implementation in ("rcs", "uk_shmi") and weights != implementation:
        warnings.warn(
            f"Implementation '{implementation}' requires '{implementation}_weights'. Overriding weights='{weights}' with '{implementation}_weights'.",
            UserWarning,
            stacklevel=2,
        )
        weights = implementation
    weight_col_name = f"{weights}_weights"

    df_combined, mutual_exclusion_rules = weighted_definitions(
        "CHARLSON", DEFINITION_FILES[implementation], implementation, CANONICAL_RULES
    )

    df_charlson = CustomComorbidityIndex(
        df=df,
        id_col=id_col,
        code_col=code_col,
        icd_version=icd_version,
        icd_version_col=icd_version_col,
        year_col=year_col if implementation == "sokolowski" else None,
        definition_data=df_combined,
        weight_col_name=weight_col_name,
        score_col_name=SCORE_COL_NAME,
        mutual_exclusion_rules=mutual_exclusion_rules,
        return_categories=return_categories,
    )

    # UK SHMI: a negative Charlson Comorbidity Index is assigned a value of zero
    if weights == "uk_shmi":
        df_charlson = df_charlson.with_columns(pl.col(SCORE_COL_NAME).clip(lower_bound=0))

    # STEP 3: Combine Age Score and Comorbidity Score
    # Only the Charlson weights include age points
    if weights != "charlson":
        age_scores = age_scores.with_columns(pl.lit(0).alias("Age Score"))

    df_charlson = df_charlson.join(
        age_scores, on=id_col, how="left", coalesce=True
    ).with_columns(
        (pl.col(SCORE_COL_NAME) + pl.col("Age Score")).alias(
            FINAL_SCORE_COL_NAME
        )
    )

    # Drop Age Score column if not needed
    if not return_categories:
        df_charlson = df_charlson.drop("Age Score")

    return df_charlson

def CharlsonComorbidity_10year_survival(
    df: pl.DataFrame = None,
    id_col: str = "id",
    code_col: str = "code",
    age_col: str = "age",
    icd_version: str = "icd10",
    icd_version_col: str | None = None,
    implementation: str = "quan",
    precalculated_df: bool = False,
) -> pl.DataFrame:
    """Estimate 10-year survival probabilities from the Charlson Comorbidity Index.

    Args:
        df (pl.DataFrame | None, optional): Source data with ICD codes, ages, and identifiers, or a precomputed Charlson score table when ``precalculated_df`` is ``True``.
        id_col (str, optional): Column name containing unique identifiers. Defaults to ``"id"``.
        code_col (str, optional): Column name containing ICD codes. Defaults to ``"code"``.
        age_col (str, optional): Column name containing patient ages. Defaults to ``"age"``.
        icd_version (str, optional): ICD version; one of ``"icd9"``, ``"icd10"``, or ``"icd9_10"``. Defaults to ``"icd10"``.
        icd_version_col (str | None, optional): Column name with ICD version labels when ``icd_version`` is ``"icd9_10"``.
        implementation (str, optional): Charlson definition set passed to :func:`CharlsonComorbidityIndex`. Defaults to ``"quan"``.
        precalculated_df (bool, optional): If ``True``, assumes ``df`` already includes a ``"Charlson Age-Comorbidity Score"`` column. Defaults to ``False``.

    Returns:
        pl.DataFrame: DataFrame with two columns: ``id_col`` and ``"10-year survival probability"``, calculated as ``0.983 ** exp(0.9 * CCI)`` from the age-adjusted Charlson score (Charlson et al. 1994).

    Raises:
        AssertionError: If required columns are missing when ``precalculated_df`` is ``True``.
    """

    if not precalculated_df:
        df = CharlsonComorbidityIndex(
            df=df,
            id_col=id_col,
            code_col=code_col,
            age_col=age_col,
            icd_version=icd_version,
            icd_version_col=icd_version_col,
            implementation=implementation,
            return_categories=False,
        )
    else:
        assert (
            FINAL_SCORE_COL_NAME in df.columns
        ), f"Input DataFrame must contain column '{FINAL_SCORE_COL_NAME}'."
        assert (
            id_col in df.columns
        ), f"Input DataFrame must contain column '{id_col}'."

    # Charlson et al. 1994: 10-year survival = 0.983 ^ exp(0.9 * CCI)
    return df.with_columns(
        (pl.lit(0.983) ** (0.9 * pl.col(FINAL_SCORE_COL_NAME)).exp())
        .round(3)
        .alias("10-year survival probability")
    ).select(id_col, "10-year survival probability")

def ElixhauserComorbidityIndex(
    df: pl.DataFrame,
    id_col: str = "id",
    code_col: str = "code",
    icd_version: str = "icd10",
    icd_version_col: str | None = None,
    year_col: str | None = None,
    implementation: str = "quan",
    weights: str = "van_walraven",
    return_categories=False,
) -> pl.DataFrame:
    """Calculate the Elixhauser Comorbidity Index (ECI) using ICD codes.

    Args:
        df (pl.DataFrame): Input data containing at least ``id_col`` and ``code_col``.
        id_col (str, optional): Column name containing unique identifiers. Defaults to ``"id"``.
        code_col (str, optional): Column name containing ICD codes. Defaults to ``"code"``.
        icd_version (str, optional): ICD version; one of ``"icd9"``, ``"icd10"``, or ``"icd9_10"``. Defaults to ``"icd10"``.
        icd_version_col (str, optional): Column name with ICD version labels when ``icd_version`` is ``"icd9_10"``. Defaults to ``None``.
        year_col (str, optional): Column name with each record's diagnosis year. Required when ``implementation`` is ``"ahrq_icd10"`` (ignored otherwise), since its POA-exemption changelog is year-specific. Defaults to ``None``.
        implementation (str, optional): Definition set to use; ``"quan"``, ``"elixhauser"``, ``"ahrq_icd9"``, or ``"ahrq_icd10"``. Defaults to ``"quan"``.
        weights (str, optional): Weighting scheme; one of ``"van_walraven"``, ``"thompson_30"``, ``"thompson_29"``, ``"ahrq"``, or ``"swiss"``. ``"ahrq"`` uses ``icd_version`` to pick ``ahrq_icd9_weights``/``ahrq_icd10_weights``; ``implementation="ahrq_icd9"``/``"ahrq_icd10"`` enforce it regardless of this argument. Defaults to ``"van_walraven"``.
        return_categories (bool, optional): If ``True``, includes indicator columns for each Elixhauser category. Defaults to ``False``.

    Returns:
        pl.DataFrame: DataFrame containing ``id_col``, the calculated score column, and, when ``return_categories`` is ``True``, category indicators.

    Raises:
        AssertionError: If ``implementation`` or ``weights`` is outside the supported values, or if ``implementation="ahrq_icd10"`` and ``year_col`` is missing.
        ValueError: If ``weights="ahrq"`` is used with ``icd_version="icd9_10"``.
    """

    icd_version = coerce_icd_version(icd_version, implementation, ICD9_ONLY, ICD10_ONLY)

    # Input validation specific to Elixhauser
    assert implementation in DEFINITION_FILES, f"implementation must be one of: {list(DEFINITION_FILES)}."
    assert weights in [
        "van_walraven",
        "thompson_30",
        "thompson_29",
        "ahrq",
        "swiss",
    ], "weights must be one of: 'van_walraven', 'thompson_30', 'thompson_29', 'ahrq', or 'swiss'."
    if implementation == "ahrq_icd10":
        assert year_col is not None and year_col in df.columns, "Implementation 'ahrq_icd10' requires a 'year_col' column (diagnosis year) in the input DataFrame." # fmt: skip

    # Determine weight column and score column names based on weights argument
    # AHRQ implementations always use their own AHRQ weights; icd_version
    # (forced above) picks between the ICD-9/ICD-10 column
    if implementation in ("ahrq_icd9", "ahrq_icd10") and weights and weights.lower() != "ahrq":
        warnings.warn(
            f"Implementation '{implementation}' requires 'ahrq' weights. Overriding weights='{weights}' with 'ahrq_weights'.",
            UserWarning,
            stacklevel=2
        )

    if implementation in ("ahrq_icd9", "ahrq_icd10") or weights == "ahrq":
        if icd_version not in ("icd9", "icd10"):
            raise ValueError("weights='ahrq' requires icd_version 'icd9' or 'icd10'.")
        weight_col_name = f"ahrq_{icd_version}_weights"
    else:
        weight_col_name = f"{weights}_weights"

    df_combined, mutual_exclusion_rules = weighted_definitions(
        "ELIXHAUSER", DEFINITION_FILES[implementation], implementation, CANONICAL_RULES
    )

    # Call the generalized function with the combined DataFrame
    df_elixhauser = CustomComorbidityIndex(
        df=df,
        id_col=id_col,
        code_col=code_col,
        icd_version=icd_version,
        icd_version_col=icd_version_col,
        year_col=year_col if implementation == "ahrq_icd10" else None,
        definition_data=df_combined,
        weight_col_name=weight_col_name,
        score_col_name=SCORE_COL_NAME,
        mutual_exclusion_rules=mutual_exclusion_rules,
        return_categories=return_categories,
    )

    return df_elixhauser

def CombinedComorbidityScore(
    df: pl.DataFrame,
    id_col: str = "id",
    code_col: str = "code",
    icd_version: str = "icd10",
    icd_version_col: str | None = None,
    return_categories=False,
) -> pl.DataFrame:
    """Calculate the Combined Comorbidity Score using ICD codes.

    Args:
        df (pl.DataFrame): Input data containing at least ``id_col`` and ``code_col``.
        id_col (str, optional): Column name containing unique identifiers. Defaults to ``"id"``.
        code_col (str, optional): Column name containing ICD codes. Defaults to ``"code"``.
        icd_version (str, optional): ICD version; one of ``"icd9"``, ``"icd10"``, or ``"icd9_10"``. Defaults to ``"icd10"``.
        icd_version_col (str, optional): Column name with ICD version labels when ``icd_version`` is ``"icd9_10"``. Defaults to ``None``.
        return_categories (bool, optional): If ``True``, includes indicator columns for each Combined Comorbidity Score category. Defaults to ``False``.

    Returns:
        pl.DataFrame: DataFrame containing ``id_col``, the calculated score column, and, when ``return_categories`` is ``True``, category indicators.

    Raises:
        AssertionError: If required columns are missing.
    """

    # No mutual exclusion rules:
    # - the published Gagne SAS code treats "Any tumor" (1) and "Metastatic cancer" (5)
    #   as independent flags and sums both.
    # - uncomplicated diabetes has weight 0 (and is not considered in the SAS code),
    #   so complicated diabetes does not need to suppress it.
    return CustomComorbidityIndex(
        df=df,
        id_col=id_col,
        code_col=code_col,
        icd_version=icd_version,
        icd_version_col=icd_version_col,
        definition_data=DEFINITION_FILE,
        weight_col_name="weights",
        score_col_name=SCORE_COL_NAME,
        return_categories=return_categories,
    )

def HospitalFrailtyRiskScore(
    df: pl.DataFrame,
    id_col: str = "id",
    code_col: str = "code",
    icd_version: str = "icd10",
    return_categories=False,
) -> pl.DataFrame:
    """Calculate the Hospital Frailty Risk Score (HFRS) using ICD-10 codes.

    Args:
        df (pl.DataFrame | pd.DataFrame): Input data containing at least ``id_col`` and ``code_col``.
        id_col (str, optional): Column name containing unique identifiers. Defaults to ``"id"``.
        code_col (str, optional): Column name containing ICD codes. Defaults to ``"code"``.
        icd_version (str, optional): ICD version; must be ``"icd10"``. Defaults to ``"icd10"``.
        return_categories (bool, optional): If ``True``, includes indicator columns for each HFRS category. Defaults to ``False``.

    Returns:
        pl.DataFrame | pd.DataFrame: DataFrame containing ``id_col``, the ``"HFRS Score"`` column, and, when ``return_categories`` is ``True``, category indicators.

    Raises:
        AssertionError: If ``icd_version`` is not ``"icd10"`` or required columns are missing.
    """

    is_pandas = isinstance(df, pd.DataFrame)
    if is_pandas:
        df = pl.from_pandas(df)

    assert icd_version == "icd10", "icd_version must be 'icd10' for HFRS."
    assert id_col in df.columns, f"Column '{id_col}' (ID) must be present in input DataFrame." # fmt: skip
    assert code_col in df.columns, f"Column '{code_col}' (ICD code) must be present in input DataFrame." # fmt: skip

    hfrs_df = CustomComorbidityIndex(
        df=df,
        id_col=id_col,
        code_col=code_col,
        icd_version=icd_version,
        definition_data=DEFINITION_FILE,
        weight_col_name="weight",
        score_col_name=SCORE_COL_NAME,
        return_categories=return_categories,
    ).with_columns(pl.col(SCORE_COL_NAME).round(1))

    return hfrs_df.to_pandas() if is_pandas else hfrs_df

def CustomComorbidityIndex(
    df: pl.DataFrame,
    id_col: str = "id",
    code_col: str = "code",
    icd_version: str = "icd10",
    icd_version_col: str | None = None,
    definition_data: Path | pl.DataFrame | None = None,
    weight_col_name: str | None = None,
    score_col_name: str | None = None,
    mutual_exclusion_rules: list[tuple[str, str]] | None = None,
    year_col: str | None = None,
    return_categories=False,
) -> pl.DataFrame:
    """Compute a custom comorbidity score from ICD-coded input data.

    Args:
        df (pl.DataFrame): Input data containing at least ``id_col`` and ``code_col``.
        id_col (str, optional): Column name containing unique identifiers. Defaults to ``"id"``.
        code_col (str, optional): Column name containing ICD codes. Defaults to ``"code"``.
        icd_version (str, optional): ICD version; one of ``"icd9"``, ``"icd10"``, or ``"icd9_10"``. Defaults to ``"icd10"``.
        icd_version_col (str, optional): Column name with ICD version labels when ``icd_version`` is ``"icd9_10"``. Defaults to ``None``.
        definition_data (Path | pl.DataFrame, optional): Path to a CSV file or DataFrame containing category definitions and weights. May optionally include an ``and_group`` column: rows sharing a ``category`` and a non-empty ``and_group`` value must all be present (AND) for that group to count, OR'd into the category's normal presence; blank ``and_group`` behaves as before. Defaults to ``None``.
        weight_col_name (str, optional): Column name with weights inside ``definition_data``. Defaults to ``None``.
        score_col_name (str, optional): Column name assigned to the calculated score in the result. Defaults to ``None``.
        mutual_exclusion_rules (list[tuple[str, str]], optional): Pairs of mutually exclusive categories where the second entry is suppressed when the first is present. Defaults to ``None``.
        year_col (str, optional): Column name with each record's definition-catalogue year (e.g. ICD-10-GM catalogue year). When provided, ``definition_data`` must include a ``year`` column, and a code only matches a category as of the definition row's year onward (``df[year_col] >= definitions["year"]``). Defaults to ``None`` (year-independent matching, the historical behavior).
        return_categories (bool, optional): If ``True``, includes indicator columns for each matched category. Defaults to ``False``.

    Returns:
        pl.DataFrame: DataFrame containing ``id_col``, the calculated score column, and, when ``return_categories`` is ``True``, category indicators.

    Patients without a usable code are kept with all categories 0.

    Raises:
        AssertionError: If required columns or definitions are missing, or if ``icd_version`` and ``icd_version_col`` are inconsistent.
        ValueError: If ``code_col`` holds numbers instead of strings, or if ``mutual_exclusion_rules`` name a category that is not in the definitions.
    """

    # Input validation
    assert icd_version in [
        "icd9",
        "icd10",
        "icd9_10",
    ], f"icd_version must be one of: 'icd9', 'icd10', or 'icd9_10'. Got '{icd_version}'."
    assert id_col in df.columns, f"Column '{id_col}' (ID) must be present in input DataFrame." # fmt: skip
    assert code_col in df.columns, f"Column '{code_col}' (ICD code) must be present in input DataFrame." # fmt: skip
    if icd_version == "icd9_10":
        assert icd_version_col is not None and icd_version_col in df.columns, f"For icd9_10, '{icd_version_col}' (ICD version) column must be present in input DataFrame." # fmt: skip

    # Load definitions
    if isinstance(definition_data, Path):
        definitions = pl.read_csv(definition_data, separator=",")
    else:
        definitions = definition_data

    assert (
        "category" in definitions.columns
    ), "'category' column not found in definition file."
    assert (
        weight_col_name in definitions.columns
    ), f"'{weight_col_name}' column not found in definition file."

    df = check_code_dtype(df, code_col)

    # Year gating only applies when both year_col and a year column are given
    year_aware = year_col is not None and "year" in definitions.columns

    def process_single_icd(df, icd_col_name):
        # Patients without a usable code stay in the result with all categories 0
        ids = df.select(id_col).unique()
        df = df.filter(pl.col(code_col).is_not_null())

        # Validation and cleaning only touch the unique codes, the diagnosis rows keep their
        # raw codes and are joined on them below. Blank codes match nothing.
        unique_codes = (
            df.select(code_col)
            .unique()
            .with_columns(clean_code_expr(code_col).alias("clean"))
            .filter(pl.col("clean") != "")
        )

        categories_list = definitions["category"].drop_nulls().unique(maintain_order=True).to_list()
        if unique_codes.is_empty():  # no codes of this ICD version: nothing present
            return ids.with_columns(pl.lit(0).cast(int).alias(c) for c in categories_list), categories_list

        # Prepare mapping DataFrame: prefix, bit, since. Every definition row is one bit in a
        # per-patient mask; a category is present if all bits of one of its groups are set,
        # where a group is a single row or the rows sharing an and_group.
        required, code_map_rows = defaultdict(lambda: defaultdict(int)), []
        for i, row in enumerate(definitions.iter_rows(named=True)):
            if row["category"] is None or row.get(icd_col_name) is None:
                continue
            bit = 1 << i
            required[row["category"]][row.get("and_group") or i] |= bit
            for code in row[icd_col_name].split("|"):
                code_map_rows.append(
                    {"prefix": code.strip(), "bit": bit, "since": row["year"] if year_aware else 0}
                )
        code_map = pl.DataFrame(code_map_rows)

        code_prefixes = code_map.select("prefix").to_series().to_list()
        longest_code = max(unique_codes["clean"].to_list(), key=len)

        # Create a DataFrame with all possible prefixes, then find the best match
        best_prefix = (
            unique_codes.with_columns(
                pl.col("clean").str.slice(0, n).alias(f"prefix_{n}")
                for n in range(1, len(longest_code) + 1)
            )
            .with_columns(
                pl.when(pl.col(f"prefix_{n}").is_in(code_prefixes))
                .then(pl.col(f"prefix_{n}"))
                .otherwise(None)
                .alias(f"prefix_{n}")
                for n in range(1, len(longest_code) + 1)
            )
            # Count ICD code as present for each comorbidity it matches
            .unpivot(
                on=[f"prefix_{n}" for n in range(1, len(longest_code) + 1)],
                index=code_col,
                value_name="best_prefix",
            )
            .drop_nulls("best_prefix")
            .select(code_col, "best_prefix")
            .unique()
        )

        # Join diagnoses with the bits of their code; only matching diagnoses are aggregated
        matches = df.join(
            best_prefix.join(
                code_map, left_on="best_prefix", right_on="prefix"
            ).select(code_col, "bit", "since"),
            on=code_col,
        )
        if year_aware:
            matches = matches.filter(pl.col(year_col) >= pl.col("since"))

        masks_by_id = matches.group_by(id_col).agg(pl.col("bit").bitwise_or().alias("mask"))
        df_presence_absence = (
            ids.join(masks_by_id, on=id_col, how="left")
            .with_columns(pl.col("mask").fill_null(0))
            .select(
                id_col,
                *(
                    pl.any_horizontal((pl.col("mask") & m) == m for m in masks.values())
                    .cast(int)
                    .alias(cat)
                    for cat, masks in required.items()
                ),
            )
        )

        # Add missing columns
        for cat in categories_list:
            if cat not in df_presence_absence.columns:
                df_presence_absence = df_presence_absence.with_columns(
                    pl.lit(0).cast(int).alias(cat)
                )
        df_presence_absence = df_presence_absence.select(
            [id_col] + categories_list
        )
        return df_presence_absence, categories_list

    if icd_version == "icd9_10":
        # Split by icd_version_col
        df_icd9 = df.filter(pl.col(icd_version_col).cast(str) == "ICD-9")
        df_icd10 = df.filter(pl.col(icd_version_col).cast(str) != "ICD-9")

        # Process each ICD version separately
        df9, cats9 = process_single_icd(df_icd9, "icd9_codes")
        df10, cats10 = process_single_icd(df_icd10, "icd10_codes")

        # Combine: outer join on id_col, take max for each category
        # Preserve definition file order instead of alphabetical sorting
        all_categories_set = set(cats9) | set(cats10)
        all_categories = [
            cat
            for cat in dict.fromkeys(definitions["category"].to_list())
            if cat in all_categories_set and cat is not None
        ]
        df_combined = df9.join(
            df10, on=id_col, how="full", suffix="_right", coalesce=True
        )
        # both halves report every category, so each has a "_right" counterpart
        df_presence_absence = df_combined.select(
            id_col,
            *(pl.max_horizontal(cat, f"{cat}_right").alias(cat) for cat in all_categories),
        )
    else:
        icd_col_name = f"{icd_version}_codes"
        df_presence_absence, all_categories = process_single_icd(
            df, icd_col_name
        )

    # STEP 2: calculate score
    category_weights = dict(
        zip(definitions["category"], definitions[weight_col_name])
    )
    category_weights = {
        k: v
        for k, v in category_weights.items()
        if k is not None and v is not None
    }

    # Apply mutual exclusion rules
    for complicated, uncomplicated in mutual_exclusion_rules or []:
        unknown = {complicated, uncomplicated} - set(df_presence_absence.columns)
        if unknown:
            raise ValueError(f"mutual_exclusion_rules refer to unknown categories: {sorted(unknown)}.") # fmt: skip
        df_presence_absence = df_presence_absence.with_columns(
            pl.when(pl.col(complicated) == 1)
            .then(0)
            .otherwise(pl.col(uncomplicated))
            .alias(uncomplicated)
        )

    # calculate scores
    # Ensure weights are applied only to existing category columns that have weights
    score_expr = pl.sum_horizontal(
        pl.col(cat) * category_weights.get(cat, 0)
        for cat in all_categories
        if cat in df_presence_absence.columns and cat in category_weights
    ).alias(score_col_name)

    cols = [id_col] + all_categories if return_categories else [id_col]
    return df_presence_absence.select(*cols, score_expr)

def coerce_icd_version(icd_version, implementation, icd9_only, icd10_only):
    """Replace an ICD version the implementation cannot use, with a warning."""
    if implementation in icd9_only:
        version, label = "icd9", "ICD-9"
    elif implementation in icd10_only:
        version, label = "icd10", "ICD-10"
    else:
        return icd_version
    if icd_version in ("icd9", "icd10", "icd9_10") and icd_version != version:
        warnings.warn(
            f"Implementation '{implementation}' only uses {label}. Setting ICD version to '{version}'.",
            UserWarning,
            stacklevel=3,
        )
        return version
    return icd_version

def weighted_definitions(prefix, definition_file, implementation, canonical_rules):
    """Definitions of a Charlson/Elixhauser implementation joined with its weights, plus its exclusion rules.

    ``prefix`` ("CHARLSON" or "ELIXHAUSER") names the weight and alias files; ``canonical_rules`` are
    (more severe, milder) pairs in canonical category names, translated to the implementation's own names.
    """
    base_path = Path(__file__).parent / "common"
    definitions = pl.read_csv(base_path / definition_file)
    weights = pl.read_csv(base_path / f"{prefix}_WEIGHTS.csv").drop("index")
    with open(base_path / f"{prefix}_CATEGORY_ALIASES.yaml") as f:
        aliases = yaml.safe_load(f).get(implementation, {})
    combined = (
        definitions.with_columns(pl.col("category").replace(aliases).alias("__weight_category__"))
        .join(weights.rename({"category": "__weight_category__"}), on="__weight_category__", how="left")
        .drop("__weight_category__")
    )
    canonical = {c: aliases.get(c, c) for c in combined["category"]}
    rules = [
        (severe, mild)
        for severe, mild in product(canonical, repeat=2)
        if (canonical[severe], canonical[mild]) in canonical_rules
    ]
    return combined, rules