Cleanup and simplify code