Heavily restructured the code