diff --git a/ast_generator/ast_generator.py b/ast_generator/ast_generator.py
index 27c8d83..3bde50a 100644
--- a/ast_generator/ast_generator.py
+++ b/ast_generator/ast_generator.py
@@ -43,7 +43,8 @@ class AstGenerator:
 
         names = get_english_words_set(['web2'], alpha=True)
         possible_names = filter(lambda x: self.settings['properties']['id-length']['max'] <= len(x) <=
-                                          self.settings['properties']['id-length']['max'] and not keyword.iskeyword(x), names)
+                                          self.settings['properties']['id-length']['max'] and not keyword.iskeyword(x),
+                                names)
 
         var_name_list = list(possible_names)
         var_name_len = len(var_name_list)
@@ -147,15 +148,15 @@ class AstGenerator:
 
         self.ast = element
 
-        # optional constants here too
-
-        self.generate_main()
+        # TODO generate constants and forward declarations
         while i < self.settings['generation-options']['max-number-of-routines']:
             if random.random() < self.settings['block-termination-probability']:
                 break
             self.generate_routine()
             i += 1
 
+        self.generate_main()
+
     def generate_main(self):
         parent = self.current_ast_element
         self.push_scope()
@@ -167,7 +168,7 @@ class AstGenerator:
         element = build_xml_element(main_args, name=GAZ_PROCEDURE_TAG)
         self.current_ast_element.append(element)
         self.current_ast_element = element
-        self.generate_block(return_stmt=True, return_value="0", return_type=GAZ_INT_KEY)
+        self.generate_block(return_stmt=True, return_value="0", return_type=GAZ_INT_KEY, block_type=GAZ_PROCEDURE_TAG)
         self.pop_scope()
         self.current_ast_element = parent
 
@@ -183,16 +184,21 @@ class AstGenerator:
         self.current_ast_element.append(element)
         self.current_ast_element = element
 
+        if block_type in [GAZ_PROCEDURE_TAG, GAZ_FUNCTION_TAG]:
+            self.generate_statements()
+        else:
+            self.generate_statements(include='declaration')
+            self.generate_statements(exclude='declaration')
+
         # Generate the loop condition increment if we are in a loop
         if block_type == GAZ_LOOP_TAG:
             self.generate_loop_condition_check(loop_var)
             self.generate_loop_condition_increment(loop_var)
 
-        self.generate_statements()
         if return_stmt:
             self.generate_return(return_type=return_type, return_value=return_value)
         if self.settings['generation-options']['generate-dead-code']:
-            self.generate_statements()
+            self.generate_statements(exclude='declaration')
         self.pop_scope()
         self.current_ast_element = parent
 
@@ -227,7 +233,8 @@ class AstGenerator:
 
         rhs = build_xml_element([], name=GAZ_RHS_TAG)
         operation.append(rhs)
-        rhs.append(self.make_literal(GAZ_INT_KEY, "'" + str(self.settings['generation-options']['max-loop-iterations']) + "'"))
+        rhs.append(
+            self.make_literal(GAZ_INT_KEY, "'" + str(self.settings['generation-options']['max-loop-iterations']) + "'"))
 
         true_block = build_xml_element([], name=GAZ_BLOCK_TAG)
         if_stmt.append(true_block)
@@ -329,7 +336,10 @@ class AstGenerator:
             self.current_ast_element.append(arg.xml)
             self.current_scope.append(arg.name, arg)
 
-    def generate_statements(self):
+    def generate_statements(self, include=None, exclude=None):
+
+        opts = ['declaration', 'routine_call', 'conditional', 'loop', 'assignment', 'out_stream', 'in_stream']
+
         # Number line
         number_line = 180  # TODO fix the numberline stuff to reflect the settings
         cutoffs = [10, 30, 50, 80, 100, 140, 180]
@@ -343,6 +353,21 @@ class AstGenerator:
             6: self.generate_in_stream,
         }
 
+        if include is not None and exclude is not None:
+            raise ValueError("Cannot specify both include and exclude")
+        elif include is not None and include in opts:
+            for i in range(len(opts)):
+                if opts[i] in include:
+                    continue
+                else:
+                    options.pop(opts.index(opts[i]))
+        elif exclude is not None and exclude in opts:
+            options.pop(opts.index(exclude))
+        elif include is None and exclude is None:
+            pass
+        else:
+            raise ValueError("Invalid include/exclude options " + str(include) + " " + str(exclude))
+
         while True:
             if random.random() < self.settings['block-termination-probability']:
                 break
@@ -353,6 +378,8 @@ class AstGenerator:
                 if cutoffs[i] < a < cutoffs[i + 1]:
                     try:
                         options[i]()
+                    except KeyError:
+                        continue
                     except ValueError:
                         break
                     break
diff --git a/gazprea_fuzzer.py b/gazprea_fuzzer.py
index 386012d..576f02f 100644
--- a/gazprea_fuzzer.py
+++ b/gazprea_fuzzer.py
@@ -55,7 +55,7 @@ class Fuzzer():
                     with redirect_stdout(y): # Workaround for fuzzer.py:49
                         try:
                             exec(f.read(), globals(), locals())
-                        except (OverflowError, ZeroDivisionError):
+                        except (OverflowError, ZeroDivisionError, ValueError):
                             os.system("rm -f fuzzer/ground_truth/{}_{}.py".format(self.file_name, i))
                             continue
             with open("fuzzer/input/{}_{}.in".format(self.file_name, i), 'w') as f: