small changes to main code