"""
Causal Forest implementation for heterogeneous treatment effect estimation

This module implements the Causal Forest algorithm for estimating conditional 
average treatment effects (CATE) based on the methodology from:
- Wager, S., & Athey, S. (2018). Estimation and inference of heterogeneous 
  treatment effects using random forests. Journal of the American Statistical 
  Association, 113(523), 1228-1242.

The implementation is inspired by and partially based on the EconML library:
- https://github.com/py-why/econml/
- Microsoft Corporation. (2019). EconML: A Python Package for ML-Based 
  Heterogeneous Treatment Effects Estimation.

Key features:
- Honest random forests for unbiased treatment effect estimation
- Bootstrap confidence intervals
- Compatible with StatsPAI outreg2 export functionality
- Both formula and direct array interfaces
"""

import numpy as np
import pandas as pd
from typing import Optional, Union, Tuple, Dict, Any, List
from sklearn.base import BaseEstimator
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier
from sklearn.linear_model import LinearRegression, LogisticRegression
from sklearn.model_selection import cross_val_predict
from sklearn.tree import DecisionTreeRegressor
import warnings

# Import our core classes
from ..core.base import BaseModel
from ..core.results import EconometricResults
from ..core.utils import parse_formula


class CausalForest(BaseModel):
    """
    Causal Forest for heterogeneous treatment effect estimation
    
    This class implements the Causal Forest algorithm, which uses random forests
    to estimate conditional average treatment effects (CATE) in a non-parametric way.
    
    The method combines ideas from:
    1. Honest estimation to avoid overfitting
    2. Double machine learning to handle confounding
    3. Random forests for flexible function approximation
    
    Parameters
    ----------
    n_estimators : int, default=100
        Number of trees in the forest
    min_samples_leaf : int, default=5
        Minimum number of samples required to be at a leaf node
    max_depth : int, default=None
        Maximum depth of trees
    max_samples : float, default=0.5
        Fraction of samples to use for each tree
    model_y : estimator, optional
        Model for outcome regression (first stage)
    model_t : estimator, optional  
        Model for treatment propensity (first stage)
    discrete_treatment : bool, default=True
        Whether treatment is discrete (binary/categorical) or continuous
    honest : bool, default=True
        Whether to use honest estimation (separate samples for splitting and effects)
    bootstrap : bool, default=True
        Whether to use bootstrap sampling for trees
    random_state : int, optional
        Random state for reproducibility
    n_jobs : int, default=1
        Number of parallel jobs
    verbose : int, default=0
        Verbosity level
        
    Attributes
    ----------
    fitted_ : bool
        Whether the model has been fitted
    params : pd.Series
        Not applicable for non-parametric methods, returns empty Series
    std_errors : pd.Series  
        Not applicable for non-parametric methods, returns empty Series
    tvalues : pd.Series
        Not applicable for non-parametric methods, returns empty Series
    pvalues : np.ndarray
        Not applicable for non-parametric methods, returns empty array
    diagnostics : dict
        Model diagnostics and fit statistics
    data_info : dict
        Information about the data used in fitting
    
    Notes
    -----
    This implementation is inspired by the EconML library's CausalForestDML
    but adapted to fit the StatsPAI architecture and interface.
    
    Examples
    --------
    >>> import numpy as np
    >>> import pandas as pd
    >>> from statspai.forest import CausalForest
    >>> 
    >>> # Generate sample data
    >>> np.random.seed(42)
    >>> n = 1000
    >>> X = np.random.normal(0, 1, (n, 3))
    >>> T = np.random.binomial(1, 0.5, n)
    >>> Y = X[:, 0] * T + X[:, 1] + np.random.normal(0, 1, n)
    >>> data = pd.DataFrame({
    ...     'Y': Y, 'T': T, 'X1': X[:, 0], 'X2': X[:, 1], 'X3': X[:, 2]
    ... })
    >>> 
    >>> # Fit Causal Forest
    >>> cf = CausalForest(n_estimators=50, random_state=42)
    >>> cf.fit('Y ~ T | X1 + X2 + X3', data=data)
    >>> 
    >>> # Estimate treatment effects
    >>> cate = cf.effect(data[['X1', 'X2', 'X3']])
    >>> print(f"Average treatment effect: {cate.mean():.3f}")
    """
    
    def __init__(
        self,
        n_estimators: int = 100,
        min_samples_leaf: int = 5,
        max_depth: Optional[int] = None,
        max_samples: float = 0.5,
        model_y: Optional[BaseEstimator] = None,
        model_t: Optional[BaseEstimator] = None,
        discrete_treatment: bool = True,
        honest: bool = True,
        bootstrap: bool = True,
        random_state: Optional[int] = None,
        n_jobs: int = 1,
        verbose: int = 0,
    ):
        self.n_estimators = n_estimators
        self.min_samples_leaf = min_samples_leaf
        self.max_depth = max_depth
        self.max_samples = max_samples
        self.model_y = model_y
        self.model_t = model_t
        self.discrete_treatment = discrete_treatment
        self.honest = honest
        self.bootstrap = bootstrap
        self.random_state = random_state
        self.n_jobs = n_jobs
        self.verbose = verbose
        
        # Initialize default models if not provided
        if self.model_y is None:
            self.model_y = RandomForestRegressor(
                n_estimators=100, random_state=random_state
            )
        if self.model_t is None:
            if discrete_treatment:
                self.model_t = RandomForestClassifier(
                    n_estimators=100, random_state=random_state
                )
            else:
                self.model_t = RandomForestRegressor(
                    n_estimators=100, random_state=random_state
                )
        
        # Initialize internal state
        self.fitted_ = False
        self._forest = None
        self._treatment_values = None
        self._feature_names = None
        
    def fit(
        self,
        formula: Optional[str] = None,
        data: Optional[pd.DataFrame] = None,
        Y: Optional[np.ndarray] = None,
        T: Optional[np.ndarray] = None,
        X: Optional[np.ndarray] = None,
        W: Optional[np.ndarray] = None,
    ) -> 'CausalForest':
        """
        Fit the Causal Forest model
        
        Parameters
        ----------
        formula : str, optional
            Formula specification in the form "Y ~ T | X1 + X2 + ... [| W1 + W2 + ...]"
            where Y is outcome, T is treatment, X are effect modifiers, W are controls
        data : pd.DataFrame, optional
            Data containing all variables if using formula interface
        Y : array-like, optional
            Outcome variable (n_samples,)
        T : array-like, optional  
            Treatment variable (n_samples,)
        X : array-like, optional
            Effect modifier variables (n_samples, n_features)
        W : array-like, optional
            Control variables for confounding adjustment (n_samples, n_controls)
            
        Returns
        -------
        self : CausalForest
            Fitted estimator
        """
        # Parse inputs
        if formula is not None and data is not None:
            Y, T, X, W = self._parse_formula_inputs(formula, data)
        elif Y is not None and T is not None and X is not None:
            Y, T, X, W = self._validate_array_inputs(Y, T, X, W)
        else:
            raise ValueError(
                "Must provide either (formula, data) or (Y, T, X) arguments"
            )
        
        # Validate inputs
        Y = np.asarray(Y).ravel()
        T = np.asarray(T).ravel()
        X = np.asarray(X)
        if W is not None:
            W = np.asarray(W)
        
        if X.ndim == 1:
            X = X.reshape(-1, 1)
        if W is not None and W.ndim == 1:
            W = W.reshape(-1, 1)
            
        n_samples = len(Y)
        if len(T) != n_samples or len(X) != n_samples:
            raise ValueError("Y, T, and X must have the same number of samples")
        if W is not None and len(W) != n_samples:
            raise ValueError("W must have the same number of samples as Y")
        
        # Store data info
        self.data_info = {
            'nobs': n_samples,
            'n_features': X.shape[1],
            'n_controls': W.shape[1] if W is not None else 0,
            'treatment_values': np.unique(T),
        }
        
        self._treatment_values = np.unique(T)
        self._feature_names = [f'X{i}' for i in range(X.shape[1])]
        self._X_original = X.copy()
        self._T_original = T.copy()
        self._Y_original = Y.copy()
        
        # Validate treatment
        if self.discrete_treatment:
            if len(self._treatment_values) < 2:
                raise ValueError("Need at least 2 treatment values for discrete treatment")
        
        # Step 1: Fit first stage models (Double ML approach)
        # This follows the EconML CausalForestDML implementation
        if self.verbose > 0:
            print("Fitting first stage models...")
            
        # Prepare features for first stage
        first_stage_features = X if W is None else np.hstack([X, W])
        
        # Fit outcome model
        if self.verbose > 0:
            print("  Fitting outcome model...")
        self.model_y.fit(first_stage_features, Y)
        Y_pred = cross_val_predict(
            self.model_y, first_stage_features, Y, cv=3, n_jobs=self.n_jobs
        )
        Y_residual = Y - Y_pred
        
        # Fit treatment model  
        if self.verbose > 0:
            print("  Fitting treatment model...")
        self.model_t.fit(first_stage_features, T)
        if self.discrete_treatment:
            T_pred = cross_val_predict(
                self.model_t, first_stage_features, T, cv=3, 
                method='predict_proba', n_jobs=self.n_jobs
            )
            if T_pred.ndim == 2 and T_pred.shape[1] == 2:
                # Binary case - use probability of positive class (class 1)
                T_pred = T_pred[:, 1]
                T_residual = T - T_pred
            elif T_pred.ndim == 1 or T_pred.shape[1] == 1:
                # Single class case
                T_pred = T_pred.ravel()
                T_residual = T - T_pred
            else:
                # Multi-class case - use one-hot encoding
                T_onehot = np.zeros((len(T), len(self._treatment_values)))
                for i, val in enumerate(self._treatment_values):
                    T_onehot[T == val, i] = 1
                T_residual = T_onehot - T_pred
        else:
            T_pred = cross_val_predict(
                self.model_t, first_stage_features, T, cv=3, n_jobs=self.n_jobs
            )
            T_residual = T - T_pred
        
        # Stash cross-fitted nuisance predictions for downstream inference
        # (used by :func:`forest_inference.calibration_test` / :func:`rate`).
        self._m_insample = Y_pred
        self._e_insample = T_pred if np.asarray(T_pred).ndim == 1 else T_pred[:, 0]

        # Step 2: Fit causal forest on residuals
        if self.verbose > 0:
            print("Fitting causal forest...")

        self._forest = self._fit_causal_forest(X, T_residual, Y_residual)
        
        # Mark as fitted
        self.fitted_ = True
        
        # Store results for compatibility with EconometricResults
        # For non-parametric methods, we don't have traditional parameters
        self.params = pd.Series([], dtype=float)
        self.std_errors = pd.Series([], dtype=float) 
        self.tvalues = pd.Series([], dtype=float)
        self.pvalues = np.array([])
        
        # Compute diagnostics
        ate = self.effect(X).mean()
        self.diagnostics = {
            'method': 'Causal Forest',
            'n_estimators': self.n_estimators,
            'n_features': X.shape[1],
            'average_treatment_effect': ate,
            'treatment_type': 'discrete' if self.discrete_treatment else 'continuous',
        }
        
        return self
    
    def _parse_formula_inputs(
        self, formula: str, data: pd.DataFrame
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]]:
        """Parse formula inputs to extract Y, T, X, W arrays"""
        # Parse formula: "Y ~ T | X1 + X2 + ... [| W1 + W2 + ...]"
        parts = formula.split('|')
        if len(parts) < 2:
            raise ValueError(
                "Formula must have format 'Y ~ T | X1 + X2 + ...' or "
                "'Y ~ T | X1 + X2 + ... | W1 + W2 + ...'"
            )
        
        # Parse outcome and treatment
        yt_part = parts[0].strip()
        if '~' not in yt_part:
            raise ValueError("Formula must contain '~' to separate outcome and treatment")
        
        y_name, t_name = yt_part.split('~')
        y_name = y_name.strip()
        t_name = t_name.strip()
        
        # Parse effect modifiers (X)
        x_part = parts[1].strip()
        x_names = [name.strip() for name in x_part.split('+')]
        
        # Parse controls (W) if provided
        w_names = []
        if len(parts) > 2:
            w_part = parts[2].strip()
            w_names = [name.strip() for name in w_part.split('+')]
        
        # Extract data
        try:
            Y = data[y_name].values
            T = data[t_name].values
            X = data[x_names].values
            W = data[w_names].values if w_names else None
        except KeyError as e:
            raise ValueError(f"Variable {e} not found in data")
        
        # Store feature names for later use
        self._feature_names = x_names
        
        return Y, T, X, W
    
    def _validate_array_inputs(
        self, Y: np.ndarray, T: np.ndarray, X: np.ndarray, W: Optional[np.ndarray]
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]]:
        """Validate direct array inputs"""
        Y = np.asarray(Y)
        T = np.asarray(T)
        X = np.asarray(X)
        if W is not None:
            W = np.asarray(W)
        return Y, T, X, W
    
    def _fit_causal_forest(
        self, X: np.ndarray, T_residual: np.ndarray, Y_residual: np.ndarray
    ) -> List[DecisionTreeRegressor]:
        """
        Fit the causal forest using honest estimation

        This is the core of the causal forest algorithm, implementing honest
        random forests as described in Wager & Athey (2018).
        """
        n_samples, n_features = X.shape
        trees = []

        # Set random seed for reproducibility
        rng = np.random.RandomState(self.random_state)

        for tree_idx in range(self.n_estimators):
            if self.verbose > 1 and tree_idx % 50 == 0:
                print(f"  Fitting tree {tree_idx + 1}/{self.n_estimators}")

            # Sample for this tree
            if self.bootstrap:
                n_tree_samples = int(self.max_samples * n_samples)
                tree_indices = rng.choice(n_samples, n_tree_samples, replace=True)
            else:
                n_tree_samples = int(self.max_samples * n_samples)
                tree_indices = rng.choice(n_samples, n_tree_samples, replace=False)
            
            # Honest estimation: split samples for tree building and effect estimation
            if self.honest:
                split_idx = len(tree_indices) // 2
                build_indices = tree_indices[:split_idx]
                estimate_indices = tree_indices[split_idx:]
            else:
                build_indices = tree_indices
                estimate_indices = tree_indices
            
            # Build tree structure using building sample
            # For simplicity, we use sklearn's DecisionTreeRegressor
            # but replace leaf predictions with causal effect estimates
            tree = DecisionTreeRegressor(
                max_depth=self.max_depth,
                min_samples_leaf=self.min_samples_leaf,
                random_state=rng.randint(0, 2**31),
                max_features='sqrt' if n_features > 1 else None,
            )
            
            # Fit tree on building sample
            if len(build_indices) > 0:
                # Use treatment residual as target for tree structure
                # Ensure T_residual is 1D for tree fitting
                T_build = T_residual[build_indices]
                if T_build.ndim > 1:
                    T_build = T_build.ravel()
                    
                tree.fit(X[build_indices], T_build)
                
                # Replace leaf values with honest causal effect estimates
                self._replace_leaf_values_with_causal_effects(
                    tree, X[estimate_indices], T_residual[estimate_indices], 
                    Y_residual[estimate_indices]
                )
            else:
                # Fallback if not enough samples
                T_tree = T_residual[tree_indices]
                if T_tree.ndim > 1:
                    T_tree = T_tree.ravel()
                tree.fit(X[tree_indices], T_tree)
            
            trees.append(tree)
        
        return trees
    
    def _replace_leaf_values_with_causal_effects(
        self,
        tree: DecisionTreeRegressor,
        X_estimate: np.ndarray, 
        T_residual_estimate: np.ndarray,
        Y_residual_estimate: np.ndarray,
    ):
        """
        Replace leaf values with honest causal effect estimates
        
        This implements the honest estimation procedure where leaf values
        are computed using a separate sample from tree construction.
        """
        if len(X_estimate) == 0:
            return
        
        # Get leaf assignments for estimation sample
        leaf_ids = tree.apply(X_estimate)
        
        # For each leaf, compute causal effect estimate
        for leaf_id in np.unique(leaf_ids):
            leaf_mask = leaf_ids == leaf_id
            if np.sum(leaf_mask) < 2:  # Need at least 2 samples
                continue
            
            # Get samples in this leaf
            t_leaf = T_residual_estimate[leaf_mask]
            y_leaf = Y_residual_estimate[leaf_mask]
            
            # Estimate causal effect in this leaf
            if self.discrete_treatment:
                if t_leaf.ndim == 1:
                    # Binary treatment
                    if np.var(t_leaf) > 1e-8:  # Check for variation
                        # Simple ratio estimator: Cov(Y,T) / Var(T)
                        causal_effect = np.cov(y_leaf, t_leaf)[0, 1] / np.var(t_leaf)
                    else:
                        causal_effect = 0.0
                else:
                    # Multi-class treatment - use first treatment effect
                    if np.var(t_leaf[:, 0]) > 1e-8:
                        causal_effect = np.cov(y_leaf, t_leaf[:, 0])[0, 1] / np.var(t_leaf[:, 0])
                    else:
                        causal_effect = 0.0
            else:
                # Continuous treatment
                if np.var(t_leaf) > 1e-8:
                    causal_effect = np.cov(y_leaf, t_leaf)[0, 1] / np.var(t_leaf)
                else:
                    causal_effect = 0.0
            
            # Update THIS leaf's value (leaf_id is the node index in tree.tree_)
            original_shape = tree.tree_.value[leaf_id].shape
            tree.tree_.value[leaf_id] = np.full(original_shape, causal_effect)
    
    def effect(self, X: np.ndarray) -> np.ndarray:
        """
        Estimate conditional average treatment effects
        
        Parameters
        ----------
        X : array-like, shape (n_samples, n_features)
            Effect modifier variables
            
        Returns
        -------
        effects : array-like, shape (n_samples,)
            Estimated conditional average treatment effects
        """
        if not self.fitted_:
            raise ValueError("Model must be fitted before predicting effects")
        
        X = np.asarray(X)
        if X.ndim == 1:
            X = X.reshape(-1, 1)
        
        # Average predictions across all trees
        predictions = []
        for tree in self._forest:
            pred = tree.predict(X)
            # Ensure pred is 1D for each tree
            if pred.ndim > 1:
                pred = pred.ravel()
            predictions.append(pred)
        
        # Return average effect across trees
        predictions = np.array(predictions)  # shape: (n_trees, n_samples)
        return np.mean(predictions, axis=0)  # shape: (n_samples,)
    
    def predict(self, data: Optional[pd.DataFrame] = None) -> np.ndarray:
        """
        Generate treatment effect predictions (required by BaseModel)
        
        Parameters
        ----------
        data : pd.DataFrame, optional
            Data containing effect modifier variables. If None, uses training data.
            
        Returns
        -------
        np.ndarray
            Predicted treatment effects
            
        Notes
        -----
        This method is required by the BaseModel interface. For Causal Forest,
        "predictions" are treatment effect estimates rather than outcome predictions.
        """
        if not self.fitted_:
            raise ValueError("Model must be fitted before making predictions")
        
        if data is None:
            if hasattr(self, '_X_original'):
                X = self._X_original
            else:
                raise ValueError("No data provided and no training data available")
        else:
            if isinstance(data, pd.DataFrame):
                if hasattr(self, '_feature_names'):
                    X = data[self._feature_names].values
                else:
                    # Use all numeric columns
                    X = data.select_dtypes(include=[np.number]).values
            else:
                X = np.asarray(data)
        
        return self.effect(X)
    
    def effect_interval(
        self, X: np.ndarray, alpha: float = 0.05
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Compute confidence intervals for treatment effects using bootstrap
        
        Parameters
        ---------- 
        X : array-like, shape (n_samples, n_features)
            Effect modifier variables
        alpha : float, default=0.05
            Significance level (1-alpha is confidence level)
            
        Returns
        -------
        lower : array-like, shape (n_samples,)
            Lower bounds of confidence intervals
        upper : array-like, shape (n_samples,)
            Upper bounds of confidence intervals
        """
        if not self.fitted_:
            raise ValueError("Model must be fitted before computing intervals")
        
        X = np.asarray(X)
        if X.ndim == 1:
            X = X.reshape(-1, 1)
        
        # Use bootstrap of little bags approach
        # Collect predictions from each tree (which are already bootstrap samples)
        predictions = []
        for tree in self._forest:
            pred = tree.predict(X)
            # Ensure pred is 1D
            if pred.ndim > 1:
                pred = pred.ravel()
            predictions.append(pred)
        
        # Convert to array and ensure consistent shapes
        try:
            predictions = np.array(predictions)  # shape: (n_trees, n_samples)
        except ValueError:
            # Handle inconsistent shapes by ensuring all predictions have same length
            min_len = min(len(p) for p in predictions)
            predictions = np.array([p[:min_len] for p in predictions])
        
        # Compute percentiles across trees for each sample
        lower_percentile = 100 * alpha / 2
        upper_percentile = 100 * (1 - alpha / 2)
        
        lower = np.percentile(predictions, lower_percentile, axis=0)
        upper = np.percentile(predictions, upper_percentile, axis=0)
        
        return lower, upper

    def average_treatment_effect(
        self,
        X: Optional[np.ndarray] = None,
        T: Optional[np.ndarray] = None,
        target_sample: str = "all",
        alpha: float = 0.05,
    ) -> Dict[str, float]:
        """GRF-style ATE/ATT/ATC/ATO aggregation of CATE predictions."""
        from .forest_inference import average_treatment_effect
        return average_treatment_effect(
            self, X=X, T=T, target_sample=target_sample, alpha=alpha,
        )

    def forest_diagnostics(
        self,
        X: Optional[np.ndarray] = None,
        T: Optional[np.ndarray] = None,
        propensity_bounds: Tuple[float, float] = (0.05, 0.95),
    ) -> Dict[str, object]:
        """Overlap and CATE-distribution diagnostics for this fitted forest."""
        from .forest_inference import forest_diagnostics
        return forest_diagnostics(
            self, X=X, T=T, propensity_bounds=propensity_bounds,
        )
    
    def summary(self) -> str:
        """
        Return a summary of the fitted model
        
        Returns
        -------
        summary : str
            Model summary string
        """
        if not self.fitted_:
            return "Causal Forest (not fitted)"
        
        ate = self.diagnostics.get('average_treatment_effect', 0)
        
        summary_lines = [
            "=" * 60,
            "Causal Forest Results",
            "=" * 60,
            f"Method:                   Causal Forest",
            f"Number of trees:          {self.n_estimators}",
            f"Min samples per leaf:     {self.min_samples_leaf}",
            f"Max depth:                {self.max_depth or 'None'}",
            f"Max samples per tree:     {self.max_samples}",
            f"Honest estimation:        {self.honest}",
            f"Treatment type:           {self.diagnostics.get('treatment_type', 'Unknown')}",
            "",
            f"Number of observations:   {self.data_info.get('nobs', 'Unknown')}",
            f"Number of features:       {self.data_info.get('n_features', 'Unknown')}",
            f"Number of controls:       {self.data_info.get('n_controls', 0)}",
            "",
            f"Average Treatment Effect: {ate:.6f}",
            "=" * 60,
            "",
            "Note: Use .effect(X) to estimate individual treatment effects",
            "      Use .effect_interval(X) for confidence intervals",
        ]
        
        return "\n".join(summary_lines)
    
    def __str__(self) -> str:
        """String representation"""
        if self.fitted_:
            return f"CausalForest(fitted=True, n_estimators={self.n_estimators})"
        else:
            return f"CausalForest(fitted=False, n_estimators={self.n_estimators})"
    
    def __repr__(self) -> str:
        """Detailed string representation"""
        return self.__str__()

    # ------------------------------------------------------------------ #
    #  GRF-inspired extensions
    # ------------------------------------------------------------------ #

    def variable_importance(self) -> pd.Series:
        """Permutation-based variable importance for the causal forest.

        For each feature j, shuffle its column in the effect-modifier
        matrix and measure how much the cross-validated CATE predictions
        degrade (MSE increase). Higher degradation → more important for
        treatment-effect heterogeneity.

        Returns a normalised importance score (sums to 1).
        """
        if not self.fitted_:
            raise ValueError("Model must be fitted before computing importance")
        X = self._X_original.copy()
        cate_baseline = self.effect(X)
        n, k = X.shape
        rng = np.random.default_rng(0)
        importance = np.empty(k)
        for j in range(k):
            X_perm = X.copy()
            X_perm[:, j] = rng.permutation(X_perm[:, j])
            cate_perm = self.effect(X_perm)
            importance[j] = float(np.mean((cate_baseline - cate_perm) ** 2))
        total = importance.sum()
        importance = importance / total if total > 0 else importance
        names = self._feature_names or [f"X{j}" for j in range(k)]
        return pd.Series(importance, index=names).sort_values(ascending=False)

    def best_linear_projection(
        self,
        X_test: Optional[np.ndarray] = None,
        alpha: float = 0.05,
        clip: float = 0.01,
    ) -> pd.DataFrame:
        r"""Best Linear Projection (BLP) of CATE on features (Semenova-Chernozhukov 2021).

        Constructs the augmented inverse-propensity-weighted (AIPW) doubly-robust
        score :math:`\Gamma_i` and regresses it on :math:`X_i` with HC1 standard
        errors:

        .. math::
            \Gamma_i = \hat{\tau}(X_i)
                + \frac{T_i - \hat{e}(X_i)}{\hat{e}(X_i)(1-\hat{e}(X_i))}
                  \bigl(Y_i - \hat{m}(X_i) - (T_i - \hat{e}(X_i))\hat{\tau}(X_i)\bigr).

        :math:`\Gamma_i` is unbiased for :math:`\tau(X_i)` under standard
        cross-fitting / overlap conditions; OLS of :math:`\Gamma_i` on
        :math:`(1, X_i)` recovers the population BLP coefficients with
        valid heteroscedasticity-robust inference (HC1).

        This replaces the earlier plug-in OLS of :math:`\hat{\tau}(X_i)` on
        :math:`X_i`, which produces anti-conservative SEs (the SE on a fitted
        model is not the SE on the population BLP).

        Parameters
        ----------
        X_test : array-like, optional
            Features to evaluate the BLP at; defaults to in-sample X.
        alpha : float, default=0.05
            Significance level for CIs (reported alongside coef/SE).
        clip : float, default=0.01
            Propensity clip (binary discrete treatment only) to prevent the
            inverse-propensity term from blowing up under near-violations of
            overlap. Counts of clipped units are exposed via ``self.diagnostics``.

        Returns
        -------
        pd.DataFrame
            Index ``["Intercept", *features]`` with columns
            ``[coef, se, t, p, ci_lower, ci_upper]``. HC1 SEs.

        References
        ----------
        Semenova V., Chernozhukov V. (2021).
        "Debiased Machine Learning of Conditional Average Treatment Effects
        and Other Causal Functions." *Econometrics Journal* 24(2): 264-289.
        DOI: 10.1093/ectj/utaa027.
        """
        if not self.fitted_:
            raise ValueError("Model must be fitted first")
        from scipy import stats as _stats

        X = X_test if X_test is not None else self._X_original
        X = np.asarray(X, dtype=float)
        n, k = X.shape

        # Use the in-sample DR construction (requires Y, T, m̂, ê).
        if X_test is not None and X_test.shape[0] != self._Y_original.shape[0]:
            warnings.warn(
                "best_linear_projection: X_test has a different sample size "
                "than the training data; falling back to the plug-in CATE "
                "regression with HC1 SE (no DR construction possible "
                "out-of-sample).",
                UserWarning, stacklevel=2,
            )
            cate = self.effect(X)
            gamma = cate
            T_used = None
        else:
            tau = self.effect(self._X_original)
            Y = self._Y_original.astype(float)
            T = self._T_original.astype(float)
            m_hat = np.asarray(self._m_insample, dtype=float)
            e_hat = np.asarray(self._e_insample, dtype=float)

            if self.discrete_treatment:
                e_clip = np.clip(e_hat, clip, 1.0 - clip)
                n_clipped = int(np.sum((e_hat < clip) | (e_hat > 1 - clip)))
                weight = (T - e_clip) / (e_clip * (1.0 - e_clip))
            else:
                # Continuous treatment: use partialled-out Robinson form.
                # Γ_i = τ̂(X_i) + (T_i - ê(X_i))(Y_i - m̂(X_i) - (T_i-ê)τ̂)/Var(T-ê)
                resid_T = T - e_hat
                var_T = float(np.mean(resid_T**2)) + 1e-12
                weight = resid_T / var_T
                n_clipped = 0

            mu_hat = m_hat + (T - e_hat) * tau
            gamma = tau + weight * (Y - mu_hat)

            self.diagnostics = dict(self.diagnostics or {})
            self.diagnostics["blp_n_clipped_propensities"] = n_clipped
            T_used = T

        # Regress Γ_i on (1, X_std) with HC1 SE.
        X_mean = X.mean(axis=0)
        X_sd = X.std(axis=0) + 1e-12
        X_std = (X - X_mean) / X_sd
        D = np.column_stack([np.ones(n), X_std])
        DtD = D.T @ D
        try:
            DtD_inv = np.linalg.inv(DtD)
        except np.linalg.LinAlgError:
            DtD_inv = np.linalg.pinv(DtD)
        beta = DtD_inv @ (D.T @ gamma)
        resid = gamma - D @ beta

        # HC1 White SE with (n / (n - p)) finite-sample correction.
        p = D.shape[1]
        omega = D * resid[:, None]
        meat = omega.T @ omega
        cov = DtD_inv @ meat @ DtD_inv * (n / max(n - p, 1))
        se = np.sqrt(np.diag(cov))
        tvals = beta / np.where(se > 0, se, np.nan)
        df = max(n - p, 1)
        pvals = 2.0 * (1.0 - _stats.t.cdf(np.abs(tvals), df=df))
        z = float(_stats.t.ppf(1.0 - alpha / 2.0, df=df))

        names = ["Intercept"] + (
            self._feature_names or [f"X{j}" for j in range(k)]
        )
        return pd.DataFrame({
            "coef": beta,
            "se": se,
            "t": tvals,
            "p": pvals,
            "ci_lower": beta - z * se,
            "ci_upper": beta + z * se,
        }, index=names)

    def ate(self, X: Optional[np.ndarray] = None) -> float:
        """Average Treatment Effect (mean CATE)."""
        return float(self.effect(X if X is not None else self._X_original).mean())

    def att(self, X: Optional[np.ndarray] = None,
            T: Optional[np.ndarray] = None) -> float:
        """Average Treatment Effect on the Treated."""
        if T is None:
            if hasattr(self, "_T_original"):
                T = self._T_original
            else:
                raise ValueError("Treatment vector needed for ATT")
        X = X if X is not None else self._X_original
        cate = self.effect(X)
        mask = np.asarray(T).ravel() == 1
        return float(cate[mask].mean())


def causal_forest(
    formula: Optional[str] = None,
    data: Optional[pd.DataFrame] = None,
    Y: Optional[np.ndarray] = None,
    T: Optional[np.ndarray] = None,
    X: Optional[np.ndarray] = None,
    W: Optional[np.ndarray] = None,
    n_estimators: int = 100,
    min_samples_leaf: int = 5,
    max_depth: Optional[int] = None,
    max_samples: float = 0.5,
    model_y: Optional[BaseEstimator] = None,
    model_t: Optional[BaseEstimator] = None,
    discrete_treatment: bool = True,
    random_state: Optional[int] = None,
    **kwargs
) -> CausalForest:
    """
    Convenience function to fit a Causal Forest model
    
    This function provides a simple interface to fit Causal Forest models,
    similar to the regress() function for OLS.
    
    Parameters
    ----------
    formula : str, optional
        Formula specification: "Y ~ T | X1 + X2 + ... [| W1 + W2 + ...]"
    data : pd.DataFrame, optional
        Data containing all variables if using formula interface
    Y : array-like, optional
        Outcome variable
    T : array-like, optional
        Treatment variable
    X : array-like, optional
        Effect modifier variables
    W : array-like, optional
        Control variables
    n_estimators : int, default=100
        Number of trees in the forest
    min_samples_leaf : int, default=5
        Minimum samples per leaf
    max_depth : int, optional
        Maximum tree depth
    max_samples : float, default=0.5
        Fraction of samples per tree
    model_y : estimator, optional
        First-stage outcome model
    model_t : estimator, optional
        First-stage treatment model
    discrete_treatment : bool, default=True
        Whether treatment is discrete
    random_state : int, optional
        Random seed
    **kwargs
        Additional arguments passed to CausalForest
        
    Returns
    -------
    result : CausalForest
        Fitted Causal Forest model
        
    Examples
    --------
    >>> import pandas as pd
    >>> import numpy as np
    >>> from statspai.forest import causal_forest
    >>> 
    >>> # Generate sample data
    >>> np.random.seed(42)
    >>> n = 500
    >>> X = np.random.normal(0, 1, (n, 2))
    >>> T = np.random.binomial(1, 0.5, n)
    >>> Y = X[:, 0] * T + X[:, 1] + np.random.normal(0, 0.5, n)
    >>> data = pd.DataFrame({
    ...     'outcome': Y, 'treatment': T, 'X1': X[:, 0], 'X2': X[:, 1]
    ... })
    >>> 
    >>> # Fit using formula interface
    >>> cf = causal_forest('outcome ~ treatment | X1 + X2', data=data, 
    ...                   n_estimators=50, random_state=42)
    >>> print(cf.summary())
    >>> 
    >>> # Estimate effects
    >>> effects = cf.effect(data[['X1', 'X2']])
    >>> print(f"Mean effect: {effects.mean():.3f}")
    """
    cf = CausalForest(
        n_estimators=n_estimators,
        min_samples_leaf=min_samples_leaf,
        max_depth=max_depth,
        max_samples=max_samples,
        model_y=model_y,
        model_t=model_t,
        discrete_treatment=discrete_treatment,
        random_state=random_state,
        **kwargs
    )
    
    cf.fit(formula=formula, data=data, Y=Y, T=T, X=X, W=W)
    try:
        from ..output._lineage import attach_provenance as _attach_prov
        _attach_prov(
            cf,
            function="sp.causal_forest",
            params={
                "formula": formula,
                "n_estimators": n_estimators,
                "min_samples_leaf": min_samples_leaf,
                "max_depth": max_depth,
                "model_y": type(model_y).__name__ if model_y is not None else None,
                "model_t": type(model_t).__name__ if model_t is not None else None,
                "discrete_treatment": discrete_treatment,
                "random_state": random_state,
            },
            data=data,
            overwrite=False,
        )
    except Exception:  # pragma: no cover
        pass
    return cf
