Clean up code to remove warnings and organize imports