Refined the main.cpp file