"""
Bayesian Structural Time Series / CausalImpact.

Implements the CausalImpact methodology of Brodersen et al. (2015) using
a structural time series model with a Kalman filter/smoother.  Two public
interfaces are provided:

* ``causal_impact()`` — Google CausalImpact-style (wide DataFrame, period
  boundaries).
* ``bsts_synth()`` — SCM-style panel-data interface matching other StatsPAI
  synth modules.

Model
-----
**Local-level + regression**::

    y_t = mu_t + X_t @ beta + eps_t,   eps_t ~ N(0, sigma_obs^2)
    mu_t = mu_{t-1} + eta_t,           eta_t ~ N(0, sigma_level^2)

**Local-linear-trend** adds a stochastic slope::

    mu_t = mu_{t-1} + delta_{t-1} + eta_t
    delta_t = delta_{t-1} + zeta_t,    zeta_t ~ N(0, sigma_slope^2)

Parameters are estimated by maximum likelihood (via ``scipy.optimize``)
on the pre-intervention period.  Counterfactual predictions in the
post-period are generated by running the Kalman filter forward with the
learned parameters, then posterior simulation provides full uncertainty
quantification.

References
----------
Brodersen, K.H., Gallusser, F., Koehler, J., Remy, N. and Scott, S.L.
(2015).  "Inferring causal impact using Bayesian structural time series
models."  *Annals of Applied Statistics*, 9(1), 247-274. [@brodersen2015inferring]
"""

from __future__ import annotations

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

import numpy as np
import pandas as pd
from scipy import optimize, stats

from ..core.results import CausalResult

# ====================================================================== #
#  Kalman filter / smoother for local-level (+ optional linear-trend)
# ====================================================================== #


def _kalman_filter(
    y: np.ndarray,
    X: np.ndarray,
    beta: np.ndarray,
    sigma_obs: float,
    sigma_level: float,
    sigma_slope: float = 0.0,
    use_trend: bool = False,
) -> Dict[str, Any]:
    """Forward pass of the Kalman filter.

    State vector:
        local-level:        [mu_t]
        local-linear-trend: [mu_t, delta_t]

    Parameters
    ----------
    y : ndarray, shape (T,)
        Observed outcome series.  ``np.nan`` entries are treated as missing.
    X : ndarray, shape (T, K)
        Regressor matrix (control-unit series).
    beta : ndarray, shape (K,)
        Regression coefficients.
    sigma_obs : float
        Observation noise standard deviation.
    sigma_level : float
        Level innovation standard deviation.
    sigma_slope : float
        Slope innovation standard deviation (local-linear-trend only).
    use_trend : bool
        Whether to include the stochastic slope component.

    Returns
    -------
    dict with keys:
        filtered_state, filtered_cov, predicted_state, predicted_cov,
        log_likelihood, prediction_errors, prediction_error_vars
    """
    T = len(y)
    dim = 2 if use_trend else 1

    # Transition matrix
    F = np.eye(dim)
    if use_trend:
        F[0, 1] = 1.0  # mu_{t} = mu_{t-1} + delta_{t-1} + eta

    # State noise covariance
    Q = np.zeros((dim, dim))
    Q[0, 0] = sigma_level ** 2
    if use_trend:
        Q[1, 1] = sigma_slope ** 2

    # Observation matrix (maps state to observation, excluding regression)
    H = np.zeros((1, dim))
    H[0, 0] = 1.0

    R = np.array([[sigma_obs ** 2]])  # observation noise variance

    # Initialisation — diffuse prior
    state = np.zeros(dim)
    if T > 0 and not np.isnan(y[0]):
        state[0] = y[0] - X[0] @ beta
    cov = np.eye(dim) * 1e6

    filtered_state = np.zeros((T, dim))
    filtered_cov = np.zeros((T, dim, dim))
    predicted_state = np.zeros((T, dim))
    predicted_cov = np.zeros((T, dim, dim))
    pred_errors = np.zeros(T)
    pred_error_vars = np.zeros(T)
    log_lik = 0.0

    for t in range(T):
        # --- Predict ---
        state_pred = F @ state
        cov_pred = F @ cov @ F.T + Q

        predicted_state[t] = state_pred
        predicted_cov[t] = cov_pred

        # Observation prediction
        y_pred = (H @ state_pred)[0] + X[t] @ beta
        S = (H @ cov_pred @ H.T + R)[0, 0]  # innovation variance
        # Guard: S must be strictly positive. Can collapse to 0 when
        # sigma_obs == sigma_level == 0 (degenerate MLE on constant Y).
        S = max(S, 1e-12)

        pred_errors[t] = y[t] - y_pred if not np.isnan(y[t]) else 0.0
        pred_error_vars[t] = S

        if np.isnan(y[t]):
            # Missing observation — no update
            filtered_state[t] = state_pred
            filtered_cov[t] = cov_pred
        else:
            # --- Update (standard Kalman gain) ---
            K = (cov_pred @ H.T) / S  # shape (dim, 1)
            innovation = y[t] - y_pred
            state = state_pred + (K * innovation).ravel()
            cov = (np.eye(dim) - K @ H) @ cov_pred

            filtered_state[t] = state
            filtered_cov[t] = cov

            # Log-likelihood contribution
            log_lik += -0.5 * (np.log(2 * np.pi * S) + innovation ** 2 / S)

        # Carry forward for next iteration
        state = filtered_state[t]
        cov = filtered_cov[t]

    return {
        "filtered_state": filtered_state,
        "filtered_cov": filtered_cov,
        "predicted_state": predicted_state,
        "predicted_cov": predicted_cov,
        "log_likelihood": log_lik,
        "prediction_errors": pred_errors,
        "prediction_error_vars": pred_error_vars,
    }


