import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import networkx as nx

np.random.seed(42)
pd.set_option("display.precision", 4)

def _ancestors(G, nodes):
    """Return all ancestors of a set of nodes in a directed graph G."""
    anc = set()
    stack = list(nodes)
    while stack:
        n = stack.pop()
        for p in G.predecessors(n):
            if p not in anc:
                anc.add(p)
                stack.append(p)
    return anc

def is_d_separated(G, X, Y, Z):
    """
    Check d-separation between sets X and Y given Z in DAG G using
    the moralized ancestral graph method.

    G : networkx.DiGraph
    X, Y, Z : iterables of node labels

    Returns True if X and Y are d-separated given Z, False otherwise.
    """
    X = set(X); Y = set(Y); Z = set(Z)

    # 1. Keep only ancestors of X ∪ Y ∪ Z (plus themselves)
    relevant = set(X) | set(Y) | set(Z)
    relevant |= _ancestors(G, relevant)
    H = G.subgraph(relevant).copy()

    # 2. Moralize: undirected graph, connect parents of each child
    M = nx.Graph()
    M.add_nodes_from(H.nodes())
    # undirected versions of all edges
    for u, v in H.edges():
        M.add_edge(u, v)
    # marry parents
    for child in H.nodes():
        parents = list(H.predecessors(child))
        for i in range(len(parents)):
            for j in range(i + 1, len(parents)):
                M.add_edge(parents[i], parents[j])

    # 3. Remove observed nodes Z
    M.remove_nodes_from(Z)

    # 4. Check if any node in X is connected to any node in Y
    for x in X:
        if x not in M:
            continue
        comp = nx.node_connected_component(M, x)
        if any(y in comp for y in Y):
            return False  # not d-separated (path exists)
    return True  # d-separated

# Quick sanity check on the three basic patterns
if __name__ == "__main__":
    G_chain = nx.DiGraph([("X","M"),("M","Y")])
    print("Chain X,Y | ∅:", is_d_separated(G_chain, {"X"},{"Y"}, set()))       # False
    print("Chain X,Y | M:", is_d_separated(G_chain, {"X"},{"Y"},{"M"}))       # True

    G_fork = nx.DiGraph([("C","X"),("C","Y")])
    print("Fork X,Y | ∅:", is_d_separated(G_fork, {"X"},{"Y"}, set()))        # False
    print("Fork X,Y | C:", is_d_separated(G_fork, {"X"},{"Y"},{"C"}))         # True

    G_col = nx.DiGraph([("X","C"),("Y","C")])
    print("Coll X,Y | ∅:", is_d_separated(G_col, {"X"},{"Y"}, set()))         # True
    print("Coll X,Y | C:", is_d_separated(G_col, {"X"},{"Y"},{"C"}))          # False

# Chain: X → M → Y
chain = nx.DiGraph([
    ("X", "M"),
    ("M", "Y"),
])

print("Are X and Y d-separated with NO conditioning?")
print(is_d_separated(chain, {"X"}, {"Y"}, set()))  # Expect False (they are connected through M)

print("\nAre X and Y d-separated GIVEN M? (conditioning on mediator)")
print(is_d_separated(chain, {"X"}, {"Y"}, {"M"}))  # Expect True (conditioning blocks the chain)

print("\nAre X and Y dependent (no conditioning)?", not is_d_separated(chain, {"X"}, {"Y"}, set()))
print("Are X and Y dependent given M?", not is_d_separated(chain, {"X"}, {"Y"}, {"M"}))

# Fork: C is a common cause of X and Y
fork = nx.DiGraph([
    ("C", "X"),
    ("C", "Y"),
])

print("Are X and Y d-separated with NO conditioning?")
print(is_d_separated(fork, {"X"}, {"Y"}, set()))  # Expect False: they share C

print("\nAre X and Y d-separated GIVEN C?")
print(is_d_separated(fork, {"X"}, {"Y"}, {"C"}))  # Expect True: conditioning on C blocks path

print("\nAre X and Y dependent (no conditioning)?", not is_d_separated(fork, {"X"}, {"Y"}, set()))
print("Are X and Y dependent given C?", not is_d_separated(fork, {"X"}, {"Y"}, {"C"}))

# Collider: C is a common effect of X and Y
collider = nx.DiGraph([
    ("X", "C"),
    ("Y", "C"),
])

print("Are X and Y d-separated with NO conditioning?")
print(is_d_separated(collider, {"X"}, {"Y"}, set()))  # Expect True: collider blocks path

print("\nAre X and Y d-separated GIVEN C?")
print(is_d_separated(collider, {"X"}, {"Y"}, {"C"}))  # Expect False: conditioning opens path

print("\nAre X and Y dependent (no conditioning)?", not is_d_separated(collider, {"X"}, {"Y"}, set()))
print("Are X and Y dependent given C?", not is_d_separated(collider, {"X"}, {"Y"}, {"C"}))

product_model = nx.DiGraph([
    ("SEG", "UI_EXP"),
    ("SEG", "AD"),
    ("H_ISS", "UI_EXP"),
    ("AD", "FAQ"),
    ("SUPPORT", "FAQ"),
    ("UI_EXP", "CHURN"),
])

print("Are SEG and H_ISS d-separated with NO conditioning?")
print(is_d_separated(product_model, {"SEG"}, {"H_ISS"}, set()))  # Expect True

print("\nAre SEG and H_ISS d-separated GIVEN UI_EXP?")
print(is_d_separated(product_model, {"SEG"}, {"H_ISS"}, {"UI_EXP"}))  # Expect False

print("\nAre SEG and H_ISS dependent (no conditioning)?", 
      not is_d_separated(product_model, {"SEG"}, {"H_ISS"}, set()))
