r"""
Marginal Structural Models (MSM) via stabilized IPTW.

Robins (1998, 2000) g-methods family for **time-varying** treatments
with time-varying confounders that are themselves affected by past
treatment (the "classical" setting where standard regression adjustment
is biased because you either block a causal path or open a collider).

Setup
-----
Panel data with one row per (unit, period). For each unit :math:`i` and
period :math:`t`:

* :math:`A_{i,t}`: treatment at t (binary or continuous)
* :math:`L_{i,t}`: time-varying confounders at t (measured *before* A_t)
* :math:`V_i`: baseline (time-invariant) covariates
* :math:`Y_{i,T}`: end-of-follow-up outcome (repeated across t for the
  pooled fit — the MSM treats the outcome as a function of cumulative
  treatment history).

Identifying assumptions
-----------------------

1. **Sequential exchangeability**: :math:`Y(a_{0:T}) \perp A_t \mid
   \bar{A}_{t-1}, \bar{L}_t, V` for all t.
2. **Positivity**: :math:`0 < P(A_t = a \mid \bar{A}_{t-1}, \bar{L}_t, V)
   < 1` for every treatment value of interest.
3. **Consistency / no interference**.

Stabilized IPTW
---------------

The stabilized weight for unit i at time t is::

    sw_{i,t} = Π_{s≤t}  P(A_s = a_s | Ā_{s-1}, V)
                         / P(A_s = a_s | Ā_{s-1}, L̄_s, V)

Numerator conditions on baseline only (so weights have mean ~1);
denominator conditions on full time-varying confounder history. The
weighted pooled regression of :math:`Y` on cumulative exposure
(and baseline V) then recovers the **marginal** causal parameter
of the MSM.

For continuous A, replace the probability ratios with the ratio of
Gaussian conditional densities (stabilized density weights).

Inference
---------

Robust cluster-sandwich standard errors are reported, clustering by
unit (the canonical choice — the pseudo-population generated by the
weights induces dependence within unit).

References
----------
Robins, J.M. (1998). "Marginal structural models." *Proceedings of
the American Statistical Association, Section on Bayesian Statistical
Science*, 1-10.

Robins, J.M., Hernán, M.A. and Brumback, B. (2000). "Marginal
structural models and causal inference in epidemiology." *Epidemiology*,
11(5), 550-560. [@robins2000marginal]

Cole, S.R. and Hernán, M.A. (2008). "Constructing Inverse Probability
Weights for Marginal Structural Models." *American Journal of
Epidemiology*, 168(6), 656-664. [@cole2008constructing]
"""

import warnings
from typing import Optional, List, Union
import numpy as np
import pandas as pd
from scipy import stats

from ..core.results import CausalResult


