"""Propensity score estimation using Inverse Probability Tilting (IPT)."""

import warnings

import numpy as np
import scipy.linalg
import scipy.optimize
import scipy.special
import statsmodels.api as sm

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


def calculate_pscore_ipt(D, X, iw, quantiles=None):
    r"""Calculate propensity scores using Inverse Probability Tilting for the ATT.

    Implements a specific variant of IPT tailored for estimating the
    Average Treatment Effect on the Treated (ATT). Instead of re-weighting both
    the treated and control groups to match the full sample, it estimates a
    propensity score model that implies a re-weighting of the control group to
    match the covariate distribution of the treated group. This is achieved by
    solving the following optimization problem for the propensity score
    parameters :math:`\gamma` given by

    .. math::
        \widehat{\gamma}^{ipt} = \arg\max_{\gamma \in \Gamma} \mathbb{E}_{n}
        \left[D X^{\prime} \gamma - (1-D) \exp(X^{\prime} \gamma)\right].

    The first-order condition of this problem implies the balancing property

    .. math::
        \sum_{i: D_i=1} w_i X_i = \sum_{i: D_i=0} w_i
        \frac{\widehat{p}(X_i)}{1-\widehat{p}(X_i)} X_i,

    where :math:`\widehat{p}(X) = \text{expit}(X'\widehat{\gamma}^{ipt})` is the estimated
    propensity score (i.e., the logistic function, :math:`\exp(v) / (1 + \exp(v))`)
    and :math:`w_i` are the observation weights. This property ensures that the
    weighted average of covariates in the control group matches the weighted
    average in the treated group, which is a key condition for identifying the ATT.

    Parameters
    ----------
    D : ndarray
        Treatment indicator (1D array).
    X : ndarray
        Covariate matrix (2D array, n_obs x n_features), must include intercept.
    iw : ndarray
        Individual weights (1D array).
    quantiles : dict[int, list[float]] | None
        Dict mapping column indices to quantiles (values between 0 and 1).
        For example, {1: [0.25, 0.5, 0.75]} adds 25th, 50th, 75th percentiles
        of the 2nd column as balance constraints. Default is None (no quantiles).

    Returns
    -------
    ndarray
        Propensity scores.

    Notes
    -----
    The general IPT framework described in [1]_ for the ATE involves solving
    two separate moment equations to find weights for the treated and control
    groups that balance covariates with the full sample. These are given by
    equations (8) and (11) in their paper

    .. math::
        \frac{1}{N} \sum_{i=1}^{N}\left\{\frac{D_{i}}
        {G\left(t\left(X_{i}\right)^{\prime} \delta_{I P T}^{1}\right)}-1\right\}
        t\left(X_{i}\right)=0

    and

    .. math::
        \frac{1}{N} \sum_{i=1}^{N}\left\{\frac{1-D_{i}}
        {1-G\left(t\left(X_{i}\right)^{\prime} \delta_{I P T}^{0}\right)}-1\right\}
        t\left(X_{i}\right)=0.

    This implementation, following [2]_, uses a single objective function tailored for
    ATT estimation. The function attempts to solve this using a trust-constr optimizer,
    falling back to a BFGS optimization of a modified loss function, and finally to a
    standard logit model if the IPT optimizations fail.

    References
    ----------

    .. [1] Graham, B., Pinto, C., and Egel, D. (2012), "Inverse Probability Tilting for Moment
        Condition Models with Missing Data," The Review of Economic Studies, 79(3), 1053-1079.
        https://doi.org/10.1093/restud/rdr047

    .. [2] Sant'Anna, P. H., and Zhao, J. (2020), "Inverse Probability Weighting with
        Missing Data," Journal of the American Statistical Association, 115(530), 1542-1552.
        https://doi.org/10.1080/01621459.2019.1635520
    """
    if np.all(iw == 0):
        warnings.warn(
            "All individual weights are zero. Propensity scores will be uninformative.",
            UserWarning,
        )
        return np.full_like(D, np.nan, dtype=float)

    unique_d = np.unique(D)
    if not np.array_equal(unique_d, [0, 1]):
        warnings.warn(
            f"Treatment indicator D contains values {unique_d}. Expected binary values [0, 1]. "
            "Results may be unreliable.",
            UserWarning,
        )

    X_processed, _ = _remove_collinear_columns(X)

    if quantiles is not None:
        X_processed = _add_quantile_constraints(X_processed, quantiles, iw)

    n_obs, k_features = X_processed.shape
    init_gamma = _get_initial_gamma(D, X_processed, iw, k_features)

    # Try trust-constr optimization first
    try:
        opt_cal_results = scipy.optimize.minimize(
            lambda g, d_arr, x_arr, iw_arr: _loss_ps_cal(g, d_arr, x_arr, iw_arr)[0],
            init_gamma.astype(np.float64),
            args=(D, X_processed, iw),
            method="trust-constr",
            jac=lambda g, d_arr, x_arr, iw_arr: _loss_ps_cal(g, d_arr, x_arr, iw_arr)[1],
            hess=lambda g, d_arr, x_arr, iw_arr: _loss_ps_cal(g, d_arr, x_arr, iw_arr)[2],
            options={"maxiter": 1000},
        )
        if opt_cal_results.success:
            gamma_cal = opt_cal_results.x
        else:
            raise RuntimeError("trust-constr did not converge")
    except (ValueError, np.linalg.LinAlgError, RuntimeError, OverflowError) as e:
        warnings.warn(f"trust-constr optimization failed: {e}. Using IPT algorithm.", UserWarning)

        # Try IPT optimization
        try:
            opt_ipt_results = scipy.optimize.minimize(
                lambda g, d_arr, x_arr, iw_arr, n: _loss_ps_ipt(g, d_arr, x_arr, iw_arr, n)[0],
                init_gamma.astype(np.float64),
                args=(D, X_processed, iw, n_obs),
                method="BFGS",
                jac=lambda g, d_arr, x_arr, iw_arr, n: _loss_ps_ipt(g, d_arr, x_arr, iw_arr, n)[1],
                options={"maxiter": 10000, "gtol": 1e-06},
            )
            if opt_ipt_results.success:
                gamma_cal = opt_ipt_results.x
            else:
                raise RuntimeError("IPT optimization did not converge") from None
        except (ValueError, np.linalg.LinAlgError, RuntimeError, OverflowError) as e_ipt:
            warnings.warn(f"IPT optimization failed: {e_ipt}. Using initial logit estimates.", UserWarning)
            gamma_cal = init_gamma

            # Validate logit fallback
            try:
                logit_model_refit = sm.GLM(D, X_processed, family=sm.families.Binomial(), freq_weights=iw)
                logit_results_refit = logit_model_refit.fit(start_params=init_gamma, maxiter=100)
                if not logit_results_refit.converged:
                    warnings.warn("Initial Logit model (used as fallback) also did not converge.", UserWarning)
            except (ValueError, np.linalg.LinAlgError, RuntimeError, OverflowError):
                warnings.warn("Checking convergence of fallback Logit model failed.", UserWarning)

    # Compute propensity scores
    pscore_linear = X_processed @ gamma_cal
    pscore = scipy.special.expit(pscore_linear)

    if np.any(np.isnan(pscore)):
        warnings.warn(
            "Propensity score model coefficients might have NA/Inf components. "
            "Multicollinearity or lack of variation in covariates is a likely reason. "
            "Resulting pscores contain NaNs.",
            UserWarning,
        )

    return pscore