def _kalman_smoother(kf: Dict[str, Any], use_trend: bool = False) -> Dict[str, Any]:
    """Rauch-Tung-Striebel backward smoother.

    Parameters
    ----------
    kf : dict
        Output of ``_kalman_filter``.
    use_trend : bool
        Must match the setting used in the filter.

    Returns
    -------
    dict with keys: smoothed_state, smoothed_cov
    """
    dim = 2 if use_trend else 1
    F = np.eye(dim)
    if use_trend:
        F[0, 1] = 1.0

    filt_s = kf["filtered_state"]
    filt_c = kf["filtered_cov"]
    pred_s = kf["predicted_state"]
    pred_c = kf["predicted_cov"]
    T = len(filt_s)

    sm_state = np.zeros_like(filt_s)
    sm_cov = np.zeros_like(filt_c)
    sm_state[-1] = filt_s[-1]
    sm_cov[-1] = filt_c[-1]

    for t in range(T - 2, -1, -1):
        # Smoother gain  G = filt_c[t] @ F.T @ P_pred^{-1}
        # via solve:  P_pred.T @ G.T = (filt_c[t] @ F.T).T
        P_pred = pred_c[t + 1] + np.eye(dim) * 1e-10
        try:
            G = np.linalg.solve(P_pred.T, (filt_c[t] @ F.T).T).T
        except np.linalg.LinAlgError:
            # Near-singular during diffuse init — fall back to pseudoinverse
            G = filt_c[t] @ F.T @ np.linalg.pinv(P_pred)
        sm_state[t] = filt_s[t] + G @ (sm_state[t + 1] - pred_s[t + 1])
        sm_cov[t] = filt_c[t] + G @ (sm_cov[t + 1] - pred_c[t + 1]) @ G.T

    return {"smoothed_state": sm_state, "smoothed_cov": sm_cov}


# ====================================================================== #
#  Parameter estimation via MLE
# ====================================================================== #


