diff --git a/source/poet.cpp b/source/poet.cpp
index fa33e6c..de83be3 100644
--- a/source/poet.cpp
+++ b/source/poet.cpp
@@ -34,35 +34,59 @@ class RbmListener : public Rbm::IListener
 	}
 };
 
-
-#define CREATE_TRAINING		0
-#define DO_TRAINING			1
-#define DO_FORWARD			1
-
-int main()
+int main(int argc, char *argv[])
 {
-#if CREATE_TRAINING	
-	arma::mat batch = RnnTextHelper::createTraining("moby_ch1.txt", RnnTextHelper::SEQ_LENGTH);
-	batch.save("poet2.training.dat", arma::arma_ascii);
-	return 0;
-#endif
+	enum Command {Nop, Create, Reset, Train, Forward};
+	Command command = Nop;
+	
+	if (argc > 1)
+	{
+		char cmd_token = *argv[1];
+		if (cmd_token == 'c')
+		{
+			command = Command::Create;
+		}
+		else if (cmd_token == 'r')
+		{
+			command = Command::Reset;
+		}
+		else if (cmd_token == 't')
+		{
+			command = Command::Train;
+		}
+		else if (cmd_token == 'f')
+		{
+			command = Command::Forward;
+		}
+	}
+	
+	if (command == Command::Create)
+	{
+		arma::mat batch = RnnTextHelper::createTraining("moby_ch1.txt", RnnTextHelper::SEQ_LENGTH);
+		batch.save("poet2.training.dat", arma::arma_ascii);
+		return 0;
+	}
 	
 	// Load project
 	RnnStack *stack = reinterpret_cast<RnnStack*>(StackCreator::fromFile(".", "poet_2v_5s"));
 	
 	// Load weights
-	bool weight_loaded = stack->loadWeights(".");
+	stack->loadWeights(".");
 	
 	// Load training
 	stack->loadTrainingBatch(".");
 	arma::mat t_vc = stack->trainingBatch();
 
-	bool do_training = not weight_loaded;
-#if DO_TRAINING
-	do_training = true;
-#endif
-			
-	if (do_training)
+	if (command == Command::Reset)
+	{
+		for (int i=0; i < stack->numLayers(); i++)
+		{
+			stack->getLayer(i)->weightsInit(0.01);
+		}
+		stack->saveWeights(".");
+	}
+	
+	if (command == Command::Train)
 	{
 		RbmListener listener;
 
@@ -90,36 +114,36 @@ int main()
 		printf("\n");
 	}
 	
-#if DO_FORWARD
-	arma::mat state;
-	arma::mat curr;
-	arma::mat next;
-	std::string curr_str;
-	std::string next_str;
-	for (int i=0; i < stack->getSeqLen(); i++)
+	if (command == Command::Forward)
 	{
-		curr = stack->to_curr(t_vc.row(i));
-		curr_str.append(1,RnnTextHelper::idx2ch(curr.index_max()));
-		arma::mat r = stack->step_forward(state, curr);
-		next = stack->to_next(r, true);
-		next_str.append(1,RnnTextHelper::idx2ch(next.index_max()));
-	}
-	cout << "Curr: " << curr_str << std::endl;
-	cout << "Next: " << next_str << std::endl;
-	
-	for (int i=stack->getSeqLen(); i < t_vc.n_rows; i++)
-	{
-		curr = next;
-		curr_str.append(1,RnnTextHelper::idx2ch(curr.index_max()));
-		arma::mat r = stack->step_forward(state, curr);
-		next = stack->to_next(r, true);
-		curr = stack->to_curr(r);
-		next_str.append(1,RnnTextHelper::idx2ch(next.index_max()));
-	}
-	cout << "Curr: " << curr_str << std::endl;
-	cout << "Next: " << next_str << std::endl;
-#endif
+		arma::mat state;
+		arma::mat curr;
+		arma::mat next;
+		std::string curr_str;
+		std::string next_str;
+		for (int i=0; i < stack->getSeqLen(); i++)
+		{
+			curr = stack->to_curr(t_vc.row(i));
+			curr_str.append(1,RnnTextHelper::idx2ch(curr.index_max()));
+			arma::mat r = stack->step_forward(state, curr);
+			next = stack->to_next(r, true);
+			next_str.append(1,RnnTextHelper::idx2ch(next.index_max()));
+		}
+		cout << "Curr: " << curr_str << std::endl;
+		cout << "Next: " << next_str << std::endl;
 
+		for (int i=stack->getSeqLen(); i < t_vc.n_rows; i++)
+		{
+			curr = next;
+			curr_str.append(1,RnnTextHelper::idx2ch(curr.index_max()));
+			arma::mat r = stack->step_forward(state, curr);
+			next = stack->to_next(r, true);
+			curr = stack->to_curr(r);
+			next_str.append(1,RnnTextHelper::idx2ch(next.index_max()));
+		}
+		cout << "Curr: " << curr_str << std::endl;
+		cout << "Next: " << next_str << std::endl;
+	}
 	printf("\nEnd of program\n");
 	return 0;
 }