"""Nuisance parameter estimation for DDD estimators with repeated cross-section data."""

from __future__ import annotations

import warnings
from typing import NamedTuple

import numpy as np
import statsmodels.api as sm

from moderndid.cupy.backend import get_backend, to_numpy
from moderndid.cupy.regression import cupy_logistic_irls, cupy_wls

from .utils import check_overlap_and_warn, get_comparison_description


class PScoreRCResult(NamedTuple):
    """Result from propensity score estimation for RCS.

    Attributes
    ----------
    propensity_scores : ndarray
        Estimated propensity scores for observations in the subgroup comparison.
    hessian_matrix : ndarray or None
        Hessian matrix from logistic regression, used for influence function.
        None when using REG method.
    keep_ps : ndarray
        Boolean array indicating which observations to keep after trimming.
        Control observations with propensity score >= trim_level are excluded.
    """

    propensity_scores: np.ndarray
    hessian_matrix: np.ndarray | None
    keep_ps: np.ndarray


class OutcomeRegRCResult(NamedTuple):
    """Result from outcome regression estimation for RCS.

    Attributes
    ----------
    y : ndarray
        Outcomes for observations in the subgroup comparison.
    out_y_cont : ndarray
        Outcome regression predictions combining pre and post periods.
    out_y_cont_pre : ndarray
        Outcome regression predictions for control group, pre-period.
    out_y_cont_post : ndarray
        Outcome regression predictions for control group, post-period.
    out_y_treat_pre : ndarray
        Outcome regression predictions for treated group, pre-period.
    out_y_treat_post : ndarray
        Outcome regression predictions for treated group, post-period.
    """

    y: np.ndarray
    out_y_cont: np.ndarray
    out_y_cont_pre: np.ndarray
    out_y_cont_post: np.ndarray
    out_y_treat_pre: np.ndarray
    out_y_treat_post: np.ndarray


class DIDRCResult(NamedTuple):
    """Result from DiD estimation for one subgroup comparison with RCS.

    Attributes
    ----------
    dr_att : float
        Doubly robust ATT estimate for the subgroup comparison.
    inf_func : ndarray
        Influence function for all observations (zeros for observations not in comparison).
    """

    dr_att: float
    inf_func: np.ndarray


def compute_all_nuisances_rc(
    y,
    post,
    subgroup,
    covariates,
    weights,
    est_method="dr",
    trim_level=0.995,
):
    """Compute all nuisance parameters for DDD estimation with repeated cross-section data.

    Parameters
    ----------
    y : ndarray
        A 1D array of outcomes from both pre- and post-treatment periods.
    post : ndarray
        A 1D array of post-treatment dummies (1 if post-treatment, 0 if pre-treatment).
    subgroup : ndarray
        A 1D array of subgroup indicators (1, 2, 3, or 4) for each observation.
    covariates : ndarray
        A 2D array of covariates including intercept.
    weights : ndarray
        A 1D array of observation weights.
    est_method : {"dr", "reg", "ipw"}, default "dr"
        Estimation method.

        - "dr": Doubly robust (both propensity score and outcome regression)
        - "reg": Outcome regression only
        - "ipw": Inverse probability weighting only

    trim_level : float, default 0.995
        Trimming level for propensity scores.

    Returns
    -------
    tuple[list[PScoreRCResult], list[OutcomeRegRCResult]]
        Lists of propensity score and outcome regression results for
        comparison subgroups [3, 2, 1].

    See Also
    --------
    compute_all_did_rc : Compute all DiD components and combine into DDD estimate.
    """
    if est_method not in ["dr", "reg", "ipw"]:
        raise ValueError(f"est_method must be 'dr', 'reg', or 'ipw', got {est_method}")

    pscores = []
    or_results = []

    for comp_subgroup in [3, 2, 1]:
        if est_method == "reg":
            ps_result = _compute_pscore_null_rc(subgroup, comp_subgroup)
        else:
            ps_result = _compute_pscore_rc(subgroup, post, covariates, weights, comp_subgroup, trim_level)
        pscores.append(ps_result)

        if est_method == "ipw":
            or_result = _compute_outcome_regression_null_rc(y, post, subgroup, comp_subgroup)
        else:
            or_result = _compute_outcome_regression_rc(
                y, post, subgroup, covariates, weights, comp_subgroup, est_method
            )
        or_results.append(or_result)

    return pscores, or_results


