More comments, takes out code which was used for testing
Fixes some problems with development
Refactored training and predict function