Bug fixes & code cleanup