diff --git a/web/src/App.tsx b/web/src/App.tsx
index abeb6c5..99745dc 100644
--- a/web/src/App.tsx
+++ b/web/src/App.tsx
@@ -4,13 +4,16 @@ import Navbar from "./components/Navbar";
 import ExpressionsPopup from "./components/ExpressionsPopup";
 import { RiQuestionFill } from "react-icons/ri";
 import { RegexMatch, RegexEngine, RegexCapture } from "regex-potata";
-import { dotFromRegex } from "./utils/viz";
+import { graphFromRegex } from "./utils/viz";
 import TestInput from "./components/TestInput";
 import Footer from "./components/Footer";
 import RegexInput from "./components/RegexInput";
 import ToolTip from "./components/ToolTip";
+import NfaVisualizer from "./components/NfaVisualizer";
+import Loader from "./components/Loader";
 
 const App = () => {
+  const [isLoading, setIsLoading] = useState(true);
   const [regexInput, setRegexInput] = useState("");
   const [testInput, setTestInput] = useState("");
   const [regexInstance, setRegexInstance] = useState<RegexEngine>();
@@ -26,6 +29,7 @@ const App = () => {
       const engine = new RegexEngine("");
       vizInstance.current = i;
       setRegexInstance(engine);
+      setTimeout(() => setIsLoading(false), 250);
     })();
   }, []);
 