def compute_all_did_rc(
    y,
    post,
    subgroup,
    covariates,
    weights,
    pscores,
    or_results,
    est_method,
    n_total,
):
    """Compute all DiD components and combine into DDD estimate for RCS.

    Computes the DiD for each of the three subgroup comparisons and combines
    them using the triple difference formula.

    Parameters
    ----------
    y : ndarray
        A 1D array of outcomes from both periods.
    post : ndarray
        A 1D array of post-treatment dummies.
    subgroup : ndarray
        A 1D array of subgroup indicators (1, 2, 3, or 4) for each observation.
    covariates : ndarray
        A 2D array of covariates including intercept.
    weights : ndarray
        A 1D array of observation weights.
    pscores : list[PScoreRCResult]
        Propensity score results for comparisons [3, 2, 1].
    or_results : list[OutcomeRegRCResult]
        Outcome regression results for comparisons [3, 2, 1].
    est_method : {"dr", "reg", "ipw"}
        Estimation method.
    n_total : int
        Total number of observations.

    Returns
    -------
    tuple[list[DIDRCResult], float, ndarray]
        A tuple containing:

        - List of DiD results for each comparison [3, 2, 1]
        - The DDD ATT estimate
        - The combined influence function (rescaled by subgroup sizes)

    See Also
    --------
    compute_all_nuisances_rc : Compute all nuisance parameters for DDD estimation.
    """
    did_results = []
    for i, comp_subgroup in enumerate([3, 2, 1]):
        did_result = _compute_did_rc(
            _y=y,
            post=post,
            subgroup=subgroup,
            covariates=covariates,
            weights=weights,
            comparison_subgroup=comp_subgroup,
            pscore_result=pscores[i],
            or_result=or_results[i],
            est_method=est_method,
            n_total=n_total,
        )
        did_results.append(did_result)

    xp = get_backend()
    dr_att_3, inf_func_3 = did_results[0].dr_att, did_results[0].inf_func
    dr_att_2, inf_func_2 = did_results[1].dr_att, did_results[1].inf_func
    dr_att_1, inf_func_1 = did_results[2].dr_att, did_results[2].inf_func

    ddd_att = dr_att_3 + dr_att_2 - dr_att_1

    n = n_total
    n3 = int(xp.sum((subgroup == 4) | (subgroup == 3)))
    n2 = int(xp.sum((subgroup == 4) | (subgroup == 2)))
    n1 = int(xp.sum((subgroup == 4) | (subgroup == 1)))

    w3 = n / n3 if n3 > 0 else 0
    w2 = n / n2 if n2 > 0 else 0
    w1 = n / n1 if n1 > 0 else 0

    inf_func = w3 * inf_func_3 + w2 * inf_func_2 - w1 * inf_func_1

    return did_results, ddd_att, inf_func


