doubleml_app.py

import streamlit as st
import numpy as np
import pandas as pd
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.model_selection import TimeSeriesSplit
from sklearn.metrics import mean_squared_error
from flaml import AutoML
import doubleml as dml
import matplotlib.pyplot as plt
import logging
import time

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Page config (no page_icon)
st.set_page_config(page_title="DoubleML Causal Analysis", layout="wide")

# Custom RMSE function
def root_mean_squared_error(y_true, y_pred):
    return np.sqrt(mean_squared_error(y_true, y_pred))

# Generic temporal transformer
class TemporalFeatures(BaseEstimator, TransformerMixin):
    """Generic temporal feature engineering for time series data"""
    
    def __init__(self, date_col="Date", include_cyclical=True, include_calendar=True, 
                 include_trend=True, include_retail_events=False):
        self.date_col = date_col
        self.include_cyclical = include_cyclical
        self.include_calendar = include_calendar
        self.include_trend = include_trend
        self.include_retail_events = include_retail_events

    def fit(self, X, y=None):
        return self

    def transform(self, X):
        df = X.copy()
        df[self.date_col] = pd.to_datetime(df[self.date_col])
        df = df.sort_values(self.date_col).reset_index(drop=True)
        
        features_added = []

        # Cyclical encodings
        if self.include_cyclical:
            df["month_sin"] = np.sin(2*np.pi * df[self.date_col].dt.month/12)
            df["month_cos"] = np.cos(2*np.pi * df[self.date_col].dt.month/12)
            df["dow_sin"] = np.sin(2*np.pi * df[self.date_col].dt.weekday/7)
            df["dow_cos"] = np.cos(2*np.pi * df[self.date_col].dt.weekday/7)
            df["doy_sin"] = np.sin(2*np.pi * df[self.date_col].dt.dayofyear/365)
            df["doy_cos"] = np.cos(2*np.pi * df[self.date_col].dt.dayofyear/365)
            features_added.extend(["month_sin", "month_cos", "dow_sin", "dow_cos", "doy_sin", "doy_cos"])

        # Calendar flags
        if self.include_calendar:
            df["is_weekend"] = (df[self.date_col].dt.weekday >= 5).astype(int)
            df["is_month_start"] = (df[self.date_col].dt.day <= 7).astype(int)
            df["is_month_end"] = (df[self.date_col].dt.day >= 24).astype(int)
            df["is_quarter_start"] = df[self.date_col].dt.month.isin([1,4,7,10]).astype(int)
            df["is_quarter_end"] = df[self.date_col].dt.month.isin([3,6,9,12]).astype(int)
            features_added.extend([
                "is_weekend", "is_month_start", "is_month_end", 
                "is_quarter_start", "is_quarter_end"
            ])

        # Retail/shopping season flags (optional)
        if self.include_retail_events:
            df["is_holiday_season"] = (
                (df[self.date_col].dt.month==12) |
                ((df[self.date_col].dt.month==11)&(df[self.date_col].dt.day>=15))
            ).astype(int)
            df["is_black_friday_week"] = (
                (df[self.date_col].dt.month==11)&(df[self.date_col].dt.day>=20)
            ).astype(int)
            df["is_nordstrom_anniversary_sale"] = (
                ((df[self.date_col].dt.month==7)&(df[self.date_col].dt.day>=15)) |
                ((df[self.date_col].dt.month==8)&(df[self.date_col].dt.day<=10))
            ).astype(int)
            df["is_sephora_spring_sale"] = (
                (df[self.date_col].dt.month==4)&
                df[self.date_col].dt.day.between(1,20)
            ).astype(int)
            df["is_walmart_july_event"] = (
                (df[self.date_col].dt.month==7)&
                df[self.date_col].dt.day.between(10,20)
            ).astype(int)
            df["is_target_july_event"] = df["is_walmart_july_event"]
            df["is_amazon_prime_days_window"] = (
                (df[self.date_col].dt.month==7)&
                df[self.date_col].dt.day.between(9,18)
            ).astype(int)
            features_added.extend([
                "is_holiday_season", "is_black_friday_week", "is_nordstrom_anniversary_sale",
                "is_sephora_spring_sale", "is_walmart_july_event", "is_target_july_event",
                "is_amazon_prime_days_window"
            ])

        # Trend features
        if self.include_trend:
            df["time_trend"] = np.arange(len(df))
            df["time_trend_sq"] = df["time_trend"] ** 2
            features_added.extend(["time_trend", "time_trend_sq"])

        return df.drop(columns=[self.date_col]), features_added

def geometric_adstock(s, alpha=0.9, lags=7):
    """Apply geometric adstock transformation"""
    w = alpha ** np.arange(lags)
    return s.rolling(lags, min_periods=1).apply(
        lambda x: np.dot(x[::-1], w[:len(x)]), raw=True
    )

def apply_lag_features(df, col, lags=[1, 7, 14, 30]):
    """Apply simple lag features"""
    for lag in lags:
        df[f"{col}_lag{lag}"] = df[col].shift(lag)
    return df

# Initialize session state
if 'df' not in st.session_state:
    st.session_state.df = None
if 'results' not in st.session_state:
    st.session_state.results = None
if 'dml_plr' not in st.session_state:
    st.session_state.dml_plr = None

# Title
st.title("DoubleML Causal Analysis Tool")
st.markdown("*Flexible time series causal inference for any domain*")
st.markdown("---")

# Sidebar
with st.sidebar:
    st.header("Navigation")
    step = st.radio(
        "Select Step:",
        [
            "1. Upload Data",
            "2. Configure Analysis",
            "3. Run Analysis",
            "4. Sensitivity Analysis",
            "5. CATE Explorer",
        ],
        key="nav_step",
    )
    st.markdown("---")
    st.info("Tip: Follow the steps in order for best results.")
    
    with st.expander("About This Tool"):
        st.markdown(
            """
            This tool implements Double Machine Learning (DoubleML) 
            for causal inference on time series data.
            
            Use cases:
            - Marketing: Ad spend → Sales
            - Healthcare: Treatment → Outcomes
            - Policy: Intervention → Impact
            - Finance: Events → Returns
            - Operations: Changes → Metrics
            """
        )

