Fixed issues and cleanup