Just changed a few things.

I thought it's better this way.