Further code tweaks