Fix: Fine tune and clean up the code.