Code cleanup

Still a lot of room for improvement...