Code cleanup and better readme