Code updates and UI improvements