# Copyright 2019-2023 The TFP CausalImpact Authors
# Copyright 2014 Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# This software is distributed on an "AS IS" BASIS,
# without warranties or conditions of any kind.


import dataclasses
import logging
import math
from typing import Dict, List, Optional, Tuple, Union

from causalimpact import posterior_processing
import causalimpact.data as cid
from causalimpact.indices import InputDateType, OutputDateType, OutputPeriodType

import numpy as np
import pandas as pd
import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow_probability.python.experimental.distributions import MultivariateNormalPrecisionFactorLinearOperator
from tensorflow_probability.python.experimental.sts_gibbs import gibbs_sampler
from tensorflow_probability.python.internal import prefer_static as ps

tfb = tfp.bijectors
tfd = tfp.distributions

TensorLike = tf.types.experimental.TensorLike
_SeedType = Union[int, Tuple[int, int], TensorLike]


@dataclasses.dataclass
class CausalImpactPosteriorSamples:
    r"""
    Posterior samples of model parameters and latent states.

    Attributes
    ----------
    observation_noise_scale : tf.Tensor
        Samples of the observation noise scale, $\sigma_{\varepsilon}$.
    level_scale : Optional[tf.Tensor]
        Samples of the local level's scale parameter.
    level : Optional[tf.Tensor]
        Samples of the local level state $\ell_t$.
    weights : Optional[tf.Tensor]
        Samples of regression coefficients $\beta$.
    seasonal_drift_scales : Optional[tf.Tensor]
        Samples of seasonal drift scales.
    seasonal_levels : Optional[tf.Tensor]
        Samples of seasonal effect states.
    """
    observation_noise_scale: tf.Tensor
    level_scale: Optional[tf.Tensor]
    level: Optional[tf.Tensor]
    weights: Optional[tf.Tensor]
    seasonal_drift_scales: Optional[tf.Tensor]
    seasonal_levels: Optional[tf.Tensor]


@dataclasses.dataclass
class CausalImpactAnalysis:
    r"""
    Results of a CausalImpact analysis.

    Attributes
    ----------
    series : pd.DataFrame
        Time-indexed DataFrame with:

        - observed: $`y_t`$
        - `posterior_mean`: $\hat{y}_t$ (posterior mean prediction)
        - `posterior_lower`, `posterior_upper`: Credible interval bounds for $\hat{y}_t$.
        - `point_effects_mean`, `point_effects_lower`, `point_effects_upper`: Mean and CI for point effects $y_t - \hat{y}_t$.
        - `cumulative_effects_mean`, `cumulative_effects_lower`, `cumulative_effects_upper`: Mean and CI for cumulative effects.

    summary : pd.DataFrame
        Summary over the post-intervention period, including:
        - Actual vs predicted outcomes
        - Absolute and relative effects
        - p-values

    posterior_samples : CausalImpactPosteriorSamples
        Posterior samples of latent variables and parameters.
    """
    series: pd.DataFrame
    summary: pd.DataFrame
    posterior_samples: CausalImpactPosteriorSamples


@dataclasses.dataclass
class DataOptions:
    r"""
    Data configuration options.

    Attributes
    ----------
    outcome_column : Optional[str]
        Name of the column in `data` containing $`y_t`$.
    standardize_data : bool
        Whether to standardize $y_t$ and covariates.
    dtype : tf.dtypes.DType
        Data type for computations.
    """
    outcome_column: Optional[str] = None
    standardize_data: bool = True
    dtype: tf.dtypes.DType = tf.float32


@dataclasses.dataclass(frozen=True)
class Seasons:
    r"""
    Seasonal modeling options.

    Attributes
    ----------
    num_seasons : int
        Number of seasons in the cycle.
    num_steps_per_season : Union[int, Tuple[int], Tuple[Tuple[int]]]
        Steps within each season. Can be a single integer or a tuple for complex patterns.
    """
    num_seasons: int
    num_steps_per_season: Union[int, Tuple[int], Tuple[Tuple[int]]] = 1


@dataclasses.dataclass
class ModelOptions:
    r"""
    Model configuration parameters.

    Attributes
    ----------
    prior_level_sd : float
        Prior standard deviation for the local level, relative to data's std. dev.
    seasons : List[Seasons]
        List of seasonal components.
    """
    prior_level_sd: float = 0.01
    seasons: List[Seasons] = dataclasses.field(default_factory=list)


