import numpy as np
import pandas as pd
from ergodic import conformance, decision_points, decision_table, discover, discover_model
from ergodic import estimate_effect, event_log
from ergodic.process import filter_variants, is_sound, seq, xor

# 1. Simulate 2,000 cases; timestamps are in hours. Complex cases are expedited
#    more often, expediting cuts 5 hours, and some standard cases escalate.
rng = np.random.default_rng(7)
n = 2000
complexity = rng.uniform(0, 1, n)
expedite = rng.uniform(size=n) < 1 / (1 + np.exp(-3 * (complexity - 0.5)))
remaining = 10 + 8 * complexity - 5 * expedite + rng.normal(0, 1, n)
escalate = ~expedite & (rng.uniform(size=n) < 0.03)
opened = rng.uniform(0, 90 * 24, n)
rows = []
for i in range(n):
    t = opened[i]
    decision = t + 1.0 + rng.uniform(0.5, 2.0)
    branch = "expedite" if expedite[i] else "standard"
    steps = [("open", t), ("triage", t + 1.0), (branch, decision)]
    steps += [("expedite", decision + 1.0)] * int(escalate[i])
    steps.append(("resolve", decision + remaining[i]))
    rows += [(f"C{i:04d}", activity, at, complexity[i]) for activity, at in steps]
frame = pd.DataFrame(rows, columns=["case", "activity", "hour", "complexity"])
log = event_log(frame, case="case", activity="activity", timestamp="hour")

# 2. Mine a model from the whole log, then from its two most frequent variants.
print(discover_model(log).tree)
model = discover_model(filter_variants(log, top=2))
print(model.tree)
print(model.tree == seq("open", "triage", xor("expedite", "standard"), "resolve"))
print(model.net, "sound:", is_sound(model.net))

# 3. Check the whole log against the model.
fit = conformance(log, model)
print(fit)
for variant in fit.variants:
    print(variant.n_cases, f"{variant.alignment_fitness:.3f}", variant.alignment)

# 4. Find the decision, and encode it as one row per case.
print(decision_points(log, model))
table = decision_table(
    log, model, covariates=["complexity", "elapsed_at_landmark", "workload_at_landmark"]
)
print(table)
positivity = table.positivity
print([round(share, 3) for share in positivity.shares], positivity.overlap_share, positivity.flagged)
by_case = table.data.to_pandas()
print(by_case.groupby("branch", observed=True)["remaining_time"].mean().round(2))

# 5. Discover the graph and estimate the effect of expediting.
graph = discover(table.data, knowledge=table.knowledge, alpha=0.01).graph
estimate = estimate_effect(
    graph.as_kind("dag"), table.data, "branch", "remaining_time",
    strategy="dml", treated="expedite", random_state=0,
)
print(estimate)