Cleaned up code to fix bug