"""Fair lending analysis via causal path decomposition.

Implements Pearl/Kilbertus causal fairness framework: decompose the total
effect of protected characteristics into direct (potentially discriminatory)
and indirect effects through legitimate business factors.
"""

from __future__ import annotations

from dataclasses import dataclass, field

from orxaq.kernel import WorldDAG

# Variables considered legitimate business factors (not discriminatory
# even if correlated with protected characteristics)
LEGITIMATE_MEDIATORS: frozenset[str] = frozenset(
    {
        "employment_length",
        "income",
        "credit_score",
        "dti",
        "loan_amount",
        "interest_rate",
        "ltv",
        "term",
        "current_balance",
    }
)

# Protected characteristics under ECOA/CFPB
PROTECTED_CHARACTERISTICS: frozenset[str] = frozenset(
    {
        "age",
    }
)


@dataclass(frozen=True)
class EffectDecomposition:
    """Decomposition of a protected characteristic's total effect."""

    protected_var: str
    outcome_var: str
    total_effect: float
    direct_effect: float
    indirect_effects: dict[str, float]  # mediator -> effect_size
    legitimate_paths: tuple[tuple[str, ...], ...]
    concerning_paths: tuple[tuple[str, ...], ...]


@dataclass(frozen=True)
class FairnessReport:
    """Fairness audit for a specific borrower decision."""

    borrower: dict[str, float]
    decision: str  # "approve" or "deny"
    counterfactual_decision: str  # decision if protected var differed
    decision_changed: bool
    direct_effect_size: float
    recommendation: str  # "clear", "review", or "flag"
    protected_var: str
    decomposition: EffectDecomposition = field(default_factory=lambda: None)  # type: ignore[assignment]


class FairLendingAnalyzer:
    """Analyze lending decisions for causal fairness.

    Uses WorldDAG causal paths to decompose effects of protected
    characteristics into direct (concerning) and indirect (through
    legitimate factors) pathways.
    """

    def __init__(self, dag: WorldDAG) -> None:
        self._dag = dag

    def decompose_effect(
        self,
        protected: str,
        outcome: str,
    ) -> EffectDecomposition:
        """Decompose effect of protected characteristic on outcome.

        Classifies each causal path as:
        - Direct: protected → outcome (no mediator, potentially discriminatory)
        - Indirect through legitimate mediator (acceptable)
        - Indirect through non-legitimate mediator (concerning)
        """
        all_paths = self._dag.all_directed_paths(protected, outcome)

        legitimate_paths: list[tuple[str, ...]] = []
        concerning_paths: list[tuple[str, ...]] = []
        indirect_effects: dict[str, float] = {}

        for path in all_paths:
            path_tuple = tuple(path)
            mediators = set(path[1:-1])  # Nodes between source and target

            if not mediators:
                # Direct path: protected → outcome
                concerning_paths.append(path_tuple)
            elif mediators <= LEGITIMATE_MEDIATORS:
                # All mediators are legitimate business factors
                legitimate_paths.append(path_tuple)
                for m in mediators:
                    # Estimate indirect effect magnitude from d-separation
                    indirect_effects[m] = indirect_effects.get(m, 0.0) + (1.0 / len(path))
            else:
                # Some mediators are not legitimate
                concerning_paths.append(path_tuple)

        # Estimate total and direct effects from path structure
        total_paths = len(all_paths) if all_paths else 1
        direct_effect = len(concerning_paths) / total_paths
        total_effect = 1.0  # Normalized to 1.0, decomposition shows split

        return EffectDecomposition(
            protected_var=protected,
            outcome_var=outcome,
            total_effect=total_effect,
            direct_effect=direct_effect,
            indirect_effects=indirect_effects,
            legitimate_paths=tuple(legitimate_paths),
            concerning_paths=tuple(concerning_paths),
        )

    def audit_decision(
        self,
        borrower: dict[str, float],
        decision_var: str = "default_flag",
        pd_threshold: float = 0.5,
    ) -> FairnessReport:
        """Audit a lending decision for fairness.

        Computes the actual decision and a counterfactual: what if the
        protected characteristic were different? If the decision changes,
        flags for review.
        """
        # Find protected characteristics in the borrower data
        protected_var = None
        for pv in PROTECTED_CHARACTERISTICS:
            if pv in borrower:
                protected_var = pv
                break

        if protected_var is None:
            return FairnessReport(
                borrower=borrower,
                decision="approve",
                counterfactual_decision="approve",
                decision_changed=False,
                direct_effect_size=0.0,
                recommendation="clear",
                protected_var="none",
            )

        # Actual decision based on PD proxy
        actual_pd = borrower.get(decision_var, 0.0)
        decision = "deny" if actual_pd >= pd_threshold else "approve"

        # Counterfactual: perturb protected variable
        counterfactual = dict(borrower)
        current_val = counterfactual.get(protected_var, 0.0)

        # Shift by 1 standard deviation as counterfactual
        if protected_var == "age":
            counterfactual[protected_var] = current_val + 10.0  # +10 years

        # Simple counterfactual decision (would need full model for precision)
        cf_pd = counterfactual.get(decision_var, 0.0)
        counterfactual_decision = "deny" if cf_pd >= pd_threshold else "approve"

        # Decompose the effect
        decomposition = self.decompose_effect(protected_var, decision_var)

        decision_changed = decision != counterfactual_decision
        direct_effect = decomposition.direct_effect

        # Recommendation
        if direct_effect > 0.3:
            recommendation = "flag"
        elif direct_effect > 0.1 or decision_changed:
            recommendation = "review"
        else:
            recommendation = "clear"

        return FairnessReport(
            borrower=borrower,
            decision=decision,
            counterfactual_decision=counterfactual_decision,
            decision_changed=decision_changed,
            direct_effect_size=direct_effect,
            recommendation=recommendation,
            protected_var=protected_var,
            decomposition=decomposition,
        )
