Updated main.cpp - some bugs