diff --git a/AlgebraicDataflowArchitectureModel/src/algorithms/CodeGenerator.java b/AlgebraicDataflowArchitectureModel/src/algorithms/CodeGenerator.java
index 3a5da26..ff8a2b0 100644
--- a/AlgebraicDataflowArchitectureModel/src/algorithms/CodeGenerator.java
+++ b/AlgebraicDataflowArchitectureModel/src/algorithms/CodeGenerator.java
@@ -3,7 +3,9 @@
 import java.util.ArrayList;
 
 import code.ast.Block;
+import code.ast.CompilationUnit;
 import code.ast.FieldDeclaration;
+import code.ast.ImportDeclaration;
 import code.ast.MethodDeclaration;
 import code.ast.TypeDeclaration;
 import code.ast.VariableDeclaration;
@@ -30,12 +32,29 @@
 
 public class CodeGenerator {
 	public static final Type typeVoid = new Type("Void", "void");
+	private static String defaultMainTypeName = "Main";
+	static String mainTypeName = defaultMainTypeName;
 
-	static public ArrayList<TypeDeclaration> doGenerate(ResourceDependencyGraph graph, DataFlowModel model) {
-		ArrayList<TypeDeclaration> codes = new ArrayList<>();
+	public static String getMainTypeName() {
+		return mainTypeName;
+	}
+
+	public static void setMainTypeName(String mainTypeName) {
+		CodeGenerator.mainTypeName = mainTypeName;
+	}
+
+	public static void resetMainTypeName() {
+		CodeGenerator.mainTypeName = defaultMainTypeName;
+	}
+
+	static public ArrayList<CompilationUnit> doGenerate(ResourceDependencyGraph graph, DataFlowModel model) {
+		ArrayList<CompilationUnit> codes = new ArrayList<>();
 		ArrayList<ResourceNode> resources = StoreResourceCheck(graph);
-
-		codes.add(new TypeDeclaration("Main"));
+		
+		TypeDeclaration mainType = new TypeDeclaration(mainTypeName);
+		CompilationUnit mainCU = new CompilationUnit(mainType);
+		mainCU.addImport(new ImportDeclaration("java.util.*"));
+		codes.add(mainCU);
 		for (ResourceNode rn : resources) {
 			boolean f = false;
 			String name = rn.getIdentifierTemplate().getResourceName().substring(0, 1).toUpperCase()
@@ -64,7 +83,7 @@
 					if (rn.getIndegree() > 1)
 						type.addField(new FieldDeclaration(
 								((ResourceNode) re.getSource()).getIdentifierTemplate().getResourceStateType(),
-								rename));
+								((ResourceNode) re.getSource()).getIdentifierTemplate().getResourceName()));
 				}
 			}
 			if (f)
@@ -72,7 +91,7 @@
 			consstr += ")";
 			FieldDeclaration field = new FieldDeclaration(new Type(name, name),
 					rn.getIdentifierTemplate().getResourceName(), consstr);
-			codes.get(0).addField(field);
+			mainType.addField(field);
 			MethodDeclaration cons = new MethodDeclaration(name, true);
 			Block block = new Block();
 			for (Edge e : rn.getOutEdges()) {
@@ -126,7 +145,7 @@
 							io = new MethodDeclaration(
 									((Term) cm.getStateTransition().getMessageExpression()).getSymbol().getImplName(),
 									false, typeVoid, params);
-							codes.get(0).addMethod(io);
+							mainType.addMethod(io);
 						}
 					}
 				}
@@ -141,7 +160,9 @@
 			}
 			type.addMethod(new MethodDeclaration("get" + type.getTypeName(),
 					rn.getIdentifierTemplate().getResourceStateType()));
-			codes.add(type);
+			CompilationUnit cu = new CompilationUnit(type);
+			cu.addImport(new ImportDeclaration("java.util.*"));
+			codes.add(cu);
 		}
 		for (Node n : graph.getNodes()) {
 			ResourceNode rn = (ResourceNode) n;
@@ -152,7 +173,7 @@
 			get.setBody(new Block());
 			get.getBody().addStatement(
 					"return " + rn.getIdentifierTemplate().getResourceName() + "." + get.getName() + "();");
-			codes.get(0).addMethod(get);
+			mainType.addMethod(get);
 		}
 		return codes;
 	}
@@ -162,7 +183,7 @@
 		for (TypeDeclaration type : codeTree) {
 			codes.add("public class " + type.getTypeName() + "{");
 			for (FieldDeclaration field : type.getFields()) {
-				if (type.getTypeName() != "Main") {
+				if (type.getTypeName() != mainTypeName) {
 					String cons = "\t" + "private " + field.getType().getImplementationTypeName() + " "
 							+ field.getName();
 					if (field.getType().equals(DataConstraintModel.typeList))
@@ -187,7 +208,7 @@
 				}
 			}
 			codes.add("");
-			if (type.getTypeName() != "Main") {
+			if (type.getTypeName() != mainTypeName) {
 				if (!type.getConstructors().isEmpty()) {
 					String cons = "\t" + "public " + type.getTypeName() + "(";
 					for (VariableDeclaration constructor : type.getConstructors()) {