Bug Fix and cleanup