Adding more error handling, fixing bugs