Code clean up..ready for testing