Code clean up and orgaization