added plotting, fixed headless mode