Data Flow¶

        # Enrich counterparties with hierarchy info
        enriched_counterparties = self._enrich_counterparties_with_hierarchy(
            counterparties,
            org_mappings,
            ratings,
            ultimate_parents,
            rating_info,
        )

        return CounterpartyLookup(
            counterparties=enriched_counterparties,
            parent_mappings=org_mappings.select(
                [
                    "child_counterparty_reference",
                    "parent_counterparty_reference",
                ]
            ),
            ultimate_parent_mappings=ultimate_parents,
            rating_inheritance=rating_info,
        ), errors

    def _build_ultimate_parent_lazy(
        self,
        org_mappings: pl.LazyFrame,
        max_depth: int = 10,
    ) -> pl.LazyFrame:
        """
        Build ultimate parent mapping using eager graph traversal.

        Collects the small edge data eagerly, resolves the full graph via dict
        traversal, and returns the result as a LazyFrame for downstream joins.

        Returns LazyFrame with columns:
        - counterparty_reference: The entity
        - ultimate_parent_reference: Its ultimate parent
        - hierarchy_depth: Number of levels traversed
        """
        edges = (
            org_mappings.select(
                [
                    "child_counterparty_reference",
                    "parent_counterparty_reference",
                ]
            )
            .unique()
            .collect()
        )

        resolved = _resolve_graph_eager(
            edges,
            child_col="child_counterparty_reference",
            parent_col="parent_counterparty_reference",
            max_depth=max_depth,
        )

        return resolved.rename(
            {
                "entity": "counterparty_reference",
                "root": "ultimate_parent_reference",
                "depth": "hierarchy_depth",
            }
        ).lazy()

    def _build_rating_inheritance_lazy(
        self,
        counterparties: pl.LazyFrame,
        ratings: pl.LazyFrame,
        ultimate_parents: pl.LazyFrame,
    ) -> pl.LazyFrame:
        """

            pl.col("approach").is_in([ApproachType.FIRB.value, ApproachType.AIRB.value])
        )
        slotting_exposures = classified.filter(pl.col("approach") == ApproachType.SLOTTING.value)
        classification_audit = self._build_audit_trail(classified)

        return ClassifiedExposuresBundle(
            all_exposures=classified,
            sa_exposures=sa_exposures,
            irb_exposures=irb_exposures,
            slotting_exposures=slotting_exposures,
            equity_exposures=data.equity_exposures,
            ciu_holdings=data.ciu_holdings,
            collateral=data.collateral,
            guarantees=data.guarantees,
            provisions=data.provisions,
            counterparty_lookup=data.counterparty_lookup,
            classification_audit=classification_audit,
            classification_errors=[],
        )

    # =========================================================================
    # Phase 1: Counterparty join (retained unchanged)
    # =========================================================================

    def _add_counterparty_attributes(
        self,
        exposures: pl.LazyFrame,
        counterparties: pl.LazyFrame,
    ) -> pl.LazyFrame:
        """
        Add counterparty attributes needed for classification.

        Joins exposures with counterparty data to get:
        - entity_type (single source of truth for exposure class)
        - annual_revenue (for SME check)
        - total_assets (for large financial sector entity threshold)
        - default_status
        - country_code
        - apply_fi_scalar (for FI scalar - LFSE/unregulated FSE)
        - is_managed_as_retail (for SME retail treatment)
        """
        cp_schema = counterparties.collect_schema()
        cp_col_names = cp_schema.names()

        select_cols = [
            pl.col("counterparty_reference"),
            pl.col("entity_type").str.to_lowercase().alias("cp_entity_type"),
            pl.col("country_code").alias("cp_country_code"),
            pl.col("annual_revenue").alias("cp_annual_revenue"),
            pl.col("total_assets").alias("cp_total_assets"),

                "_ben_ref_facility": pl.String,
                "_ead_facility": pl.Float64,
                "_currency_facility": pl.String,
                "_maturity_facility": pl.Date,
            }
        )

    # Counterparty: aggregated per counterparty_reference
    cp_lookup = (
        exposures.group_by("counterparty_reference")
        .agg(
            [
                pl.col("ead_gross").sum().alias("_ead_cp"),
                pl.col("currency").first().alias("_currency_cp"),
                pl.col("maturity_date").first().alias("_maturity_cp"),
            ]
        )
        .rename({"counterparty_reference": "_ben_ref_cp"})
    )

    return direct_lookup, facility_lookup, cp_lookup


def _join_collateral_to_lookups(
    collateral: pl.LazyFrame,
    direct_lookup: pl.LazyFrame,
    facility_lookup: pl.LazyFrame,
    cp_lookup: pl.LazyFrame,
) -> pl.LazyFrame:
    """
    Join all lookup columns (EAD, currency, maturity) onto collateral in one pass.

    Replaces the separate _resolve_pledge_percentages and join_exposure_currency
    join passes — each used 3 left joins, so 6 total. Now we do 3 joins total,
    halving the plan size from lookup subtree duplication.

    When beneficiary_type is absent, falls back to a single direct join.
    """
    from rwa_calc.engine.crm.constants import DIRECT_BENEFICIARY_TYPES

    coll_schema = collateral.collect_schema()

    if "beneficiary_type" not in coll_schema.names():
        # Direct-only join — single pass with all columns
        return collateral.join(
            direct_lookup.select(
                pl.col("_ben_ref_direct"),
                pl.col("_ead_direct").alias("_beneficiary_ead"),
                pl.col("_currency_direct").alias("exposure_currency"),
                pl.col("_maturity_direct").alias("exposure_maturity"),
            ),
            left_on="beneficiary_reference",
            right_on="_ben_ref_direct",
            how="left",
        )

    bt_lower = pl.col("beneficiary_type").str.to_lowercase()

    # 3 left joins — each adds level-specific EAD, currency, maturity columns
    collateral = (
        collateral.join(
            direct_lookup,
            left_on="beneficiary_reference",
            right_on="_ben_ref_direct",
            how="left",
        )
        .join(
            facility_lookup,
            left_on="beneficiary_reference",
            right_on="_ben_ref_facility",
            how="left",
        )
        .join(
            cp_lookup,

    # Optimisation: if all floors are the same (CRR case), return a scalar
    all_values = {
        floors.corporate,
        floors.corporate_sme,
        floors.retail_mortgage,
        floors.retail_other,
        floors.retail_qrre_transactor,
        floors.retail_qrre_revolver,
    }
    if len(all_values) == 1:
        return pl.lit(float(all_values.pop()))

    # Basel 3.1: differentiated floors by exposure class
    exp_class = pl.col("exposure_class").cast(pl.String).fill_null("CORPORATE").str.to_uppercase()

    # QRRE transactor/revolver distinction (CRE30.55):
    # Transactors (repay in full each period) get 0.03% floor;
    # revolvers (carry balance) get 0.10% floor.
    if has_transactor_col:
        qrre_floor = (
            pl.when(pl.col("is_qrre_transactor").fill_null(False))
            .then(pl.lit(float(floors.retail_qrre_transactor)))
            .otherwise(pl.lit(float(floors.retail_qrre_revolver)))
        )
    else:
        # Conservative default: revolver floor (0.10% under Basel 3.1)
        qrre_floor = pl.lit(float(floors.retail_qrre_revolver))

    return (
        pl.when(exp_class.str.contains("QRRE"))
        .then(qrre_floor)
        .when(exp_class.str.contains("MORTGAGE") | exp_class.str.contains("RESIDENTIAL"))
        .then(pl.lit(float(floors.retail_mortgage)))
        .when(exp_class.str.contains("RETAIL"))
        .then(pl.lit(float(floors.retail_other)))
        .when(exp_class == "CORPORATE_SME")
        .then(pl.lit(float(floors.corporate_sme)))
        .otherwise(pl.lit(float(floors.corporate)))
    )


def _lgd_floor_expression(
    config: CalculationConfig,
    *,
    has_seniority: bool = False,
) -> pl.Expr:
    """
    Build Polars expression for LGD floor (no collateral_type column).

    Under CRR: No LGD floors (returns 0.0).
    Under Basel 3.1 (CRE30.41): Differentiated floors for A-IRB:
        - Unsecured (senior): 25%
        - Unsecured (subordinated): 50%
        - Financial collateral: 0%
        - Receivables: 10%
        - CRE: 10%, RRE: 5%
        - Other physical: 15%

    Without a collateral_type column, defaults to unsecured floor (25%/50%).
    When has_seniority=True, checks seniority column for subordinated (50%).

    Returns a Polars expression evaluating to the per-row LGD floor value.
    """
    if config.is_crr:
        return pl.lit(0.0)

    floors = config.lgd_floors

    if has_seniority:
        is_subordinated = (
            pl.col("seniority").fill_null("senior").str.to_lowercase().str.contains("sub")
        )
        return (
            pl.when(is_subordinated)
            .then(pl.lit(float(floors.subordinated_unsecured)))
            .otherwise(pl.lit(float(floors.unsecured)))
        )

    # Default to unsecured floor (25%) — most conservative for senior
    return pl.lit(float(floors.unsecured))


def _lgd_floor_expression_with_collateral(
    config: CalculationConfig,
    *,
    has_seniority: bool = False,
) -> pl.Expr:
    """
    Build Polars expression for per-collateral-type LGD floor when collateral_type
    column is available.

    This is used when the dataframe has a collateral_type column, allowing
    precise per-row LGD floors based on the primary collateral type.

    When has_seniority=True, subordinated unsecured exposures get the higher
    50% floor instead of 25% (CRE30.41).
    """
    if config.is_crr:
        return pl.lit(0.0)

    floors = config.lgd_floors
    coll = pl.col("collateral_type").fill_null("unsecured").str.to_lowercase()

    # Determine unsecured floor: 50% for subordinated, 25% for senior (CRE30.41)