Clean up
Made some messy branching, then fixed it