Code clean-up and file re-organization