@dataclasses.dataclass
class InferenceOptions:
    r"""
    Options for posterior inference (MCMC).

    Attributes
    ----------
    num_results : int
        Number of posterior draws.
    num_warmup_steps : Optional[int]
        Number of warmup (burn-in) steps.
    """
    num_results: int = 900
    num_warmup_steps: Optional[int] = None

    def __post_init__(self):
        # If not set, use about 1/9 of num_results as warmup
        if self.num_warmup_steps is None:
            self.num_warmup_steps = math.ceil(self.num_results / 9)


def fit_causalimpact(data: pd.DataFrame,
                     pre_period: Tuple[InputDateType, InputDateType],
                     post_period: Tuple[InputDateType, InputDateType],
                     alpha: float = 0.05,
                     seed: Optional[_SeedType] = None,
                     data_options: Optional[DataOptions] = None,
                     model_options: Optional[ModelOptions] = None,
                     inference_options: Optional[InferenceOptions] = None,
                     **kwargs) -> CausalImpactAnalysis:
    r"""
    Fit the CausalImpact model and compute posterior effects.

    Steps:
    1. Process data, identify pre/post intervention periods.
    2. Fit an STS model using MCMC (Gibbs sampling) on pre-period data.
    3. Obtain posterior predictions $\hat{y}_t$ over the full period.
    4. Compute point effects $y_t - \hat{y}_t$ and cumulative effects.

    Parameters
    ----------
    data : pd.DataFrame
        DataFrame with the outcome time series $y_t$ and optionally covariates.
    pre_period : tuple
        (start, end) of the pre-intervention period.
    post_period : tuple
        (start, end) of the post-intervention period.
    alpha : float, optional
        Significance level for credible intervals (default 0.05 for 95% CI).
    seed : int, tuple, or TensorLike, optional
        Random seed for reproducibility.
    data_options : DataOptions, optional
        Configuration for data handling.
    model_options : ModelOptions, optional
        Model hyperparameters.
    inference_options : InferenceOptions, optional
        MCMC sampling configurations.
    **kwargs
        Experimental arguments, if any.

    Returns
    -------
    CausalImpactAnalysis
        Analysis results, including posterior samples and effect estimates.
    """
    tf_log_level = tf.get_logger().level
    tf.get_logger().setLevel(logging.ERROR)
    try:
        data_options = data_options if data_options is not None else DataOptions()
        model_options = model_options if model_options is not None else ModelOptions()
        inference_options = inference_options if inference_options is not None else InferenceOptions()

        experimental_model = kwargs.pop("experimental_model", None)
        experimental_tf_function_cache_key_addition = kwargs.pop("experimental_tf_function_cache_key_addition", 0)
        if kwargs:
            raise TypeError(f"Received unknown {kwargs=}")

        ci_data = cid.CausalImpactData(
            data=data,
            pre_intervention_period=pre_period,
            post_intervention_period=post_period,
            target_col_name=data_options.outcome_column,
            standardize_data=data_options.standardize_data,
            dtype=data_options.dtype)

        posterior_samples, posterior_means, posterior_trajectories = _train_causalimpact_sts(
            ci_data=ci_data,
            prior_level_sd=model_options.prior_level_sd,
            seed=seed,
            num_results=inference_options.num_results,
            num_warmup_steps=inference_options.num_warmup_steps,
            model=experimental_model,
            dtype=data_options.dtype,
            seasons=model_options.seasons,
            experimental_tf_function_cache_key_addition=experimental_tf_function_cache_key_addition
        )

        series, summary = _compute_impact(
            posterior_means=posterior_means,
            posterior_trajectories=posterior_trajectories,
            ci_data=ci_data,
            alpha=alpha)

        if posterior_samples.seasonal_levels.shape[-1] > 0:
            seasonal_levels = []
            idx = 0
            for season in model_options.seasons:
                seasonal_levels.append(posterior_samples.seasonal_levels[..., idx])
                idx += season.num_seasons - 1
            seasonal_levels = tf.stack(seasonal_levels, axis=-1)
        else:
            seasonal_levels = posterior_samples.seasonal_levels

        ci_posterior_samples = CausalImpactPosteriorSamples(
            observation_noise_scale=posterior_samples.observation_noise_scale,
            level_scale=posterior_samples.level_scale,
            level=posterior_samples.level,
            weights=(posterior_samples.weights
                     if posterior_samples.weights.shape[1] > 0 else None),
            seasonal_drift_scales=(
                posterior_samples.seasonal_drift_scales
                if posterior_samples.seasonal_drift_scales.shape[-1] > 0 else None),
            seasonal_levels=seasonal_levels
        )
        return CausalImpactAnalysis(series, summary, ci_posterior_samples)
    finally:
        tf.get_logger().setLevel(tf_log_level)