def msm(
    data: pd.DataFrame,
    y: str,
    treat: str,
    id: str,
    time: str,
    time_varying: List[str],
    baseline: Optional[List[str]] = None,
    exposure: str = 'cumulative',
    treat_type: str = 'auto',
    trim: float = 0.01,
    trim_per_period: bool = False,
    alpha: float = 0.05,
    family: str = 'gaussian',
) -> CausalResult:
    """
    Estimate a Marginal Structural Model via stabilized IPTW.

    Parameters
    ----------
    data : pd.DataFrame
        Long-format panel data, one row per (unit, period).
    y : str
        Outcome column. Can be constant within unit (end-of-follow-up
        outcome repeated) or time-varying.
    treat : str
        Treatment column. Binary (0/1) or continuous.
    id : str
        Unit identifier.
    time : str
        Period identifier. Numeric, monotonically increasing within unit.
    time_varying : list of str
        Time-varying confounders. Must be measured *before* treatment
        at period t (lagged variables are the user's responsibility —
        pass pre-treatment values of L at t).
    baseline : list of str, optional
        Time-invariant / baseline covariates V.
    exposure : {'cumulative', 'current', 'ever'}, default 'cumulative'
        How treatment history enters the MSM:

        - ``'cumulative'``: sum of A_s for s ≤ t (dose-response slope)
        - ``'current'``: A_t itself (assumes no carryover)
        - ``'ever'``: 1 if any A_s=1 for s ≤ t (binary, A binary only)
    treat_type : {'auto', 'binary', 'continuous'}, default 'auto'
        Treatment type for weight construction. 'auto' detects from data.
    trim : float, default 0.01
        Truncation quantile applied to weights (symmetric: trims extreme
        weights at ``trim`` and ``1 - trim``). Set to 0 to disable.
    trim_per_period : bool, default False
        When True, apply ``trim`` at each period's per-period ratio
        *before* taking the cumulative product. This is a more
        aggressive stabilisation strategy used in some longitudinal
        literature (Cole & Hernán 2008 §3) — it prevents a single
        outlier period from dominating the cumulative product. When
        False (the default), trimming is applied once on the final
        cumulative-product weight, which is the most common convention
        and matches the original Robins (1998) formulation.
    alpha : float, default 0.05
        Significance level.
    family : {'gaussian', 'binomial'}, default 'gaussian'
        Outcome family for the weighted pooled regression.

    Returns
    -------
    CausalResult
        ``estimate`` is the coefficient on the exposure term in the MSM.
        For continuous A with exposure='cumulative', this is the marginal
        dose-response slope. For binary A with exposure='ever', it is
        the marginal effect of ever being treated.

    Examples
    --------
    >>> # Time-varying treatment with time-varying confounders
    >>> sp.msm(panel, y='cd4_count', treat='hiv_therapy',
    ...        id='pid', time='visit',
    ...        time_varying=['cd4_lag', 'viral_load_lag'],
    ...        baseline=['age', 'sex'])

    >>> # Dose-response of cumulative exposure
    >>> sp.msm(panel, y='bp', treat='drug_dose',
    ...        id='pid', time='month',
    ...        time_varying=['bp_lag', 'weight_lag'],
    ...        exposure='cumulative')

    Notes
    -----
    * The ``time_varying`` list must contain variables measured *before*
      treatment at period t. If you have contemporaneous L, lag it first.
    * MSM is identified under sequential exchangeability + positivity;
      near-violations of positivity (extreme weights) are the main source
      of bias. The returned ``model_info`` reports weight diagnostics.
    """
    if exposure not in ('cumulative', 'current', 'ever'):
        raise ValueError(
            f"exposure must be 'cumulative', 'current', or 'ever'; "
            f"got '{exposure}'"
        )
    if family not in ('gaussian', 'binomial'):
        raise ValueError(
            f"family must be 'gaussian' or 'binomial'; got '{family}'"
        )

    baseline = list(baseline or [])
    required = [y, treat, id, time] + time_varying + baseline
    missing = [c for c in required if c not in data.columns]
    if missing:
        raise ValueError(f"Columns not found in data: {missing}")

    df = data[required].dropna().sort_values([id, time]).reset_index(drop=True)
    if df.empty:
        raise ValueError("No rows remain after dropping missing values.")

    if treat_type == 'auto':
        treat_type = 'binary' if set(df[treat].unique()).issubset({0, 1}) else 'continuous'

    if exposure == 'ever' and treat_type != 'binary':
        raise ValueError("exposure='ever' requires binary treatment.")

    # Build stabilized weights. If per-period trimming is requested,
    # we trim each period's density-ratio before the cumulative
    # product is taken (Cole & Hernán 2008 §3).
    sw = stabilized_weights(
        df,
        treat=treat,
        id=id,
        time=time,
        time_varying=time_varying,
        baseline=baseline,
        treat_type=treat_type,
        trim_per_period=(trim if (trim_per_period and trim and trim > 0) else 0.0),
    )

    # Post-cumulative trimming (the conventional default — once on the
    # final cumulative weight). Only applied when trim_per_period=False,
    # since per-period trimming already stabilised each step above.
    w = sw.copy()
    if trim and trim > 0 and not trim_per_period:
        lo = np.quantile(w, trim)
        hi = np.quantile(w, 1 - trim)
        w = np.clip(w, lo, hi)

    # Build the exposure summary A_t (or cumulative-A, ever-A) per row
    df = df.copy()
    grouped = df.groupby(id, sort=False)[treat]
    if exposure == 'current':
        df['_exposure'] = df[treat].values
        exposure_label = treat
    elif exposure == 'cumulative':
        df['_exposure'] = grouped.cumsum().values
        exposure_label = f'cum_{treat}'
    else:  # ever
        df['_exposure'] = (grouped.cummax() > 0).astype(float).values
        exposure_label = f'ever_{treat}'

    # Pooled weighted regression of Y on exposure + baseline
    feat = ['_exposure'] + baseline
    X_mat = np.column_stack([np.ones(len(df)), df[feat].values.astype(float)])
    y_vec = df[y].values.astype(float)
    cluster_ids = df[id].values

    if family == 'gaussian':
        beta, se = _wls_cluster(X_mat, y_vec, w, cluster_ids)
    else:
        beta, se = _weighted_logit_cluster(X_mat, y_vec, w, cluster_ids)

    # Coefficient of interest is position 1 (after constant)
    coef = float(beta[1])
    se_coef = float(se[1])
    z = coef / se_coef if se_coef > 0 else 0.0
    pvalue = float(2 * (1 - stats.norm.cdf(abs(z))))
    z_crit = stats.norm.ppf(1 - alpha / 2)
    ci = (coef - z_crit * se_coef, coef + z_crit * se_coef)

    # Coefficient table (detail)
    param_names = ['const', exposure_label] + baseline
    detail = pd.DataFrame({
        'term': param_names,
        'estimate': beta,
        'se': se,
        'z': np.where(se > 0, beta / se, 0.0),
        'pvalue': 2 * (1 - stats.norm.cdf(np.abs(np.where(se > 0, beta / se, 0.0)))),
    })

    model_info = {
        'estimator': 'Marginal Structural Model (stabilized IPTW)',
        'exposure_type': exposure,
        'treat_type': treat_type,
        'family': family,
        'n_units': int(df[id].nunique()),
        'n_periods': int(df[time].nunique()),
        'sw_mean': float(np.mean(sw)),
        'sw_max': float(np.max(sw)),
        'sw_min': float(np.min(sw)),
        'sw_trimmed': bool(trim and trim > 0),
        'trim_per_period': bool(trim_per_period),
        'coef_table': detail.copy(),
        'cluster_var': id,
    }

    _result = CausalResult(
        method='Marginal Structural Model (IPTW)',
        estimand='marginal ' + exposure,
        estimate=coef,
        se=se_coef,
        pvalue=pvalue,
        ci=ci,
        alpha=alpha,
        n_obs=len(df),
        detail=detail,
        model_info=model_info,
        _citation_key='msm',
    )
    try:
        from ..output._lineage import attach_provenance as _attach_prov
        _attach_prov(
            _result,
            function="sp.msm",
            params={
                "y": y, "treat": treat,
                "id": id, "time": time,
                "time_varying": list(time_varying),
                "baseline": list(baseline) if baseline else None,
                "exposure": exposure, "treat_type": treat_type,
                "trim": trim, "trim_per_period": trim_per_period,
                "alpha": alpha, "family": family,
            },
            data=data,
            overwrite=False,
        )
    except Exception:  # pragma: no cover
        pass
    return _result