def _loss_ps_cal(gamma, D, X, iw):
    """Loss function for calibrated propensity score estimation using trust.

    Parameters
    ----------
    gamma : ndarray
        Coefficient vector for propensity score model.
    D : ndarray
        Treatment indicator (1D array).
    X : ndarray
        Covariate matrix (2D array, n_obs x n_features), includes intercept.
    iw : ndarray
        Individual weights (1D array).

    Returns
    -------
    tuple
        (value, gradient, hessian)
    """
    n_obs, k_features = X.shape

    if np.any(np.isnan(gamma)):
        return np.inf, np.full(k_features, np.nan), np.full((k_features, k_features), np.nan)

    ps_ind = X @ gamma
    ps_ind_clipped = np.clip(ps_ind, -500, 500)
    exp_ps_ind = np.exp(ps_ind_clipped)

    value = -np.mean(np.where(D, ps_ind, -exp_ps_ind) * iw)

    grad_terms = np.where(D[:, np.newaxis], 1.0, -exp_ps_ind[:, np.newaxis]) * iw[:, np.newaxis] * X
    gradient = -np.mean(grad_terms, axis=0)

    hess_M_vector = np.where(D, 0.0, -exp_ps_ind) * iw
    hessian_term_matrix = X * hess_M_vector[:, np.newaxis]
    hessian = -(X.T @ hessian_term_matrix) / n_obs
    return value, gradient, hessian


def _loss_ps_ipt(gamma, D, X, iw, n_obs):
    """Loss function for inverse probability tilting propensity score estimation.

    Parameters
    ----------
    gamma : ndarray
        Coefficient vector for propensity score model.
    D : ndarray
        Treatment indicator (1D array).
    X : ndarray
        Covariate matrix (2D array, n_obs x n_features), includes intercept.
    iw : ndarray
        Individual weights (1D array).
    n_obs : int
        Number of observations.

    Returns
    -------
    tuple
        (value, gradient, hessian)
    """
    k_features = X.shape[1]
    if np.any(np.isnan(gamma)):
        return np.inf, np.full(k_features, np.nan), np.full((k_features, k_features), np.nan)

    log_n_minus_1 = np.log(n_obs - 1)
    cn = -(n_obs - 1)
    bn = -n_obs + (n_obs - 1) * log_n_minus_1
    an = -(n_obs - 1) * (1 - log_n_minus_1 + 0.5 * (log_n_minus_1**2))
    v_star = log_n_minus_1

    v = X @ gamma
    v_clipped = np.clip(v, -500, 500)

    phi = np.where(v < v_star, -v - np.exp(v_clipped), an + bn * v + 0.5 * cn * (v**2))
    phi1 = np.where(v < v_star, -1.0 - np.exp(v_clipped), bn + cn * v)
    phi2 = np.where(v < v_star, -np.exp(v_clipped), cn)
    value = -np.sum(iw * (1 - D) * phi + v)

    grad_vec_term = iw * (1 - D) * phi1 + 1.0
    gradient = -(X.T @ grad_vec_term)

    hess_M_ipt_vector = (1 - D) * iw * phi2
    hessian_term_matrix = X * hess_M_ipt_vector[:, np.newaxis]
    hessian = -(hessian_term_matrix.T @ X)
    return value, gradient, hessian


