diff --git a/_scripts/generate.py b/_scripts/generate.py
index 0259855..1a71cab 100644
--- a/_scripts/generate.py
+++ b/_scripts/generate.py
@@ -9,7 +9,6 @@ from datetime import datetime
 
 from org_rw import OrgTime, dom
 from org_rw import dump as dump_org
-from org_rw import dumps as dumps_org
 from org_rw import load as load_org
 from org_rw import token_list_to_raw
 
@@ -53,15 +52,21 @@ def main(src_top, dest_top):
     for doc in docs:
         relpath = os.path.relpath(doc.path, src_top)
         changed = False
-        headlines = list(doc.getTopHeadlines())
+        headlines = list(doc.getAllHeadlines())
         related = None
 
         i = len(headlines)
         while i > 0:
             i -= 1
             headline = headlines[i]
-            if headline.title.strip().lower() == "related":
-                assert related is None
+            if headline.title.strip().lower() == "related" and headline.depth == 1:
+                if related is not None:
+                    print(
+                        "Found duplicated related: {} vs {}".format(
+                            related.id, headline.id
+                        )
+                    )
+                    assert related is None
                 related = headline
                 headlines.pop(i)
 
@@ -160,9 +165,19 @@ def render(headline, doc):
     content = []
     render_tree(dom, content)
 
+    if headline.state is None:
+        state = ""
+    else:
+        state = f'<span class="state todo-{headline.is_todo} state-{headline.state}">{headline.state}</span>'
+
+    if headline.is_todo:
+        todo_state = "todo"
+    else:
+        todo_state = "done"
     return f"""
-<article id="{html.escape(headline.id)}">
+<article id="{html.escape(headline.id)}" class="{todo_state}">
   <h1 class="title">
+    {state}
     <a href=\"org-protocol://org-id?id={html.escape(headline.id)}\">
       {html.escape(headline.title)}
     </a>
@@ -176,7 +191,7 @@ def render(headline, doc):
 def save_changes(doc):
     assert doc.path is not None
     with open(doc.path, "wt") as f:
-        dumps_org(doc, f)
+        dump_org(doc, f)
 
 
 if __name__ == "__main__":