Refactored plot() to remove unnecessary code