def _get_initial_gamma(D, X, iw, k_features):
    """Get initial gamma values for optimization."""
    xp = get_backend()
    if xp is not np:
        try:
            beta, _ = cupy_logistic_irls(
                xp.asarray(D, dtype=xp.float64),
                xp.asarray(X, dtype=xp.float64),
                xp.asarray(iw, dtype=xp.float64),
            )
            return to_numpy(beta)
        except (RuntimeError, np.linalg.LinAlgError):
            return np.zeros(k_features)

    try:
        logit_model = sm.GLM(D, X, family=sm.families.Binomial(), freq_weights=iw)
        logit_results = logit_model.fit(maxiter=100)
        init_gamma = logit_results.params
        if not logit_results.converged:
            warnings.warn(
                "Initial Logit model for IPT did not converge. Using pseudo-inverse for initial gamma.", UserWarning
            )
            try:
                init_gamma = np.linalg.pinv(X.T @ (iw[:, np.newaxis] * X)) @ (X.T @ (iw * D))
            except np.linalg.LinAlgError:
                warnings.warn("Pseudo-inverse for initial gamma failed. Using zeros.", UserWarning)
                init_gamma = np.zeros(k_features)
        return init_gamma
    except (np.linalg.LinAlgError, ValueError) as e:
        warnings.warn(f"Initial Logit model failed: {e}. Using zeros for initial gamma.", UserWarning)
        return np.zeros(k_features)


def _remove_collinear_columns(X):
    """Remove collinear columns from covariate matrix."""
    n_obs, n_cols = X.shape

    if n_cols == 1:
        return X, []

    _, R, P = scipy.linalg.qr(X, mode="economic", pivoting=True)

    diag_R = np.abs(np.diag(R))
    tol = diag_R[0] * max(n_obs, n_cols) * np.finfo(float).eps
    rank = np.sum(diag_R > tol)

    if rank == n_cols:
        return X, []

    independent_cols = P[:rank]
    removed_cols = P[rank:]

    if len(removed_cols) > 0:
        warnings.warn(
            f"Removed {len(removed_cols)} collinear column(s) from covariate matrix. "
            f"Column indices removed: {removed_cols.tolist()}",
            UserWarning,
        )

    return X[:, independent_cols], removed_cols.tolist()


def _add_quantile_constraints(X, quantiles, iw):
    """Add quantile balance constraints to covariate matrix."""
    quantile_features = []

    for col_idx, q_list in quantiles.items():
        if col_idx >= X.shape[1]:
            warnings.warn(f"Column index {col_idx} exceeds number of columns ({X.shape[1]}). Skipping.", UserWarning)
            continue

        if col_idx == 0 and np.all(X[:, 0] == X[0, 0]):
            warnings.warn("Skipping quantile constraints for intercept column.", UserWarning)
            continue

        col_values = X[:, col_idx]

        for q in q_list:
            if not 0 < q < 1:
                warnings.warn(f"Quantile {q} must be between 0 and 1. Skipping.", UserWarning)
                continue

            threshold = _weighted_quantile(col_values, q, iw)

            q_indicator = (col_values <= threshold).astype(float)
            quantile_features.append(q_indicator)

    if quantile_features:
        X_extended = np.column_stack([X, *quantile_features])
        return X_extended
    return X


def _weighted_quantile(values, q, weights):
    """Compute weighted quantile."""
    sorted_idx = np.argsort(values)
    sorted_values = values[sorted_idx]
    sorted_weights = weights[sorted_idx]

    # Cumulative weights
    cum_weights = np.cumsum(sorted_weights)
    total_weight = cum_weights[-1]

    # Find quantile position
    quantile_weight = q * total_weight
    idx = np.searchsorted(cum_weights, quantile_weight)

    if idx == 0:
        return sorted_values[0]
    if idx >= len(sorted_values):
        return sorted_values[-1]

    # Linear interpolation between adjacent values
    w_below = cum_weights[idx - 1]
    w_above = cum_weights[idx]

    if w_above == w_below:
        return sorted_values[idx]

    # Interpolation factor
    alpha = (quantile_weight - w_below) / (w_above - w_below)
    return sorted_values[idx - 1] + alpha * (sorted_values[idx] - sorted_values[idx - 1])
