Major logic changes to try and clean up ugly code