"""
Triply Robust Panel (TROP) estimator.

Implements the TROP estimator from Athey, Imbens, Qu & Viviano (2025).
TROP combines three robustness components:
1. Nuclear norm regularized factor model (interactive fixed effects)
2. Exponential distance-based unit weights
3. Exponential time decay weights

The estimator uses leave-one-out cross-validation for tuning parameter
selection and provides robust treatment effect estimates under factor
confounding.

References
----------
Athey, S., Imbens, G. W., Qu, Z., & Viviano, D. (2025). Triply Robust Panel
Estimators. *Working Paper*. https://arxiv.org/abs/2508.21536
"""

import logging
import warnings
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import pandas as pd

logger = logging.getLogger(__name__)

from diff_diff._backend import (
    HAS_RUST_BACKEND,
    _rust_loocv_grid_search,
)
from diff_diff.trop_global import TROPGlobalMixin
from diff_diff.trop_local import TROPLocalMixin, _setup_trop_data
from diff_diff.trop_results import (
    _LAMBDA_INF,
    _PrecomputedStructures,
    TROPResults,
)
from diff_diff.utils import safe_inference, warn_if_not_converged


class TROP(TROPLocalMixin, TROPGlobalMixin):
    """
    Triply Robust Panel (TROP) estimator.

    Implements the exact methodology from Athey, Imbens, Qu & Viviano (2025).
    TROP combines three robustness components:

    1. **Nuclear norm regularized factor model**: Estimates interactive fixed
       effects L_it via matrix completion with nuclear norm penalty ||L||_*

    2. **Exponential distance-based unit weights**: ω_j = exp(-λ_unit × d(j,i))
       where d(j,i) is the RMSE of outcome differences between units

    3. **Exponential time decay weights**: θ_s = exp(-λ_time × :math:`|s-t|`)
       weighting pre-treatment periods by proximity to treatment

    Tuning parameters (λ_time, λ_unit, λ_nn) are selected via leave-one-out
    cross-validation on control observations.

    Parameters
    ----------
    method : str, default='local'
        Estimation method to use:

        - 'local': Per-observation model fitting following Algorithm 2 of
          Athey et al. (2025). Computes observation-specific weights and fits
          a model for each treated observation, averaging the individual
          treatment effects. More flexible but computationally intensive.

        - 'global': Computationally efficient adaptation using the (1-W)
          masking principle from Eq. 2. Fits a single model on control
          observations with global weights, then computes per-observation
          treatment effects as residuals:
          tau_it = Y_it - mu - alpha_i - beta_t - L_it for treated cells.
          ATT is the mean of these effects. For the paper's full
          per-treated-cell estimator, use ``method='local'``.

    lambda_time_grid : list, optional
        Grid of time weight decay parameters. 0.0 = uniform weights (disabled).
        Must not contain inf. Default: [0, 0.1, 0.5, 1, 2, 5].
    lambda_unit_grid : list, optional
        Grid of unit weight decay parameters. 0.0 = uniform weights (disabled).
        Must not contain inf. Default: [0, 0.1, 0.5, 1, 2, 5].
    lambda_nn_grid : list, optional
        Grid of nuclear norm regularization parameters. inf = factor model
        disabled (L=0). Default: [0, 0.01, 0.1, 1].
    max_iter : int, default=100
        Maximum iterations for nuclear norm optimization.
    tol : float, default=1e-6
        Convergence tolerance for optimization.
    alpha : float, default=0.05
        Significance level for confidence intervals.
    n_bootstrap : int, default=200
        Number of bootstrap replications for variance estimation. Must be >= 2.
    seed : int, optional
        Random seed for reproducibility.

    Attributes
    ----------
    results_ : TROPResults
        Estimation results after calling fit().
    is_fitted_ : bool
        Whether the model has been fitted.

    Examples
    --------
    >>> from diff_diff import TROP
    >>> trop = TROP()
    >>> results = trop.fit(
    ...     data,
    ...     outcome='outcome',
    ...     treatment='treated',
    ...     unit='unit',
    ...     time='period',
    ... )
    >>> results.print_summary()

    References
    ----------
    Athey, S., Imbens, G. W., Qu, Z., & Viviano, D. (2025). Triply Robust
    Panel Estimators. *Working Paper*. https://arxiv.org/abs/2508.21536
    """

    def __init__(
        self,
        method: str = "local",
        lambda_time_grid: Optional[List[float]] = None,
        lambda_unit_grid: Optional[List[float]] = None,
        lambda_nn_grid: Optional[List[float]] = None,
        max_iter: int = 100,
        tol: float = 1e-6,
        alpha: float = 0.05,
        n_bootstrap: int = 200,
        seed: Optional[int] = None,
    ):
        # Validate method parameter
        valid_methods = ("local", "global")
        if method not in valid_methods:
            raise ValueError(f"method must be one of {valid_methods}, got '{method}'")
        self.method = method

        # Default grids from paper
        self.lambda_time_grid = lambda_time_grid or [0.0, 0.1, 0.5, 1.0, 2.0, 5.0]
        self.lambda_unit_grid = lambda_unit_grid or [0.0, 0.1, 0.5, 1.0, 2.0, 5.0]
        self.lambda_nn_grid = lambda_nn_grid or [0.0, 0.01, 0.1, 1.0, 10.0]

        if n_bootstrap < 2:
            raise ValueError(
                "n_bootstrap must be >= 2 for TROP (bootstrap variance "
                "estimation is always used)"
            )

        self.max_iter = max_iter
        self.tol = tol
        self.alpha = alpha
        self.n_bootstrap = n_bootstrap
        self.seed = seed

        # Validate that time/unit grids do not contain inf.
        # Per Athey et al. (2025) Eq. 3, λ_time=0 and λ_unit=0 give uniform
        # weights (exp(-0 × dist) = 1). Using inf is a misunderstanding of
        # the paper's convention. Only λ_nn=∞ is valid (disables factor model).
        for grid_name, grid_vals in [
            ("lambda_time_grid", self.lambda_time_grid),
            ("lambda_unit_grid", self.lambda_unit_grid),
        ]:
            if any(np.isinf(v) for v in grid_vals):
                raise ValueError(
                    f"{grid_name} must not contain inf. Use 0.0 for uniform "
                    f"weights (disabled) per Athey et al. (2025) Eq. 3: "
                    f"exp(-0 × dist) = 1 for all distances."
                )

        # Internal state
        self.results_: Optional[TROPResults] = None
        self.is_fitted_: bool = False
        self._optimal_lambda: Optional[Tuple[float, float, float]] = None

        # Pre-computed structures (set during fit)
        self._precomputed: Optional[_PrecomputedStructures] = None

    # =========================================================================
    # Parameter search (used by local method's fit() path)
    # =========================================================================

    def _univariate_loocv_search(
        self,
        Y: np.ndarray,
        D: np.ndarray,
        control_mask: np.ndarray,
        control_unit_idx: np.ndarray,
        n_units: int,
        n_periods: int,
        param_name: str,
        grid: List[float],
        fixed_params: Dict[str, float],
    ) -> Tuple[float, float]:
        """
        Search over one parameter with others fixed.

        Following paper's footnote 2, this performs a univariate grid search
        for one tuning parameter while holding others fixed. The fixed_params
        use 0.0 for disabled time/unit weights and _LAMBDA_INF for disabled
        factor model:
        - lambda_nn = inf: Skip nuclear norm regularization (L=0)
        - lambda_time = 0.0: Uniform time weights (exp(-0×dist)=1)
        - lambda_unit = 0.0: Uniform unit weights (exp(-0×dist)=1)

        Parameters
        ----------
        Y : np.ndarray
            Outcome matrix (n_periods x n_units).
        D : np.ndarray
            Treatment indicator matrix (n_periods x n_units).
        control_mask : np.ndarray
            Boolean mask for control observations.
        control_unit_idx : np.ndarray
            Indices of control units.
        n_units : int
            Number of units.
        n_periods : int
            Number of periods.
        param_name : str
            Name of parameter to search: 'lambda_time', 'lambda_unit', or 'lambda_nn'.
        grid : List[float]
            Grid of values to search over.
        fixed_params : Dict[str, float]
            Fixed values for other parameters. May include _LAMBDA_INF for lambda_nn.

        Returns
        -------
        Tuple[float, float]
            (best_value, best_score) for the searched parameter.
        """
        best_score = np.inf
        best_value = grid[0] if grid else 0.0

        for value in grid:
            params = {**fixed_params, param_name: value}

            lambda_time = params.get("lambda_time", 0.0)
            lambda_unit = params.get("lambda_unit", 0.0)
            lambda_nn = params.get("lambda_nn", 0.0)

            # Convert λ_nn=∞ → large finite value (factor model disabled, L≈0)
            # λ_time and λ_unit use 0.0 for uniform weights per Eq. 3 (no inf conversion needed)
            if np.isinf(lambda_nn):
                lambda_nn = 1e10

            try:
                score = self._loocv_score_obs_specific(
                    Y,
                    D,
                    control_mask,
                    control_unit_idx,
                    lambda_time,
                    lambda_unit,
                    lambda_nn,
                    n_units,
                    n_periods,
                )
                if score < best_score:
                    best_score = score
                    best_value = value
            except (np.linalg.LinAlgError, ValueError):
                continue

        return best_value, best_score

    def _cycling_parameter_search(
        self,
        Y: np.ndarray,
        D: np.ndarray,
        control_mask: np.ndarray,
        control_unit_idx: np.ndarray,
        n_units: int,
        n_periods: int,
        initial_lambda: Tuple[float, float, float],
        max_cycles: int = 10,
    ) -> Tuple[float, float, float]:
        """
        Cycle through parameters until convergence (coordinate descent).

        Following paper's footnote 2 (Stage 2), this iteratively optimizes
        each tuning parameter while holding the others fixed, until convergence.

        Parameters
        ----------
        Y : np.ndarray
            Outcome matrix (n_periods x n_units).
        D : np.ndarray
            Treatment indicator matrix (n_periods x n_units).
        control_mask : np.ndarray
            Boolean mask for control observations.
        control_unit_idx : np.ndarray
            Indices of control units.
        n_units : int
            Number of units.
        n_periods : int
            Number of periods.
        initial_lambda : Tuple[float, float, float]
            Initial values (lambda_time, lambda_unit, lambda_nn).
        max_cycles : int, default=10
            Maximum number of coordinate descent cycles.

        Returns
        -------
        Tuple[float, float, float]
            Optimized (lambda_time, lambda_unit, lambda_nn).
        """
        lambda_time, lambda_unit, lambda_nn = initial_lambda
        prev_score = np.inf

        for cycle in range(max_cycles):
            # Optimize λ_unit (fix λ_time, λ_nn)
            lambda_unit, _ = self._univariate_loocv_search(
                Y,
                D,
                control_mask,
                control_unit_idx,
                n_units,
                n_periods,
                "lambda_unit",
                self.lambda_unit_grid,
                {"lambda_time": lambda_time, "lambda_nn": lambda_nn},
            )

            # Optimize λ_time (fix λ_unit, λ_nn)
            lambda_time, _ = self._univariate_loocv_search(
                Y,
                D,
                control_mask,
                control_unit_idx,
                n_units,
                n_periods,
                "lambda_time",
                self.lambda_time_grid,
                {"lambda_unit": lambda_unit, "lambda_nn": lambda_nn},
            )

            # Optimize λ_nn (fix λ_unit, λ_time)
            lambda_nn, score = self._univariate_loocv_search(
                Y,
                D,
                control_mask,
                control_unit_idx,
                n_units,
                n_periods,
                "lambda_nn",
                self.lambda_nn_grid,
                {"lambda_unit": lambda_unit, "lambda_time": lambda_time},
            )

            # Check convergence
            if abs(score - prev_score) < 1e-6:
                logger.debug(
                    "Cycling search converged after %d cycles with score %.6f", cycle + 1, score
                )
                break
            prev_score = score

        return lambda_time, lambda_unit, lambda_nn

    # =========================================================================
    # Main fit method
    # =========================================================================

    def fit(
        self,
        data: pd.DataFrame,
        outcome: str,
        treatment: str,
        unit: str,
        time: str,
        survey_design=None,
    ) -> TROPResults:
        """
        Fit the TROP model.

        Parameters
        ----------
        data : pd.DataFrame
            Panel data with observations for multiple units over multiple
            time periods.
        outcome : str
            Name of the outcome variable column.
        treatment : str
            Name of the treatment indicator column (0/1).

            IMPORTANT: This should be an ABSORBING STATE indicator, not a
            treatment timing indicator. For each unit, D=1 for ALL periods
            during and after treatment:

            - D[t, i] = 0 for all t < g_i (pre-treatment periods)
            - D[t, i] = 1 for all t >= g_i (treatment and post-treatment)

            where g_i is the treatment start time for unit i.

            For staggered adoption, different units can have different g_i.
            The ATT averages over ALL D=1 cells per Equation 1 of the paper.
        unit : str
            Name of the unit identifier column.
        time : str
            Name of the time period column.
        survey_design : SurveyDesign, optional
            Survey design specification. Supports pweight, strata, PSU, and
            FPC. Full-design surveys (strata/PSU/FPC) use Rao-Wu rescaled
            bootstrap; Rust backend is pweight-only (Python fallback for
            full design). Survey weights enter ATT aggregation only.

        Returns
        -------
        TROPResults
            Object containing the ATT estimate, standard error,
            factor estimates, and tuning parameters. The lambda_*
            attributes show the selected grid values. For lambda_time and
            lambda_unit, 0.0 means uniform weights; inf is not accepted.
            For lambda_nn, inf is converted to 1e10 (factor model disabled).

        Raises
        ------
        ValueError
            If required columns are missing or non-pweight survey design.
        """
        # Validate inputs
        required_cols = [outcome, treatment, unit, time]
        missing = [c for c in required_cols if c not in data.columns]
        if missing:
            raise ValueError(f"Missing columns: {missing}")

        # Resolve survey design
        from diff_diff.survey import (
            _resolve_survey_for_fit,
            _validate_unit_constant_survey,
        )

        resolved_survey, _survey_weights, _survey_wt, survey_metadata = _resolve_survey_for_fit(
            survey_design, data, "analytical"
        )
        # Reject replicate-weight designs — TROP uses Rao-Wu bootstrap
        if resolved_survey is not None and resolved_survey.uses_replicate_variance:
            raise NotImplementedError(
                "TROP does not yet support replicate-weight survey designs. "
                "Use a TSL-based survey design (strata/psu/fpc)."
            )
        # Validate weight_type is pweight (keep restriction), but allow
        # strata/PSU/FPC — those are handled via Rao-Wu rescaled bootstrap.
        if resolved_survey is not None and resolved_survey.weight_type != "pweight":
            raise ValueError(
                "TROP requires pweight survey weights. "
                f"Got weight_type='{resolved_survey.weight_type}'."
            )
        if resolved_survey is not None:
            _validate_unit_constant_survey(data, unit, survey_design)

        # Dispatch based on estimation method
        if self.method == "global":
            return self._fit_global(
                data,
                outcome,
                treatment,
                unit,
                time,
                resolved_survey=resolved_survey,
                survey_metadata=survey_metadata,
                survey_design=survey_design,
            )

        # Below is the local method (default)
        _ctx = _setup_trop_data(
            data, outcome, treatment, unit, time, resolved_survey, survey_design
        )
        n_units = _ctx["n_units"]
        n_periods = _ctx["n_periods"]
        idx_to_unit = _ctx["idx_to_unit"]
        idx_to_period = _ctx["idx_to_period"]
        unit_weight_arr = _ctx["unit_weight_arr"]
        Y = _ctx["Y"]
        D = _ctx["D"]
        n_treated_obs = _ctx["n_treated_obs"]
        treated_unit_idx = _ctx["treated_unit_idx"]
        control_unit_idx = _ctx["control_unit_idx"]
        n_pre_periods = _ctx["n_pre_periods"]
        n_post_periods = _ctx["n_post_periods"]

        # Step 1: Grid search with LOOCV for tuning parameters
        best_lambda = None
        best_score = np.inf

        # Control observations mask (for LOOCV)
        control_mask = D == 0

        # Pre-compute structures that are reused across LOOCV iterations
        self._precomputed = self._precompute_structures(Y, D, control_unit_idx, n_units, n_periods)

        # Use Rust backend for parallel LOOCV grid search (10-50x speedup)
        if HAS_RUST_BACKEND and _rust_loocv_grid_search is not None:
            try:
                # Prepare inputs for Rust function
                control_mask_u8 = control_mask.astype(np.uint8)
                time_dist_matrix = self._precomputed["time_dist_matrix"].astype(np.int64)

                lambda_time_arr = np.array(self.lambda_time_grid, dtype=np.float64)
                lambda_unit_arr = np.array(self.lambda_unit_grid, dtype=np.float64)
                lambda_nn_arr = np.array(self.lambda_nn_grid, dtype=np.float64)

                result = _rust_loocv_grid_search(
                    Y,
                    D.astype(np.float64),
                    control_mask_u8,
                    time_dist_matrix,
                    lambda_time_arr,
                    lambda_unit_arr,
                    lambda_nn_arr,
                    self.max_iter,
                    self.tol,
                )
                # Unpack result - 7 values including optional first_failed_obs
                best_lt, best_lu, best_ln, best_score, n_valid, n_attempted, first_failed_obs = (
                    result
                )
                # Only accept finite scores - infinite means all fits failed
                if np.isfinite(best_score):
                    best_lambda = (best_lt, best_lu, best_ln)
                # else: best_lambda stays None, triggering defaults fallback
                # Emit warnings consistent with Python implementation
                if n_valid == 0:
                    # Include failed observation coordinates if available (Issue 2 fix)
                    obs_info = ""
                    if first_failed_obs is not None:
                        t_idx, i_idx = first_failed_obs
                        obs_info = f" First failure at observation ({t_idx}, {i_idx})."
                    warnings.warn(
                        f"LOOCV: All {n_attempted} fits failed for "
                        f"\u03bb=({best_lt}, {best_lu}, {best_ln}). "
                        f"Returning infinite score.{obs_info}",
                        UserWarning,
                    )
                elif n_attempted > 0 and (n_attempted - n_valid) > 0.1 * n_attempted:
                    n_failed = n_attempted - n_valid
                    # Include failed observation coordinates if available
                    obs_info = ""
                    if first_failed_obs is not None:
                        t_idx, i_idx = first_failed_obs
                        obs_info = f" First failure at observation ({t_idx}, {i_idx})."
                    warnings.warn(
                        f"LOOCV: {n_failed}/{n_attempted} fits failed for "
                        f"\u03bb=({best_lt}, {best_lu}, {best_ln}). "
                        f"This may indicate numerical instability.{obs_info}",
                        UserWarning,
                    )
            except Exception as e:
                # Fall back to Python implementation on error
                logger.debug("Rust LOOCV grid search failed, falling back to Python: %s", e)
                warnings.warn(
                    f"Rust backend failed for LOOCV grid search; "
                    f"falling back to Python. Performance may be reduced. "
                    f"Error: {e}",
                    UserWarning,
                    stacklevel=2,
                )
                best_lambda = None
                best_score = np.inf

        # Fall back to Python implementation if Rust unavailable or failed
        # Uses two-stage approach per paper's footnote 2:
        # Stage 1: Univariate searches for initial values
        # Stage 2: Cycling (coordinate descent) until convergence
        if best_lambda is None:
            # Stage 1: Univariate searches with extreme fixed values
            # Following paper's footnote 2 for initial bounds

            # λ_time search: fix λ_unit=0, λ_nn=∞ (disabled - no factor adjustment)
            lambda_time_init, _ = self._univariate_loocv_search(
                Y,
                D,
                control_mask,
                control_unit_idx,
                n_units,
                n_periods,
                "lambda_time",
                self.lambda_time_grid,
                {"lambda_unit": 0.0, "lambda_nn": _LAMBDA_INF},
            )

            # λ_nn search: fix λ_time=0 (uniform time weights), λ_unit=0
            lambda_nn_init, _ = self._univariate_loocv_search(
                Y,
                D,
                control_mask,
                control_unit_idx,
                n_units,
                n_periods,
                "lambda_nn",
                self.lambda_nn_grid,
                {"lambda_time": 0.0, "lambda_unit": 0.0},
            )

            # λ_unit search: fix λ_nn=∞, λ_time=0
            lambda_unit_init, _ = self._univariate_loocv_search(
                Y,
                D,
                control_mask,
                control_unit_idx,
                n_units,
                n_periods,
                "lambda_unit",
                self.lambda_unit_grid,
                {"lambda_nn": _LAMBDA_INF, "lambda_time": 0.0},
            )

            # Stage 2: Cycling refinement (coordinate descent)
            lambda_time, lambda_unit, lambda_nn = self._cycling_parameter_search(
                Y,
                D,
                control_mask,
                control_unit_idx,
                n_units,
                n_periods,
                (lambda_time_init, lambda_unit_init, lambda_nn_init),
            )

            # Compute final score for the optimized parameters
            try:
                best_score = self._loocv_score_obs_specific(
                    Y,
                    D,
                    control_mask,
                    control_unit_idx,
                    lambda_time,
                    lambda_unit,
                    lambda_nn,
                    n_units,
                    n_periods,
                )
                # Only accept finite scores - infinite means all fits failed
                if np.isfinite(best_score):
                    best_lambda = (lambda_time, lambda_unit, lambda_nn)
                # else: best_lambda stays None, triggering defaults fallback
            except (np.linalg.LinAlgError, ValueError):
                # If even the optimized parameters fail, best_lambda stays None
                pass

        if best_lambda is None:
            warnings.warn("All tuning parameter combinations failed. Using defaults.", UserWarning)
            best_lambda = (1.0, 1.0, 0.1)
            best_score = np.nan

        self._optimal_lambda = best_lambda
        lambda_time, lambda_unit, lambda_nn = best_lambda

        # Store original λ_nn for results (only λ_nn needs original→effective conversion).
        # λ_time and λ_unit use 0.0 for uniform weights directly per Eq. 3.
        original_lambda_nn = lambda_nn

        # Convert λ_nn=∞ → large finite value (factor model disabled, L≈0)
        if np.isinf(lambda_nn):
            lambda_nn = 1e10

        # effective_lambda with converted λ_nn for ALL downstream computation
        # (variance estimation uses the same parameters as point estimation)
        effective_lambda = (lambda_time, lambda_unit, lambda_nn)

        # Step 2: Final estimation - per-observation model fitting following Algorithm 2
        # For each treated (i,t): compute observation-specific weights, fit model, compute tau_{it}
        treatment_effects = {}
        tau_values = []
        tau_weights = []  # parallel to tau_values for survey-weighted ATT
        alpha_estimates = []
        beta_estimates = []
        L_estimates = []

        # Use pre-computed treated observations
        treated_observations = self._precomputed["treated_observations"]
        nonconverg_tracker: list = []
        n_fits_attempted = 0

        for t, i in treated_observations:
            unit_id = idx_to_unit[i]
            time_id = idx_to_period[t]

            # Skip observations where outcome is missing -- record NaN but
            # don't fit the model or include in tau_values (avoids NaN poisoning)
            if not np.isfinite(Y[t, i]):
                treatment_effects[(unit_id, time_id)] = np.nan
                continue

            # Compute observation-specific weights for this (i, t)
            weight_matrix = self._compute_observation_weights(
                Y, D, i, t, lambda_time, lambda_unit, control_unit_idx, n_units, n_periods
            )

            # Fit model with these weights
            n_fits_attempted += 1
            alpha_hat, beta_hat, L_hat = self._estimate_model(
                Y,
                control_mask,
                weight_matrix,
                lambda_nn,
                n_units,
                n_periods,
                _nonconvergence_tracker=nonconverg_tracker,
            )

            # Compute treatment effect: tau_{it} = Y_{it} - alpha_i - beta_t - L_{it}
            tau_it = Y[t, i] - alpha_hat[i] - beta_hat[t] - L_hat[t, i]

            treatment_effects[(unit_id, time_id)] = tau_it
            tau_values.append(tau_it)
            if unit_weight_arr is not None:
                tau_weights.append(unit_weight_arr[i])

            # Store for averaging
            alpha_estimates.append(alpha_hat)
            beta_estimates.append(beta_hat)
            L_estimates.append(L_hat)

        if nonconverg_tracker:
            warn_if_not_converged(
                False,
                f"TROP local per-treated-observation fit: "
                f"{len(nonconverg_tracker)} of {n_fits_attempted} "
                f"fits did not converge",
                self.max_iter,
                self.tol,
            )

        # Count valid treated observations
        n_valid_treated = len(tau_values)
        if n_valid_treated == 0:
            warnings.warn(
                "All treated outcomes are NaN/missing. Cannot estimate ATT.",
                UserWarning,
            )
        elif n_valid_treated < n_treated_obs:
            warnings.warn(
                f"Only {n_valid_treated} of {n_treated_obs} treated outcomes are finite. "
                "df and n_treated_obs reflect valid observations only.",
                UserWarning,
            )

        # Average ATT (survey-weighted when applicable)
        if unit_weight_arr is not None and tau_values:
            att = float(np.average(tau_values, weights=tau_weights))
        else:
            att = np.mean(tau_values) if tau_values else np.nan

        # Average parameter estimates for output (representative)
        alpha_hat = np.mean(alpha_estimates, axis=0) if alpha_estimates else np.zeros(n_units)
        beta_hat = np.mean(beta_estimates, axis=0) if beta_estimates else np.zeros(n_periods)
        L_hat = np.mean(L_estimates, axis=0) if L_estimates else np.zeros((n_periods, n_units))

        # Compute effective rank
        _, s, _ = np.linalg.svd(L_hat, full_matrices=False)
        if s[0] > 0:
            effective_rank = np.sum(s) / s[0]
        else:
            effective_rank = 0.0

        # Step 4: Variance estimation
        # Use effective_lambda (converted values) to ensure SE is computed with same
        # parameters as point estimation. This fixes the variance inconsistency issue.
        se, bootstrap_dist = self._bootstrap_variance(
            data,
            outcome,
            treatment,
            unit,
            time,
            effective_lambda,
            Y=Y,
            D=D,
            control_unit_idx=control_unit_idx,
            survey_design=survey_design,
            unit_weight_arr=unit_weight_arr,
            resolved_survey=resolved_survey,
        )

        # Compute test statistics
        df_trop = max(1, n_valid_treated - 1)
        t_stat, p_value, conf_int = safe_inference(att, se, alpha=self.alpha, df=df_trop)

        # Create results dictionaries
        unit_effects_dict = {idx_to_unit[i]: alpha_hat[i] for i in range(n_units)}
        time_effects_dict = {idx_to_period[t]: beta_hat[t] for t in range(n_periods)}

        # Store results
        self.results_ = TROPResults(
            att=att,
            se=se,
            t_stat=t_stat,
            p_value=p_value,
            conf_int=conf_int,
            n_obs=len(data),
            n_treated=len(treated_unit_idx),
            n_control=len(control_unit_idx),
            n_treated_obs=int(n_valid_treated),
            unit_effects=unit_effects_dict,
            time_effects=time_effects_dict,
            treatment_effects=treatment_effects,
            lambda_time=lambda_time,
            lambda_unit=lambda_unit,
            lambda_nn=original_lambda_nn,
            factor_matrix=L_hat,
            effective_rank=effective_rank,
            loocv_score=best_score,
            alpha=self.alpha,
            n_pre_periods=n_pre_periods,
            n_post_periods=n_post_periods,
            n_bootstrap=self.n_bootstrap,
            bootstrap_distribution=bootstrap_dist if len(bootstrap_dist) > 0 else None,
            survey_metadata=survey_metadata,
        )

        self.is_fitted_ = True
        return self.results_

    # =========================================================================
    # sklearn-like API
    # =========================================================================

    def get_params(self) -> Dict[str, Any]:
        """Get estimator parameters."""
        return {
            "method": self.method,
            "lambda_time_grid": self.lambda_time_grid,
            "lambda_unit_grid": self.lambda_unit_grid,
            "lambda_nn_grid": self.lambda_nn_grid,
            "max_iter": self.max_iter,
            "tol": self.tol,
            "alpha": self.alpha,
            "n_bootstrap": self.n_bootstrap,
            "seed": self.seed,
        }

    def set_params(self, **params) -> "TROP":
        """Set estimator parameters."""
        for key, value in params.items():
            if key == "method" and value not in ("local", "global"):
                raise ValueError(f"method must be one of ('local', 'global'), got '{value}'")
            if hasattr(self, key):
                setattr(self, key, value)
            else:
                raise ValueError(f"Unknown parameter: {key}")
        return self


