Made minor changes and fixed a bug (maybe?)