Cleanup and separate classes into files