def trop(
    data: pd.DataFrame,
    outcome: str,
    treatment: str,
    unit: str,
    time: str,
    survey_design=None,
    **kwargs,
) -> TROPResults:
    """
    Convenience function for TROP estimation.

    Parameters
    ----------
    data : pd.DataFrame
        Panel data.
    outcome : str
        Outcome variable column name.
    treatment : str
        Treatment indicator column name (0/1).

        IMPORTANT: This should be an ABSORBING STATE indicator, not a treatment
        timing indicator. For each unit, D=1 for ALL periods during and after
        treatment (D[t,i]=0 for t < g_i, D[t,i]=1 for t >= g_i where g_i is
        the treatment start time for unit i).
    unit : str
        Unit identifier column name.
    time : str
        Time period column name.
    survey_design : SurveyDesign, optional
        Survey design specification. Supports pweight, strata, PSU, and FPC.
    **kwargs
        Additional arguments passed to TROP constructor.

    Returns
    -------
    TROPResults
        Estimation results.

    Examples
    --------
    >>> from diff_diff import trop
    >>> results = trop(data, 'y', 'treated', 'unit', 'time')
    >>> print(f"ATT: {results.att:.3f}")
    """
    estimator = TROP(**kwargs)
    return estimator.fit(data, outcome, treatment, unit, time, survey_design=survey_design)