def _compute_did_rc(
    _y,
    post,
    subgroup,
    covariates,
    weights,
    comparison_subgroup,
    pscore_result,
    or_result,
    est_method,
    n_total,
):
    """Compute doubly robust DiD for one subgroup comparison with RCS."""
    xp = get_backend()
    mask = (subgroup == 4) | (subgroup == comparison_subgroup)
    sub_subgroup = subgroup[mask]
    sub_post = post[mask]
    sub_covariates = covariates[mask]
    sub_weights = weights[mask]

    pscore = pscore_result.propensity_scores
    hessian = pscore_result.hessian_matrix
    keep_ps = pscore_result.keep_ps.astype(float)

    sub_y = or_result.y
    out_y_cont = or_result.out_y_cont
    out_y_cont_pre = or_result.out_y_cont_pre
    out_y_cont_post = or_result.out_y_cont_post
    out_y_treat_pre = or_result.out_y_treat_pre
    out_y_treat_post = or_result.out_y_treat_post

    pa4 = (sub_subgroup == 4).astype(float)
    pa_comp = (sub_subgroup == comparison_subgroup).astype(float)

    w_treat_pre = keep_ps * sub_weights * pa4 * (1 - sub_post)
    w_treat_post = keep_ps * sub_weights * pa4 * sub_post

    if est_method == "reg":
        w_cont_pre = None
        w_cont_post = None
        w_reg_control = keep_ps * sub_weights * pa4
    else:
        if xp is np:
            with np.errstate(divide="ignore", invalid="ignore"):
                w_cont_pre = keep_ps * sub_weights * pscore * pa_comp * (1 - sub_post) / (1 - pscore)
                w_cont_post = keep_ps * sub_weights * pscore * pa_comp * sub_post / (1 - pscore)
        else:
            w_cont_pre = keep_ps * sub_weights * pscore * pa_comp * (1 - sub_post) / (1 - pscore)
            w_cont_post = keep_ps * sub_weights * pscore * pa_comp * sub_post / (1 - pscore)
        w_cont_pre = xp.nan_to_num(w_cont_pre)
        w_cont_post = xp.nan_to_num(w_cont_post)
        w_reg_control = None

    w_d = keep_ps * sub_weights * pa4
    w_dt1 = keep_ps * sub_weights * pa4 * sub_post
    w_dt0 = keep_ps * sub_weights * pa4 * (1 - sub_post)

    mean_w_treat_pre = xp.mean(w_treat_pre)
    mean_w_treat_post = xp.mean(w_treat_post)
    mean_w_d = xp.mean(w_d)
    mean_w_dt1 = xp.mean(w_dt1)
    mean_w_dt0 = xp.mean(w_dt0)

    if mean_w_treat_pre == 0 or mean_w_treat_post == 0:
        raise ValueError(
            f"No effectively treated observations (subgroup 4) in comparison with subgroup {comparison_subgroup}."
        )

    if est_method == "reg":
        mean_w_reg_control = xp.mean(w_reg_control)

        eta_treat_pre = w_treat_pre * sub_y / mean_w_treat_pre
        eta_treat_post = w_treat_post * sub_y / mean_w_treat_post

        reg_att_control = w_reg_control * (out_y_cont_post - out_y_cont_pre)
        eta_reg_control = xp.mean(reg_att_control) / mean_w_reg_control

        att_treat_pre = xp.mean(eta_treat_pre)
        att_treat_post = xp.mean(eta_treat_post)

        dr_att = (att_treat_post - att_treat_pre) - eta_reg_control

        att_cont_pre = None
        att_cont_post = None
        att_d_post = None
        att_dt1_post = None
        att_d_pre = None
        att_dt0_pre = None
        eta_cont_pre = None
        eta_cont_post = None
        eta_d_post = None
        eta_dt1_post = None
        eta_d_pre = None
        eta_dt0_pre = None
        mean_w_cont_pre = None
        mean_w_cont_post = None

    else:
        # DR or IPW estimator
        mean_w_cont_pre = xp.mean(w_cont_pre)
        mean_w_cont_post = xp.mean(w_cont_post)

        if mean_w_cont_pre == 0 or mean_w_cont_post == 0:
            raise ValueError(f"No effectively control observations (subgroup {comparison_subgroup}) after weighting.")

        eta_treat_pre = w_treat_pre * (sub_y - out_y_cont) / mean_w_treat_pre
        eta_treat_post = w_treat_post * (sub_y - out_y_cont) / mean_w_treat_post
        eta_cont_pre = w_cont_pre * (sub_y - out_y_cont) / mean_w_cont_pre
        eta_cont_post = w_cont_post * (sub_y - out_y_cont) / mean_w_cont_post

        eta_d_post = w_d * (out_y_treat_post - out_y_cont_post) / mean_w_d
        eta_dt1_post = w_dt1 * (out_y_treat_post - out_y_cont_post) / mean_w_dt1
        eta_d_pre = w_d * (out_y_treat_pre - out_y_cont_pre) / mean_w_d
        eta_dt0_pre = w_dt0 * (out_y_treat_pre - out_y_cont_pre) / mean_w_dt0

        att_treat_pre = xp.mean(eta_treat_pre)
        att_treat_post = xp.mean(eta_treat_post)
        att_cont_pre = xp.mean(eta_cont_pre)
        att_cont_post = xp.mean(eta_cont_post)
        att_d_post = xp.mean(eta_d_post)
        att_dt1_post = xp.mean(eta_dt1_post)
        att_d_pre = xp.mean(eta_d_pre)
        att_dt0_pre = xp.mean(eta_dt0_pre)

        dr_att = (
            (att_treat_post - att_treat_pre)
            - (att_cont_post - att_cont_pre)
            + (att_d_post - att_dt1_post)
            - (att_d_pre - att_dt0_pre)
        )

        w_reg_control = None
        mean_w_reg_control = None
        eta_reg_control = None
        reg_att_control = None

    inf_func_sub = _compute_inf_func_rc(
        sub_y=sub_y,
        sub_post=sub_post,
        sub_covariates=sub_covariates,
        sub_weights=sub_weights,
        pa4=pa4,
        pa_comp=pa_comp,
        pscore=pscore,
        hessian=hessian,
        out_y_cont=out_y_cont,
        out_y_cont_pre=out_y_cont_pre,
        out_y_cont_post=out_y_cont_post,
        out_y_treat_pre=out_y_treat_pre,
        out_y_treat_post=out_y_treat_post,
        w_treat_pre=w_treat_pre,
        w_treat_post=w_treat_post,
        w_cont_pre=w_cont_pre,
        w_cont_post=w_cont_post,
        w_d=w_d,
        w_dt1=w_dt1,
        w_dt0=w_dt0,
        eta_treat_pre=eta_treat_pre,
        eta_treat_post=eta_treat_post,
        eta_cont_pre=eta_cont_pre,
        eta_cont_post=eta_cont_post,
        eta_d_post=eta_d_post,
        eta_dt1_post=eta_dt1_post,
        eta_d_pre=eta_d_pre,
        eta_dt0_pre=eta_dt0_pre,
        att_treat_pre=att_treat_pre,
        att_treat_post=att_treat_post,
        att_cont_pre=att_cont_pre,
        att_cont_post=att_cont_post,
        att_d_post=att_d_post,
        att_dt1_post=att_dt1_post,
        att_d_pre=att_d_pre,
        att_dt0_pre=att_dt0_pre,
        est_method=est_method,
        w_reg_control=w_reg_control,
        eta_reg_control=eta_reg_control if est_method == "reg" else None,
        reg_att_control=reg_att_control,
    )

    inf_func = xp.zeros(n_total)
    mask_indices = xp.where(mask)[0]
    inf_func[mask_indices] = inf_func_sub

    return DIDRCResult(dr_att=float(dr_att), inf_func=inf_func)


