Code improvements and cleanup