print("Are SEG and H_ISS dependent given UI_EXP?",
      not is_d_separated(product_model, {"SEG"}, {"H_ISS"}, {"UI_EXP"}))

print("Are AD and SUPPORT d-separated with NO conditioning?")
print(is_d_separated(product_model, {"AD"}, {"SUPPORT"}, set()))  # Expect True

print("\nAre AD and SUPPORT d-separated GIVEN FAQ?")
print(is_d_separated(product_model, {"AD"}, {"SUPPORT"}, {"FAQ"}))  # Expect False

print("\nAre AD and SUPPORT dependent (no conditioning)?", 
      not is_d_separated(product_model, {"AD"}, {"SUPPORT"}, set()))
print("Are AD and SUPPORT dependent given FAQ?",
      not is_d_separated(product_model, {"AD"}, {"SUPPORT"}, {"FAQ"}))

print("Are H_ISS and CHURN d-separated with NO conditioning?")
print(is_d_separated(product_model, {"H_ISS"}, {"CHURN"}, set()))  # Expect False

print("\nAre H_ISS and CHURN d-separated GIVEN UI_EXP?")
print(is_d_separated(product_model, {"H_ISS"}, {"CHURN"}, {"UI_EXP"}))  # Expect True

print("\nAre H_ISS and CHURN dependent (no conditioning)?", 
      not is_d_separated(product_model, {"H_ISS"}, {"CHURN"}, set()))
print("Are H_ISS and CHURN dependent given UI_EXP?",
      not is_d_separated(product_model, {"H_ISS"}, {"CHURN"}, {"UI_EXP"}))

marketing_model = nx.DiGraph([
    ("revenue_prev", "email_campaign"),
    ("revenue_prev", "revenue_next"),
    ("email_campaign", "revenue_next"),
])

# To see if there is a noncausal backdoor path from email_campaign to revenue_next,
# we conceptually remove the direct arrow email_campaign → revenue_next
marketing_severed = nx.DiGraph([
    ("revenue_prev", "email_campaign"),
    ("revenue_prev", "revenue_next"),
])

print("Are email_campaign and revenue_next d-separated in the severed graph?")
print(is_d_separated(marketing_severed, {"email_campaign"}, {"revenue_next"}, set()))

print("\nAre they still dependent (backdoor path exists)?",
      not is_d_separated(marketing_severed, {"email_campaign"}, {"revenue_next"}, set()))

df = pd.DataFrame({
    "revenue_prev":   [200, 200, 200, 800, 800, 800],
    "email_campaign": [0,   0,   1,   0,   1,   1],
    "revenue_next":   [210, 220, 245, 820, 860, 870],
})
df

naive_effect = (df.loc[df.email_campaign == 1, "revenue_next"].mean() -
                 df.loc[df.email_campaign == 0, "revenue_next"].mean())
print("Naive difference (campaign - no campaign):", naive_effect)

avg = (df
        .groupby(["email_campaign", "revenue_prev"])
        ["revenue_next"].mean())

effects_by_stratum = avg.loc[1] - avg.loc[0]
print("Effect by revenue_prev group:")
print(effects_by_stratum)

ate = effects_by_stratum.mean()
print("\nAdjusted ATE (average across strata):", ate)

n = 20000
feature_flag = np.random.binomial(1, 0.5, size=n)

# Potential outcomes
rating0 = np.random.normal(loc=0.0, scale=1.0, size=n)  # baseline
true_effect = 0.3
rating1 = rating0 + true_effect
rating_latent = np.where(feature_flag == 1, rating1, rating0)

df_sel = pd.DataFrame({
    "feature_flag": feature_flag,
    "rating0": rating0,
    "rating1": rating1,
    "rating_latent": rating_latent,
})

def sigmoid(x):
    return 1 / (1 + np.exp(-x))

alpha = -2.0   # baseline response log-odds
beta_t = 0.9   # treatment effect on response
beta_y = 0.7   # satisfaction effect on response

response_prob = sigmoid(alpha + beta_t * df_sel["feature_flag"] + beta_y * df_sel["rating_latent"])
df_sel["responded"] = np.random.binomial(1, response_prob, size=n)
df_sel["rating_observed"] = np.where(df_sel["responded"] == 1, df_sel["rating_latent"], np.nan)

df_sel.head()

# True ATE from potential outcomes
true_ate = (df_sel["rating1"] - df_sel["rating0"]).mean()

# ATE using full data (no selection)
full_ate = (df_sel.loc[df_sel.feature_flag == 1, "rating_latent"].mean() -
            df_sel.loc[df_sel.feature_flag == 0, "rating_latent"].mean())

# ATE using only responders
responders = df_sel[df_sel["responded"] == 1]
resp_ate = (responders.loc[responders.feature_flag == 1, "rating_observed"].mean() -
            responders.loc[responders.feature_flag == 0, "rating_observed"].mean())

print(f"True ATE (by design): {true_ate:.3f}")
print(f"ATE using ALL users: {full_ate:.3f}")
print(f"Naive ATE among RESPONDERS only: {resp_ate:.3f}")

baseline_by_group = (responders
                      .groupby("feature_flag")["rating0"]
                      .mean())
baseline_by_group

plt.figure(figsize=(8,4))
plt.hist(responders.loc[responders.feature_flag == 0, "rating0"], bins=30, alpha=0.7, label="Control")
plt.hist(responders.loc[responders.feature_flag == 1, "rating0"], bins=30, alpha=0.7, label="Treatment")
plt.xlabel("Baseline rating (rating0)")
plt.ylabel("Count among responders")
plt.title("Baseline satisfaction among responders by feature flag")
plt.legend()
plt.show()