def _compute_inf_func_rc(
    sub_y,
    sub_post,
    sub_covariates,
    sub_weights,
    pa4,
    pa_comp,
    pscore,
    hessian,
    out_y_cont,
    out_y_cont_pre,
    out_y_cont_post,
    out_y_treat_pre,
    out_y_treat_post,
    w_treat_pre,
    w_treat_post,
    w_cont_pre,
    w_cont_post,
    w_d,
    w_dt1,
    w_dt0,
    eta_treat_pre,
    eta_treat_post,
    eta_cont_pre,
    eta_cont_post,
    eta_d_post,
    eta_dt1_post,
    eta_d_pre,
    eta_dt0_pre,
    att_treat_pre,
    att_treat_post,
    att_cont_pre,
    att_cont_post,
    att_d_post,
    att_dt1_post,
    att_d_pre,
    att_dt0_pre,
    est_method,
    w_reg_control=None,
    eta_reg_control=None,
    reg_att_control=None,
):
    """Compute influence function for the DiD comparison with RCS."""
    xp = get_backend()
    n_sub = len(sub_weights)

    mean_w_treat_pre = xp.mean(w_treat_pre)
    mean_w_treat_post = xp.mean(w_treat_post)
    mean_w_d = xp.mean(w_d)
    mean_w_dt1 = xp.mean(w_dt1)
    mean_w_dt0 = xp.mean(w_dt0)

    if est_method == "reg":
        mean_w_reg_control = xp.mean(w_reg_control)

        reg_att_treat_pre = w_treat_pre * sub_y
        reg_att_treat_post = w_treat_post * sub_y

        inf_treat_pre = (reg_att_treat_pre - w_treat_pre * att_treat_pre) / mean_w_treat_pre
        inf_treat_post = (reg_att_treat_post - w_treat_post * att_treat_post) / mean_w_treat_post
        inf_treat = inf_treat_post - inf_treat_pre

        weights_ols_pre = sub_weights * pa_comp * (1 - sub_post)
        wols_x_pre = weights_ols_pre[:, None] * sub_covariates
        wols_eX_pre = (weights_ols_pre * (sub_y - out_y_cont_pre))[:, None] * sub_covariates
        XpX_pre = (wols_x_pre.T @ sub_covariates) / n_sub

        s = xp.linalg.svd(XpX_pre, compute_uv=False)
        cond_pre = s[0] / s[-1] if s[-1] > 0 else float("inf")
        XpX_inv_pre = xp.linalg.pinv(XpX_pre) if cond_pre > 1 / xp.finfo(float).eps else xp.linalg.inv(XpX_pre)
        asy_lin_rep_ols_pre = wols_eX_pre @ XpX_inv_pre

        weights_ols_post = sub_weights * pa_comp * sub_post
        wols_x_post = weights_ols_post[:, None] * sub_covariates
        wols_eX_post = (weights_ols_post * (sub_y - out_y_cont_post))[:, None] * sub_covariates
        XpX_post = (wols_x_post.T @ sub_covariates) / n_sub

        s = xp.linalg.svd(XpX_post, compute_uv=False)
        cond_post = s[0] / s[-1] if s[-1] > 0 else float("inf")
        XpX_inv_post = xp.linalg.pinv(XpX_post) if cond_post > 1 / xp.finfo(float).eps else xp.linalg.inv(XpX_post)
        asy_lin_rep_ols_post = wols_eX_post @ XpX_inv_post

        inf_control_1 = reg_att_control - w_reg_control * eta_reg_control
        M1 = xp.mean((w_reg_control)[:, None] * sub_covariates, axis=0)
        inf_control_2_post = asy_lin_rep_ols_post @ M1
        inf_control_2_pre = asy_lin_rep_ols_pre @ M1
        inf_control = (inf_control_1 + inf_control_2_post - inf_control_2_pre) / mean_w_reg_control

        inf_func = inf_treat - inf_control
        return inf_func

    mean_w_cont_pre = xp.mean(w_cont_pre)
    mean_w_cont_post = xp.mean(w_cont_post)

    inf_treat_pre = eta_treat_pre - w_treat_pre * att_treat_pre / mean_w_treat_pre
    inf_treat_post = eta_treat_post - w_treat_post * att_treat_post / mean_w_treat_post

    inf_cont_pre = eta_cont_pre - w_cont_pre * att_cont_pre / mean_w_cont_pre
    inf_cont_post = eta_cont_post - w_cont_post * att_cont_post / mean_w_cont_post

    cont_moment_pre = (
        xp.mean((w_cont_pre * (sub_y - out_y_cont - att_cont_pre))[:, None] * sub_covariates, axis=0) / mean_w_cont_pre
    )
    cont_moment_post = (
        xp.mean((w_cont_post * (sub_y - out_y_cont - att_cont_post))[:, None] * sub_covariates, axis=0)
        / mean_w_cont_post
    )

    score_ps = (sub_weights * (pa4 - pscore))[:, None] * sub_covariates
    asy_lin_rep_ps = score_ps @ hessian
    inf_cont_ps = asy_lin_rep_ps @ (cont_moment_post - cont_moment_pre)

    if est_method == "ipw":
        inf_treat_or = xp.zeros(n_sub)
        inf_cont_or = xp.zeros(n_sub)
        inf_or_adj = xp.zeros(n_sub)
    else:
        weights_ols_pre = sub_weights * pa_comp * (1 - sub_post)
        weighted_x_pre = weights_ols_pre[:, None] * sub_covariates
        weighted_resid_x_pre = (weights_ols_pre * (sub_y - out_y_cont_pre))[:, None] * sub_covariates
        gram_pre = (weighted_x_pre.T @ sub_covariates) / n_sub

        s = xp.linalg.svd(gram_pre, compute_uv=False)
        cond_pre = s[0] / s[-1] if s[-1] > 0 else float("inf")
        gram_inv_pre = xp.linalg.pinv(gram_pre) if cond_pre > 1 / xp.finfo(float).eps else xp.linalg.inv(gram_pre)
        asy_lin_rep_ols_pre = weighted_resid_x_pre @ gram_inv_pre

        weights_ols_post = sub_weights * pa_comp * sub_post
        weighted_x_post = weights_ols_post[:, None] * sub_covariates
        weighted_resid_x_post = (weights_ols_post * (sub_y - out_y_cont_post))[:, None] * sub_covariates
        gram_post = (weighted_x_post.T @ sub_covariates) / n_sub

        s = xp.linalg.svd(gram_post, compute_uv=False)
        cond_post = s[0] / s[-1] if s[-1] > 0 else float("inf")
        gram_inv_post = xp.linalg.pinv(gram_post) if cond_post > 1 / xp.finfo(float).eps else xp.linalg.inv(gram_post)
        asy_lin_rep_ols_post = weighted_resid_x_post @ gram_inv_post

        weights_ols_pre_treat = sub_weights * pa4 * (1 - sub_post)
        weighted_x_pre_treat = weights_ols_pre_treat[:, None] * sub_covariates
        weighted_resid_x_pre_treat = (weights_ols_pre_treat * (sub_y - out_y_treat_pre))[:, None] * sub_covariates

        sum_pre_treat = xp.sum(weights_ols_pre_treat)
        if sum_pre_treat > 0:
            gram_pre_treat = (weighted_x_pre_treat.T @ sub_covariates) / n_sub
            s = xp.linalg.svd(gram_pre_treat, compute_uv=False)
            cond_pre_treat = s[0] / s[-1] if s[-1] > 0 else float("inf")
            if cond_pre_treat > 1 / xp.finfo(float).eps:
                gram_inv_pre_treat = xp.linalg.pinv(gram_pre_treat)
            else:
                gram_inv_pre_treat = xp.linalg.inv(gram_pre_treat)
            asy_lin_rep_ols_pre_treat = weighted_resid_x_pre_treat @ gram_inv_pre_treat
        else:
            asy_lin_rep_ols_pre_treat = xp.zeros((n_sub, sub_covariates.shape[1]))

        weights_ols_post_treat = sub_weights * pa4 * sub_post
        weighted_x_post_treat = weights_ols_post_treat[:, None] * sub_covariates
        weighted_resid_x_post_treat = (weights_ols_post_treat * (sub_y - out_y_treat_post))[:, None] * sub_covariates

        sum_post_treat = xp.sum(weights_ols_post_treat)
        if sum_post_treat > 0:
            gram_post_treat = (weighted_x_post_treat.T @ sub_covariates) / n_sub
            s = xp.linalg.svd(gram_post_treat, compute_uv=False)
            cond_post_treat = s[0] / s[-1] if s[-1] > 0 else float("inf")
            if cond_post_treat > 1 / xp.finfo(float).eps:
                gram_inv_post_treat = xp.linalg.pinv(gram_post_treat)
            else:
                gram_inv_post_treat = xp.linalg.inv(gram_post_treat)
            asy_lin_rep_ols_post_treat = weighted_resid_x_post_treat @ gram_inv_post_treat
        else:
            asy_lin_rep_ols_post_treat = xp.zeros((n_sub, sub_covariates.shape[1]))

        treat_moment_post = -xp.mean((w_treat_post * sub_post)[:, None] * sub_covariates, axis=0) / mean_w_treat_post
        treat_moment_pre = -xp.mean((w_treat_pre * (1 - sub_post))[:, None] * sub_covariates, axis=0) / mean_w_treat_pre

        inf_treat_or_post = asy_lin_rep_ols_post @ treat_moment_post
        inf_treat_or_pre = asy_lin_rep_ols_pre @ treat_moment_pre
        inf_treat_or = inf_treat_or_post + inf_treat_or_pre

        cont_reg_moment_post = -xp.mean((w_cont_post * sub_post)[:, None] * sub_covariates, axis=0) / mean_w_cont_post
        cont_reg_moment_pre = (
            -xp.mean((w_cont_pre * (1 - sub_post))[:, None] * sub_covariates, axis=0) / mean_w_cont_pre
        )

        inf_cont_or_post = asy_lin_rep_ols_post @ cont_reg_moment_post
        inf_cont_or_pre = asy_lin_rep_ols_pre @ cont_reg_moment_pre
        inf_cont_or = inf_cont_or_post + inf_cont_or_pre

        mom_post = xp.mean(((w_d / mean_w_d) - (w_dt1 / mean_w_dt1))[:, None] * sub_covariates, axis=0)
        mom_pre = xp.mean(((w_d / mean_w_d) - (w_dt0 / mean_w_dt0))[:, None] * sub_covariates, axis=0)
        inf_or_post = (asy_lin_rep_ols_post_treat - asy_lin_rep_ols_post) @ mom_post
        inf_or_pre = (asy_lin_rep_ols_pre_treat - asy_lin_rep_ols_pre) @ mom_pre
        inf_or_adj = inf_or_post - inf_or_pre

    inf_treat = inf_treat_post - inf_treat_pre + inf_treat_or
    inf_cont = inf_cont_post - inf_cont_pre + inf_cont_ps + inf_cont_or

    inf_eff1 = eta_d_post - w_d * att_d_post / mean_w_d
    inf_eff2 = eta_dt1_post - w_dt1 * att_dt1_post / mean_w_dt1
    inf_eff3 = eta_d_pre - w_d * att_d_pre / mean_w_d
    inf_eff4 = eta_dt0_pre - w_dt0 * att_dt0_pre / mean_w_dt0
    inf_eff = (inf_eff1 - inf_eff2) - (inf_eff3 - inf_eff4)

    inf_func = inf_treat - inf_cont + inf_eff + inf_or_adj

    return inf_func


