import numpy as np

from ergodic import dag, tabular, test_graph

# 1. Simulate 1,000 weeks from a known linear graph.
rng = np.random.default_rng(0)
n = 1000
season = rng.normal(size=n)
ads = 0.8 * season + rng.normal(size=n)
visits = 0.9 * ads + 0.3 * season + rng.normal(size=n)
price = rng.normal(size=n)
sales = 0.7 * visits - 0.6 * price + rng.normal(size=n)
data = tabular(
    {"season": season, "ads": ads, "visits": visits, "price": price, "sales": sales}
)
truth = dag(
    [
        "season -> ads",
        "season -> visits",
        "ads -> visits",
        "visits -> sales",
        "price -> sales",
    ]
)

# 2. Test the true graph: one claim per pair without an edge.
result = test_graph(truth, data)
print(result)
print(result.to_frame()[["x", "y", "given", "p_value", "p_holm"]].round(3))

# 3. Drop an edge the data need.
missing = test_graph(truth.remove_edge("season", "visits"), data)
print(missing)
worst = missing.implications[0]
print(f"{worst.x} _||_ {worst.y} | {', '.join(worst.given)}: p_holm = {worst.p_holm:.2g}")

# 4. Reverse season -> ads: a different graph with the same independences.
twin = truth.remove_edge("season", "ads").add_edge("ads -> season").as_kind("dag")
print(test_graph(twin, data))