@tf.function(autograph=False, jit_compile=False)
def _run_gibbs_sampler(
        sts_model: Optional[tfp.sts.StructuralTimeSeries],
        outcome_ts: TensorLike,
        outcome_sd: TensorLike,
        design_matrix: Optional[TensorLike],
        num_results: int,
        num_warmup_steps: int,
        observation_noise_scale: TensorLike,
        level_scale: TensorLike,
        seasonal_drift_scales: TensorLike,
        weights: TensorLike,
        level: TensorLike,
        slope: TensorLike,
        seed: TensorLike,
        dtype,
        seasons: List[Seasons],
        experimental_tf_function_cache_key_addition: int):
    r"""
    Run Gibbs sampling for STS parameters and states.

    Given a structural time series model and observed data (with possible missing values),
    this function draws samples from the posterior distribution over parameters and latent states
    using Gibbs sampling.

    Parameters
    ----------
    sts_model : tfp.sts.StructuralTimeSeries, optional
        The STS model. If None, a default model is constructed.
    outcome_ts : TensorLike
        Time series of $y_t$, possibly with NaNs for post-intervention.
    outcome_sd : TensorLike
        Standard deviation of the pre-intervention data, used to scale priors.
    design_matrix : Optional[TensorLike]
        Covariates, if any.
    num_results : int
        Number of MCMC samples to draw.
    num_warmup_steps : int
        Warmup (burn-in) steps before collecting samples.
    observation_noise_scale : TensorLike
        Initial guess for observation noise scale $\sigma_{\varepsilon}$.
    level_scale : TensorLike
        Prior scale for the local level component.
    seasonal_drift_scales : TensorLike
        Drift scales for seasonal effects.
    weights : TensorLike
        Initial regression weights (if covariates are included).
    level : TensorLike
        Initial level state.
    slope : TensorLike
        Initial slope state (if any, else zero).
    seed : TensorLike
        PRNG seed.
    dtype :
        Data type for computations.
    seasons : List[Seasons]
        Seasonal configuration.
    experimental_tf_function_cache_key_addition : int
        Experimental parameter to control tracing.

    Returns
    -------
    posterior_samples : gibbs_sampler.GibbsSamplerState
        Samples of parameters and states.
    posterior_means : tf.Tensor
        Posterior mean predictions $\hat{y}_t$.
    posterior_trajectories : tf.Tensor
        Posterior predictive samples for $\hat{y}_t$.
    """
    if not sts_model:
        sts_model = _build_default_gibbs_model(
            design_matrix=design_matrix,
            outcome_ts=outcome_ts,
            level_scale=level_scale,
            outcome_sd=outcome_sd,
            dtype=dtype,
            seasons=seasons)

    sample_seed, forecast_seed = tfp.random.split_seed(seed)
    posterior_samples = gibbs_sampler.fit_with_gibbs_sampling(
        sts_model,
        outcome_ts,
        num_results=num_results,
        num_warmup_steps=num_warmup_steps,
        initial_state=gibbs_sampler.GibbsSamplerState(
            observation_noise_scale=observation_noise_scale,
            level_scale=level_scale,
            slope_scale=tf.zeros([], dtype=dtype),
            weights=weights,
            level=level,
            slope=slope,
            seed=None,
            seasonal_drift_scales=seasonal_drift_scales,
            seasonal_levels=tf.zeros(
                shape=gibbs_sampler.get_seasonal_latents_shape(
                    outcome_ts.time_series, sts_model),
                dtype=dtype)),
        default_pseudo_observations=tf.ones([], dtype=dtype) * 0.01,
        seed=sample_seed,
        experimental_use_dynamic_cholesky=True,
        experimental_use_weight_adjustment=True)

    posterior_means, posterior_trajectories = _get_posterior_means_and_trajectories(
        sts_model=sts_model,
        posterior_samples=posterior_samples,
        seed=forecast_seed)

    return posterior_samples, posterior_means, posterior_trajectories


