# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
# ---

# %% [markdown]
# # US COVID-19 Mortality: Panel Structural Break Analysis
#
# Panel structural break detection with multiple entities (US states) sharing
# common break dates. The framework tests whether all entities experienced
# regime changes at the same time points.
#
# Synthetic data mimics state-level COVID mortality rates, with a break at
# the onset of widespread vaccination (early 2021).
#
# Reference: Ditzen, Karavias & Westerlund (2024), arXiv:2110.14550

# %%
import numpy as np
from xtbreak import test, estimate

# %% [markdown]
# ## Synthetic panel dataset
#
# N=8 states over T=52 weeks, with a common break at week 30
# (vaccination inflection point).

# %%
np.random.seed(42)

N = 8       # Number of states (entities)
T = 52      # Number of time periods (weeks)
break_t = 30  # True common break location

# Generate entity and time indices (stacked panel format)
entity = np.repeat(np.arange(N), T)
time = np.tile(np.arange(1, T + 1), N)

# State-specific intercepts (fixed effects)
state_effects = np.random.uniform(5.0, 15.0, N)

# Generate dependent variable: mortality rate per 100k
# We use a constant-only breaking regressor to clearly detect the mean shift
y = np.zeros(N * T)
x_regressor = np.ones((N * T, 1))  # Constant regressor (tests mean shift)

for i in range(N):
    start = i * T
    end = (i + 1) * T
    # Pre-break: high mortality regime
    y[start:start + break_t] = (
        state_effects[i] + 12.0 + np.random.randn(break_t) * 2.0
    )
    # Post-break: lower mortality (vaccination effect)
    y[start + break_t:end] = (
        state_effects[i] + 5.0 + np.random.randn(T - break_t) * 2.0
    )

print(f"Panel dimensions: N={N} entities, T={T} periods")
print(f"Total observations: {len(y)}")
print(f"True break location: t={break_t}")
print()

# %% [markdown]
# ## Step 1: Test for a common structural break
#
# The panel SupF test examines H0 (no break) vs H1 (1 common break)
# across all entities.

# %%
test_result = test(
    y, x_regressor,
    breaks=1,
    trimming=0.15,
    entity=entity,
    time=time,
    fixed_effects=True,
    vce="ssr",
)
print("=" * 55)
print("Panel Structural Break Test")
print("=" * 55)
print(f"  Entities:        {N} states")
print(f"  Time periods:    {T} weeks")
print(f"  Hypothesis:      H0 (0 breaks) vs H1 (1 common break)")
print(f"  SupF statistic:  {test_result.statistic:.4f}")
print(f"  Break indices:   {test_result.breaks}")
print()

# %% [markdown]
# ## Step 2: Estimate the common break date
#
# The estimator finds the break location that minimizes the global SSR
# across all panel entities simultaneously.

# %%
est_result = estimate(
    y, x_regressor,
    breaks=1,
    trimming=0.15,
    entity=entity,
    time=time,
    fixed_effects=True,
    vce="ssr",
)

break_est = est_result.break_estimates[0]
print("=" * 55)
print("Panel Break Estimation")
print("=" * 55)
print(f"  Estimated break week:  {break_est.time_value}")
print(f"  True break week:       {break_t}")
print(f"  Break index:           {break_est.index}")
print(f"  95% CI (time):         [{break_est.ci_time_value[0]}, "
      f"{break_est.ci_time_value[1]}]")
print(f"  Global SSR:            {est_result.ssr:.4f}")
print()

# %% [markdown]
# ## Step 3: Regime statistics

# %%
print("=" * 55)
print("Regime Statistics (Panel)")
print("=" * 55)
for regime in est_result.regime_statistics:
    print(f"  Regime {regime.regime_index}:")
    print(f"    Periods:      {regime.start_index}–{regime.end_index}")
    print(f"    Observations: {regime.n_observations} (across all entities)")
    print(f"    Coefficient:  {regime.coefficients[0]:.4f}")
    print(f"    Std. error:   {regime.standard_errors[0]:.4f}")
    print(f"    R²:           {regime.r_squared:.6f}")
    print()

# %% [markdown]
# ## Step 4: Summary

# %%
print(est_result.summary())

# %% [markdown]
# ## Notes
#
# - **Common break dates**: The panel framework assumes all entities share
#   the same break timing, while allowing entity-specific fixed effects.
#
# - **Pooling across entities**: Information from all N entities is used
#   jointly, yielding higher power and narrower confidence intervals than
#   separate time-series analyses.
#
# - **Fixed effects**: Entity-specific intercepts are absorbed when
#   `fixed_effects=True`.
#
# - **Stacked format**: Data is in long format with explicit `entity`
#   and `time` index arrays.