def _compute_pscore_rc(subgroup, _post, covariates, weights, comparison_subgroup, trim_level=0.995):
    """Compute propensity scores for a subgroup comparison with RCS."""
    if comparison_subgroup not in [1, 2, 3]:
        raise ValueError(f"comparison_subgroup must be 1, 2, or 3, got {comparison_subgroup}")

    mask = (subgroup == 4) | (subgroup == comparison_subgroup)
    sub_covariates = covariates[mask]
    sub_weights = weights[mask]
    sub_subgroup = subgroup[mask]

    pa4 = (sub_subgroup == 4).astype(float)

    xp = get_backend()
    if xp is not np:
        try:
            beta, ps_fit = cupy_logistic_irls(
                xp.asarray(pa4, dtype=xp.float64),
                xp.asarray(sub_covariates, dtype=xp.float64),
                xp.asarray(sub_weights, dtype=xp.float64),
            )
            if xp.any(xp.isnan(beta)):
                comparison_desc = get_comparison_description(comparison_subgroup)
                raise ValueError(
                    f"Propensity score model has NA coefficients.\n"
                    f"  Comparison: {comparison_desc} units.\n"
                    f"  This is likely due to multicollinearity among covariates."
                )
        except (np.linalg.LinAlgError, RuntimeError) as e:
            raise ValueError(
                f"Failed to estimate propensity scores for subgroup {comparison_subgroup} due to singular matrix."
            ) from e

        check_overlap_and_warn(to_numpy(ps_fit), comparison_subgroup)

        keep_ps = xp.ones(len(pa4), dtype=bool)
        keep_ps[pa4 == 0] = ps_fit[pa4 == 0] < trim_level
        ps_fit = xp.minimum(ps_fit, 1 - 1e-6)

        n_sub = len(sub_weights)
        W = sub_weights * ps_fit * (1 - ps_fit)
        hessian_ps = sub_covariates.T @ (W[:, None] * sub_covariates)
        hessian_matrix = xp.linalg.inv(hessian_ps) * n_sub

        return PScoreRCResult(propensity_scores=ps_fit, hessian_matrix=hessian_matrix, keep_ps=keep_ps)

    try:
        pscore_model = sm.GLM(pa4, sub_covariates, family=sm.families.Binomial(), freq_weights=sub_weights)
        pscore_results = pscore_model.fit(disp=False)

        if not pscore_results.converged:
            comparison_desc = get_comparison_description(comparison_subgroup)
            warnings.warn(
                f"Propensity score model did not converge.\n"
                f"  Comparison: {comparison_desc} units.\n"
                f"  Consider using fewer covariates or checking for separation issues.",
                UserWarning,
            )

        if np.any(np.isnan(pscore_results.params)):
            comparison_desc = get_comparison_description(comparison_subgroup)
            raise ValueError(
                f"Propensity score model has NA coefficients.\n"
                f"  Comparison: {comparison_desc} units.\n"
                f"  This is likely due to multicollinearity among covariates."
            )

        ps_fit = pscore_results.predict(sub_covariates)

    except np.linalg.LinAlgError as e:
        raise ValueError(
            f"Failed to estimate propensity scores for subgroup {comparison_subgroup} due to singular matrix."
        ) from e

    check_overlap_and_warn(ps_fit, comparison_subgroup)

    keep_ps = np.ones(len(pa4), dtype=bool)
    keep_ps[pa4 == 0] = ps_fit[pa4 == 0] < trim_level
    ps_fit = np.minimum(ps_fit, 1 - 1e-6)

    n_sub = len(sub_weights)
    hessian_matrix = pscore_results.cov_params() * n_sub

    return PScoreRCResult(propensity_scores=ps_fit, hessian_matrix=hessian_matrix, keep_ps=keep_ps)