def _build_default_gibbs_model(
        design_matrix: Optional[tf.Tensor],
        outcome_ts: tfp.sts.MaskedTimeSeries,
        level_scale: tf.Tensor,
        outcome_sd: tf.Tensor,
        dtype,
        seasons: List[Seasons],
):
    """
    Construct a default STS model for Gibbs sampling.

    Includes:
    - A local level component with an InverseGamma prior on the level variance.
    - Optional regression using covariates (if provided).
    - Optional seasonal components.

    Parameters
    ----------
    design_matrix : Optional[tf.Tensor]
        Covariates for regression.
    outcome_ts : tfp.sts.MaskedTimeSeries
        Observed series with possible NaNs.
    level_scale : tf.Tensor
        Scale for the local level prior.
    outcome_sd : tf.Tensor
        Standard deviation of pre-period data.
    dtype :
        Data type.
    seasons : List[Seasons]
        Seasonal configurations.

    Returns
    -------
    sts_model : tfp.sts.StructuralTimeSeries
        The constructed STS model suitable for Gibbs sampling.
    """
    local_level_prior_sample_size = tf.constant(32., dtype=dtype)
    level_concentration = tf.cast(local_level_prior_sample_size / 2., dtype=dtype)
    level_variance_prior_scale = level_scale * level_scale * (local_level_prior_sample_size / 2.)
    level_variance_prior = tfd.InverseGamma(
        concentration=level_concentration, scale=level_variance_prior_scale)
    level_variance_prior.upper_bound = outcome_sd

    if design_matrix is not None:
        observation_noise_variance_prior = tfd.InverseGamma(
            concentration=tf.constant(25., dtype=dtype),
            scale=tf.math.square(outcome_sd) * tf.constant(5., dtype=dtype))
    else:
        observation_noise_variance_prior = tfd.InverseGamma(
            concentration=tf.constant(0.005, dtype=dtype),
            scale=tf.math.square(outcome_sd) * tf.constant(0.005, dtype=dtype))
    observation_noise_variance_prior.upper_bound = outcome_sd * tf.constant(1.2, dtype=dtype)

    if design_matrix is not None:
        design_shape = ps.shape(design_matrix)
        num_outputs = design_shape[-2]
        num_dimensions = design_shape[-1]
        sparse_weights_nonzero_prob = tf.minimum(
            tf.constant(1., dtype=dtype), 3. / num_dimensions)
        x_transpose_x = tf.matmul(design_matrix, design_matrix, transpose_a=True)
        weights_prior_precision = 0.01 * tf.linalg.set_diag(
            0.5 * x_transpose_x, tf.linalg.diag_part(x_transpose_x)) / num_outputs
        precision_factor = tf.linalg.cholesky(weights_prior_precision)
        weights_prior = MultivariateNormalPrecisionFactorLinearOperator(
            precision_factor=tf.linalg.LinearOperatorFullMatrix(precision_factor),
            precision=tf.linalg.LinearOperatorFullMatrix(weights_prior_precision))
    else:
        sparse_weights_nonzero_prob = None
        weights_prior = None

    initial_level_prior = tfd.Normal(
        loc=tf.cast(outcome_ts.time_series[..., 0], dtype=dtype),
        scale=outcome_sd)

    seasonal_components = []
    seasonal_variance_prior = tfd.InverseGamma(
        concentration=0.005, scale=5e-7 * tf.square(outcome_sd))
    seasonal_variance_prior.upper_bound = outcome_sd
    for seasonal_options in seasons:
        seasonal_components.append(
            tfp.sts.Seasonal(
                num_seasons=seasonal_options.num_seasons,
                num_steps_per_season=np.array(seasonal_options.num_steps_per_season),
                allow_drift=True,
                constrain_mean_effect_to_zero=True,
                drift_scale_prior=tfd.TransformedDistribution(
                    bijector=tfb.Invert(tfb.Square()),
                    distribution=seasonal_variance_prior),
                initial_effect_prior=tfd.Normal(loc=0., scale=outcome_sd)))

    return gibbs_sampler.build_model_for_gibbs_fitting(
        outcome_ts,
        design_matrix=design_matrix,
        weights_prior=weights_prior,
        level_variance_prior=level_variance_prior,
        slope_variance_prior=None,
        observation_noise_variance_prior=observation_noise_variance_prior,
        initial_level_prior=initial_level_prior,
        sparse_weights_nonzero_prob=sparse_weights_nonzero_prob,
        seasonal_components=seasonal_components)


