More code tweaking