Code clean up. Fixes bugs.