Fixed NN bug
Tweaking edge cases