# Step 1: Upload Data
if step == "1. Upload Data":
    st.header("Step 1: Upload Your Dataset")
    
    col1, col2 = st.columns([2, 1])
    
    with col1:
        uploaded_file = st.file_uploader("Choose a CSV file", type="csv")
    
    with col2:
        st.info(
            "**Requirements:**\n"
            "- Time series data\n"
            "- Date column\n"
            "- Outcome variable\n"
            "- Treatment variable"
        )
    
    if uploaded_file is not None:
        df = pd.read_csv(uploaded_file)
        st.session_state.df = df
        
        st.success(f"Dataset uploaded successfully! {len(df)} rows × {len(df.columns)} columns")
        
        tab1, tab2, tab3 = st.tabs(["Preview", "Summary Stats", "Data Quality"])
        
        with tab1:
            st.dataframe(df.head(20), use_container_width=True)
        
        with tab2:
            col1, col2 = st.columns(2)
            with col1:
                st.write("Numeric Columns:")
                st.dataframe(df.describe())
            with col2:
                st.write("Data Types:")
                st.dataframe(pd.DataFrame(df.dtypes, columns=['Data Type']))
        
        with tab3:
            missing = df.isnull().sum()
            if missing.sum() > 0:
                st.warning(
                    f"Missing values detected in {(missing > 0).sum()} columns:"
                )
                missing_df = pd.DataFrame({
                    'Column': missing[missing > 0].index,
                    'Missing Count': missing[missing > 0].values,
                    'Missing %': (missing[missing > 0].values / len(df) * 100).round(2)
                })
                st.dataframe(missing_df)
            else:
                st.success("No missing values detected")
            
            # Check for potential date columns
            potential_dates = [
                col for col in df.columns
                if 'date' in col.lower() or 'time' in col.lower()
            ]
            if potential_dates:
                st.info(
                    f"Potential date columns detected: {', '.join(potential_dates)}"
                )
        
    else:
        st.info("Please upload a CSV file to begin")

# Step 2: Configure Analysis
elif step == "2. Configure Analysis":
    st.header("Step 2: Configure Your Analysis")
    
    if st.session_state.df is None:
        st.warning("Please upload data first (Step 1)")
    else:
        import pandas as pd
        import numpy as np
        from datetime import timedelta
        
        df = st.session_state.df
        
        # Core Variables Section
        st.subheader("Core Variables")
        col1, col2, col3 = st.columns(3)
        
        with col1:
            date_col = st.selectbox("Date Column", df.columns.tolist(), key="date_col")
        
        with col2:
            outcome_col = st.selectbox(
                "Outcome Variable (Y)",
                [col for col in df.columns if col != date_col],
                key="outcome_col",
            )
        
        with col3:
            treatment_col = st.selectbox(
                "Treatment Variable (D)",
                [col for col in df.columns if col not in [date_col, outcome_col]],
                key="treatment_col",
            )
        
        # Confounders Section
        st.markdown("---")
        st.subheader("Confounding Variables")
        
        available_cols = [
            col for col in df.columns if col not in [date_col, outcome_col, treatment_col]
        ]
        confounder_cols = st.multiselect(
            "Select confounding variables (X) - variables that affect both treatment and outcome",
            available_cols,
            help="These are the observed covariates that may confound the causal relationship",
            key="confounder_cols",
        )
        
        # Treatment Transformation Section
        st.markdown("---")
        st.subheader("Treatment Transformation (Optional)")
        
        col1, col2 = st.columns([1, 2])
        
        with col1:
            treatment_transform = st.selectbox(
                "Transform treatment variable?",
                ["None", "Adstock (Geometric Decay)", "Simple Lags"],
                help="Apply transformations to capture delayed/cumulative effects",
            )
        
        with col2:
            if treatment_transform == "Adstock (Geometric Decay)":
                st.info("Adstock models carryover effects (common in marketing/advertising)")
                col2a, col2b = st.columns(2)
                with col2a:
                    alpha = st.slider(
                        "Decay Rate (α)", 0.0, 1.0, 0.9, 0.05,
                        help="How quickly effects decay (higher = slower decay)",
                    )
                with col2b:
                    lags = st.slider(
                        "Number of Lags", 1, 30, 7,
                        help="How many periods to look back",
                    )
            elif treatment_transform == "Simple Lags":
                st.info("Simple lags use past values as features")
                lag_periods = st.multiselect(
                    "Lag periods", [1, 7, 14, 30, 60, 90], default=[1, 7]
                )
        
        # Feature Engineering Section
        st.markdown("---")
        st.subheader("Feature Engineering")
        
        st.write("Select which temporal features to automatically generate:")
        
        col1, col2, col3, col4 = st.columns(4)
        
        with col1:
            include_cyclical = st.checkbox(
                "Cyclical Features", value=True,
                help="Sin/cos encodings for month, day of week, day of year (repeat yearly)",
            )
        
        with col2:
            include_calendar = st.checkbox(
                "Calendar Features", value=True,
                help="Weekend, month start/end, quarter flags",
            )
        
        with col3:
            include_trend = st.checkbox(
                "Trend Features", value=True,
                help="Linear and quadratic time trends (depends on dataset length)",
            )
        
        with col4:
            include_retail = st.checkbox(
                "Retail Events", value=False,
                help="Holiday season, Black Friday, Prime Day, etc.",
            )

        # Retail events helper
        def get_retail_events(year: int) -> pd.DataFrame:
            """Compute key US retail events for a given year."""
            thanksgiving = pd.date_range(
                f"{year}-11-01", f"{year}-11-30", freq="W-THU"
            )[3]
            black_friday = thanksgiving + timedelta(days=1)
            cyber_monday = thanksgiving + timedelta(days=4)
            prime_day_start = pd.Timestamp(f"{year}-07-15")
            prime_day_end = prime_day_start + timedelta(days=1)
            back_to_school_start = pd.Timestamp(f"{year}-08-01")
            back_to_school_end = pd.Timestamp(f"{year}-09-15")

            events = [
                {"event": "Black Friday", "date": black_friday},
                {"event": "Cyber Monday", "date": cyber_monday},
                {"event": "Prime Day Start", "date": prime_day_start},
                {"event": "Prime Day End", "date": prime_day_end},
                {"event": "Back to School Start", "date": back_to_school_start},
                {"event": "Back to School End", "date": back_to_school_end},
            ]
            return pd.DataFrame(events)

        if include_retail:
            years = sorted(pd.to_datetime(df[date_col]).dt.year.unique())
            retail_flags = []
            for y in years:
                retail_flags.append(get_retail_events(y))
            all_events = pd.concat(retail_flags)
            st.write("Retail events automatically detected:")
            st.dataframe(all_events)

        # Model Configuration Section
        st.markdown("---")
        st.subheader("Model Configuration")
        
        col1, col2, col3 = st.columns(3)
        
        with col1:
            n_folds = st.slider(
                "CV Folds", 2, 10, 5,
                help="Number of time series cross-validation folds",
            )
        
        with col2:
            time_budget = st.slider(
                "FLAML Time Budget (s)", 30, 600, 120, 30,
                help="Time budget per model (outcome & treatment)",
            )
        
        with col3:
            estimators = st.multiselect(
                "ML Estimators",
                ["lgbm", "xgboost", "histgb", "xgb_limitdepth",
                 "catboost", "rf", "extra_tree"],
                default=["lgbm", "xgboost", "histgb"],
                help="Machine learning algorithms to try",
            )
        
        # Advanced Options
        with st.expander("Advanced Options"):
            col1, col2 = st.columns(2)
            
            with col1:
                score_type = st.selectbox(
                    "DML Score Type",
                    ["partialling out", "IV-type"],
                    help="Orthogonalization approach for DML",
                )
                
                n_rep = st.slider(
                    "Number of Repetitions", 1, 10, 1,
                    help="Repeated sample splitting for stability",
                )
            
            with col2:
                dml_procedure = st.selectbox(
                    "DML Procedure",
                    ["dml1", "dml2"],
                    index=1,
                    help="DML1: no cross-fitting, DML2: with cross-fitting",
                )
                
                trim_threshold = st.slider(
                    "Propensity Score Trimming", 0.0, 0.2, 0.01, 0.01,
                    help="Trim extreme propensity scores (0 = no trimming)",
                )
        
        # Save Configuration
        st.markdown("---")
        if st.button("Confirm Configuration", type="primary", use_container_width=True):
            
            if not confounder_cols:
                st.warning("Consider adding confounding variables for more robust estimates")
            
            config = {
                'date_col': date_col,
                'outcome_col': outcome_col,
                'treatment_col': treatment_col,
                'confounder_cols': confounder_cols,
                'treatment_transform': treatment_transform,
                'n_folds': n_folds,
                'time_budget': time_budget,
                'estimators': estimators,
                'include_cyclical': include_cyclical,
                'include_calendar': include_calendar,
                'include_trend': include_trend,
                'include_retail': include_retail,
                'score_type': score_type,
                'n_rep': n_rep,
                'dml_procedure': dml_procedure,
                'trim_threshold': trim_threshold
            }
            
            if treatment_transform == "Adstock (Geometric Decay)":
                config['alpha'] = alpha
                config['lags'] = lags
            elif treatment_transform == "Simple Lags":
                config['lag_periods'] = lag_periods
            
            st.session_state.variable_config = config
            
            st.success("Configuration saved! Here's your setup:")
            
            summary_col1, summary_col2 = st.columns(2)
            with summary_col1:
                st.write("Variables:")
                st.write(f"- Outcome: `{outcome_col}`")
                st.write(f"- Treatment: `{treatment_col}`")
                st.write(f"- Confounders: {len(confounder_cols)}")
            with summary_col2:
                st.write("Configuration:")
                st.write(f"- Transform: {treatment_transform}")
                st.write(f"- CV Folds: {n_folds}")
                st.write(f"- Time Budget: {time_budget}s")
            
            st.info("Proceed to Step 3 to run the analysis")

