Bug fixes & code cleaning