Cleaned up code
