Modified main.cpp