Bug fixes + cleanup