GH-18 Restructured & cleaned up code