General fixes, code cleanup