# Step 3: Run Analysis
elif step == "3. Run Analysis":
    st.header("Step 3: Run Causal Analysis")
    
    if st.session_state.df is None:
        st.warning("Please upload data first (Step 1)")
    elif 'variable_config' not in st.session_state:
        st.warning("Please configure your analysis first (Step 2)")
    else:
        config = st.session_state.variable_config
        
        # Configuration Summary
        st.subheader("Configuration Summary")
        
        col1, col2, col3, col4 = st.columns(4)
        
        with col1:
            st.metric("Outcome", config['outcome_col'])
            st.metric("Treatment", config['treatment_col'])
        
        with col2:
            st.metric("Confounders", len(config['confounder_cols']))
            st.metric("Transform", config['treatment_transform'].split()[0])
        
        with col3:
            st.metric("CV Folds", config['n_folds'])
            st.metric("Estimators", len(config['estimators']))
        
        with col4:
            st.metric("Time Budget", f"{config['time_budget']}s")
            st.metric("DML Score", config['score_type'].split()[0])
        
        st.markdown("---")
        
        if st.button("Run Analysis", type="primary", use_container_width=True):
            with st.spinner("Running analysis... This may take several minutes."):
                try:
                    # Progress tracking
                    progress_bar = st.progress(0)
                    status_text = st.empty()
                    
                    # Step 1: Load and preprocess
                    status_text.text("Step 1/7: Loading and preprocessing data...")
                    progress_bar.progress(5)
                    
                    df = st.session_state.df.copy()
                    
                    # Handle missing values in treatment
                    if df[config['treatment_col']].isnull().any():
                        df[config['treatment_col']].fillna(method='ffill', inplace=True)
                    
                    # Step 2: Apply treatment transformation
                    status_text.text("Step 2/7: Applying treatment transformation...")
                    progress_bar.progress(15)
                    
                    treatment_var_name = config['treatment_col']
                    
                    if config['treatment_transform'] == "Adstock (Geometric Decay)":
                        df["transformed_treatment"] = geometric_adstock(
                            df[config['treatment_col']], 
                            alpha=config['alpha'], 
                            lags=config['lags']
                        )
                        treatment_var_name = "transformed_treatment"
                    elif config['treatment_transform'] == "Simple Lags":
                        df = apply_lag_features(df, config['treatment_col'], config['lag_periods'])
                        treatment_var_name = config['treatment_col']
                    else:
                        df["transformed_treatment"] = df[config['treatment_col']]
                        treatment_var_name = "transformed_treatment"
                    
                    # Drop rows with missing outcome or treatment
                    df.dropna(subset=[config['outcome_col'], treatment_var_name], inplace=True)
                    df.reset_index(drop=True, inplace=True)
                    
                    # Step 3: Feature engineering
                    status_text.text("Step 3/7: Engineering temporal features...")
                    progress_bar.progress(25)
                    
                    transformer = TemporalFeatures(
                        date_col=config['date_col'],
                        include_cyclical=config['include_cyclical'],
                        include_calendar=config['include_calendar'],
                        include_trend=config['include_trend'],
                        include_retail_events=config['include_retail']
                    )
                    
                    df_temp, temp_feats = transformer.fit_transform(df)
                    
                    # Combine all features
                    all_x = config['confounder_cols'] + temp_feats
                    
                    # Build model dataframe
                    df_model = pd.concat([
                        df_temp[temp_feats],
                        df[config['confounder_cols']],
                        df[[config['outcome_col'], treatment_var_name]]
                    ], axis=1).loc[:, lambda d: ~d.columns.duplicated()]
                    
                    # Drop any remaining NaNs
                    df_model.dropna(inplace=True)
                    
                    st.info(f"Analysis dataset: {len(df_model)} observations, {len(all_x)} features")
                    
                    # Step 4: Set up cross-validation
                    status_text.text("Step 4/7: Setting up time series cross-validation...")
                    progress_bar.progress(35)
                    
                    my_splitter = TimeSeriesSplit(n_splits=config['n_folds'])
                    
                    X_all = df_model[all_x]
                    y_outcome = df_model[config['outcome_col']]
                    d_treatment = df_model[treatment_var_name]
                    
                    # Step 5: Fit outcome model
                    status_text.text("Step 5/7: Fitting outcome model (ml_l)...")
                    progress_bar.progress(45)
                    
                    automl_l = AutoML()
                    automl_l.fit(
                        X_train=X_all, y_train=y_outcome,
                        task="regression", metric="rmse", 
                        time_budget=config['time_budget'],
                        split_type=my_splitter, 
                        estimator_list=config['estimators'],
                        verbose=0
                    )
                    
                    # Step 6: Fit treatment model
                    status_text.text("Step 6/7: Fitting treatment model (ml_m)...")
                    progress_bar.progress(60)
                    
                    automl_m = AutoML()
                    automl_m.fit(
                        X_train=X_all, y_train=d_treatment,
                        task="regression", metric="rmse", 
                        time_budget=config['time_budget'],
                        split_type=my_splitter, 
                        estimator_list=config['estimators'],
                        verbose=0
                    )
                    
                    # Step 7: Generate CV predictions
                    status_text.text(
                        "Step 7/7: Generating cross-validated predictions and computing causal effect..."
                    )
                    progress_bar.progress(75)
                    
                    cv_preds_l = np.full_like(y_outcome, fill_value=np.nan, dtype=float)
                    cv_preds_m = np.full_like(d_treatment, fill_value=np.nan, dtype=float)
                    
                    best_estimator_l = automl_l.model.estimator
                    best_estimator_m = automl_m.model.estimator
                    
                    for fold_idx, (train_idx, test_idx) in enumerate(my_splitter.split(X_all)):
                        best_estimator_l.fit(X_all.iloc[train_idx], y_outcome.iloc[train_idx])
                        cv_preds_l[test_idx] = best_estimator_l.predict(X_all.iloc[test_idx])
                        best_estimator_m.fit(X_all.iloc[train_idx], d_treatment.iloc[train_idx])
                        cv_preds_m[test_idx] = best_estimator_m.predict(X_all.iloc[test_idx])
                    
                    mask = ~np.isnan(cv_preds_l)
                    
                    cv_loss_l = root_mean_squared_error(y_outcome[mask], cv_preds_l[mask])
                    cv_loss_m = root_mean_squared_error(d_treatment[mask], cv_preds_m[mask])
                    
                    # Prepare DoubleML (use real X for benchmarking)
                    y_for_doubleml = y_outcome.iloc[mask]
                    d_for_doubleml = d_treatment.iloc[mask]
                    
                    # Align X with mask
                    X_for_doubleml = df_model.loc[y_for_doubleml.index, all_x]
                    
                    # Combine outcome, treatment, and confounders
                    df_for_doubleml = pd.concat(
                        [
                            pd.Series(y_for_doubleml, name=config['outcome_col']),
                            pd.Series(d_for_doubleml, name=treatment_var_name),
                            X_for_doubleml,
                        ],
                        axis=1,
                    )

                    dml_data = dml.DoubleMLData(
                        df_for_doubleml,
                        y_col=config['outcome_col'],
                        d_cols=treatment_var_name,
                        x_cols=all_x,
                    )

                    progress_bar.progress(90)
                    
                    ml_l_dummy = dml.utils.DMLDummyRegressor()
                    ml_m_dummy = dml.utils.DMLDummyRegressor()
                    
                    pred_dict = {
                        treatment_var_name: {
                            "ml_l": cv_preds_l[mask].reshape(-1, 1),
                            "ml_m": cv_preds_m[mask].reshape(-1, 1),
                        }
                    }
                    
                    dml_plr = dml.DoubleMLPLR(
                        dml_data,
                        ml_l=ml_l_dummy,
                        ml_m=ml_m_dummy,
                        n_folds=config['n_folds'],
                        n_rep=config['n_rep'],
                        score=config['score_type']
                    )
                    
                    dml_plr.fit(external_predictions=pred_dict)
                    
                    progress_bar.progress(100)
                    status_text.text("Analysis complete!")
                    
                    # Store for Step 4 benchmarking
                    st.session_state.results = {
                        'dml_plr': dml_plr,
                        'cv_loss_l': cv_loss_l,
                        'cv_loss_m': cv_loss_m,
                        'n_obs': dml_data.n_obs,
                        'mask_sum': mask.sum(),
                        'total_obs': len(df_model),
                        'confounder_cols': config['confounder_cols'],
                        'x_cols_for_benchmark': all_x,
                        'treatment_var_name': treatment_var_name,
                        'best_estimator_l': automl_l.best_estimator,
                        'best_estimator_m': automl_m.best_estimator,
                        'temporal_features': temp_feats
                    }
                    st.session_state.dml_plr = dml_plr
                    
                    time.sleep(0.5)
                    progress_bar.empty()
                    status_text.empty()
                    
                    st.success("Analysis completed successfully!")
                    
                    # Display Results
                    st.markdown("---")
                    st.header("Causal Effect Estimates")
                    
                    # Main metrics
                    col1, col2, col3 = st.columns(3)
                    
                    with col1:
                        st.metric(
                            "Causal Effect (θ)", f"{dml_plr.coef[0]:.6f}",
                            help="Average treatment effect on the outcome",
                        )
                        st.metric("Standard Error", f"{dml_plr.se[0]:.6f}")
                    
                    with col2:
                        ci_lower = dml_plr.confint().iloc[0, 0]
                        ci_upper = dml_plr.confint().iloc[0, 1]
                        st.metric("95% CI Lower", f"{ci_lower:.6f}")
                        st.metric("95% CI Upper", f"{ci_upper:.6f}")
                    
                    with col3:
                        p_val = dml_plr.pval[0]
                        st.metric("P-value", f"{p_val:.6f}")
                        if p_val < 0.05:
                            st.success("Statistically significant (p < 0.05)")
                        else:
                            st.warning("Not statistically significant (p ≥ 0.05)")
                    
                    # Interpretation
                    st.subheader("Interpretation")
                    effect = dml_plr.coef[0]
                    if effect > 0:
                        st.info(
                            f"A one-unit increase in **{config['treatment_col']}** "
                            f"is associated with a **{effect:.4f}** unit increase in "
                            f"**{config['outcome_col']}** (on average, all else equal)."
                        )
                    else:
                        st.info(
                            f"A one-unit increase in **{config['treatment_col']}** "
                            f"is associated with a **{abs(effect):.4f}** unit decrease in "
                            f"**{config['outcome_col']}** (on average, all else equal)."
                        )
                    
                    # Summary table
                    st.subheader("Detailed Summary")
                    summary_df = dml_plr.summary
                    st.dataframe(summary_df, use_container_width=True)
                    
                    # Model Performance
                    st.markdown("---")
                    st.subheader("Nuisance Model Performance")
                    
                    perf_col1, perf_col2, perf_col3 = st.columns(3)
                    
                    with perf_col1:
                        st.metric(
                            "Outcome Model RMSE", f"{cv_loss_l:.4f}",
                            help="Lower is better - measures prediction accuracy for outcome",
                        )
                    
                    with perf_col2:
                        st.metric(
                            "Treatment Model RMSE", f"{cv_loss_m:.4f}",
                            help="Lower is better - measures prediction accuracy for treatment",
                        )
                    
                    with perf_col3:
                        st.metric("Best Outcome Estimator", automl_l.best_estimator)
                        st.metric("Best Treatment Estimator", automl_m.best_estimator)
                    
                    # Data Processing Details
                    with st.expander("Data Processing Details"):
                        st.write(f"Original observations: {len(df_model)}")
                        st.write(f"Valid CV predictions: {mask.sum()} ({100*mask.sum()/len(df_model):.1f}%)")
                        st.write(f"Used for DoubleML: {dml_data.n_obs}")
                        st.write(f"Time series folds: {config['n_folds']}")
                        st.write(f"Total features used: {len(all_x)}")
                        st.write(f"Temporal features added: {len(temp_feats)}")
                        
                        if config['treatment_transform'] != "None":
                            st.write(
                                f"Treatment transformation: {config['treatment_transform']}"
                            )
                    
                    # Feature Importance (if available)
                    with st.expander("Feature Information"):
                        st.write("Confounding Variables:")
                        st.write(
                            ", ".join(config['confounder_cols'])
                            if config['confounder_cols']
                            else "None"
                        )
                        
                        st.write("\nTemporal Features Added:")
                        st.write(", ".join(temp_feats))
                    
                    st.markdown("---")
                    st.info(
                        "Proceed to Step 4 for sensitivity analysis to test robustness to unobserved confounding"
                    )
                    
                except Exception as e:
                    st.error(f"Error during analysis: {str(e)}")
                    st.exception(e)
                    progress_bar.empty()
                    status_text.empty()