@@ -39,7 +43,7 @@ const App = () => {
 
   useEffect(() => {
     if (regexInstance) {
-      const dot = dotFromRegex(regexInstance);
+      const dot = graphFromRegex(regexInstance);
       const elem = vizInstance.current?.renderSVGElement(dot);
 
       if (elem) {
@@ -55,6 +59,10 @@ const App = () => {
     }
   }, [testInput, regexInstance]);
 
+  if (isLoading) {
+    return <Loader />;
+  }
+
   return (
     <div
       className="min-h-screen min-w-screen 
@@ -91,18 +99,7 @@ const App = () => {
           </div>
           <div className="space-y-10">
             <div className="font-semibold">NFA Visualizer</div>
-            <div
-              className="pt-12 pb-8 w-full overflow-scroll
-              rounded-md border-[1px] border-slate-800"
-            >
-              {svg && (
-                <svg
-                  height={svg?.height.baseVal.value}
-                  width={svg?.width.baseVal.value}
-                  dangerouslySetInnerHTML={{ __html: svg.innerHTML }}
-                ></svg>
-              )}
-            </div>
+            <NfaVisualizer svg={svg} />
           </div>
         </div>
       </div>
@@ -111,6 +108,12 @@ const App = () => {
         open={isPopupOpen}
         onClose={() => setIsPopupOpen(false)}
       />
+
+      {/* Transition mask */}
+      <div
+        className="fixed w-screen h-screen z-50 opacity-0
+        animate-fade pointer-events-none bg-slate-900"
+      ></div>
     </div>
   );
 };
diff --git a/web/src/components/Loader.tsx b/web/src/components/Loader.tsx
new file mode 100644
index 0000000..d1d9514
--- /dev/null
+++ b/web/src/components/Loader.tsx
@@ -0,0 +1,18 @@
+const Loader = () => (
+  <div
+    className="w-screen h-screen bg-slate-900
+    flex items-center justify-center"
+  >
+    <div
+      className="absolute right-1/2 bottom-1/2
+      transform translate-x-1/2 translate-y-1/2 "
+    >
+      <div
+        className="border-solid border-cyan-300 border-4
+        border-t-transparent animate-spin rounded-full h-20 w-20"
+      ></div>
+    </div>
+  </div>
+);
+
+export default Loader;
diff --git a/web/src/components/NfaVisualizer.tsx b/web/src/components/NfaVisualizer.tsx
new file mode 100644
index 0000000..92c4ae5
--- /dev/null
+++ b/web/src/components/NfaVisualizer.tsx
@@ -0,0 +1,21 @@
+import { useEffect, useRef } from "react";
+
+const NfaVisualizer = ({ svg }: { svg?: SVGSVGElement }) => {
+  const containerRef = useRef<HTMLDivElement>(null);
+
+  useEffect(() => {
+    if (svg) {
+      containerRef.current?.replaceChildren(svg);
+    }
+  }, [svg]);
+
+  return (
+    <div
+      ref={containerRef}
+      className="pt-12 pb-8 w-full overflow-scroll
+      rounded-md border-[1px] border-slate-800"
+    ></div>
+  );
+};
+
+export default NfaVisualizer;
diff --git a/web/src/components/TestInput.tsx b/web/src/components/TestInput.tsx
index 46e3575..8029f73 100644
--- a/web/src/components/TestInput.tsx
+++ b/web/src/components/TestInput.tsx
@@ -16,7 +16,7 @@ const TestInput = ({ input, matches, captures, onInput }: InputProps) => {
 
   useEffect(() => {
     if (matches.length) {
-      setHighlightExtension(getMatchHighlight(matches));
+      setHighlightExtension(getMatchHighlight(matches, captures));
     } else {
       setHighlightExtension(undefined);
     }
@@ -33,12 +33,12 @@ const TestInput = ({ input, matches, captures, onInput }: InputProps) => {
   return (
     <ReactCodeMirror
       value={input}
-      height="200px"
-      className="leading-10"
       basicSetup={{
         lineNumbers: false,
         foldGutter: false,
         highlightActiveLine: false,
+        bracketMatching: false,
+        closeBrackets: false,
       }}
       extensions={[hoverExtension!, highlightExtension!].filter((v) => v)}
       onChange={(value) => onInput(value)}
diff --git a/web/src/index.css b/web/src/index.css
index 1f016a4..4a288a6 100644
--- a/web/src/index.css
+++ b/web/src/index.css
@@ -3,11 +3,11 @@
 @tailwind utilities;
 
 .cm-editor {
-  @apply border-[1px] border-slate-800 rounded-md p-3 !bg-transparent;
+  @apply h-52 border-[1px] border-slate-800 rounded-md p-3 !bg-transparent;
 }
 
 .cm-content {
-  @apply !font-sans !leading-6;
+  @apply !font-sans !leading-8;
 }
 
 .cm-focused {
diff --git a/web/src/utils/extensions.ts b/web/src/utils/extensions.ts
index 7ac897f..ef9ffd6 100644
--- a/web/src/utils/extensions.ts
+++ b/web/src/utils/extensions.ts
@@ -12,13 +12,23 @@ const matchDecoration = Decoration.mark({
   inclusiveEnd: false,
 });
 
-function getMatchHighlight(matches: RegexMatch[]) {
+const groupDecoration = Decoration.mark({
+  class: "underline underline-offset-8 decoration-2 decoration-cyan-300",
+  inclusiveStart: true,
+  inclusiveEnd: false,
+});
+
+function getMatchHighlight(matches: RegexMatch[], captures: RegexCapture[]) {
   const decorationBuilder = new RangeSetBuilder<Decoration>();
 
   for (const match of matches) {
     decorationBuilder.add(match.start, match.end, matchDecoration);
   }
 
+  for (const capture of captures.slice(1)) {
+    decorationBuilder.add(capture.start, capture.end, groupDecoration);
+  }
+
   const plugin = ViewPlugin.define(
     () => ({
       decorations: decorationBuilder.finish(),
@@ -42,7 +52,7 @@ function groupHoverTooltip(captures: RegexCapture[]) {
           create() {
             const dom = document.createElement("div");
             dom.innerHTML = `Group <span class="font-semibold">${capture.name()}</span>`;
-            dom.classList.add("px-4", "rounded-md", "!bg-slate-600");
+            dom.classList.add("py-2", "px-4", "rounded-md", "!bg-slate-600");
             return { dom };
           },
         };
diff --git a/web/src/utils/viz.ts b/web/src/utils/viz.ts
index 5a25746..8b707df 100644
--- a/web/src/utils/viz.ts
+++ b/web/src/utils/viz.ts
@@ -1,28 +1,53 @@
+import { Graph } from "@viz-js/viz";
 import { RegexEngine } from "regex-potata";
 
-function dotFromRegex(regex: RegexEngine) {
+function graphFromRegex(regex: RegexEngine) {
   const states = regex.nfaStates();
   const endState = states.length - 1;
-  const transitions = Array.from(states)
-    .flatMap((state) =>
-      regex
-        .nfaTransition(state)
-        ?.map((t) => `${state} -> ${t.end} [label="${t.toString()}"]\n`)
-    )
-    .join("\n");
-  const dot = `
-    digraph { 
-      bgcolor=none; 
-      graph [rankdir=LR]; 
-      node [shape=circle, color=white, penwidth=2, fontcolor=white, fontname="Arial"];
-      edge [color="#67e8f9", fontcolor=white, fontname="Arial"];
-      ${endState} [shape=doublecircle, color="#67e8f9"];
-      "" [shape=none]
-      "" -> 0
-      ${transitions}
-    }`;
 
-  return dot;
+  const config: Graph = {
+    graphAttributes: {
+      bgcolor: "none",
+      rankdir: "LR",
+    },
+    nodeAttributes: {
+      shape: "circle",
+      color: "white",
+      penwidth: 2,
+      fontcolor: "white",
+      fontname: "Arial",
+    },
+    edgeAttributes: {
+      color: "#67e8f9",
+      fontcolor: "white",
+      fontname: "Arial",
+    },
+    nodes: [
+      { name: "", attributes: { shape: "none" } },
+      {
+        name: endState.toString(),
+        attributes: { shape: "doublecircle", color: "#67e8f9" },
+      },
+    ],
+    edges: [{ tail: "", head: "0" }],
+    subgraphs: [],
+  };
+
+  for (const state of states) {
+    const transitions = regex.nfaTransition(state);
+
+    if (transitions) {
+      for (const transition of transitions) {
+        config.edges!.push({
+          tail: state.toString(),
+          head: transition.end.toString(),
+          attributes: { label: transition.toString() },
+        });
+      }
+    }
+  }
+
+  return config;
 }
 
-export { dotFromRegex };
+export { graphFromRegex };
diff --git a/web/tailwind.config.js b/web/tailwind.config.js
index 8679225..5e5add2 100644
--- a/web/tailwind.config.js
+++ b/web/tailwind.config.js
@@ -5,7 +5,21 @@ export default {
   },
   content: ["./index.html", "./src/**/*.{js,ts,jsx,tsx}"],
   theme: {
-    extend: {},
+    extend: {
+      keyframes: {
+        fade: {
+          "0%,50%": {
+            opacity: 1,
+          },
+          "100%": {
+            opacity: 0,
+          },
+        },
+      },
+      animation: {
+        fade: "fade 0.8s ease-in-out",
+      },
+    },
   },
   plugins: [],
 };