Cleaned up some code. Added some error checks