Fix bugs in main code file