diff --git a/src/diagram.typ b/src/diagram.typ
index 66480f6..67d97c0 100644
--- a/src/diagram.typ
+++ b/src/diagram.typ
@@ -51,20 +51,37 @@
   return (total-depth, nodes)
 }
 
+#let FLOW-COLOR = (
+  "source",
+  "target",
+  "gradient"
+)
+
 #let diagram(
   elmts,
   height,
-  min-spacing: 1cm
+  width: auto,
+  h-space: 2cm,
+  min-spacing: 1cm,
+  node-width: 0.5cm,
+  flow-color: "source",
+  roundiness: 50%,
+  align-sources: false,
+  align-sinks: false
 ) = {
   let columns = ()
   let nodes = (:)
   let flows = ()
 
+  let order = ()
   for elmt in elmts {
     if elmt.type == "node" {
+      order.push(elmt.name)
       nodes.insert(elmt.name, elmt)
     } else if elmt.type == "flow" {
       flows.push(elmt)
+      order.push(elmt.from)
+      order.push(elmt.to)
     }
   }
 
@@ -100,23 +117,23 @@
 
   columns = ((),) * (total-depth + 1)
   let column-totals = (0,) * columns.len()
-  for (key, node) in nodes.pairs() {
-    columns.at(node.max-depth).push(key)
-    column-totals.at(node.max-depth) += node.total.max
+  let pairs = nodes.pairs()
+  // Order nodes by reference order
+  pairs = pairs.sorted(key: p => {
+    order.position(n => n == p.first())
+  })
+  for (key, node) in pairs {
+    let col = node.max-depth
+    if not node.is-dst and align-sources {
+      col = 0
+    }
+    if not node.is-src and align-sinks {
+      col = columns.len() - 1
+    }
+    columns.at(col).push(key)
+    column-totals.at(col) += node.total.max
   }
 
-  //total-depth
-  [
-    *nodes*: #nodes
-
-    *total-depth*: #total-depth
-
-    *columns*: #columns
-
-    *column-totals*: #column-totals
-  ]
-  //flows
-
   let h-factor = calc.min(
     ..columns.zip(column-totals).map(((col, total)) => {
       let available-h = height - (col.len() - 1) * min-spacing
@@ -124,12 +141,21 @@
     })
   )
 
+  let get-color(nid) = {
+    let r = nodes.keys().position(n => n == nid) / nodes.len()
+    return color.hsv(r * 270deg, 60%, 90%)
+  }
   let colors = (red, orange, yellow, green, blue, purple).map(c => c.lighten(50%))
 
+  let h-space = h-space
+  if width != auto and columns.len() > 1 {
+    let remaining-w = width - columns.len() * node-width
+    h-space = remaining-w / (columns.len() - 1)
+  }
+
   canvas({
-    let c = 0
+    let x = 0pt
     for (i, column) in columns.enumerate() {
-      let x = i * 3
       let total = column.map(n => nodes.at(n).total.max).sum()
       let total-h = total * h-factor
       let remaining-h = height - total-h
@@ -142,39 +168,111 @@
         let h = h-factor * node.total.max
         draw.rect(
           (x, y),
-          (x + 1, y - h),
+          (x + node-width, y - h),
           name: nid,
-          fill: colors.at(c)
+          fill: get-color(nid),
+          stroke: none
         )
-        c += 1
 
+        let pos = nid
+        let algn = center
+        let txt = [
+          #node.display-name\
+          #node.total.max
+        ]
+        let args = (
+          anchor: "center"
+        )
         if node.is-src and not node.is-dst {
-          draw.content(
-            nid + ".west",
-            anchor: "east",
-            node.display-name,
-            padding: 5pt
-          )
+          pos = nid + ".west"
+          args.anchor = "east"
+          algn = right
         } else if node.is-src and node.is-dst {
-          draw.content(
-            nid + ".center",
-            anchor: "center",
-            node.display-name,
-            padding: 5pt,
-            frame: "rect",
-            fill: white.transparentize(20%)
-          )
+          pos = nid + ".center"
+          args.anchor = "center"
+          algn = center
+          args.insert("frame", "rect")
+          args.insert("fill", white.transparentize(50%))
+          args.insert("stroke", none)
         } else if not node.is-src and node.is-dst {
-          draw.content(
-            nid + ".east",
-            anchor: "west",
-            node.display-name,
-            padding: 5pt
-          )
+          pos = nid + ".east"
+          args.anchor = "west"
+          algn = left
         }
+        draw.content(
+          pos,
+          ..args,
+          align(algn, txt),
+          padding: 5pt
+        )
 
         y -= spacing + h
       }
+
+      x += node-width + h-space
     }
+
+    let offsets = (:)
+    for nid in nodes.keys() {
+      offsets.insert(nid, (
+        "in": 0pt,
+        "out": 0pt
+      ))
+    }
+
+    draw.on-layer(-1, {
+
+      for flow in flows {
+        let tl = flow.from + ".north-east"
+        let tr = flow.to + ".north-west"
+        let off-out = offsets.at(flow.from).out
+        let off-in = offsets.at(flow.to).in
+        let h = flow.amount * h-factor
+
+        offsets.at(flow.from).out += h
+        offsets.at(flow.to).in += h
+
+        tl = (rel: (0, -off-out), to: tl)
+        tr = (rel: (0, -off-in), to: tr)
+        let br = (rel: (0, -h), to: tr)
+        let bl = (rel: (0, -h), to: tl)
+
+        let fill = gray.lighten(50%).transparentize(50%)
+
+        let col1 = get-color(flow.from).desaturate(50%).lighten(50%).transparentize(50%)
+        let col2 = get-color(flow.to).desaturate(50%).lighten(50%).transparentize(50%)
+        
+        if flow-color == "source" {
+          fill = col1
+        } else if flow-color == "target" {
+          fill = col2
+        } else if flow-color == "gradient" {
+          fill = gradient.linear(col1, col2)
+        } else {
+          fill = flow-color
+        }
+
+        let dx = roundiness * h-space
+        draw.merge-path(
+          fill: fill,
+          stroke: none,
+          {
+            draw.bezier(
+              tl,
+              tr,
+              (rel: (dx, 0), to: tl),
+              (rel: (-dx, 0), to: tr),
+            )
+            draw.bezier(
+              br,
+              bl,
+              (rel: (-dx, 0), to: br),
+              (rel: (dx, 0), to: bl),
+            )
+          }
+        )
+      }
+    })
   })
+
 }
\ No newline at end of file