diff --git a/Makefile b/Makefile
index 454c55e..b5fdd9f 100644
--- a/Makefile
+++ b/Makefile
@@ -1,2 +1,5 @@
 all:
-	go run main.go tests/test.md
+	go run *.go tests/test.md
+
+clean:
+	rm -vf main
diff --git a/main.go b/main.go
index 0c8712d..e0dbcc6 100644
--- a/main.go
+++ b/main.go
@@ -31,53 +31,23 @@ func get_heading_level(heading_text string) int {
 	return level
 }
 
-func print_toc(headings []heading) {
-	tab_count := 0
-	for index, _ := range headings {
-		cur_level := get_heading_level(headings[index].text)
-
-		headings[index].level = cur_level
-
-		if index > 0 {
-			prev_level := get_heading_level(headings[index-1].text)
-
-			if cur_level > prev_level {
-				tab_count += (cur_level - prev_level)
-			} else if cur_level < prev_level {
-				tab_count -= (prev_level - cur_level)
-			}
-			for i := 0; i < tab_count; i++ {
-				fmt.Printf("\t")
-			}
-		}
-		fmt.Printf("%v\n", headings[index])
-	}
-}
-
-func count_levels(headings []heading) map[string]int {
-	level_count := map[string]int{
-		"header1": 0,
-		"header2": 0,
-		"header3": 0,
-		"header4": 0,
-		"header5": 0,
-		"header6": 0,
-	}
+func count_levels(headings []heading) map[int]int {
+	level_count := make(map[int]int)
 
 	for index, _ := range headings {
 		switch headings[index].level {
 		case 1:
-			level_count["header1"]++
+			level_count[1]++
 		case 2:
-			level_count["header2"]++
+			level_count[2]++
 		case 3:
-			level_count["header3"]++
+			level_count[3]++
 		case 4:
-			level_count["header4"]++
+			level_count[4]++
 		case 5:
-			level_count["header5"]++
+			level_count[5]++
 		case 6:
-			level_count["header6"]++
+			level_count[6]++
 		}
 	}
 
@@ -99,13 +69,16 @@ func main() {
 
 	for index, value := range file_lines {
 		if is_heading(value) {
-			headings = append(headings, heading{value, 0, (index + 1)})
+			headings = append(
+				headings, heading{value, get_heading_level(value), (index + 1)})
 		}
 	}
 
-	print_toc(headings)
+	fmt.Printf("%v\n", headings)
 
-	header_levels := count_levels(headings)
+	header_level_count := count_levels(headings)
 
-	fmt.Printf("%v\n", header_levels)
+	for key, value := range header_level_count {
+		fmt.Printf("Header %d occurs %d times.\n", key, value)
+	}
 }