def _train_causalimpact_sts(
        *,
        ci_data: cid.CausalImpactData,
        prior_level_sd,
        seed: _SeedType,
        num_results: int,
        num_warmup_steps: int,
        model: Optional[tfp.sts.StructuralTimeSeries] = None,
        dtype,
        seasons: List[Seasons],
        experimental_tf_function_cache_key_addition: int = 0,
) -> Tuple[gibbs_sampler.GibbsSamplerState, TensorLike, TensorLike]:
    r"""
    Train the STS model via Gibbs sampling and return posterior predictions.

    Steps:
    - Extend $y_t$ into the post-period with NaNs.
    - Run Gibbs sampling to get posterior samples.
    - Return posterior samples and $\hat{y}_t$ predictions.

    Parameters
    ----------
    ci_data : cid.CausalImpactData
        Processed data with pre/post intervention periods.
    prior_level_sd : float
        Prior level standard deviation (scaled).
    seed : _SeedType
        Seed for reproducibility.
    num_results : int
        Number of posterior draws.
    num_warmup_steps : int
        Warmup steps (burn-in).
    model : Optional[tfp.sts.StructuralTimeSeries]
        Custom STS model, if any.
    dtype :
        Data type.
    seasons : List[Seasons]
        Seasonal configurations.
    experimental_tf_function_cache_key_addition : int
        Experimental parameter for tracing.

    Returns
    -------
    posterior_samples : gibbs_sampler.GibbsSamplerState
        Posterior samples of parameters and states.
    y_hat_means : tf.Tensor
        Posterior mean predictions $\hat{y}_t$.
    y_hat_samples : tf.Tensor
        Posterior predictive samples of $\hat{y}_t$.
    """
    if isinstance(seed, int):
        seed = (0, seed)
    seed = tfp.random.sanitize_seed(seed)

    X = (tf.convert_to_tensor(ci_data.normalized_whole_period_features.values, dtype=dtype)
         if ci_data.normalized_whole_period_features is not None else None)

    post_len = ci_data.normalized_after_pre_intervention_data.shape[0]
    y_pre = ci_data.pre_intervention_target_ts.time_series
    y_extended = tf.concat([y_pre, tf.fill([post_len], tf.constant(float("nan"), dtype=dtype))], axis=0)
    is_missing_extended = tf.concat([ci_data.pre_intervention_target_ts.is_missing,
                                     tf.fill([post_len], True)], axis=0)
    extended_target_ts = tfp.sts.MaskedTimeSeries(
        time_series=y_extended, is_missing=is_missing_extended)

    y_pre_std = tf.convert_to_tensor(
        np.nanstd(ci_data.pre_intervention_target_ts.time_series, ddof=1), dtype=dtype)

    r2 = 0.8
    if X is not None:
        observation_noise_scale = tf.cast(tf.math.sqrt(1 - r2), dtype=dtype) * y_pre_std
    else:
        observation_noise_scale = y_pre_std

    level_scale = tf.ones([], dtype=dtype) * prior_level_sd * y_pre_std
    seasonal_drift_scales = 0.01 * y_pre_std * tf.ones(shape=[len(seasons)], dtype=dtype)

    if ci_data.normalized_whole_period_features is None:
        weights_init = tf.zeros([0], dtype=dtype)
    else:
        weights_init = tf.zeros(ci_data.normalized_whole_period_features.shape[-1:], dtype=dtype)

    level_init = tf.zeros_like(extended_target_ts.time_series)
    slope_init = tf.zeros_like(extended_target_ts.time_series)

    posterior_samples, y_hat_means, y_hat_samples = _run_gibbs_sampler(
        sts_model=model,
        outcome_ts=extended_target_ts,
        outcome_sd=y_pre_std,
        design_matrix=X,
        num_results=num_results,
        num_warmup_steps=num_warmup_steps,
        observation_noise_scale=observation_noise_scale,
        level_scale=level_scale,
        seasonal_drift_scales=seasonal_drift_scales,
        weights=weights_init,
        level=level_init,
        slope=slope_init,
        seed=seed,
        dtype=dtype,
        seasons=seasons,
        experimental_tf_function_cache_key_addition=experimental_tf_function_cache_key_addition
    )

    return posterior_samples, y_hat_means, y_hat_samples


def _get_posterior_means_and_trajectories(sts_model, posterior_samples, seed):
    r"""
    Extract posterior mean and trajectory samples of $\hat{y}_t$.

    Parameters
    ----------
    sts_model : tfp.sts.StructuralTimeSeries
        The STS model used.
    posterior_samples : gibbs_sampler.GibbsSamplerState
        Posterior samples of parameters and states.
    seed : int or tuple
        Seed for sampling predictive trajectories.

    Returns
    -------
    y_hat_mean : tf.Tensor
        Mean $\hat{y}_t$ predictions.
    y_hat_samples : tf.Tensor
        Sampled trajectories of $\hat{y}_t$.
    """
    predictive_distributions = gibbs_sampler.one_step_predictive(
        sts_model,
        posterior_samples,
        thin_every=1,
        use_zero_step_prediction=True)

    y_hat_mean = predictive_distributions.mean()
    y_hat_samples = tf.transpose(predictive_distributions.components_distribution.sample(seed=seed))
    return y_hat_mean, y_hat_samples


