import numpy as np
import pandas as pd

from ergodic import discover, tabular
from ergodic.discovery import GCM, Linear

# 1. A linear world with uniform noise: a -> b -> c -> d and a -> d.
rng = np.random.default_rng(0)
n = 2_000
a = rng.uniform(-1, 1, n)
b = 0.8 * a + rng.uniform(-1, 1, n)
c = 0.7 * b + rng.uniform(-1, 1, n)
d = 0.5 * a + 0.6 * c + rng.uniform(-1, 1, n)
linear_world = tabular(pd.DataFrame({"a": a, "b": b, "c": c, "d": d}))

# 2. PC returns the equivalence class; LiNGAM uses the non-Gaussian noise.
print("pc:    ", discover(linear_world, method="pc").graph)
print("lingam:", discover(linear_world, method="lingam").graph)

# 3. A nonlinear chain: dose -> level -> response.
rng = np.random.default_rng(0)
dose = rng.uniform(0, 3, 1_000)
level = np.exp(dose) + rng.normal(0, 1, 1_000)
response = 10 * level / (level + 5) + rng.normal(0, 0.5, 1_000)
chain = tabular(pd.DataFrame({"dose": dose, "level": level, "response": response}))

# 4. Is dose independent of response given level? It is, in truth.
for test in [Linear(chain), GCM(chain, learner="knn")]:
    p = test.test("dose", "response", ["level"]).p_value
    print(f"{test.name}: p = {p:.3f}")

# 5. Run PC with each test.
print("linear:", discover(chain, method="pc", test="linear").graph)
print("gcm:   ", discover(chain, method="pc", test=GCM(learner="knn")).graph)