Cleaned up code a bit and fixed #6 and fixed #8