"""Causal PD (Probability of Default) model.

Uses WorldDAG structure and backdoor adjustment for causal effect estimation.
Distinguishes observational prediction from interventional queries.

Implementation: Gradient-descent logistic regression (stdlib-only, ADR-007).
Adjustment sets identified via AdjustmentSetFinder (backdoor criterion).
"""

from __future__ import annotations

import math

from orxaq.intelligence.intervention import AdjustmentSetFinder
from orxaq.kernel import WorldDAG


def _sigmoid(x: float) -> float:
    """Numerically stable sigmoid function."""
    if x >= 0:
        return 1.0 / (1.0 + math.exp(-x))
    exp_x = math.exp(x)
    return exp_x / (1.0 + exp_x)


def _dot(weights: list[float], features: list[float]) -> float:
    return sum(w * f for w, f in zip(weights, features))


class _LogisticRegression:
    """Gradient-descent logistic regression — stdlib-only, no external deps."""

    LR = 0.05
    N_ITER = 300

    def __init__(self) -> None:
        self.coef_: list[float] = []
        self.intercept_: float = 0.0
        self.coef_variance_: list[float] = []
        self.means_: list[float] = []
        self.stds_: list[float] = []

    def _standardize(
        self, x_mat: list[list[float]]
    ) -> tuple[list[list[float]], list[float], list[float]]:
        """Z-score standardize columns. Returns (x_std, means, stds)."""
        n = len(x_mat)
        if n == 0:
            return x_mat, [], []
        p = len(x_mat[0])
        means = [sum(x_mat[i][j] for i in range(n)) / n for j in range(p)]
        vars_ = [
            sum((x_mat[i][j] - means[j]) ** 2 for i in range(n)) / max(n - 1, 1)
            for j in range(p)
        ]
        stds = [math.sqrt(v) if v > 1e-12 else 1.0 for v in vars_]
        x_std = [[(x_mat[i][j] - means[j]) / stds[j] for j in range(p)] for i in range(n)]
        return x_std, means, stds

    def fit(self, x_mat: list[list[float]], y: list[float]) -> _LogisticRegression:
        n = len(x_mat)
        if n == 0:
            return self
        x_std, self.means_, self.stds_ = self._standardize(x_mat)
        p = len(x_mat[0])
        coef = [0.0] * p
        intercept = 0.0

        for _ in range(self.N_ITER):
            grad_coef = [0.0] * p
            grad_int = 0.0
            for i in range(n):
                pred = _sigmoid(intercept + _dot(coef, x_std[i]))
                err = pred - y[i]
                grad_int += err
                for j in range(p):
                    grad_coef[j] += err * x_std[i][j]
            intercept -= self.LR * grad_int / n
            coef = [coef[j] - self.LR * grad_coef[j] / n for j in range(p)]

        self.coef_ = coef
        self.intercept_ = intercept
        # Diagonal Fisher information approx for variance estimates
        self.coef_variance_ = self._fisher_variance(x_std, coef, intercept)
        return self

    def _fisher_variance(
        self,
        x_std: list[list[float]],
        coef: list[float],
        intercept: float,
    ) -> list[float]:
        """Diagonal of (X^T W X)^{-1} where W = p*(1-p)."""
        n, p = len(x_std), len(coef)
        diag = [1e-8] * p
        for i in range(n):
            pred = _sigmoid(intercept + _dot(coef, x_std[i]))
            w = max(pred * (1.0 - pred), 1e-10)
            for j in range(p):
                diag[j] += w * x_std[i][j] * x_std[i][j]
        return [1.0 / d for d in diag]

    def predict_proba(self, x: list[float]) -> float:
        x_std = [
            (x[j] - self.means_[j]) / self.stds_[j] if self.stds_[j] > 1e-12 else 0.0
            for j in range(len(x))
        ]
        return _sigmoid(self.intercept_ + _dot(self.coef_, x_std))


