"""
Bridge: Long-term Surrogate ≡ Proximal Causal Inference (Imbens, Kallus,
Mao & Wang 2025, JRSS-B 87(2) 362-388; arXiv 2202.07234).

Both paths target the long-term ATT under unobserved confounding via
short-term proxy variables. Surrogate Index assumes the proxy fully
mediates the long-term effect. PCI assumes the proxy resolves
confounding through bridge functions. Reporting both estimates
disciplines the long-term policy claim.
"""

from __future__ import annotations

from typing import List, Optional

import numpy as np
from ..core._bootstrap import bootstrap_se as _bootstrap_se
import pandas as pd

from .core import BridgeResult, _agreement_test, _dr_combine, _register


@_register("surrogate_pci")
def surrogate_pci_bridge(
    data: pd.DataFrame,
    long_term: str,
    short_term: List[str],
    treat: str,
    covariates: Optional[List[str]] = None,
    alpha: float = 0.05,
    n_boot: int = 200,
    seed: int = 0,
) -> BridgeResult:
    """
    Compare Surrogate Index (path A) against PCI bridge (path B) on
    the long-term ATT, using short-term proxies.

    Parameters
    ----------
    data : pd.DataFrame
    long_term : str
        Long-term outcome column.
    short_term : list of str
        Short-term proxy / surrogate variables.
    treat : str
        Binary treatment column.
    covariates : list of str, optional
    alpha : float
    n_boot : int, default 200
    seed : int
    """
    cov = list(covariates or [])
    df = data[[long_term, treat] + list(short_term) + cov].dropna() \
        .reset_index(drop=True)
    Y = df[long_term].to_numpy(float)
    D = df[treat].to_numpy(int)
    S = df[list(short_term)].to_numpy(float)
    X = df[cov].to_numpy(float) if cov else np.zeros((len(df), 0))
    n = len(df)
    rng = np.random.default_rng(seed)

    # ---------- Path A: Surrogate Index (Athey et al.) ---------- #
    # 1. Predict long-term Y from short-term S (and X) on controls only
    # 2. Apply prediction to treated to get expected long-term Y under
    #    treatment via the surrogate index, then take treated - control mean
    def _surrogate(Yi, Di, Si, Xi):
        from sklearn.linear_model import LinearRegression
        Z = np.hstack([Si, Xi])
        idx_c = Di == 0
        if idx_c.sum() < Z.shape[1] + 2:
            raise ValueError("Too few controls for surrogate regression.")
        model = LinearRegression().fit(Z[idx_c], Yi[idx_c])
        # Predict for treated
        idx_t = Di == 1
        if idx_t.sum() == 0:
            raise ValueError("No treated units.")
        # Surrogate index prediction for treated minus observed control mean
        # in long-term Y
        y_treated_pred = model.predict(Z[idx_t]).mean()
        y_control_obs = Yi[idx_c].mean()
        return float(y_treated_pred - y_control_obs)

    try:
        att_surr = _surrogate(Y, D, S, X)
    except Exception:
        att_surr = np.nan

    # ---------- Path B: PCI bridge (two-model counterfactual) ---------- #
    # Imbens-Kallus-Mao-Wang (2025) show that under proxy completeness the bridge
    # function b(W, X) = E[Y | W, X, D = 0] identifies the ATT via
    #     ATT = E[b(W, X) | D = 1] - E[Y | D = 0].
    # We go one step further and fit a symmetric bridge on BOTH arms —
    # m_d(W, X) = E[Y | W, X, D = d] — then report
    #     ATT_{PCI} = E[m_1(W, X) - m_0(W, X) | D = 1].
    # This uses a *different* identifying assumption from Path A
    # (Path A only requires E[Y | W, X, D = 0], i.e. surrogacy +
    # control-arm model correctness).  Path B instead leans on the
    # treated-arm model being well-specified.  The agreement test
    # between the two paths is therefore a genuine bridge: if Path A
    # (control-arm extrapolation) and Path B (treated-arm
    # counterfactual on W, X) disagree, at least one identifying
    # assumption is violated.  Plain OLS on (D, W, X), which the
    # previous implementation returned, would have collapsed Path B
    # into Path A under linear outcomes.
    def _pci(Yi, Di, Si, Xi):
        from sklearn.linear_model import LinearRegression
        Z = np.hstack([Si, Xi]) if Xi.shape[1] > 0 else Si
        idx_t = Di == 1
        idx_c = Di == 0
        if idx_t.sum() < Z.shape[1] + 2 or idx_c.sum() < Z.shape[1] + 2:
            return np.nan
        m1 = LinearRegression().fit(Z[idx_t], Yi[idx_t])
        m0 = LinearRegression().fit(Z[idx_c], Yi[idx_c])
        # Evaluate both bridge functions on the treated subpopulation.
        Zt = Z[idx_t]
        return float(np.mean(m1.predict(Zt) - m0.predict(Zt)))

    att_pci = _pci(Y, D, S, X)

    # Bootstrap SE
    boot_s = np.full(n_boot, np.nan)
    boot_p = np.full(n_boot, np.nan)
    for b in range(n_boot):
        idx = rng.integers(0, n, size=n)
        try:
            boot_s[b] = _surrogate(Y[idx], D[idx], S[idx], X[idx])
        except Exception:
            pass
        try:
            boot_p[b] = _pci(Y[idx], D[idx], S[idx], X[idx])
        except Exception:
            pass
    se_s = _bootstrap_se(boot_s, label="bridge.surrogate_pci")
    se_p = _bootstrap_se(boot_p, label="bridge.surrogate_pci")

    if np.isnan(att_surr):
        att_surr, se_s = att_pci, se_p
    if np.isnan(att_pci):
        att_pci, se_p = att_surr, se_s

    diff, diff_se, diff_p = _agreement_test(
        att_surr, se_s, att_pci, se_p
    )
    est_dr, se_dr = _dr_combine(
        att_surr, se_s, att_pci, se_p, diff_p
    )

    return BridgeResult(
        kind="surrogate_pci",
        path_a_name="Surrogate Index",
        path_b_name="PCI linear bridge",
        estimate_a=float(att_surr),
        estimate_b=float(att_pci),
        se_a=se_s,
        se_b=se_p,
        diff=diff,
        diff_se=diff_se,
        diff_p=diff_p,
        estimate_dr=est_dr,
        se_dr=se_dr,
        n_obs=n,
        detail={"n_short_term": len(short_term)},
        reference="Imbens, Kallus, Mao & Wang (2025), JRSS-B 87(2); arXiv 2202.07234",
    )
