Major changes. Made code more modular