UPDATE Cleaned up code