class MarginalStructuralModel:
    """Class wrapper around :func:`msm` for programmatic access."""

    def __init__(self, **kwargs):
        self._kwargs = kwargs
        self.result_: Optional[CausalResult] = None

    def fit(self, data: pd.DataFrame) -> 'MarginalStructuralModel':
        self.result_ = msm(data=data, **self._kwargs)
        return self

    def summary(self) -> str:
        if self.result_ is None:
            raise RuntimeError("Call .fit(data) first.")
        return self.result_.summary()


def stabilized_weights(
    data: pd.DataFrame,
    treat: str,
    id: str,
    time: str,
    time_varying: List[str],
    baseline: Optional[List[str]] = None,
    treat_type: str = 'auto',
    trim_per_period: float = 0.0,
) -> np.ndarray:
    """
    Compute stabilized IPTW weights for time-varying treatments.

    Returns one weight per row (long-format). For treatment :math:`A_t`
    with history :math:`\\bar{A}_{t-1}`, time-varying confounders
    :math:`\\bar{L}_t` (must be pre-treatment), and baseline V:

    * Numerator: conditional density/probability of :math:`A_t` given
      :math:`\\bar{A}_{t-1}` and V (no L).
    * Denominator: conditional density/probability of :math:`A_t` given
      :math:`\\bar{A}_{t-1}`, :math:`\\bar{L}_t`, and V.

    The row weight is the *cumulative product* of per-period ratios up
    to and including period t.

    Parameters
    ----------
    data : pd.DataFrame
        Sorted long-format panel (one row per unit-time).
    treat : str
        Treatment column (binary 0/1 or continuous).
    id, time : str
        Unit and time identifiers.
    time_varying : list of str
        Time-varying confounders (already lagged to pre-treatment).
    baseline : list of str, optional
        Baseline covariates.
    treat_type : {'auto', 'binary', 'continuous'}
        Overrides auto-detection.
    trim_per_period : float, default 0.0
        If > 0, trim each period's density ratio at the symmetric
        quantile ``[trim_per_period, 1 - trim_per_period]`` *before*
        taking the cumulative product. A common value is 0.01. Set to
        0 to disable per-period trimming (the default).

    Returns
    -------
    np.ndarray
        Stabilized weight for each row.
    """
    baseline = list(baseline or [])
    df = data.sort_values([id, time]).reset_index(drop=True)

    if treat_type == 'auto':
        treat_type = 'binary' if set(df[treat].unique()).issubset({0, 1}) else 'continuous'

    A = df[treat].values.astype(float)
    L = df[time_varying].values.astype(float) if time_varying else np.zeros((len(df), 0))
    V = df[baseline].values.astype(float) if baseline else np.zeros((len(df), 0))

    # Lagged treatment — 0 at first period within each unit
    df['_A_lag'] = df.groupby(id, sort=False)[treat].shift(1).fillna(0.0)
    A_lag = df['_A_lag'].values.astype(float)

    # Numerator: P(A_t | A_{t-1}, V)
    X_num = np.column_stack([A_lag.reshape(-1, 1), V])
    # Denominator: P(A_t | A_{t-1}, V, L_t)
    X_den = np.column_stack([A_lag.reshape(-1, 1), V, L])

    if treat_type == 'binary':
        p_num = _logit_proba(X_num, A)
        p_den = _logit_proba(X_den, A)
        # Probability of observed A
        lik_num = np.where(A == 1, p_num, 1 - p_num)
        lik_den = np.where(A == 1, p_den, 1 - p_den)
    else:
        lik_num = _gauss_density(X_num, A)
        lik_den = _gauss_density(X_den, A)

    eps = 1e-10
    ratio = lik_num / np.maximum(lik_den, eps)

    # Optional per-period trimming BEFORE the cumulative product.
    if trim_per_period and trim_per_period > 0:
        lo = float(np.quantile(ratio, trim_per_period))
        hi = float(np.quantile(ratio, 1 - trim_per_period))
        ratio = np.clip(ratio, lo, hi)

    # Cumulative product within unit (ordered by time via sort above)
    df['_ratio'] = ratio
    sw = df.groupby(id, sort=False)['_ratio'].cumprod().values

    return sw


