Added TODO
Fixed better code
Removed unused imports