Fix indentation. Fix minor bug. Refactor a bit for readability and intuitiveness.