Improvements and bug fixes for main and graph