def _estimate_beta_ridge(
    y_pre: np.ndarray,
    X_pre: np.ndarray,
    n_lambdas: int = 20,
) -> np.ndarray:
    """Ridge regression for regressor coefficients, with GCV lambda selection.

    This implements a simplified spike-and-slab prior: ridge regularisation
    with the penalty chosen by generalised cross-validation.

    Parameters
    ----------
    y_pre : ndarray, shape (T0,)
        Pre-period outcome.
    X_pre : ndarray, shape (T0, K)
        Pre-period regressor matrix.
    n_lambdas : int
        Number of lambda candidates on the log-grid.

    Returns
    -------
    beta : ndarray, shape (K,)
    """
    T0, K = X_pre.shape
    if K == 0:
        return np.array([])

    # De-mean for numerical stability
    y_mean = y_pre.mean()
    X_mean = X_pre.mean(axis=0)
    y_c = y_pre - y_mean
    X_c = X_pre - X_mean

    XtX = X_c.T @ X_c
    Xty = X_c.T @ y_c

    # Lambda grid
    lam_max = max(np.abs(Xty).max(), 1e-4)
    lambdas = np.logspace(np.log10(lam_max * 1e-4), np.log10(lam_max * 10), n_lambdas)

    best_gcv = np.inf
    best_beta = np.zeros(K)

    for lam in lambdas:
        A = XtX + lam * np.eye(K)
        try:
            beta_c = np.linalg.solve(A, Xty)
        except np.linalg.LinAlgError:
            continue

        # GCV score
        H = X_c @ np.linalg.solve(A, X_c.T)
        resid = y_c - X_c @ beta_c
        denom = (1.0 - np.diag(H).mean()) ** 2
        if denom < 1e-12:
            continue
        gcv = float(np.mean(resid ** 2) / denom)

        if gcv < best_gcv:
            best_gcv = gcv
            best_beta = beta_c

    # Convert back
    intercept = y_mean - X_mean @ best_beta  # noqa: F841 — absorbed into level
    return best_beta


def _neg_log_likelihood(
    params: np.ndarray,
    y_pre: np.ndarray,
    X_pre: np.ndarray,
    beta: np.ndarray,
    use_trend: bool,
) -> float:
    """Negative log-likelihood for hyperparameter optimisation.

    Parameters
    ----------
    params : ndarray
        Log-transformed [sigma_obs, sigma_level (, sigma_slope)].
    y_pre, X_pre, beta : arrays
        Pre-period data and regression coefficients.
    use_trend : bool
        Whether the model includes a stochastic slope.
    """
    sigma_obs = np.exp(params[0])
    sigma_level = np.exp(params[1])
    sigma_slope = np.exp(params[2]) if use_trend else 0.0

    kf = _kalman_filter(
        y_pre, X_pre, beta,
        sigma_obs=sigma_obs,
        sigma_level=sigma_level,
        sigma_slope=sigma_slope,
        use_trend=use_trend,
    )
    ll = kf["log_likelihood"]
    if not np.isfinite(ll):
        return 1e12
    return -ll


def _fit_model(
    y_pre: np.ndarray,
    X_pre: np.ndarray,
    use_trend: bool = False,
) -> Dict[str, Any]:
    """Fit BSTS model on pre-period data.

    Steps:
    1. Estimate beta via ridge regression (GCV).
    2. Estimate sigma_obs, sigma_level (, sigma_slope) via MLE.
    3. Run smoother on pre-period.

    Returns
    -------
    dict with fitted parameters and smoother output.
    """
    beta = _estimate_beta_ridge(y_pre, X_pre)

    # Residual variance as initial guess
    resid = y_pre - X_pre @ beta if len(beta) > 0 else y_pre.copy()
    resid_std = max(float(np.std(resid)), 1e-6)

    # MLE for noise variances
    n_hyper = 3 if use_trend else 2
    x0 = np.log(np.full(n_hyper, resid_std * 0.5))

    result = optimize.minimize(
        _neg_log_likelihood,
        x0,
        args=(y_pre, X_pre, beta, use_trend),
        method="L-BFGS-B",
        options={"maxiter": 500, "ftol": 1e-8},
    )

    sigma_obs = np.exp(result.x[0])
    sigma_level = np.exp(result.x[1])
    sigma_slope = np.exp(result.x[2]) if use_trend else 0.0

    # Final filter + smoother on pre-period
    kf = _kalman_filter(
        y_pre, X_pre, beta,
        sigma_obs=sigma_obs,
        sigma_level=sigma_level,
        sigma_slope=sigma_slope,
        use_trend=use_trend,
    )
    sm = _kalman_smoother(kf, use_trend=use_trend)

    return {
        "beta": beta,
        "sigma_obs": sigma_obs,
        "sigma_level": sigma_level,
        "sigma_slope": sigma_slope,
        "use_trend": use_trend,
        "kf_pre": kf,
        "sm_pre": sm,
        "mle_result": result,
    }


