Code clearn up and bug fixes