def _logit_proba(X, y):
    """Logistic regression predicted ``P(Y=1|X)``.

    Zero-variance columns (e.g. the all-zero lagged-treatment column in a
    single-period panel) are dropped before fitting — they only duplicate
    the intercept that :func:`statsmodels.api.add_constant` already adds, and
    leaving them in makes the design singular. If the fit still fails (e.g.
    perfect separation), the function falls back to the marginal mean **and
    warns loudly**: that fallback removes the confounding adjustment carried
    by ``X``, so silently swallowing the failure would hide a correctness
    regression (CLAUDE.md §7).
    """
    X = np.asarray(X, dtype=float)
    if X.ndim == 1:
        X = X.reshape(-1, 1)
    if X.shape[1]:
        keep = np.std(X, axis=0) > 0
        X = X[:, keep] if keep.any() else X[:, :0]
    try:
        import statsmodels.api as sm
        design = sm.add_constant(X, has_constant='add')
        fit = sm.Logit(y, design).fit(disp=0, maxiter=200, warn_convergence=False)
        return np.clip(fit.predict(design), 1e-6, 1 - 1e-6)
    except Exception as exc:
        warnings.warn(
            "MSM treatment model failed to fit "
            f"({type(exc).__name__}: {exc}); falling back to the marginal "
            "mean for this term, which drops its confounding adjustment. "
            "Check for perfect separation or collinear confounders.",
            RuntimeWarning,
            stacklevel=2,
        )
        mean = float(np.clip(np.mean(y), 1e-6, 1 - 1e-6))
        return np.full(len(y), mean)