# ====================================================================== #
#  Posterior simulation for counterfactual
# ====================================================================== #


def _simulate_counterfactual(
    model: Dict[str, Any],
    X_post: np.ndarray,
    n_simulations: int = 1000,
    rng: np.random.Generator | None = None,
) -> np.ndarray:
    """Draw from the posterior predictive distribution of the counterfactual.

    Starting from the last filtered state of the pre-period, simulate the
    state forward through the post-period using the learned transition
    dynamics and add observation noise.

    Parameters
    ----------
    model : dict
        Output of ``_fit_model``.
    X_post : ndarray, shape (T1, K)
        Post-period regressor matrix.
    n_simulations : int
        Number of posterior draws.
    rng : numpy Generator, optional

    Returns
    -------
    draws : ndarray, shape (n_simulations, T1)
        Posterior counterfactual draws.
    """
    if rng is None:
        rng = np.random.default_rng()

    T1 = X_post.shape[0]
    beta = model["beta"]
    sigma_obs = model["sigma_obs"]
    sigma_level = model["sigma_level"]
    sigma_slope = model["sigma_slope"]
    use_trend = model["use_trend"]

    dim = 2 if use_trend else 1
    F = np.eye(dim)
    if use_trend:
        F[0, 1] = 1.0

    Q = np.zeros((dim, dim))
    Q[0, 0] = sigma_level ** 2
    if use_trend:
        Q[1, 1] = sigma_slope ** 2

    # Last filtered state and covariance from pre-period
    kf_pre = model["kf_pre"]
    state_mean = kf_pre["filtered_state"][-1].copy()
    state_cov = kf_pre["filtered_cov"][-1].copy()

    draws = np.zeros((n_simulations, T1))

    for s in range(n_simulations):
        # Draw initial state from filtered posterior
        state = rng.multivariate_normal(state_mean, state_cov + np.eye(dim) * 1e-10)

        for t in range(T1):
            # State transition with noise
            state_noise = rng.multivariate_normal(np.zeros(dim), Q)
            state = F @ state + state_noise

            # Observation
            reg_component = X_post[t] @ beta if len(beta) > 0 else 0.0
            obs_noise = rng.normal(0.0, sigma_obs)
            draws[s, t] = state[0] + reg_component + obs_noise

    return draws


# ====================================================================== #
#  Public API: causal_impact
# ====================================================================== #


