import numpy as np
from ergodic import recover, timeseries
from ergodic.latent.likelihoods import spectral_whittle

# 1. Simulate 1,440 ticks of a sensor. A switch nobody logged is on from
#    tick 500 to tick 999, and while it is on the noise is three times larger.
rng = np.random.default_rng(0)
switch = np.zeros(1440, dtype=int)
switch[500:1000] = 1
signal = 100 + rng.normal(0, 0.1, 1440) * np.where(switch == 1, 3.0, 1.0)
data = timeseries({"tick": np.arange(1440), "signal": signal}, time="tick")

# 2. Recover the switch with the spectral likelihood.
result = recover(
    data, "signal", model=spectral_whittle(), draws=400, warmup=400, chains=2, seed=0
)
print(result)
print(result.chain)

# 3. Which mechanism moved: the posterior delta per block, then the verdicts.
for block, (low, mean, high) in result.deltas().items():
    print(f"{block:<7}{mean:+.2f}  [{low:+.2f}, {high:+.2f}]")
print(result.verdicts())

# 4. When the switch was on: one state per 96-tick frame.
print(result.path())
ticks = result.tick_path()
agreement = np.mean(ticks == switch[: ticks.size])
print(f"agreement with the planted switch: {max(agreement, 1 - agreement):.2f}")