elif step == "4. Sensitivity Analysis":
    st.header("Step 4: Sensitivity Analysis")
    
    if st.session_state.get('results') is None or st.session_state.results.get('dml_plr') is None:
        st.warning("Please run the causal analysis first (Step 3)")
        st.info(
            "The sensitivity analysis tests how robust your results are to unobserved confounding."
        )
    else:
        dml_plr = st.session_state.results['dml_plr']

        treatment_from_app = st.session_state.results.get('treatment_var_name', None)
        if getattr(dml_plr, "d_cols", None) and len(dml_plr.d_cols) > 0:
            treatment_names = list(dml_plr.d_cols)
        else:
            treatment_names = [treatment_from_app or "treatment"]
            dml_plr.d_cols = treatment_names

        try:
            import plotly.graph_objs as go
            _PlotlyFigure = go.Figure
        except Exception:
            _PlotlyFigure = tuple()

        from matplotlib.figure import Figure as _MplFigure
        from matplotlib.axes import Axes as _MplAxes
        import matplotlib.pyplot as plt

        def _render_doubleml_plot(plot_obj):
            if _PlotlyFigure and isinstance(plot_obj, _PlotlyFigure):
                st.plotly_chart(plot_obj, use_container_width=True)
                return
            if isinstance(plot_obj, _MplFigure) or hasattr(plot_obj, "savefig"):
                fig = plot_obj
                try:
                    fig.tight_layout()
                except Exception:
                    pass
                st.pyplot(fig)
                plt.close(fig)
                return
            if isinstance(plot_obj, _MplAxes) and hasattr(plot_obj, "figure"):
                fig = plot_obj.figure
                try:
                    fig.tight_layout()
                except Exception:
                    pass
                st.pyplot(fig)
                plt.close(fig)
                return
            fig = plt.gcf()
            if isinstance(fig, _MplFigure):
                try:
                    fig.tight_layout()
                except Exception:
                    pass
                st.pyplot(fig)
                plt.close(fig)
            else:
                st.info("Plot object not recognized; nothing to display.")

        st.info(
            """
            Sensitivity analysis assesses how robust your causal estimates are to potential unobserved confounding.
            
            Even after controlling for observed confounders, there may be unmeasured variables that affect both treatment 
            and outcome. This analysis shows how strong such confounding would need to be to change your conclusions.
            """
        )
        
        # Sensitivity Parameters
        st.subheader("Sensitivity Parameters")
        
        default_cf_y = st.session_state.get('cf_y', 0.05)
        default_cf_d = st.session_state.get('cf_d', 0.05)
        default_rho = st.session_state.get('rho', 1.0)
        default_level = st.session_state.get('level', 0.95)

        col1, col2 = st.columns(2)
        
        with col1:
            cf_y = st.number_input(
                "Partial R² with outcome (cf_y)",
                min_value=0.0,
                max_value=1.0,
                value=default_cf_y,
                step=0.01,
                format="%.4f",
                help="Strength of association between unobserved confounder and outcome (0-1)",
            )
            st.caption("Represents how much of the outcome variance the unobserved confounder explains")
        
        with col2:
            cf_d = st.number_input(
                "Partial R² with treatment (cf_d)",
                min_value=0.0,
                max_value=1.0,
                value=default_cf_d,
                step=0.01,
                format="%.4f",
                help="Strength of association between unobserved confounder and treatment (0-1)",
            )
            st.caption("Represents how much of the treatment variance the unobserved confounder explains")
        
        col3, col4 = st.columns(2)
        
        with col3:
            rho = st.slider(
                "Correlation (ρ)",
                min_value=-1.0,
                max_value=1.0,
                value=default_rho,
                step=0.1,
                help="Correlation between confounding effects on outcome and treatment",
            )
        
        with col4:
            level = st.slider(
                "Confidence Level",
                min_value=0.80,
                max_value=0.99,
                value=default_level,
                step=0.01,
                format="%.2f"
            )
        
        st.session_state.cf_y = cf_y
        st.session_state.cf_d = cf_d
        st.session_state.rho = rho
        st.session_state.level = level

        # Preset scenarios
        st.subheader("Preset Scenarios")
        scenario = st.selectbox(
            "Or choose a preset scenario:",
            ["Custom", "Weak Confounding", "Moderate Confounding", "Strong Confounding", "Extreme Confounding"]
        )
        if scenario == "Weak Confounding":
            cf_y, cf_d = 0.01, 0.01
        elif scenario == "Moderate Confounding":
            cf_y, cf_d = 0.05, 0.05
        elif scenario == "Strong Confounding":
            cf_y, cf_d = 0.10, 0.10
        elif scenario == "Extreme Confounding":
            cf_y, cf_d = 0.20, 0.20
        
        # Benchmarking
        st.markdown("---")
        st.subheader("Benchmarking Against Observed Variables")
        st.write(
            "Compare the sensitivity parameters to observed confounders to understand "
            "what level of unobserved confounding would be needed to overturn your results."
        )
        available_vars = st.session_state.results.get('confounder_cols', [])
        if available_vars:
            default_benchmarks = st.session_state.get(
                'benchmark_vars', available_vars[:min(3, len(available_vars))]
            )
            benchmark_vars = st.multiselect(
                "Select variables to use as benchmarks:",
                available_vars,
                default=default_benchmarks,
                help="These variables will be used to calibrate the sensitivity analysis",
            )
            st.session_state.benchmark_vars = benchmark_vars
        else:
            st.warning("No confounding variables available for benchmarking")
            benchmark_vars = []
        
        # Run Analysis Button
        st.markdown("---")
        if st.button("Run Sensitivity Analysis", type="primary", use_container_width=True):
            with st.spinner("Running sensitivity analysis..."):
                try:
                    dml_plr.sensitivity_analysis(cf_y=cf_y, cf_d=cf_d, rho=rho, level=level)
                    st.success("Sensitivity analysis complete!")
                    
                    # Summary
                    st.markdown("---")
                    st.subheader("Sensitivity Summary")

                    params = getattr(dml_plr, "sensitivity_params", None)
                    if params is not None:
                        try:
                            tname = dml_plr.d_cols[0] if hasattr(dml_plr, "d_cols") else "treatment"
                        except Exception:
                            tname = "treatment"

                        row = {
                            "theta lower": float(params["theta"]["lower"][0]),
                            "theta": float(getattr(dml_plr, "coef", [np.nan])[0]),
                            "theta upper": float(params["theta"]["upper"][0]),
                            "CI lower": float(params["ci"]["lower"][0]),
                            "CI upper": float(params["ci"]["upper"][0]),
                            "RV (%)": float(params.get("rv", [np.nan])[0]) * 100
                            if params.get("rv") is not None
                            else np.nan,
                            "RVa (%)": float(params.get("rva", [np.nan])[0]) * 100
                            if params.get("rva") is not None
                            else np.nan,
                        }
                        summary_df = pd.DataFrame([row], index=[tname])
                        st.dataframe(summary_df, use_container_width=True)
                    else:
                        summary = getattr(dml_plr, "sensitivity_summary", None)
                        if isinstance(summary, pd.DataFrame):
                            st.dataframe(summary, use_container_width=True)
                        elif summary is not None:
                            st.markdown(f"```text\n{summary}\n```")
                        else:
                            st.info("No sensitivity summary was returned by DoubleML.")

                    # Plots
                    st.markdown("---")
                    st.subheader("Sensitivity Plots")
                    tab1, tab2 = st.tabs(["Effect Bounds (θ)", "Confidence Interval Bounds"])

                    st.caption(f"Treatment: `{treatment_names[0]}`")

                    with tab1:
                        st.write("Shows how the point estimate changes with different levels of confounding:")
                        plt.close('all')
                        obj_theta = dml_plr.sensitivity_plot(value='theta')
                        _render_doubleml_plot(obj_theta)
                        st.caption(
                            "The plot shows the estimated causal effect across different values of confounding strength. "
                            "The shaded region represents the sensitivity bounds."
                        )
                    
                    with tab2:
                        st.write(
                            f"Shows how the {level*100:.0f}% confidence interval changes with different levels of confounding:"
                        )
                        plt.close('all')
                        obj_ci = dml_plr.sensitivity_plot(value='ci', level=level)
                        _render_doubleml_plot(obj_ci)
                        st.caption(
                            "The plot shows the confidence interval bounds across different values of confounding strength. "
                            "The line where the lower bound crosses zero indicates the confounding needed to make the effect insignificant."
                        )

                    # Benchmarking Results
                    if benchmark_vars:
                        st.markdown("---")
                        st.subheader("Benchmarking Results")
                        st.write(
                            "These results compare the sensitivity parameters to observed confounders, "
                            "helping you understand what strength of unobserved confounding would be needed."
                        )
                        for var in benchmark_vars:
                            with st.expander(f"Benchmark: {var}"):
                                try:
                                    bench_result = dml_plr.sensitivity_benchmark(
                                        benchmarking_set=[var]
                                    )
                                    if hasattr(bench_result, 'to_frame'):
                                        st.dataframe(
                                            bench_result.to_frame().T,
                                            use_container_width=True,
                                        )
                                    elif isinstance(bench_result, dict):
                                        st.write("Benchmark Parameters:")
                                        for key, value in bench_result.items():
                                            try:
                                                st.write(f"- {key}: {float(value):.4f}")
                                            except Exception:
                                                st.write(f"- {key}: {value}")
                                    else:
                                        st.write(bench_result)
                                    st.info(
                                        f"This shows the confounding strength of {var}. "
                                        "If unobserved confounding is similar in strength to this variable, "
                                        "these would be the sensitivity parameters (cf_y and cf_d)."
                                    )
                                except Exception as e:
                                    st.error(f"Error benchmarking {var}: {str(e)}")
                    
                    # Additional Insights
                    st.markdown("---")
                    st.subheader("Key Insights")
                    with st.expander("How to interpret these results"):
                        st.markdown(
                            """
                            Sensitivity Analysis Guidelines:
                            
                            1. Small cf_y and cf_d values (< 0.05): results are robust to weak unobserved confounding  
                            2. Moderate values (0.05 - 0.15): results require careful interpretation  
                            3. Large values (> 0.15): results are sensitive to unobserved confounding  
                            
                            Benchmarking helps you answer:
                            - How strong would an unobserved confounder need to be compared to my observed confounders?
                            - Is it plausible that such a confounder exists in my context?
                            
                            Best practices:
                            - Compare cf_y and cf_d to R² values from your observed confounders  
                            - Consider domain knowledge about potential unobserved confounders  
                            - Report sensitivity results alongside main estimates  
                            - If results are sensitive, consider collecting additional covariates  
                            """
                        )
                    
                    with st.expander("Learn More"):
                        st.markdown(
                            """
                            Resources:
                            - DoubleML Documentation: https://docs.doubleml.org/  
                            - Chernozhukov et al. (2018): "Double/debiased machine learning for treatment and structural parameters"  
                            - Cinelli & Hazlett (2020): "Making sense of sensitivity: Extending omitted variable bias"  
                            
                            If you use this tool in research, please cite the DoubleML package and relevant methodology papers.
                            """
                        )
                    
                except Exception as e:
                    st.error(f"Error during sensitivity analysis: {str(e)}")
                    st.exception(e)

