Clarified the code a bit and fixed bug in last commit