added code to Main class