from edsl import Survey
from edsl.surveys.survey_helpers.survey_mermaid import SurveyMermaidVisualization

survey = Survey([q1, q2, q3])
survey = survey.add_rule("q1", "{{ q1.answer }} == 'skip'", "q3")

# Create visualization
viz = SurveyMermaidVisualization(survey)

# Get Mermaid markdown code
mermaid_code = viz.to_mermaid()
print(mermaid_code)

flowchart TD
    Start(("Start"))
    q1["<b>q1</b><br/>Type: free_text<br/>First question?"]
    q2["<b>q2</b><br/>Type: free_text<br/>Second question?"]
    q3["<b>q3</b><br/>Type: free_text<br/>Third question?"]
    EndOfSurvey(("End"))

    Start --> q1
    q1 --> q2
    q2 --> q3
    q3 --> EndOfSurvey

    %% Skip and navigation rules
    q1 -->|"{{ q1.answer }} == 'skip'"| q3

viz = SurveyMermaidVisualization(
    survey,
    max_text_length=40,      # Truncate question text at 40 chars
    show_options=True,       # Show MC options in nodes
    show_piping=True,        # Show piping dependency arrows
    show_default_flow=True   # Show sequential flow arrows
)

# Just display the visualization object
viz = SurveyMermaidVisualization(survey)
viz  # Renders as HTML with Mermaid diagram

# Or get HTML representation directly
html = viz._repr_html_()

# Show rules in table format
survey.show_rules()

# Access rule collection directly
for rule in survey.rule_collection:
    print(f"Q{rule.current_q} -> Q{rule.next_q}: {rule.expression}")

# Get the DAG
dag = survey.dag

# DAG structure: {child_index: {parent_indices}}
print(dag)
# Example: {2: {0, 1}} means question 2 depends on questions 0 and 1

# Check dependencies
print(survey.question_name_to_index)  # Name to index mapping

```mermaid
flowchart TD
    Start(("Start"))
    q1["Question 1"]
    q2["Question 2"]
    EndOfSurvey(("End"))

    Start --> q1
    q1 --> q2
    q2 --> EndOfSurvey
```

from edsl import Survey, QuestionMultipleChoice, QuestionFreeText
from edsl.surveys.survey_helpers.survey_mermaid import SurveyMermaidVisualization

# Create branching survey
q1 = QuestionMultipleChoice(
    question_name="path",
    question_text="Choose your path",
    question_options=["A", "B"]
)
q2a = QuestionFreeText(question_name="path_a", question_text="Path A question")
q2b = QuestionFreeText(question_name="path_b", question_text="Path B question")
q3 = QuestionFreeText(question_name="conclusion", question_text="Final thoughts?")

survey = (Survey([q1, q2a, q2b, q3])
    .add_rule("path", "{{ path.answer }} == 'A'", "path_a")
    .add_rule("path", "{{ path.answer }} == 'B'", "path_b")
    .add_skip_rule("path_a", "{{ path.answer }} != 'A'")
    .add_skip_rule("path_b", "{{ path.answer }} != 'B'"))

# Generate visualization
viz = SurveyMermaidVisualization(survey)
print(viz.to_mermaid())

# Save to file
mermaid_code = viz.to_mermaid()
with open("survey_flow.md", "w") as f:
    f.write("```mermaid\n")
    f.write(mermaid_code)
    f.write("\n```")

# Or save as standalone .mmd file
with open("survey_flow.mmd", "w") as f:
    f.write(mermaid_code)