re-structured the code