Cleaned up code, fixed null pointer exception, added some error handling