def _compute_pscore_null_rc(subgroup, comparison_subgroup):
    """Compute null propensity scores for REG method with RCS."""
    xp = get_backend()
    mask = (subgroup == 4) | (subgroup == comparison_subgroup)
    n_sub = int(xp.sum(mask))

    return PScoreRCResult(
        propensity_scores=xp.ones(n_sub),
        hessian_matrix=None,
        keep_ps=xp.ones(n_sub, dtype=bool),
    )


def _compute_outcome_regression_rc(y, post, subgroup, covariates, weights, comparison_subgroup, _est_method):
    """Compute outcome regression for a subgroup comparison with RCS."""
    if comparison_subgroup not in [1, 2, 3]:
        raise ValueError(f"comparison_subgroup must be 1, 2, or 3, got {comparison_subgroup}")

    mask = (subgroup == 4) | (subgroup == comparison_subgroup)
    sub_y = y[mask]
    sub_post = post[mask]
    sub_subgroup = subgroup[mask]
    sub_covariates = covariates[mask]
    sub_weights = weights[mask]

    n_features = sub_covariates.shape[1]

    out_y_cont_pre = _fit_ols_cell(
        y=sub_y,
        post=sub_post,
        d=(sub_subgroup == 4).astype(int),
        covariates=sub_covariates,
        weights=sub_weights,
        pre=True,
        treat=False,
        n_features=n_features,
        comparison_subgroup=comparison_subgroup,
    )

    out_y_cont_post = _fit_ols_cell(
        y=sub_y,
        post=sub_post,
        d=(sub_subgroup == 4).astype(int),
        covariates=sub_covariates,
        weights=sub_weights,
        pre=False,
        treat=False,
        n_features=n_features,
        comparison_subgroup=comparison_subgroup,
    )

    out_y_treat_pre = _fit_ols_cell(
        y=sub_y,
        post=sub_post,
        d=(sub_subgroup == 4).astype(int),
        covariates=sub_covariates,
        weights=sub_weights,
        pre=True,
        treat=True,
        n_features=n_features,
        comparison_subgroup=comparison_subgroup,
    )

    out_y_treat_post = _fit_ols_cell(
        y=sub_y,
        post=sub_post,
        d=(sub_subgroup == 4).astype(int),
        covariates=sub_covariates,
        weights=sub_weights,
        pre=False,
        treat=True,
        n_features=n_features,
        comparison_subgroup=comparison_subgroup,
    )

    out_y_cont = sub_post * out_y_cont_post + (1 - sub_post) * out_y_cont_pre

    return OutcomeRegRCResult(
        y=sub_y,
        out_y_cont=out_y_cont,
        out_y_cont_pre=out_y_cont_pre,
        out_y_cont_post=out_y_cont_post,
        out_y_treat_pre=out_y_treat_pre,
        out_y_treat_post=out_y_treat_post,
    )