def _gauss_density(X, y):
    """
    Gaussian conditional density p(y | X) from OLS residuals.
    Returns the density evaluated at the observed y_i.
    """
    design = np.column_stack([np.ones(X.shape[0]), X])
    try:
        beta, *_ = np.linalg.lstsq(design, y, rcond=None)
        fitted = design @ beta
        resid = y - fitted
        # Use N - k df for sigma
        k = design.shape[1]
        denom = max(len(y) - k, 1)
        sigma = float(np.sqrt(np.sum(resid**2) / denom))
        sigma = max(sigma, 1e-6)
        dens = stats.norm.pdf(y, loc=fitted, scale=sigma)
        return np.clip(dens, 1e-12, None)
    except Exception:
        mean = float(np.mean(y))
        sigma = max(float(np.std(y, ddof=1)), 1e-6)
        return np.clip(stats.norm.pdf(y, loc=mean, scale=sigma), 1e-12, None)


def _wls_cluster(X, y, w, cluster_ids):
    """
    Weighted least squares with cluster-robust (CR1) variance.

    Returns (beta, se) where se is the cluster-robust SE vector.
    """
    w = np.asarray(w, dtype=float)
    W = np.sqrt(w).reshape(-1, 1)
    Xw = X * W
    yw = y * w**0.5  # y * sqrt(w)
    beta, *_ = np.linalg.lstsq(Xw, yw, rcond=None)
    resid = y - X @ beta
    # Score contributions u_i = w_i * x_i * resid_i
    u = (w * resid)[:, None] * X

    clusters = np.unique(cluster_ids)
    G = len(clusters)
    p = X.shape[1]
    n = len(y)
    u_cluster = np.zeros((G, p))
    for g, c in enumerate(clusters):
        mask = cluster_ids == c
        u_cluster[g] = u[mask].sum(axis=0)

    meat = u_cluster.T @ u_cluster
    bread_inv = np.linalg.pinv((X * w[:, None]).T @ X)
    # CR1 small-sample adjustment
    adj = (G / (G - 1)) * ((n - 1) / (n - p)) if G > 1 and n > p else 1.0
    vcov = adj * (bread_inv @ meat @ bread_inv)
    se = np.sqrt(np.maximum(np.diag(vcov), 0.0))
    return beta, se