def _compute_impact(
        posterior_means,
        posterior_trajectories,
        ci_data: cid.CausalImpactData,
        alpha: float = 0.05,
) -> Tuple[pd.DataFrame, pd.DataFrame]:
    r"""
    Compute time-series effects and summary statistics.

    - Derive credible intervals for $\hat{y}_t$.
    - Compute point effects $y_t - \hat{y}_t$.
    - Compute cumulative effects.
    - Summarize results.

    Parameters
    ----------
    posterior_means : tf.Tensor
        Posterior mean predictions of $\hat{y}_t$.
    posterior_trajectories : tf.Tensor
        Posterior predictive samples of $\hat{y}_t$.
    ci_data : cid.CausalImpactData
        Data structure with pre/post intervention info.
    alpha : float
        Significance level for intervals.

    Returns
    -------
    series : pd.DataFrame
        Detailed time series with observed, predicted, and effect columns.
    summary : pd.DataFrame
        Aggregate summary over the post-intervention period.
    """
    if not 0 < alpha < 1:
        raise ValueError("`alpha` must be between 0 and 1.")

    observed_pre = ci_data.pre_intervention_data[ci_data.target_col]
    observed_post = ci_data.after_pre_intervention_data[ci_data.target_col]
    observed_post = observed_post.loc[(observed_post.index >= ci_data.post_intervention_period[0]) &
                                      (observed_post.index <= ci_data.post_intervention_period[1])]
    observed_full = pd.concat([observed_pre, observed_post], axis=0)

    quantiles = (alpha / 2.0, 1.0 - (alpha / 2.0))

    posterior_trajectories, posterior_trajectory_summary = (
        _sample_posterior_predictive(
            posterior_means=posterior_means,
            posterior_trajectories=posterior_trajectories,
            ci_data=ci_data,
            quantiles=quantiles))

    trajectory_dict = _compute_impact_trajectories(
        posterior_trajectories,
        observed_full,
        treatment_start=ci_data.post_intervention_period[0])

    series = _compute_impact_estimates(
        posterior_trajectory_summary=posterior_trajectory_summary,
        trajectory_dict=trajectory_dict,
        observed_ts_full=observed_full,
        ci_data=ci_data,
        quantiles=quantiles)

    summary = _compute_summary(
        posterior_trajectory_summary=posterior_trajectory_summary,
        trajectory_dict=trajectory_dict,
        observed_ts_post=observed_post,
        post_period=ci_data.post_intervention_period,
        quantiles=quantiles,
        alpha=alpha)

    return series, summary


def _sample_posterior_predictive(
        posterior_means: TensorLike,
        posterior_trajectories: TensorLike,
        ci_data: cid.CausalImpactData,
        quantiles: Tuple[float, float]
) -> Tuple[pd.DataFrame, pd.DataFrame]:
    r"""
    Package posterior predictions into DataFrames with mean and quantiles.

    Parameters
    ----------
    posterior_means : TensorLike
        Posterior mean $\hat{y}_t$.
    posterior_trajectories : TensorLike
        Posterior samples of $\hat{y}_t$.
    ci_data : cid.CausalImpactData
        Data structure with full data info.
    quantiles : tuple
        Quantiles for intervals (e.g., (0.025, 0.975)).

    Returns
    -------
    posterior_trajectories : pd.DataFrame
        DataFrame of sampled trajectories.
    posterior_trajectory_summary : pd.DataFrame
        DataFrame with posterior mean and quantile columns.
    """
    posterior_means = posterior_processing.process_posterior_quantities(
        ci_data, posterior_means, ["posterior_mean"])
    posterior_trajectories = _package_posterior_trajectories(
        posterior_trajectories, ci_data)
    posterior_quantiles = posterior_processing.calculate_trajectory_quantiles(
        posterior_trajectories, "posterior", quantiles)

    posterior_trajectory_summary = posterior_means.join(posterior_quantiles)
    return posterior_trajectories, posterior_trajectory_summary


def _package_posterior_trajectories(
        posterior_trajectories: TensorLike,
        ci_data: cid.CausalImpactData) -> pd.DataFrame:
    r"""
    Convert sampled $\hat{y}_t$ trajectories into a DataFrame.

    Columns are `sample_1, sample_2, ...`

    Parameters
    ----------
    posterior_trajectories : TensorLike
        Posterior samples of predictions.
    ci_data : cid.CausalImpactData
        Data for indexing.

    Returns
    -------
    pd.DataFrame
        DataFrame with each column a sample trajectory.
    """
    col_names = [f"sample_{i + 1}" for i in range(posterior_trajectories.shape[0])]
    return posterior_processing.process_posterior_quantities(
        ci_data, posterior_trajectories, col_names)


