import numpy as np
import pandas as pd

from ergodic import DomainKnowledge, discover, granger
from ergodic.data import timeseries

# 1. Simulate 500 weeks: ads -> visits at lag 1, visits -> sales at lag 2.
rng = np.random.default_rng(0)
weeks = 500
ads, visits, sales = np.zeros(weeks), np.zeros(weeks), np.zeros(weeks)
noise = rng.normal(size=(weeks, 3))
for t in range(2, weeks):
    ads[t] = 0.6 * ads[t - 1] + noise[t, 0]
    visits[t] = 0.5 * visits[t - 1] + 0.5 * ads[t - 1] + noise[t, 1]
    sales[t] = 0.5 * sales[t - 1] + 0.4 * visits[t - 2] + noise[t, 2]
frame = pd.DataFrame({"week": np.arange(weeks), "ads": ads, "visits": visits, "sales": sales})
data = timeseries(frame, time="week")

# 2. Run PCMCI with links up to two weeks back.
result = discover(data, method="pcmci", max_lag=2)
print(result)

# 3. Each link is (source, target, lag), with the p-value that kept it.
for (source, target, lag), p in sorted(result.links.items()):
    print(f"{source}[t-{lag}] -> {target}[t]  p={p:.2g}")

# 4. Granger tests answer a narrower, predictive question.
print("granger:", sorted(granger(data, max_lag=2).significant))

# 5. Add the order of the funnel: ads, then visits, then sales.
knowledge = DomainKnowledge().with_tiers([["ads"], ["visits"], ["sales"]])
informed = discover(data, method="pcmci", max_lag=2, knowledge=knowledge)
print(sorted(informed.links) == sorted(result.links))
pruned = int(informed.diagnostics["n_knowledge_pruned"])
tests = int(informed.diagnostics["n_ci_tests"]), int(result.diagnostics["n_ci_tests"])
print(f"{pruned} links pruned; {tests[0]} tests instead of {tests[1]}")