def causal_impact(
    data: pd.DataFrame,
    pre_period: Tuple[Any, Any],
    post_period: Tuple[Any, Any],
    outcome: Optional[str] = None,
    covariates: Optional[List[str]] = None,
    model: str = "local_level",
    n_simulations: int = 1000,
    alpha: float = 0.05,
    seed: Optional[int] = None,
) -> CausalResult:
    """Google CausalImpact-style causal inference for time series.

    Fits a Bayesian structural time series model on the pre-intervention
    period and produces a counterfactual prediction for the post-period.
    The treatment effect is the difference between observed and
    counterfactual, with full posterior uncertainty.

    Parameters
    ----------
    data : pd.DataFrame
        Time-indexed DataFrame.  If *outcome* is ``None``, the first column
        is the outcome and all remaining columns are controls.
    pre_period : tuple of (start, end)
        Pre-intervention period boundaries (inclusive).  Values are matched
        against the DataFrame index.
    post_period : tuple of (start, end)
        Post-intervention period boundaries (inclusive).
    outcome : str, optional
        Column name of the outcome variable.  If ``None``, uses the first
        column.
    covariates : list of str, optional
        Column names to use as controls.  If ``None``, all columns except
        the outcome are used.
    model : ``'local_level'`` or ``'local_linear_trend'``
        State-space model type.  ``'local_level'`` uses a random-walk
        latent state; ``'local_linear_trend'`` adds a stochastic slope.
    n_simulations : int, default 1000
        Number of posterior draws for uncertainty quantification.
    alpha : float, default 0.05
        Significance level for confidence/credible intervals.
    seed : int, optional
        Random seed for reproducibility.

    Returns
    -------
    CausalResult
        Unified result object with:
        - ``estimate`` — average treatment effect on the treated (ATT)
        - ``detail`` — per-period effects DataFrame
        - ``model_info`` — counterfactual trajectories, cumulative effects,
          regression coefficients, posterior draws, and diagnostics

    Notes
    -----
    The implementation follows Brodersen et al. (2015).  Regression
    coefficients are estimated via ridge regression with GCV-selected
    penalty (an empirical-Bayes analogue of the spike-and-slab prior).
    Hyperparameters (observation/level/slope noise) are estimated by
    maximum likelihood through the Kalman filter.

    Examples
    --------
    >>> import pandas as pd
    >>> import statspai as sp
    >>> # Wide-format: columns = [outcome, control1, control2, ...]
    >>> result = sp.synth.causal_impact(
    ...     data, pre_period=(1, 70), post_period=(71, 100)
    ... )
    >>> print(result.summary())

    References
    ----------
    Brodersen, K.H., Gallusser, F., Koehler, J., Remy, N. and Scott, S.L.
    (2015).  "Inferring causal impact using Bayesian structural time series
    models."  *Annals of Applied Statistics*, 9(1), 247-274. [@brodersen2015inferring]
    """
    # ------------------------------------------------------------------
    # Input validation
    # ------------------------------------------------------------------
    if not isinstance(data, pd.DataFrame):
        raise TypeError("`data` must be a pandas DataFrame.")
    if data.empty:
        raise ValueError("`data` is empty.")

    use_trend = model == "local_linear_trend"
    if model not in ("local_level", "local_linear_trend"):
        raise ValueError(
            f"Unknown model '{model}'. Choose 'local_level' or 'local_linear_trend'."
        )

    # Resolve outcome / covariates
    if outcome is None:
        outcome = data.columns[0]
    if covariates is None:
        covariates = [c for c in data.columns if c != outcome]

    if outcome not in data.columns:
        raise ValueError(f"Outcome column '{outcome}' not found in data.")
    for c in covariates:
        if c not in data.columns:
            raise ValueError(f"Covariate column '{c}' not found in data.")

    # ------------------------------------------------------------------
    # Period slicing
    # ------------------------------------------------------------------
    idx = data.index
    pre_mask = (idx >= pre_period[0]) & (idx <= pre_period[1])
    post_mask = (idx >= post_period[0]) & (idx <= post_period[1])

    if pre_mask.sum() == 0:
        raise ValueError("No observations found in the pre-period.")
    if post_mask.sum() == 0:
        raise ValueError("No observations found in the post-period.")

    # Ensure pre ends before post starts
    pre_end_idx = idx[pre_mask][-1]
    post_start_idx = idx[post_mask][0]
    if pre_end_idx >= post_start_idx:
        raise ValueError(
            "Pre-period must end before post-period starts. "
            f"Got pre_end={pre_end_idx}, post_start={post_start_idx}."
        )

    y_pre = data.loc[pre_mask, outcome].values.astype(float)
    y_post = data.loc[post_mask, outcome].values.astype(float)
    X_pre = data.loc[pre_mask, covariates].values.astype(float) if covariates else np.zeros((len(y_pre), 0))
    X_post = data.loc[post_mask, covariates].values.astype(float) if covariates else np.zeros((len(y_post), 0))

    # Handle missing covariates: fill NaN with column means from pre-period
    if X_pre.shape[1] > 0:
        col_means = np.nanmean(X_pre, axis=0)
        for k in range(X_pre.shape[1]):
            nan_mask_pre = np.isnan(X_pre[:, k])
            X_pre[nan_mask_pre, k] = col_means[k]
            nan_mask_post = np.isnan(X_post[:, k])
            X_post[nan_mask_post, k] = col_means[k]

    # ------------------------------------------------------------------
    # Fit model on pre-period
    # ------------------------------------------------------------------
    fitted = _fit_model(y_pre, X_pre, use_trend=use_trend)

    # ------------------------------------------------------------------
    # One-step-ahead prediction errors (model checking)
    # ------------------------------------------------------------------
    one_step_errors = fitted["kf_pre"]["prediction_errors"]

    # ------------------------------------------------------------------
    # Posterior simulation of counterfactual in post-period
    # ------------------------------------------------------------------
    rng = np.random.default_rng(seed)
    cf_draws = _simulate_counterfactual(
        fitted, X_post, n_simulations=n_simulations, rng=rng,
    )

    cf_mean = cf_draws.mean(axis=0)
    cf_lower = np.percentile(cf_draws, 100 * alpha / 2, axis=0)
    cf_upper = np.percentile(cf_draws, 100 * (1 - alpha / 2), axis=0)

    # ------------------------------------------------------------------
    # Point-wise treatment effects
    # ------------------------------------------------------------------
    effects = y_post - cf_mean
    effects_draws = y_post[np.newaxis, :] - cf_draws  # (n_sim, T1)

    # Cumulative effects
    cum_effects = np.cumsum(effects)
    cum_draws = np.cumsum(effects_draws, axis=1)
    cum_lower = np.percentile(cum_draws, 100 * alpha / 2, axis=0)
    cum_upper = np.percentile(cum_draws, 100 * (1 - alpha / 2), axis=0)

    # Summary statistics
    att = float(np.mean(effects))
    att_draws = effects_draws.mean(axis=1)  # (n_sim,)
    se = float(np.std(att_draws, ddof=1))

    # Two-sided Bayesian p-value: P(effect has opposite sign)
    if att >= 0:
        pvalue = float(np.mean(att_draws <= 0))
    else:
        pvalue = float(np.mean(att_draws >= 0))
    pvalue = max(pvalue, 1.0 / (n_simulations + 1))
    # Convert one-sided to two-sided
    pvalue = min(2.0 * pvalue, 1.0)

    ci = (
        float(np.percentile(att_draws, 100 * alpha / 2)),
        float(np.percentile(att_draws, 100 * (1 - alpha / 2))),
    )

    # Relative effect
    cf_period_mean = float(np.mean(cf_mean))
    relative_effect = att / cf_period_mean if abs(cf_period_mean) > 1e-10 else np.nan

    # ------------------------------------------------------------------
    # Also produce a full-series counterfactual (pre + post)
    # ------------------------------------------------------------------
    sm_pre = fitted["sm_pre"]
    cf_pre_mean = sm_pre["smoothed_state"][:, 0]
    if X_pre.shape[1] > 0:
        cf_pre_mean = cf_pre_mean + X_pre @ fitted["beta"]

    full_times = data.index[pre_mask | post_mask]
    full_observed = np.concatenate([y_pre, y_post])
    full_cf_mean = np.concatenate([cf_pre_mean, cf_mean])
    full_cf_lower = np.concatenate([cf_pre_mean, cf_lower])  # pre = point estimate
    full_cf_upper = np.concatenate([cf_pre_mean, cf_upper])

    # ------------------------------------------------------------------
    # Detail DataFrame
    # ------------------------------------------------------------------
    post_times = data.index[post_mask]
    effects_df = pd.DataFrame({
        "time": post_times,
        "observed": y_post,
        "counterfactual": cf_mean,
        "counterfactual_lower": cf_lower,
        "counterfactual_upper": cf_upper,
        "effect": effects,
        "effect_lower": np.percentile(effects_draws, 100 * alpha / 2, axis=0),
        "effect_upper": np.percentile(effects_draws, 100 * (1 - alpha / 2), axis=0),
        "cumulative_effect": cum_effects,
        "cumulative_lower": cum_lower,
        "cumulative_upper": cum_upper,
    })

    # Trajectory DataFrame (full series for plotting)
    trajectory_df = pd.DataFrame({
        "time": full_times,
        "observed": full_observed,
        "counterfactual": full_cf_mean,
        "counterfactual_lower": full_cf_lower,
        "counterfactual_upper": full_cf_upper,
    })

    # ------------------------------------------------------------------
    # Model info
    # ------------------------------------------------------------------
    model_info: Dict[str, Any] = {
        "model_type": model,
        "sigma_obs": fitted["sigma_obs"],
        "sigma_level": fitted["sigma_level"],
        "sigma_slope": fitted["sigma_slope"] if use_trend else None,
        "regression_coefficients": dict(zip(covariates, fitted["beta"])) if covariates else {},
        "n_pre_periods": int(pre_mask.sum()),
        "n_post_periods": int(post_mask.sum()),
        "n_covariates": len(covariates),
        "n_simulations": n_simulations,
        "counterfactual_mean": cf_mean,
        "counterfactual_lower": cf_lower,
        "counterfactual_upper": cf_upper,
        "cumulative_effect": cum_effects,
        "cumulative_effect_ci": (cum_lower, cum_upper),
        "relative_effect": relative_effect,
        "posterior_draws": cf_draws,
        "one_step_ahead_prediction_errors": one_step_errors,
        "effects_by_period": effects_df,
        "trajectory": trajectory_df,
        "Y_observed": full_observed,
        "Y_counterfactual": full_cf_mean,
        "times": full_times,
        "pre_period": pre_period,
        "post_period": post_period,
        "mle_converged": fitted["mle_result"].success,
    }

    return CausalResult(
        method="Bayesian Structural Time Series (CausalImpact)",
        estimand="ATT",
        estimate=att,
        se=se,
        pvalue=pvalue,
        ci=ci,
        alpha=alpha,
        n_obs=len(data),
        detail=effects_df,
        model_info=model_info,
        _citation_key="causal_impact",
    )


