Improve

improve and add a todo