elif step == "5. CATE Explorer":
    st.header("Step 5: CATE (Heterogeneous Effects) Explorer")

    if st.session_state.get('results') is None or st.session_state.results.get('dml_plr') is None:
        st.warning("Please run Step 3 (analysis) first.")
        st.stop()

    dml_plr = st.session_state.results['dml_plr']
    config = st.session_state.get('variable_config', {})

    try:
        data_used = dml_plr._dml_data.data.copy()
    except Exception:
        st.error("Could not access DoubleML's data. Re-run Step 3 so we can pull the same rows.")
        st.stop()

    st.info("Select 1 or 2 variables to explore how the treatment effect varies (CATE).")

    var_options = []
    if 'treatment_col' in config:
        var_options.append(config['treatment_col'])
    var_options += config.get('confounder_cols', [])
    var_options += st.session_state.results.get('temporal_features', [])
    var_options = [v for v in var_options if v in data_used.columns]

    import numpy as np
    import pandas as pd
    import patsy
    import matplotlib.pyplot as plt

    dim = st.radio("Basis dimension", ["1D", "2D"], horizontal=True)
    range_policy = st.radio(
        "X-range policy",
        ["5–95% quantiles (safer)", "Full observed range"],
        index=0,
        help="Quantiles avoid extrapolation in sparse tails; Full uses min–max of observed values.",
    )

    if dim == "1D":
        x_var = st.selectbox("Choose variable for the CATE axis (X)", var_options)

        # Sample size diagnostics
        st.markdown("---")
        st.subheader("Sample Size Diagnostics")

        n_obs = len(data_used)
        recommended_max_df = max(3, n_obs // 40)
        safe_df = max(3, n_obs // 50)

        col1, col2 = st.columns(2)
        with col1:
            st.metric("Total Observations", n_obs)
        with col2:
            st.metric(
                "Safe maximum df",
                safe_df,
                help="Best practice: stay at or below this df for stable estimates with your sample size.",
            )

        min_df = 3
        max_safe_df = min(5, safe_df)

        df_spline = st.slider(
            "Spline df (flexibility)",
            min_value=min_df,
            max_value=max(5, max_safe_df),
            value=3,
            help=f"With {n_obs} observations, df>4 risks overfitting.",
        )

        obs_per_df = n_obs / df_spline
        st.metric("Observations per spline parameter", f"{obs_per_df:.0f}")

        if df_spline > safe_df:
            st.error(f"df={df_spline} is too high for {n_obs} observations. Reduce to {safe_df} or lower.")
        elif df_spline > 4:
            st.warning(
                f"df={df_spline} with {n_obs} observations: moderate overfitting risk. "
                "Consider df=3–4."
            )
        elif obs_per_df < 30:
            st.error(f"Only {obs_per_df:.0f} observations per parameter: severe overfitting risk.")
        elif obs_per_df < 50:
            st.warning(
                f"Only {obs_per_df:.0f} observations per parameter. Results may be unstable."
            )
        else:
            st.success(f"{obs_per_df:.0f} observations per parameter looks reasonable.")

        # Distribution check on X (raw)
        x_temp = pd.to_numeric(data_used[x_var], errors="coerce").dropna()
        with st.expander("View Data Distribution"):
            fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))

            lo_temp, hi_temp = np.nanquantile(x_temp, [0.05, 0.95])
            ax1.hist(x_temp, bins=30, edgecolor='black', alpha=0.7)
            ax1.axvline(lo_temp, color='red', linestyle='--', label='5th percentile', linewidth=2)
            ax1.axvline(hi_temp, color='red', linestyle='--', label='95th percentile', linewidth=2)
            ax1.set_xlabel(x_var)
            ax1.set_ylabel("Count")
            ax1.set_title("Data Distribution")
            ax1.legend()
            ax1.grid(True, alpha=0.3)

            bins_check = pd.cut(x_temp, bins=10)
            bin_counts = bins_check.value_counts().sort_index()
            ax2.bar(range(len(bin_counts)), bin_counts.values, edgecolor='black')
            ax2.axhline(n_obs/20, color='red', linestyle='--', label='5% of data', linewidth=2)
            ax2.set_xlabel("Decile")
            ax2.set_ylabel("Count")
            ax2.set_title("Observations per Decile")
            ax2.legend()
            ax2.grid(True, alpha=0.3)

            plt.tight_layout()
            st.pyplot(fig)
            plt.close()

            min_bin_count = bin_counts.min()
            if min_bin_count < n_obs * 0.03:
                st.warning(
                    f"Sparse region detected (min bin count: {min_bin_count}). "
                    "Confidence intervals will widen at extremes."
                )

        st.warning(
            f"You have {n_obs} observations. Good for broad trends; marginal for complex shapes. "
            "Recommendation: interpret as directional; prefer df=3–4."
        )

    else:
        c1, c2 = st.columns(2)
        with c1:
            x0_var = st.selectbox("X-axis variable", var_options, key="cate_x0")
        with c2:
            x1_var = st.selectbox(
                "Y-axis variable",
                [v for v in var_options if v != st.session_state.get('cate_x0')],
                key="cate_x1",
            )

        n_obs = len(data_used)
        st.warning(
            f"2D CATE with {n_obs} observations needs substantially more data than 1D. "
            "Use df=3–4 max per axis, expect wide confidence intervals, and treat as exploratory."
        )

        df_spline = st.slider(
            "Spline df per axis (flexibility)",
            3, min(5, max(3, n_obs // 100)), 3
        )
        grid_size = st.slider(
            "Grid density (per axis)", 30, 120, 60,
            help="Higher = smoother surface, slower compute",
        )

    n_boot = st.slider(
        "Bootstrap reps for CIs", 200, 3000, 2000, 200,
        help="More reps → tighter uniform CIs (slower). Recommended: 2000+",
    )

    st.markdown("---")
    run = st.button("Compute CATE", type="primary", use_container_width=True)

    if run:
        with st.spinner("Estimating CATE..."):

            if dim == "1D":
                x = pd.to_numeric(data_used[x_var], errors="coerce").to_numpy()

                design = patsy.dmatrix(
                    f"bs(x, df={df_spline}, degree=3)",
                    {"x": x},
                    return_type="dataframe",
                )

                cate = dml_plr.cate(design)

                if range_policy.startswith("5–95%"):
                    lo, hi = np.nanquantile(x, [0.05, 0.95])
                else:
                    lo, hi = np.nanmin(x), np.nanmax(x)
                
                x_grid = np.linspace(lo, hi, 200)
                st.caption(f"Plotting range for {x_var}: [{lo:.4g}, {hi:.4g}]")
                design_grid = patsy.build_design_matrices(
                    [design.design_info], {"x": x_grid}
                )[0]
                design_grid = pd.DataFrame(design_grid, columns=design.columns)

                df_cate = cate.confint(
                    design_grid, level=0.95, joint=True, n_rep_boot=n_boot
                )

                fig, ax = plt.subplots(figsize=(8, 5))
                ax.plot(x_grid, df_cate["effect"], label="Estimated Effect")
                ax.fill_between(
                    x_grid,
                    df_cate["2.5 %"],
                    df_cate["97.5 %"],
                    alpha=0.25,
                    label="Confidence Interval",
                )
                ax.set_xlabel(x_var)
                ax.set_ylabel("Effect and 95%-CI")
                ax.set_title("CATE")
                ax.legend(loc="best")
                st.pyplot(fig)

                out = pd.DataFrame({
                    x_var: x_grid,
                    "effect": df_cate["effect"],
                    "ci_lo": df_cate["2.5 %"],
                    "ci_hi": df_cate["97.5 %"],
                })
                st.download_button(
                    "Download CATE (CSV)",
                    data=out.to_csv(index=False).encode(),
                    file_name=f"cate_1d_{x_var}.csv",
                    mime="text/csv",
                )

            else:
                x0 = pd.to_numeric(data_used[x0_var], errors="coerce").to_numpy()
                x1 = pd.to_numeric(data_used[x1_var], errors="coerce").to_numpy()

                design = patsy.dmatrix(
                    f"te(bs(x0, df={df_spline}, degree=3), "
                    f"bs(x1, df={df_spline}, degree=3))",
                    {"x0": x0, "x1": x1},
                    return_type="dataframe",
                )
                cate = dml_plr.cate(design)

                lo0, hi0 = np.nanquantile(x0, [0.05, 0.95])
                lo1, hi1 = np.nanquantile(x1, [0.05, 0.95])
                g0 = np.linspace(lo0, hi0, grid_size)
                g1 = np.linspace(lo1, hi1, grid_size)
                X0, X1 = np.meshgrid(g0, g1)
                new = {"x0": X0.ravel(), "x1": X1.ravel()}

                design_grid = patsy.build_design_matrices(
                    [design.design_info], new
                )[0]
                design_grid = pd.DataFrame(design_grid, columns=design.columns)

                df_cate = cate.confint(
                    design_grid, level=0.95, joint=True, n_rep_boot=n_boot
                )

                import plotly.graph_objects as go
                eff = np.asarray(df_cate["effect"]).reshape(X0.shape)
                lo_vals = np.asarray(df_cate["2.5 %"]).reshape(X0.shape)
                hi_vals = np.asarray(df_cate["97.5 %"]).reshape(X0.shape)

                fig = go.Figure(
                    data=[
                        go.Surface(
                            x=X0,
                            y=X1,
                            z=eff,
                            name="Effect",
                            colorscale="RdBu",
                            reversescale=True,
                            showscale=True,
                        ),
                        go.Surface(
                            x=X0,
                            y=X1,
                            z=hi_vals,
                            name="Upper 95% CI",
                            showscale=False,
                            opacity=0.35,
                            colorscale="RdBu",
                            reversescale=True,
                        ),
                        go.Surface(
                            x=X0,
                            y=X1,
                            z=lo_vals,
                            name="Lower 95% CI",
                            showscale=False,
                            opacity=0.35,
                            colorscale="RdBu",
                            reversescale=True,
                        ),
                    ]
                )
                fig.update_traces(
                    contours_z=dict(show=True, usecolormap=True, project_z=True)
                )
                fig.update_layout(
                    title="2D CATE surface",
                    scene=dict(
                        xaxis_title=x0_var,
                        yaxis_title=x1_var,
                        zaxis_title="Effect",
                    ),
                    width=950,
                    margin=dict(r=20, b=10, l=10, t=40),
                )
                st.plotly_chart(fig, use_container_width=True)

                out = pd.DataFrame({
                    x0_var: X0.ravel(),
                    x1_var: X1.ravel(),
                    "effect": df_cate["effect"],
                    "ci_lo": df_cate["2.5 %"],
                    "ci_hi": df_cate["97.5 %"],
                })
                st.download_button(
                    "Download CATE Surface (CSV)",
                    data=out.to_csv(index=False).encode(),
                    file_name=f"cate_2d_{x0_var}_{x1_var}.csv",
                    mime="text/csv",
                )

    st.markdown("---")
    with st.expander("Notes"):
        st.markdown(
            """
            - Raw feature scales (no standardization) are used to match DoubleML tutorials.  
            - The prediction grid stays within the 5–95% quantiles to avoid extrapolation artifacts.  
            - Use df=3–4 for roughly 400–500 rows; higher df risks overfitting.  
            - Increase bootstrap reps for tighter uniform confidence intervals (slower). Recommended: 2000+.  
            """
        )

# Footer
st.markdown("---")
st.markdown(
    """
    <div style='text-align: center; color: #666;'>
        <p><strong>DoubleML Causal Analysis Tool</strong></p>
        <p>Built by Growth by Science | Powered by DoubleML &amp; FLAML</p>
    </div>
    """,
    unsafe_allow_html=True,
)