AI Fixes and general code cleanup