def _compute_impact_trajectories(
        posterior_trajectories: pd.DataFrame, y_full: pd.Series,
        treatment_start: OutputDateType) -> Dict[str, pd.DataFrame]:
    r"""
    Compute point and cumulative effect trajectories.

    Point effects:
    $$y_t - \hat{y}_{t, \text{sample}}$$

    Cumulative effects:
    $$\sum_{\tau=T_{\text{start}}}^{t} (y_\tau - \hat{y}_{\tau,\text{sample}})$$

    Parameters
    ----------
    posterior_trajectories : pd.DataFrame
        DataFrame of $\hat{y}_t$ samples.
    y_full : pd.Series
        Full observed data series.
    treatment_start : OutputDateType
        Start of the post-intervention period.

    Returns
    -------
    dict of pd.DataFrame
        Contains "predictions", "point_effects", "cumulative_effects".
    """
    point_effects = posterior_trajectories.sub(y_full, axis=0).mul(-1)
    cum_base = point_effects.copy()
    cum_base.loc[cum_base.index < treatment_start] = 0
    cumulative_effects = cum_base.cumsum(axis=0)

    return {
        "predictions": posterior_trajectories,
        "point_effects": point_effects,
        "cumulative_effects": cumulative_effects
    }


def _compute_impact_estimates(posterior_trajectory_summary: pd.DataFrame,
                              trajectory_dict: Dict[str, pd.DataFrame],
                              observed_ts_full: pd.Series,
                              ci_data: cid.CausalImpactData,
                              quantiles: Tuple[float, float]) -> pd.DataFrame:
    r"""
    Construct a time series DataFrame with effects and intervals.

    - point_effects_mean: $E[y_t - \hat{y}_t]$
    - cumulative_effects_mean: $\sum_{\tau=T_{\text{start}}}^{t} (y_\tau - \hat{y}_\tau)$

    Parameters
    ----------
    posterior_trajectory_summary : pd.DataFrame
        Summary with posterior_mean and intervals of $\hat{y}_t$.
    trajectory_dict : dict
        Dict with "point_effects" and "cumulative_effects".
    observed_ts_full : pd.Series
        Full observed series (pre + post).
    ci_data : cid.CausalImpactData
        Data info.
    quantiles : tuple
        Interval quantiles.

    Returns
    -------
    pd.DataFrame
        DataFrame with observed, predicted, effects, and intervals.
    """
    point_effects_mean = (observed_ts_full - posterior_trajectory_summary["posterior_mean"])
    point_effects_mean = point_effects_mean.to_frame(name="point_effects_mean")

    zero_inds = point_effects_mean.index < ci_data.post_intervention_period[0]
    cum_effects_mean_base = point_effects_mean.copy()
    cum_effects_mean_base.loc[zero_inds] = 0
    cum_effects_mean = cum_effects_mean_base.cumsum()
    cum_effects_mean.columns = ["cumulative_effects_mean"]

    point_effects_quantiles = posterior_processing.calculate_trajectory_quantiles(
        trajectory_dict["point_effects"], "point_effects", quantiles)
    cum_effects_quantiles = posterior_processing.calculate_trajectory_quantiles(
        trajectory_dict["cumulative_effects"], "cumulative_effects", quantiles)

    impact_estimates = pd.concat([
        observed_ts_full.to_frame(name="observed"),
        posterior_trajectory_summary,
        point_effects_mean,
        point_effects_quantiles,
        cum_effects_mean,
        cum_effects_quantiles
    ], axis=1)

    impact_estimates.loc[
        ((impact_estimates.index > ci_data.pre_intervention_period[1]) &
         (impact_estimates.index < ci_data.post_intervention_period[0])) |
        (impact_estimates.index > ci_data.post_intervention_period[1]),
        impact_estimates.columns.difference(
            ["observed", "posterior_mean", "posterior_lower", "posterior_upper"]
        )] = np.nan

    impact_estimates.loc[
        np.isnan(impact_estimates["observed"]),
        impact_estimates.columns.difference(
            ["observed", "posterior_mean", "posterior_lower", "posterior_upper"]
        )] = np.nan

    impact_estimates = impact_estimates.reindex(ci_data.data.index, copy=False, fill_value=np.nan)
    impact_estimates["observed"] = ci_data.data[ci_data.target_col]

    impact_estimates["pre_period_start"] = ci_data.pre_intervention_period[0]
    impact_estimates["pre_period_end"] = ci_data.pre_intervention_period[1]
    impact_estimates["post_period_start"] = ci_data.post_intervention_period[0]
    impact_estimates["post_period_end"] = ci_data.post_intervention_period[1]

    return impact_estimates