# ====================================================================== #
#  Public API: bsts_synth (panel-data SCM interface)
# ====================================================================== #


def bsts_synth(
    data: pd.DataFrame,
    outcome: str,
    unit: str,
    time: str,
    treated_unit: Any,
    treatment_time: Any,
    covariates: Optional[List[str]] = None,
    model: str = "local_level",
    n_simulations: int = 1000,
    alpha: float = 0.05,
    seed: Optional[int] = None,
) -> CausalResult:
    """BSTS synthetic control with a panel-data interface.

    Converts long-format panel data into the wide format expected by
    :func:`causal_impact`, using control-unit outcome series as
    covariates/regressors.  This provides a CausalImpact-style analysis
    that integrates seamlessly with StatsPAI's other synthetic-control
    methods.

    Parameters
    ----------
    data : pd.DataFrame
        Long-format panel data with columns for unit, time, and outcome.
    outcome : str
        Outcome variable column name.
    unit : str
        Unit identifier column name.
    time : str
        Time period column name.
    treated_unit : any
        Identifier of the treated unit.
    treatment_time : any
        First treatment period (inclusive).
    covariates : list of str, optional
        Additional time-varying covariates to include alongside control
        unit series.  Each covariate is averaged across control units
        per time period and appended as an extra regressor.
    model : ``'local_level'`` or ``'local_linear_trend'``
        State-space model type.
    n_simulations : int, default 1000
        Number of posterior draws.
    alpha : float, default 0.05
        Significance level.
    seed : int, optional
        Random seed.

    Returns
    -------
    CausalResult
        Unified result object.  ``model_info`` contains additional keys
        ``treated_unit``, ``treatment_time``, ``donor_units``.

    Examples
    --------
    >>> import statspai as sp
    >>> result = sp.synth.bsts_synth(
    ...     data, outcome='gdp', unit='country', time='year',
    ...     treated_unit='West Germany', treatment_time=1990,
    ... )
    >>> print(result.summary())

    See Also
    --------
    causal_impact : Wide-format CausalImpact interface.
    """
    # ------------------------------------------------------------------
    # Input validation
    # ------------------------------------------------------------------
    if not isinstance(data, pd.DataFrame):
        raise TypeError("`data` must be a pandas DataFrame.")
    for col in [outcome, unit, time]:
        if col not in data.columns:
            raise ValueError(f"Column '{col}' not found in data.")
    if treated_unit not in data[unit].values:
        raise ValueError(f"Treated unit '{treated_unit}' not found in '{unit}' column.")

    # ------------------------------------------------------------------
    # Reshape: long → wide
    # ------------------------------------------------------------------
    panel = data.pivot_table(index=time, columns=unit, values=outcome, aggfunc="mean")
    panel = panel.sort_index()

    if treated_unit not in panel.columns:
        raise ValueError(
            f"Treated unit '{treated_unit}' has no outcome data after pivoting."
        )

    all_times = panel.index
    donors = [u for u in panel.columns if u != treated_unit]

    if len(donors) == 0:
        raise ValueError("No control/donor units found.")

    # Drop donors with too many missing values in pre-period
    pre_times = all_times[all_times < treatment_time]
    keep_donors = []
    for d in donors:
        missing_frac = panel.loc[pre_times, d].isna().mean()
        if missing_frac < 0.5:
            keep_donors.append(d)
    if len(keep_donors) == 0:
        raise ValueError(
            "All donor units have >50% missing data in the pre-period."
        )
    donors = keep_donors

    # Build wide DataFrame: outcome = treated, covariates = donors
    wide = pd.DataFrame(index=all_times)
    wide[outcome] = panel[treated_unit].values
    for d in donors:
        wide[str(d)] = panel[d].values

    # Add averaged covariates if supplied
    extra_cov_names: List[str] = []
    if covariates:
        for cov in covariates:
            if cov not in data.columns:
                raise ValueError(f"Covariate column '{cov}' not found in data.")
            cov_panel = data.pivot_table(
                index=time, columns=unit, values=cov, aggfunc="mean",
            ).sort_index()
            # Average across control units
            ctrl_cols = [u for u in cov_panel.columns if u != treated_unit]
            if ctrl_cols:
                col_name = f"_cov_{cov}"
                wide[col_name] = cov_panel[ctrl_cols].mean(axis=1).values
                extra_cov_names.append(col_name)

    # Forward/backward fill remaining NaN then fill zeros
    wide = wide.ffill().bfill().fillna(0.0)

    # ------------------------------------------------------------------
    # Period boundaries
    # ------------------------------------------------------------------
    pre_start = all_times[0]
    # Pre-period ends at the last time strictly before treatment
    pre_end_candidates = all_times[all_times < treatment_time]
    if len(pre_end_candidates) == 0:
        raise ValueError("No pre-treatment periods found.")
    pre_end = pre_end_candidates[-1]

    post_start = treatment_time
    post_end = all_times[-1]

    cov_cols = [str(d) for d in donors] + extra_cov_names

    # ------------------------------------------------------------------
    # Delegate to causal_impact
    # ------------------------------------------------------------------
    result = causal_impact(
        data=wide,
        pre_period=(pre_start, pre_end),
        post_period=(post_start, post_end),
        outcome=outcome,
        covariates=cov_cols,
        model=model,
        n_simulations=n_simulations,
        alpha=alpha,
        seed=seed,
    )

    # Augment model_info with panel-specific metadata
    result.model_info["treated_unit"] = treated_unit
    result.model_info["treatment_time"] = treatment_time
    result.model_info["donor_units"] = donors
    result.model_info["n_donors"] = len(donors)
    result.model_info["interface"] = "bsts_synth"

    return result