class CausalPDModel:
    """Causal probability of default model.

    Uses WorldDAG backdoor adjustment to estimate causal effects of
    variables on default. Supports both observational prediction and
    interventional "what if" queries.

    Estimation method: backdoor-adjustment logistic regression.
    Features selected via AdjustmentSetFinder on the credit DAG.
    """

    def __init__(self, dag: WorldDAG) -> None:
        self._dag = dag
        self._weights: dict[str, float] = {}
        self._intercept: float = 0.0
        self._fitted: bool = False
        self._outcome = "default_flag"
        self._feature_names: list[str] = []
        self._coef_variance: dict[str, float] = {}
        self._lr: _LogisticRegression | None = None
        self._feature_means: list[float] = []
        self._feature_stds: list[float] = []

    @property
    def is_fitted(self) -> bool:
        return self._fitted


    def _normalize(self, x: list[float]) -> list[float]:
        """Z-score normalize a feature vector using fitted mean/std."""
        if not self._feature_means:
            return x
        return [(x[j] - self._feature_means[j]) / self._feature_stds[j] for j in range(len(x))]

    def fit(self, data: list[dict[str, float]]) -> None:
        """Fit backdoor-adjustment logistic regression on credit DAG.

        1. Identifies predictors (ancestors of default_flag in the DAG).
        2. Uses AdjustmentSetFinder to select adjustment covariates.
        3. Fits logistic regression via gradient descent (stdlib-only).
        4. Stores coefficients as causal weights for predict_pd().
        """
        if not data:
            msg = "Cannot fit on empty data"
            raise ValueError(msg)

        predictors = self._get_predictors()

        # Expand feature set: predictors + adjustment sets
        feature_set: set[str] = set(predictors)
        finder = AdjustmentSetFinder(self._dag)
        for treatment in predictors:
            adj = finder.find_minimal(treatment, self._outcome)
            if adj is not None:
                feature_set.update(adj.variables)

        # Remove outcome itself
        feature_set.discard(self._outcome)
        self._feature_names = sorted(feature_set)

        if not self._feature_names:
            self._fitted = True
            return

        # Build design matrix (raw values — _LogisticRegression handles standardization)
        x_mat = [[row.get(f, 0.0) for f in self._feature_names] for row in data]
        y = [row.get(self._outcome, 0.0) for row in data]

        lr = _LogisticRegression()
        lr.fit(x_mat, y)
        self._lr = lr
        self._intercept = lr.intercept_
        self._weights = dict(zip(self._feature_names, lr.coef_))
        self._coef_variance = dict(zip(self._feature_names, lr.coef_variance_))
        self._fitted = True

    def predict_pd(self, borrower: dict[str, float]) -> float:
        """Predict P(default=1 | X=borrower). Observational prediction."""
        if not self._fitted:
            msg = "Model not fitted. Call fit() first."
            raise RuntimeError(msg)

        if not self._feature_names:
            return 0.01  # no features — return base rate

        x = [borrower.get(f, 0.0) for f in self._feature_names]
        if self._lr is not None:
            return self._lr.predict_proba(x)
        # Fallback (shouldn't be reached after fit)
        linear = self._intercept + _dot(list(self._weights.values()), x)
        return _sigmoid(linear)

    def interventional_pd(
        self,
        borrower: dict[str, float],
        treatment: str,
        value: float,
    ) -> float:
        """Predict P(default=1 | do(treatment=value), X=borrower).

        Uses backdoor adjustment: set treatment to value, keep adjustment
        variables at observed levels, use fitted logistic weights.
        """
        if not self._fitted:
            msg = "Model not fitted. Call fit() first."
            raise RuntimeError(msg)

        # Get adjustment set for this treatment via backdoor criterion
        finder = AdjustmentSetFinder(self._dag)
        adj_result = finder.find_minimal(treatment, self._outcome)
        adjustment_set: set[str] = set(adj_result.variables) if adj_result is not None else set()

        modified = dict(borrower)
        modified[treatment] = value

        if not self._feature_names:
            return 0.01

        x = []
        for feat in self._feature_names:
            if feat == treatment:
                x.append(value)
            elif feat in adjustment_set:
                # Keep adjustment variables at observed values (block backdoor)
                x.append(borrower.get(feat, 0.0))
            else:
                x.append(modified.get(feat, 0.0))

        if self._lr is not None:
            return self._lr.predict_proba(x)
        linear = self._intercept + _dot(list(self._weights.values()), x)
        return _sigmoid(linear)

    def feature_importance(self) -> dict[str, float]:
        """Return causal feature importance (absolute effect on default).

        Unlike permutation importance, this reflects CAUSAL contribution
        from the backdoor-adjusted logistic regression coefficients.
        """
        if not self._fitted:
            msg = "Model not fitted. Call fit() first."
            raise RuntimeError(msg)

        return dict(sorted(self._weights.items(), key=lambda x: abs(x[1]), reverse=True))

    def coef_standard_errors(self) -> dict[str, float]:
        """Return standard errors for fitted coefficients (delta method)."""
        if not self._fitted:
            msg = "Model not fitted. Call fit() first."
            raise RuntimeError(msg)
        return {k: math.sqrt(v) for k, v in self._coef_variance.items()}

    def _get_predictors(self) -> list[str]:
        """Get variables with directed paths to default_flag."""
        ancestors = self._dag.ancestors(self._outcome)
        return [v for v in ancestors if v != self._outcome]

    def _estimate_causal_effect(
        self,
        var_id: str,
        data: list[dict[str, float]],
    ) -> float:
        """Estimate causal effect of var_id on default using backdoor adjustment.

        Kept for backward compatibility with edge-case tests.
        Uses stratified regression within backdoor adjustment strata.
        """
        adjustment_set = self._dag.backdoor_adjustment_set(var_id, self._outcome)

        n = len(data)
        if n == 0:
            return 0.0

        var_values = [row.get(var_id, 0.0) for row in data]
        outcome_values = [row.get(self._outcome, 0.0) for row in data]

        var_mean = sum(var_values) / n
        var_std = math.sqrt(sum((v - var_mean) ** 2 for v in var_values) / n)

        if var_std < 1e-10:
            return 0.0

        if not adjustment_set:
            out_mean = sum(outcome_values) / n
            cov = sum((var_values[i] - var_mean) * (outcome_values[i] - out_mean) for i in range(n))
            return cov / (n * var_std * var_std)

        # Backdoor adjustment: stratified estimation
        adj_var = next(iter(adjustment_set))
        adj_values = [row.get(adj_var, 0.0) for row in data]
        adj_sorted = sorted(adj_values)
        t1 = adj_sorted[n // 3] if n >= 3 else adj_sorted[0]
        t2 = adj_sorted[2 * n // 3] if n >= 3 else adj_sorted[-1]

        effects = []
        for low, high in [(-float("inf"), t1), (t1, t2), (t2, float("inf"))]:
            stratum_idx = [
                i
                for i in range(n)
                if low <= adj_values[i] < high or (high == float("inf") and adj_values[i] >= low)
            ]
            if len(stratum_idx) < 2:
                continue

            s_var = [var_values[i] for i in stratum_idx]
            s_out = [outcome_values[i] for i in stratum_idx]
            s_n = len(stratum_idx)
            s_var_mean = sum(s_var) / s_n
            s_out_mean = sum(s_out) / s_n
            s_var_var = sum((v - s_var_mean) ** 2 for v in s_var) / s_n

            if s_var_var < 1e-10:
                continue

            s_cov = sum((s_var[j] - s_var_mean) * (s_out[j] - s_out_mean) for j in range(s_n))
            effects.append(s_cov / (s_n * s_var_var))

        if not effects:
            return 0.0

        return sum(effects) / len(effects)