def _compute_summary(posterior_trajectory_summary: pd.DataFrame,
                     trajectory_dict: Dict[str, pd.DataFrame],
                     observed_ts_post: pd.Series, post_period: OutputPeriodType,
                     quantiles: Tuple[float, float], alpha: float) -> pd.DataFrame:
    r"""
    Compute summary statistics over the post-period:

    - Average and cumulative predicted $\hat{y}_t$
    - Absolute effects $y_t - \hat{y}_t$
    - Relative effects $(y_t/\hat{y}_t - 1)$
    - p-value for extremeness of observed outcome

    Parameters
    ----------
    posterior_trajectory_summary : pd.DataFrame
        Posterior mean and intervals of $\hat{y}_t$.
    trajectory_dict : dict
        Contains trajectories for "point_effects" and "cumulative_effects".
    observed_ts_post : pd.Series
        Observed data in the post-period.
    post_period : tuple
        (start, end) of the post-intervention period.
    quantiles : tuple
        Interval quantiles (e.g., (0.025, 0.975)).
    alpha : float
        Significance level.

    Returns
    -------
    pd.DataFrame
        Summary of actual, predicted, and effect statistics over post-period.
    """
    posterior_mean = posterior_trajectory_summary.loc[
        (posterior_trajectory_summary.index >= post_period[0]) &
        (posterior_trajectory_summary.index <= post_period[1]), "posterior_mean"]
    trajectory_dict = {
        k: v.loc[(v.index >= post_period[0]) & (v.index <= post_period[1])]
        for k, v in trajectory_dict.items()
    }

    avg_pred = posterior_mean.mean()
    cum_pred = posterior_mean.sum()

    pred_mean_samples = trajectory_dict["predictions"].mean(axis=0)
    pred_sum_samples = trajectory_dict["predictions"].sum(axis=0)
    avg_pred_lower, avg_pred_upper = pred_mean_samples.quantile(quantiles)
    cum_pred_lower, cum_pred_upper = pred_sum_samples.quantile(quantiles)

    avg_effect = observed_ts_post.mean() - avg_pred
    avg_eff_samples = trajectory_dict["point_effects"].mean(axis=0)
    avg_eff_lower, avg_eff_upper = avg_eff_samples.quantile(quantiles)

    cum_effect = observed_ts_post.sum() - cum_pred
    cum_eff_samples = trajectory_dict["point_effects"].sum(axis=0)
    cum_eff_lower, cum_eff_upper = cum_eff_samples.quantile(quantiles)

    rel_eff_samples = (observed_ts_post.sum() / pred_sum_samples - 1.)
    rel_eff_lower, rel_eff_upper = rel_eff_samples.quantile(quantiles)

    summary_dict = {
        "actual": {"average": observed_ts_post.mean(), "cumulative": observed_ts_post.sum()},
        "predicted": {"average": avg_pred, "cumulative": cum_pred},
        "predicted_lower": {"average": avg_pred_lower, "cumulative": cum_pred_lower},
        "predicted_upper": {"average": avg_pred_upper, "cumulative": cum_pred_upper},
        "predicted_sd": {"average": pred_mean_samples.std(), "cumulative": pred_sum_samples.std()},
        "abs_effect": {"average": avg_effect, "cumulative": cum_effect},
        "abs_effect_lower": {"average": avg_eff_lower, "cumulative": cum_eff_lower},
        "abs_effect_upper": {"average": avg_eff_upper, "cumulative": cum_eff_upper},
        "abs_effect_sd": {"average": avg_eff_samples.std(), "cumulative": cum_eff_samples.std()},
        "rel_effect": {"average": rel_eff_samples.mean(), "cumulative": rel_eff_samples.mean()},
        "rel_effect_lower": {"average": rel_eff_lower, "cumulative": rel_eff_lower},
        "rel_effect_upper": {"average": rel_eff_upper, "cumulative": rel_eff_upper},
        "rel_effect_sd": {"average": rel_eff_samples.std(), "cumulative": rel_eff_samples.std()}
    }
    summary_df = pd.DataFrame(summary_dict)

    observed_cumulative_outcome = observed_ts_post.sum()
    sampled_cumulative_outcomes = pd.concat([pred_sum_samples, pd.Series(observed_cumulative_outcome)], axis=0)
    prop_obs_less = (observed_cumulative_outcome <= sampled_cumulative_outcomes).mean()
    prop_obs_greater = (observed_cumulative_outcome >= sampled_cumulative_outcomes).mean()
    p_value = min(prop_obs_less, prop_obs_greater)
    summary_df["p_value"] = p_value
    summary_df["alpha"] = alpha

    return summary_df
