import graphviz

# Build the decision-tree flow chart
dot = graphviz.Digraph("z_vs_t_decision_tree", format="png")
dot.attr(rankdir="TB", fontsize="10", fontname="Helvetica")

# Nodes
dot.node("start", "Start", shape="ellipse", style="filled", fillcolor="#e0f7fa")

dot.node("samples", "Number of samples?", shape="diamond", style="filled", fillcolor="#fff9c4")

# One-sample branch
dot.node("one_n", "n ≥ 30 AND σ known?", shape="diamond", style="filled", fillcolor="#ffe0b2")
dot.node("one_z", "1-Sample Z", shape="box", style="filled", fillcolor="#c8e6c9")
dot.node("one_t", "1-Sample t", shape="box", style="filled", fillcolor="#c5cae9")

# Two-sample branch
dot.node("two_n", "Both n ≥ 30?", shape="diamond", style="filled", fillcolor="#ffe0b2")
dot.node("two_sigma", "σ₁ and σ₂ known?", shape="diamond", style="filled", fillcolor="#ffe0b2")
dot.node("two_z", "2-Sample Z", shape="box", style="filled", fillcolor="#c8e6c9")

dot.node("equal_var", "Equal variances\nplausible?", shape="diamond", style="filled", fillcolor="#ffe0b2")
dot.node("pooled_t", "Pooled t\n(df = n₁ + n₂ − 2)", shape="box", style="filled", fillcolor="#c5cae9")
dot.node("welch_t", "Welch t\n(df*)", shape="box", style="filled", fillcolor="#c5cae9")

# Edges
dot.edge("start", "samples")
dot.edge("samples", "one_n", label="1")
dot.edge("samples", "two_n", label="2")

# One-sample path
dot.edge("one_n", "one_z", label="Yes")
dot.edge("one_n", "one_t", label="No")

# Two-sample path
dot.edge("two_n", "two_sigma", label="Yes")
dot.edge("two_n", "equal_var", label="No")

dot.edge("two_sigma", "two_z", label="Yes")
dot.edge("two_sigma", "equal_var", label="No")

dot.edge("equal_var", "pooled_t", label="Yes")
dot.edge("equal_var", "welch_t", label="No")

# Render to file
file_path = "/mnt/data/decision_tree_z_t.png"
dot.render(filename="/mnt/data/decision_tree_z_t", cleanup=True)

file_path