Code cleanup and organisation