def _weighted_logit_cluster(X, y, w, cluster_ids):
    """
    Probability-weighted (IPTW) logistic regression with cluster-robust SE.

    Semantics: IPTW weights are **probability weights** (a.k.a. sampling
    weights / pweights in Stata), not frequency weights. A row with
    weight ``w_i`` represents ``w_i`` expected units in the target
    pseudo-population, not ``w_i`` replicated observations. This matters
    for inference: the correct sandwich uses the weighted score, with
    cluster-robust meat and an information-matrix bread.

    We implement this directly (IRLS) rather than via statsmodels'
    ``freq_weights`` (integer replication) or ``var_weights`` (inverse-
    variance scaling). Neither of those matches the IPTW definition
    without careful reinterpretation — doing it by hand keeps the
    semantics unambiguous and the sandwich coherent.
    """
    n, p = X.shape
    beta = np.zeros(p)
    # Initialize the intercept at the weighted log-odds of y so IRLS has
    # a reasonable starting point. Other coefficients start at 0.
    y_bar_w = float(np.sum(w * y) / max(np.sum(w), 1e-12))
    y_bar_w = float(np.clip(y_bar_w, 1e-4, 1 - 1e-4))
    beta[0] = np.log(y_bar_w / (1 - y_bar_w))

    max_iter = 50
    tol = 1e-8
    for _ in range(max_iter):
        # Stabilize the logistic link against overflow on extreme eta.
        # CRITICAL: the working response z must be built from the SAME eta
        # that produced p_hat — otherwise the Newton step mixes clipped
        # probabilities with unclipped linear predictors and can diverge
        # for extreme rows. Use eta_clip throughout the update.
        eta_clip = np.clip(X @ beta, -30, 30)
        p_hat = 1 / (1 + np.exp(-eta_clip))
        p_hat = np.clip(p_hat, 1e-8, 1 - 1e-8)
        W_diag = w * p_hat * (1 - p_hat)
        # Working response: z = eta + (y - p) / (p*(1-p))
        z = eta_clip + (y - p_hat) / np.clip(p_hat * (1 - p_hat), 1e-8, None)
        # Solve weighted LS: (X' W X) beta = X' W z
        XtWX = (X * W_diag[:, None]).T @ X
        XtWz = (X * W_diag[:, None]).T @ z
        try:
            beta_new = np.linalg.solve(XtWX, XtWz)
        except np.linalg.LinAlgError:
            beta_new = np.linalg.lstsq(XtWX, XtWz, rcond=None)[0]
        if np.max(np.abs(beta_new - beta)) < tol:
            beta = beta_new
            break
        beta = beta_new

    # Cluster-robust sandwich:
    #   Var(beta) = (X' W X)^{-1} [Σ_g U_g U_g'] (X' W X)^{-1}
    # where U_g = Σ_{i ∈ g} w_i * (y_i - p_i) * x_i is the score.
    eta = np.clip(X @ beta, -30, 30)
    p_hat = np.clip(1 / (1 + np.exp(-eta)), 1e-8, 1 - 1e-8)
    resid = y - p_hat
    u = (w * resid)[:, None] * X
    clusters = np.unique(cluster_ids)
    G = len(clusters)
    u_cluster = np.zeros((G, p))
    for g, c in enumerate(clusters):
        mask = cluster_ids == c
        u_cluster[g] = u[mask].sum(axis=0)
    meat = u_cluster.T @ u_cluster
    W_diag = w * p_hat * (1 - p_hat)
    bread_inv = np.linalg.pinv((X * W_diag[:, None]).T @ X)
    adj = (G / (G - 1)) * ((n - 1) / (n - p)) if G > 1 and n > p else 1.0
    vcov = adj * (bread_inv @ meat @ bread_inv)
    se = np.sqrt(np.maximum(np.diag(vcov), 0.0))
    return beta, se


# Citation
CausalResult._CITATIONS['msm'] = (
    "@article{robins2000marginal,\n"
    "  title={Marginal Structural Models and Causal Inference in "
    "Epidemiology},\n"
    "  author={Robins, James M. and Hern{\\'a}n, Miguel A. and "
    "Brumback, Babette},\n"
    "  journal={Epidemiology},\n"
    "  volume={11},\n"
    "  number={5},\n"
    "  pages={550--560},\n"
    "  year={2000}\n"
    "}"
)