def _fit_ols_cell(y, post, d, covariates, weights, pre, treat, n_features, comparison_subgroup):
    """Fit OLS for a specific (D, T) cell."""
    if pre and treat:
        subs = (d == 1) & (post == 0)
    elif not pre and treat:
        subs = (d == 1) & (post == 1)
    elif pre and not treat:
        subs = (d == 0) & (post == 0)
    else:
        subs = (d == 0) & (post == 1)

    xp = get_backend()
    n_subs = int(xp.sum(subs))

    if n_subs == 0 or n_subs < n_features:
        return xp.full(len(y), float("nan"))

    sub_y = y[subs]
    sub_x = covariates[subs]
    sub_weights = weights[subs]

    try:
        if xp is not np:
            beta, _ = cupy_wls(xp.asarray(sub_y), xp.asarray(sub_x), xp.asarray(sub_weights))
            if xp.any(xp.isnan(beta)):
                subg_desc = "Treated-Eligible" if treat else get_comparison_description(comparison_subgroup)
                period_desc = "pre-treatment" if pre else "post-treatment"
                raise ValueError(
                    f"Outcome regression model has NA coefficients.\n"
                    f"  Cell: {subg_desc} units in {period_desc} period.\n"
                    f"  This is likely due to multicollinearity among covariates."
                )
            fitted_values = xp.asarray(covariates, dtype=xp.float64) @ beta
        else:
            wls_model = sm.WLS(sub_y, sub_x, weights=sub_weights)
            results = wls_model.fit()
            coefficients = results.params

            if np.any(np.isnan(coefficients)):
                subg_desc = "Treated-Eligible" if treat else get_comparison_description(comparison_subgroup)
                period_desc = "pre-treatment" if pre else "post-treatment"
                raise ValueError(
                    f"Outcome regression model has NA coefficients.\n"
                    f"  Cell: {subg_desc} units in {period_desc} period.\n"
                    f"  This is likely due to multicollinearity among covariates."
                )
            fitted_values = covariates @ coefficients

        return fitted_values

    except (np.linalg.LinAlgError, ValueError) as e:
        if "NA coefficients" in str(e):
            raise
        return xp.full(len(y), float("nan"))


def _compute_outcome_regression_null_rc(y, _post, subgroup, comparison_subgroup):
    """Compute null outcome regression for IPW method with RCS."""
    xp = get_backend()
    mask = (subgroup == 4) | (subgroup == comparison_subgroup)
    sub_y = y[mask]
    n_sub = len(sub_y)

    zeros = xp.zeros(n_sub)

    return OutcomeRegRCResult(
        y=sub_y,
        out_y_cont=zeros,
        out_y_cont_pre=zeros,
        out_y_cont_post=zeros,
        out_y_treat_pre=zeros,
        out_y_treat_post=zeros,
    )
