edited main function